Merge pull request #947 from linsword13/gke-wm

Add a workflow manager to support running hpc apps in k8s cluster

Author: douglasjacobsen

lib/ramble/ramble/test/workflow_manager_functionality/gke_mpi_workflow_manager.py

@@ -0,0 +1,95 @@
+# 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.
+
+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_gke_mpi_workflow(request):
+    workspace_name = request.node.name
+    test_config = """
+ramble:
+  env_vars:
+    set:
+      OMP_NUM_THREADS: '{n_threads}'
+  variants:
+    workflow_manager: gke-mpi
+  variables:
+    mpi_command: mpirun -n {n_ranks}
+    processes_per_node: 1
+    n_nodes: 2
+    container_image: docker.pkg.dev/myproject/myimage
+    extra_metadata: |
+      a: 1
+      b: 2
+    extra_container_config_files: |
+      {experiment_run_dir}/app_config.txt
+  applications:
+    hostname:
+      workloads:
+        parallel:
+          experiments:
+            generated: {}
+"""
+    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])
+
+        run_path = os.path.join(ws.experiment_dir, "hostname", "parallel", "generated")
+        files = [f for f in os.listdir(run_path) if os.path.isfile(os.path.join(run_path, f))]
+        assert "batch_submit" in files
+        assert "batch_query" in files
+        assert "batch_cancel" in files
+        assert "gke_mpi.yaml" in files
+        assert "kustomization.yaml" in files
+        assert "launcher_execute_script" in files
+        assert "worker_execute_script" in files
+        assert "batch_print_deployment" in files
+        with open(os.path.join(run_path, "batch_submit")) as f:
+            content = f.read()
+            assert f"kubectl apply --kustomize {run_path}" in content
+        with open(os.path.join(run_path, "batch_query")) as f:
+            content = f.read()
+            assert "kubectl describe mpijobs hostname-parallel-generated" in content
+        with open(os.path.join(run_path, "batch_cancel")) as f:
+            content = f.read()
+            assert "kubectl delete mpijobs hostname-parallel-generated" in content
+        with open(os.path.join(run_path, "gke_mpi.yaml")) as f:
+            content = f.read()
+            assert "kind: MPIJob" in content
+            assert "name: hostname-parallel-generated" in content
+            assert "replicas: 2" in content
+            assert "image: docker.pkg.dev/myproject/myimage" in content
+        with open(os.path.join(run_path, "kustomization.yaml")) as f:
+            content = f.read()
+            assert "files:" in content
+            assert os.path.join(run_path, "app_config.txt") in content
+        with open(os.path.join(run_path, "launcher_execute_script")) as f:
+            content = f.read()
+            assert "hostname" in content
+        with open(os.path.join(run_path, "worker_execute_script")) as f:
+            content = f.read()
+            assert "sshd" in content
+        with open(os.path.join(run_path, "batch_print_deployment")) as f:
+            content = f.read()
+            assert "kubectl kustomize" in content

var/ramble/repos/builtin/workflow_managers/gke-mpi/batch_cancel.tpl

@@ -0,0 +1,2 @@
+#!/bin/bash
+kubectl delete mpijobs {job_name}

var/ramble/repos/builtin/workflow_managers/gke-mpi/batch_print_deployment.tpl

@@ -0,0 +1,2 @@
+#!/bin/bash
+kubectl kustomize {experiment_run_dir}

var/ramble/repos/builtin/workflow_managers/gke-mpi/batch_query.tpl

@@ -0,0 +1,16 @@
+#!/bin/bash
+echo "========================"
+echo "Print out mpi job status"
+echo "========================"
+echo ""
+kubectl describe mpijobs {job_name}
+
+lname=$(kubectl get pods | grep '{job_name}-launcher' | awk '{print $1}')
+if [ ! -z "$lname" ]; then
+    echo " "
+    echo "=========================="
+    echo "Print out the launcher log"
+    echo "=========================="
+    echo " "
+    kubectl logs $lname | tee {experiment_run_dir}/launcher.log
+fi

var/ramble/repos/builtin/workflow_managers/gke-mpi/batch_submit.tpl

@@ -0,0 +1,2 @@
+#!/bin/bash
+kubectl apply --kustomize {experiment_run_dir}

var/ramble/repos/builtin/workflow_managers/gke-mpi/gke_mpi.yaml.tpl

@@ -0,0 +1,48 @@
+apiVersion: kubeflow.org/v2beta1
+kind: MPIJob
+metadata:
+  name: {job_name}
+{extra_metadata_section}
+spec:
+  slotsPerWorker: {cores_per_node}
+  runPolicy:
+    cleanPodPolicy: Running
+  mpiReplicaSpecs:
+    Launcher:
+      replicas: 1
+      template:
+        spec:
+          hostPID: true
+          hostIPC: true
+          dnsPolicy: ClusterFirstWithHostNet
+          volumes:
+          - name: config
+            configMap:
+              name: gke-mpi-config
+          containers:
+          - image: {container_image}
+            name: mpi-launcher
+            volumeMounts:
+            - name: config
+              mountPath: /config
+            command: ["bash", "{launcher_script_path}"]
+            securityContext:
+              privileged: true
+    Worker:
+      replicas: {n_nodes}
+      template:
+        spec:
+          containers:
+          - image: {container_image}
+            name: mpi-worker
+            securityContext:
+              privileged: true
+            volumeMounts:
+            - name: config
+              mountPath: /config
+            command: ["bash", "{worker_script_path}"]
+          volumes:
+          - name: config
+            configMap:
+              name: gke-mpi-config
+            

var/ramble/repos/builtin/workflow_managers/gke-mpi/kustomization.yaml.tpl

@@ -0,0 +1,3 @@
+resources:
+- {gke_mpi_yaml}
+{config_map_gen_section}

var/ramble/repos/builtin/workflow_managers/gke-mpi/launcher_execute_script.tpl

@@ -0,0 +1,6 @@
+#!/bin/bash
+mkdir -p {container_work_dir} && cd {container_work_dir}
+# important to resolve symlink
+cp --remove-destination -r -L /config/* .
+
+{unformatted_command_without_logs}

var/ramble/repos/builtin/workflow_managers/gke-mpi/worker_execute_script.tpl

@@ -0,0 +1,5 @@
+#!/bin/bash
+mkdir -p {container_work_dir} && cd {container_work_dir}
+# important to resolve symlink
+cp --remove-destination -r -L /config/* .
+/usr/sbin/sshd -De -f /etc/ssh/sshd_config

var/ramble/repos/builtin/workflow_managers/gke-mpi/workflow_manager.py

@@ -0,0 +1,186 @@
+# 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.
+
+import os
+import shutil
+import textwrap
+
+from ramble.wmkit import *
+
+from spack.util.executable import Executable, ProcessError
+
+
+class GkeMpi(WorkflowManagerBase):
+    """GKE workflow manager that uses the MPI operator"""
+
+    name = "gke-mpi"
+
+    maintainers("linsword13")
+
+    tags("workflow", "gke", "mpi")
+
+    workflow_manager_variable(
+        name="job_name",
+        default="{application_name}-{workload_name}-{experiment_name}",
+        description="GKE job name",
+    )
+
+    workflow_manager_variable(
+        name="extra_metadata",
+        default="",
+        description="Extra line-separated key:val pairs for the metadata section",
+    )
+
+    workflow_manager_variable(
+        name="cores_per_node",
+        default="{processes_per_node}",
+        description="Cores per node",
+    )
+
+    workflow_manager_variable(
+        name="container_image",
+        default="",
+        description="url to the container image",
+    )
+
+    workflow_manager_variable(
+        name="extra_container_config_files",
+        default="",
+        description="extra line-separated list of config files to be mapped to containers",
+    )
+
+    workflow_manager_variable(
+        name="container_work_dir",
+        default="/config",
+        description="working directory inside the container",
+    )
+
+    workflow_manager_variable(
+        name="launcher_execute_script_template",
+        default="launcher_execute_script.tpl",
+        description="execute script template for the launcher",
+    )
+
+    register_template(
+        name="launcher_execute_script",
+        src_path="{launcher_execute_script_template}",
+    )
+
+    workflow_manager_variable(
+        name="worker_execute_script_template",
+        default="worker_execute_script.tpl",
+        description="execute script template for the workers",
+    )
+
+    register_template(
+        name="worker_execute_script",
+        src_path="{worker_execute_script_template}",
+    )
+
+    register_template(
+        name="gke_mpi_yaml",
+        src_path="gke_mpi.yaml.tpl",
+        dest_path="gke_mpi.yaml",
+        extra_vars_func="gke_mpi_yaml_vars",
+    )
+
+    def _gke_mpi_yaml_vars(self):
+        expander = self.app_inst.expander
+        extra_metadata_str = expander.expand_var_name("extra_metadata")
+        launcher_script = expander.expand_var_name("launcher_execute_script")
+        worker_script = expander.expand_var_name("worker_execute_script")
+        if extra_metadata_str:
+            extra_metadata_section = textwrap.indent(
+                extra_metadata_str, " " * 2
+            )
+        else:
+            extra_metadata_section = ""
+
+        return {
+            "extra_metadata_section": extra_metadata_section,
+            "launcher_script_path": os.path.join(
+                "/config", os.path.basename(launcher_script)
+            ),
+            "worker_script_path": os.path.join(
+                "/config", os.path.basename(worker_script)
+            ),
+        }
+
+    register_template(
+        name="kustomization.yaml",
+        src_path="kustomization.yaml.tpl",
+        dest_path="kustomization.yaml",
+        extra_vars_func="kustomization_yaml_vars",
+    )
+
+    def _kustomization_yaml_vars(self):
+        files = ["{launcher_execute_script}", "{worker_execute_script}"]
+        expander = self.app_inst.expander
+        extra_files_str = expander.expand_var_name(
+            "extra_container_config_files"
+        )
+        if extra_files_str:
+            files.extend(extra_files_str.split("\n"))
+        file_lines = "\n".join(
+            [expander.expand_var(f.lstrip("- ")) for f in files]
+        )
+        lines = [
+            "configMapGenerator:",
+            "- name: gke-mpi-config",
+            "  files:",
+            textwrap.indent(file_lines, "  - "),
+            "generatorOptions:",
+            # For some reason kustomization does not apply the generated name properly.
+            # So disable the suffix as a workaround.
+            "  disableNameSuffixHash: true",
+        ]
+        config_map_gen_section = "\n".join(lines)
+        return {
+            "config_map_gen_section": config_map_gen_section,
+        }
+
+    register_template(
+        name="batch_submit",
+        src_path="batch_submit.tpl",
+        dest_path="batch_submit",
+    )
+
+    register_template(
+        name="batch_query",
+        src_path="batch_query.tpl",
+        dest_path="batch_query",
+    )
+
+    register_template(
+        name="batch_cancel",
+        src_path="batch_cancel.tpl",
+        dest_path="batch_cancel",
+    )
+
+    # A convenience for printing the deployment config
+    register_template(
+        name="batch_print_deployment",
+        src_path="batch_print_deployment.tpl",
+        dest_path="batch_print_deployment",
+    )
+
+    def _prepare_analysis(self, workspace):
+        if workspace.dry_run:
+            return
+        expander = self.app_inst.expander
+        query_script = expander.expand_var_name("batch_query")
+        query_cmd = Executable(query_script)
+        try:
+            query_cmd(output=os.devnull)
+        except ProcessError as e:
+            logger.warn(f"batch_query returns error {e}")
+        run_dir = expander.expand_var_name("experiment_run_dir")
+        launcher_log = os.path.join(run_dir, "launcher.log")
+        if os.path.exists(launcher_log):
+            log_file = expander.expand_var_name("log_file")
+            shutil.copy2(launcher_log, log_file)