Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a slurm workflow manager #789

Merged
merged 8 commits into from
Jan 10, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 14 additions & 0 deletions etc/ramble/defaults/base_workflow_manager_repos.yaml
Original file line number Diff line number Diff line change
@@ -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
14 changes: 14 additions & 0 deletions etc/ramble/defaults/workflow_manager_repos.yaml
Original file line number Diff line number Diff line change
@@ -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
51 changes: 46 additions & 5 deletions lib/ramble/ramble/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
##########################################
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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
Expand Down
2 changes: 2 additions & 0 deletions lib/ramble/ramble/cmd/common/info.py
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@
"package_manager_requirements": None,
# Package manager specific:
"package_manager_variables": None,
# Workflow manager specific:
"workflow_manager_variables": "wm_vars",
}


Expand Down
7 changes: 7 additions & 0 deletions lib/ramble/ramble/cmd/style.py
Original file line number Diff line number Diff line change
Expand Up @@ -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": [
Expand Down
4 changes: 4 additions & 0 deletions lib/ramble/ramble/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion lib/ramble/ramble/language/language_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
47 changes: 47 additions & 0 deletions lib/ramble/ramble/language/workflow_manager_language.py
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions lib/ramble/ramble/namespace.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,5 +62,6 @@ class namespace:

# For variants
package_manager = "package_manager"
workflow_manager = "workflow_manager"

metadata = "metadata"
30 changes: 30 additions & 0 deletions lib/ramble/ramble/repository.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,9 +62,11 @@
"applications",
"modifiers",
"package_managers",
"workflow_managers",
"base_applications",
"base_modifiers",
"base_package_managers",
"base_workflow_managers",
],
)

Expand Down Expand Up @@ -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",
Expand All @@ -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",
},
}


Expand All @@ -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)
Expand All @@ -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),
}

#####################################
Expand Down
32 changes: 32 additions & 0 deletions lib/ramble/ramble/schema/base_workflow_manager_repos.py
Original file line number Diff line number Diff line change
@@ -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,
}
32 changes: 32 additions & 0 deletions lib/ramble/ramble/schema/workflow_manager_repos.py
Original file line number Diff line number Diff line change
@@ -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,
}
Loading
Loading