def get_spec_with_expansion(filepath, override_vars=None):
"""
Return a MerlinSpec with overrides and expansion, without
creating a MerlinStudy.
"""
expanded_spec_text = expand_spec_no_study(filepath, override_vars)
return MerlinSpec.load_spec_from_string(expanded_spec_text)
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.merlin_spec_filepath = os.path.join(self.tmpdir, "basic_ensemble.yaml")
with open(self.merlin_spec_filepath, "w+") as _file:
_file.write(MERLIN_SPEC)
self.spec = MerlinSpec.load_specification(self.merlin_spec_filepath)
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
self.merlin_spec_filepath = os.path.join(self.tmpdir, "no_merlin.yaml")
with open(self.merlin_spec_filepath, "w+") as _file:
_file.write(NO_MERLIN)
self.spec = MerlinSpec.load_specification(self.merlin_spec_filepath)
def get_expanded_spec(self):
"""
Get a new yaml spec file with defaults, cli overrides, and variable expansions.
Useful for provenance.
"""
# get specification including defaults and cli-overridden user variables
full_spec_text = dump_with_overrides(self.original_spec, self.override_vars)
new_spec = MerlinSpec.load_spec_from_string(full_spec_text)
# expand user variables
new_spec_text = expand_by_line(
new_spec.dump(), MerlinStudy.get_user_vars(new_spec)
)
# expand reserved words
new_spec_text = expand_by_line(new_spec_text, self.special_vars)
return MerlinSpec.load_spec_from_string(new_spec_text)
def expand_spec_no_study(filepath, override_vars=None):
"""
Get the expanded text of a spec without creating
a MerlinStudy. Expansion is limited to user variables
(the ones defined inside the yaml spec or at the command
line).
"""
error_override_vars(override_vars, filepath)
spec = MerlinSpec.load_specification(filepath)
full_spec = dump_with_overrides(spec, override_vars)
spec = MerlinSpec.load_spec_from_string(full_spec)
uvars = []
if "variables" in spec.environment:
uvars.append(spec.environment["variables"])
if "labels" in spec.environment:
uvars.append(spec.environment["labels"])
evaluated_uvars = determine_user_variables(*uvars)
return expand_by_line(full_spec.split("\n"), evaluated_uvars)
def expanded_spec(self):
"""
Determines, writes to yaml, and loads into memory an expanded
specification.
"""
# Write expanded yaml spec
self.expanded_filepath = os.path.join(
self.info, self.spec.name.replace(" ", "_") + ".yaml"
)
# If we are restarting, we don't need to re-expand, just need to read
# in the previously expanded spec
if self.restart_dir is None:
self.write_expanded_spec(self.expanded_filepath)
return MerlinSpec.load_specification(self.expanded_filepath)
def __init__(
self,
filepath,
override_vars=None,
restart_dir=None,
samples_file=None,
dry_run=False,
no_errors=False,
pgen_file=None,
pargs=None,
):
self.original_spec = MerlinSpec.load_specification(filepath)
self.override_vars = override_vars
error_override_vars(self.override_vars, self.original_spec.path)
self.samples_file = samples_file
self.label_clash_error()
self.dry_run = dry_run
self.no_errors = no_errors
# If we load from a file, record that in the object for provenance
# downstream
if self.samples_file is not None:
self.original_spec.merlin["samples"]["file"] = self.samples_file
self.original_spec.merlin["samples"]["generate"]["cmd"] = ""
self.restart_dir = restart_dir
self.special_vars = {
"SPECROOT": self.original_spec.specroot,
"MERLIN_TIMESTAMP": self.timestamp,
"MERLIN_INFO": self.info,
"MERLIN_WORKSPACE": self.workspace,
"OUTPUT_PATH": self.output_path,
"MERLIN_SUCCESS": str(int(ReturnCode.OK)),
"MERLIN_RESTART": str(int(ReturnCode.RESTART)),
"MERLIN_SOFT_FAIL": str(int(ReturnCode.SOFT_FAIL)),
"MERLIN_HARD_FAIL": str(int(ReturnCode.HARD_FAIL)),
"MERLIN_RETRY": str(int(ReturnCode.RETRY)),
}
self.pgen_file = pgen_file
self.pargs = pargs
self.dag = None
self.load_dag()
def stop_workers(args):
"""
CLI command for stopping all workers.
:param `args`: parsed CLI arguments
"""
print(banner_small)
worker_names = []
if args.spec:
spec_path = verify_filepath(args.spec)
spec = MerlinSpec.load_specification(spec_path)
worker_names = spec.get_worker_names()
for worker_name in worker_names:
if "$" in worker_name:
LOG.warning(
f"Worker '{worker_name}' is unexpanded. Target provenance spec instead?"
)
router.stop_workers(args.task_server, worker_names, args.queues, args.workers)
def write_expanded_spec(self, dest):
"""
Write a new yaml spec file with defaults and variable expansions.
Useful for provenance.
:param `dest`: destination for fully expanded yaml file
"""
# specification text including defaults and overridden user variables
full_spec = dump_with_overrides(self.spec, self.override_vars)
with open(dest, "w") as dumped_file:
dumped_file.write(full_spec)
# update spec so that user_vars update will be accurate
self.spec = MerlinSpec.load_specification(dest)
# expand user variables
self.write_expand_by_line(dest, self.user_vars)
# expand reserved words
self.write_expand_by_line(dest, self.special_vars)
def get_expanded_spec(self):
"""
Get a new yaml spec file with defaults, cli overrides, and variable expansions.
Useful for provenance.
"""
# get specification including defaults and cli-overridden user variables
new_env = replace_override_vars(self.original_spec.environment,
self.override_vars)
new_spec = deepcopy(self.original_spec)
new_spec.environment = new_env
# expand user variables
new_spec_text = expand_by_line(new_spec.dump(),
MerlinStudy.get_user_vars(new_spec))
# expand reserved words
new_spec_text = expand_by_line(new_spec_text, self.special_vars)
result = MerlinSpec.load_spec_from_string(new_spec_text)
return expand_env_vars(result)
def get_adapter_config(self, override_type=None):
spec = MerlinSpec.load_specification(self.spec.path)
adapter_config = dict(spec.batch)
if "type" not in adapter_config.keys():
adapter_config["type"] = "local"
# The type may be overriden, preserve the batch type
adapter_config["batch_type"] = adapter_config["type"]
if override_type is not None:
adapter_config["type"] = override_type
# if a dry run was ordered by the yaml spec OR the cli flag, do a dry run.
adapter_config["dry_run"] = self.dry_run or adapter_config["dry_run"]
# Add the version if using flux to switch the command in the step
if adapter_config["batch_type"] == "flux":
adapter_config["flux_command"] = self.flux_command
LOG.debug(f"Adapter config = {adapter_config}")
return adapter_config
def expanded_spec(self):
"""
Determines, writes to yaml, and loads into memory an expanded
specification.
"""
# If we are restarting, we don't need to re-expand, just need to read
# in the previously expanded spec
if self.restart_dir is not None:
return self.get_expanded_spec()
result = self.get_expanded_spec()
expanded_name = result.description["name"].replace(
" ", "_") + ".expanded.yaml"
# Set expanded filepath
expanded_filepath = os.path.join(self.info, expanded_name)
# expand provenance spec filename
if contains_token(self.original_spec.name) or contains_shell_ref(
self.original_spec.name):
name = f"{result.description['name'].replace(' ', '_')}_{self.timestamp}"
name = expand_line(name, {}, env_vars=True)
if "/" in name:
raise ValueError(
f"Expanded value '{name}' for field 'name' in section 'description' is not a valid filename."
)
expanded_workspace = os.path.join(self.output_path, name)
if result.merlin["samples"]:
sample_file = result.merlin["samples"]["file"]
if sample_file.startswith(self.workspace):
new_samples_file = sample_file.replace(
self.workspace, expanded_workspace)
result.merlin["samples"]["generate"][
"cmd"] = result.merlin["samples"]["generate"][
"cmd"].replace(self.workspace, expanded_workspace)
result.merlin["samples"]["file"] = new_samples_file
shutil.move(self.workspace, expanded_workspace)
self.workspace = expanded_workspace
self.info = os.path.join(self.workspace, "merlin_info")
self.special_vars["MERLIN_INFO"] = self.info
expanded_filepath = os.path.join(self.info, expanded_name)
new_spec_text = expand_by_line(result.dump(),
MerlinStudy.get_user_vars(result))
result = MerlinSpec.load_spec_from_string(new_spec_text)
result = expand_env_vars(result)
# pgen
if self.pgen_file:
env = result.get_study_environment()
result.globals = self.load_pgen(self.pgen_file, self.pargs, env)
# copy the --samplesfile (if any) into merlin_info
if self.samples_file:
shutil.copyfile(
self.samples_file,
os.path.join(self.info, os.path.basename(self.samples_file)),
)
# write expanded spec for provenance
with open(expanded_filepath, "w") as f:
f.write(result.dump())
# write original spec for provenance
result = MerlinSpec.load_spec_from_string(result.dump())
result.path = expanded_filepath
name = result.description["name"].replace(" ", "_")
self.write_original_spec(name)
# write partially-expanded spec for provenance
partial_spec = deepcopy(self.original_spec)
if "variables" in result.environment:
partial_spec.environment["variables"] = result.environment[
"variables"]
if "labels" in result.environment:
partial_spec.environment["labels"] = result.environment["labels"]
partial_spec_path = os.path.join(self.info, name + ".partial.yaml")
with open(partial_spec_path, "w") as f:
f.write(partial_spec.dump())
LOG.info(f"Study workspace is '{self.workspace}'.")
return result