diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 1ef957b601f5..c4cfbdf20a3d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -14,7 +14,11 @@ # limitations under the License. default_install_hook_types: [pre-commit, commit-msg] -exclude: ^(src/grpc_generated|.*\.patch$|.*/connect/.*\.py|components/src/dynamo/planner/plugins/proto/v1/plugin_pb2(_grpc)?\.pyi?$) +# container/deps/**/patches/ holds verbatim snapshots of upstream engine source +# that a Dockerfile COPYs over the installed package. They are maintained by +# diffing against upstream, so reformatting them destroys that diff, and the +# lint findings in them are upstream's rather than ours. +exclude: ^(src/grpc_generated|.*\.patch$|.*/connect/.*\.py|container/deps/.*/patches/.*|components/src/dynamo/planner/plugins/proto/v1/plugin_pb2(_grpc)?\.pyi?$) repos: - repo: https://github.com/timothycrosley/isort rev: 5.12.0 diff --git a/components/src/dynamo/planner/core/base.py b/components/src/dynamo/planner/core/base.py index d0905719e8db..94f9a1beb77d 100644 --- a/components/src/dynamo/planner/core/base.py +++ b/components/src/dynamo/planner/core/base.py @@ -725,7 +725,8 @@ async def _collect_traffic(self) -> Optional[TrafficObservation]: if self.prometheus_traffic_client.scrape_gap_recent(self.model_name): logger.warning( "Metrics gap detected around the request counter " - "(raw num_req=%.1f); skipping throughput tick", m.num_req + "(raw num_req=%.1f); skipping throughput tick", + m.num_req, ) return None m.request_duration = self.prometheus_traffic_client.get_avg_request_duration( diff --git a/components/src/dynamo/planner/core/load_scaling.py b/components/src/dynamo/planner/core/load_scaling.py index faff7c59b177..ceefac086519 100644 --- a/components/src/dynamo/planner/core/load_scaling.py +++ b/components/src/dynamo/planner/core/load_scaling.py @@ -227,8 +227,13 @@ def _confirm_proposal( logger.info( "Confirmation buffer [%s]: last %d proposals=%s commit=%d " "observed=%d -> CONFIRMED %d (up, min of last %d)", - label, len(buffer), list(buffer), last_suggested, observed, - confirmed, ticks_up, + label, + len(buffer), + list(buffer), + last_suggested, + observed, + confirmed, + ticks_up, ) setattr(self, last_suggested_attr, confirmed) self._last_up_tick[last_suggested_attr] = tick @@ -238,8 +243,13 @@ def _confirm_proposal( logger.info( "Confirmation buffer [%s]: last %d proposals=%s commit=%d " "observed=%d -> HOLD (filling %d/%d)", - label, len(buffer), list(buffer), last_suggested, observed, - len(buffer), buffer.maxlen, + label, + len(buffer), + list(buffer), + last_suggested, + observed, + len(buffer), + buffer.maxlen, ) return last_suggested @@ -256,15 +266,24 @@ def _confirm_proposal( "Confirmation buffer [%s]: last %d proposals=%s commit=%d " "observed=%d -> HOLD (down suppressed, cooldown %d/%d " "ticks since scale-up)", - label, buffer.maxlen, list(buffer), last_suggested, - observed, tick - last_up, cooldown, + label, + buffer.maxlen, + list(buffer), + last_suggested, + observed, + tick - last_up, + cooldown, ) return last_suggested confirmed = max(buffer) logger.info( "Confirmation buffer [%s]: last %d proposals=%s commit=%d " "observed=%d -> CONFIRMED %d (down, max of buffer)", - label, buffer.maxlen, list(buffer), last_suggested, observed, + label, + buffer.maxlen, + list(buffer), + last_suggested, + observed, confirmed, ) setattr(self, last_suggested_attr, confirmed) @@ -280,7 +299,12 @@ def _confirm_proposal( logger.info( "Confirmation buffer [%s]: last %d proposals=%s commit=%d " "observed=%d -> HOLD (%s)", - label, buffer.maxlen, list(buffer), last_suggested, observed, detail, + label, + buffer.maxlen, + list(buffer), + last_suggested, + observed, + detail, ) return last_suggested @@ -493,7 +517,8 @@ def _decay_reactive_floor(self) -> None: self._reactive_floor_bump_tick = self._load_tick_counter logger.info( "Reactive prefill floor decayed to %d (no force-up for %d ticks)", - self._reactive_floor_p, decay, + self._reactive_floor_p, + decay, ) def _advance_load_disagg(self, obs: FpmObservations) -> Optional[ScalingDecision]: @@ -1070,10 +1095,12 @@ def _decode_load_decision( # Unpadded projection (pure N/(N-1) redistribution of # CURRENT demand) — input to the K* residency fixed point, # which must not compound with the spike pad (see check 3). - post_itl_raw = self._decode_regression.estimate_scheduled_decode_itl( - group, - decode_scale=consolidation, - include_queued_decode=True, + post_itl_raw = ( + self._decode_regression.estimate_scheduled_decode_itl( + group, + decode_scale=consolidation, + include_queued_decode=True, + ) ) if post_itl_raw is not None: sum_post_itl_raw_s += post_itl_raw @@ -1209,10 +1236,12 @@ def _decode_load_decision( sum_itl_2_s = 0.0 n_2 = 0 for _, group in groups_list: - itl_2 = self._decode_regression.estimate_scheduled_decode_itl( - group, - decode_scale=scale_2, - include_queued_decode=True, + itl_2 = ( + self._decode_regression.estimate_scheduled_decode_itl( + group, + decode_scale=scale_2, + include_queued_decode=True, + ) ) if itl_2 is not None: sum_itl_2_s += itl_2 diff --git a/components/src/dynamo/planner/core/perf_model/rust_adapter.py b/components/src/dynamo/planner/core/perf_model/rust_adapter.py index 1579d5d8fa11..4f7232e829a0 100644 --- a/components/src/dynamo/planner/core/perf_model/rust_adapter.py +++ b/components/src/dynamo/planner/core/perf_model/rust_adapter.py @@ -14,9 +14,9 @@ import json import logging import math -from dataclasses import dataclass import statistics from collections import deque +from dataclasses import dataclass from typing import Any, Optional from dynamo.common.forward_pass_metrics import ( @@ -664,7 +664,8 @@ def find_engine_capacity_rps( # silently refuses to answer; debugging the 1e9 prefill sentinel. logger.info( "RUST_CAPACITY[%s]: req=%s -> None", - self._worker_type, request_kwargs, + self._worker_type, + request_kwargs, ) return None # DEEPINFRA: per-query I/O log for the Rust shim. Surfaces the exact @@ -672,9 +673,13 @@ def find_engine_capacity_rps( # so we can map where the 1e9 sentinel kicks in. logger.info( "RUST_CAPACITY[%s]: req=%s -> rps=%s ttft_ms=%s itl_ms=%s e2e_ms=%s eligible=%s", - self._worker_type, request_kwargs, - result.rps, result.ttft_ms, result.itl_ms, - result.e2e_latency_ms, result.eligible, + self._worker_type, + request_kwargs, + result.rps, + result.ttft_ms, + result.itl_ms, + result.e2e_latency_ms, + result.eligible, ) rps = result.rps itl_ms = result.itl_ms @@ -691,8 +696,10 @@ def find_engine_capacity_rps( logger.info( "RUST_CAPACITY[%s]: FALLBACK_OVERRIDE shim_rps=%s -> " "rps=%.2f ttft_ms=%.2f (min_wt_s=%s slope=%s)", - self._worker_type, rps, - fallback.rps, fallback.ttft_ms or 0, + self._worker_type, + rps, + fallback.rps, + fallback.ttft_ms or 0, self._fallback_min_wt_s, self._fallback_per_token_slope_s(), ) @@ -706,9 +713,7 @@ def find_engine_capacity_rps( # single-request iter latency, then rps = 1 / iter_ttft. effective_isl = max( 1, - int(math.ceil( - isl * (1.0 - _clamp_kv_hit_rate(kv_hit_rate)) - )), + int(math.ceil(isl * (1.0 - _clamp_kv_hit_rate(kv_hit_rate)))), ) synth_fpm = ForwardPassMetrics( version=FPM_VERSION, @@ -723,13 +728,13 @@ def find_engine_capacity_rps( ), ) try: - ttft_batch1_s = self._rust_model.get_queued_prefill_time( - [synth_fpm] - ) + ttft_batch1_s = self._rust_model.get_queued_prefill_time([synth_fpm]) except _RUST_SHIM_FALLBACK_EXCEPTIONS as e: logger.warning( "RUST_CAPACITY[prefill]: batch-1 query failed: %s " - "(keeping shim rps=%.2f)", e, rps, + "(keeping shim rps=%.2f)", + e, + rps, ) ttft_batch1_s = None if ttft_batch1_s is not None and ttft_batch1_s > 0: @@ -739,8 +744,12 @@ def find_engine_capacity_rps( "RUST_CAPACITY[prefill]: batch-1 override: " "shim rps=%.2f (batch≈%.1f, ttft=%.2fms) -> " "rps=%.2f (effective_isl=%d, ttft=%.2fms)", - rps, rps * result.ttft_ms / 1000.0, result.ttft_ms, - rps_batch1, effective_isl, ttft_batch1_ms, + rps, + rps * result.ttft_ms / 1000.0, + result.ttft_ms, + rps_batch1, + effective_isl, + ttft_batch1_ms, ) rps = rps_batch1 # keep result.ttft_ms as the SLA-eligibility signal (unchanged) diff --git a/components/src/dynamo/planner/core/state_machine.py b/components/src/dynamo/planner/core/state_machine.py index 602679d634b3..c0e91c99405d 100644 --- a/components/src/dynamo/planner/core/state_machine.py +++ b/components/src/dynamo/planner/core/state_machine.py @@ -164,9 +164,7 @@ def __init__( # logic. Substituted in by LoadScalingMixin before each tick's # decision; expired after _FPM_REAL_TTL_SECONDS so a genuinely # drained worker eventually surfaces as idle. - self._last_real_fpm: dict[ - tuple[str, int], tuple[Any, float] - ] = {} + self._last_real_fpm: dict[tuple[str, int], tuple[Any, float]] = {} # Most recent observed KV hit rate from the router. Runtime metadata like # this is intentionally last-value only, not fed through the traffic load diff --git a/components/src/dynamo/planner/core/throughput_scaling.py b/components/src/dynamo/planner/core/throughput_scaling.py index bbf57943d8a8..d18a1f9b9262 100644 --- a/components/src/dynamo/planner/core/throughput_scaling.py +++ b/components/src/dynamo/planner/core/throughput_scaling.py @@ -302,9 +302,7 @@ def _prefill_replicas_erlang( logger.warning("Traffic shape provider failed: %s", e) if shape is not None and shape.isl_scv is not None: hit_scv = ( - shape.one_minus_hit_scv - if shape.one_minus_hit_scv is not None - else 0.5 + shape.one_minus_hit_scv if shape.one_minus_hit_scv is not None else 0.5 ) # eff = isl * (1 - hit), independent marginals scv_eff = (1.0 + shape.isl_scv) * (1.0 + hit_scv) - 1.0 @@ -329,8 +327,11 @@ def _prefill_replicas_erlang( "Erlang-C prefill: TTFT budget infeasible (service=%.0fms + " "overhead=%.0fms > sla=%.0fms); best-effort N=%d at " "rho_ceiling=%.2f", - service_s * 1000, cfg.prefill_ttft_overhead_ms, cfg.ttft_ms, - n_rho_ceiling, cfg.prefill_rho_ceiling, + service_s * 1000, + cfg.prefill_ttft_overhead_ms, + cfg.ttft_ms, + n_rho_ceiling, + cfg.prefill_rho_ceiling, ) self._diag_engine_rps_prefill = 1.0 / service_s return max(n_rho_ceiling, cfg.min_endpoint) @@ -342,7 +343,9 @@ def _prefill_replicas_erlang( logger.warning( "Erlang-C prefill: no N<=%d meets wait budget %.0fms; " "falling back to rho ceiling N=%d", - _ERLANG_MAX_N, wait_budget_s * 1000, n_rho_ceiling, + _ERLANG_MAX_N, + wait_budget_s * 1000, + n_rho_ceiling, ) n_queue = n_rho_ceiling @@ -376,10 +379,20 @@ def _prefill_replicas_erlang( "Prefill[erlang_c]: %.2f rps, eff_tokens=%.0f (isl=%.1f hit=%.3f), " "S=%.1fms (%s), Ca2=%.1f Cs2=%.1f (%s), offered=%.2f, " "wait_budget=%.0fms -> N=%d (queue=%d rho_ceil=%d min=%d%s)", - demand_rps, eff_tokens, isl, hit, - service_s * 1000, service_source, - cfg.prefill_arrival_scv, cs2, cs2_source, offered, - wait_budget_s * 1000, result, n_queue, n_rho_ceiling, + demand_rps, + eff_tokens, + isl, + hit, + service_s * 1000, + service_source, + cfg.prefill_arrival_scv, + cs2, + cs2_source, + offered, + wait_budget_s * 1000, + result, + n_queue, + n_rho_ceiling, cfg.min_endpoint, " hysteresis_hold" if hysteresis_held else "", ) diff --git a/components/src/dynamo/planner/monitoring/traffic_metrics.py b/components/src/dynamo/planner/monitoring/traffic_metrics.py index 80d611fee054..bc31888fa686 100644 --- a/components/src/dynamo/planner/monitoring/traffic_metrics.py +++ b/components/src/dynamo/planner/monitoring/traffic_metrics.py @@ -527,7 +527,9 @@ def scrape_gap_recent( logger.warning( "Scrape-gap check: series has %.0f/%.0f expected samples in " "the last %ds — metrics gap in progress or just ended", - worst, expected, lookback_s, + worst, + expected, + lookback_s, ) return True return False @@ -651,9 +653,7 @@ def match(labels: dict) -> bool: buckets = self._filtered_bucket_sums(f"{metric}_bucket", window, match) exact_sum = self._filtered_scalar_sum(f"{metric}_sum", window, match) exact_count = self._filtered_scalar_sum(f"{metric}_count", window, match) - exact_mean = ( - exact_sum / exact_count if exact_sum and exact_count else None - ) + exact_mean = exact_sum / exact_count if exact_sum and exact_count else None moments = _histogram_moments( buckets, log_spaced=True, calibrate_mean=exact_mean ) diff --git a/components/src/dynamo/planner/plugins/builtins/local_planner.py b/components/src/dynamo/planner/plugins/builtins/local_planner.py index 38f2ab282ea5..8f2f10943dca 100644 --- a/components/src/dynamo/planner/plugins/builtins/local_planner.py +++ b/components/src/dynamo/planner/plugins/builtins/local_planner.py @@ -256,10 +256,11 @@ def _log_rust_diagnostics(self) -> None: model handle hasn't been built. """ import json as _json + for name, attr in ( ("prefill", "_prefill_regression"), - ("decode", "_decode_regression"), - ("agg", "_agg_regression"), + ("decode", "_decode_regression"), + ("agg", "_agg_regression"), ): model = getattr(self._state, attr, None) if model is None: @@ -274,7 +275,8 @@ def _log_rust_diagnostics(self) -> None: try: log.info( "RUST_DIAG[%s]: %s", - name, _json.dumps(diag, default=str, sort_keys=True), + name, + _json.dumps(diag, default=str, sort_keys=True), ) except (TypeError, ValueError) as e: log.warning("RUST_DIAG[%s]: serialize failed: %s", name, e) diff --git a/components/src/dynamo/planner/tests/unit/test_confirm_proposal.py b/components/src/dynamo/planner/tests/unit/test_confirm_proposal.py index d668fba67f8c..207b9fd84bf7 100644 --- a/components/src/dynamo/planner/tests/unit/test_confirm_proposal.py +++ b/components/src/dynamo/planner/tests/unit/test_confirm_proposal.py @@ -16,6 +16,13 @@ from dynamo.planner.core.load_scaling import LoadScalingMixin +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + class _Gate(LoadScalingMixin): """Minimal host for the mixin: only the fields the gate touches.""" @@ -187,7 +194,9 @@ def test_reactive_floor_disabled_with_zero(): class _Trend(LoadScalingMixin): """Minimal host for the decode consolidation peak/trend pad.""" - def __init__(self, horizon: int = 360, pad_max: float = 2.0, peak_window: int = 360): + def __init__( + self, horizon: int = 360, pad_max: float = 2.0, peak_window: int = 360 + ): from collections import deque as _deque self._config = SimpleNamespace( @@ -251,7 +260,7 @@ def test_peak_pad_covers_wave(): def test_peak_pad_expires_outside_window(): t = _Trend(horizon=0, peak_window=50, pad_max=5.0) - t.feed([3_000_000] * 10) # old peak + t.feed([3_000_000] * 10) # old peak t.feed([1_000_000] * 100) # peak now outside the 50-tick window assert t._decode_consolidation_pad() == 1.0 @@ -359,9 +368,7 @@ def test_guards_combine_by_max_not_product(): class _Tolerator(LoadScalingMixin): - from dynamo.planner.core.state_machine import ( - PlannerScalingState as _PSS, - ) + from dynamo.planner.core.state_machine import PlannerScalingState as _PSS _reconcile_fpm_worker_count = staticmethod( _PSS.__dict__["_reconcile_fpm_worker_count"].__func__ diff --git a/components/src/dynamo/planner/tests/unit/test_erlang_sizing.py b/components/src/dynamo/planner/tests/unit/test_erlang_sizing.py index 1c212c0a0c4e..eb0de5a9dc45 100644 --- a/components/src/dynamo/planner/tests/unit/test_erlang_sizing.py +++ b/components/src/dynamo/planner/tests/unit/test_erlang_sizing.py @@ -21,6 +21,13 @@ from dynamo.planner.core.types import TrafficShape from dynamo.planner.monitoring.traffic_metrics import _histogram_moments +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + # --------------------------------------------------------------------------- # erlang_c # --------------------------------------------------------------------------- @@ -76,7 +83,9 @@ def test_histogram_moments_calibration_scales_mean(): assert mean == pytest.approx(200.0) # SCV is scale-invariant: matches the uncalibrated shape mean0, m20, _ = _histogram_moments(buckets, log_spaced=True) - assert (m2 - mean**2) / mean**2 == pytest.approx((m20 - mean0**2) / mean0**2) + assert (m2 - mean**2) / mean**2 == pytest.approx( + (m20 - mean0**2) / mean0**2 + ) def test_histogram_moments_empty_and_degenerate(): @@ -163,14 +172,17 @@ def test_erlang_path_infeasible_budget_uses_rho_ceiling(): def test_erlang_path_min_endpoint_floor(): - state = _make_state(slope_service=0.01, shape=None, min_endpoint=3, measure_shape=False) + state = _make_state( + slope_service=0.01, shape=None, min_endpoint=3, measure_shape=False + ) n = state._prefill_replicas_erlang(1.0, 500.0, 0.0, aic_engine_rps=100.0) assert n >= 3 def test_down_hysteresis_holds_boundary_dither(): - state = _make_state(slope_service=0.0636, shape=None, measure_shape=False, - down_pad=1.25) + state = _make_state( + slope_service=0.0636, shape=None, measure_shape=False, down_pad=1.25 + ) # demand oscillating a few percent around an integer boundary n_high = state._prefill_replicas_erlang(46.0, 5724.0, 0.484, 37.0) dithered = [ @@ -185,8 +197,9 @@ def test_down_hysteresis_holds_boundary_dither(): def test_down_hysteresis_disabled_with_pad_one(): - state = _make_state(slope_service=0.0636, shape=None, measure_shape=False, - down_pad=1.0) + state = _make_state( + slope_service=0.0636, shape=None, measure_shape=False, down_pad=1.0 + ) n_high = state._prefill_replicas_erlang(46.0, 5724.0, 0.484, 37.0) n_dip = state._prefill_replicas_erlang(40.0, 5724.0, 0.484, 37.0) # pad=1.0: padded run equals the plain run, so any strictly lower diff --git a/components/src/dynamo/planner/tests/unit/test_kstar_affine.py b/components/src/dynamo/planner/tests/unit/test_kstar_affine.py index bfed4ed51ca1..f9ea22a35bec 100644 --- a/components/src/dynamo/planner/tests/unit/test_kstar_affine.py +++ b/components/src/dynamo/planner/tests/unit/test_kstar_affine.py @@ -9,6 +9,13 @@ from dynamo.planner.core.load_scaling import _kstar_affine +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + def _itl(alpha, beta, k): return alpha + beta * k @@ -42,9 +49,7 @@ def test_divergence_detected(): k_curr, k_naive = 500_000.0, 900_000.0 itl_curr = _itl(alpha, beta, k_curr) # 21ms q_expected = k_naive * beta / itl_curr # ~1.71 - k_star, q = _kstar_affine( - k_curr, k_naive, itl_curr, _itl(alpha, beta, k_naive) - ) + k_star, q = _kstar_affine(k_curr, k_naive, itl_curr, _itl(alpha, beta, k_naive)) assert math.isinf(k_star) assert q == pytest.approx(q_expected) @@ -56,9 +61,7 @@ def test_closed_form_exceeds_truncated_iteration(): itl_curr = _itl(alpha, beta, k_curr) k1 = k_naive * _itl(alpha, beta, k_naive) / itl_curr k2 = k_naive * _itl(alpha, beta, k1) / itl_curr # 2-iter estimate - k_star, q = _kstar_affine( - k_curr, k_naive, itl_curr, _itl(alpha, beta, k_naive) - ) + k_star, q = _kstar_affine(k_curr, k_naive, itl_curr, _itl(alpha, beta, k_naive)) assert k_star > k2 > k_naive diff --git a/components/src/dynamo/planner/tests/unit/test_num_req_gate.py b/components/src/dynamo/planner/tests/unit/test_num_req_gate.py index acd3e1e7c3ea..f219f10e208c 100644 --- a/components/src/dynamo/planner/tests/unit/test_num_req_gate.py +++ b/components/src/dynamo/planner/tests/unit/test_num_req_gate.py @@ -8,17 +8,24 @@ surge — a real 3x+ spike arrives with a full complement of scrape samples. """ +import pytest + from dynamo import prometheus_names # Local dev envs may carry an older prometheus_names than the repo bindings # (which define REQUESTS_STARTED_TOTAL = "requests_started_total"). if not hasattr(prometheus_names.frontend_service, "REQUESTS_STARTED_TOTAL"): - prometheus_names.frontend_service.REQUESTS_STARTED_TOTAL = ( - "requests_started_total" - ) + prometheus_names.frontend_service.REQUESTS_STARTED_TOTAL = "requests_started_total" from dynamo.planner.monitoring.traffic_metrics import PrometheusAPIClient +pytestmark = [ + pytest.mark.gpu_0, + pytest.mark.pre_merge, + pytest.mark.unit, + pytest.mark.planner, +] + def _client(query_result=None, raises=False): c = PrometheusAPIClient.__new__(PrometheusAPIClient) diff --git a/components/src/dynamo/trtllm/publisher.py b/components/src/dynamo/trtllm/publisher.py index d830cab49b2e..230c9252c888 100644 --- a/components/src/dynamo/trtllm/publisher.py +++ b/components/src/dynamo/trtllm/publisher.py @@ -684,7 +684,9 @@ def pause_fpm(self) -> None: """ if self._fpm_paused: return - logging.info("Pausing FPM emission for drain (Python gate; publisher stays bound)") + logging.info( + "Pausing FPM emission for drain (Python gate; publisher stays bound)" + ) self._fpm_paused = True def resume_fpm(self) -> None: diff --git a/components/src/dynamo/trtllm/request_handlers/handler_base.py b/components/src/dynamo/trtllm/request_handlers/handler_base.py index a7fca7289ada..8390678b3995 100644 --- a/components/src/dynamo/trtllm/request_handlers/handler_base.py +++ b/components/src/dynamo/trtllm/request_handlers/handler_base.py @@ -1114,8 +1114,18 @@ async def _generate_locally_impl( f"Using dynamo router dp_rank={dp_rank} for TRTLLM attention DP scheduling" ) - # Priority is a float in [0.0, 1.0]; health checks use 1.0. Default is 0.5. - priority = request.get("priority", DEFAULT_REQUEST_PRIORITY) + # Priority is a float in [0.0, 1.0]; default 0.5. The top-level key is + # the canary health-check pin and takes precedence. Real requests carry + # it in the routing hints on dynamo's unbounded higher-is-urgent scale, + # so map that into the range TRT-LLM validates. + priority = request.get("priority") + if priority is None: + routed = routing.get("priority") if routing else None + priority = ( + DEFAULT_REQUEST_PRIORITY + if routed is None + else min(1.0, max(0.0, 0.5 + 0.1 * float(routed))) + ) cache_salt = request_cache_salt(request) try: diff --git a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py index f14e407b2bf4..9a1e04493e6e 100644 --- a/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py +++ b/components/src/dynamo/trtllm/tests/test_trtllm_handler_base.py @@ -716,9 +716,11 @@ class TestHealthCheckPriority: """Verify generate_locally forwards the correct priority to generate_async. Health check requests (built by TrtllmHealthCheckPayload) must reach - the TRT-LLM engine at priority=1.0. Regular inference requests - (built by the Rust frontend as PreprocessedRequest, which has no - priority field) must fall back to DEFAULT_REQUEST_PRIORITY (0.5). + the TRT-LLM engine at priority=1.0, and that top-level pin wins over any + routing hint. Regular inference requests carry priority in the routing + hints on dynamo's unbounded higher-is-urgent scale, mapped into the + [0.0, 1.0] range TRT-LLM validates; with no hint at all they fall back to + DEFAULT_REQUEST_PRIORITY (0.5). """ def _make_handler(self) -> HandlerBase: @@ -803,6 +805,55 @@ async def test_regular_request_gets_default_priority(self): _, kwargs = handler.engine.llm.generate_async.call_args assert kwargs["priority"] == DEFAULT_REQUEST_PRIORITY + @pytest.mark.asyncio + @pytest.mark.parametrize( + "routed,expected", + [(1, 0.6), (2, 0.7), (5, 1.0), (50, 1.0), (0, 0.5)], + ) + async def test_routing_priority_mapped_into_engine_range(self, routed, expected): + """routing.priority is unbounded and higher-is-urgent; TRT-LLM takes + [0.0, 1.0]. Values at or above 5 saturate.""" + handler = self._make_handler() + generation_result = self._make_mock_generation_result() + handler.engine.llm.generate_async = MagicMock(return_value=generation_result) + + request = { + "token_ids": [1, 2, 3], + "stop_conditions": {"max_tokens": 10}, + "sampling_options": {"temperature": 0.7}, + "routing": {"priority": routed}, + } + + chunks = [ + c async for c in handler.generate_locally(request, self._make_context()) + ] + assert len(chunks) > 0 + + handler.engine.llm.generate_async.assert_called_once() + _, kwargs = handler.engine.llm.generate_async.call_args + assert kwargs["priority"] == pytest.approx(expected) + + @pytest.mark.asyncio + async def test_health_check_priority_beats_routing_hint(self): + """The canary pin is top-level and must not be overridden by routing.""" + handler = self._make_handler() + generation_result = self._make_mock_generation_result() + handler.engine.llm.generate_async = MagicMock(return_value=generation_result) + + request = TrtllmHealthCheckPayload( + disaggregation_mode=DisaggregationMode.AGGREGATED, + ).to_dict() + request["routing"] = {"priority": 0} + + chunks = [ + c async for c in handler.generate_locally(request, self._make_context()) + ] + assert len(chunks) > 0 + + handler.engine.llm.generate_async.assert_called_once() + _, kwargs = handler.engine.llm.generate_async.call_args + assert kwargs["priority"] == 1.0 + @pytest.mark.asyncio async def test_routing_cache_salt_forwarded_to_generate_async(self): handler = self._make_handler()