Swarm and configuration#

DomynLLMSwarm#

The context manager that brings an endpoint up, submits work to it, and tears it down.

class domyn_swarm.core.swarm.DomynLLMSwarm(*, cfg, name=<factory>, endpoint=None, delete_on_exit=False, serving_handle=None, preflight=True, swarm_dir=<factory>, watchdog_db_path=<factory>)[source]#

Bases: BaseModel

Context manager orchestrating a distributed LLM serving swarm.

Provides a unified interface for deploying, managing and interacting with a large language model serving cluster on a compute backend such as Slurm or Lepton. It handles the whole lifecycle from resource allocation to cleanup, with state persistence and job submission built in.

A swarm consists of a load balancer (nginx) distributing requests, several vLLM server instances serving the model, a head node coordinating jobs and running user scripts, and persisted state enabling recovery and reconnection.

On Slurm the swarm is deployed as a job array with roles assigned by SLURM_NODEID: node 0 runs the load balancer and the user driver, nodes 1 to N run vLLM servers. Cloud platforms are reached through the deployment abstractions instead.

State is persisted automatically - deployment metadata and resource handles, configuration, platform identifiers such as job IDs and node assignments, and endpoint URLs - which is what allows a swarm to be recovered after a failure or reattached from another process via from_state.

Variables:

  • cfg (domyn_swarm.config.swarm.DomynLLMSwarmConfig) – Deployment parameters, resource requirements and platform-specific settings.

  • endpoint (str | None) – Public URL of the deployed load balancer. Set once deployment succeeds through the context manager.

  • delete_on_exit (bool | None) – Whether to clean up all allocated resources when leaving the context manager. Useful for temporary deployments.

  • serving_handle (domyn_swarm.platform.protocols.ServingHandle | None) – Platform-specific handle for the serving deployment, carrying metadata such as job IDs, node assignments and status.

  • model – Name or path of the model being served. May be set after initialization to switch models.

  • preflight (bool) – Whether to check the config’s filesystem paths before deploying. Off for sites where they are visible only from the compute nodes.

Raises:
Parameters:

  • cfg (DomynLLMSwarmConfig)

  • name (str)

  • endpoint (str | None)

  • delete_on_exit (bool | None)

  • serving_handle (ServingHandle | None)

  • preflight (bool)

  • swarm_dir (EnvPath)

  • watchdog_db_path (EnvPath)

Note

The swarm must be used as a context manager for resources to be managed correctly. Startup waits up to cfg.wait_endpoint_s for the endpoint. Paths passed to job submission resolve relative to the execution environment, ENDPOINT and MODEL are set automatically for submitted jobs, and checkpoint directories are created as needed. Cleanup failures are logged but do not prevent the context from exiting.

Example

Basic deployment, cleaned up on exit:

cfg = DomynLLMSwarmConfig.read("config.yaml")
with DomynLLMSwarm(cfg=cfg) as swarm:
    # Reachable at swarm.endpoint
    swarm.submit_job(
        my_job,
        run=JobRunSpec(
            input_path=Path("data.parquet"),
            output_path=Path("results.parquet"),
        ),
    )

Persistent deployment, reattached later from another process:

swarm = DomynLLMSwarm(cfg=cfg, delete_on_exit=False)
with swarm:
    pass
# Resources remain allocated
swarm = DomynLLMSwarm.from_state("my-deployment-abc123")

Detached job submission:

with DomynLLMSwarm(cfg=cfg) as swarm:
    handle = swarm.submit_job(
        job=LongRunningJob(),
        run=JobRunSpec(
            input_path=Path("large_dataset.parquet"),
            output_path=Path("results.parquet"),
            detach=True,
        ),
    )
    print(handle.pid, handle.external_id)

Running a script on the head node:

with DomynLLMSwarm(cfg=cfg) as swarm:
    swarm.submit_script(Path("analysis.py"), extra_args=["--mode", "evaluation"])

See also

SwarmJob: Base class for jobs executable within the swarm. DomynLLMSwarmConfig: Configuration schema and validation.

cfg: DomynLLMSwarmConfig#
name: str#
endpoint: str | None#
delete_on_exit: bool | None#
serving_handle: ServingHandle | None#
preflight: bool#
swarm_dir: utils.EnvPath#
watchdog_db_path: utils.EnvPath#
property model: str#

The model name, either from the config or the job submission. If not set, defaults to the config’s model.

property platform: str#

The platform this swarm deploys to, e.g. "slurm" or "lepton".

Derived from cfg.backend.type, which is the authoritative value and is present in every persisted record.

model_post_init(_DomynLLMSwarm__context)[source]#

Wire up in-memory collaborators. Performs no I/O.

Parameters:

_DomynLLMSwarm__context (Any)

Return type:

None

ensure_directories()[source]#

Create this swarm’s on-disk layout. Idempotent.

Called before deploying (the sbatch templates reference these paths and Slurm rejects a submission whose output directory is missing) and before writing checkpoints on a swarm reattached via from_state(), which does not pass through the context manager.

Return type:

None

classmethod from_state(deployment_name)[source]#

Initialize a swarm from a saved state.

Parameters:

deployment_name (str) – Deployment name.

Returns:

Loaded swarm.

Return type:

DomynLLMSwarm

classmethod from_record(swarm_payload, cfg_payload, handle_payload)[source]#

Build a complete swarm from persisted state.

Parameters:

  • swarm_payload (dict) – The persisted swarm column, without cfg.

  • cfg_payload (dict) – The persisted cfg column.

  • handle_payload (dict) – The persisted serving_handle column.

Returns:

A swarm with its handle adopted and its compute backend built.

Raises:

ValueError – If the record’s serving handle is missing metadata the platform needs; see validate_serving_handle().

Return type:

DomynLLMSwarm

submit_job(job, *, run)[source]#

Launch a serialized SwarmJob inside the current SLURM swarm allocation.

The job object is converted to keyword arguments via SwarmJob.to_kwargs(), transmitted to the head node (where SLURM_NODEID == 0), reconstructed by domyn_swarm.jobs.cli.run, and executed under srun.

Parameters:

  • job (SwarmJob) – The job instance to execute.

  • run (JobRunSpec) – How the job should run: input/output paths, sharding, checkpointing, engine selection, and resources. See JobRunSpec for the full field list and their semantics.

Returns:

Compute job handle with normalized status and metadata.

Return type:

JobHandle

Raises:

Note

The constructed command is logged with rich for transparency, e.g.:

srun --jobid=<...> --nodelist=<...> --ntasks=1 --overlap ...
    python -m domyn_swarm.jobs.cli.run --job-class=<module:Class> ...

Example

swarm.submit_job(
    my_job,
    run=JobRunSpec(
        input_path=Path("batch.parquet"),
        output_path=Path("predictions.parquet"),
        num_shards=4,
    ),
)
submit_script(script_path, detach=False, extra_args=None, job_resources=None)[source]#

Submit a Python script to the compute backend for execution.

This method validates the script path, composes the runtime environment, and submits the script for execution via the configured deployment backend.

Parameters:

  • script_path (Path) – Path to the Python script to be executed.

  • detach (bool, optional) – If True, run the script in detached mode. If False, run synchronously. Defaults to False.

  • extra_args (list[str] | None, optional) – Additional command-line arguments to pass to the script. Defaults to None.

  • job_resources (dict | None) – Scheduler resource overrides for the job step. They take precedence over the compute backend and plan defaults.

Returns:

Submitted job handle with normalized status and metadata.

Return type:

JobHandle

Raises:

  • FileNotFoundError – If the script file does not exist (only checked for SLURM platform).

  • RuntimeError – If the script submission to the compute backend fails.

Example

>>> swarm = Swarm(...)
>>> # Submit script synchronously
>>> swarm.submit_script(Path("my_script.py"))

>>> # Submit script in detached mode with arguments
>>> handle = swarm.submit_script(
...     Path("my_script.py"), detach=True, extra_args=["--config", "config.yaml"]
... )
>>> handle.pid
wait_job(handle, *, stream_logs=True)[source]#

Wait for a submitted compute job to reach a terminal state.

Parameters:

  • handle (JobHandle) – Job handle to wait on.

  • stream_logs (bool) – Whether to stream backend logs while waiting.

Returns:

Normalized terminal job status.

Return type:

JobStatus

cancel_job(handle)[source]#

Cancel a submitted compute job.

Parameters:

handle (JobHandle) – Job handle to cancel.

Returns:

Final normalized status after cancellation.

Return type:

JobStatus

refresh_job_status(job_id)[source]#

Refresh a persisted job status via backend probe (best effort).

Parameters:

job_id (str) – Internal persisted job identifier.

Returns:

refresh_source and refresh_error.

Return type:

Updated job record payload, including transient refresh metadata

cleanup()[source]#

Release the swarm’s resources and drop its state record.

Idempotent: calling it again after a successful teardown is a no-op, so overlapping paths (a failed startup, __exit__, an explicit down()) cannot issue a second cancel against the platform.

The record is dropped only once teardown succeeds – if down raises, the record survives so the deployment stays reachable for a retry.

down()[source]#

Manually clean up the swarm resources.

status()[source]#

Get the current status of the swarm.

Safe to call at any point in the lifecycle: a swarm that has not been deployed reports ServingPhase.UNKNOWN rather than raising.

Returns:

The backend-reported serving status, falling back to the endpoint this swarm already knows about when the backend reports none.

Return type:

ServingStatus

model_config = {}#

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

serving_status()[source]#

Ask the serving backend for this swarm’s live status.

Unlike status(), which reports what this object already knows, this queries the backend directly. Safe at any point in the lifecycle: a swarm that was never deployed reports UNKNOWN.

Returns:

The backend-reported serving status.

Return type:

ServingStatus

DomynLLMSwarmConfig#

For the full field-by-field table, see Configuration.

class domyn_swarm.config.swarm.DomynLLMSwarmConfig(*, model, name, revision=None, replicas=1, gpus_per_replica=4, gpus_per_node=4, replicas_per_node=None, nodes=None, cpus_per_task=None, mem_per_cpu=None, wait_endpoint_s=1200, image=<factory>, args='', port=8000, home_directory=<factory>, backend, env=None, watchdog=<factory>)[source]#

Bases: BaseModel

Validated configuration for one serving swarm.

Field descriptions are the canonical user-facing reference for individual settings. Use read and persist for YAML-backed configuration and build_plan to resolve the configured serving and compute backends.

Parameters:

  • model (str)

  • name (Annotated[str, StringConstraints(strip_whitespace=True, to_upper=None, to_lower=True, strict=None, min_length=None, max_length=38, pattern=None, ascii_only=None)])

  • revision (str | None)

  • replicas (int)

  • gpus_per_replica (int)

  • gpus_per_node (Annotated[int, Ge(ge=1), Le(le=4)])

  • replicas_per_node (int | None)

  • nodes (int | None)

  • cpus_per_task (Annotated[int | None, Ge(ge=1)])

  • mem_per_cpu (str | None)

  • wait_endpoint_s (int)

  • image (str | EnvPath)

  • args (str)

  • port (int)

  • home_directory (EnvPath)

  • backend (Annotated[LeptonConfig | SlurmConfig, FieldInfo(annotation=NoneType, required=True, discriminator='type')] | None)

  • env (dict[str, str] | None)

  • watchdog (WatchdogConfig)

model: str#
name: Annotated[str, StringConstraints(strip_whitespace=True, to_upper=None, to_lower=True, strict=None, min_length=None, max_length=38, pattern=None, ascii_only=None)]#
revision: str | None#
replicas: int#
gpus_per_replica: int#
gpus_per_node: int#
replicas_per_node: int | None#
nodes: int | None#
cpus_per_task: int | None#
mem_per_cpu: str | None#
wait_endpoint_s: int#
image: str | EnvPath#
args: str#
port: int#
home_directory: EnvPath#
backend: Annotated[LeptonConfig | SlurmConfig, FieldInfo(annotation=NoneType, required=True, discriminator='type')] | None#
env: dict[str, str] | None#
watchdog: WatchdogConfig#
get_deployment_plan()[source]#
Return type:

DeploymentPlan | None

build_plan()[source]#
Return type:

DeploymentPlan

classmethod not_empty(v)[source]#
Parameters:

v (Annotated[LeptonConfig | SlurmConfig, FieldInfo(annotation=NoneType, required=True, discriminator='type')] | None)

Return type:

Annotated[LeptonConfig | SlurmConfig, FieldInfo(annotation=NoneType, required=True, discriminator=’type’)]

classmethod read(path)[source]#
Parameters:

path (str)

Return type:

DomynLLMSwarmConfig

persist(path)[source]#
Parameters:

path (str | Path)

Return type:

None

classmethod validate_resource_allocations(data)[source]#

Validate and auto-compute all derived resource allocation fields.

Parameters:

data (Any)

Return type:

DomynLLMSwarmConfig

model_post_init(context, /)#

This function is meant to behave like a BaseModel method to initialize private attributes.

It takes context as an argument since that’s what pydantic-core passes when calling it.

Parameters:

  • self (BaseModel) – The BaseModel instance.

  • context (Any) – The context.

Return type:

None

On this page
Show Source