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

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

python/sglang/srt/model_executor/cpu_graph_runner.py

@@ -137,9 +137,9 @@ def get_batch_sizes_to_capture(model_runner: ModelRunner):
     # Users can customize the batch sizes supported by cpu_graph, such as:
     # --cuda-graph-bs-decode 1 2 4 8 16
     capture_bs = get_exec().graph.cuda_graph_config.decode.bs
-    assert (
-        max(capture_bs) <= get_exec().graph.torch_compile_max_bs
-    ), f"{capture_bs=}, {get_exec().graph.torch_compile_max_bs=}"
+    assert max(capture_bs) <= get_exec().graph.torch_compile_max_bs, (
+        f"{capture_bs=}, {get_exec().graph.torch_compile_max_bs=}"
+    )
     capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size]
     capture_bs = list(sorted(set(capture_bs)))
     assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}"
@@ -619,21 +619,21 @@ class CPUGraphRunner:
         self.captured_req_width = 1
 
         assert not get_lora().enable_lora, "CPUGraphRunner does not support LoRA yet."
-        assert (
-            not self.enable_two_batch_overlap
-        ), "CPUGraphRunner does not support two batch overlap yet."
-        assert (
-            not self.require_mlp_tp_gather
-        ), "CPUGraphRunner does not support MLP TP gather yet."
-        assert (
-            not self.require_mlp_sync
-        ), "CPUGraphRunner does not support MLP sync yet."
-        assert (
-            not self.require_gathered_buffer
-        ), "CPUGraphRunner does not support gathered buffer yet."
-        assert (
-            model_runner.spec_algorithm.is_none()
-        ), "CPUGraphRunner does not support speculative inference yet."
+        assert not self.enable_two_batch_overlap, (
+            "CPUGraphRunner does not support two batch overlap yet."
+        )
+        assert not self.require_mlp_tp_gather, (
+            "CPUGraphRunner does not support MLP TP gather yet."
+        )
+        assert not self.require_mlp_sync, (
+            "CPUGraphRunner does not support MLP sync yet."
+        )
+        assert not self.require_gathered_buffer, (
+            "CPUGraphRunner does not support gathered buffer yet."
+        )
+        assert model_runner.spec_algorithm.is_none(), (
+            "CPUGraphRunner does not support speculative inference yet."
+        )
 
         assert self.dp_size == 1, "CPUGraphRunner does not support DP yet."
         assert self.pp_size == 1, "CPUGraphRunner does not support PP yet."
@@ -952,9 +952,9 @@ class CPUGraphRunner:
         forward_batch: ForwardBatch,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
-        assert (
-            pp_proxy_tensors is None
-        ), "PPProxyTensors is not supported in CPUGraphRunner yet."
+        assert pp_proxy_tensors is None, (
+            "PPProxyTensors is not supported in CPUGraphRunner yet."
+        )
 
         replay_context = (
             model_capture_mode if self.is_encoder_decoder else empty_context

python/sglang/srt/model_executor/forward_batch_info.py

@@ -1420,9 +1420,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
                 # branch handles decode rows padded to a 1-token extend.
                 if hybrid_ssm or self.seq_lens.shape[0] == 0:
                     dev = self.seq_lens.device
-                    assert (
-                        self.seq_lens.shape[0] == 0
-                    ), "extend-idle conversion expects an empty rank"
+                    assert self.seq_lens.shape[0] == 0, (
+                        "extend-idle conversion expects an empty rank"
+                    )
                     self.extend_num_tokens = num_tokens
                     self.extend_seq_lens = torch.tensor(
                         [num_tokens], dtype=torch.int32, device=dev

python/sglang/srt/model_executor/hook_manager.py

@@ -46,14 +46,13 @@ def register_forward_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -
 
         if not matched:
             logger.warning(
-                f"No modules matched hook spec '{spec_name}' "
-                f"patterns={target_patterns}"
+                f"No modules matched hook spec '{spec_name}' patterns={target_patterns}"
             )
             continue
 
         for module_name, module in matched:
             _ = module.register_forward_hook(hook)
-            logger.info(f"Registered forward hook '{spec_name}' " f"on {module_name}")
+            logger.info(f"Registered forward hook '{spec_name}' on {module_name}")
 
 
 def resolve_callable(path: Optional[str]) -> Optional[Callable]:

python/sglang/srt/model_executor/input_buffers.py

@@ -58,7 +58,6 @@ INDEX_SEMANTIC_BUFFERS = frozenset(
 
 @dataclass
 class ForwardInputBuffers:
-
     def reset_index_buffers(self) -> None:
         """Zero the index-semantic buffers this set declares."""
         for f in fields(self):
@@ -87,14 +86,14 @@ class ForwardInputBuffers:
 
             if isinstance(buffer, dict):
                 for sub_name, sub_buffer in buffer.items():
-                    assert isinstance(
-                        sub_buffer, torch.Tensor
-                    ), f"Field {name}.{sub_name} is expected to be a torch.Tensor, but got {type(sub_buffer)}."
+                    assert isinstance(sub_buffer, torch.Tensor), (
+                        f"Field {name}.{sub_name} is expected to be a torch.Tensor, but got {type(sub_buffer)}."
+                    )
                     buffer[sub_name] = share_input_buffer(
                         f"{name}.{sub_name}", sub_buffer
                     )
             else:
-                assert isinstance(
-                    buffer, torch.Tensor
-                ), f"Field {name} is expected to be a torch.Tensor, a dict of torch.Tensor, or a dataclass of torch.Tensor, but got {type(buffer)}."
+                assert isinstance(buffer, torch.Tensor), (
+                    f"Field {name} is expected to be a torch.Tensor, a dict of torch.Tensor, or a dataclass of torch.Tensor, but got {type(buffer)}."
+                )
                 setattr(self, name, share_input_buffer(name, buffer))

python/sglang/srt/model_executor/model_runner.py

@@ -310,8 +310,7 @@ class ModelRunner:
     def sampling_observer(self, observer: Optional[SamplingObserver]) -> None:
         if observer is not None and not self.supports_sampling_observer():
             raise ValueError(
-                "sampling observers are not supported by the configured "
-                "sampling path"
+                "sampling observers are not supported by the configured sampling path"
             )
         self._sampling_observer = observer
 
@@ -481,9 +480,9 @@ class ModelRunner:
         )
 
         if self.ps.pp_size > 1:
-            assert (
-                self.support_pp
-            ), "Pipeline Parallel is not compatible with this model."
+            assert self.support_pp, (
+                "Pipeline Parallel is not compatible with this model."
+            )
 
         # For weight updates
         self.init_weight_updater()

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -58,23 +58,23 @@ logger = logging.getLogger(__name__)
 def _align_pipeline_layers(layers: list, layer_model) -> list:
     has_start_layer = hasattr(layer_model, "start_layer")
     has_end_layer = hasattr(layer_model, "end_layer")
-    assert (
-        has_start_layer == has_end_layer
-    ), "pipeline layer ranges must define start_layer and end_layer together"
+    assert has_start_layer == has_end_layer, (
+        "pipeline layer ranges must define start_layer and end_layer together"
+    )
     start_layer = layer_model.start_layer if has_start_layer else 0
     end_layer = layer_model.end_layer if has_end_layer else len(layer_model.layers)
-    assert isinstance(start_layer, int) and isinstance(
-        end_layer, int
-    ), "pipeline layer ranges must define integer start_layer and end_layer"
+    assert isinstance(start_layer, int) and isinstance(end_layer, int), (
+        "pipeline layer ranges must define integer start_layer and end_layer"
+    )
     assert 0 <= start_layer <= end_layer <= len(layer_model.layers), (
         f"invalid pipeline layer range [{start_layer}, {end_layer}) for "
         f"{len(layer_model.layers)} layers"
     )
     if len(layers) == len(layer_model.layers):
         return layers
-    assert (
-        len(layers) <= end_layer - start_layer
-    ), f"found {len(layers)} layers in PP range [{start_layer}, {end_layer})"
+    assert len(layers) <= end_layer - start_layer, (
+        f"found {len(layers)} layers in PP range [{start_layer}, {end_layer})"
+    )
     return (
         [None] * start_layer + layers + [None] * (len(layer_model.layers) - end_layer)
     )

python/sglang/srt/model_executor/model_runner_components/load_model_utils.py

@@ -93,8 +93,10 @@ def maybe_trigger_remote_instance_nccl_send_group(
     ``--speculative-draft-draft-load-format`` needs its own send group, and the
     target's format cannot answer for it."""
     if (
-        load_format or get_model().load_format
-    ) == LoadFormat.REMOTE_INSTANCE and get_model().remote_instance_weight_loader_backend == RemoteInstanceWeightLoaderBackend.NCCL:
+        (load_format or get_model().load_format) == LoadFormat.REMOTE_INSTANCE
+        and get_model().remote_instance_weight_loader_backend
+        == RemoteInstanceWeightLoaderBackend.NCCL
+    ):
         if tp_rank == 0:
             instance_ip = NetworkAddress.resolve_host(socket.gethostname())
             t = threading.Thread(

python/sglang/srt/model_executor/model_runner_components/ngram_embedding_manager.py

@@ -50,9 +50,9 @@ class NgramEmbeddingManager:
                 device=device,
             )
             chunked_prefill_size = get_schedule().chunked_prefill_size
-            assert (
-                chunked_prefill_size is not None and chunked_prefill_size > 0
-            ), "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)"
+            assert chunked_prefill_size is not None and chunked_prefill_size > 0, (
+                "Ngram embedding requires chunked prefill to be enabled (chunked_prefill_size > 0)"
+            )
             for module in model.modules():
                 if isinstance(module, NgramEmbedding):
                     module.init_buffers(

python/sglang/srt/model_executor/model_runner_components/startup_weight_load.py

@@ -520,8 +520,7 @@ class StartupWeightLoadManager:
         if changed_names:
             preview = ", ".join(changed_names[:8])
             raise RuntimeError(
-                "Startup weight commit changed graph-visible tensor storage: "
-                f"{preview}"
+                f"Startup weight commit changed graph-visible tensor storage: {preview}"
             )
         unchanged_names = manifest.unchanged_parameter_names(
             CAPTURE_SAFE_WEIGHT_SENTINEL

python/sglang/srt/model_executor/model_runner_components/weight_exporter.py

@@ -32,15 +32,15 @@ class WeightExporter:
         group_name,
         backend="nccl",
     ):
-        assert (
-            torch.distributed.is_initialized()
-        ), "Default torch process group must be initialized"
+        assert torch.distributed.is_initialized(), (
+            "Default torch process group must be initialized"
+        )
         assert group_name != "", "Group name cannot be empty"
 
         ports_list = ports.split(",")
-        assert (
-            len(ports_list) == self.tp_size
-        ), f"Expected {self.tp_size} ports, but got {len(ports_list)} ports."
+        assert len(ports_list) == self.tp_size, (
+            f"Expected {self.tp_size} ports, but got {len(ports_list)} ports."
+        )
         group_port = ports_list[self.tp_rank]
         group_name = f"{group_name}_{group_port}_{self.tp_rank}"
 
@@ -78,15 +78,15 @@ class WeightExporter:
         ports,
         group_name,
     ):
-        assert (
-            torch.distributed.is_initialized()
-        ), "Default torch process group must be initialized"
+        assert torch.distributed.is_initialized(), (
+            "Default torch process group must be initialized"
+        )
         assert group_name != "", "Group name cannot be empty"
 
         ports_list = ports.split(",")
-        assert (
-            len(ports_list) == self.tp_size
-        ), f"Expected {self.tp_size} ports, but got {len(ports_list)} ports."
+        assert len(ports_list) == self.tp_size, (
+            f"Expected {self.tp_size} ports, but got {len(ports_list)} ports."
+        )
         group_port = ports_list[self.tp_rank]
         group_name = f"{group_name}_{group_port}_{self.tp_rank}"
 

python/sglang/srt/model_executor/model_runner_components/weight_updater.py

@@ -83,9 +83,9 @@ class WeightUpdater:
         weights/parameters online, and broadcasts them to the inference
         engine through the `_model_update_group` process group.
         """
-        assert (
-            torch.distributed.is_initialized()
-        ), "Default torch process group must be initialized"
+        assert torch.distributed.is_initialized(), (
+            "Default torch process group must be initialized"
+        )
         assert group_name != "", "Group name cannot be empty"
 
         rank = rank_offset + self.tp_rank

python/sglang/srt/model_executor/pool_configurator.py

@@ -462,9 +462,9 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
 
         self._full_layers_num = len(model_config.full_attention_layer_ids)
         self._swa_layers_num = len(model_config.swa_attention_layer_ids)
-        assert (
-            self._swa_layers_num > 0
-        ), "Hybrid SWA model must have at least one SWA layer"
+        assert self._swa_layers_num > 0, (
+            "Hybrid SWA model must have at least one SWA layer"
+        )
 
         self._swa_full_tokens_ratio = get_schedule().swa_full_tokens_ratio
         self._sliding_window_size = kvc.sliding_window_size
@@ -1014,9 +1014,9 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
     def calculate_pool_sizes(
         self, available_bytes: int, page_size: int
     ) -> MemoryPoolConfig:
-        assert (
-            page_size % 128 == 0
-        ), "page_size must be multiple of 128 for compressed attention"
+        assert page_size % 128 == 0, (
+            "page_size must be multiple of 128 for compressed attention"
+        )
 
         if self.requested_max_running_requests_per_worker is not None:
             c128_state_fixed_bytes = self._get_c128_state_fixed_bytes(
@@ -1044,9 +1044,9 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
     def calculate_pool_sizes_from_max_tokens(
         self, max_total_num_tokens: int, page_size: int
     ) -> MemoryPoolConfig:
-        assert (
-            page_size % 128 == 0
-        ), "page_size must be multiple of 128 for compressed attention"
+        assert page_size % 128 == 0, (
+            "page_size must be multiple of 128 for compressed attention"
+        )
         sizes = self._compute_dsv4_sizes(max_total_num_tokens, page_size)
         return self._to_config(sizes)
 

python/sglang/srt/model_executor/runner/base_runner.py

@@ -254,9 +254,9 @@ class BaseRunner(ABC):
 
         if should_run_flashinfer_autotune(self.model_runner):
             buffers, batch_size = self._autotune_buffers()
-            assert (
-                buffers is not None
-            ), "_autotune_buffers() must return a reusable buffer set for autotune"
+            assert buffers is not None, (
+                "_autotune_buffers() must return a reusable buffer set for autotune"
+            )
             self._flashinfer_autotune(buffers=buffers, batch_size=batch_size)
             maybe_flashinfer_autotune_extend(self, decode_num_tokens=batch_size)
 
@@ -427,9 +427,9 @@ class BaseRunner(ABC):
         )
         if mr.spec_algorithm.is_speculative() and not _is_pd_prefill_target:
             if mr.is_draft_worker:
-                assert (
-                    mr.spec_algorithm.supports_target_verify_for_draft()
-                ), "This should not happen"
+                assert mr.spec_algorithm.supports_target_verify_for_draft(), (
+                    "This should not happen"
+                )
             capture_forward_mode = ForwardMode.TARGET_VERIFY
             num_tokens_per_req = mr.decode_num_tokens_per_req()
         if extend_num_tokens_per_req is not None:

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -326,9 +326,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         if model_runner.spec_algorithm.is_speculative():
             if self.model_runner.is_draft_worker:
                 # Draft workers can use TARGET_VERIFY mode.
-                if (
-                    not self.model_runner.spec_algorithm.supports_target_verify_for_draft()
-                ):
+                if not self.model_runner.spec_algorithm.supports_target_verify_for_draft():
                     raise RuntimeError("This should not happen")
             self.capture_forward_mode = ForwardMode.TARGET_VERIFY
         elif self.is_dllm:
@@ -493,7 +491,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 self.capture()
         except RuntimeError as e:
             raise Exception(
-                f"Capture cuda graph failed: {e}\n" f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
+                f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
             )
 
     def _record_in_graph_metadata_prep_done(self):
@@ -1167,9 +1165,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
 
         # Sanity-check: --debug-cuda-graph requires breakable backend.
         if get_exec().graph.debug_cuda_graph:
-            assert isinstance(
-                self.backend, BreakableCudaGraphBackend
-            ), "Breakable CUDA graph is required for --debug-cuda-graph"
+            assert isinstance(self.backend, BreakableCudaGraphBackend), (
+                "Breakable CUDA graph is required for --debug-cuda-graph"
+            )
 
         forward_batch, attn_backend, pp_proxy_tensors = self.capture_prepare(
             bs, stream_idx=stream_idx, num_tokens=num_tokens
@@ -1529,7 +1527,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             if self.model_runner.is_draft_worker:
                 raise RuntimeError("This should not happen.")
             else:
-
                 capture_mode = (
                     CaptureHiddenMode.NULL
                     if self.model_runner.spec_algorithm.is_standalone()

python/sglang/srt/model_executor/runner/flashinfer_autotune.py

@@ -277,11 +277,15 @@ def flashinfer_autotune_context(model_runner: ModelRunner, *, run_lm_head: bool)
         from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
 
         skip_ops = get_flashinfer_autotune_skip_ops(mr)
-        with _autotune_process_group(sync_group), autotune(
-            True,
-            cache=str(autotune_cache),
-            skip_ops=skip_ops,
-        ), autotune_dummy_run_mode(run_lm_head=run_lm_head):
+        with (
+            _autotune_process_group(sync_group),
+            autotune(
+                True,
+                cache=str(autotune_cache),
+                skip_ops=skip_ops,
+            ),
+            autotune_dummy_run_mode(run_lm_head=run_lm_head),
+        ):
             yield
     torch.cuda.current_stream().wait_stream(mr.forward_stream)
     logger.info("FlashInfer autotune completed.")

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -543,9 +543,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
         # the contract only when the backend is Breakable; FullCG and
         # TC_PIECEWISE use the eager init_forward_metadata path.
         if isinstance(self.backend, BreakableCudaGraphBackend):
-            self.use_captured_attn_metadata = (
-                model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
-            )
+            self.use_captured_attn_metadata = model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
         else:
             self.use_captured_attn_metadata = False
         self.attn_metadata_buffers: Optional[Dict[int, object]] = (
@@ -625,9 +623,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
     def _prefill_logits_buffer_rows(self, forward_batch: ForwardBatch) -> int:
         if not forward_batch.return_logprob:
             return forward_batch.batch_size
-        assert (
-            self._uses_eager_prefill_tail()
-        ), "Prefill return_logprob requires an eager logits tail."
+        assert self._uses_eager_prefill_tail(), (
+            "Prefill return_logprob requires an eager logits tail."
+        )
 
         global_num_tokens = forward_batch.global_num_tokens_for_logprob_cpu
         if global_num_tokens is not None:

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -131,11 +131,14 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         size = shape_key.size
         if self._shared_output_buffer is None:
             self._shared_output_buffer = self._alloc_full_buffer(warmup_out, size)
-        with graph_pool_capture_scope(), BreakableCUDAGraphCapture(
-            cuda_graph=graph,
-            pool=self._pool,
-            stream=self._capture_stream,
-            barrier_fn=self._tp_group.barrier,
+        with (
+            graph_pool_capture_scope(),
+            BreakableCUDAGraphCapture(
+                cuda_graph=graph,
+                pool=self._pool,
+                stream=self._capture_stream,
+                barrier_fn=self._tp_group.barrier,
+            ),
         ):
             self._precarve.mint()
             out = captured_fn()

python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py

@@ -259,9 +259,9 @@ class DedupedCudaGraphRegistry:
     def replay(self, graph: DedupedCudaGraph, stream: int) -> None:
         assert cuda_rt is not None
         group = graph.group
-        assert (
-            group is not None
-        ), "captured CUDA graph does not belong to this dedup state"
+        assert group is not None, (
+            "captured CUDA graph does not belong to this dedup state"
+        )
 
         raw_graph = graph.raw_graph
         graph_exec = group.graph_exec

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -317,9 +317,9 @@ class BreakableCUDAGraphCapture:
         capture_error_mode: str = "global",
         barrier_fn: Callable[[], None] | None = None,
     ):
-        assert isinstance(
-            cuda_graph, BreakableCUDAGraph
-        ), "cuda_graph must be a BreakableCUDAGraph"
+        assert isinstance(cuda_graph, BreakableCUDAGraph), (
+            "cuda_graph must be a BreakableCUDAGraph"
+        )
         self.cuda_graph = cuda_graph
         self._pool = pool if pool is not None else (0, 0)
         self._stream = stream

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py

@@ -33,8 +33,7 @@ def _cudaGetErrorString(error):
 def checkCudaErrors(result):
     if rt is None:
         raise RuntimeError(
-            "cuda.bindings is not available. "
-            "Install it with: pip install cuda-python"
+            "cuda.bindings is not available. Install it with: pip install cuda-python"
         )
     if result[0] != rt.cudaError_t.cudaSuccess:
         raise RuntimeError(