def test_upload_file(self):
# check file in remote
p1 = common_util.generate_short_id()
p2 = common_util.generate_short_id()
remote_file = (Path("/tmp") / p1 / p2 /
Path(self.file_a).name).as_posix()
self.run_upload_file(remote_file)
def test_makedirs():
ssh_config = load_ssh_psw_config()
with ssh_utils.sftp_client(**ssh_config) as client:
p1 = common_util.generate_short_id()
p2 = common_util.generate_short_id()
r_path = (Path("/tmp") / p1 / p2).as_posix()
print(f"made {r_path}")
assert not ssh_utils.exists(client, r_path)
ssh_utils.makedirs(client, r_path)
assert ssh_utils.exists(client, r_path)
def upload_dir(self):
with ssh_utils.sftp_client(**self.ssh_config) as client:
# check file in remote
p1 = common_util.generate_short_id()
p2 = common_util.generate_short_id()
remote_dir_path = (Path("/tmp") / p1 / p2)
remote_dir = remote_dir_path.as_posix()
ssh_utils.upload_dir(client, self.data_dir, remote_dir)
return remote_dir
def test_makedirs():
ssh_config = get_ssh_test_config(use_password=True, use_rsa_file=False)[0]
with ssh_utils.sftp_client(**ssh_config) as client:
p1 = common_util.generate_short_id()
p2 = common_util.generate_short_id()
r_path = (Path("/tmp") / p1 / p2).as_posix()
print(f"made {r_path}")
assert not ssh_utils.exists(client, r_path)
ssh_utils.makedirs(client, r_path)
assert ssh_utils.exists(client, r_path)
def test_upload():
ssh_config = load_ssh_psw_config()
# generate temp file
fd, fp = tempfile.mkstemp()
os.close(fd)
# upload
with ssh_utils.sftp_client(**ssh_config) as client:
p1 = common_util.generate_short_id()
p2 = common_util.generate_short_id()
r_path = (Path("/tmp") / p1 / p2 / Path(fp).name).as_posix()
# check file in remote
ssh_utils.upload_file(client, fp, r_path)
assert ssh_utils.exists(client, r_path)
def test_upload():
ssh_config = get_ssh_test_config(use_password=True, use_rsa_file=False)[0]
# generate temp file
fd, fp = tempfile.mkstemp()
os.close(fd)
# upload
with ssh_utils.sftp_client(**ssh_config) as client:
p1 = common_util.generate_short_id()
p2 = common_util.generate_short_id()
r_path = (Path("/tmp") / p1 / p2 / Path(fp).name).as_posix()
# check file in remote
ssh_utils.copy_from_local_to_remote(client, fp, r_path)
assert ssh_utils.exists(client, r_path)
def run_generate_job_specs(template, output):
yaml_file = template
output_path = Path(output)
# 1. validation
# 1.1. checkout output
if output_path.exists():
raise FileExistsError(output)
# load file
config_dict = load_yaml(yaml_file)
# 1.3. check values should be array
assert "params" in config_dict
params = config_dict['params']
for k, v in params.items():
if not isinstance(v, list):
raise ValueError(f"Value of param '{k}' should be list")
# 1.4. check command exists
assert "command" in config_dict
# assert "data_dir" in config_dict
# assert "working_dir" in config_dict
# 2. combine params to generate jobs
job_param_names = params.keys()
param_values = [params[_] for _ in job_param_names]
def make_job_dict(job_param_values):
job_params_dict = dict(zip(job_param_names, job_param_values))
job_dict = {
"name": common_util.generate_short_id(),
"params": job_params_dict
}
copy_item(config_dict, job_dict, 'resource')
copy_item(config_dict, job_dict, 'data_dir')
copy_item(config_dict, job_dict, 'working_dir')
return job_dict
jobs = [make_job_dict(_) for _ in itertools.product(*param_values)]
# 3. merge to bath spec
batch_spec = {
"jobs": jobs,
'name': config_dict.get('name', common_util.generate_short_id()),
"version": config_dict.get('version', current_version)
}
copy_item(config_dict, batch_spec, 'server')
copy_item(config_dict, batch_spec, 'scheduler')
copy_item(config_dict, batch_spec, 'backend')
# 4. write to file
os.makedirs(output_path.parent, exist_ok=True)
with open(output_path, 'w', newline='\n') as f:
f.write(json.dumps(batch_spec, indent=4))
return batch_spec
def make_job_dict(job_param_values):
job_params_dict = dict(zip(job_param_names, job_param_values))
job_dict = {
"name": common_util.generate_short_id(),
"params": job_params_dict,
"execution": config_dict['execution']
}
_copy_item(config_dict, job_dict, 'resource')
return job_dict
def run_batch_config(config_dict, batches_data_dir):
# add batch name
batch_name = config_dict.get('name')
if batch_name is None:
batch_name = common_util.generate_short_id()
logger.debug(f"generated batch name {batch_name}")
# add job name
jobs_dict = config_dict['jobs']
for job_dict in jobs_dict:
if job_dict.get('name') is None:
job_name = common_util.generate_short_id()
logger.debug(f"generated job name {job_name}")
job_dict['name'] = job_name
app = BatchApplication.load(config_dict, batches_data_dir)
app.start()
def make_job_dict(job_param_values):
job_params_dict = dict(zip(job_param_names, job_param_values))
job_dict = {
"name": common_util.generate_short_id(),
"params": job_params_dict
}
copy_item(config_dict, job_dict, 'resource')
copy_item(config_dict, job_dict, 'data_dir')
copy_item(config_dict, job_dict, 'working_dir')
return job_dict
def run_batch(config_dict, batches_data_dir=None):
# add batch name
if config_dict.get('name') is None:
batch_name = common_util.generate_short_id()
logger.debug(f"generated batch name {batch_name}")
config_dict['name'] = batch_name
# add job name
jobs_dict = config_dict['jobs']
for job_dict in jobs_dict:
if job_dict.get('name') is None:
job_name = common_util.generate_short_id()
logger.debug(f"generated job name {job_name}")
job_dict['name'] = job_name
batches_data_dir = get_batches_data_dir(batches_data_dir)
batches_data_dir = Path(batches_data_dir)
batch = load_batch(config_dict, batches_data_dir)
logger.info(f"batches_data_path: {batches_data_dir.absolute()}")
logger.info(f"batch name: {batch.name}")
# prepare batch data dir
if batch.data_dir_path().exists():
logger.info(f"batch {batch.name} already exists, run again")
else:
os.makedirs(batch.data_dir_path(), exist_ok=True)
# write batch config
batch_spec_file_path = batch.spec_file_path()
with open(batch_spec_file_path, 'w', newline='\n') as f:
json.dump(config_dict, f, indent=4)
# create executor manager
backend_config = config_dict.get('backend')
if backend_config is None: # set default backend
backend_config = {'type': 'local', 'conf': {}}
backend_type = backend_config['type']
if backend_type == 'remote':
remote_backend_config = backend_config['conf']
machines = [
SSHRemoteMachine(_) for _ in remote_backend_config['machines']
]
executor_manager = RemoteSSHExecutorManager(machines)
elif backend_type == 'local':
executor_manager = LocalExecutorManager()
else:
raise ValueError(f"unknown backend {backend_type}")
# set context
c = Context(executor_manager, batch)
set_context(c)
# write pid file
with open(batch.pid_file_path(), 'w', newline='\n') as f:
f.write(str(os.getpid()))
# create web app
logger.info(f"start daemon server at: {batch.daemon_conf.portal}")
create_batch_manage_webapp().listen(batch.daemon_conf.port)
# start scheduler
Scheduler(batch.daemon_conf.exit_on_finish, 5000).start()
# run io loop
ioloop.IOLoop.instance().start()
