Custom Worker Selection

Experimental. Implement custom worker selection in a normal Rust crate and link it into a custom frontend or Endpoint Picker Provider (EPP) image. Dynamo continues to own worker discovery, queueing, eligibility, validation, accounting, reservations, and metrics.

This is a compile-time extension point. It does not load policies dynamically or through a C ABI.

Selection Boundary

Custom code replaces only scoring and picking:

Frontend or EPP
  -> SelectionService
  -> policy-class queue and scheduler
  -> host-owned worker eligibility
  -> WorkerScorer(s)
  -> WorkerPicker
  -> host validation, accounting, and reservation

Hard filters run before the custom policy. They enforce request allowlists, exact worker and data-parallel rank pins, required taints, and busy-threshold overload state. See Router Filtering for the complete eligibility behavior.

Implement a Policy

Implement one or more WorkerScorer traits and one WorkerPicker trait. Scorers contribute finite, lower-is-better costs for every eligible row. Dynamo sums those contributions before calling the picker. The picker returns one row index, which Dynamo validates before it books the request.

The following policy selects the worker with the fewest active requests:

1
use dynamo_kv_router::{
2
    KvRouterConfig, RoutingPartitionRef, WorkerCandidate, WorkerInputView, WorkerInputs,
3
    WorkerPicker, WorkerScorer, WorkerSelectionContext, WorkerSelectionPolicy,
4
    WorkerSelectionPolicyError,
5
};
6
7
struct ActiveRequestScorer;
8
9
impl WorkerScorer for ActiveRequestScorer {
10
    fn required_worker_inputs(&self) -> WorkerInputs {
11
        WorkerInputs::LOAD
12
    }
13
14
    fn score(
15
        &mut self,
16
        _context: &WorkerSelectionContext<'_>,
17
        candidate: &WorkerCandidate,
18
    ) -> Result<f64, WorkerSelectionPolicyError> {
19
        let load = candidate
20
            .load()
21
            .ok_or_else(|| WorkerSelectionPolicyError::failed("load inputs unavailable"))?;
22
        Ok(load.active_requests() as f64)
23
    }
24
}
25
26
struct MinimumCostPicker;
27
28
impl WorkerPicker for MinimumCostPicker {
29
    fn pick(
30
        &mut self,
31
        _context: &WorkerSelectionContext<'_>,
32
        input: WorkerInputView<'_>,
33
    ) -> Result<usize, WorkerSelectionPolicyError> {
34
        input
35
            .candidates()
36
            .iter()
37
            .enumerate()
38
            .min_by(|(_, left), (_, right)| left.cost().total_cmp(&right.cost()))
39
            .map(|(row, _)| row)
40
            .ok_or_else(|| WorkerSelectionPolicyError::failed("no eligible worker"))
41
    }
42
}
43
44
fn active_request_policy(
45
    config: &KvRouterConfig,
46
    worker_type: &'static str,
47
    _partition: RoutingPartitionRef<'_>,
48
) -> WorkerSelectionPolicy {
49
    WorkerSelectionPolicy::new(
50
        config.clone(),
51
        worker_type,
52
        vec![Box::new(ActiveRequestScorer)],
53
        Box::new(MinimumCostPicker),
54
    )
55
}

The factory runs when Dynamo constructs a decode or prefill worker set, not for each request. It receives the model and routing group through RoutingPartitionRef. Policy state belongs to that scheduler queue actor and is called serially.

Request Worker Inputs

Return only the signal groups that the scorer or picker reads. Dynamo creates the union requested by the composed policy and skips unused per-worker calculations.

A scorer reads its requested values from WorkerCandidate. A picker receives index-aligned columns through WorkerInputView. Candidate row order is unspecified; inspect each candidate’s worker identity or signals instead of relying on position.

Link a Custom Frontend

Keep the policy in the crate’s library target. Add a binary target that constructs the normal DistributedRuntime and EngineConfig, then inject the policy factory into Dynamo’s complete HTTP frontend:

1
HttpFrontend::default()
2
    .worker_selection_policy_factory(active_request_policy)
3
    .run(distributed_runtime, engine_config)
4
    .await?;

The binary reuses Dynamo’s model watcher, tokenizer, request pipeline, worker registry, scheduler, HTTP routes, validation, accounting, and metrics. Only the factory differs from the default frontend.

Build and run that binary in the custom image. python3 -m dynamo.frontend starts the frontend compiled into the installed Dynamo Python extension and cannot discover an external statically linked Rust crate. Without an injected factory, it uses the concrete default worker selector.

Link a Custom EPP

Build a SelectionService with the same factory and move it into EppRouter:

1
let service = SelectionServiceBuilder::new(kv_router_config)
2
    .worker_selection_policy_factory(active_request_policy)
3
    .build()
4
    .await?;
5
6
let epp_config = EppStandaloneConfig::from_env()?;
7
let router = EppRouter::from_selection_service(epp_config, service).await?;

EppRouter takes exclusive ownership of the service so its workers, peer membership, and background tasks share the EPP lifecycle. The standard EPP path remains unchanged when the caller does not supply a prebuilt service.

Policy Contract

• Return finite scorer contributions. Dynamo rejects non-finite contributions and accumulated costs.
• Return an index into WorkerInputView::candidates(). Dynamo rejects an out-of-range index before accounting or reservation.
• Treat row order as unspecified.
• Request only the signal groups the policy reads.
• Use host eligibility for hard exclusion. Scorers bias costs; the picker chooses among eligible rows.
• Keep blocking I/O out of score and pick; both execute in the scheduler queue actor.
• Do not panic in score or pick. Return WorkerSelectionPolicyError; a panic stops the scheduler queue actor.

For Dynamo’s default cost model, see Routing Concepts. For the embedded service lifecycle and reservation API, see Standalone Selection Service.