# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
# pyre-strict
import ipaddress
import json
import logging
import os
import shlex
import subprocess
import sys
from typing import Any, Dict, FrozenSet, List, Optional, Sequence
from monarch._rust_bindings.monarch_hyperactor.channel import ChannelTransport
from monarch._rust_bindings.monarch_hyperactor.config import configure
from monarch._rust_bindings.monarch_hyperactor.proc import ProcId
from monarch._src.actor.bootstrap import attach_to_workers
from monarch._src.job._batch_env import in_batch_job
from monarch._src.job._slurm_batch import _worker_bootstrap
from monarch._src.job.job import BatchJob, JobState, JobTrait
from monarch._src.job.service_identity import (
allocate_service_proc_ids,
serialize_service_proc_ids,
service_proc_addrs,
SERVICE_PROC_IDS_ENV,
)
logger: logging.Logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
logger.addHandler(logging.StreamHandler(sys.stderr))
logger.propagate = False
# terminal states that indicate the job is no longer active
_SLURM_TERMINAL_STATES: FrozenSet[str] = frozenset(
["FAILED", "CANCELLED", "TIMEOUT", "PREEMPTED", "COMPLETED"]
)
[docs]
class SlurmJob(JobTrait):
"""
A job scheduler that uses SLURM command line tools to schedule jobs.
This implementation:
1. Uses sbatch to submit SLURM jobs that start monarch workers
2. Queries job status with squeue to get allocated hostnames
3. Uses the hostnames to connect to the started workers
Two launch modes are supported:
- External controller (default): ``apply()`` submits the workers; a separate
controller process attaches via ``state()``.
- Batch (``apply(client_script=...)``): one allocation runs both the workers
and the client; the allocation is released when the client exits.
"""
def __init__(
self,
meshes: Dict[str, int],
python_exe: str = "python",
slurm_args: Sequence[str] = (),
monarch_port: int = 22222,
job_name: str = "monarch_job",
ntasks_per_node: int = 1,
time_limit: Optional[str] = None,
partition: Optional[str] = None,
log_dir: Optional[str] = None,
exclusive: bool = True,
gpus_per_node: Optional[int] = None,
cpus_per_task: Optional[int] = None,
mem: Optional[str] = None,
job_start_timeout: Optional[int] = None,
account: Optional[str] = None,
qos: Optional[str] = None,
out_of_cluster: bool = False,
attach_to: Optional[str] = None,
bind_to: Optional[str] = None,
) -> None:
"""
Args:
meshes: Dictionary mapping mesh names to number of nodes
python_exe: Python executable to use for worker processes
slurm_args: Additional SLURM arguments to pass to sbatch
monarch_port: Port for TCP communication between workers
job_name: Name for the SLURM job
ntasks_per_node: Number of tasks per node
time_limit: Maximum runtime in HH:MM:SS format. If None, uses SLURM's default time limit.
partition: SLURM partition to submit to
log_dir: Directory for SLURM log files
exclusive: Whether to request exclusive node access (no other jobs can run on the nodes).
Defaults to True for predictable performance and resource isolation,
but may increase queue times and waste resources if nodes are underutilized.
gpus_per_node: Number of GPUs to request per node. If None, no GPU resources are requested.
job_start_timeout: Maximum time in seconds to wait for the SLURM job to start running.
This should account for potential queueing delays. If None (default), waits indefinitely.
account: SLURM account to charge the job to (``#SBATCH --account``). If None, uses the cluster default.
qos: SLURM quality-of-service to request (``#SBATCH --qos``). If None, uses the cluster default.
out_of_cluster: Whether the client runs outside the cluster network.
The client attaches through the first allocated worker so
workers can route messages back through that gateway.
attach_to: ZMQ-style address of the worker gateway for out-of-cluster
access. When omitted in out-of-cluster mode, the first
allocated worker's Monarch address is used.
bind_to: Optional IP address on which workers listen, without a port.
When set, workers advertise their Slurm hostname while
binding this address through the ``dial_to@bind_to`` alias
format. The bind address uses ``monarch_port``.
"""
configure(default_transport=ChannelTransport.TcpWithHostname)
self._meshes = meshes
self._python_exe = python_exe
self._slurm_args = slurm_args
self._port = monarch_port
self._job_name = job_name
self._ntasks_per_node = ntasks_per_node
self._time_limit = time_limit
self._partition = partition
self._log_dir: str = log_dir if log_dir is not None else os.getcwd()
self._exclusive = exclusive
self._gpus_per_node = gpus_per_node
self._cpus_per_task = cpus_per_task
self._mem = mem
self._job_start_timeout = job_start_timeout
self._account = account
self._qos = qos
self._out_of_cluster = out_of_cluster
self._attach_to = attach_to
try:
self._bind_to = (
str(ipaddress.ip_address(bind_to)) if bind_to is not None else None
)
except ValueError as error:
raise ValueError(
"bind_to must be an IPv4 or IPv6 address without a port"
) from error
# Track the single SLURM job ID and all allocated hostnames
self._slurm_job_id: Optional[str] = None
self._all_hostnames: List[str] = []
self._service_proc_ids: list[ProcId] = []
super().__init__()
@staticmethod
def _allocate_service_proc_ids(num_nodes: int) -> list[ProcId]:
return allocate_service_proc_ids(num_nodes)
def __setstate__(self, state: Dict[str, Any]) -> None:
# In batch mode the in-allocation client rehydrates this object via
# pickle.load (job_load), which bypasses __init__. Re-apply the
# transport config here so the client's host_agent serves on a TCP
# address remote workers can dial back, instead of an abstract unix
# socket whose namespace is local to the original (head) node.
self.__dict__.update(state)
if "_bind_to" not in state:
self._bind_to = None
configure(default_transport=ChannelTransport.TcpWithHostname)
def _resolve_attach_to(self) -> str | None:
if self._attach_to is not None:
return self._attach_to
if self._out_of_cluster:
return f"tcp://{self._all_hostnames[0]}:{self._port}"
return None
def _requires_sidecar_gateway(self) -> bool:
return self._out_of_cluster or self._attach_to is not None
def _prepare_client_gateway(self) -> None:
if not self._jobs_active():
raise RuntimeError("SLURM job is no longer active")
if not self._all_hostnames:
job_id = self._resolved_job_id()
if job_id is None:
raise RuntimeError("SLURM job ID is not set")
total_nodes = sum(self._meshes.values())
self._all_hostnames = self._wait_for_job_start(
job_id, total_nodes, timeout=self._job_start_timeout
)
self._attach_client(self._resolve_attach_to())
[docs]
def add_mesh(self, name: str, num_nodes: int) -> None:
self._meshes[name] = num_nodes
def _resolved_job_id(self) -> Optional[str]:
"""The submitted job id, falling back to ``$SLURM_JOB_ID`` in batch mode
where the unpickled instance never had ``_slurm_job_id`` set. Gated on
batch mode so an external controller doesn't adopt an unrelated id from
its own surrounding SLURM job."""
if self._slurm_job_id is not None:
return self._slurm_job_id
return os.environ.get("SLURM_JOB_ID") if in_batch_job() else None
def _cleanup_log_context(self) -> dict[str, Any]:
return {
"job_type": type(self).__name__,
"job_id": self._resolved_job_id(),
}
def _create(self, client_script: Optional[str]) -> None:
"""Submit a single SLURM job for all meshes.
In batch mode (``client_script`` set) the job runs the client inside the
allocation alongside the workers, and a ``BatchJob`` wrapper is cached to
``.monarch/job_state.pkl`` so the client reconnects to this allocation
instead of submitting a new one.
"""
total_nodes = sum(self._meshes.values())
self._ensure_service_proc_ids(total_nodes)
if client_script is not None:
# Dumped before submit: the job id isn't known yet, and the client
# resolves it from $SLURM_JOB_ID anyway.
BatchJob(self).dump(".monarch/job_state.pkl")
self._slurm_job_id = self._submit_slurm_job(
total_nodes, client_script=client_script
)
def _submit_slurm_job(
self, num_nodes: int, client_script: Optional[str] = None
) -> str:
"""Submit a SLURM job for all nodes."""
self._ensure_service_proc_ids(num_nodes)
unique_job_name = f"{self._job_name}_{os.getpid()}"
# Create log directory if it doesn't exist
os.makedirs(self._log_dir, exist_ok=True)
log_path_out = os.path.join(self._log_dir, f"slurm_%j_{unique_job_name}.out")
log_path_err = os.path.join(self._log_dir, f"slurm_%j_{unique_job_name}.err")
# Build SBATCH directives
sbatch_directives = [
"#!/bin/bash",
f"#SBATCH --job-name={unique_job_name}",
f"#SBATCH --ntasks-per-node={self._ntasks_per_node}",
f"#SBATCH --nodes={num_nodes}",
f"#SBATCH --output={log_path_out}",
f"#SBATCH --error={log_path_err}",
]
if self._time_limit is not None:
sbatch_directives.append(f"#SBATCH --time={self._time_limit}")
if self._gpus_per_node is not None:
sbatch_directives.append(f"#SBATCH --gpus-per-node={self._gpus_per_node}")
if self._cpus_per_task is not None:
sbatch_directives.append(f"#SBATCH --cpus-per-task={self._cpus_per_task}")
if self._mem is not None:
sbatch_directives.append(f"#SBATCH --mem={self._mem}")
if self._exclusive:
sbatch_directives.append("#SBATCH --exclusive")
if self._partition is not None:
sbatch_directives.append(f"#SBATCH --partition={self._partition}")
if self._account is not None:
sbatch_directives.append(f"#SBATCH --account={self._account}")
if self._qos is not None:
sbatch_directives.append(f"#SBATCH --qos={self._qos}")
if (
not self._exclusive
and self._partition is not None
and self._gpus_per_node is not None
):
gpus_per_task = self._gpus_per_node // self._ntasks_per_node
assert self._partition, (
"Slurm partition must be set for jobs that share nodes with other jobs"
)
self.share_node(
tasks_per_node=self._ntasks_per_node,
gpus_per_task=gpus_per_task,
partition=self._partition,
)
# Add any additional slurm args as directives
for arg in self._slurm_args:
if arg.startswith("-"):
sbatch_directives.append(f"#SBATCH {arg}")
batch_script = "\n".join(sbatch_directives)
batch_script += (
f"\nexport {SERVICE_PROC_IDS_ENV}="
f"{shlex.quote(serialize_service_proc_ids(self._service_proc_ids))}\n"
)
if client_script is None:
# Workers only; an external controller attaches and manages the
# lifetime. Shares worker bootstrap generation with the batch runner.
worker_cmd = _worker_bootstrap(self._port, self._bind_to)
batch_script += f"\nsrun {self._python_exe} -c '{worker_cmd}'\n"
else:
# Batch mode: the in-allocation runner seeds the workers, runs the
# client (MONARCH_BATCH_JOB=1 so its cached BatchJob reconnects to
# this allocation), and tears the workers down in a finally.
bind_to_arg = (
f" --bind-to {shlex.quote(self._bind_to)}"
if self._bind_to is not None
else ""
)
batch_script += (
f"\n{self._python_exe} -m monarch._src.job._slurm_batch "
f"--port {self._port}{bind_to_arg} {shlex.quote(client_script)}\n"
)
logger.info(f"Submitting SLURM job with {num_nodes} nodes")
try:
result = subprocess.run(
["sbatch"],
input=batch_script,
capture_output=True,
text=True,
check=True,
)
# Parse the job ID from sbatch output (typically "Submitted batch job 12345")
job_id = None
for line in result.stdout.strip().split("\n"):
if "Submitted batch job" in line:
job_id = line.split()[-1]
break
if not job_id:
raise RuntimeError(
f"Failed to parse job ID from sbatch output: {result.stdout}"
)
logger.info(
f"SLURM job {job_id} submitted. Logs will be written to: {self._log_dir}/slurm_{job_id}_{unique_job_name}.out"
)
return job_id
except subprocess.CalledProcessError as e:
raise RuntimeError(f"Failed to submit SLURM job: {e.stderr}") from e
def _ensure_service_proc_ids(self, num_nodes: int) -> None:
if len(self._service_proc_ids) != num_nodes:
self._service_proc_ids = [
proc_id
for mesh_nodes in self._meshes.values()
for proc_id in self._allocate_service_proc_ids(mesh_nodes)
]
def _get_job_info_json(self, job_id: str) -> Optional[Dict[str, Any]]:
"""Get job information using squeue --json."""
try:
result = subprocess.run(
["squeue", "--job", job_id, "--json"],
capture_output=True,
text=True,
check=True,
)
if result.stdout.strip():
data = json.loads(result.stdout)
jobs = data.get("jobs", [])
return jobs[0] if jobs else None
return None
except subprocess.CalledProcessError as e:
logger.warning(f"Error checking job {job_id} status: {e.stderr}")
return None
except (json.JSONDecodeError, KeyError) as e:
logger.warning(f"Error parsing JSON response for job {job_id}: {e}")
return None
def _wait_for_job_start(
self, job_id: str, expected_nodes: int, timeout: Optional[int] = None
) -> List[str]:
"""
Wait for the SLURM job to start and return the allocated hostnames.
Requires Slurm 20.02+ for squeue --json support.
"""
import time
start_time = time.time()
try:
while timeout is None or time.time() - start_time < timeout:
job_info = self._get_job_info_json(job_id)
if not job_info:
raise RuntimeError(f"SLURM job {job_id} not found in queue")
job_state = job_info.get("job_state", [])
if "RUNNING" in job_state:
# Extract hostnames from job_resources.nodes.allocation
job_resources = job_info.get("job_resources", {})
nodes_info = job_resources.get("nodes", {})
allocation = nodes_info.get("allocation", [])
hostnames = [node["name"] for node in allocation]
logger.info(
f"SLURM job {job_id} is running on {len(hostnames)} nodes: {hostnames}"
)
if len(hostnames) != expected_nodes:
raise RuntimeError(
f"Expected {expected_nodes} nodes but got {len(hostnames)}. "
f"Partial allocation not supported."
)
return hostnames
elif any(state in job_state for state in _SLURM_TERMINAL_STATES):
raise RuntimeError(
f"SLURM job {job_id} failed with status: {job_state}"
)
else:
logger.debug(f"SLURM job {job_id} status: {job_state}, waiting...")
time.sleep(2) # Check every 2 seconds
raise RuntimeError(f"Timeout waiting for SLURM job {job_id} to start")
except Exception:
# Cleanup on failure - reuse _kill() logic
logger.error(f"Failed to start SLURM job {job_id}, cancelling job")
self._kill()
raise
def _state(self) -> JobState:
self._prepare_client_gateway()
if not self._jobs_active():
raise RuntimeError("SLURM job is no longer active")
# Wait for job to start and get hostnames if not already done
if not self._all_hostnames:
job_id = self._resolved_job_id()
if job_id is None:
raise RuntimeError("SLURM job ID is not set")
total_nodes = sum(self._meshes.values())
self._all_hostnames = self._wait_for_job_start(
job_id, total_nodes, timeout=self._job_start_timeout
)
attach_to = self._attach_to
if self._out_of_cluster and attach_to is None:
attach_to = f"tcp://{self._all_hostnames[0]}:{self._port}"
self._attach_client(attach_to)
# Distribute the allocated hostnames among meshes
host_meshes = {}
hostname_idx = 0
for mesh_name, num_nodes in self._meshes.items():
next_hostname_idx = hostname_idx + num_nodes
mesh_hostnames = self._all_hostnames[hostname_idx:next_hostname_idx]
service_proc_ids = self._service_proc_ids[hostname_idx:next_hostname_idx]
hostname_idx = next_hostname_idx
workers = [f"tcp://{hostname}:{self._port}" for hostname in mesh_hostnames]
host_mesh = attach_to_workers(
name=mesh_name,
ca="trust_all_connections",
workers=service_proc_addrs(workers, service_proc_ids),
)
host_meshes[mesh_name] = host_mesh
return JobState(host_meshes)
[docs]
def can_run(self, spec: "JobTrait") -> bool:
"""Check if this job can run the given spec."""
return (
isinstance(spec, SlurmJob)
and spec._meshes == self._meshes
and spec._python_exe == self._python_exe
and spec._port == self._port
and spec._slurm_args == self._slurm_args
and spec._job_name == self._job_name
and spec._ntasks_per_node == self._ntasks_per_node
and spec._time_limit == self._time_limit
and spec._partition == self._partition
and spec._gpus_per_node == self._gpus_per_node
and spec._cpus_per_task == self._cpus_per_task
and spec._mem == self._mem
and spec._job_start_timeout == self._job_start_timeout
and spec._account == self._account
and spec._qos == self._qos
and spec._out_of_cluster == self._out_of_cluster
and spec._attach_to == self._attach_to
and spec._bind_to == self._bind_to
and self._jobs_active()
)
def _jobs_active(self) -> bool:
"""Check if SLURM job is still active by querying squeue."""
# A batch job reloaded inside its own allocation pickles as not_running,
# so fall back to $SLURM_JOB_ID + squeue instead of the active flag.
if not self.active and not in_batch_job():
return False
job_id = self._resolved_job_id()
if job_id is None:
return False
job_info = self._get_job_info_json(job_id)
if not job_info:
logger.warning(f"SLURM job {job_id} not found in queue")
return False
job_state = job_info.get("job_state", [])
if any(state in job_state for state in _SLURM_TERMINAL_STATES):
logger.warning(f"SLURM job {job_id} has status: {job_state}")
return False
return True
[docs]
def share_node(
self, tasks_per_node: int, gpus_per_task: int, partition: str
) -> None:
"""
Share a node with other jobs.
"""
try:
import clusterscope
except ImportError:
raise RuntimeError(
"please install clusterscope to use share_node. `pip install clusterscope`"
)
self._exclusive = False
slurm_args = clusterscope.job_gen_task_slurm(
partition=partition,
gpus_per_task=gpus_per_task,
tasks_per_node=tasks_per_node,
)
self._cpus_per_task = slurm_args["cpus_per_task"]
self._mem = slurm_args["memory"]
def _kill(self) -> None:
"""Cancel the SLURM job.
No-op in batch mode: ``BatchJob`` registers this as an ``atexit`` hook on
the in-allocation client, but the sbatch runner owns teardown there, so
the client must not scancel its own allocation. The external-client
path scancels as before.
"""
if in_batch_job():
return
job_id = self._resolved_job_id()
if job_id is not None:
try:
subprocess.run(
["scancel", job_id],
capture_output=True,
text=True,
check=True,
)
logger.info(f"Cancelled SLURM job {job_id}")
except subprocess.CalledProcessError as e:
logger.warning(f"Failed to cancel SLURM job {job_id}: {e.stderr}")
self._slurm_job_id = None
self._all_hostnames.clear()
