Jobs#

SwarmJob#

Subclass this and implement transform_items to define a job. Batching, bounded concurrency, retries and checkpointing are provided by the framework.

Note

Import SwarmJob from the domyn_swarm package root. The domyn_swarm.jobs.base module is a deprecated alias for domyn_swarm.jobs.api.base, which is what is documented here.

class domyn_swarm.jobs.api.base.SwarmJob(*, config=None, client=None, **overrides)[source]#

Bases: ABC

Abstract base class for distributed LLM processing jobs in the Domyn swarm framework.

This class provides a robust foundation for running large-scale language model tasks with built-in reliability features including automatic checkpointing, retry mechanisms, and concurrent processing capabilities.

Key Features:

Automatic Checkpointing: Periodically saves progress to enable recovery from failures
Concurrent Processing: Configurable parallelism with rate limiting
and timeout handling
Retry Logic: Built-in exponential backoff for handling transient failures
Provider Agnostic: Supports multiple LLM providers (OpenAI, vLLM, etc.)
via pluggable clients
Callback System: Extensible event hooks for monitoring and custom behavior
DataFrame Integration: Native pandas DataFrame support for batch processing

Architecture:

The class follows a template method pattern where subclasses implement the core transform_items() method while inheriting all reliability and concurrency infrastructure. Processing flows through: DataFrame → batching → transform_items → results → checkpointing.

Configuration:

Configuration lives on JobConfig, not on the job itself. A subclass declares it by setting a class-level config – an instance of JobConfig, or of a subclass when the job adds fields of its own. No constructor is needed. Every configuration field is readable and writable as an attribute of the job (self.threshold below) and overridable at construction time (MyJob(model="gpt-4", threshold=0.9)). Provider request parameters (temperature, top_p, and the like) are not configuration fields: they live in config.request_params and are read back through self.kwargs.

Example

class MyJobConfig(JobConfig):
    threshold: float = 0.5


class MyJob(SwarmJob):
    config = MyJobConfig(
        input_column_name="prompt", output_cols="answer", max_concurrency=8
    )

    async def transform_items(self, items: list[Any]) -> list[Any]:
        # Process items using self.client
        results = []
        for item in items:
            response = await self.client.chat.completions.create(
                model=self.model, messages=item, **self.kwargs
            )
            results.append(response.choices[0].message.content)
        return results


job = MyJob(model="gpt-4")
results_df = await job.run(input_df, tag="experiment_1")

Variables:

config (domyn_swarm.jobs.api.config.JobConfig) – The job’s resolved configuration – see JobConfig.
client – Initialized async LLM client instance.
results – Final processed DataFrame after job completion.

Raises:

RuntimeError – When ENDPOINT environment variable is missing
ValueError – When required model name is not provided
NotImplementedError – When required abstract methods are not implemented

Parameters:

config (JobConfig)
overrides (Any)

Note

Subclasses must implement the transform_items() method which processes a list of items and returns results in the same order. The framework handles all infrastructure concerns including error handling, checkpointing, progress tracking, and result aggregation. The transform() method is deprecated and should not be used.

api_version: int = 2#

config_class#: alias of JobConfig

config: JobConfig#

property kwargs: dict[str, Any]#

Provider request parameters.

Returns:: The configured request_params.

register_callback(event, fn)[source]#

Register a named callback (e.g., ‘on_batch_done’).

Parameters:

event (str)
fn (Callable)

Return type:

None

get_callback(event)[source]#

Parameters:: event (str)
Return type:: Callable | None

async run(df, tag, checkpoint_dir='.checkpoints')[source]#

Run the job end-to-end with checkpointing support.

Parameters:

df (DataFrame)
tag (str)
checkpoint_dir (str | Path)

Return type:

DataFrame

async batched(seq, fn)[source]#

Run a batched async pipeline over seq using fn.

Supports retrying and invokes the ‘on_batch_done’ callback if registered.

Parameters:

seq (list)
fn (Callable)

Return type:

list

async transform(df)[source]#

Parameters:: df (DataFrame)

to_kwargs()[source]#

Return the configuration to serialize into --job-kwargs.

Most fields are emitted only when a caller actually set them, so a merely derived value is derived again on the far side instead of arriving as if it had been chosen. Two are unconditional: name, which __init__ always assigns, and a non-empty request_params, which a job can add to by mutating self.kwargs in place – a mutation exclude_unset cannot see, because it performs no assignment. endpoint and model are supplied separately by the runner, and the client is not serializable, so neither appears here.

Returns:: A JSON-serializable configuration payload.
Return type:: dict

abstractmethod async transform_items(items)[source]#

Pure transform: items -> results (same order). No I/O or checkpointing.

Parameters:: items (list[Any])
Return type:: list[Any]

async transform_streaming(items, *, on_flush, checkpoint_every)[source]#

Run a streaming transform without retaining all outputs in memory.

Parameters:

items (list[Any]) – Input items to process.
on_flush (Callable[[list[int], list[Any]], Awaitable[None]]) – Callback invoked with (idxs, outputs) for flushed batches.
checkpoint_every (int) – Number of items between flushes.

run_job_unified#

async domyn_swarm.jobs.execution.dispatch.run_job_unified(job_factory, data, *, input_col, output_cols, nshards=1, store_uri=None, checkpoint_every=16, data_backend=None, native_backend=None, checkpointing=True, runner='pandas', engine=None, ray_address=None, output_path=None, shard_output=False, shard_mode='id', global_resume=False)[source]#

Run a SwarmJob with backend-aware execution and checkpointing.

Parameters:

job_factory (Callable[[], Any]) – Callable that returns a SwarmJob instance.
data (Any) – Input dataset in backend-native form (pandas, polars, arrow, ray dataset, etc.).
input_col (str) – Column name to read inputs from.
output_cols (list[str] | None) – Optional list of output column names (None for dict outputs).
nshards (int) – Number of shards to split input into for non-ray execution.
store_uri (str | None) – Base checkpoint store URI (required when checkpointing is enabled).
checkpoint_every (int) – Flush interval in items.
data_backend (str | None) – Backend name override (defaults to the job’s data_backend or “pandas”). Continues to select the DataBackend used for IO and conversions; it no longer affects the routing decision, which is made by engine.
native_backend (bool | None) – Override for native execution (required for ray).
checkpointing (bool) – Whether to read/write checkpoint state.
runner (str) – Deprecated. Use engine instead. Runner implementation to use for non-ray backends (“pandas” or “arrow”), consulted only when engine is not given.
engine (str | None) – Execution engine to use (“pandas”, “arrow”, “polars” or “ray”). Takes precedence over the deprecated runner/data_backend pair. Defaults to the value resolved from that pair when not given.
ray_address (str | None) – Optional Ray cluster address (only used for ray backend).
output_path (Path | None) – Optional output path used to enable direct shard writes when using the pandas runner and directory outputs.
shard_output (bool) – If True and output_path is a directory, write one parquet file per shard directly from checkpoint outputs and return None, rather than assembling the merged result in memory and returning it. Read by the pandas and polars engines; ignored by the Arrow and Ray engines. On the pandas and polars engines this also makes global_resume a no-op: shards are cut from the full input and each shard’s own checkpoint store already supplies resume, which is equivalent to global resume as long as nshards stays fixed across resumes of the same job.
shard_mode (str) – Sharding strategy (“id” for stable id hashing, “index” for legacy order).
global_resume (bool) – Resume by filtering inputs with global done ids across shards. Ignored on the pandas and polars engines when shard_output is True (see shard_output).

Returns:

Backend-native result for non-ray runs, or the Ray runner result. Returns None when the pandas or polars paths write outputs directly to a directory.

Raises:

TypeError – If the job does not implement the streaming API.
ValueError – If required id columns are missing, if checkpointing is misconfigured, or if the resolved engine requires a data backend other than the one resolved (e.g. engine="ray" with a non-ray backend, or engine="polars" with a non-polars backend).
RuntimeError – If the backend cannot be resolved.

Return type:

Any