Slurm

This contains the TorchX Slurm scheduler which can be used to run TorchX components on a Slurm cluster.

class torchx.schedulers.slurm_scheduler.SlurmScheduler(session_name: str, docker_client: DockerClient | None = None)

Bases: MultiWorkspaceMixin, Scheduler[SlurmOpts]

SlurmScheduler is a TorchX scheduling interface to slurm. TorchX expects that slurm CLI tools are locally installed and job accounting is enabled.

Each app def is scheduled using a heterogenous job via sbatch. Each replica of each role has a unique shell script generated with it’s resource allocations and args and then sbatch is used to launch all of them together.

Logs are available in combined form via torchx log, the programmatic API as well as in the job launch directory as slurm-<jobid>-<role>-<replica_id>.out. If TorchX is running in a different directory than where the job was created the logs won’t be able to be found.

Some of the config options passed to it are added as SBATCH arguments to each replica. See https://slurm.schedmd.com/sbatch.html#SECTION_OPTIONS for info on the arguments.

Slurm jobs inherit the currently active conda or virtualenv and run in the current working directory. This matches the behavior of the local_cwd scheduler.

For more info see:

• https://slurm.schedmd.com/sbatch.html

• https://slurm.schedmd.com/heterogeneous_jobs.html

$ torchx run --scheduler slurm utils.echo --msg hello
slurm://torchx_user/1234
$ torchx status slurm://torchx_user/1234
$ less slurm-1234.out
...

Config Options

    usage:
        [account=ACCOUNT],[partition=PARTITION],[time=TIME],[comment=COMMENT],[constraint=CONSTRAINT],[mail-user=MAIL-USER],[mail-type=MAIL-TYPE],[job_dir=JOB_DIR],[qos=QOS],[workspace_type=WORKSPACE_TYPE],[image_repo=IMAGE_REPO],[quiet=QUIET]

    optional arguments:
        account=ACCOUNT (str, None)
            The account to use for the slurm job.
        partition=PARTITION (str, None)
            The partition to run the job in.
        time=TIME (str, None)
            The maximum time the job is allowed to run for. Formats: "minutes",
    "minutes:seconds", "hours:minutes:seconds", "days-hours",
    "days-hours:minutes" or "days-hours:minutes:seconds"
        comment=COMMENT (str, None)
            Comment to set on the slurm job.
        constraint=CONSTRAINT (str, None)
            Constraint to use for the slurm job.
        mail-user=MAIL-USER (str, None)
            User to mail on job end.
        mail-type=MAIL-TYPE (str, None)
            What events to mail users on.
        job_dir=JOB_DIR (str, None)
            The directory to place the job script, logs and, with the ``dir``
    workspace builder (the default), the workspace copy; that builder
    requires it not to exist yet, any other case creates it if missing.
    To enable log iteration, jobs will be tracked in
    ``.torchxslurmjobdirs``.
        qos=QOS (str, None)
            Quality of Service (QoS) to assign to the job.
        workspace_type=WORKSPACE_TYPE (str, dir)
            which workspace builder patches the image: one of dir, docker
        image_repo=IMAGE_REPO (str, None)
            (remote jobs) the image repository to use when pushing patched images, must have push access. Ex: example.com/your/container
        quiet=QUIET (bool, False)
            whether to suppress verbose output for image building. Defaults to ``False``.

Compatibility

Feature	Scheduler Support
Fetch Logs	✔️
Distributed Jobs	✔️
Cancel Job	✔️
Describe Job	Partial support. SlurmScheduler will return job and replica status but does not provide the complete original AppSpec.
Workspaces / Patching	Two builders, picked with the ``workspace_type`` run option. ``dir`` (the default) copies the workspace into ``job_dir`` when that option is set. ``docker`` builds a patched image, pushes it to ``image_repo`` and runs every replica through the cluster’s container plugin (``srun –container-image``).
Mounts	❌
Elasticity	❌

If a partition has less than 1GB of RealMemory configured we disable memory requests to workaround https://github.com/aws/aws-parallelcluster/issues/2198.

describe(app_id: str) → DescribeAppResponse | None

Returns app description, or None if it no longer exists.

list(cfg: Mapping[str, str | int | float | bool | list[str] | dict[str, str] | None] | None = None) → List[ListAppResponse]

Lists jobs on this scheduler.

log_iter(app_id: str, role_name: str, k: int = 0, regex: str | None = None, since: datetime | None = None, until: datetime | None = None, should_tail: bool = False, streams: Stream | None = None) → Iterable[str]

Returns an iterator over log lines for the k-th replica of role_name.

Important

Not all schedulers support log iteration, tailing, or time-based cursors. Check the specific scheduler docs.

Lines include trailing whitespace (\n). When should_tail=True, the iterator blocks until the app reaches a terminal state.

Parameters:

• k – replica (node) index

• regex – optional filter pattern

• since – start cursor (scheduler-dependent)

• until – end cursor (scheduler-dependent)

• should_tail – if True, follow output like tail -f

• streams – stdout, stderr, or combined

Raises:

NotImplementedError – if the scheduler does not support log iteration

schedule(dryrun_info: AppDryRunInfo[SlurmBatchRequest, SlurmOpts]) → str

Submits a previously dry-run request. Returns the app_id.

workspace_builders() → Mapping[str, WorkspaceMixin[Any]]

Returns {name: builder}. The first entry is the default.

Return the same builder instances on every call (build the mapping once, in __init__): the push step looks the builder up again by name and expects the instance that ran the build.

torchx.schedulers.slurm_scheduler.create_scheduler(session_name: str, docker_client: DockerClient | None = None, **kwargs: Any) → SlurmScheduler

class torchx.schedulers.slurm_scheduler.SlurmBatchRequest(cmd: list[str], replicas: dict[str, SlurmReplicaRequest], job_dir: str | None, max_retries: int, images_to_push: tuple[str, Any] | None = None)

Holds parameters used to launch a slurm job via sbatch.

images_to_push: tuple[str, Any] | None = None

the docker workspace builder’s (name, images); None when nothing was built into an image

Type:

what SlurmScheduler.schedule() pushes before sbatch

materialize() → str

materialize returns the contents of the script that can be passed to sbatch to run the job.

class torchx.schedulers.slurm_scheduler.SlurmReplicaRequest(name: str, entrypoint: str, args: list[str], srun_opts: dict[str, str], sbatch_opts: dict[str, str | None], env: dict[str, str])

Holds parameters for a single replica running on slurm and can be materialized down to a bash script.

classmethod from_role(name: str, role: Role, cfg: SlurmOpts, nomem: bool, container_image: str | None = None) → SlurmReplicaRequest

from_role creates a SlurmReplicaRequest for the specific role and name. container_image, when given, is passed to srun as --container-image so the cluster’s container plugin runs the replica inside it.

materialize() → tuple[list[str], list[str]]

materialize returns the sbatch and srun groups for this role. They should be combined using : per slurm heterogenous groups.