def _init(n: int, outdir: Path, j: int, env_name: str, session: Executor):
outdir.mkdir(exist_ok=True, parents=True)
for i in range(1, j + 1):
session.submit(
_init_worker,
env_name=env_name,
n=int(ceil(n / j)),
outfile=outdir / f".op:1:startup-shard-{i:02d}.txt",
)
def _reset(
benchmarks: List[str],
n: int,
outdir: Path,
env_name: str,
j: int,
session: Executor,
):
outdir.mkdir(exist_ok=True, parents=True)
num_measurements_per_benchmark = int(ceil(n / len(benchmarks)))
for i, benchmarks_chunk in enumerate(chunkify(benchmarks, j), start=1):
session.submit(
_reset_worker,
num_measurements_per_benchmark=num_measurements_per_benchmark,
benchmarks=benchmarks_chunk,
env_name=env_name,
outfile=outdir / f".op:2:reset-shard-{i:02d}.txt",
)
def _step(
n: int,
benchmarks: List[str],
env_name: str,
seed: int,
j: int,
outdir: Path,
session: Executor,
):
outdir.mkdir(exist_ok=True, parents=True)
num_measurements_per_benchmark = int(ceil(n / len(benchmarks)))
for i, benchmarks_chunk in enumerate(chunkify(benchmarks, j), start=1):
session.submit(
_step_worker,
num_measurements_per_benchmark=num_measurements_per_benchmark,
seed=seed + (i * len(benchmarks_chunk)),
benchmarks=benchmarks_chunk,
env_name=env_name,
step_outfile=outdir / f".op:3:step-shard-{i:02d}.txt",
batched_outfile=outdir / f".op:3:step-batched-shard-{i:02d}.txt",
)
def _observations(
observation_spaces: List[str],
benchmarks: List[str],
n: int,
j: int,
session: Executor,
outdir: Path,
env_name: str,
):
outdir.mkdir(exist_ok=True, parents=True)
num_measurements_per_benchmark = int(ceil(n / len(benchmarks)))
for i, benchmarks_chunk in enumerate(chunkify(benchmarks, j), start=1):
for observation_space in observation_spaces:
session.submit(
_observations_worker,
observation_space=observation_space,
num_measurements_per_benchmark=num_measurements_per_benchmark,
benchmarks=benchmarks_chunk,
env_name=env_name,
outfile=outdir / f".observation:{observation_space}-shard-{i:02d}.txt",
)