Agents & Orchestration

How work reaches a machine. This covers the the agent fabric fabric that AI Flow uses to route each task to a healthy agent, and how the selectors you write in a flow resolve to a real worker.

What’s shipped vs. what’s designed. the agent fabric’s design/spec-v2.md describes a richer typed protocol and a separate resilience package. Those are not wired into the running code. This doc describes the shipped behavior: a single opaque-bytes gRPC stream and the utils/ resilience primitives. The roadmap items are flagged at the end.

The picture

        ┌───────────────────────────────────────────────┐
        │              Orchestrator (gRPC server)         │
        │                                                 │
        │   AgentSelector (least CPU/mem, tag, model, gpu)│
        │   per-agent outbound queue (1000) + health map  │
        └───┬───────────────┬───────────────┬────────────┘
   KeepAlive│    KeepAlive   │    KeepAlive  │   (bidirectional streams)
            ▼                ▼               ▼
        Agent A          Agent B         Agent C
        heartbeat ↑      heartbeat ↑     heartbeat ↑   (CPU/mem/GPU/disk, capacity)
        run task  ↓      run task  ↓     run task  ↓

• Orchestrator = the gRPC server. Agents connect to it; it selects one for each task.
• Agent = a gRPC client worker node. It also can run a small server for direct agent-to-agent data pipes (high-throughput streaming that bypasses the orchestrator).

Each aiflow.svc process is a cluster node that can play either role; a deployment runs a small number of orchestrators and a fleet of agents.

The wire protocol

A single bidirectional streaming RPC carries everything:

service OrchestratorService {
    rpc KeepAlive(stream AgentMessage) returns (stream ServerMessage) {}
}

The same stream multiplexes three kinds of message, discriminated by which field is populated:

Kind	Field	Payload
Heartbeat	heartbeat	JSON-encoded AgentParameters (metrics + metadata). The first message on a stream must be a heartbeat.
Data	message	Opaque bytes (a JSON-encoded typed message).
Chunked data	chunkid + chunkend	Parts of a payload > 3.5 MiB, reassembled in order.

Typed message concepts (type/priority/TTL/control-action) live in Go and are serialized into the opaque bytes field — the wire itself stays simple.

Agent lifecycle

1. Dial & register. The agent dials the orchestrator and sends an initial heartbeat carrying its AgentParameters: Status: active, TypeParameter: remote, and metadata — tags, capabilities (the ML models it has loaded), tenants, exclusive, max_concurrent_tasks, and its addresses. The orchestrator validates service name & version against its own (this gates rollouts) and allocates a 1000-slot outbound queue.

2. Heartbeat. A dedicated goroutine sends heartbeats on a jittered (±10%) interval, each carrying fresh metrics and the agent’s current task count. A per-send mutex (streamSendMu) is held only across a single stream.Send, so a heartbeat can slip between chunks of a large transfer — a busy agent is never wrongly flagged unhealthy.

3. Health classification. A per-agent watchdog compares the last heartbeat time to 2× the heartbeat interval. If stale, it marks the agent degraded (still selectable) when transport is otherwise alive, or unhealthy otherwise.

4. Reconnect. On disconnect the agent retries with backoff. A connection-generation counter prevents a stale disconnect (from an old connection) from evicting a freshly reconnected agent.

5. Drain / stop. Drain(timeout) marks agents draining and broadcasts a drain control message; the engine completes in-flight work before shutdown.

Agent states: Initializing → Connecting → Ready → Draining / Degraded / Disconnected / Stopped. Agent types: local, remote, build.

Metrics per heartbeat: CPU cores & usage %, memory total/used/available/cached, data inflow/outflow, and (when present) GPU count/usage/VRAM/temperature and disk total/used/free. GPU metrics come from shelling out to nvidia-smi; disk from statfs.

Agent selection

For each Task node, the engine asks the orchestrator’s selector for a worker. The default selector, LeastCpuMemoryUseAgentSelector, keeps a hot in-memory cache (updated every heartbeat) plus a DB-backed store (written on connect/disconnect/health change).

Selection order (each step filters to healthy, tenant-matching, capacity-having agents):

1. Direct ID — id / workerid targets one agent exactly (skips the capacity filter).
2. Tag — tag vs the agent’s tags; round-robin among matches.
3. Model — model vs the agent’s loaded capabilities; round-robin.
4. GPU — gpu / freegpu; require a GPU with enough free VRAM %.
5. CPU / memory — first agent meeting both free-resource floors.
6. Fallback — score every candidate (100 − CPU%) + FreeMem% [+ 0.5 × FreeGPUMem%], and round-robin the best.

These map directly to the selectors: you write in a flow (see [authoring-flows.md §10](/blocks/ai/AI Flow/authoring-flows/#10-selectors—task-groups)).

Capacity & fairness

• Reservation (on the engine dispatch path) increments the agent’s task count under a lock and rejects at max_concurrent_tasks — TOCTOU-safe, so two dispatches can’t overfill an agent.
• Transient vs permanent errors are load-bearing:
  • Transient (NoFreeAgent — matching agents exist but are full/unhealthy) → the engine queues the task and retries.
  • Permanent (NoMatchingAgent — nothing could ever match this tag/model/id) → fail fast.

Task groups & affinity

• Task groups (group:) bind the first task’s agent and reuse it for the rest of the group, within one run.
• Affinity pins across runs (single-agent / definition / payload), with bindings in mlflow_affinity_binding. On agent failure, bindings are cleared and the run resumes or restarts.

See [authoring-flows.md §11](/blocks/ai/AI Flow/authoring-flows/#11-affinity).

Resilience (shipped: utils/)

Primitive	Behavior	Default
Circuit breaker	3-state; opens after N send failures, one probe to close after the reset window.	5 failures / 30 s reset
Flow control	Message cap; on overflow, wait 100 ms and retry (no hard backpressure signal).	1000 in flight
Retry	Exponential 2^i s backoff (no jitter).	—
Timeout	Buffered done-channel bound on requeue/write ops.	5 s

On a transient send failure the orchestrator requeues the message rather than dropping it.

Overflow queue (engine side)

When selection returns a transient “no free agent”, the flow engine persists the task as a mlflow_work_item and marks the instance Queued. A drain loop retries pending items (priority then age), continue-ing past starved tenants so one tenant can’t block others. Orphaned instances are re-queued on orchestrator restart. Depth is orchestrator.queue.maxdepth (0 = unlimited).

Security

• Transport: gRPC with mTLS when both server and client certificates are present; otherwise it falls back to insecure with a warning. Certs load from config keys cacert / servercert / serverkey (+ optional clientcert / clientkey).
• Version gating: the orchestrator rejects agents whose service name/version don’t match — the mechanism behind safe rolling upgrades.

Metrics to watch

Meter	Metric	Meaning
kis.agent	connected (gauge)	Currently connected agents.
kis.agent	connect.count / disconnect.count	Connection churn.
kis.agent	heartbeat.count	Heartbeat volume.
kis.agent	state_change.count	State transitions (tagged from→to).
kis.agent	backpressure.count	Flow-control overflow events.
kis.selector	duration (histogram)	Selection latency.
kis.selector	success.count / fail.count	Selection outcomes.
agent (OTel gauges)	CPU/mem/GPU/disk, data.inflow/outflow	Per-agent load.

Rising backpressure.count or selector.fail.count, or a growing overflow queue, means you need more agents (or agents with the right tags/models). See [operations.md](/blocks/ai/AI Flow/operations/) and the agent fabric’s docs/tuning-guide.md for tuning.

Roadmap (designed, not yet wired)

• The typed protobuf and additional RPCs (Health, GetAgents, broadcast, etc.) in spec-v2.md — the wire is still the v1 opaque-bytes KeepAlive stream.
• The resilience package (blocking backpressure, jittered retry, config-driven reconnect) and the types/config.go config structs — present and tested but not consumed; the runtime still uses utils/ with positional args and hard-coded constants.
• Redis-backed agent store, multi-orchestrator HA, persistent request queue in the orchestrator.