ContextQMD
Back to Cog

coglet

crates/coglet/README.md

0.19.311.1 KB
Original Source

coglet

Core Rust library for the coglet prediction server. Pure Rust with no Python dependencies - the Python bindings live in coglet-python.

Architecture

                                    coglet
    ┌─────────────────────────────────────────────────────────────────┐
    │                                                                 │
    │  ┌─────────────────────────────────────────────────────────┐   │
    │  │                    transport/http                        │   │
    │  │  ┌──────────────┐  ┌─────────────────────────────────┐  │   │
    │  │  │   server.rs  │  │           routes.rs             │  │   │
    │  │  │  Axum setup  │  │ /health, /predictions, /cancel  │  │   │
    │  │  └──────────────┘  └─────────────────────────────────┘  │   │
    │  └───────────────────────────────┬─────────────────────────┘   │
    │                                  │                              │
    │  ┌───────────────────────────────▼─────────────────────────┐   │
    │  │                     service.rs                          │   │
    │  │  PredictionService: health, permits, state, webhooks    │   │
    │  └───────────────────────────────┬─────────────────────────┘   │
    │                                  │                              │
    │         ┌────────────────────────┼────────────────┐            │
    │         │                        │                │            │
    │         ▼                        ▼                ▼            │
    │  ┌─────────────┐    ┌────────────────────┐    ┌──────────┐   │
    │  │ permit/     │    │   orchestrator.rs  │    │webhook.rs│   │
    │  │ PermitPool  │    │   Parent-side:     │    │ Sender   │   │
    │  │ Slot alloc  │    │   spawn, route     │    │ Retry    │   │
    │  └─────────────┘    └─────────┬──────────┘    └──────────┘   │
    │                               │                                │
    │  ┌────────────────────────────▼────────────────────────────┐   │
    │  │                      bridge/                            │   │
    │  │  ┌──────────────┐  ┌─────────────┐  ┌────────────────┐  │   │
    │  │  │ protocol.rs  │  │  codec.rs   │  │ transport.rs   │  │   │
    │  │  │ Message types│  │ JSON lines  │  │ Unix sockets   │  │   │
    │  │  └──────────────┘  └─────────────┘  └────────────────┘  │   │
    │  └─────────────────────────────────────────────────────────┘   │
    │                                                                 │
    │  ┌─────────────────────────────────────────────────────────┐   │
    │  │                      worker.rs                          │   │
    │  │  Child-side: PredictHandler trait, run_worker loop      │   │
    │  └─────────────────────────────────────────────────────────┘   │
    │                                                                 │
    └─────────────────────────────────────────────────────────────────┘

Directory Structure

coglet/
└── src/
    ├── lib.rs              # Public API exports
    │
    │   # Core Types
    ├── health.rs           # Health, SetupStatus, SetupResult
    ├── prediction.rs       # Prediction state machine
    ├── predictor.rs        # PredictionResult, PredictionError, PredictionOutput
    ├── version.rs          # VersionInfo
    │
    │   # Service Layer
    ├── service.rs          # PredictionService - lifecycle, state, webhooks
    ├── webhook.rs          # WebhookSender, webhook types
    │
    │   # Orchestrator (Parent Process)
    ├── orchestrator.rs     # spawn_worker, OrchestratorHandle, event loop
    │
    │   # Worker (Child Process)
    ├── worker.rs           # run_worker, PredictHandler trait, SetupError
    │
    │   # Concurrency Control
    ├── permit/
    │   ├── mod.rs
    │   ├── pool.rs         # PermitPool - slot permit management
    │   └── slot.rs         # PredictionSlot - permit + prediction binding
    │
    │   # IPC Bridge
    ├── bridge/
    │   ├── mod.rs
    │   ├── protocol.rs     # ControlRequest, ControlResponse, SlotRequest, SlotResponse
    │   ├── codec.rs        # JsonCodec - newline-delimited JSON
    │   └── transport.rs    # Unix socket transport, ChildTransportInfo
    │
    │   # HTTP Transport
    └── transport/
        ├── mod.rs
        └── http/
            ├── mod.rs
            ├── server.rs   # ServerConfig, serve()
            └── routes.rs   # Route handlers, request/response types

Key Components

PredictionService (service.rs)

Single owner of prediction state. Manages:

  • Health state (Unknown → Starting → Ready/SetupFailed)
  • PermitPool + Orchestrator reference
  • Active predictions (DashMap — single source of truth)
  • Cancellation (cancel tokens + orchestrator delegation)
  • Webhooks fire from Prediction mutation methods (no dual state)
rust
let service = PredictionService::new_no_pool()
    .with_health(Health::Starting)
    .with_version(version);

// Later, after worker is ready:
service.set_orchestrator(pool, handle).await;
service.set_health(Health::Ready).await;

Orchestrator (orchestrator.rs)

Parent-side worker lifecycle management.

spawn_worker(config)
    │
    ├─▶ Create Unix socket transport (N slots)
     ├─▶ Spawn: python -c "import coglet; coglet.server._run_worker()"
    ├─▶ Send Init message via stdin
    ├─▶ Wait for worker to connect sockets
    ├─▶ Wait for Ready message (with timeout)
    ├─▶ Populate PermitPool with slot writers
    ├─▶ Spawn event loop task
    └─▶ Return OrchestratorReady {pool, schema, handle}

Event loop handles:

  • ControlResponse::Idle - Slot ready for next prediction
  • ControlResponse::Failed - Slot poisoned, mark unavailable
  • SlotResponse::Log/Output/Done/Failed - Route to prediction
  • Worker crash - Fail all in-flight predictions

Worker (worker.rs)

Child-side event loop. Implements PredictHandler trait.

run_worker(handler, config)
    │
    ├─▶ Connect to slot sockets (from env)
    ├─▶ Setup control channel (stdin/stdout)
    ├─▶ Run handler.setup() with log routing
    ├─▶ Send Ready {slots, schema}
    ├─▶ Enter event loop:
    │       - ControlRequest::Cancel → handler.cancel(slot)
    │       - ControlRequest::Shutdown → exit
    │       - SlotRequest::Predict → spawn prediction task
    └─▶ Exit on shutdown or all slots poisoned

PermitPool (permit/pool.rs)

Slot-based concurrency control.

rust
let pool = PermitPool::new(max_concurrency);

// Add slot with its socket writer
pool.add_permit(slot_id, writer);

// Acquire permit (returns None if at capacity)
let permit = pool.try_acquire()?;

// Send prediction request
permit.send(SlotRequest::Predict { id, input }).await?;

// Return permit when done
drop(permit);

Bridge Protocol (bridge/protocol.rs)

Message types for parent-worker communication.

Control Channel:

  • ControlRequest: Init, Cancel, Shutdown
  • ControlResponse: Ready, Log, Idle, Failed, Cancelled, ShuttingDown

Slot Channel:

  • SlotRequest: Predict
  • SlotResponse: Log, Output, Done, Failed, Cancelled

All messages are JSON with {"type": "..."} discriminator.

Behaviors

Health States

Unknown ──▶ Starting ──┬──▶ Ready ◀──▶ Busy
                       │
                       └──▶ SetupFailed ──▶ Defunct
  • Unknown: Initial state, health-check returns status in body
  • Starting: Setup in progress
  • Ready: Accepting predictions
  • Busy: Ready but all slots in use (HTTP 409 on new predictions)
  • SetupFailed: setup() raised exception
  • Defunct: Unrecoverable error

Prediction States

Starting ──▶ Processing ──┬──▶ Succeeded
                          ├──▶ Failed
                          └──▶ Canceled

Cancellation

  1. HTTP DELETE /predictions/{id} or PUT /predictions/{id}/cancel
  2. Parent sends ControlRequest::Cancel { slot }
  3. Worker calls handler.cancel(slot)
  4. For sync: SIGUSR1 raises KeyboardInterrupt
  5. For async: future.cancel() on the asyncio task
  6. Prediction returns with SlotResponse::Cancelled

Shutdown

Graceful (SIGTERM with await_explicit_shutdown):

  1. Stop accepting new predictions
  2. Wait for in-flight to complete
  3. Send ControlRequest::Shutdown
  4. Worker responds ShuttingDown, exits
  5. Parent exits

Immediate (SIGTERM without flag):

  1. Send ControlRequest::Shutdown
  2. Cancel in-flight predictions
  3. Exit

Worker crash:

  1. Control channel closes
  2. Event loop detects, fails all in-flight predictions
  3. Health → Defunct

Slot Poisoning

If a slot socket has an error (write fails, etc.), the slot is marked poisoned. It won't receive new predictions. If all slots are poisoned, worker exits.

rust
enum SlotOutcome {
    Idle(SlotId),              // Ready for next prediction
    Poisoned { slot, error },  // Slot is dead
}