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, 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.

Raises:
Parameters:

  • cfg (DomynLLMSwarmConfig)

  • name (str)

  • endpoint (str | None)

  • delete_on_exit (bool | None)

  • serving_handle (ServingHandle | None)

  • 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, input_path="data.parquet", output_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(),
        input_path="large_dataset.parquet",
        output_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#
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.

model_post_init(_DomynLLMSwarm__context)[source]#

Post-init to set up the deployment backend.

Parameters:

_DomynLLMSwarm__context (Any)

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

submit_job(job, *, input_path, output_path, num_threads=1, shard_output=False, detach=False, limit=None, mail_user=None, checkpoint_dir=None, checkpoint_interval=None, no_resume=False, no_checkpointing=False, runner='pandas', shard_mode='id', global_resume=False, job_resources=None, checkpoint_tag=None, ray_address=None)[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#

jobSwarmJob

The job instance to execute.

input_pathutils.EnvPath | str

Parquet file produced by the upstream pipeline stage.

output_pathutils.EnvPath | str

Destination Parquet file to be written by job.

num_threadsint, default 1

Number of CPU threads the job may use in the worker process.

shard_outputbool, default False

If True and output_path is a directory, emit one parquet file per shard using checkpoint outputs as the source of truth (supported by the polars runner).

shard_modestr, default “id”

Sharding strategy for num_threads > 1 (“id” for stable id hashing, “index” for legacy row order sharding).

global_resumebool, default False

When resuming a sharded job, filter inputs using global done ids across shards.

detachbool, default False

If True, start the job in a new process group and return immediately; if False (default), the call blocks until completion.

limitint or None, optional

Maximum number of rows to read from input_path — handy for dry-runs and debugging. When None (default) the entire dataset is processed.

Returns#

JobHandle

Compute job handle with normalized status and metadata.

Raises#

RuntimeError

The swarm manager is not ready (self.jobid or self.endpoint is None).

FileNotFoundError

input_path does not exist.

subprocess.CalledProcessError

Propagated when the synchronous srun command exits with a non-zero status code.

Notes#

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> ...

Examples#

>>> swarm.submit_job(
...     my_job,
...     input_path=Path("batch.parquet"),
...     output_path=Path("predictions.parquet"),
...     num_threads=4,
... )
Parameters:

  • job (SwarmJob)

  • input_path (Path)

  • output_path (Path)

  • num_threads (int)

  • shard_output (bool)

  • detach (bool)

  • limit (int | None)

  • mail_user (str | None)

  • checkpoint_dir (str | Path | None)

  • checkpoint_interval (int | None)

  • no_resume (bool)

  • no_checkpointing (bool)

  • runner (str)

  • shard_mode (str)

  • global_resume (bool)

  • job_resources (dict | None)

  • checkpoint_tag (str | None)

  • ray_address (str | None)

Return type:

JobHandle

submit_script(script_path, detach=False, extra_args=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.

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]#
down()[source]#

Manually clean up the swarm resources.

model_config = {}#

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

status()[source]#

Get the current status of the swarm.

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

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