# SPDX-FileCopyrightText: 2025-2026 Domyn
# SPDX-License-Identifier: Apache-2.0
"""
Light-weight framework for driver scripts that run **inside** a Domyn swarm.
Every class:
1. Reads the load-balancer URL from the `ENDPOINT` env-var (injected by
`DomynLLMSwarm` on the head node).
2. Creates a single `openai.AsyncOpenAI` client pointing to that URL
(`base_url=ENDPOINT`, `api_key="-"`).
3. Provides `.run(df)` - a *synchronous* wrapper around an async
coroutine so users don't have to think about `asyncio` unless they
want to.
4. Implements `.to_kwargs()` ⇒ JSON-serialisable dict so the object can
be reconstructed by `domyn_swarm.jobs.cli.run` inside the allocation.
Sub-classes included:
* `CompletionJob` → one prompt → one text completion
* `ChatCompletionJob` → list-of-messages → one assistant reply
"""
import abc
from collections.abc import Awaitable, Callable
import dataclasses
from enum import Enum
import inspect
import logging
import os
from pathlib import Path
import threading
from typing import Any, ClassVar
import warnings
from deprecated import deprecated
from openai import AsyncOpenAI
import pandas as pd
from tqdm import tqdm
from domyn_swarm.checkpoint.manager import CheckpointManager
from domyn_swarm.config.settings import get_settings
from domyn_swarm.helpers.logger import setup_logger
from .batching import BatchExecutor
logger = setup_logger(__name__, level=logging.INFO)
settings = get_settings()
class OutputJoinMode(str, Enum):
APPEND = "append"
REPLACE = "replace"
IO_ONLY = "io_only"
[docs]
class SwarmJob(abc.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.
Example:
::
class MyLLMJob(SwarmJob):
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 = MyLLMJob(model="gpt-4", max_concurrency=5)
results_df = await job.run(input_df, tag="experiment_1")
Attributes:
api_version: API version for compatibility tracking (default: 2)
endpoint: LLM service endpoint URL (from ENDPOINT env var or parameter)
model: Model identifier (e.g., "gpt-4", "claude-3-sonnet")
provider: LLM provider name ("openai", "anthropic", etc.)
input_column_name: DataFrame column containing input data
id_column_name: Optional column name used for stable row ids
output_cols: DataFrame column(s) for storing results
checkpoint_interval: Items processed between automatic checkpoints
max_concurrency: Maximum concurrent requests allowed
retries: Maximum retry attempts for failed requests
timeout: Request timeout in seconds
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
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
_REQUEST_KWARG_BLOCKLIST: ClassVar = {
"api_version",
"checkpoint_interval",
"client",
"client_kwargs",
"default_output_cols",
"endpoint",
"input_column_name",
"max_concurrency",
"model",
"name",
"output_cols",
"output_column_name",
"output_mode",
"provider",
"retries",
"results",
"system_prompt",
"timeout",
}
def __init__(
self,
*,
name: str | None = None,
endpoint: str | None = None,
model: str = "",
provider: str = "openai",
input_column_name: str = "messages",
id_column_name: str | None = None,
output_column_name: str | list | None = None,
output_cols: str | list | None = None,
checkpoint_interval: int = 16,
max_concurrency: int = 2,
retries: int = 5,
timeout: float = 600,
client=None,
client_kwargs: dict | None = None,
output_mode: OutputJoinMode = OutputJoinMode.APPEND,
default_output_cols: list[str] | None = None,
data_backend: str | None = None,
native_backend: bool = False,
backend_read_kwargs: dict | None = None,
backend_write_kwargs: dict | None = None,
native_batch_size: int | None = None,
**extra_kwargs,
):
"""Initialize the job with parameters and an optional LLM client.
Args:
name: Optional job name (for logging).
endpoint: Optional LLM endpoint URL (overrides `ENDPOINT` env var).
model: Model name to use (e.g., "gpt-4").
provider: LLM provider (default: "openai").
input_column_name: Name of the input column in the DataFrame.
id_column_name: Optional column name for stable row identifiers.
output_column_name: [DEPRECATED] Name of the output column(s) in the DataFrame.
Use output_cols instead.
output_cols: Name of the output column(s) in the DataFrame.
checkpoint_interval: Number of items to process before checkpointing.
max_concurrency: Maximum number of concurrent requests to process.
retries: Number of retries for failed requests.
timeout: Request timeout in seconds.
client: Optional pre-initialized LLM client (e.g., `AsyncOpenAI`).
client_kwargs: Additional kwargs for the LLM client.
output_mode: How to join outputs to the input DataFrame.
default_output_cols: Default output columns if none are specified.
**extra_kwargs: Additional parameters to pass to the job constructor.
Raises:
RuntimeError: If ENDPOINT environment variable is not set.
ValueError: If model name is not specified.
"""
self.name = name or self.__class__.__name__
self.endpoint = endpoint or os.getenv("ENDPOINT")
if not self.endpoint:
raise RuntimeError("ENDPOINT environment variable is not set")
if not model:
raise ValueError("Model name must be specified")
# Handle deprecated output_column_name parameter
if output_column_name is not None and output_cols is not None:
warnings.warn(
"Both 'output_column_name' and 'output_cols' parameters are provided. "
"The 'output_column_name' parameter is deprecated and "
"will be ignored in favor of 'output_cols'.",
DeprecationWarning,
stacklevel=2,
)
if output_column_name is not None:
warnings.warn(
"The 'output_column_name' parameter is "
"deprecated and will be removed in a future version. "
"Use 'output_cols' instead.",
DeprecationWarning,
stacklevel=2,
)
self.output_column_name = output_column_name
self.output_cols = output_column_name
elif output_cols is not None:
self.output_cols = output_cols
else:
self.output_cols = "result"
self.model = model
self.provider = provider
self.input_column_name = input_column_name
self.id_column_name = id_column_name
self.checkpoint_interval = checkpoint_interval
self.max_concurrency = max_concurrency
self.retries = retries
self.timeout = timeout
self.kwargs = {**extra_kwargs.get("kwargs", extra_kwargs)}
self.output_mode = output_mode
self.default_output_cols = (
default_output_cols
if default_output_cols is not None
else ([self.output_cols] if isinstance(self.output_cols, str) else self.output_cols)
)
self.data_backend = data_backend
self.native_backend = native_backend
self.backend_read_kwargs = backend_read_kwargs
self.backend_write_kwargs = backend_write_kwargs
self.native_batch_size = native_batch_size
headers = {}
token = settings.api_token or settings.vllm_api_key or settings.singularityenv_vllm_api_key
if token:
logger.info("Using API_TOKEN from environment for authentication")
headers["Authorization"] = f"Bearer {token.get_secret_value()}"
self.client = client or AsyncOpenAI(
base_url=f"{self.endpoint}/v1",
api_key="-",
organization="-",
project="-",
timeout=timeout,
default_headers=headers,
**(client_kwargs or {}),
)
self._callbacks: dict[str, Callable] = {}
self.results = None
[docs]
def register_callback(self, event: str, fn: Callable) -> None:
"""Register a named callback (e.g., 'on_batch_done')."""
self._callbacks[event] = fn
[docs]
def get_callback(self, event: str) -> Callable | None:
return self._callbacks.get(event)
[docs]
async def run(
self,
df: pd.DataFrame,
tag: str,
checkpoint_dir: str | Path = ".checkpoints",
) -> pd.DataFrame:
"""
Run the job end-to-end with checkpointing support.
"""
checkpoint_dir = Path(checkpoint_dir)
checkpoint_dir.mkdir(parents=True, exist_ok=True)
path = checkpoint_dir / f"{self.__class__.__name__}_{tag}.parquet"
manager = CheckpointManager(
path,
df,
expected_output_cols=self.output_cols,
input_col=self.input_column_name,
)
todo_df = manager.filter_todo()
idx_map = todo_df.index.to_numpy()
async def flush(out_list, new_ids):
thread_name = threading.current_thread().name
manager.flush(out_list, new_ids, self.output_cols, idx_map)
tqdm.write(
f"[{thread_name}] Checkpoint flushed {len(new_ids)} "
f"rows, new total: {len(manager.done_df)}"
)
self.register_callback("on_batch_done", flush)
try:
items = todo_df[self.input_column_name].tolist()
await self.batched(items, self._call_unit)
finally:
self._callbacks.clear()
self.results = manager.finalize()
return self.results
[docs]
async def batched(self, seq: list, fn: Callable) -> list:
"""
Run a batched async pipeline over `seq` using `fn`.
Supports retrying and invokes the 'on_batch_done' callback if registered.
"""
executor = BatchExecutor(self.max_concurrency, self.checkpoint_interval, self.retries)
return await executor.run(
seq,
fn,
on_batch_done=self.get_callback("on_batch_done"),
progress=True,
)
[docs]
def to_kwargs(self) -> dict:
"""
Serialize the job's constructor parameters (for remote reconstruction).
"""
if dataclasses.is_dataclass(self):
return dataclasses.asdict(self)
return {
k: v
for k, v in self.__dict__.items()
if isinstance(v, str | int | float | bool | list | dict | type(None))
and k not in {"endpoint", "model", "client", "_callbacks", "results"}
}
def _request_kwargs(self) -> dict:
if not self.kwargs:
return {}
return {
k: v
for k, v in self.kwargs.items()
if not k.startswith("_") and k not in self._REQUEST_KWARG_BLOCKLIST
}
async def _call_unit(self, item: Any) -> Any:
"""
Bridge: run `transform_items` on a single element and return the single result.
Ensures the contract (len(out) == 1).
"""
out = self.transform_items([item])
if inspect.isawaitable(out):
out = await out
if not isinstance(out, list) or len(out) != 1:
raise RuntimeError(
"transform_items(items) must return a list of the same length as `items`."
)
return out[0]
