Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,7 @@ link no catalog and reject a configured policy type at startup.
| Policy type | Behavior |
|---|---|
| `default` | Dynamo's built-in selector and cost model. Reserved; always available. |
| `thunderagent` | Honors the paired ThunderAgent classifier's worker/rank preference, then falls back to the least-loaded eligible worker. Enable both roles for program-aware admission and repacking; see [ThunderAgent Program Scheduler](../../../../use-cases/agents/thunderagent-program-scheduler.md#native-frontend-plugin). |
| `dynamo-two-tier-cost-fn` | Ranks on two tiers instead of one additive cost: active-request load first, then device-KV prefix overlap. Prefers the worker holding the largest prefix overlap unless load is badly imbalanced. Thresholds and selection order ported from the experimental SGLang router's `cache_aware_zmq` policy. Thresholds are tunable; the defaults reproduce it exactly. |

Write the instance into the same YAML file that `--router-policy-config` already points at:
Expand Down
2 changes: 1 addition & 1 deletion docs/fern/pages/use-cases/agents/priority-scheduling.md
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ The router queue only matters when requests are held before dispatch. If a reque
(strict_priority, due_at, configured_policy_key)
```

The strict tier is compared first. Within that tier, every request with a classifier-provided `due_at` is scheduled ahead of every request without one (an unset due time sorts last) before FCFS, LCFS, or Weighted Shortest Processing Time (WSPT) applies its configured policy key. Sustained deadline-bearing traffic can therefore starve non-deadline requests in the same class. A `due_at` is set only by a router request classifier.
The strict tier is compared first. Within that tier, every request with an internal `due_at` is scheduled ahead of every request without one (an unset due time sorts last) before FCFS, LCFS, or Weighted Shortest Processing Time (WSPT) applies its configured policy key. Sustained deadline-bearing traffic can therefore starve non-deadline requests in the same class.

The default policy is `fcfs`, which uses the priority value as a positive arrival-time bump. Higher values move the request earlier in the queue. Negative priority values are clamped to zero for router queueing, so a request cannot be pushed behind normal first-come, first-served ordering by sending a negative priority.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,89 @@ title: ThunderAgent Program Scheduler
subtitle: Program-level scheduling with tool-boundary pause/resume on top of KV-aware routing
---

> **Experimental — not a released component.** Run it from a source checkout, not from a `pip install ai-dynamo`. The CLI flags, session headers, and lifecycle hooks are all unstable and will change.
**Experimental.** NVIDIA Dynamo includes a native ThunderAgent plugin for program-aware admission and worker placement. It runs inside the frontend and is enabled through router-policy YAML. It ports the scheduler from the [ThunderAgent](https://arxiv.org/abs/2602.13692) paper (Kang et al., 2026).

## Native Frontend Plugin

Use this path for aggregated serving through `dynamo.frontend`. It requires a Dynamo build that includes the `thunderagent` builtin plugin; older wheels and images do not recognize that type. The plugin is included in the builtin catalog without a separate ThunderAgent build feature. Admission behavior remains unchanged until the plugin is selected in YAML.

Custom catalogs that already register `thunderagent` must remove that registration; duplicate type names fail startup.

Save the following as `thunderagent.yaml`:

```yaml
request_classifier:
type: thunderagent
worker_selection:
aggregated: thunderagent
instances:
- name: thunderagent
type: thunderagent
```

Start the frontend against your existing aggregated workers, using the same discovery configuration as the workers:

```bash
python3 -m dynamo.frontend \
--router-mode kv \
--router-policy-config thunderagent.yaml \
--router-session-affinity-ttl-secs 1800 \
--router-session-affinity-mode soft
```

Send a stable `X-Dynamo-Session-ID` on every turn of an agent's reasoning and tool-use chain. Dynamo also accepts the [native agent headers](session-ids.mdx#native-agent-headers). Requests without session context pass through admission unchanged. At the end of a session, a dedicated minimal request with `X-Dynamo-Session-Final: true` releases the retained program state; otherwise idle retention expires it.

Select both plugin roles. The classifier owns the program table and delays admission when a session is busy or paused. When it releases a request, it supplies a preferred worker/rank. The worker selector honors that target if it is eligible, falling back to the least-loaded eligible candidate otherwise. Soft session affinity allows this repacking; hard affinity would constrain the request before the selector can move it. The `Sent` event records the worker actually selected.

```mermaid
flowchart LR
A[Request and session ID] --> B[ThunderAgent classifier]
B -->|Released request and preferred worker| C[Dynamo policy queue]
C --> D[ThunderAgent worker selector]
D --> E[Aggregated worker]
E -->|Progress and lifecycle events| B
```

The classifier uses each worker/rank's published GPU KV-block capacity and a logical token estimate from active and retained programs. It observes context growth during generation, then pauses and repacks programs between turns. This estimate is not physical KV occupancy and does not deduplicate shared prefixes or include host/offload capacity. Running generation is allowed to finish.

### Native Plugin Parameters

Set admission parameters under `request_classifier.parameters`. Leave worker-selection parameters empty. Unknown parameter names and invalid values fail startup.

| Parameter | Default | Meaning |
|---|---|---|
| `pause_threshold` | `0.95` | Per-worker usage fraction that starts pausing programs. |
| `pause_target` | `0.80` | Usage fraction the pause cycle drains toward. |
| `resume_hysteresis` | `0.10` | Headroom below the pause threshold required for normal resume. |
| `resume_timeout_seconds` | `1800` | Maximum deferral before forced release on a live worker. This does not reject the request. |
| `session_retention_seconds` | `1800` | Retention period for an idle, continuing program. |
| `scheduler_interval_seconds` | `5` | Interval between pressure and resume decisions. |
| `acting_token_weight` | `1` | Capacity weight while the agent is doing tool work. |
| `acting_decay_tau_seconds` | `1` | Acting-token decay half-life used for timeout fallback placement. |
| `buffer_per_program` | `100` | Token headroom reserved for each active program. |
| `max_tracked_requests` | `10000` | Independent limits on tracked requests and retained programs. Requests fail classification at the request limit. New programs fail at the program limit if no idle program can be evicted. |

The default deferral timeout permits waits of roughly 30 minutes under pressure. Set client deadlines to match the workload's acceptable wait. Router policy classes and queues still apply after classification releases a request.

Program state is local to each frontend's router instance. This is not a cluster-wide admission budget; keep a session's turns on the same frontend when using multiple replicas. The native plugin has been exercised with aggregated vLLM workers. Standalone EPP and disaggregated deployment guidance are outside this example.

The native plugin's classifier and worker selector register together in Dynamo's builtin catalog. It does not expose the Python prototype's soft-demotion controls, priority boosts, host-cache budget, or separate status endpoints described below.

### Native Plugin Logs

The plugin logs its interval, deferral timeout, and tracking limit at INFO when admission is enabled. A WARN event reports the number of programs forcibly resumed after reaching the deferral timeout. These requests continue processing; they are not rejected. Timeout warnings are aggregated into one event per scheduler tick.

To see scheduler state changes, enable DEBUG for the `thunderagent` log target before starting the frontend:

```bash
export DYN_LOG='info,thunderagent=debug'
```

Each changed tick reports counts of active, paused, and marked-for-pause programs, waiting requests, and tracked requests. Unchanged ticks are silent. Log snapshots use stored counts without scanning programs or requests, and log emission happens after releasing the scheduler lock.

## Python Router Prototype

The remaining sections describe the older standalone Python implementation. Run that prototype from a source checkout. Its CLI flags and observability endpoints differ from the native plugin above.

`dynamo.thunderagent_router` is a standalone Dynamo router that schedules at the granularity of an agent run — the whole `LLM turn → tool call → next turn` loop — instead of individual requests. It wraps Dynamo's native KV router and adds a program-level scheduler with tool-boundary pause/resume on top of KV-aware routing, porting the scheduler from the [ThunderAgent](https://arxiv.org/abs/2602.13692) paper (Kang et al., 2026).

Expand Down
4 changes: 4 additions & 0 deletions lib/bindings/kvbm/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 4 additions & 0 deletions lib/bindings/python/Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion lib/bindings/python/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ crate-type = ["cdylib", "rlib"]

[features]
default = ["custom-policy"]
# Links the worker-selection policies Dynamo ships plus the replaceable catalog selected by the
# Links the router plugins Dynamo ships plus the replaceable catalog selected by the
# dependency alias below. Enabled by default so a deployment can select a shipped policy through
# router-policy YAML without rebuilding. A private image replaces only the catalog dependency, so
# its policies are added alongside the shipped ones rather than displacing them.
Expand Down
6 changes: 5 additions & 1 deletion lib/router-plugins/builtin/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,17 @@ authors.workspace = true
license.workspace = true
homepage.workspace = true
repository.workspace = true
description = "Worker-selection policies Dynamo ships and always links"
description = "Router plugins shipped with Dynamo"

[dependencies]
fastrand = "2"
parking_lot.workspace = true
dynamo-kv-router = { workspace = true, features = ["standalone-selection"] }
serde.workspace = true
async-trait.workspace = true
thiserror.workspace = true
tokio.workspace = true
tracing.workspace = true

[dev-dependencies]
tokio = { workspace = true, features = ["macros", "rt", "time"] }
Expand Down
16 changes: 9 additions & 7 deletions lib/router-plugins/builtin/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,15 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

//! Worker-selection policies Dynamo ships.
//! Router plugins Dynamo ships.
//!
//! Routing hosts always link the default through `default_registry`. The optional custom
//! catalog adds the named default and two-tier providers through `register`. The default
//! catalog adds the named default, two-tier, and ThunderAgent providers through `register`. The default
//! itself uses the same public candidate inputs and scorer/picker dispatch as external policies.
//! Sequence tracking, eligibility, and admission remain in dynamo-kv-router.

mod default;
mod thunderagent;
mod two_tier_cost_fn;
pub use default::{DefaultWorkerSelector, default_factory, default_policy};

Expand All @@ -17,22 +18,23 @@ pub fn default_registry() -> RouterPluginRegistry {
RouterPluginRegistry::default().with_default_factory(default_factory())
}

use dynamo_kv_router::plugins::{RouterPluginRegistry, WorkerSelectionPolicyRegistryError};
use dynamo_kv_router::plugins::{RouterPluginRegistry, RouterPluginRegistryError};

/// Register the named providers Dynamo ships, without changing the host's default factory.
///
/// `default` is reserved by the registry for Dynamo's built-in worker selector, so no policy here
/// can shadow it. A later catalog that reuses one of these type names fails registration rather
/// than overriding it.
pub fn register(
registry: &mut RouterPluginRegistry,
) -> Result<(), WorkerSelectionPolicyRegistryError> {
pub fn register(registry: &mut RouterPluginRegistry) -> Result<(), RouterPluginRegistryError> {
default::register(registry)?;
two_tier_cost_fn::register(registry)
two_tier_cost_fn::register(registry)?;
thunderagent::register(registry)?;
Ok(())
}

#[cfg(test)]
mod tests {
use dynamo_kv_router::plugins::WorkerSelectionPolicyRegistryError;
use dynamo_kv_router::plugins::worker_selection::WorkerSelectionPolicyFactory;
use dynamo_kv_router::{KvRouterConfig, RoutingPartitionRef, WorkerType};

Expand Down
147 changes: 147 additions & 0 deletions lib/router-plugins/builtin/src/thunderagent/config.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,147 @@
// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

use std::time::Duration;

use serde::Deserialize;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum ConfigError {
#[error("invalid ThunderAgent configuration: {0}")]
Invalid(&'static str),
}

#[derive(Debug, Clone, PartialEq, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct ThunderAgentConfig {
pub pause_threshold: f64,
pub pause_target: f64,
pub resume_hysteresis: f64,
pub resume_timeout_seconds: f64,
pub session_retention_seconds: f64,
pub scheduler_interval_seconds: f64,
pub acting_token_weight: f64,
pub acting_decay_tau_seconds: f64,
pub buffer_per_program: usize,
pub max_tracked_requests: usize,
}

impl Default for ThunderAgentConfig {
fn default() -> Self {
Self {
pause_threshold: 0.95,
pause_target: 0.80,
resume_hysteresis: 0.10,
resume_timeout_seconds: 1_800.0,
session_retention_seconds: 1_800.0,
scheduler_interval_seconds: 5.0,
acting_token_weight: 1.0,
acting_decay_tau_seconds: 1.0,
buffer_per_program: 100,
max_tracked_requests: 10_000,
}
}
}

impl ThunderAgentConfig {
pub fn validate(&self) -> Result<(), ConfigError> {
validate_fraction(self.pause_threshold, "pause_threshold must be in [0, 1]")?;
validate_fraction(self.pause_target, "pause_target must be in [0, 1]")?;
if self.pause_target > self.pause_threshold {
return Err(ConfigError::Invalid(
"pause_target must not exceed pause_threshold",
));
}
validate_fraction(
self.resume_hysteresis,
"resume_hysteresis must be in [0, 1]",
)?;
if self.resume_hysteresis > self.pause_threshold {
return Err(ConfigError::Invalid(
"resume_hysteresis must not exceed pause_threshold",
));
}
validate_duration(
self.resume_timeout_seconds,
"resume_timeout_seconds must be positive",
)?;
validate_duration(
self.session_retention_seconds,
"session_retention_seconds must be positive",
)?;
validate_duration(
self.scheduler_interval_seconds,
"scheduler_interval_seconds must be positive",
)?;
if !self.acting_token_weight.is_finite() || self.acting_token_weight <= 0.0 {
return Err(ConfigError::Invalid(
"acting_token_weight must be finite and positive",
));
}
if !self.acting_decay_tau_seconds.is_finite() || self.acting_decay_tau_seconds <= 0.0 {
return Err(ConfigError::Invalid(
"acting_decay_tau_seconds must be finite and positive",
));
}
if self.max_tracked_requests == 0 {
return Err(ConfigError::Invalid(
"max_tracked_requests must be positive",
));
}
Ok(())
}

pub(crate) fn resume_timeout(&self) -> Duration {
Duration::from_secs_f64(self.resume_timeout_seconds)
}

pub(crate) fn session_retention(&self) -> Duration {
Duration::from_secs_f64(self.session_retention_seconds)
}

pub(crate) fn scheduler_interval(&self) -> Duration {
Duration::from_secs_f64(self.scheduler_interval_seconds)
}
}

fn validate_fraction(value: f64, message: &'static str) -> Result<(), ConfigError> {
if value.is_finite() && (0.0..=1.0).contains(&value) {
Ok(())
} else {
Err(ConfigError::Invalid(message))
}
}

fn validate_duration(value: f64, message: &'static str) -> Result<(), ConfigError> {
match Duration::try_from_secs_f64(value) {
Ok(duration) if !duration.is_zero() => Ok(()),
_ => Err(ConfigError::Invalid(message)),
}
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn validates_control_values() {
assert!(ThunderAgentConfig::default().validate().is_ok());
assert!(
ThunderAgentConfig {
pause_target: 0.96,
..Default::default()
}
.validate()
.is_err()
);
assert!(
ThunderAgentConfig {
scheduler_interval_seconds: 0.0,
..Default::default()
}
.validate()
.is_err()
);
}
}
Loading
Loading