feat(scheduling): workload- and capacity-aware routing for heterogeneous rollout buckets

.gitignore

@@ -1,4 +1,6 @@
 IMPLEMENTATION_SUMMARY.md
+docs/project_structure.md
+docs/rollout_scheduling_improvements.md
 
 # Python and test caches
 __pycache__/

config.py

@@ -301,6 +301,56 @@ class SchedulingConfig:
 
     max_queue_length: int = 100
 
+    # Load metric for least-connections style selection: "requests" (legacy
+    # request-count based), "tokens" (estimated in-flight token load,
+    # prompt + EMA-expected output), or "kv_tokens" (the same token load
+    # normalized by each TP bucket's estimated KV-cache capacity in tokens).
+    # Applies to load_balance / length_aware / la_mlfq schedulers; cmlfq
+    # keeps its own workload-aware routing.
+    load_metric: str = "requests"
+
+    # Explicit KV capacity override (tokens) per TP degree, e.g. {1: 120000,
+    # 2: 230000, 4: 450000}. When empty and load_metric is "kv_tokens",
+    # capacities are auto-estimated at config-load time from the hardware /
+    # model-arch configs and the rollout gpu_memory_utilization.
+    kv_capacity_tokens_by_tp: dict = field(default_factory=dict)
+
+    # Per-GPU bytes reserved for activations / workspace / fragmentation when
+    # auto-estimating KV capacity.
+    kv_activation_reserve_gib: float = 4.0
+
+    # Closed-loop /metrics feedback (engine-layer poller feeding the
+    # schedulers). Disabled by default; when enabled, observed vLLM gauges
+    # (running/waiting, gpu cache occupancy, preemptions) drive an additive
+    # bias correction of the local load estimate plus occupancy admission
+    # with hysteresis. Preemption counters are cumulative: events penalize
+    # an instance for the TTL window instead of blacklisting forever.
+    enable_metrics_feedback: bool = False
+    metrics_poll_interval_s: float = 3.0
+    metrics_request_timeout_s: float = 1.0
+    metrics_staleness_ttl_s: float = 10.0
+    metrics_admission_enter: float = 0.90
+    metrics_admission_exit: float = 0.75
+    metrics_preemption_penalty_ttl_s: float = 60.0
+    metrics_bias_alpha: float = 0.3
+
+    # Cross-rank load aggregation: every training rank publishes its
+    # per-instance in-flight accounting to this shared directory (atomic
+    # renames + heartbeat TTL); schedulers rank instances by the summed
+    # cluster load instead of their own partial view. Empty disables.
+    shared_load_dir: str = ""
+    shared_load_ttl_s: float = 30.0
+    shared_load_heartbeat_s: float = 10.0
+    shared_load_cache_ttl_s: float = 1.0
+
+    # Prefix affinity: bind prompt_id -> instance (bounded LRU) so repeated
+    # prompts (n_samples, multi-turn replays, shared system prompts) land
+    # on the instance whose prefix cache already holds them. Affinity only
+    # reorders preference among candidates that survived readiness / queue /
+    # capacity / admission filters; it never overrides them.
+    prefix_affinity: bool = False
+    prefix_affinity_max_entries: int = 8192
+
     enable_fallback: bool = True
 
     adaptive_routing: bool = True
@@ -950,6 +1000,57 @@ def __post_init__(self):
             else:
                 self.master_addr = "localhost"
 
+
+        self._maybe_estimate_kv_capacities()
+
+
+    def _maybe_estimate_kv_capacities(self) -> None:
+        """Auto-fill scheduling.kv_capacity_tokens_by_tp when needed.
+
+        Only runs for heterogeneous rollouts using the capacity-normalized
+        ``kv_tokens`` load metric without an explicit capacity table. The
+        estimate follows vLLM's budget rule: per GPU, the utilization
+        fraction of HBM minus sharded weights and an activation reserve,
+        divided by the per-token KV bytes of the bucket's TP layout.
+        """
+        sched = self.heterogeneous_rollout.scheduling
+        if not self.heterogeneous_rollout.enabled:
+            return
+        if getattr(sched, "load_metric", "") != "kv_tokens":
+            return
+        if sched.kv_capacity_tokens_by_tp:
+            return
+
+        from RL_Framework.infra.scheduling.base import estimate_kv_capacity_tokens
+
+        arch = self.model_arch
+        weights_bytes = float(arch.num_params) * max(1, int(arch.dtype_bytes))
+        reserve = max(0.0, float(sched.kv_activation_reserve_gib)) * 1024**3
+        tp_degrees = sorted({int(i.tp) for i in self.heterogeneous_rollout.instances if int(i.tp) > 0})
+        if not tp_degrees:
+            return
+
+        capacities: dict[int, int] = {}
+        for tp in tp_degrees:
+            capacity = estimate_kv_capacity_tokens(
+                tp_degree=tp,
+                mem_capacity_bytes=self.hardware.mem_capacity,
+                gpu_memory_utilization=self.heterogeneous_rollout.gpu_memory_utilization,
+                weights_bytes=weights_bytes,
+                n_layers=arch.n_layers,
+                n_kv_heads=arch.n_kv_heads,
+                n_heads=arch.n_heads,
+                d_model=arch.d_model,
+                head_dim=arch.head_dim,
+                dtype_bytes=arch.dtype_bytes,
+                activation_reserve_bytes=reserve,
+            )
+            if capacity > 0:
+                capacities[tp] = capacity
+
+        if capacities:
+            sched.kv_capacity_tokens_by_tp = capacities
+
     # ----------------------------------------------------------------
 
     # ----------------------------------------------------------------

configs/scheduling_adaptive.yaml

@@ -0,0 +1,20 @@
+# Portable rollout-scheduling preset. Model, endpoints and GPU placement
+# belong to the deployment config or the validation command, not this file.
+# Example (existing vLLM services):
+# python examples/validate_rollout_scheduling.py --config configs/scheduling_adaptive.yaml \
+#   --model <served-model-name> --endpoint <instance-id>,<tp>,<http://host:port> \
+#   --endpoint <instance-id>,<tp>,<http://host:port> --output <report.json>
+# To use for training: inherit this file via base_config from a full experiment
+# config, and supply model_path and heterogeneous_rollout.instances there.
+heterogeneous_rollout:
+  enabled: true
+  scheduling:
+    scheduler_type: load_balance
+    load_balance_strategy: least_connections
+    load_metric: kv_tokens
+    enable_metrics_feedback: true
+    prefix_affinity: true
+    # For multi-rank runs only, use a unique run directory visible to ALL ranks:
+    # shared_load_dir: <shared-filesystem>/<unique-run>/rollout-load
+    # Measured per-instance capacities from /metrics take precedence over
+    # analytic estimates; the latter require correct hardware/model_arch inputs.

docs/configuration_reference.md

@@ -78,6 +78,66 @@ phase names allow the same converter to compare baseline runs.
 See the [cluster manual](manual.md) for recommended combinations and launch
 examples.
 
+## Rollout Scheduling Options
+
+These options live under `heterogeneous_rollout.scheduling` and control how
+requests are routed across the heterogeneous TP buckets.
+
+### Load metric
+
+| Option | Meaning |
+| --- | --- |
+| `load_metric` | Selection signal for least-connections routing: `requests` (legacy request-count, default), `tokens` (estimated in-flight token load = prompt + EMA-expected output), or `kv_tokens` (that load normalized by each TP bucket's KV-cache capacity, i.e. occupancy ratio) |
+| `kv_capacity_tokens_by_tp` | Explicit KV capacity per TP degree, e.g. `{1: 120000, 4: 450000}`. When empty and `load_metric: kv_tokens`, capacities are auto-estimated from the hardware / model-arch configs and the rollout `gpu_memory_utilization`; at engine startup they are additionally calibrated against vLLM's own profiled `kv_cache_size_tokens` from `/metrics` |
+| `kv_activation_reserve_gib` | Per-GPU reserve subtracted when auto-estimating KV capacity (default 4) |
+| `load_balance_strategy` | `least_connections` (default), `round_robin`, or `weighted` |
+
+Under `kv_tokens`, instances whose TP degree has no configured capacity are
+excluded from selection (with a warning) rather than comparing raw token
+counts against ratios. All strategies rotate among equal-load endpoints to
+avoid starving later-registered instances.
+
+The expected-output component of the estimate is a per-category EMA learned
+from observed completion lengths; it survives weight-sync rebinding, and is
+clamped by each request's `max_new_tokens` (callers that pass exact
+`input_tokens` — all bundled workflows do — skip the chars-based fallback).
+
+### Closed-loop /metrics feedback
+
+| Option | Meaning |
+| --- | --- |
+| `enable_metrics_feedback` | Master switch (default false). An engine-layer poller scrapes each instance's `/metrics`; observed gauges drive an additive bias correction of the local load estimate plus occupancy admission |
+| `metrics_poll_interval_s` | Scrape interval per instance (default 3; measured interference is <1% at 5Hz) |
+| `metrics_request_timeout_s` | Per-scrape HTTP timeout (default 1) |
+| `metrics_staleness_ttl_s` | Snapshots older than this are ignored (default 10) |
+| `metrics_admission_enter` / `metrics_admission_exit` | Occupancy hysteresis band: block new routes above `enter` (default 0.90), unblock below `exit` (default 0.75) |
+| `metrics_preemption_penalty_ttl_s` | Preemption counters are cumulative; a detected increase penalizes the instance for this window (default 60) instead of blacklisting forever |
+| `metrics_bias_alpha` | EMA weight for the additive bias `bias = EMA(observed - local)` (default 0.3) |
+
+The poller runs at the engine layer (it survives scheduler rebinding during
+weight sync), staggers endpoints with random phases (multi-rank safety), and
+any feed outage degrades transparently back to the open-loop estimate.
+
+### Cross-rank load aggregation
+
+| Option | Meaning |
+| --- | --- |
+| `shared_load_dir` | Shared directory where every training rank publishes its per-instance in-flight accounting (atomic renames + heartbeat TTL). Schedulers then rank instances by the summed cluster load instead of their own partial view, eliminating multi-rank herd. Empty disables |
+| `shared_load_ttl_s` / `shared_load_heartbeat_s` | Liveness window / refresh rate for rank files (defaults 30 / 10) |
+| `shared_load_cache_ttl_s` | Aggregate-read cache window (default 1); a rank's own writes invalidate immediately |
+
+### Prefix affinity
+
+| Option | Meaning |
+| --- | --- |
+| `prefix_affinity` | Bind `prompt_id` -> instance (bounded LRU) so repeated prompts (GRPO `n_samples`, multi-turn replays, shared system prompts) land on the instance whose prefix cache already holds them (default false) |
+| `prefix_affinity_max_entries` | LRU bound (default 8192) |
+
+Affinity only reorders preference among candidates that survived the
+readiness / queue / capacity / admission filters — a sticky instance at its
+queue cap or blocked by occupancy admission is bypassed and the mapping
+remaps. The affinity table survives weight-sync rebinding.
+
 ## Rollout Weight Reloads
 
 When `rollout_weight_sync_mode` is `restart`, the trainer publishes one reload