# SPDX-FileCopyrightText: 2025-2026 Domyn
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
from collections.abc import Callable
import contextlib
import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING, Any
import uuid
from pydantic import (
BaseModel,
Field,
PrivateAttr,
computed_field,
)
from ulid import ULID
from domyn_swarm import utils
from domyn_swarm.backends.compute.slurm import SlurmComputeBackend
from domyn_swarm.config.plan import DeploymentContext
from domyn_swarm.config.settings import get_settings
from domyn_swarm.config.slurm import SlurmConfig
from domyn_swarm.config.swarm import DomynLLMSwarmConfig
from domyn_swarm.deploy.deployment import Deployment
from domyn_swarm.helpers.io import to_path
from domyn_swarm.helpers.logger import setup_logger
from domyn_swarm.helpers.swarm import generate_swarm_name
from domyn_swarm.platform.protocols import (
JobHandle,
JobProbe,
JobStatus,
ServingHandle,
ServingPhase,
ServingStatus,
coerce_job_status,
)
from ..core.state.state_manager import SwarmStateManager
if TYPE_CHECKING:
from domyn_swarm.jobs import SwarmJob
logger = setup_logger(__name__, level=logging.INFO)
settings = get_settings()
[docs]
class DomynLLMSwarm(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`.
Attributes:
cfg: Deployment parameters, resource requirements and platform-specific
settings.
endpoint: Public URL of the deployed load balancer. Set once deployment
succeeds through the context manager.
delete_on_exit: Whether to clean up all allocated resources when leaving
the context manager. Useful for temporary deployments.
serving_handle: 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:
RuntimeError: If resource allocation fails; the message carries
diagnostic information.
subprocess.CalledProcessError: Propagated from a failed job submission.
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 = Field(
default_factory=lambda data: data.get(
"name", generate_swarm_name(data["cfg"].name, data["cfg"].backend.type)
),
description="Unique name for this swarm deployment",
)
endpoint: str | None = None # LB endpoint, set after job submission
delete_on_exit: bool | None = (
False # Delete the resources for this cluster at the end of the job
)
serving_handle: ServingHandle | None = None # ServingHandle, set after deployment
_platform: str = PrivateAttr("")
swarm_dir: utils.EnvPath = Field(
description="Directory where swarm-related files are stored",
default_factory=lambda data: data["cfg"].home_directory / "swarms" / data["name"],
)
watchdog_db_path: utils.EnvPath = Field(
description="Path to the watchdog SQLite database file",
default_factory=lambda data: data["swarm_dir"] / "watchdog.db",
)
@computed_field
@property
def model(self) -> str:
"""
The model name, either from the config or the job submission.
If not set, defaults to the config's model.
"""
return self.cfg.model
@model.setter
def model(self, value: str):
"""
Setter for the model name. This allows setting the model after
the swarm has been created, e.g., when loading from a state file.
"""
self.cfg.model = value
[docs]
def model_post_init(self, __context: Any) -> None:
"""Post-init to set up the deployment backend."""
swarm_dirs = [
self.swarm_dir,
self.swarm_dir / "serving",
self.swarm_dir / "jobs",
self.swarm_dir / "checkpoints",
self.swarm_dir / "logs" / "endpoint",
self.swarm_dir / "logs" / "replicas",
self.swarm_dir / "logs" / "slurm",
]
for d in swarm_dirs:
os.makedirs(d, exist_ok=True)
self._cleaned = False
self._state_mgr = SwarmStateManager(self)
plan = self.cfg.get_deployment_plan()
if plan is None and self.cfg.backend is not None:
plan = self.cfg.build_plan()
if plan is not None:
extras = plan.extras | {"swarm_directory": str(self.swarm_dir)}
self._plan = plan
self._platform = plan.platform
self._deployment = Deployment(serving=plan.serving, compute=plan.compute, extras=extras)
return
def __enter__(self):
# If instantiating from state, the swarm will be already deployed,
# so just return self
if self.serving_handle is not None:
return self
assert self._deployment is not None
serving_spec = dict(self._plan.serving_spec, swarm_directory=self.swarm_dir)
ctx = DeploymentContext(
serving_spec=serving_spec,
job_resources=self._plan.job_resources,
extras=self._plan.extras | {"swarm_directory": str(self.swarm_dir)},
timeout_s=self._plan.timeout_s or self.cfg.wait_endpoint_s,
shared_env=self._plan.shared_env,
image=self._plan.image,
)
logger.info(f"Creating deployment [cyan]{self.name}[/cyan] on {self._platform}...")
handle = self._deployment.up(self.name, ctx)
# We save the handle before waiting to be ready, so we can clean up
self.serving_handle = handle
handle = self._deployment.wait_ready(timeout_s=ctx.timeout_s or self.cfg.wait_endpoint_s)
# Update the handle and deployment with the ready state
self.serving_handle = handle
self.endpoint = handle.url
self._deployment.compute = self._make_compute_backend(handle)
# Persist the state after successful deployment
self._persist(self.name)
return self
def __exit__(self, exc_type, exc, tb):
if self.delete_on_exit:
self.cleanup()
def _persist(self, deployment_name: str):
"""Save the state.
Args:
deployment_name (str): Deployment name.
"""
self.cfg.persist(self.swarm_dir / "config.yaml")
self._state_mgr.save(deployment_name)
[docs]
@classmethod
def from_state(cls, deployment_name: str) -> DomynLLMSwarm:
"""Initialize a swarm from a saved state.
Args:
deployment_name (str): Deployment name.
Returns:
DomynLLMSwarm: Loaded swarm.
"""
return SwarmStateManager.load(deployment_name)
def _delete_record(self) -> None:
"""Delete swarm from the DB
Args:
deployment_name (str): Deployment name.
"""
self._state_mgr.delete_record(self.name)
def _record_job_submission(
self,
*,
name: str,
command: list[str],
resources: dict | None,
kind: str,
status: JobStatus,
external_id: str | None = None,
error: str | None = None,
) -> str | None:
"""Persist job metadata in the local swarm DB.
Args:
name: Job name.
command: Command argv list (no secrets).
resources: Resource dict (if any).
kind: Job kind (e.g., "step", "script").
status: Job status.
external_id: Optional backend external identifier.
error: Optional error message.
Returns:
Job record ID if created, else None.
"""
try:
job_id = SwarmStateManager.create_job(
deployment_name=self.name,
provider=self._platform,
kind=kind,
status=status,
external_id=external_id,
name=name,
command=command,
resources=resources,
error=error,
)
return job_id
except Exception as exc:
logger.warning("Failed to persist job record: %s", exc)
return None
def _update_job_submission(
self,
*,
job_id: str | None,
status: JobStatus,
external_id: str | None = None,
log_paths: dict | None = None,
error: str | None = None,
) -> None:
"""Update job submission state in the local swarm DB.
Args:
job_id: Internal persisted job ID.
status: Job status.
external_id: Optional backend external identifier.
log_paths: Optional backend log path dictionary.
error: Optional error message.
"""
if not job_id:
return
try:
SwarmStateManager.update_job(
job_id,
status=status,
external_id=external_id,
log_paths=log_paths,
error=error,
)
except Exception as exc:
logger.warning("Failed to update job record %s: %s", job_id, exc)
def _submit_with_tracking(
self,
*,
name: str,
command: list[str],
resources: dict | None,
kind: str,
submit: Callable[[], object | None],
) -> JobHandle:
"""Submit a compute job and persist lifecycle transitions.
Args:
name: Job name.
command: Command argv list to persist.
resources: Optional resource requirements.
kind: Job kind (for persistence).
submit: Backend submission callable.
Returns:
Submitted job handle.
Raises:
RuntimeError: If backend submission returns ``None``.
Exception: Re-raises backend submission errors.
"""
job_record_id = self._record_job_submission(
name=name,
command=command,
resources=resources,
kind=kind,
status=JobStatus.PENDING,
)
try:
raw_handle = submit()
except Exception as exc:
self._update_job_submission(
job_id=job_record_id,
status=JobStatus.FAILED,
error=str(exc),
)
raise
if raw_handle is None:
error_msg = "Failed to submit job to compute backend."
self._update_job_submission(
job_id=job_record_id,
status=JobStatus.FAILED,
error=error_msg,
)
raise RuntimeError(error_msg)
if isinstance(raw_handle, JobHandle):
job_handle = raw_handle
else:
raw_meta = getattr(raw_handle, "meta", None)
job_handle = JobHandle(
id=str(getattr(raw_handle, "id", name)),
status=coerce_job_status(getattr(raw_handle, "status", JobStatus.PENDING)),
meta=dict(raw_meta) if isinstance(raw_meta, dict) else {},
)
normalized_status = coerce_job_status(job_handle.status)
job_handle.status = normalized_status
self._update_job_submission(
job_id=job_record_id,
status=normalized_status,
external_id=job_handle.external_id,
log_paths=job_handle.meta.get("log_paths"),
)
if job_record_id:
job_handle.meta["job_id"] = job_record_id
return job_handle
[docs]
def submit_job(
self,
job: SwarmJob,
*,
input_path: Path,
output_path: Path,
num_threads: int = 1,
shard_output: bool = False,
detach: bool = False,
limit: int | None = None,
mail_user: str | None = None,
checkpoint_dir: str | Path | None = None,
checkpoint_interval: int | None = None,
no_resume: bool = False,
no_checkpointing: bool = False,
runner: str = "pandas",
shard_mode: str = "id",
global_resume: bool = False,
job_resources: dict | None = None,
checkpoint_tag: str | None = None,
ray_address: str | None = None,
) -> JobHandle:
"""
Launch a serialized :class:`~domyn_swarm.SwarmJob` inside the current
SLURM swarm allocation.
The *job* object is converted to keyword arguments via
:py:meth:`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.
input_path : utils.EnvPath | str
Parquet file produced by the upstream pipeline stage.
output_path : utils.EnvPath | str
Destination Parquet file to be written by *job*.
num_threads : int, default 1
Number of CPU threads the job may use in the worker process.
shard_output : bool, 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_mode : str, default "id"
Sharding strategy for `num_threads` > 1 ("id" for stable id hashing, "index" for
legacy row order sharding).
global_resume : bool, default False
When resuming a sharded job, filter inputs using global done ids across shards.
detach : bool, default False
If *True*, start the job in a new process group and return immediately;
if *False* (default), the call blocks until completion.
limit : int 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,
... )
"""
if checkpoint_dir is None:
checkpoint_dir = self.swarm_dir / "checkpoints"
input_parquet = to_path(input_path)
output_parquet = to_path(output_path)
from domyn_swarm.jobs import JobBuilder
job_class = JobBuilder.to_class_path(job)
job_kwargs = JobBuilder.to_kwargs_json(job)
python_interpreter, image, resources, env = self._compose_runtime()
resources = self._merge_resources(resources, None, job_resources)
env = self._augment_job_env(env, job, job_class, ray_address=ray_address)
exe = self._build_job_command(
job=job,
job_class=job_class,
job_kwargs=job_kwargs,
input_parquet=input_parquet,
output_parquet=output_parquet,
num_threads=num_threads,
checkpoint_dir=checkpoint_dir,
checkpoint_interval=checkpoint_interval,
runner=runner,
no_resume=no_resume,
no_checkpointing=no_checkpointing,
checkpoint_tag=checkpoint_tag,
shard_output=shard_output,
shard_mode=shard_mode,
global_resume=global_resume,
limit=limit,
ray_address=env.get("DOMYN_SWARM_RAY_ADDRESS"),
python_interpreter=str(python_interpreter),
)
job_name = job.name.lower() if job.name else f"{self.name}-job"
ulid = str(ULID())
job_name = f"{self.name}-{job_name}-{ulid.lower()}"
logger.info(
f"Submitting {job.__class__.__name__} [cyan]{job_name}[/cyan] job "
f"to swarm {self.name} on {self._platform}"
)
return self._submit_with_tracking(
name=job_name[:36], # type: ignore[arg-type]
command=[*map(str, exe)],
resources=resources,
kind="step",
submit=lambda: self._deployment.run(
name=job_name[:36], # type: ignore[arg-type]
image=image,
command=exe,
env=env,
resources=resources,
detach=detach,
),
)
def _augment_job_env(
self,
env: dict,
job: SwarmJob,
job_class: str,
*,
ray_address: str | None,
) -> dict:
"""Augment the base environment for job submission.
Args:
env: Base environment variables.
job: Job instance being submitted.
job_class: Fully qualified job class name.
ray_address: Optional Ray address override.
Returns:
Environment dictionary for the job.
"""
env = dict(env)
env.update(
{
"ENDPOINT": self.endpoint,
"MODEL": self.model,
"JOB_CLASS": job_class,
}
)
token = settings.api_token or settings.vllm_api_key or settings.singularityenv_vllm_api_key
if token:
env["DOMYN_SWARM_API_TOKEN"] = token.get_secret_value()
env["VLLM_API_KEY"] = token.get_secret_value()
if getattr(job, "data_backend", None) == "ray":
resolved_ray_address = (
ray_address
or env.get("DOMYN_SWARM_RAY_ADDRESS")
or env.get("RAY_ADDRESS")
or os.environ.get("DOMYN_SWARM_RAY_ADDRESS")
or os.environ.get("RAY_ADDRESS")
)
if not resolved_ray_address:
raise ValueError(
"Ray backend requires an explicit ray address. Provide --ray-address or set "
"DOMYN_SWARM_RAY_ADDRESS/RAY_ADDRESS in the swarm environment."
)
env["DOMYN_SWARM_RAY_ADDRESS"] = resolved_ray_address
return env
def _build_job_command(
self,
*,
job: SwarmJob,
job_class: str,
job_kwargs: str,
input_parquet: Path,
output_parquet: Path,
num_threads: int,
checkpoint_dir: str | Path,
checkpoint_interval: int | None,
runner: str,
no_resume: bool,
no_checkpointing: bool,
checkpoint_tag: str | None,
shard_output: bool,
shard_mode: str,
global_resume: bool,
limit: int | None,
ray_address: str | None,
python_interpreter: str,
) -> list[str]:
"""Build the job runner command to execute inside the swarm.
Args:
job: Job instance being submitted.
job_class: Fully qualified job class name.
job_kwargs: Serialized job kwargs JSON.
input_parquet: Input dataset path.
output_parquet: Output dataset path.
num_threads: Worker thread count.
checkpoint_dir: Checkpoint directory.
checkpoint_interval: Checkpoint interval override.
runner: Runner implementation name.
no_resume: Whether to ignore existing checkpoints.
no_checkpointing: Whether to disable checkpointing.
checkpoint_tag: Optional checkpoint tag override.
shard_output: Whether to write shard outputs to a directory.
shard_mode: Sharding strategy.
global_resume: Whether to enable global resume.
limit: Optional input row limit.
ray_address: Optional Ray address override.
python_interpreter: Python interpreter path to use.
Returns:
Command list for job execution.
"""
exe = [
python_interpreter,
"-m",
"domyn_swarm.jobs.cli.run",
f"--job-class={job_class}",
f"--model={self.model}",
f"--input-parquet={input_parquet}",
f"--output-parquet={output_parquet}",
f"--endpoint={self.endpoint}",
f"--nthreads={num_threads}",
f"--checkpoint-dir={checkpoint_dir}",
f"--checkpoint-interval={checkpoint_interval or job.checkpoint_interval}",
f"--runner={runner}",
"--job-kwargs",
job_kwargs,
]
if no_resume:
exe.append("--no-resume")
if no_checkpointing:
exe.append("--no-checkpointing")
if checkpoint_tag:
exe.append(f"--checkpoint-tag={checkpoint_tag}")
if shard_output:
exe.append("--shard-output")
if shard_mode:
exe.append(f"--shard-mode={shard_mode}")
if global_resume:
exe.append("--global-resume")
if limit:
exe.append(f"--limit={limit}")
if ray_address:
exe.append(f"--ray-address={ray_address}")
return exe
def _compose_runtime(self, extra_env: dict | None = None):
"""Build interpreter/image/resources/env in one place."""
self._deployment.ensure_ready()
assert self._deployment.compute is not None, "Compute backend not initialized"
compute = self._deployment.compute
python_interpreter = compute.default_python(self.cfg)
image = (
self._plan.image
if self._plan and self._plan.image
else compute.default_image(self.cfg.backend)
)
resources = self._merge_resources(
compute.default_resources(self.cfg.backend),
self._plan.job_resources if self._plan else None,
None,
)
env = {
"ENDPOINT": self.endpoint,
"MODEL": self.model,
}
if self._plan and self._plan.shared_env:
env.update(self._plan.shared_env)
# Global backend env from config
if getattr(self.cfg, "backend", None) and getattr(self.cfg.backend, "env", None):
env.update(self.cfg.backend.env) # type: ignore[attr-defined]
# Backend-specific overrides
overrides = compute.default_env(self.cfg)
if overrides:
env.update(overrides)
# Call-site extras
if extra_env:
env.update(extra_env)
return str(python_interpreter), image, resources, env
@staticmethod
def _merge_resources(
base: dict | None,
plan_resources: dict | None,
overrides: dict | None,
) -> dict | None:
if not base and not plan_resources and not overrides:
return None
merged: dict = {}
if base:
merged.update(base)
if plan_resources:
merged.update(plan_resources)
if overrides:
merged.update(overrides)
return merged
[docs]
def submit_script(
self,
script_path: Path,
detach: bool = False,
extra_args: list[str] | None = None,
) -> JobHandle:
"""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.
Args:
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:
JobHandle: Submitted job handle with normalized status and metadata.
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
"""
# Basic validation
if self._platform == "slurm" and not script_path.is_file():
raise FileNotFoundError(f"Script not found: {script_path}")
# Compose runtime (interpreter/image/resources/env) once
python_interpreter, image, resources, env = self._compose_runtime()
# Build the command
args = extra_args or []
command = [python_interpreter, str(script_path), *args]
script_name = f"{self.name.lower()}-script"
return self._submit_with_tracking(
name=script_name,
command=[*map(str, command)],
resources=resources,
kind="script",
submit=lambda: self._deployment.run(
name=script_name, # type: ignore[arg-type]
image=image,
command=command,
env=env,
resources=resources,
detach=detach,
),
)
def _require_compute_backend(self):
"""Return the initialized compute backend.
Raises:
RuntimeError: If compute backend is not initialized.
"""
compute = self._deployment.compute
if compute is None:
raise RuntimeError("Compute backend is not initialized for this swarm.")
return compute
[docs]
def wait_job(self, handle: JobHandle, *, stream_logs: bool = True) -> JobStatus:
"""Wait for a submitted compute job to reach a terminal state.
Args:
handle: Job handle to wait on.
stream_logs: Whether to stream backend logs while waiting.
Returns:
Normalized terminal job status.
"""
compute = self._require_compute_backend()
status = compute.wait(handle, stream_logs=stream_logs)
normalized_status = coerce_job_status(status)
handle.status = normalized_status
return normalized_status
[docs]
def cancel_job(self, handle: JobHandle) -> JobStatus:
"""Cancel a submitted compute job.
Args:
handle: Job handle to cancel.
Returns:
Final normalized status after cancellation.
"""
compute = self._require_compute_backend()
compute.cancel(handle)
normalized_status = coerce_job_status(
getattr(handle, "status", JobStatus.CANCELLED),
default=JobStatus.CANCELLED,
)
if normalized_status != JobStatus.CANCELLED:
normalized_status = JobStatus.CANCELLED
handle.status = normalized_status
return normalized_status
[docs]
def refresh_job_status(self, job_id: str) -> dict[str, Any]:
"""Refresh a persisted job status via backend probe (best effort).
Args:
job_id: Internal persisted job identifier.
Returns:
Updated job record payload, including transient refresh metadata:
``refresh_source`` and ``refresh_error``.
"""
record = SwarmStateManager.get_job(job_id)
external_id = record.get("external_id")
status = coerce_job_status(record.get("status", JobStatus.PENDING))
handle = JobHandle(
id=str(external_id or job_id),
status=status,
meta={
"job_id": job_id,
"external_id": external_id,
},
)
if not external_id:
record["refresh_source"] = "db"
record["refresh_error"] = "Missing external_id; backend probe skipped."
return record
compute = self._require_compute_backend()
probe: JobProbe
try:
probe = compute.probe(handle)
except Exception as exc:
refresh_error = str(exc)
logger.warning("Job status probe failed for %s: %s", job_id, refresh_error)
with contextlib.suppress(Exception):
SwarmStateManager.update_job(job_id, error=refresh_error)
refreshed = SwarmStateManager.get_job(job_id)
refreshed["refresh_source"] = "backend"
refreshed["refresh_error"] = refresh_error
return refreshed
try:
SwarmStateManager.update_job(
job_id,
status=probe.status,
raw_status=probe.raw_status,
external_id=handle.external_id or external_id,
error=probe.error,
)
except Exception as exc:
logger.warning("Failed to persist probed status for %s: %s", job_id, exc)
refreshed = SwarmStateManager.get_job(job_id)
refreshed["refresh_source"] = probe.source
refreshed["refresh_error"] = probe.error
return refreshed
[docs]
def cleanup(self):
if self._deployment and self.serving_handle:
self._deployment.down(self.serving_handle)
self._delete_record()
self._cleaned = True
[docs]
def down(self):
"""Manually clean up the swarm resources."""
self.cleanup()
def _deployment_name(self) -> str:
unique_id = uuid.uuid4()
short_id = str(unique_id)[:8]
return f"{self.cfg.name}-{short_id}"
def _make_compute_backend(self, handle: ServingHandle):
if self._plan and self._plan.platform == "slurm":
lb_jobid = handle.meta.get("lb_jobid")
lb_node = handle.meta.get("lb_node")
if not lb_jobid or not lb_node:
raise RuntimeError("LB Job ID/Node missing in Slurm handle.")
assert self.cfg.backend is not None and isinstance(self.cfg.backend, SlurmConfig)
return SlurmComputeBackend(cfg=self.cfg.backend, lb_jobid=lb_jobid, lb_node=lb_node)
elif self._plan and self._plan.platform == "lepton":
return self._plan.compute
else:
raise RuntimeError(f"Unsupported platform for compute backend: {self._plan.platform}")
[docs]
def status(self) -> ServingStatus:
"""Get the current status of the swarm."""
if self._deployment:
s = self._deployment.status()
if s is not None:
return s
return ServingStatus(phase=ServingPhase.UNKNOWN, url=self.endpoint)
