277 lines
9.9 KiB
Python
277 lines
9.9 KiB
Python
"""Unit tests for MLX Metal profiling patch (hardware_backend/mlx/profiler.py).
|
|
|
|
Covers:
|
|
- apply_metal_profiler_patches() replaces torch.profiler.profile
|
|
- MLX path: MetalTorchProfiler.start/stop produces a .gputrace file
|
|
- MPS path: MetalTorchProfiler wraps torch.mps.profiler.metal_capture
|
|
- RuntimeError from start_capture is caught and returned as success=False
|
|
- SchedulerProfilerManager._start_profile returns success=False gracefully
|
|
when Metal capture fails (no MTL_CAPTURE_ENABLED)
|
|
|
|
Skips on non-Apple-Silicon platforms and when ``mlx`` is missing.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import platform
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from sglang.srt.runtime_context import get_parallel
|
|
from sglang.test.ci.ci_register import register_mlx_ci
|
|
|
|
register_mlx_ci(est_time=5, suite="stage-a-unit-test-mlx")
|
|
|
|
_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 TestApplyMetalProfilerPatches(unittest.TestCase):
|
|
"""apply_metal_profiler_patches() replaces torch.profiler.profile."""
|
|
|
|
def setUp(self):
|
|
import torch
|
|
|
|
self._original_profile = getattr(
|
|
torch.profiler.profile, "_sglang_original_profile", None
|
|
)
|
|
|
|
def tearDown(self):
|
|
import torch
|
|
|
|
if self._original_profile is not None:
|
|
torch.profiler.profile = self._original_profile
|
|
|
|
def test_patch_replaces_profile(self):
|
|
import torch
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import (
|
|
MetalTorchProfiler,
|
|
apply_metal_profiler_patches,
|
|
)
|
|
|
|
apply_metal_profiler_patches()
|
|
self.assertTrue(getattr(torch.profiler.profile, "_sglang_metal_patched", False))
|
|
p = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA])
|
|
self.assertIsInstance(p, MetalTorchProfiler)
|
|
|
|
def test_patch_is_idempotent(self):
|
|
import torch
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import (
|
|
apply_metal_profiler_patches,
|
|
)
|
|
|
|
apply_metal_profiler_patches()
|
|
first = torch.profiler.profile
|
|
apply_metal_profiler_patches()
|
|
self.assertIs(torch.profiler.profile, first)
|
|
|
|
def test_no_cuda_activity_uses_original(self):
|
|
import torch
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import (
|
|
MetalTorchProfiler,
|
|
apply_metal_profiler_patches,
|
|
)
|
|
|
|
apply_metal_profiler_patches()
|
|
p = torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CPU])
|
|
self.assertNotIsInstance(p, MetalTorchProfiler)
|
|
|
|
|
|
@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
|
|
class TestMetalCaptureProfilerMLX(unittest.TestCase):
|
|
"""MLX path: start_mlx produces a .gputrace and stop_capture is called."""
|
|
|
|
def test_start_mlx_success(self):
|
|
import mlx.core as mx
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
trace_path = Path(tmp) / "test.gputrace"
|
|
with (
|
|
patch.object(mx.metal, "start_capture"),
|
|
patch.object(mx.metal, "stop_capture"),
|
|
):
|
|
profiler, result = MetalCaptureProfiler.start_mlx(trace_path)
|
|
|
|
self.assertTrue(result.success)
|
|
self.assertIsNotNone(profiler)
|
|
self.assertEqual(profiler.label, "MLX")
|
|
self.assertTrue(profiler.standalone)
|
|
|
|
def test_start_mlx_runtime_error_returns_failure(self):
|
|
import mlx.core as mx
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
trace_path = Path(tmp) / "test.gputrace"
|
|
with patch.object(
|
|
mx.metal,
|
|
"start_capture",
|
|
side_effect=RuntimeError("Capture layer is not inserted"),
|
|
):
|
|
profiler, result = MetalCaptureProfiler.start_mlx(trace_path)
|
|
|
|
self.assertIsNone(profiler)
|
|
self.assertFalse(result.success)
|
|
self.assertIn("MTL_CAPTURE_ENABLED", result.message)
|
|
|
|
def test_stop_calls_stop_capture(self):
|
|
import mlx.core as mx
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
trace_path = Path(tmp) / "test.gputrace"
|
|
with (
|
|
patch.object(mx.metal, "start_capture"),
|
|
patch.object(mx.metal, "stop_capture") as mock_stop,
|
|
):
|
|
profiler, _ = MetalCaptureProfiler.start_mlx(trace_path)
|
|
profiler.stop()
|
|
mock_stop.assert_called_once()
|
|
|
|
|
|
@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
|
|
class TestMetalCaptureProfilerMPS(unittest.TestCase):
|
|
"""MPS path: start_mps wraps torch.mps.profiler.metal_capture."""
|
|
|
|
def test_start_mps_success(self):
|
|
import torch
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
|
|
|
|
mock_ctx = MagicMock()
|
|
mock_ctx.__enter__ = MagicMock(return_value=mock_ctx)
|
|
mock_ctx.__exit__ = MagicMock(return_value=False)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
trace_path = Path(tmp) / "test.gputrace"
|
|
with patch.object(
|
|
torch.mps.profiler, "metal_capture", return_value=mock_ctx
|
|
):
|
|
profiler, result = MetalCaptureProfiler.start_mps(trace_path)
|
|
|
|
self.assertTrue(result.success)
|
|
self.assertIsNotNone(profiler)
|
|
self.assertEqual(profiler.label, "MPS")
|
|
self.assertFalse(profiler.standalone)
|
|
|
|
def test_start_mps_runtime_error_returns_failure(self):
|
|
import torch
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import MetalCaptureProfiler
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
trace_path = Path(tmp) / "test.gputrace"
|
|
with patch.object(
|
|
torch.mps.profiler,
|
|
"metal_capture",
|
|
side_effect=RuntimeError("MPS profiler unavailable"),
|
|
):
|
|
profiler, result = MetalCaptureProfiler.start_mps(trace_path)
|
|
|
|
self.assertIsNone(profiler)
|
|
self.assertFalse(result.success)
|
|
self.assertIn("MTL_CAPTURE_ENABLED", result.message)
|
|
|
|
|
|
@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
|
|
class TestSchedulerProfilerManagerMPS(unittest.TestCase):
|
|
"""SchedulerProfilerManager._start_profile handles Metal capture failures.
|
|
|
|
apply_metal_profiler_patches() dispatches on use_mlx(), which reads
|
|
SGLANG_USE_MLX and is cached for the process (tensor_bridge.use_mlx).
|
|
This class exists to exercise the MPS strategy, so it pins use_mlx at
|
|
its point of use in profiler.py rather than the ambient environment;
|
|
an env-only pin would not reliably override an already-cached value.
|
|
"""
|
|
|
|
def setUp(self):
|
|
patcher = patch(
|
|
"sglang.srt.hardware_backend.mlx.profiler.use_mlx", return_value=False
|
|
)
|
|
patcher.start()
|
|
self.addCleanup(patcher.stop)
|
|
|
|
def _make_manager(self, output_dir):
|
|
from sglang.srt.managers.scheduler_components.profiler_manager import (
|
|
SchedulerProfilerManager,
|
|
)
|
|
|
|
mgr = SchedulerProfilerManager(dp_tp_cpu_group=None, get_forward_ct=lambda: 0)
|
|
mgr._init_profile(output_dir, None, None, None, None, None, False, "test")
|
|
return mgr
|
|
|
|
# MetalCaptureProfiler has two strategies: start_mlx drives
|
|
# mx.metal.start_capture, start_mps drives torch.mps.profiler.metal_capture.
|
|
# This manager takes the MPS one, so that is the symbol to stand in for --
|
|
# patching mx.metal here leaves the real Metal capture running, which fails
|
|
# with "Capture layer is not inserted" unless MTL_CAPTURE_ENABLED=1 is set
|
|
# in the environment.
|
|
def test_start_profile_failure_does_not_crash(self):
|
|
import torch
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import (
|
|
apply_metal_profiler_patches,
|
|
)
|
|
|
|
apply_metal_profiler_patches()
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
mgr = self._make_manager(tmp)
|
|
with patch.object(
|
|
torch.mps.profiler,
|
|
"metal_capture",
|
|
side_effect=RuntimeError("Capture layer is not inserted"),
|
|
):
|
|
result = mgr._start_profile()
|
|
|
|
self.assertFalse(result.success)
|
|
self.assertIn("Capture layer is not inserted", result.message)
|
|
self.assertFalse(mgr.profile_in_progress)
|
|
self.assertIsNone(mgr.torch_profiler)
|
|
|
|
def test_start_profile_success_with_mock_capture(self):
|
|
from unittest.mock import MagicMock
|
|
from unittest.mock import patch as mock_patch
|
|
|
|
import torch
|
|
|
|
from sglang.srt.hardware_backend.mlx.profiler import (
|
|
apply_metal_profiler_patches,
|
|
)
|
|
|
|
apply_metal_profiler_patches()
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
mgr = self._make_manager(tmp)
|
|
capture_ctx = MagicMock()
|
|
with (
|
|
mock_patch.object(
|
|
torch.mps.profiler, "metal_capture", return_value=capture_ctx
|
|
),
|
|
mock_patch("torch.distributed.barrier"),
|
|
get_parallel().override(tp_rank=0, dp_size=1, pp_size=1, moe_ep_size=1),
|
|
):
|
|
result = mgr._start_profile()
|
|
self.assertTrue(result.success, result.message)
|
|
self.assertTrue(mgr.profile_in_progress)
|
|
capture_ctx.__enter__.assert_called_once()
|
|
mgr._stop_profile()
|
|
self.assertFalse(mgr.profile_in_progress)
|
|
capture_ctx.__exit__.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|