Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
1 change: 1 addition & 0 deletions lib/bindings/python/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ media-ffmpeg = ["dynamo-llm/media-ffmpeg"]
kv-indexer = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-indexer"]
kv-indexer-metrics = ["kv-indexer", "dynamo-kv-router/metrics"]
slot-tracker = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-slot-tracker"]
select-service = ["dep:clap", "dep:tracing-subscriber", "dynamo-kv-router/standalone-selection"]
mocker-kvbm-offload = ["dynamo-mocker/kvbm-offload"]
aic-forward-pass = ["dynamo-mocker/aic-forward-pass", "dep:aiconfigurator-core"]
nvtx = ["dynamo-runtime/nvtx"]
Expand Down
9 changes: 9 additions & 0 deletions lib/bindings/python/rust/lib.rs
Comment thread
PeaBrane marked this conversation as resolved.
Resolved
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_function(wrap_pyfunction!(fetch_model, m)?)?;
m.add_function(wrap_pyfunction!(run_kv_indexer, m)?)?;
m.add_function(wrap_pyfunction!(run_slot_tracker, m)?)?;
m.add_function(wrap_pyfunction!(run_select_service, m)?)?;
m.add_function(wrap_pyfunction!(llm::entrypoint::make_engine, m)?)?;
m.add_function(wrap_pyfunction!(llm::replay::run_mocker_trace_replay, m)?)?;
m.add_function(wrap_pyfunction!(
Expand Down Expand Up @@ -290,6 +291,14 @@ fn run_slot_tracker(py: Python<'_>, argv: Option<Vec<String>>) -> PyResult<()> {
.map_err(standalone_to_pyerr)
}

#[pyfunction(name = "run_select_service")]
#[pyo3(signature = (argv=None))]
fn run_select_service(py: Python<'_>, argv: Option<Vec<String>>) -> PyResult<()> {
let argv = argv.unwrap_or_default();
py.allow_threads(move || llm::kv::run_select_service_cli(argv))
.map_err(standalone_to_pyerr)
}

/// Log a message from Python with file and line info
#[pyfunction]
#[pyo3(text_signature = "(level, message, module, file, line)")]
Expand Down
65 changes: 62 additions & 3 deletions lib/bindings/python/rust/llm/kv.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,21 @@ use tokio_stream::StreamExt;
use super::local_model::RoutingConstraints;
use super::*;
use crate::Endpoint;
#[cfg(any(feature = "kv-indexer", feature = "slot-tracker"))]
#[cfg(any(
feature = "kv-indexer",
feature = "slot-tracker",
feature = "select-service"
))]
use clap::Parser;
use dynamo_kv_router::config::{KvRouterConfig, RouterConfigOverride};
use dynamo_kv_router::config::{
KvRouterConfig, RouterConfigOverride, kv_router_config_from_dynamo_env,
};
use dynamo_kv_router::protocols::compute_block_hash_for_seq;
use dynamo_kv_router::protocols::*;
#[cfg(feature = "kv-indexer")]
use dynamo_kv_router::services::indexer::{self, IndexerConfig};
#[cfg(feature = "select-service")]
use dynamo_kv_router::services::selection::{self, SelectionServiceConfig};
#[cfg(feature = "slot-tracker")]
use dynamo_kv_router::services::slot_tracker::{self, SlotTrackerConfig};
use rs::pipeline::{AsyncEngine, SingleIn};
Expand Down Expand Up @@ -191,7 +199,58 @@ where
}
}

#[cfg(any(feature = "kv-indexer", feature = "slot-tracker"))]
#[cfg(feature = "select-service")]
#[derive(Parser)]
#[command(
name = "python -m dynamo.select_service",
about = "Runtime-free Dynamo worker selection service"
)]
struct SelectServiceCli {
/// HTTP server port
#[arg(long, default_value_t = 8092)]
port: u16,

/// Number of KV indexer worker threads
#[arg(long, default_value_t = 4)]
threads: usize,
}

pub fn run_select_service_cli<I, T>(args: I) -> anyhow::Result<()>
where
I: IntoIterator<Item = T>,
T: Into<OsString>,
{
#[cfg(feature = "select-service")]
{
let cli = SelectServiceCli::try_parse_from(
std::iter::once(OsString::from("python -m dynamo.select_service"))
.chain(args.into_iter().map(Into::into)),
)?;

init_standalone_logging();

let rt = tokio::runtime::Runtime::new()?;
rt.block_on(selection::run_server(SelectionServiceConfig {
port: cli.port,
threads: cli.threads,
kv_router_config: kv_router_config_from_dynamo_env(),
}))
}

#[cfg(not(feature = "select-service"))]
{
let _ = args;
anyhow::bail!(
"dynamo.select_service is not available in this build; reinstall with --features select-service"
)
}
}

#[cfg(any(
feature = "kv-indexer",
feature = "slot-tracker",
feature = "select-service"
))]
fn init_standalone_logging() {
let _ = tracing_subscriber::fmt()
.with_env_filter(
Expand Down
1 change: 1 addition & 0 deletions lib/bindings/python/src/dynamo/llm/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
from dynamo._core import register_model as register_model
from dynamo._core import run_input
from dynamo._core import run_kv_indexer as run_kv_indexer
from dynamo._core import run_select_service as run_select_service
from dynamo._core import run_slot_tracker as run_slot_tracker
from dynamo._core import unregister_model as unregister_model

Expand Down
2 changes: 2 additions & 0 deletions lib/bindings/python/src/dynamo/select_service/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
7 changes: 7 additions & 0 deletions lib/bindings/python/src/dynamo/select_service/__main__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from dynamo.select_service.main import main

if __name__ == "__main__":
raise SystemExit(main())
19 changes: 19 additions & 0 deletions lib/bindings/python/src/dynamo/select_service/main.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import sys
from collections.abc import Sequence

from dynamo.llm import run_select_service


def main(argv: Sequence[str] | None = None) -> int:
args = list(sys.argv[1:] if argv is None else argv)
try:
run_select_service(args)
except Exception as exc:
if "-h" in args or "--help" in args:
print(exc)
return 0
raise
return 0
Comment thread
PeaBrane marked this conversation as resolved.
Resolved
1 change: 1 addition & 0 deletions lib/kv-router/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ profile = []
shard-metrics = []
standalone-indexer = ["dep:axum", "dep:serde_json", "dep:reqwest", "dep:zmq"]
standalone-slot-tracker = ["dep:axum", "dep:serde_json", "dep:zmq"]
standalone-selection = ["standalone-indexer"]

[dependencies]
# repo
Expand Down
9 changes: 8 additions & 1 deletion lib/kv-router/src/services/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,11 @@

//! Standalone services built from brokerless transport primitives.

#[cfg(any(feature = "standalone-indexer", feature = "standalone-slot-tracker"))]
#[cfg(any(
feature = "standalone-indexer",
feature = "standalone-slot-tracker",
feature = "standalone-selection"
))]
pub(crate) mod zmq;

#[cfg(feature = "standalone-indexer")]
Expand All @@ -12,5 +16,8 @@ pub mod indexer;
#[cfg(feature = "standalone-indexer")]
pub mod shared_cache;

#[cfg(feature = "standalone-selection")]
pub mod selection;

#[cfg(feature = "standalone-slot-tracker")]
pub mod slot_tracker;
99 changes: 99 additions & 0 deletions lib/kv-router/src/services/selection/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
# Dynamo Selection Service

`dynamo-select-service` is a runtime-free HTTP service built on `dynamo-kv-router`.
It selects workers and tracks reservations, but it never forwards model requests
and never owns model responses.

## Readiness

- `GET /health` returns process liveness.
- `GET /ready` returns `200` only after at least one worker is schedulable.
- `GET /ready` returns `503` with worker lifecycle details while no workers are schedulable.
- `/select`, `/select_and_reserve`, `/reservations`, `/potential_loads`, and
`/overlap_scores` return `503` while the target worker pool is not ready.

## Worker Lifecycle

Workers are accepted into a single `WorkerCatalog` and become schedulable only
after the reconciler has published consistent derived state to the scheduler,
slot tracker, and indexer/listener view.

Lifecycle values:

- `incomplete`: accepted, but missing metadata required for valid booking.
- `schedulable`: metadata is complete and the scheduler can safely book it.
- `draining`: removed from the scheduler view before cleanup.
- `unschedulable`: cleanup has completed and the worker is no longer selectable.

When queueing is enabled, `max_num_batched_tokens` is required before a worker
can become schedulable. When KV events are enabled, KV event listener endpoints
are also required.

## Selection APIs

### `POST /select`

Selects a worker through the existing `SchedulerQueue` admission path with
`update_states=false`.

This endpoint is query-only. It does not book active-load state. If
`selection_id` is provided, the service accepts and echoes it for observability,
but it does not dedupe, replay, or otherwise make retries idempotent. Retrying
`/select` may produce a fresh selection.

### `POST /select_and_reserve`

Selects a worker through the same `SchedulerQueue` admission path and books
active-load state before returning.

`reservation_id` is used as the scheduler and slot-tracker request ID. If it is
omitted, the service generates one. Duplicate IDs use the existing slot-tracker
behavior and may return `409 Conflict`; V1 does not provide retry idempotency,
distributed atomicity, or exactly-once booking semantics.

### `POST /reservations`

Books an explicit post-select reservation on the supplied worker and dp-rank.
`reservation_id` is the booking request ID. Duplicate IDs use existing
slot-tracker behavior and may return `409 Conflict`.

Clients must treat reservation creation as at-most-once unless they can tolerate
duplicate/conflict handling.

### Lifecycle

- `POST /reservations/{reservation_id}/prefill_complete`
- `DELETE /reservations/{reservation_id}`

Lifecycle endpoints operate on `reservation_id`.

## Request Inputs

For `/select`, `/select_and_reserve`, `/potential_loads`, and `/overlap_scores`:

- If `token_ids` is present, the service computes block hashes and sequence
hashes internally and ignores supplied hashes. `isl_tokens` is derived from
`token_ids.len()`.
- Otherwise, callers must provide `block_hashes`, `sequence_hashes`, and
`isl_tokens`.

For `POST /reservations`:

- If `token_ids` is present, the service computes sequence hashes internally.
- Otherwise, callers must provide `sequence_hashes` and `isl_tokens`.

Signed hash values are reinterpreted bit-for-bit as unsigned router hashes,
matching the standalone indexer and slot-tracker APIs.

## Retry Semantics

V1 intentionally has no internal idempotency ledger.

- `/select` retries may return a different worker.
- `/select_and_reserve` retries with the same `reservation_id` may return
`409 Conflict` if the first booking succeeded.
- `POST /reservations` retries with the same `reservation_id` may return
`409 Conflict` if the first booking succeeded.
- Clients that need replayable retry semantics should keep their own request
ledger. The service core is structured so a core idempotency ledger can be
added later without changing the scheduler queue contract.
Loading
Loading