nemo-evaluator-launcher 0.1.0__py3-none-any.whl

nemo_evaluator_launcher/executors/slurm/executor.py

@@ -0,0 +1,996 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""SLURM executor implementation for nemo-evaluator-launcher.
+
+Handles submitting evaluation jobs to a SLURM cluster via SSH and sbatch scripts.
+"""
+
+import copy
+import os
+import re
+import shlex
+import subprocess
+import tempfile
+import time
+import warnings
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import yaml
+from omegaconf import DictConfig, OmegaConf
+
+from nemo_evaluator_launcher.common.execdb import (
+    ExecutionDB,
+    JobData,
+    generate_invocation_id,
+    generate_job_id,
+)
+from nemo_evaluator_launcher.common.helpers import (
+    get_api_key_name,
+    get_endpoint_url,
+    get_eval_factory_command,
+    get_eval_factory_dataset_size_from_run_config,
+    get_health_url,
+    get_timestamp_string,
+)
+from nemo_evaluator_launcher.common.mapping import (
+    get_task_from_mapping,
+    load_tasks_mapping,
+)
+from nemo_evaluator_launcher.executors.base import (
+    BaseExecutor,
+    ExecutionState,
+    ExecutionStatus,
+)
+from nemo_evaluator_launcher.executors.registry import register_executor
+
+
+@register_executor("slurm")
+class SlurmExecutor(BaseExecutor):
+    @staticmethod
+    def execute_eval(cfg: DictConfig, dry_run: bool = False) -> str:
+        """Submit evaluation jobs to a SLURM cluster using the provided configuration.
+
+        Args:
+            cfg: The configuration object for the evaluation run.
+            dry_run: If True, prepare scripts and save them without submission.
+
+        Returns:
+            str: The invocation ID for the evaluation run.
+
+        Raises:
+            AssertionError: If deployment type is 'none'.
+            RuntimeError: If remote directory creation or sbatch submission fails.
+        """
+
+        # Generate invocation ID
+        invocation_id = generate_invocation_id()
+
+        local_runsub_paths = []
+        remote_runsub_paths = []
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            timestamp = get_timestamp_string(include_microseconds=False)
+            rundir_name = timestamp + "-" + invocation_id
+            remote_rundir = Path(cfg.execution.output_dir) / rundir_name
+            local_rundir = Path(tmpdirname) / rundir_name
+            local_rundir.mkdir()
+
+            # Preload mapping for image resolution
+            tasks_mapping = load_tasks_mapping()
+            eval_images: list[str] = []
+
+            for idx, task in enumerate(cfg.evaluation.tasks):
+                # calculate job_id
+                job_id = f"{invocation_id}.{idx}"
+
+                # preapre locally
+                remote_task_subdir = remote_rundir / task.name
+                local_task_subdir = local_rundir / task.name
+                local_task_subdir.mkdir()  # this ensures the task.name hasn't been used
+                (local_task_subdir / "logs").mkdir()
+                (local_task_subdir / "artifacts").mkdir()
+
+                # resolve eval image and pass directly via task override
+                task_definition = get_task_from_mapping(task.name, tasks_mapping)
+                eval_image = task_definition["container"]
+                if "container" in task:
+                    eval_image = task["container"]
+
+                eval_images.append(eval_image)
+
+                # generate and write down sbatch script
+                sbatch_script_content_str = _create_slurm_sbatch_script(
+                    cfg=cfg,
+                    task=task,
+                    eval_image=eval_image,
+                    remote_task_subdir=remote_task_subdir,
+                    invocation_id=invocation_id,
+                    job_id=job_id,
+                )
+                local_runsub_path = local_task_subdir / "run.sub"
+                remote_runsub_path = remote_task_subdir / "run.sub"
+                with open(local_runsub_path, "w") as f:
+                    f.write(sbatch_script_content_str.rstrip("\n") + "\n")
+
+                local_runsub_paths.append(local_runsub_path)
+                remote_runsub_paths.append(remote_runsub_path)
+
+            if dry_run:
+                print("\n\n=============================================\n\n")
+                print("DRY RUN: SLURM scripts prepared")
+                for idx, local_runsub_path in enumerate(local_runsub_paths):
+                    print(f"\n\n =========== Task {idx} ===================== \n\n")
+                    with open(local_runsub_path, "r") as f:
+                        print(f.read())
+                print("\nTo submit jobs, run the executor without --dry-run")
+                return invocation_id
+
+            socket = str(Path(tmpdirname) / "socket")
+            socket_or_none = _open_master_connection(
+                username=cfg.execution.username,
+                hostname=cfg.execution.hostname,
+                socket=socket,
+            )
+            _make_remote_execution_output_dir(
+                dirpath=cfg.execution.output_dir,
+                username=cfg.execution.username,
+                hostname=cfg.execution.hostname,
+                socket=socket_or_none,
+            )
+            _rsync_upload_rundirs(
+                local_sources=[local_rundir],
+                remote_target=cfg.execution.output_dir,
+                username=cfg.execution.username,
+                hostname=cfg.execution.hostname,
+            )
+            slurm_job_ids = _sbatch_remote_runsubs(
+                remote_runsub_paths=remote_runsub_paths,
+                username=cfg.execution.username,
+                hostname=cfg.execution.hostname,
+                socket=socket_or_none,
+            )
+            _close_master_connection(
+                username=cfg.execution.username,
+                hostname=cfg.execution.hostname,
+                socket=socket_or_none,
+            )
+
+            # save launched jobs metadata
+            db = ExecutionDB()
+            for idx, (slurm_job_id, remote_runsub_path) in enumerate(
+                zip(slurm_job_ids, remote_runsub_paths)
+            ):
+                db.write_job(
+                    job=JobData(
+                        invocation_id=invocation_id,
+                        job_id=generate_job_id(invocation_id, idx),
+                        timestamp=time.time(),
+                        executor="slurm",
+                        data={
+                            "slurm_job_id": slurm_job_id,
+                            "remote_rundir_path": str(remote_runsub_path.parent),
+                            "hostname": cfg.execution.hostname,
+                            "username": cfg.execution.username,
+                            "eval_image": eval_images[idx],
+                        },
+                        config=OmegaConf.to_object(cfg),
+                    )
+                )
+            return invocation_id
+
+    @staticmethod
+    def get_status(id: str) -> List[ExecutionStatus]:
+        """Get the status of a specific SLURM job or all jobs in an invocation group.
+
+        Args:
+            id: Unique job identifier or invocation identifier.
+
+        Returns:
+            List containing the execution status for the job(s).
+        """
+        db = ExecutionDB()
+
+        # If id looks like an invocation_id (8 hex digits, no dot), get all jobs for it
+        if len(id) == 8 and "." not in id:
+            jobs = db.get_jobs(id)
+            if not jobs:
+                return []
+            return SlurmExecutor._get_status_for_invocation(jobs)
+
+        # Otherwise, treat as job_id
+        else:
+            job_data = db.get_job(id)
+            if job_data is None or job_data.executor != "slurm":
+                return []
+            return [SlurmExecutor._get_status_for_job(id, job_data)]
+
+    @staticmethod
+    def _get_status_for_job(id: str, job_data: JobData) -> ExecutionStatus:
+        slurm_job_id = job_data.data.get("slurm_job_id")
+        if not slurm_job_id:
+            return ExecutionStatus(id=id, state=ExecutionState.FAILED)
+
+        try:
+            return SlurmExecutor._query_slurm_for_status_and_progress(
+                slurm_job_ids=[slurm_job_id],
+                remote_rundir_paths=[Path(job_data.data.get("remote_rundir_path"))],
+                username=job_data.data["username"],
+                hostname=job_data.data["hostname"],
+                job_id_to_execdb_id={slurm_job_id: id},
+            )[0]
+        except Exception:
+            return ExecutionStatus(id=id, state=ExecutionState.FAILED)
+
+    @staticmethod
+    def _get_status_for_invocation(jobs: dict) -> List[ExecutionStatus]:
+        slurm_job_ids = []
+        remote_rundir_paths = []
+        job_id_to_execdb_id = {}
+        username = None
+        hostname = None
+
+        for job_id, job_data in jobs.items():
+            if job_data.executor != "slurm":
+                continue
+            slurm_job_id = job_data.data.get("slurm_job_id")
+            if slurm_job_id:
+                slurm_job_ids.append(slurm_job_id)
+                remote_rundir_paths.append(
+                    Path(job_data.data.get("remote_rundir_path"))
+                )
+                job_id_to_execdb_id[slurm_job_id] = job_id
+                username = job_data.data.get("username")
+                hostname = job_data.data.get("hostname")
+
+        if not slurm_job_ids or not remote_rundir_paths or not username or not hostname:
+            return [
+                ExecutionStatus(id=job_id, state=ExecutionState.FAILED)
+                for job_id in jobs.keys()
+            ]
+
+        try:
+            return SlurmExecutor._query_slurm_for_status_and_progress(
+                slurm_job_ids=slurm_job_ids,
+                remote_rundir_paths=remote_rundir_paths,
+                username=username,
+                hostname=hostname,
+                job_id_to_execdb_id=job_id_to_execdb_id,
+            )
+        except Exception:
+            return [
+                ExecutionStatus(id=job_id, state=ExecutionState.FAILED)
+                for job_id in jobs.keys()
+            ]
+
+    @staticmethod
+    def _query_slurm_for_status_and_progress(
+        slurm_job_ids: List[str],
+        remote_rundir_paths: List[Path],
+        username: str,
+        hostname: str,
+        job_id_to_execdb_id: dict,
+    ) -> List[ExecutionStatus]:
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            socket = str(Path(tmpdirname) / "socket")
+            socket_or_none = _open_master_connection(
+                username=username,
+                hostname=hostname,
+                socket=socket,
+            )
+            # get slurm job status for initial jobs:
+            slurm_jobs_status = _query_slurm_jobs_status(
+                slurm_job_ids=slurm_job_ids,
+                username=username,
+                hostname=hostname,
+                socket=socket_or_none,
+            )
+            # handle slurm status for autoresumed jobs:
+            autoresumed_slurm_job_ids = _read_autoresumed_slurm_job_ids(
+                slurm_job_ids=slurm_job_ids,
+                remote_rundir_paths=remote_rundir_paths,
+                username=username,
+                hostname=hostname,
+                socket=socket_or_none,
+            )
+            latest_slurm_job_ids = {
+                slurm_job_id: slurm_job_id_list[-1]
+                for slurm_job_id, slurm_job_id_list in autoresumed_slurm_job_ids.items()
+                if len(slurm_job_id_list) > 0 and slurm_job_id_list[-1] != slurm_job_id
+            }
+            latest_slurm_jobs_status = _query_slurm_jobs_status(
+                slurm_job_ids=list(latest_slurm_job_ids.values()),
+                username=username,
+                hostname=hostname,
+                socket=socket_or_none,
+            )
+            # get progress:
+            progress_list = _get_progress(
+                remote_rundir_paths=remote_rundir_paths,
+                username=username,
+                hostname=hostname,
+                socket=socket_or_none,
+            )
+            _close_master_connection(
+                username=username,
+                hostname=hostname,
+                socket=socket_or_none,
+            )
+        statuses = []
+        for i, slurm_job_id in enumerate(slurm_job_ids):
+            slurm_status = slurm_jobs_status[slurm_job_id]
+            if slurm_job_id in latest_slurm_job_ids:
+                latest_slurm_job_id = latest_slurm_job_ids[slurm_job_id]
+                slurm_status = latest_slurm_jobs_status[latest_slurm_job_id]
+            progress = progress_list[i]
+            progress = progress if progress is not None else "unknown"
+            execution_state = SlurmExecutor._map_slurm_state_to_execution_state(
+                slurm_status
+            )
+            execdb_job_id = job_id_to_execdb_id.get(slurm_job_id)
+            if execdb_job_id:
+                statuses.append(
+                    ExecutionStatus(
+                        id=execdb_job_id,
+                        state=execution_state,
+                        progress=progress,
+                    )
+                )
+        return statuses
+
+    @staticmethod
+    def _map_slurm_state_to_execution_state(slurm_status: str) -> ExecutionState:
+        """Map SLURM state to ExecutionState.
+
+        Args:
+            slurm_status: SLURM status string.
+
+        Returns:
+            Corresponding ExecutionState.
+        """
+        if slurm_status in ["COMPLETED"]:
+            return ExecutionState.SUCCESS
+        elif slurm_status in [
+            "PENDING",
+            "RESV_DEL_HOLD",
+            "REQUEUE_FED",
+            "REQUEUE_HOLD",
+            "REQUEUED",
+            "REVOKED",
+        ]:
+            return ExecutionState.PENDING
+        elif slurm_status in ["RUNNING", "CONFIGURING", "SUSPENDED", "COMPLETING"]:
+            return ExecutionState.RUNNING
+        elif slurm_status in ["PREEMPTED", "TIMEOUT", "NODE_FAIL"]:
+            return ExecutionState.PENDING  # autoresume
+        elif slurm_status in ["CANCELLED"]:
+            return ExecutionState.KILLED
+        elif slurm_status in ["FAILED"]:
+            return ExecutionState.FAILED
+        else:
+            return ExecutionState.FAILED
+
+    @staticmethod
+    def kill_job(job_id: str) -> None:
+        """Kill a SLURM job.
+
+        Args:
+            job_id: The job ID to kill.
+        """
+        db = ExecutionDB()
+        job_data = db.get_job(job_id)
+
+        if job_data is None:
+            raise ValueError(f"Job {job_id} not found")
+
+        if job_data.executor != "slurm":
+            raise ValueError(
+                f"Job {job_id} is not a slurm job (executor: {job_data.executor})"
+            )
+
+        killed_something = False
+
+        result = _kill_slurm_job(
+            slurm_job_ids=[job_data.data.get("slurm_job_id")],
+            username=job_data.data.get("username"),
+            hostname=job_data.data.get("hostname"),
+            socket=job_data.data.get("socket"),
+        )
+
+        if result.returncode == 0:
+            killed_something = True
+
+        # Mark job as killed in database if we killed something
+        if killed_something:
+            job_data.data["killed"] = True
+            db.write_job(job_data)
+        else:
+            raise RuntimeError(
+                f"Could not find or kill job {job_id} (slurm_job_id: {job_data.data.get('slurm_job_id')})"
+            )
+
+
+def _create_slurm_sbatch_script(
+    cfg: DictConfig,
+    task: DictConfig,
+    eval_image: str,
+    remote_task_subdir: Path,
+    invocation_id: str,
+    job_id: str,
+) -> str:
+    """Generate the contents of a SLURM sbatch script for a given evaluation task.
+
+    Args:
+        cfg: The configuration object for the evaluation run.
+        task: The evaluation task configuration.
+        remote_task_subdir: The remote directory path for the `run.sub` file.
+        invocation_id: The invocation ID for this evaluation run.
+        job_id: The complete job ID string.
+
+    Returns:
+        str: The contents of the sbatch script.
+    """
+    # get task from mapping, overrides, urls
+    tasks_mapping = load_tasks_mapping()
+    task_definition = get_task_from_mapping(task.name, tasks_mapping)
+    health_url = get_health_url(cfg, get_endpoint_url(cfg, task, task_definition))
+
+    # TODO(public release): convert to template
+    s = "#!/bin/bash\n"
+
+    # SBATCH headers
+    s += "#SBATCH --time {}\n".format(cfg.execution.walltime)
+    s += "#SBATCH --account {}\n".format(cfg.execution.account)
+    s += "#SBATCH --partition {}\n".format(cfg.execution.partition)
+    s += "#SBATCH --nodes {}\n".format(cfg.execution.num_nodes)
+    s += "#SBATCH --ntasks-per-node {}\n".format(cfg.execution.ntasks_per_node)
+    if cfg.execution.get("gpus_per_node", None) is not None:
+        s += "#SBATCH --gpus-per-node {}\n".format(cfg.execution.gpus_per_node)
+    if hasattr(cfg.execution, "gres"):
+        s += "#SBATCH --gres {}\n".format(cfg.execution.gres)
+    job_name = "{account}-{subproject}.{details}".format(
+        account=cfg.execution.account,
+        subproject=cfg.execution.subproject,
+        details=remote_task_subdir.name,
+    )
+    s += "#SBATCH --job-name {}\n".format(job_name)
+    s += "#SBATCH --exclusive\n"
+    s += "#SBATCH --output {}\n".format(remote_task_subdir / "logs" / "slurm-%A.out")
+    s += "\n"
+    s += f'TASK_DIR="{str(remote_task_subdir)}"\n'
+    s += "\n"
+
+    # collect all env vars
+    env_vars = copy.deepcopy(dict(cfg.evaluation.get("env_vars", {})))
+    env_vars.update(task.get("env_vars", {}))
+    api_key_name = get_api_key_name(cfg)
+    if api_key_name:
+        assert "API_KEY" not in env_vars
+        env_vars["API_KEY"] = api_key_name
+
+    # check if the environment variables are set
+    for env_var in env_vars.values():
+        if os.getenv(env_var) is None:
+            raise ValueError(f"Trying to pass an unset environment variable {env_var}.")
+
+    # check if required env vars are defined:
+    for required_env_var in task_definition.get("required_env_vars", []):
+        if required_env_var not in env_vars.keys():
+            raise ValueError(
+                f"{task.name} task requires environment variable {required_env_var}."
+                " Specify it in the task subconfig in the 'env_vars' dict as the following"
+                f" pair {required_env_var}: YOUR_ENV_VAR_NAME"
+            )
+
+    # save env vars:
+    for env_var_dst, env_var_src in env_vars.items():
+        s += f"export {env_var_dst}={os.getenv(env_var_src)}\n"
+    all_env_vars = {
+        **cfg.execution.get("env_vars", {}).get("deployment", {}),
+        **cfg.execution.get("env_vars", {}).get("evaluation", {}),
+    }
+    if cfg.deployment.get("env_vars"):
+        warnings.warn(
+            "cfg.deployment.env_vars will be deprecated in future versions. "
+            "Use cfg.execution.env_vars.deployment instead.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        all_env_vars.update(cfg.deployment["env_vars"])
+    for env_var_dst, env_var_value in all_env_vars.items():
+        s += f"export {env_var_dst}={env_var_value}\n"
+    if env_vars:
+        s += "\n"
+
+    # auto resume after timeout
+    s += _AUTORESUME_HANDLER
+    s += "\n\n"
+
+    # echo the current SLURM_JOB_ID
+    s += "# save the current job id\n"
+    s += "echo $SLURM_JOB_ID >> {}\n\n".format(
+        remote_task_subdir / ".slurm_job_id.list"
+    )
+
+    # shell options
+    s += "set -e  # exit immediately if any command exits with a non-zero status\n"
+    s += "set -u  # treat unset variables as an error when substituting\n"
+    s += "set -x  # print commands and their arguments as they are executed\n"
+    s += "\n"
+
+    # prepare deployment mounts
+    deployment_mounts_list = []
+    if cfg.deployment.type != "none":
+        if checkpoint_path := cfg.deployment.get("checkpoint_path"):
+            deployment_mounts_list.append(f"{checkpoint_path}:/checkpoint:ro")
+        if cache_path := cfg.deployment.get("cache_path"):
+            deployment_mounts_list.append(f"{cache_path}:/cache")
+        for source_mnt, target_mnt in (
+            cfg.execution.get("mounts", {}).get("deployment", {}).items()
+        ):
+            deployment_mounts_list.append(f"{source_mnt}:{target_mnt}")
+
+        # add deployment srun command
+        s += "# deployment server\n"
+        s += "srun --mpi pmix --overlap "
+        s += "--container-image {} ".format(cfg.deployment.image)
+        if deployment_mounts_list:
+            s += "--container-mounts {} ".format(",".join(deployment_mounts_list))
+        if not cfg.execution.get("mounts", {}).get("mount_home", True):
+            s += "--no-container-mount-home "
+        s += "--output {} ".format(remote_task_subdir / "logs" / "server-%A.out")
+        deployment_env_var_names = list(
+            cfg.execution.get("env_vars", {}).get("deployment", {})
+        )
+        if cfg.deployment.get("env_vars"):
+            warnings.warn(
+                "cfg.deployment.env_vars will be deprecated in future versions. "
+                "Use cfg.execution.env_vars.deployment instead.",
+                category=DeprecationWarning,
+                stacklevel=2,
+            )
+            deployment_env_var_names.extend(list(cfg.deployment["env_vars"]))
+        if deployment_env_var_names:
+            s += f"--container-env {','.join(deployment_env_var_names)} "
+        s += "{} &\n\n".format(cfg.deployment.command)  # run asynchronously
+        s += (
+            "SERVER_PID=$!  # capture the PID of the server background srun process\n\n"
+        )
+
+        # wait for the server to initialize
+        s += _WAIT_FOR_SERVER_HANDLER.format(health_url=health_url)
+        s += "\n\n"
+
+    # prepare evaluation mounts
+    evaluation_mounts_list = [
+        "{}:/results".format(remote_task_subdir / "artifacts"),
+    ]
+    for source_mnt, target_mnt in (
+        cfg.execution.get("mounts", {}).get("evaluation", {}).items()
+    ):
+        evaluation_mounts_list.append(f"{source_mnt}:{target_mnt}")
+
+    # add evaluation srun command
+    s += "# evaluation client\n"
+    s += "srun --mpi pmix --overlap "
+    s += "--container-image {} ".format(eval_image)
+    evaluation_env_var_names = list(
+        cfg.execution.get("env_vars", {}).get("evaluation", {})
+    )
+    if evaluation_env_var_names:
+        s += "--container-env {} ".format(",".join(evaluation_env_var_names))
+    if not cfg.execution.get("mounts", {}).get("mount_home", True):
+        s += "--no-container-mount-home "
+    s += "--container-mounts {} ".format(",".join(evaluation_mounts_list))
+    s += "--output {} ".format(remote_task_subdir / "logs" / "client-%A.out")
+    s += "bash -c '"
+    s += get_eval_factory_command(cfg, task, task_definition)
+    s += "'\n\n"
+
+    # terminate the server after all evaluation clients finish
+    if cfg.deployment.type != "none":
+        s += "kill $SERVER_PID  # terminate the server to finish gracefully\n\n"
+
+    # auto-export
+    if cfg.execution.get("auto_export", {}).get("destinations", []):
+        s += _generate_auto_export_section(cfg, job_id)
+
+    return s
+
+
+def _generate_auto_export_section(
+    cfg: DictConfig,
+    job_id: str,  # Complete job_id string
+) -> str:
+    """Generate simple auto-export section for sbatch script."""
+    auto_export_config = cfg.execution.get("auto_export", {})
+    destinations = auto_export_config.get("destinations", [])
+
+    if not destinations:
+        return ""
+
+    s = "\n# Auto-export on success\n"
+    s += "EVAL_EXIT_CODE=$?\n"
+    s += "if [ $EVAL_EXIT_CODE -eq 0 ]; then\n"
+    s += "    echo 'Evaluation completed successfully. Starting auto-export...'\n"
+    s += "    set +e\n"  # per exporter failure allowed
+    s += "    set +x\n"
+    s += '    cd "$TASK_DIR/artifacts"\n'
+    auto_export_cfg = OmegaConf.to_container(
+        cfg.execution.get("auto_export", {}), resolve=True
+    )
+    yaml_str = yaml.safe_dump(
+        {"execution": {"auto_export": auto_export_cfg}}, sort_keys=False
+    )
+    s += "    cat > export_config.yml << 'EOF'\n"
+    s += yaml_str
+    s += "EOF\n"
+    for dest in destinations:
+        s += f"    echo 'Exporting to {dest}...'\n"
+        s += f"    nemo-evaluator-launcher export {job_id} --dest {dest} || echo 'Export to {dest} failed'\n"
+
+    s += "    echo 'Auto-export completed.'\n"
+    s += "else\n"
+    s += "    echo 'Evaluation failed with exit code $EVAL_EXIT_CODE. Skipping auto-export.'\n"
+    s += "fi\n"
+
+    return s
+
+
+def _open_master_connection(
+    username: str,
+    hostname: str,
+    socket: str,
+) -> str | None:
+    ssh_command = f"ssh -MNf -S {socket} {username}@{hostname}"
+    completed_process = subprocess.run(args=shlex.split(ssh_command))
+    if completed_process.returncode == 0:
+        return socket
+    return None
+
+
+def _close_master_connection(
+    username: str,
+    hostname: str,
+    socket: str | None,
+) -> None:
+    if socket is None:
+        return
+    ssh_command = f"ssh -O exit -S {socket} {username}@{hostname}"
+    completed_process = subprocess.run(
+        args=shlex.split(ssh_command), capture_output=True
+    )
+    if completed_process.returncode != 0:
+        raise RuntimeError(
+            "failed to close the master connection\n{}".format(
+                completed_process.stderr.decode("utf-8")
+            )
+        )
+
+
+def _make_remote_execution_output_dir(
+    dirpath: str,
+    username: str,
+    hostname: str,
+    socket: str | None,
+) -> None:
+    mkdir_command = f"mkdir -p {dirpath}"
+    ssh_command = ["ssh"]
+    if socket is not None:
+        ssh_command.append(f"-S {socket}")
+    ssh_command.append(f"{username}@{hostname}")
+    ssh_command.append(mkdir_command)
+    ssh_command = " ".join(ssh_command)
+    completed_process = subprocess.run(args=shlex.split(ssh_command))
+    if completed_process.returncode != 0:
+        raise RuntimeError(
+            "failed to make a remote execution output dir\n{}".format(
+                completed_process.stderr.decode("utf-8")
+            )
+        )
+
+
+def _rsync_upload_rundirs(
+    local_sources: List[Path],
+    remote_target: str,
+    username: str,
+    hostname: str,
+) -> None:
+    """Upload local run directories to a remote host using rsync over SSH.
+
+    Args:
+        local_sources: List of local Path objects to upload.
+        remote_target: Remote directory path as a string.
+        hostname: SSH hostname.
+        username: SSH username.
+
+    Raises:
+        RuntimeError: If rsync fails.
+    """
+    for local_source in local_sources:
+        assert local_source.is_dir()
+    remote_destination_str = f"{username}@{hostname}:{remote_target}"
+    local_sources_str = " ".join(map(str, local_sources))
+    rsync_upload_command = f"rsync -qcaz {local_sources_str} {remote_destination_str}"
+    completed_process = subprocess.run(args=shlex.split(rsync_upload_command))
+    if completed_process.returncode != 0:
+        raise RuntimeError(
+            "failed to upload local sources\n{}".format(
+                completed_process.stderr.decode("utf-8")
+            )
+        )
+
+
+def _sbatch_remote_runsubs(
+    remote_runsub_paths: List[Path],
+    username: str,
+    hostname: str,
+    socket: str | None,
+) -> List[str]:
+    sbatch_commands = [
+        "sbatch {}".format(remote_runsub_path)
+        for remote_runsub_path in remote_runsub_paths
+    ]
+    sbatch_commands = " ; ".join(sbatch_commands)
+
+    ssh_command = ["ssh"]
+    if socket is not None:
+        ssh_command.append(f"-S {socket}")
+    ssh_command.append(f"{username}@{hostname}")
+    ssh_command.append(sbatch_commands)
+    ssh_command = " ".join(ssh_command)
+
+    completed_process = subprocess.run(
+        args=shlex.split(ssh_command), capture_output=True
+    )
+    if completed_process.returncode != 0:
+        raise RuntimeError(
+            "failed to submit sbatch scripts for execution\n{}".format(
+                completed_process.stderr.decode("utf-8")
+            )
+        )
+
+    sbatch_output = completed_process.stdout.decode("utf-8")
+    slurm_job_ids = re.findall(r"(?<=Submitted batch job )\d+", sbatch_output)
+    return slurm_job_ids
+
+
+def _query_slurm_jobs_status(
+    slurm_job_ids: List[str],
+    username: str,
+    hostname: str,
+    socket: str | None,
+) -> Dict[str, str]:
+    """Query SLURM for job statuses using sacct command.
+
+    Args:
+        slurm_job_ids: List of SLURM job IDs to query.
+        username: SSH username.
+        hostname: SSH hostname.
+        socket: control socket location or None
+
+    Returns:
+        Dict mapping from slurm_job_id to returned slurm status.
+    """
+    if len(slurm_job_ids) == 0:
+        return {}
+    sacct_command = "sacct -j {} --format='JobID,State%32' --noheader -P".format(
+        ",".join(slurm_job_ids)
+    )
+    ssh_command = ["ssh"]
+    if socket is not None:
+        ssh_command.append(f"-S {socket}")
+    ssh_command.append(f"{username}@{hostname}")
+    ssh_command.append(sacct_command)
+    ssh_command = " ".join(ssh_command)
+    completed_process = subprocess.run(
+        args=shlex.split(ssh_command), capture_output=True
+    )
+    if completed_process.returncode != 0:
+        raise RuntimeError(
+            "failed to query slurm job status\n{}".format(
+                completed_process.stderr.decode("utf-8")
+            )
+        )
+    sacct_output = completed_process.stdout.decode("utf-8")
+    sacct_output_lines = sacct_output.strip().split("\n")
+    slurm_jobs_status = {}
+    for slurm_job_id in slurm_job_ids:
+        slurm_job_status = _parse_slurm_job_status(slurm_job_id, sacct_output_lines)
+        slurm_jobs_status[slurm_job_id] = slurm_job_status
+    return slurm_jobs_status
+
+
+def _kill_slurm_job(
+    slurm_job_ids: List[str], username: str, hostname: str, socket: str | None
+) -> None:
+    """Kill a SLURM job.
+
+    Args:
+        slurm_job_ids: List of SLURM job IDs to kill.
+        username: SSH username.
+        hostname: SSH hostname.
+        socket: control socket location or None
+    """
+    if len(slurm_job_ids) == 0:
+        return {}
+    kill_command = "scancel {}".format(",".join(slurm_job_ids))
+    ssh_command = ["ssh"]
+    if socket is not None:
+        ssh_command.append(f"-S {socket}")
+    ssh_command.append(f"{username}@{hostname}")
+    ssh_command.append(kill_command)
+    ssh_command = " ".join(ssh_command)
+    completed_process = subprocess.run(
+        args=shlex.split(ssh_command), capture_output=True
+    )
+    if completed_process.returncode != 0:
+        raise RuntimeError(
+            "failed to kill slurm job\n{}".format(
+                completed_process.stderr.decode("utf-8")
+            )
+        )
+    return completed_process
+
+
+def _parse_slurm_job_status(slurm_job_id: str, sacct_output_lines: List[str]) -> str:
+    """Parse SLURM job status from sacct output for a specific job.
+
+    Args:
+        slurm_job_id: The SLURM job ID to parse.
+        sacct_output_lines: Lines from sacct output.
+
+    Returns:
+        SLURM status string.
+    """
+    for line in sacct_output_lines:
+        if line.startswith(f"{slurm_job_id}|"):
+            state = line.split("|")[1]
+            state = state.strip()
+            if state:
+                state_split = state.split()
+                if len(state_split) > 0:
+                    return state_split[0]
+    return "UNKNOWN"
+
+
+def _read_autoresumed_slurm_job_ids(
+    slurm_job_ids: List[str],
+    remote_rundir_paths: List[Path],
+    username: str,
+    hostname: str,
+    socket: str | None,
+) -> Dict[str, List[str]]:
+    assert len(slurm_job_ids) == len(remote_rundir_paths)
+    slurm_job_id_list_paths = [
+        str(remote_rundir_path / ".slurm_job_id.list")
+        for remote_rundir_path in remote_rundir_paths
+    ]
+    slurm_job_id_list_strs = _read_files_from_remote(
+        slurm_job_id_list_paths, username, hostname, socket
+    )
+    assert len(slurm_job_id_list_strs) == len(slurm_job_ids)
+    autoresumed_slurm_job_ids = {}
+    for i, slurm_job_id_list_str in enumerate(slurm_job_id_list_strs):
+        slurm_job_id = slurm_job_ids[i]
+        slurm_job_id_list = slurm_job_id_list_str.split()
+        autoresumed_slurm_job_ids[slurm_job_id] = slurm_job_id_list
+    return autoresumed_slurm_job_ids
+
+
+def _read_files_from_remote(
+    filepaths: List[Path],
+    username: str,
+    hostname: str,
+    socket: str | None,
+) -> List[str]:
+    cat_commands = [
+        "echo _START_OF_FILE_ ; cat {} 2>/dev/null ; echo _END_OF_FILE_ ".format(
+            filepath
+        )
+        for filepath in filepaths
+    ]
+    cat_commands = " ; ".join(cat_commands)
+    ssh_command = ["ssh"]
+    if socket is not None:
+        ssh_command.append(f"-S {socket}")
+    ssh_command.append(f"{username}@{hostname}")
+    ssh_command.append(cat_commands)
+    ssh_command = " ".join(ssh_command)
+    completed_process = subprocess.run(
+        args=shlex.split(ssh_command), capture_output=True
+    )
+    if completed_process.returncode != 0:
+        raise RuntimeError(
+            "failed to read files from remote\n{}".format(
+                completed_process.stderr.decode("utf-8")
+            )
+        )
+    cat_outputs = completed_process.stdout.decode("utf-8")
+    cat_outputs = cat_outputs.replace("\n", " ")
+    matches = re.findall(r"(?<=_START_OF_FILE_)(.*?)(?=_END_OF_FILE_)", cat_outputs)
+    outputs = [match.strip() for match in matches]
+    return outputs
+
+
+def _get_progress(
+    remote_rundir_paths: List[Path],
+    username: str,
+    hostname: str,
+    socket: str | None,
+) -> List[Optional[float]]:
+    remote_progress_paths = [
+        remote_rundir_path / "artifacts" / "progress"
+        for remote_rundir_path in remote_rundir_paths
+    ]
+    remote_run_config_paths = [
+        remote_rundir_path / "artifacts" / "run_config.yml"
+        for remote_rundir_path in remote_rundir_paths
+    ]
+    progress_strs = _read_files_from_remote(
+        remote_progress_paths, username, hostname, socket
+    )
+    if any(map(bool, progress_strs)):
+        run_config_strs = _read_files_from_remote(
+            remote_run_config_paths, username, hostname, socket
+        )
+    else:
+        run_config_strs = [""] * len(progress_strs)
+    progress_list = []
+    for progress_str, run_config_str in zip(progress_strs, run_config_strs):
+        if not progress_str or not run_config_str:
+            progress_list.append(None)
+            continue
+        run_config = yaml.safe_load(run_config_str)
+        dataset_size = get_eval_factory_dataset_size_from_run_config(run_config)
+        if dataset_size is not None:
+            progress = int(progress_str) / dataset_size
+        else:
+            progress = int(progress_str)
+        progress_list.append(progress)
+    return progress_list
+
+
+_AUTORESUME_HANDLER = """
+_this_script=$0
+_prev_slurm_job_id=$1
+# Handle automatic resumption after some failed state.
+if [[ "$_prev_slurm_job_id" != "" ]]; then
+    _prev_state=`sacct -j $_prev_slurm_job_id -P -n -o State | head -n 1`
+    _prev_info="previous SLURM_JOB_ID $_prev_slurm_job_id finished with '$_prev_state' state."
+    if [[ $_prev_state == 'TIMEOUT' || $_prev_state == 'PREEMPTED' || $_prev_state == 'NODE_FAIL' ]]; then
+        echo "$_prev_info RESUMING..."
+    else
+        echo "$_prev_info EXIT!"
+        if [[ $_prev_state == 'COMPLETED' ]]; then
+            exit 0
+        else
+            exit 1
+        fi
+    fi
+fi
+# Schedule next execution of this script  with the current $SLURM_JOB_ID as an argument.
+# "afternotok" means next execution will be invoked only if the current execution terminates in some failed state.
+sbatch --dependency=afternotok:$SLURM_JOB_ID $_this_script $SLURM_JOB_ID
+""".strip()
+
+
+_WAIT_FOR_SERVER_HANDLER = """
+date
+# wait for the server to initialize
+bash -c 'while [[ "$(curl -s -o /dev/null -w "%{{http_code}}" {health_url})" != "200" ]]; do kill -0 '"$SERVER_PID"' 2>/dev/null || {{ echo "Server process '"$SERVER_PID"' died"; exit 1; }}; sleep 5; done'
+date
+""".strip()