Bundle Scheduler rank/size fields into a frozen ParallelState (#25444)

test/registered/unit/observability/test_forward_pass_metrics.py

@@ -7,12 +7,38 @@ import unittest
 from unittest.mock import patch
 
 from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.observability.scheduler_metrics_mixin import (
     PrefillStats,
     SchedulerMetricsMixin,
 )
 
 
+def _make_ps(**overrides) -> ParallelState:
+    """Build a ParallelState with reasonable defaults for tests; override fields via kwargs."""
+    defaults = dict(
+        tp_rank=0,
+        tp_size=1,
+        pp_rank=0,
+        pp_size=1,
+        dp_rank=None,
+        dp_size=1,
+        attn_tp_rank=0,
+        attn_tp_size=1,
+        attn_cp_rank=0,
+        attn_cp_size=1,
+        attn_dp_rank=0,
+        attn_dp_size=1,
+        moe_ep_rank=0,
+        moe_ep_size=1,
+        moe_dp_rank=None,
+        moe_dp_size=1,
+        gpu_id=0,
+    )
+    defaults.update(overrides)
+    return ParallelState(**defaults)
+
+
 class _FakeReq:
     def __init__(
         self,
@@ -221,10 +247,7 @@ class TestForwardPassMetrics(unittest.TestCase):
             forward_pass_metrics_ipc_name=None,
             kv_events_config=None,
         )
-        scheduler.attn_tp_rank = 0
-        scheduler.dp_rank = 2
-        scheduler.pp_rank = 0
-        scheduler.pp_size = 1
+        scheduler.ps = _make_ps(attn_tp_rank=0, dp_rank=2, pp_rank=0, pp_size=1)
         scheduler.enable_kv_cache_events = False
 
         with patch(
@@ -252,10 +275,7 @@ class TestForwardPassMetrics(unittest.TestCase):
             forward_pass_metrics_ipc_name=None,
             kv_events_config=None,
         )
-        scheduler.attn_tp_rank = 0
-        scheduler.dp_rank = 0
-        scheduler.pp_rank = 0
-        scheduler.pp_size = 2
+        scheduler.ps = _make_ps(attn_tp_rank=0, dp_rank=0, pp_rank=0, pp_size=2)
         scheduler.enable_kv_cache_events = False
 
         with patch(