[CI][RFC] Replace black-jupyter with ruff-format (#37210)

Co-authored-by: Alison Shao <a.shao@wustl.edu>

test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py

@@ -60,13 +60,17 @@ class TestInitProfileBatchMode(CustomTestCase):
         env = {_BATCH_CAPTURE: "1"}
         if profiler_dir is not None:
             env["SGLANG_TORCH_PROFILER_DIR"] = profiler_dir
-        with mock.patch.dict(os.environ, env, clear=False), mock.patch.object(
-            mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank)
-        ), mock.patch.object(mod, "profile") as mock_profile, mock.patch(
-            "torch.profiler.schedule"
-        ) as mock_schedule, mock.patch(
-            "torch.cuda.memory._record_memory_history"
-        ) as mock_record_history:
+        with (
+            mock.patch.dict(os.environ, env, clear=False),
+            mock.patch.object(
+                mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank)
+            ),
+            mock.patch.object(mod, "profile") as mock_profile,
+            mock.patch("torch.profiler.schedule") as mock_schedule,
+            mock.patch(
+                "torch.cuda.memory._record_memory_history"
+            ) as mock_record_history,
+        ):
             os.environ.pop(_CAPTURE_TRACE, None)  # original flag off
             if profiler_dir is None:
                 os.environ.pop("SGLANG_TORCH_PROFILER_DIR", None)
@@ -110,19 +114,16 @@ class TestInitProfileBatchMode(CustomTestCase):
         # No SGLANG_TORCH_PROFILER_DIR -> falls back to the envs default base dir.
         # Patch makedirs so the test never writes to the cwd.
         fake_self = _make_fake_self([1])
-        with mock.patch.dict(
-            os.environ, {_BATCH_CAPTURE: "1"}, clear=False
-        ), mock.patch.object(
-            mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0)
-        ), mock.patch.object(
-            mod, "profile"
-        ), mock.patch(
-            "torch.profiler.schedule"
-        ), mock.patch(
-            "torch.cuda.memory._record_memory_history"
-        ), mock.patch.object(
-            mod.os, "makedirs"
-        ) as mock_makedirs:
+        with (
+            mock.patch.dict(os.environ, {_BATCH_CAPTURE: "1"}, clear=False),
+            mock.patch.object(
+                mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0)
+            ),
+            mock.patch.object(mod, "profile"),
+            mock.patch("torch.profiler.schedule"),
+            mock.patch("torch.cuda.memory._record_memory_history"),
+            mock.patch.object(mod.os, "makedirs") as mock_makedirs,
+        ):
             os.environ.pop("SGLANG_TORCH_PROFILER_DIR", None)
             os.environ.pop(_CAPTURE_TRACE, None)
             DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self)
@@ -143,12 +144,14 @@ class TestInitProfileOriginalMode(CustomTestCase):
         with tempfile.TemporaryDirectory() as tmp:
             environ = dict(env)
             environ["SGLANG_TORCH_PROFILER_DIR"] = tmp
-            with mock.patch.dict(os.environ, environ, clear=False), mock.patch.object(
-                mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0)
-            ), mock.patch.object(mod, "profile") as mock_profile, mock.patch(
-                "torch.profiler.schedule"
-            ) as mock_schedule, mock.patch(
-                "torch.cuda.memory._record_memory_history"
+            with (
+                mock.patch.dict(os.environ, environ, clear=False),
+                mock.patch.object(
+                    mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0)
+                ),
+                mock.patch.object(mod, "profile") as mock_profile,
+                mock.patch("torch.profiler.schedule") as mock_schedule,
+                mock.patch("torch.cuda.memory._record_memory_history"),
             ):
                 for k in (_CAPTURE_TRACE, _BATCH_CAPTURE):
                     if k not in environ:
@@ -176,18 +179,18 @@ class TestInitProfileOriginalMode(CustomTestCase):
 class TestOnTraceReadyNaming(CustomTestCase):
     def _build_on_trace_ready(self, *, capture_bs, rank, tmp):
         fake_self = _make_fake_self(capture_bs)
-        with mock.patch.dict(
-            os.environ,
-            {"SGLANG_TORCH_PROFILER_DIR": tmp, _BATCH_CAPTURE: "1"},
-            clear=False,
-        ), mock.patch.object(
-            mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank)
-        ), mock.patch.object(
-            mod, "profile"
-        ) as mock_profile, mock.patch(
-            "torch.profiler.schedule"
-        ), mock.patch(
-            "torch.cuda.memory._record_memory_history"
+        with (
+            mock.patch.dict(
+                os.environ,
+                {"SGLANG_TORCH_PROFILER_DIR": tmp, _BATCH_CAPTURE: "1"},
+                clear=False,
+            ),
+            mock.patch.object(
+                mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank)
+            ),
+            mock.patch.object(mod, "profile") as mock_profile,
+            mock.patch("torch.profiler.schedule"),
+            mock.patch("torch.cuda.memory._record_memory_history"),
         ):
             os.environ.pop(_CAPTURE_TRACE, None)
             DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self)

test/registered/unit/model_executor/runner/test_hidden_state_graph_recapture.py

@@ -118,9 +118,12 @@ class TestHiddenStateGraphRecapture(CustomTestCase):
         ):
             runner = self._make_runner(runner_cls, CaptureHiddenMode.NULL)
 
-            with self.subTest(runner_cls=runner_cls), self.assertRaisesRegex(
-                RuntimeError,
-                "exceeds the fixed (CUDA|CPU) graph capture mode",
+            with (
+                self.subTest(runner_cls=runner_cls),
+                self.assertRaisesRegex(
+                    RuntimeError,
+                    "exceeds the fixed (CUDA|CPU) graph capture mode",
+                ),
             ):
                 runner._validate_capture_hidden_mode(
                     self._make_forward_batch(CaptureHiddenMode.LAST)

test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py

@@ -139,8 +139,11 @@ class TestCaptureOneWithProfiling(CustomTestCase):
             rf_names.append(name)
             return contextlib.nullcontext()
 
-        with mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"), mock.patch(
-            "torch.profiler.record_function", side_effect=_fake_record_function
+        with (
+            mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"),
+            mock.patch(
+                "torch.profiler.record_function", side_effect=_fake_record_function
+            ),
         ):
             backend.capture_one(ShapeKey(size=size), forward_fn)
 

test/registered/unit/model_executor/test_model_runner_decode_rows.py

@@ -35,14 +35,18 @@ class TestModelRunnerDecodeRows(unittest.TestCase):
             spec = SimpleNamespace(
                 speculative_adaptive=True, speculative_adaptive_config=f.name
             )
-            with patch(
-                "sglang.srt.model_executor.model_runner.get_spec", return_value=spec
-            ), patch(
-                "sglang.srt.model_executor.model_runner.max_speculative_num_draft_tokens",
-                return_value=6,
-            ), patch(
-                "sglang.srt.model_executor.model_runner.get_batch_sizes_to_capture",
-                side_effect=_alignment_8_capture_bs,
+            with (
+                patch(
+                    "sglang.srt.model_executor.model_runner.get_spec", return_value=spec
+                ),
+                patch(
+                    "sglang.srt.model_executor.model_runner.max_speculative_num_draft_tokens",
+                    return_value=6,
+                ),
+                patch(
+                    "sglang.srt.model_executor.model_runner.get_batch_sizes_to_capture",
+                    side_effect=_alignment_8_capture_bs,
+                ),
             ):
                 self.assertEqual(runner.max_decode_logits_rows(), 72)
 

test/registered/unit/model_executor/test_pool_configurator.py

@@ -813,8 +813,9 @@ class TestDSAIndexerAllocationPolicy(CustomTestCase):
         mr.model_config.hf_config.index_topk_freq = 4
         mr.model_config.hf_config.index_skip_topk_offset = 3
 
-        with get_memory().override(enable_hierarchical_cache=True), mock_cpu_env(
-            kv_size=1
+        with (
+            get_memory().override(enable_hierarchical_cache=True),
+            mock_cpu_env(kv_size=1),
         ):
             from sglang.srt.model_executor.pool_configurator import (
                 DefaultPoolConfigurator,