Merge pull request #789 from linsword13/workflow

Add a slurm workflow manager
GoogleCloudPlatform · Jan 10, 2025 · bccd9cb · bccd9cb
2 parents 9be8e3f + ab82b5d
commit bccd9cb
Show file tree

Hide file tree

Showing 25 changed files with 710 additions and 21 deletions.
diff --git a/etc/ramble/defaults/base_workflow_manager_repos.yaml b/etc/ramble/defaults/base_workflow_manager_repos.yaml
@@ -0,0 +1,14 @@
+# -------------------------------------------------------------------------
+# This is the default ramble repository configuration. It includes the
+# builtin ramble base workflow manager repository.
+#
+# Users can override these settings by editing the following files.
+#
+# Per-ramble-instance settings (overrides defaults):
+#   $RAMBLE_ROOT/etc/ramble/base_workflow_manager_repos.yaml
+#
+# Per-user settings (overrides default and site settings):
+#   ~/.ramble/base_workflow_manager_repos.yaml
+# -------------------------------------------------------------------------
+base_workflow_manager_repos:
+  - $ramble/var/ramble/repos/builtin
diff --git a/etc/ramble/defaults/workflow_manager_repos.yaml b/etc/ramble/defaults/workflow_manager_repos.yaml
@@ -0,0 +1,14 @@
+# -------------------------------------------------------------------------
+# This is the default ramble repository configuration. It includes the
+# builtin ramble application repository.
+#
+# Users can override these settings by editing the following files.
+#
+# Per-ramble-instance settings (overrides defaults):
+#   $RAMBLE_ROOT/etc/ramble/repos.yaml
+#
+# Per-user settings (overrides default and site settings):
+#   ~/.ramble/repos.yaml
+# -------------------------------------------------------------------------
+workflow_manager_repos:
+  - $ramble/var/ramble/repos/builtin
diff --git a/lib/ramble/ramble/application.py b/lib/ramble/ramble/application.py
@@ -65,7 +65,8 @@
 from enum import Enum
 
 experiment_status = Enum(
-    "experiment_status", ["UNKNOWN", "SETUP", "RUNNING", "COMPLETE", "SUCCESS", "FAILED"]
+    "experiment_status",
+    ["UNKNOWN", "SETUP", "SUBMITTED", "RUNNING", "COMPLETE", "SUCCESS", "FAILED", "CANCELLED"],
 )
 
 _NULL_CONTEXT = "null"
@@ -181,6 +182,8 @@ def __init__(self, file_path):
         self.license_path = ""
         self.license_file = ""
 
+        self.workflow_manager = None
+
         ramble.util.directives.define_directive_methods(self)
 
     def experiment_lock(self):
@@ -236,7 +239,10 @@ def set_variants(self, variants):
                              experiment.
         """
         self.variants = variants.copy()
+        self._set_package_manager()
+        self._set_workflow_manager()
 
+    def _set_package_manager(self):
         if namespace.package_manager in self.variants:
             pkgman_name = self.expander.expand_var(
                 self.variants[namespace.package_manager], typed=True
@@ -277,6 +283,24 @@ def set_variants(self, variants):
                         }
                     )
 
+    def _set_workflow_manager(self):
+        if namespace.workflow_manager in self.variants:
+            workflow_name = self.expander.expand_var(
+                self.variants[namespace.workflow_manager], typed=True
+            )
+
+            if workflow_name is not None:
+                try:
+                    wfman_type = ramble.repository.ObjectTypes.workflow_managers
+                    self.workflow_manager = ramble.repository.get(workflow_name, wfman_type).copy()
+                    self.workflow_manager.set_application(self)
+                except ramble.repository.UnknownObjectError:
+                    logger.die(
+                        f"{workflow_name} is not a valid workflow manager. "
+                        "Valid workflow managers can be listed via:\n"
+                        "\tramble list --type workflow_managers"
+                    )
+
     def build_phase_order(self):
         if self._pipeline_graphs is not None:
             return
@@ -479,6 +503,10 @@ def build_used_variables(self, workspace):
             for var in self.package_manager.package_manager_variables.values():
                 self.variables[var.name] = var.default
 
+        if self.workflow_manager is not None:
+            for var in self.workflow_manager.wm_vars.values():
+                self.variables[var.name] = var.default
+
         ##########################################
         # Expand used variables to track all usage
         ##########################################
@@ -942,6 +970,9 @@ def _set_default_experiment_variables(self):
         for mod_inst in self._modifier_instances:
             var_sets.append(mod_inst.mode_variables())
 
+        if self.workflow_manager is not None:
+            var_sets.append(self.workflow_manager.wm_vars)
+
         for var_set in var_sets:
             for var, val in var_set.items():
                 if var not in self.variables.keys():
@@ -1715,10 +1746,13 @@ def format_context(context_match, context_format):
                     success = True
         success = success and criteria_list.passed()
 
-        if success:
-            self.set_status(status=experiment_status.SUCCESS)
-        else:
-            self.set_status(status=experiment_status.FAILED)
+        status = experiment_status.SUCCESS if success else experiment_status.FAILED
+        # When workflow_manager is present, only use app_status when workflow is completed.
+        if self.workflow_manager is not None:
+            wm_status = self.workflow_manager.get_status(workspace)
+            if not wm_status == experiment_status.COMPLETE:
+                status = wm_status
+        self.set_status(status)
 
         self._init_result()
 
@@ -2295,6 +2329,13 @@ def _get_template_config(
                     tpl_config,
                     obj_type=ramble.repository.ObjectTypes.package_managers,
                 )
+        if self.workflow_manager is not None:
+            for tpl_config in self.workflow_manager.templates.values():
+                yield _get_template_config(
+                    self.workflow_manager,
+                    tpl_config,
+                    obj_type=ramble.repository.ObjectTypes.workflow_managers,
+                )
 
     def _render_object_templates(self, extra_vars):
         run_dir = self.expander.experiment_run_dir

diff --git a/lib/ramble/ramble/cmd/common/info.py b/lib/ramble/ramble/cmd/common/info.py
@@ -55,6 +55,8 @@
     "package_manager_requirements": None,
     # Package manager specific:
     "package_manager_variables": None,
+    # Workflow manager specific:
+    "workflow_manager_variables": "wm_vars",
 }
 
 

diff --git a/lib/ramble/ramble/cmd/style.py b/lib/ramble/ramble/cmd/style.py
@@ -94,6 +94,13 @@ def is_object(f):
         "F403": [r"^from ramble.pkgmankit import \*$"],
         **common_object_exemptions,
     },
+    # exemptions applied only to workflow_manager.py files.
+    r"workflow_manager.py$": {
+        # Allow 'from ramble.modkit import *' in workflow_managers,
+        # but no other wildcards
+        "F403": [r"^from ramble.wmkit import \*$"],
+        **common_object_exemptions,
+    },
     # exemptions applied to all files.
     r".py$": {
         "E501": [

diff --git a/lib/ramble/ramble/config.py b/lib/ramble/ramble/config.py
@@ -69,9 +69,11 @@
 import ramble.schema.repos
 import ramble.schema.modifier_repos
 import ramble.schema.package_manager_repos
+import ramble.schema.workflow_manager_repos
 import ramble.schema.base_application_repos
 import ramble.schema.base_modifier_repos
 import ramble.schema.base_package_manager_repos
+import ramble.schema.base_workflow_manager_repos
 
 from ramble.error import RambleError
 from ramble.util.logger import logger
@@ -99,9 +101,11 @@
     "repos": ramble.schema.repos.schema,
     "modifier_repos": ramble.schema.modifier_repos.schema,
     "package_manager_repos": ramble.schema.package_manager_repos.schema,
+    "workflow_manager_repos": ramble.schema.workflow_manager_repos.schema,
     "base_application_repos": ramble.schema.base_application_repos.schema,
     "base_modifier_repos": ramble.schema.base_modifier_repos.schema,
     "base_package_manager_repos": ramble.schema.base_package_manager_repos.schema,
+    "base_workflow_manager_repos": ramble.schema.base_workflow_manager_repos.schema,
 }
 
 # Same as above, but including keys for workspaces

diff --git a/lib/ramble/ramble/language/language_base.py b/lib/ramble/ramble/language/language_base.py
@@ -27,7 +27,7 @@
 #: them
 reserved_names = []
 
-namespaces = ["ramble.app", "ramble.mod", "ramble.pkg_man", "ramble.package_manager"]
+namespaces = ["ramble.app", "ramble.mod", "ramble.pkg_man", "ramble.package_manager", "ramble.wm"]
 
 
 class DirectiveMeta(type):

diff --git a/lib/ramble/ramble/language/workflow_manager_language.py b/lib/ramble/ramble/language/workflow_manager_language.py
@@ -0,0 +1,47 @@
+# Copyright 2022-2025 The Ramble Authors
+#
+# Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+# https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+# <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
+# option. This file may not be copied, modified, or distributed
+# except according to those terms.
+
+from typing import Optional
+
+import ramble.language.shared_language
+
+
+class WorkflowManagerMeta(ramble.language.shared_language.SharedMeta):
+    _directive_names = set()
+    _directives_to_be_executed = []
+
+
+workflow_manager_directive = WorkflowManagerMeta.directive
+
+
+@workflow_manager_directive("wm_vars")
+def workflow_manager_variable(
+    name: str,
+    default,
+    description: str,
+    values: Optional[list] = None,
+):
+    """Define a variable for this wm
+    Args:
+        name: Name of variable
+        default: Default value if the variable is not defined
+        description: Description of the variable
+        values: Optional list of suggested values for this variable
+    """
+
+    def _define_wm_variable(wm):
+        import ramble.workload
+
+        wm.wm_vars[name] = ramble.workload.WorkloadVariable(
+            name,
+            default=default,
+            description=description,
+            values=values,
+        )
+
+    return _define_wm_variable
diff --git a/lib/ramble/ramble/namespace.py b/lib/ramble/ramble/namespace.py
@@ -62,5 +62,6 @@ class namespace:
 
     # For variants
     package_manager = "package_manager"
+    workflow_manager = "workflow_manager"
 
     metadata = "metadata"
diff --git a/lib/ramble/ramble/repository.py b/lib/ramble/ramble/repository.py
@@ -62,9 +62,11 @@
         "applications",
         "modifiers",
         "package_managers",
+        "workflow_managers",
         "base_applications",
         "base_modifiers",
         "base_package_managers",
+        "base_workflow_managers",
     ],
 )
 
@@ -99,6 +101,14 @@
         "accepted_configs": ["package_manager_repo.yaml", unified_config],
         "singular": "package manager",
     },
+    ObjectTypes.workflow_managers: {
+        "file_name": "workflow_manager.py",
+        "dir_name": "workflow_managers",
+        "abbrev": "wm",
+        "config_section": "workflow_manager_repos",
+        "accepted_configs": ["workflow_manager_repo.yaml", unified_config],
+        "singular": "workflow manager",
+    },
     ObjectTypes.base_applications: {
         "file_name": "base_application.py",
         "dir_name": "base_applications",
@@ -123,6 +133,14 @@
         "accepted_configs": ["base_package_manager_repo.yaml", unified_config],
         "singular": "base package manager",
     },
+    ObjectTypes.base_workflow_managers: {
+        "file_name": "base_workflow_manager.py",
+        "dir_name": "base_workflow_managers",
+        "abbrev": "base_wm",
+        "config_section": "base_workflow_manager_repos",
+        "accepted_configs": ["base_workflow_manager_repo.yaml", unified_config],
+        "singular": "base workflow manager",
+    },
 }
 
 
@@ -141,6 +159,11 @@ def _package_managers(repo_dirs=None):
     return _gen_path(repo_dirs=repo_dirs, obj_type=ObjectTypes.package_managers)
 
 
+def _workflow_managers(repo_dirs=None):
+    """Get the workflow managers singleton RepoPath instance for Ramble."""
+    return _gen_path(repo_dirs=repo_dirs, obj_type=ObjectTypes.workflow_managers)
+
+
 def _base_apps(repo_dirs=None):
     """Get the base applications singleton RepoPath instance for Ramble."""
     return _gen_path(repo_dirs=repo_dirs, obj_type=ObjectTypes.base_applications)
@@ -156,13 +179,20 @@ def _base_package_managers(repo_dirs=None):
     return _gen_path(repo_dirs=repo_dirs, obj_type=ObjectTypes.base_package_managers)
 
 
+def _base_workflow_managers(repo_dirs=None):
+    """Get the base workflow managers singleton RepoPath instance for Ramble."""
+    return _gen_path(repo_dirs=repo_dirs, obj_type=ObjectTypes.base_workflow_managers)
+
+
 paths = {
     ObjectTypes.applications: llnl.util.lang.Singleton(_apps),
     ObjectTypes.modifiers: llnl.util.lang.Singleton(_mods),
     ObjectTypes.package_managers: llnl.util.lang.Singleton(_package_managers),
+    ObjectTypes.workflow_managers: llnl.util.lang.Singleton(_workflow_managers),
     ObjectTypes.base_applications: llnl.util.lang.Singleton(_base_apps),
     ObjectTypes.base_modifiers: llnl.util.lang.Singleton(_base_mods),
     ObjectTypes.base_package_managers: llnl.util.lang.Singleton(_base_package_managers),
+    ObjectTypes.base_workflow_managers: llnl.util.lang.Singleton(_base_workflow_managers),
 }
 
 #####################################

diff --git a/lib/ramble/ramble/schema/base_workflow_manager_repos.py b/lib/ramble/ramble/schema/base_workflow_manager_repos.py
@@ -0,0 +1,32 @@
+# Copyright 2022-2025 The Ramble Authors
+#
+# Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+# https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+# <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
+# option. This file may not be copied, modified, or distributed
+# except according to those terms.
+
+"""Schema for base_workflow_manager_repos.yaml configuration file.
+.. literalinclude:: _ramble_root/lib/ramble/ramble/schema/base_workflow_manager_repos.py
+   :lines: 13-
+"""
+
+
+#: Properties for inclusion in other schemas
+properties = {
+    "base_workflow_manager_repos": {
+        "type": "array",
+        "default": [],
+        "items": {"type": "string"},
+    },
+}
+
+
+#: Full schema with metadata
+schema = {
+    "$schema": "http://json-schema.org/schema#",
+    "title": "Ramble base workflow manager repository configuration file schema",
+    "type": "object",
+    "additionalProperties": False,
+    "properties": properties,
+}
diff --git a/lib/ramble/ramble/schema/workflow_manager_repos.py b/lib/ramble/ramble/schema/workflow_manager_repos.py
@@ -0,0 +1,32 @@
+# Copyright 2022-2025 The Ramble Authors
+#
+# Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+# https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+# <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
+# option. This file may not be copied, modified, or distributed
+# except according to those terms.
+
+"""Schema for workflow_manager_repos.yaml configuration file.
+.. literalinclude:: _ramble_root/lib/ramble/ramble/schema/workflow_manager_repos.py
+   :lines: 13-
+"""
+
+
+#: Properties for inclusion in other schemas
+properties = {
+    "workflow_manager_repos": {
+        "type": "array",
+        "default": [],
+        "items": {"type": "string"},
+    },
+}
+
+
+#: Full schema with metadata
+schema = {
+    "$schema": "http://json-schema.org/schema#",
+    "title": "Ramble workflow manager repository configuration file schema",
+    "type": "object",
+    "additionalProperties": False,
+    "properties": properties,
+}