Omni-Agent: Asynchronous Distributed Runtime for ML Workloads

A fault-tolerant asynchronous distributed runtime for executing concurrent ML workloads across heterogeneous compute clusters. Integrates with Ray for dynamic cloud clusters and SLURM for HPC environments.

Core Abstractions

OmniAgentRuntime
├── PriorityTaskQueue       Priority-based scheduling, preemption support
├── DistributedWorker[]     Ray actors, persistent, checkpoint-aware
├── CheckpointManager       Incremental, zero-copy state persistence
├── RecoveryPolicy          Exponential backoff, jitter, configurable retries
└── AsyncPrefetcher         Memory-pinned background data loading

Key Features

Zero-copy state migration: Uses Apache Arrow IPC format for cross-node state transfer and torch.share_memory_() for intra-node handoff — no redundant tensor copies during worker reassignment.

Adaptive prefetching: AsyncPrefetcher runs data loading on a background thread, pins tensors to page-locked memory, and uses non-blocking CUDA streams for H2D transfer. GPU never waits for data.

Fault recovery: Configurable RecoveryPolicy with exponential backoff + jitter. Tasks resume from latest checkpoint rather than restarting from scratch — critical for long-horizon training runs.

Priority scheduling: PriorityTaskQueue dispatches higher-priority tasks first. Useful for prioritizing evaluation tasks over data generation during active training.

Usage

import asyncio
from omni_agent.runtime.executor import OmniAgentRuntime, TaskSpec
from omni_agent.fault_tolerance.recovery import RecoveryPolicy, RecoveryStrategy

async def main():
    runtime = OmniAgentRuntime(
        n_workers=8,
        checkpoint_dir="./checkpoints",
        recovery_policy=RecoveryPolicy(
            strategy=RecoveryStrategy.JITTERED_BACKOFF,
            max_retries=3,
            base_delay_s=5.0,
        ),
        max_concurrent_tasks=32,
    )

    # Submit tasks
    for i in range(100):
        spec = TaskSpec(
            fn=my_training_fn,
            args=(dataset_shard_i,),
            priority=1 if i < 10 else 0,   # Prioritize first 10 shards
            max_retries=3,
            timeout_s=3600.0,
            checkpoint_every_n=100,
        )
        await runtime.submit(spec)

    # Run until all tasks complete
    final_states = await runtime.run_until_complete()
    print(runtime.summary())
    await runtime.shutdown()

asyncio.run(main())

Prefetcher Usage

from omni_agent.data.prefetcher import AsyncPrefetcher

prefetcher = AsyncPrefetcher(
    loader=dataloader,
    n_prefetch=2,
    device="cuda:0",
)

for batch in prefetcher:
    # batch tensors are already on GPU
    # H2D transfer overlapped with previous forward pass
    loss = model(batch)
    loss.backward()

Checkpointing

Tasks receive a checkpoint_fn callback — call it periodically to save state:

def my_training_fn(*args, start_step=0, checkpoint_fn=None, **kwargs):
    for step in range(start_step, 10000):
        # ... training logic ...
        if step % 100 == 0 and checkpoint_fn:
            checkpoint_fn(step, {"model": model.state_dict(), "step": step})
    return {"final_loss": loss}

On worker failure, the task resumes from start_step automatically.

Recovery Strategies

Strategy	Behavior
IMMEDIATE	Retry instantly
FIXED_DELAY	Wait base_delay_s between retries
EXPONENTIAL_BACKOFF	Delay doubles each retry
JITTERED_BACKOFF	Exponential + random jitter (recommended for distributed)

License

MIT