diff --git a/hpcflow/sdk/core/command_files.py b/hpcflow/sdk/core/command_files.py index 5b1572afb..2c265c7bd 100644 --- a/hpcflow/sdk/core/command_files.py +++ b/hpcflow/sdk/core/command_files.py @@ -453,6 +453,13 @@ def to_dict(self): dct = super().to_dict() return dct + def __eq__(self, other: object) -> bool: + if not isinstance(other, self.__class__): + return False + if self.to_dict() == other.to_dict(): + return True + return False + def _get_members(self, ensure_contents=False, use_file_label=False): out = super()._get_members(ensure_contents) if use_file_label: diff --git a/hpcflow/sdk/core/utils.py b/hpcflow/sdk/core/utils.py index f2db24728..73b32ce8c 100644 --- a/hpcflow/sdk/core/utils.py +++ b/hpcflow/sdk/core/utils.py @@ -13,6 +13,7 @@ from typing import Mapping from ruamel.yaml import YAML +from ruamel.yaml.comments import CommentedSeq import sentry_sdk from hpcflow.sdk.core.errors import FromSpecMissingObjectError, InvalidIdentifier @@ -311,6 +312,12 @@ def read_YAML_file(path: PathLike): return read_YAML(Path(path)) +def FSlist(l): + cs = CommentedSeq(l) + cs.fa.set_flow_style() + return cs + + def read_JSON_string(string: str): return json.loads(string) diff --git a/hpcflow/sdk/core/workflow.py b/hpcflow/sdk/core/workflow.py index 10302d097..2e90898bd 100644 --- a/hpcflow/sdk/core/workflow.py +++ b/hpcflow/sdk/core/workflow.py @@ -25,8 +25,12 @@ read_JSON_string, read_YAML, read_YAML_file, + FSlist, replace_items, ) +from ruamel.yaml import YAML +from io import StringIO +from json import dumps from hpcflow.sdk.core.errors import ( InvalidInputSourceTaskReference, LoopAlreadyExistsError, @@ -181,6 +185,135 @@ def from_YAML_file(cls, path: PathLike) -> app.WorkflowTemplate: """ return cls._from_data(read_YAML_file(path)) + @classmethod + def to_input_dict(cls, data): + """Called by to_yaml and to_json. + Returns a python dict with only the necessary elements to save the template.""" + d_workflow = {"name": data.name} + + # Task_Schemas + + # Tasks + l_tasks = [] + for task in data.tasks: + d_task = {} + # Schemas + l_schemas = [] + for schema in task.schemas: + l_schemas.append(schema.name) + d_task["schemas"] = l_schemas + + # Element sets + l_element_sets = [] + for element_set in task.element_sets: + d_element_set = {} + + # Inputs + d_inputs = {} + for input in element_set.inputs: + if isinstance(input.value, list): + d_inputs[input.parameter.typ] = input.value + else: + d_inputs[input.parameter.typ] = input.value + if d_inputs: + d_element_set["inputs"] = d_inputs + + # Sequences + l_sequences = [] + for sequence in element_set.sequences: + l_sequences.extend( + [ + { + "path": sequence.path, + "values": sequence.values, + "nesting_order": sequence.nesting_order, + } + ] + ) + if l_sequences: + d_element_set["sequences"] = l_sequences + + # Resources + l_resources = {} + for resource in element_set.resources: + if resource.num_cores: + l_resources[resource.normalised_resources_path] = { + "num_cores": resource.num_cores + } + if l_resources: + d_element_set["resources"] = l_resources + + # Input files + l_input_files = [] + for input_file in element_set.input_files: + l_input_files.extend( + [ + { + "file": input_file.file.label, + "path": input_file._path, + } + ] + ) + if l_input_files: + d_element_set["input_files"] = l_input_files + + l_element_sets.append(d_element_set) + d_task["element_sets"] = l_element_sets + + l_tasks.append(d_task) + if l_tasks: + d_workflow["tasks"] = l_tasks + + return d_workflow + + @classmethod + def to_yaml(cls, dumper, data): + """Called by to_yaml_string.""" + d_workflow = cls.to_input_dict(data) + + for i, task in enumerate(d_workflow["tasks"]): + d_workflow["tasks"][i]["schemas"] = FSlist(task["schemas"]) + + for j, element_set in enumerate(task["element_sets"]): + if "inputs" in element_set: + for name, values in element_set["inputs"].items(): + if isinstance(values, list): + d_workflow["tasks"][i]["element_sets"][j]["inputs"][ + name + ] = FSlist(values) + + if "sequences" in element_set: + for k, sequence in enumerate(element_set["sequences"]): + d_workflow["tasks"][i]["element_sets"][j]["sequences"][k][ + "values" + ] = FSlist(sequence["values"]) + + return dumper.represent_mapping("tag:yaml.org,2002:map", d_workflow) + + def to_yaml_string(self): + """Get templeate as a formatted YAML string.""" + wt_yaml = YAML() + yaml_string = StringIO() + wt_yaml.register_class(type(self)) + wt_yaml.dump(self, yaml_string) + return yaml_string.getvalue() + + def to_yaml_file(self, yaml_file_path): + """Save to a YAML file. + + Parameters + ---------- + string + The path to the yaml file in which the template will be saved. + + """ + if not any(x in yaml_file_path for x in [".yml", ".yaml"]): + yaml_file_path = yaml_file_path + ".yml" + wt_yaml = YAML() + wt_yaml.register_class(type(self)) + with open(yaml_file_path, "w") as output_file: + wt_yaml.dump(self, output_file) + @classmethod def from_JSON_string(cls, string: str) -> app.WorkflowTemplate: """Load from a JSON string. @@ -205,6 +338,26 @@ def from_JSON_file(cls, path: PathLike) -> app.WorkflowTemplate: """ return cls._from_data(read_JSON_file(path)) + def to_json_string(self): + """Get templeate as a formatted JSON string.""" + d_workflow = self.to_input_dict(self) + json_string = dumps(d_workflow, indent=4) + return json_string + + def to_json_file(self, json_file_path): + """Save to a JSON file. + + Parameters + ---------- + string + The path to the json file in which the template will be saved. + + """ + if not ".json" in json_file_path: + json_file_path = json_file_path + ".json" + with open(json_file_path, "w") as output_file: + output_file.write(self.to_json_string()) + @classmethod def from_file( cls, diff --git a/hpcflow/tests/unit/test_workflow.py b/hpcflow/tests/unit/test_workflow.py index de75eef14..0ec5e0f7b 100644 --- a/hpcflow/tests/unit/test_workflow.py +++ b/hpcflow/tests/unit/test_workflow.py @@ -1,4 +1,5 @@ import copy +import os from textwrap import dedent import pytest @@ -9,6 +10,7 @@ WorkflowNotFoundError, ) from hpcflow.sdk.core.test_utils import make_workflow +from ruamel.yaml import YAML def modify_workflow_metadata_on_disk(workflow): @@ -117,6 +119,69 @@ def workflow_w1(null_config, tmp_path, schema_s3, param_p1): return hf.Workflow.from_template(wkt, path=tmp_path) +@pytest.fixture +def wkt_yml(): + return dedent( + """ + name: simple_workflow + tasks: + - schemas: [dummy_task_1] + element_sets: + - inputs: + p2: 201 + p5: 501 + sequences: + - path: inputs.p1 + values: [101, 102] + nesting_order: 0 + resources: + any: + num_cores: 8 + """ + ) + + +@pytest.fixture +def wkt_json(): + return dedent( + """ + { + "name": "simple_workflow", + "tasks": [ + { + "schemas": [ + "dummy_task_1" + ], + "element_sets": [ + { + "inputs": { + "p2": 201, + "p5": 501 + }, + "sequences": [ + { + "path": "inputs.p1", + "values": [ + 101, + 102 + ], + "nesting_order": 0 + } + ], + "resources": { + "any": { + "num_cores": 8 + } + } + } + ] + } + ] + } + """ + ) + + def test_make_empty_workflow(empty_workflow): assert empty_workflow.path is not None @@ -231,6 +296,146 @@ def test_WorkflowTemplate_from_YAML_string_with_and_without_element_sets_equival assert wkt_1 == wkt_2 +def test_WorkflowTemplate_to_YAML_string_format(null_config, wkt_yml): + wkt = hf.WorkflowTemplate.from_YAML_string(wkt_yml) + yaml_string = wkt.to_yaml_string() + assert wkt_yml == "\n" + yaml_string + + +def test_WorkflowTemplate_to_YAML_string_functional(null_config, wkt_yml): + wkt = hf.WorkflowTemplate.from_YAML_string(wkt_yml) + yaml_string = wkt.to_yaml_string() + wkt_2 = hf.WorkflowTemplate.from_YAML_string(yaml_string) + assert wkt == wkt_2 + + +def test_WorkflowTemplate_to_YAML_file_functional(null_config, wkt_yml): + wkt = hf.WorkflowTemplate.from_YAML_string(wkt_yml) + yaml_file_path = "to_yaml_test.yml" + wkt.to_yaml_file(yaml_file_path) + saved_yaml = "" + with open(yaml_file_path, "r") as output_file: + saved_yaml = output_file.read() + os.remove(yaml_file_path) + wkt_3 = hf.WorkflowTemplate.from_YAML_string(saved_yaml) + assert wkt == wkt_3 + + +def test_WorkflowTemplate_to_YAML_complex(null_config): + wkt_yml_name = dedent( + """ + name: test_wk + """ + ) + wkt_yml_command_files = dedent( + """ + command_files: + - label: file1 + name: + name: file1.txt + - label: file2 + name: + name: file2.txt + - label: file3 + name: + name: file3.txt + """ + ) + wkt_yml_task_schemas = dedent( + """ + task_schemas: + - objective: t1 + inputs: + - parameter: p1 + - parameter: p2 + outputs: + - parameter: p3 + actions: + - environments: + - scope: + type: any + environment: null_env + commands: + - command: doSomething < <> <> --out <> + input_file_generators: + file1: + from_inputs: [p1, p2] + output_file_parsers: + p3: + from_files: [file2] + - objective: t2 + inputs: + - parameter: p2 + - parameter: p3 + - parameter: p4 + outputs: + - parameter: p4 + actions: + - environments: + - scope: + type: any + environment: null_env + commands: + - command: doSomething2 <> <> <> --out <> + output_file_parsers: + p4: + from_files: [file3] + """ + ) + wkt_yml_tasks = dedent( + """ + tasks: + - schemas: [t1] + element_sets: + - inputs: + p1: 101 + input_files: + - file: file1 + path: file1.txt + - inputs: + p2: 201 + sequences: + - path: inputs.p1 + values: [101, 102] + nesting_order: 0 + - path: inputs.p2.b + values: [201] + nesting_order: 1 + resources: + any: + num_cores: 8 + processing: + num_cores: 1 + input_file_generator[file=file1]: + num_cores: 2 + - schemas: [t2] + inputs: + p4: [1, 2, 3] + """ + ) + + wkt_yml = wkt_yml_name + wkt_yml_command_files + wkt_yml_task_schemas + wkt_yml_tasks + wkt_1 = hf.WorkflowTemplate.from_YAML_string(wkt_yml) + + yaml_file_path = "to_yaml_test.yml" + wkt_1.to_yaml_file(yaml_file_path) + + saved_yaml = "" + with open(yaml_file_path, "r") as output_file: + saved_yaml = output_file.read() + # Command_files and task_schemas currently not suported - inserting manually + saved_yaml = ( + wkt_yml_name + + wkt_yml_command_files + + wkt_yml_task_schemas + + saved_yaml.replace(wkt_yml_name.strip("\n"), "") + ) + os.remove(yaml_file_path) + wkt_2 = hf.WorkflowTemplate.from_YAML_string(saved_yaml) + + assert wkt_1 == wkt_2 + + def test_store_has_pending_during_add_task(workflow_w1, schema_s2, param_p3): t2 = hf.Task(schemas=schema_s2, inputs=[hf.InputValue(param_p3, 301)]) with workflow_w1.batch_update(): @@ -291,3 +496,28 @@ def test_WorkflowTemplate_from_JSON_string_without_element_sets(null_config): """ ) hf.WorkflowTemplate.from_JSON_string(wkt_json) + + +def test_WorkflowTemplate_to_JSON_string_format(null_config, wkt_json): + wkt = hf.WorkflowTemplate.from_JSON_string(wkt_json) + json_string = wkt.to_json_string() + assert wkt_json == "\n" + json_string + "\n" + + +def test_WorkflowTemplate_to_JSON_string_functional(null_config, wkt_json): + wkt = hf.WorkflowTemplate.from_JSON_string(wkt_json) + json_string = wkt.to_json_string() + wkt_2 = hf.WorkflowTemplate.from_JSON_string(json_string) + assert wkt == wkt_2 + + +def test_WorkflowTemplate_to_JSON_file_functional(null_config, wkt_json): + wkt = hf.WorkflowTemplate.from_JSON_string(wkt_json) + json_file_path = "to_json_test.json" + wkt.to_json_file(json_file_path) + saved_json = "" + with open(json_file_path, "r") as output_file: + saved_json = output_file.read() + os.remove(json_file_path) + wkt_3 = hf.WorkflowTemplate.from_YAML_string(saved_json) + assert wkt == wkt_3