def test_kv_to_flag(self):
self.assertEqual(pu.key_value_to_flag("x", "a"), "-x a")
self.assertEqual(pu.key_value_to_flag("y", True), "-y")
self.assertEqual(pu.key_value_to_flag("y", False), "")
self.assertEqual(pu.key_value_to_flag("yy", True), "--yy")
self.assertEqual(pu.key_value_to_flag("zz", "c"), "--zz c")
eq = True
self.assertEqual(pu.key_value_to_flag("x", "a", eq), "-x=a")
self.assertEqual(pu.key_value_to_flag("y", True, eq), "-y=true")
self.assertEqual(pu.key_value_to_flag("y", False, eq), "-y=false")
self.assertEqual(pu.key_value_to_flag("yy", True, eq), "--yy=true")
self.assertEqual(pu.key_value_to_flag("zz", "c", eq), "--zz=c")
def test_kv_to_flag(self):
self.assertEqual(pu.key_value_to_flag('x', 'a'), '-x a')
self.assertEqual(pu.key_value_to_flag('y', True), '-y')
self.assertEqual(pu.key_value_to_flag('y', False), '')
self.assertEqual(pu.key_value_to_flag('yy', True), '--yy')
self.assertEqual(pu.key_value_to_flag('zz', 'c'), '--zz c')
eq = True
self.assertEqual(pu.key_value_to_flag('x', 'a', eq), '-x=a')
self.assertEqual(pu.key_value_to_flag('y', True, eq), '-y=true')
self.assertEqual(pu.key_value_to_flag('y', False, eq), '-y=false')
self.assertEqual(pu.key_value_to_flag('yy', True, eq), '--yy=true')
self.assertEqual(pu.key_value_to_flag('zz', 'c', eq), '--zz=c')
def _create_cmd(self, step, img, cid):
container_args = self._get_container_kwargs(step, img, cid)
container_args.pop("detach")
cmd = ["docker create"]
cmd.append(f"--name {container_args.pop('name')}")
cmd.append(f"--workdir {container_args.pop('working_dir')}")
entrypoint = container_args.pop("entrypoint", None)
if entrypoint:
cmd.append(f"--entrypoint {' '.join(entrypoint)}")
# append volume and environment flags
for vol in container_args.pop("volumes"):
cmd.append(f"-v {vol}")
for env_key, env_val in container_args.pop("environment").items():
cmd.append(f"-e {env_key}={env_val}")
command = container_args.pop("command")
image = container_args.pop("image")
# anything else is treated as a flag
for k, v in container_args.items():
cmd.append(pu.key_value_to_flag(k, v))
# append the image and the commands
cmd.append(image)
if command:
cmd.append(" ".join(command))
return " ".join(cmd)
def _create_cmd(self, step, img, cid):
container_args = self._get_container_kwargs(step, img, cid)
container_args.pop('detach')
cmd = ['docker create']
cmd.append(f"--name {container_args.pop('name')}")
cmd.append(f"--workdir {container_args.pop('working_dir')}")
entrypoint = container_args.pop('entrypoint', None)
if entrypoint:
cmd.append(f"--entrypoint {' '.join(entrypoint)}")
# append volume and environment flags
for vol in container_args.pop('volumes'):
cmd.append(f'-v {vol}')
for env_key, env_val in container_args.pop('environment').items():
cmd.append(f'-e {env_key}={env_val}')
command = ' '.join(container_args.pop('command', []))
image = container_args.pop('image')
# anything else is treated as a flag
for k, v in container_args.items():
cmd.append(pu.key_value_to_flag(k, v))
# append the image and the commands
cmd.append(f'{image} {command}')
return ' '.join(cmd)
def _get_resman_kwargs(self, step):
default_options = ["nodes", "nodelist", "ntasks", "ntasks-per-node"]
resman_options = []
for k, v in self._config.resman_opts.get(step.id, {}).items():
if k not in default_options:
flag = pu.key_value_to_flag(k, v)
if flag:
resman_options.extend(flag.split())
return resman_options
def _get_container_options(self):
container_args = {
'userns': True,
'pwd': '/workspace',
'bind': [f'{self._config.workspace_dir}:/workspace']
}
self._update_with_engine_config(container_args)
options = []
for k, v in container_args.items():
if isinstance(v, list):
for item in v:
options.append(pu.key_value_to_flag(k, item))
else:
options.append(pu.key_value_to_flag(k, v))
options = ' '.join(options).split(' ')
log.debug(f'container options: {options}\n')
return options
def _get_container_options(self):
container_args = {
"userns": True,
"pwd": "/workspace",
"bind": [f"{self._config.workspace_dir}:/workspace"],
}
self._update_with_engine_config(container_args)
options = []
for k, v in container_args.items():
if isinstance(v, list):
for item in v:
options.append(pu.key_value_to_flag(k, item))
else:
options.append(pu.key_value_to_flag(k, v))
options = " ".join(options).split(" ")
log.debug(f"container options: {options}\n")
return options
def _submit_batch_job(self, cmd, step):
job_name = pu.sanitized_name(step.id, self._config.wid)
temp_dir = "/tmp/popper/slurm/"
os.makedirs(temp_dir, exist_ok=True)
job_script = os.path.join(temp_dir, f"{job_name}.sh")
out_file = os.path.join(temp_dir, f"{job_name}.out")
# create/truncate log
with open(out_file, "w"):
pass
with open(job_script, "w") as f:
f.write("#!/bin/bash\n")
f.write("\n".join(cmd))
sbatch_cmd = f"sbatch --wait --job-name {job_name} --output {out_file}"
sbatch_cmd = sbatch_cmd.split()
for k, v in self._config.resman_opts.get(step.id, {}).items():
sbatch_cmd.append(pu.key_value_to_flag(k, v))
sbatch_cmd.append(job_script)
log.info(f'[{step.id}] {" ".join(sbatch_cmd)}')
if self._config.dry_run:
return 0
self._spawned_jobs.add(job_name)
# start a tail (background) process on the output file
self._start_out_stream(out_file)
# submit the job and wait
_, ecode, output = HostRunner._exec_cmd(sbatch_cmd, logging=False)
# kill the tail process
self._stop_out_stream()
self._spawned_jobs.remove(job_name)
return ecode