hpcflow-new2 0.2.0a162__py3-none-any.whl → 0.2.0a164__py3-none-any.whl

hpcflow/_version.py

@@ -1 +1 @@
-__version__ = "0.2.0a162"
+__version__ = "0.2.0a164"

hpcflow/data/scripts/main_script_test_direct_in_direct_out_env_spec.py

@@ -0,0 +1,7 @@
+def main_script_test_direct_in_direct_out_env_spec(p1, env_spec):
+
+    # process
+    p2 = p1 + 100
+
+    # return outputs
+    return {"p2": env_spec}

hpcflow/sdk/app.py

@@ -729,49 +729,28 @@ class BaseApp(metaclass=Singleton):
 
     @TimeIt.decorator
     def _load_scripts(self):
-        from setuptools import find_packages
 
         # TODO: load custom directories / custom functions (via decorator)
+        scripts_package = f"{self.package_name}.{self.scripts_dir}"
 
-        app_module = import_module(self.package_name)
-        root_scripts_dir = self.scripts_dir
-
-        packages = find_packages(
-            where=str(Path(app_module.__path__[0], *root_scripts_dir.split(".")))
-        )
-        packages = [root_scripts_dir] + [root_scripts_dir + "." + i for i in packages]
-        packages = [self.package_name + "." + i for i in packages]
-        num_root_dirs = len(root_scripts_dir.split(".")) + 1
+        try:
+            ctx = resources.as_file(resources.files(scripts_package))
+        except AttributeError:
+            # < python 3.9; `resource.path` deprecated since 3.11
+            ctx = resources.path(scripts_package, "")
 
         scripts = {}
-        for pkg in packages:
-            try:
-                contents = (
-                    resource.name
-                    for resource in resources.files(pkg).iterdir()
-                    if resource.is_file()
-                )
-                _is_rsrc = lambda pkg, name: resources.files(pkg).joinpath(name).is_file()
-
-            except AttributeError:
-                # < python 3.9; `resource.contents` deprecated since 3.11
-                contents = resources.contents(pkg)
-                _is_rsrc = lambda pkg, name: resources.is_resource(pkg, name)
-
-            script_names = (
-                name for name in contents if name != "__init__.py" and _is_rsrc(pkg, name)
-            )
-
-            for i in script_names:
-                script_key = "/".join(pkg.split(".")[num_root_dirs:] + [i])
-                try:
-                    script_ctx = resources.as_file(resources.files(pkg).joinpath(i))
-                except AttributeError:
-                    # < python 3.9; `resource.path` deprecated since 3.11
-                    script_ctx = resources.path(pkg, i)
-
-                with script_ctx as script:
-                    scripts[script_key] = script
+        with ctx as path:
+            for dirpath, _, filenames in os.walk(path):
+                dirpath = Path(dirpath)
+                if dirpath.name == "__pycache__":
+                    continue
+                for filename in filenames:
+                    if filename == "__init__.py":
+                        continue
+                    val = dirpath.joinpath(filename)
+                    key = str(val.relative_to(path).as_posix())
+                    scripts[key] = Path(val)
 
         return scripts
 
@@ -1380,15 +1359,19 @@ class BaseApp(metaclass=Singleton):
             variables=variables,
             status=status,
         )
-        return wk.submit(
+        submitted_js = wk.submit(
             JS_parallelism=JS_parallelism,
             wait=wait,
             add_to_known=add_to_known,
-            return_idx=return_idx,
+            return_idx=True,
             tasks=tasks,
             cancel=cancel,
             status=status,
         )
+        if return_idx:
+            return (wk, submitted_js)
+        else:
+            return wk
 
     def _make_demo_workflow(
         self,
@@ -1553,15 +1536,19 @@ class BaseApp(metaclass=Singleton):
             store_kwargs=store_kwargs,
             variables=variables,
         )
-        return wk.submit(
+        submitted_js = wk.submit(
             JS_parallelism=JS_parallelism,
             wait=wait,
             add_to_known=add_to_known,
-            return_idx=return_idx,
+            return_idx=True,
             tasks=tasks,
             cancel=cancel,
             status=status,
         )
+        if return_idx:
+            return (wk, submitted_js)
+        else:
+            return wk
 
     def _submit_workflow(
         self,

hpcflow/sdk/cli.py

@@ -173,7 +173,7 @@ def _make_API_CLI(app):
             status=status,
         )
         if print_idx:
-            click.echo(out)
+            click.echo(out[1])
 
     @click.command(context_settings={"ignore_unknown_options": True})
     @click.argument("py_test_args", nargs=-1, type=click.UNPROCESSED)

hpcflow/sdk/core/actions.py

@@ -16,6 +16,7 @@ from watchdog.utils.dirsnapshot import DirectorySnapshotDiff
 from hpcflow.sdk import app
 from hpcflow.sdk.core import ABORT_EXIT_CODE
 from hpcflow.sdk.core.errors import (
+    ActionEnvironmentMissingNameError,
     MissingCompatibleActionEnvironment,
     OutputFileParserNoOutputError,
     UnknownScriptDataKey,
@@ -140,6 +141,7 @@ class ElementActionRun:
         element_action,
         index: int,
         data_idx: Dict,
+        commands_idx: List[int],
         start_time: Union[datetime, None],
         end_time: Union[datetime, None],
         snapshot_start: Union[Dict, None],
@@ -156,6 +158,7 @@ class ElementActionRun:
         self._element_action = element_action
         self._index = index  # local index of this run with the action
         self._data_idx = data_idx
+        self._commands_idx = commands_idx
         self._start_time = start_time
         self._end_time = end_time
         self._submission_idx = submission_idx
@@ -221,6 +224,10 @@ class ElementActionRun:
     def data_idx(self):
         return self._data_idx
 
+    @property
+    def commands_idx(self):
+        return self._commands_idx
+
     @property
     def metadata(self):
         return self._metadata
@@ -458,14 +465,18 @@ class ElementActionRun:
             self._output_files = self.app.ElementOutputFiles(element_action_run=self)
         return self._output_files
 
+    @property
+    def env_spec(self) -> Dict[str, Any]:
+        return self.resources.environments[self.action.get_environment_name()]
+
     @TimeIt.decorator
     def get_resources(self):
         """Resolve specific resources for this EAR, considering all applicable scopes and
         template-level resources."""
         return self.element_iteration.get_resources(self.action)
 
-    def get_environment_label(self) -> str:
-        return self.action.get_environment_label()
+    def get_environment_spec(self) -> str:
+        return self.action.get_environment_spec()
 
     def get_environment(self) -> app.Environment:
         return self.action.get_environment()
@@ -542,6 +553,9 @@ class ElementActionRun:
             else:
                 out[key] = val_i
 
+        if self.action.script_pass_env_spec:
+            out["env_spec"] = self.env_spec
+
         return out
 
     def get_input_values_direct(self, label_dict: bool = True):
@@ -562,6 +576,10 @@ class ElementActionRun:
             typ = i.path[len("inputs.") :]
             if typ in input_types:
                 inputs[typ] = i.value
+
+        if self.action.script_pass_env_spec:
+            inputs["env_spec"] = self.env_spec
+
         return inputs
 
     def get_OFP_output_files(self) -> Dict[str, Union[str, List[str]]]:
@@ -585,6 +603,10 @@ class ElementActionRun:
         inputs = {}
         for inp_typ in self.action.output_file_parsers[0].inputs or []:
             inputs[inp_typ] = self.get(f"inputs.{inp_typ}")
+
+        if self.action.script_pass_env_spec:
+            inputs["env_spec"] = self.env_spec
+
         return inputs
 
     def get_OFP_outputs(self) -> Dict[str, Union[str, List[str]]]:
@@ -626,7 +648,7 @@ class ElementActionRun:
 
         # write the script if it is specified as a app data script, otherwise we assume
         # the script already exists in the working directory:
-        snip_path = self.action.get_snippet_script_path(self.action.script)
+        snip_path = self.action.get_snippet_script_path(self.action.script, self.env_spec)
         if snip_path:
             script_name = snip_path.name
             source_str = self.action.compose_source(snip_path)
@@ -677,32 +699,35 @@ class ElementActionRun:
             "stdout"/"stderr").
         """
         self.app.persistence_logger.debug("EAR.compose_commands")
+        env_spec = self.env_spec
+
         for ifg in self.action.input_file_generators:
             # TODO: there should only be one at this stage if expanded?
-            ifg.write_source(self.action)
+            ifg.write_source(self.action, env_spec)
 
         for ofp in self.action.output_file_parsers:
             # TODO: there should only be one at this stage if expanded?
             if ofp.output is None:
                 raise OutputFileParserNoOutputError()
-            ofp.write_source(self.action)
+            ofp.write_source(self.action, env_spec)
 
         if self.action.script:
             self.write_source(js_idx=jobscript.index, js_act_idx=JS_action_idx)
 
         command_lns = []
-        env_label = self.action.get_environment_label()
-        env = jobscript.submission.environments.get(env_label)
+        env = jobscript.submission.environments.get(**env_spec)
         if env.setup:
             command_lns += list(env.setup)
 
         shell_vars = {}  # keys are cmd_idx, each value is a list of tuples
         for cmd_idx, command in enumerate(self.action.commands):
-            cmd_str, shell_vars_i = command.get_command_line(
-                EAR=self, shell=jobscript.shell, env=env
-            )
-            shell_vars[cmd_idx] = shell_vars_i
-            command_lns.append(cmd_str)
+            if cmd_idx in self.commands_idx:
+                # only execute commands that have no rules, or all valid rules:
+                cmd_str, shell_vars_i = command.get_command_line(
+                    EAR=self, shell=jobscript.shell, env=env
+                )
+                shell_vars[cmd_idx] = shell_vars_i
+                command_lns.append(cmd_str)
 
         commands = "\n".join(command_lns) + "\n"
 
@@ -953,13 +978,24 @@ class ActionEnvironment(JSONLike):
         ),
     )
 
-    environment: str  # app.Environment
+    environment: Union[str, Dict[str, Any]]
     scope: Optional[app.ActionScope] = None
 
     def __post_init__(self):
         if self.scope is None:
             self.scope = self.app.ActionScope.any()
 
+        orig_env = copy.deepcopy(self.environment)
+        if isinstance(self.environment, str):
+            self.environment = {"name": self.environment}
+
+        if "name" not in self.environment:
+            raise ActionEnvironmentMissingNameError(
+                f"The action-environment environment specification must include a string "
+                f"`name` key, or be specified as string that is that name. Provided "
+                f"environment key was {orig_env!r}."
+            )
+
 
 class ActionRule(JSONLike):
     """Class to represent a rule/condition that must be True if an action is to be
@@ -997,6 +1033,7 @@ class ActionRule(JSONLike):
 
         self.rule = rule
         self.action = None  # assigned by parent action
+        self.command = None  # assigned by parent command
 
     def __eq__(self, other):
         if type(other) is not self.__class__:
@@ -1089,6 +1126,7 @@ class Action(JSONLike):
         script_data_out: Optional[str] = None,
         script_data_files_use_opt: Optional[bool] = False,
         script_exe: Optional[str] = None,
+        script_pass_env_spec: Optional[bool] = False,
         abortable: Optional[bool] = False,
         input_file_generators: Optional[List[app.InputFileGenerator]] = None,
         output_file_parsers: Optional[List[app.OutputFileParser]] = None,
@@ -1117,6 +1155,7 @@ class Action(JSONLike):
             script_data_files_use_opt if not self.script_is_python else True
         )
         self.script_exe = script_exe.lower() if script_exe else None
+        self.script_pass_env_spec = script_pass_env_spec
         self.environments = environments or [
             self.app.ActionEnvironment(environment="null_env")
         ]
@@ -1413,7 +1452,10 @@ class Action(JSONLike):
             commands=self.commands,
         )
 
-    def get_environment_label(self) -> str:
+    def get_environment_name(self) -> str:
+        return self.get_environment_spec()["name"]
+
+    def get_environment_spec(self) -> Dict[str, Any]:
         if not self._from_expand:
             raise RuntimeError(
                 f"Cannot choose a single environment from this action because it is not "
@@ -1422,7 +1464,7 @@ class Action(JSONLike):
         return self.environments[0].environment
 
     def get_environment(self) -> app.Environment:
-        return self.app.envs.get(self.get_environment_label())
+        return self.app.envs.get(**self.get_environment_spec())
 
     @staticmethod
     def is_snippet_script(script: str) -> bool:
@@ -1442,7 +1484,9 @@ class Action(JSONLike):
             return script
 
     @classmethod
-    def get_snippet_script_str(cls, script) -> str:
+    def get_snippet_script_str(
+        cls, script, env_spec: Optional[Dict[str, Any]] = None
+    ) -> str:
         if not cls.is_snippet_script(script):
             raise ValueError(
                 f"Must be an app-data script name (e.g. "
@@ -1450,14 +1494,24 @@ class Action(JSONLike):
             )
         pattern = r"\<\<script:(.*:?)\>\>"
         match_obj = re.match(pattern, script)
-        return match_obj.group(1)
+        out = match_obj.group(1)
+
+        if env_spec:
+            out = re.sub(
+                pattern=r"\<\<env:(.*?)\>\>",
+                repl=lambda match_obj: env_spec[match_obj.group(1)],
+                string=out,
+            )
+        return out
 
     @classmethod
-    def get_snippet_script_path(cls, script_path) -> Path:
+    def get_snippet_script_path(
+        cls, script_path, env_spec: Optional[Dict[str, Any]] = None
+    ) -> Path:
         if not cls.is_snippet_script(script_path):
             return False
 
-        path = cls.get_snippet_script_str(script_path)
+        path = cls.get_snippet_script_str(script_path, env_spec)
         if path in cls.app.scripts:
             path = cls.app.scripts.get(path)
 
@@ -1527,7 +1581,9 @@ class Action(JSONLike):
                     input_file_generators=[ifg],
                     environments=[self.get_input_file_generator_action_env(ifg)],
                     rules=main_rules + ifg.get_action_rules(),
+                    script_pass_env_spec=ifg.script_pass_env_spec,
                     abortable=ifg.abortable,
+                    # TODO: add script_data_in etc? and to OFP?
                 )
                 act_i._task_schema = self.task_schema
                 if ifg.input_file not in inp_files:
@@ -1558,6 +1614,7 @@ class Action(JSONLike):
                     output_file_parsers=[ofp],
                     environments=[self.get_output_file_parser_action_env(ofp)],
                     rules=list(self.rules) + ofp.get_action_rules(),
+                    script_pass_env_spec=ofp.script_pass_env_spec,
                     abortable=ofp.abortable,
                 )
                 act_i._task_schema = self.task_schema
@@ -1616,6 +1673,7 @@ class Action(JSONLike):
                 script_data_in=self.script_data_in,
                 script_data_out=self.script_data_out,
                 script_exe=self.script_exe,
+                script_pass_env_spec=self.script_pass_env_spec,
                 environments=[self.get_commands_action_env()],
                 abortable=self.abortable,
                 rules=main_rules,
@@ -1909,9 +1967,17 @@ class Action(JSONLike):
                 return True
 
     @TimeIt.decorator
-    def test_rules(self, element_iter) -> List[bool]:
+    def test_rules(self, element_iter) -> Tuple[bool, List[int]]:
         """Test all rules against the specified element iteration."""
-        return [i.test(element_iteration=element_iter) for i in self.rules]
+        rules_valid = [rule.test(element_iteration=element_iter) for rule in self.rules]
+        action_valid = all(rules_valid)
+        commands_idx = []
+        if action_valid:
+            for cmd_idx, cmd in enumerate(self.commands):
+                if any(not i.test(element_iteration=element_iter) for i in cmd.rules):
+                    continue
+                commands_idx.append(cmd_idx)
+        return action_valid, commands_idx
 
     def get_required_executables(self) -> Tuple[str]:
         """Return executable labels required by this action."""

hpcflow/sdk/core/command_files.py

@@ -131,6 +131,7 @@ class InputFileGenerator(JSONLike):
     inputs: List[app.Parameter]
     script: str = None
     environment: app.Environment = None
+    script_pass_env_spec: Optional[bool] = False
     abortable: Optional[bool] = False
     rules: Optional[List[app.ActionRule]] = None
 
@@ -188,11 +189,11 @@ class InputFileGenerator(JSONLike):
         out = out.format(script_str=script_str, main_block=main_block)
         return out
 
-    def write_source(self, action):
+    def write_source(self, action, env_spec: Dict[str, Any]):
 
         # write the script if it is specified as a snippet script, otherwise we assume
         # the script already exists in the working directory:
-        snip_path = action.get_snippet_script_path(self.script)
+        snip_path = action.get_snippet_script_path(self.script, env_spec)
         if snip_path:
             source_str = self.compose_source(snip_path)
             with Path(snip_path.name).open("wt", newline="\n") as fp:
@@ -255,6 +256,7 @@ class OutputFileParser(JSONLike):
     inputs: List[str] = None
     outputs: List[str] = None
     options: Dict = None
+    script_pass_env_spec: Optional[bool] = False
     abortable: Optional[bool] = False
     save_files: Union[List[str], bool] = True
     clean_up: Optional[List[str]] = None
@@ -342,14 +344,14 @@ class OutputFileParser(JSONLike):
         out = out.format(script_str=script_str, main_block=main_block)
         return out
 
-    def write_source(self, action):
+    def write_source(self, action, env_spec: Dict[str, Any]):
         if self.output is None:
             # might be used just for saving files:
             return
 
         # write the script if it is specified as a snippet script, otherwise we assume
         # the script already exists in the working directory:
-        snip_path = action.get_snippet_script_path(self.script)
+        snip_path = action.get_snippet_script_path(self.script, env_spec)
         if snip_path:
             source_str = self.compose_source(snip_path)
             with Path(snip_path.name).open("wt", newline="\n") as fp:

hpcflow/sdk/core/commands.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import partial
 from pathlib import Path
 import re
@@ -6,15 +6,24 @@ from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 import numpy as np
 
+from hpcflow.sdk import app
 from hpcflow.sdk.core.element import ElementResources
 from hpcflow.sdk.core.errors import NoCLIFormatMethodError
-from hpcflow.sdk.core.json_like import JSONLike
+from hpcflow.sdk.core.json_like import ChildObjectSpec, JSONLike
 from hpcflow.sdk.core.parameters import ParameterValue
 
 
 @dataclass
 class Command(JSONLike):
     _app_attr = "app"
+    _child_objects = (
+        ChildObjectSpec(
+            name="rules",
+            class_name="ActionRule",
+            is_multiple=True,
+            parent_ref="command",
+        ),
+    )
 
     command: Optional[str] = None
     executable: Optional[str] = None
@@ -23,6 +32,7 @@ class Command(JSONLike):
     stdout: Optional[str] = None
     stderr: Optional[str] = None
     stdin: Optional[str] = None
+    rules: Optional[List[app.ActionRule]] = field(default_factory=lambda: [])
 
     def __repr__(self) -> str:
         out = []
@@ -40,6 +50,8 @@ class Command(JSONLike):
             out.append(f"stderr={self.stderr!r}")
         if self.stdin:
             out.append(f"stdin={self.stdin!r}")
+        if self.rules:
+            out.append(f"rules={self.rules!r}")
 
         return f"{self.__class__.__name__}({', '.join(out)})"
 
@@ -107,6 +119,7 @@ class Command(JSONLike):
 
         file_regex = r"(\<\<file:{}\>\>?)"
         exe_script_regex = r"\<\<(executable|script):(.*?)\>\>"
+        env_specs_regex = r"\<\<env:(.*?)\>\>"
 
         # substitute executables:
         cmd_str = re.sub(
@@ -118,6 +131,12 @@ class Command(JSONLike):
         # executable command might itself contain variables defined in `variables`, and/or
         # an `<<args>>` variable::
         for var_key, var_val in (self.variables or {}).items():
+            # substitute any `<<env:>>` specifiers
+            var_val = re.sub(
+                pattern=env_specs_regex,
+                repl=lambda match_obj: EAR.env_spec[match_obj.group(1)],
+                string=var_val,
+            )
             cmd_str = cmd_str.replace(f"<<{var_key}>>", var_val)
             if "<<args>>" in cmd_str:
                 args_str = " ".join(self.arguments or [])

hpcflow/sdk/core/element.py

@@ -199,6 +199,7 @@ class ElementResources(JSONLike):
     scheduler_args: Optional[Dict] = None
     shell_args: Optional[Dict] = None
     os_name: Optional[str] = None
+    environments: Optional[Dict] = None
 
     # SGE scheduler specific:
     SGE_parallel_env: str = None
@@ -241,16 +242,24 @@ class ElementResources(JSONLike):
             return hash(tuple((keys, vals)))
 
         exclude = ("time_limit",)
-        sub_dicts = ("scheduler_args", "shell_args")
         dct = {k: copy.deepcopy(v) for k, v in self.__dict__.items() if k not in exclude}
-        if "options" in dct.get("scheduler_args", []):
-            dct["scheduler_args"]["options"] = _hash_dict(
-                dct["scheduler_args"]["options"]
-            )
 
-        for k in sub_dicts:
-            if k in dct:
-                dct[k] = _hash_dict(dct[k])
+        scheduler_args = dct["scheduler_args"]
+        shell_args = dct["shell_args"]
+        envs = dct["environments"]
+
+        if isinstance(scheduler_args, dict):
+            if "options" in scheduler_args:
+                dct["scheduler_args"]["options"] = _hash_dict(scheduler_args["options"])
+            dct["scheduler_args"] = _hash_dict(dct["scheduler_args"])
+
+        if isinstance(shell_args, dict):
+            dct["shell_args"] = _hash_dict(shell_args)
+
+        if isinstance(envs, dict):
+            for k, v in envs.items():
+                dct["environments"][k] = _hash_dict(v)
+            dct["environments"] = _hash_dict(dct["environments"])
 
         return _hash_dict(dct)
 
@@ -878,6 +887,28 @@ class ElementIteration:
         # an EAR?" which would then allow us to test a resources-based action rule.
 
         resource_specs = copy.deepcopy(self.get("resources"))
+
+        env_spec = action.get_environment_spec()
+        env_name = env_spec["name"]
+
+        # set default env specifiers, if none set:
+        if "any" not in resource_specs:
+            resource_specs["any"] = {}
+        if "environments" not in resource_specs["any"]:
+            resource_specs["any"]["environments"] = {env_name: copy.deepcopy(env_spec)}
+
+        for scope, dat in resource_specs.items():
+            if "environments" in dat:
+                # keep only relevant user-provided environment specifiers:
+                resource_specs[scope]["environments"] = {
+                    k: v for k, v in dat["environments"].items() if k == env_name
+                }
+                # merge user-provided specifiers into action specifiers:
+                resource_specs[scope]["environments"][env_name] = {
+                    **resource_specs[scope]["environments"].get(env_name, {}),
+                    **copy.deepcopy(env_spec),
+                }
+
         resources = {}
         for scope in action.get_possible_scopes()[::-1]:
             # loop in reverse so higher-specificity scopes take precedence:

hpcflow/sdk/core/errors.py

@@ -95,6 +95,10 @@ class MissingActionEnvironment(Exception):
     pass
 
 
+class ActionEnvironmentMissingNameError(Exception):
+    pass
+
+
 class FromSpecMissingObjectError(Exception):
     pass
 
@@ -391,3 +395,15 @@ class UnknownScriptDataKey(ValueError):
 
 class MissingVariableSubstitutionError(KeyError):
     pass
+
+
+class EnvironmentPresetUnknownEnvironmentError(ValueError):
+    pass
+
+
+class UnknownEnvironmentPresetError(ValueError):
+    pass
+
+
+class MultipleEnvironmentsError(ValueError):
+    pass