diff --git a/lib/ramble/ramble/test/conftest.py b/lib/ramble/ramble/test/conftest.py index b5e686390..a9dd8a3fb 100644 --- a/lib/ramble/ramble/test/conftest.py +++ b/lib/ramble/ramble/test/conftest.py @@ -148,6 +148,12 @@ def mock_pkg_mans_repo_path(): yield ramble.repository.Repo(ramble.paths.mock_builtin_path, obj_type) +@pytest.fixture(scope="function") +def mock_wms_repo_path(): + obj_type = ramble.repository.ObjectTypes.workflow_managers + yield ramble.repository.Repo(ramble.paths.mock_builtin_path, obj_type) + + @pytest.fixture(scope="function") def mutable_apps_repo_path(): obj_type = ramble.repository.ObjectTypes.applications @@ -166,6 +172,12 @@ def mutable_pkg_mans_repo_path(): yield ramble.repository.Repo(ramble.paths.builtin_path, obj_type) +@pytest.fixture(scope="function") +def mutable_wms_repo_path(): + obj_type = ramble.repository.ObjectTypes.workflow_managers + yield ramble.repository.Repo(ramble.paths.builtin_path, obj_type) + + @pytest.fixture(scope="function") def mock_applications(mock_apps_repo_path): """Use the 'builtin.mock' repository for applications instead of 'builtin'""" @@ -187,18 +199,25 @@ def mock_modifiers(mock_mods_repo_path): @pytest.fixture(scope="function") -def mock_package_managers(mock_mods_repo_path): +def mock_package_managers(mock_pkg_mans_repo_path): """Use the 'builtin.mock' repository for package managers of 'builtin'""" obj_type = ramble.repository.ObjectTypes.package_managers with ramble.repository.use_repositories( - mock_mods_repo_path, object_type=obj_type - ) as mock_mods_repo: - yield mock_mods_repo + mock_pkg_mans_repo_path, object_type=obj_type + ) as mock_pkg_mans_repo: + yield mock_pkg_mans_repo + + +@pytest.fixture(scope="function") +def mock_workflow_managers(mock_wms_repo_path): + """Use the 'builtin.mock' repository for package managers of 'builtin'""" + obj_type = ramble.repository.ObjectTypes.workflow_managers + with ramble.repository.use_repositories(mock_wms_repo_path, object_type=obj_type) as mock_repo: + yield mock_repo @pytest.fixture(scope="function") def mutable_applications(mutable_apps_repo_path): - """Use the 'builtin.mock' repository for applications instead of 'builtin'""" obj_type = ramble.repository.ObjectTypes.applications with ramble.repository.use_repositories( mutable_apps_repo_path, object_type=obj_type @@ -208,7 +227,6 @@ def mutable_applications(mutable_apps_repo_path): @pytest.fixture(scope="function") def mutable_modifiers(mutable_mods_repo_path): - """Use the 'builtin.mock' repository for modifiers instead of 'builtin'""" obj_type = ramble.repository.ObjectTypes.modifiers with ramble.repository.use_repositories( mutable_mods_repo_path, object_type=obj_type @@ -216,14 +234,21 @@ def mutable_modifiers(mutable_mods_repo_path): yield mods_repo -@pytest.fixture(scope="function") -def mutable_package_managers(mutable_mods_repo_path): - """Use the 'builtin.mock' repository for package_mangers instead of 'builtin'""" +def mutable_package_managers(mutable_pkg_mans_repo_path): obj_type = ramble.repository.ObjectTypes.package_managers with ramble.repository.use_repositories( - mutable_mods_repo_path, object_type=obj_type - ) as mods_repo: - yield mods_repo + mutable_pkg_mans_repo_path, object_type=obj_type + ) as pkg_mans_repo: + yield pkg_mans_repo + + +@pytest.fixture(scope="function") +def mutable_workflow_managers(mutable_wms_repo_path): + obj_type = ramble.repository.ObjectTypes.workflow_managers + with ramble.repository.use_repositories( + mutable_wms_repo_path, object_type=obj_type + ) as wms_repo: + yield wms_repo @pytest.fixture(scope="function") @@ -245,7 +270,7 @@ def mutable_mock_mods_repo(mock_mods_repo_path): @pytest.fixture(scope="function") -def mutable_mock_pkg_mans_repo(mock_mods_repo_path): +def mutable_mock_pkg_mans_repo(mock_pkg_mans_repo_path): """Function-scoped mock package managers, for tests that need to modify them.""" obj_type = ramble.repository.ObjectTypes.package_managers mock_repo = ramble.repository.Repo(ramble.paths.mock_builtin_path, object_type=obj_type) @@ -253,6 +278,15 @@ def mutable_mock_pkg_mans_repo(mock_mods_repo_path): yield mock_repo_path +@pytest.fixture(scope="function") +def mutable_mock_wms_repo(mock_wms_repo_path): + """Function-scoped mock package managers, for tests that need to modify them.""" + obj_type = ramble.repository.ObjectTypes.workflow_managers + mock_repo = ramble.repository.Repo(ramble.paths.mock_builtin_path, object_type=obj_type) + with ramble.repository.use_repositories(mock_repo, object_type=obj_type) as mock_repo_path: + yield mock_repo_path + + @pytest.fixture(scope="function") def default_config(): """Isolates the default configuration from the user configs. diff --git a/lib/ramble/ramble/test/data/config/base_workflow_manager_repos.yaml b/lib/ramble/ramble/test/data/config/base_workflow_manager_repos.yaml new file mode 100644 index 000000000..70fd089b5 --- /dev/null +++ b/lib/ramble/ramble/test/data/config/base_workflow_manager_repos.yaml @@ -0,0 +1,2 @@ +base_workflow_manager_repos: + - $ramble/var/ramble/repos/builtin diff --git a/lib/ramble/ramble/test/data/config/workflow_manager_repos.yaml b/lib/ramble/ramble/test/data/config/workflow_manager_repos.yaml new file mode 100644 index 000000000..dd5a5ed44 --- /dev/null +++ b/lib/ramble/ramble/test/data/config/workflow_manager_repos.yaml @@ -0,0 +1,2 @@ +workflow_manager_repos: + - $ramble/var/ramble/repos/builtin diff --git a/lib/ramble/ramble/test/workflow_manager_functionality/slurm_workflow_manager.py b/lib/ramble/ramble/test/workflow_manager_functionality/slurm_workflow_manager.py new file mode 100644 index 000000000..a2aae7605 --- /dev/null +++ b/lib/ramble/ramble/test/workflow_manager_functionality/slurm_workflow_manager.py @@ -0,0 +1,91 @@ +# Copyright 2022-2025 The Ramble Authors +# +# Licensed under the Apache License, Version 2.0 or the MIT license +# , at your +# option. This file may not be copied, modified, or distributed +# except according to those terms. + +import os + +import pytest + +import ramble.workspace +from ramble.main import RambleCommand + +workspace = RambleCommand("workspace") + +pytestmark = pytest.mark.usefixtures( + "mutable_config", + "mutable_mock_workspace_path", +) + + +def test_slurm_workflow(): + workspace_name = "test_slurm_workflow" + + test_config = """ +ramble: + variants: + workflow_manager: '{wm_name}' + variables: + # This batch_submit is overridden with slurm workflow manager + batch_submit: echo {wm_name} + mpi_command: mpirun -n {n_ranks} -hostfile hostfile + processes_per_node: 1 + wm_name: ['None', 'slurm'] + applications: + hostname: + workloads: + local: + experiments: + test_{wm_name}: + variables: + n_nodes: 1 + extra_sbatch_headers: "#SBATCH --gpus-per-task={n_threads}" +""" + with ramble.workspace.create(workspace_name) as ws: + ws.write() + config_path = os.path.join(ws.config_dir, ramble.workspace.config_file_name) + with open(config_path, "w+") as f: + f.write(test_config) + ws._re_read() + workspace("setup", "--dry-run", global_args=["-D", ws.root]) + + # assert the batch_submit is overridden, pointing to the generated script + all_exec_file = os.path.join(ws.root, "all_experiments") + with open(all_exec_file) as f: + content = f.read() + assert "echo None" in content + assert "echo slurm" not in content + assert os.path.join("hostname", "local", "test_slurm", "batch_submit") in content + + # Assert on no workflow manager + path = os.path.join(ws.experiment_dir, "hostname", "local", "test_None") + files = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))] + assert "slurm_execute_experiment" not in files + assert "batch_submit" not in files + assert "batch_query" not in files + assert "batch_cancel" not in files + + # Assert on slurm workflow manager + path = os.path.join(ws.experiment_dir, "hostname", "local", "test_slurm") + files = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))] + assert "slurm_execute_experiment" in files + assert "batch_submit" in files + assert "batch_query" in files + assert "batch_cancel" in files + with open(os.path.join(path, "batch_submit")) as f: + content = f.read() + assert "slurm_execute_experiment" in content + assert ".slurm_job" in content + with open(os.path.join(path, "slurm_execute_experiment")) as f: + content = f.read() + assert "scontrol show hostnames" in content + assert "#SBATCH --gpus-per-task=1" in content + with open(os.path.join(path, "batch_query")) as f: + content = f.read() + assert "squeue" in content + with open(os.path.join(path, "batch_cancel")) as f: + content = f.read() + assert "scancel" in content