diff --git a/python/sglang/benchmark/offline_throughput.py b/python/sglang/benchmark/offline_throughput.py index 9612ecc39..35395a1b4 100644 --- a/python/sglang/benchmark/offline_throughput.py +++ b/python/sglang/benchmark/offline_throughput.py @@ -284,7 +284,15 @@ def throughput_test_once( dir = os.getenv("SGLANG_TORCH_PROFILER_DIR") if not profile_steps: known_files = set(os.listdir(dir)) + # With --profile-steps the scheduler auto-stops mid-run after N steps, so + # a second stop here raises "not in progress"; a run shorter than N steps + # never hit the target and still needs this explicit stop. Either way we + # must stop before monitor_trace_file, which loops forever waiting for a + # trace that would otherwise never be finalized. + try: backend.stop_profile() + except RuntimeError: + pass monitor_trace_file(known_files, dir) if backend_name == "runtime": diff --git a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py index af1f7c7c6..58ab0ff7f 100644 --- a/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py +++ b/python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py @@ -86,6 +86,16 @@ class SchedulerMlxOverlapMixin: """Mixin that adds MLX overlap scheduling to :class:`Scheduler`.""" def _finalize_mlx_pending_job(self: Scheduler, pending: MlxPendingJob): + # Account for this completed forward step. The standard scheduler does + # this inside run_batch(), but the MLX overlap loop bypasses run_batch, + # so without this forward_ct never advances on MLX. That stalls the + # watchdog liveness counter and, more importantly, breaks step-bounded + # profiling: _profile_batch_predicate auto-starts/stops based on + # forward_ct, so `--profile-steps` (and the server /start_profile + # num_steps path) only takes effect once the counter moves here. + self.forward_ct += 1 + self.profiler_manager._profile_batch_predicate(pending.schedule_batch) + result = self.tp_worker.finalize_mlx_result( pending.prefills, pending.extends, diff --git a/test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py b/test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py new file mode 100644 index 000000000..7d53834f6 --- /dev/null +++ b/test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py @@ -0,0 +1,78 @@ +"""Unit tests for the MLX overlap scheduler mixin (hardware_backend/mlx/scheduler_mixin.py). + +Covers: + - _finalize_mlx_pending_job advances forward_ct once per completed step + - _finalize_mlx_pending_job calls the profiler batch predicate with the + finalized batch, so step-bounded profiling (``--profile-steps`` / + ``/start_profile`` num_steps) can auto-stop on the MLX overlap loop, which + bypasses the standard Scheduler.run_batch(). + +Skips on non-Apple-Silicon platforms and when ``mlx`` is missing (importing +scheduler_mixin requires ``mlx.core``). +""" + +from __future__ import annotations + +import importlib.util +import platform +import unittest +from unittest.mock import MagicMock + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64" +_HAS_MLX = importlib.util.find_spec("mlx") is not None +_SKIP_REASON = "requires Apple Silicon and mlx" + + +@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON) +class TestFinalizeMlxPendingJob(unittest.TestCase): + """forward_ct accounting + profiler predicate wiring in the overlap loop.""" + + def _make_scheduler(self): + scheduler = MagicMock() + scheduler.forward_ct = 0 + result = MagicMock() + result.next_token_ids = None + scheduler.tp_worker.finalize_mlx_result.return_value = result + return scheduler + + def test_finalize_advances_forward_ct_and_runs_predicate(self): + from sglang.srt.hardware_backend.mlx.scheduler_mixin import ( + SchedulerMlxOverlapMixin, + ) + + scheduler = self._make_scheduler() + pending = MagicMock() + + SchedulerMlxOverlapMixin._finalize_mlx_pending_job(scheduler, pending) + + # Standard run_batch() advances forward_ct and runs the profiler + # predicate; the MLX overlap loop must do the same here. + self.assertEqual(scheduler.forward_ct, 1) + scheduler.profiler_manager._profile_batch_predicate.assert_called_once_with( + pending.schedule_batch + ) + # The rest of finalization still runs. + scheduler.process_batch_result.assert_called_once() + + def test_forward_ct_advances_once_per_step(self): + from sglang.srt.hardware_backend.mlx.scheduler_mixin import ( + SchedulerMlxOverlapMixin, + ) + + scheduler = self._make_scheduler() + + for expected in (1, 2, 3): + SchedulerMlxOverlapMixin._finalize_mlx_pending_job(scheduler, MagicMock()) + self.assertEqual(scheduler.forward_ct, expected) + + self.assertEqual( + scheduler.profiler_manager._profile_batch_predicate.call_count, 3 + ) + + +if __name__ == "__main__": + unittest.main()