Fix isort, yapf, and ruff formatting issues

Author: xingliu14

examples/offline_inference.py

@@ -1,11 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import tpu_inference.envs as envs
 import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+import tpu_inference.envs as envs
 from tpu_inference.core import disagg_utils
 
 

examples/offline_lora_inference.py

@@ -3,12 +3,13 @@
 
 import time
 
-import tpu_inference.envs as envs
 import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.lora.request import LoRARequest
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+import tpu_inference.envs as envs
+
 
 def create_parser():
     parser = FlexibleArgumentParser()

examples/offline_safety_model_inference.py

@@ -18,11 +18,11 @@
     --max-num_batched_tokens=4096
 """
 
-import tpu_inference.envs as envs
 import vllm.envs as vllm_envs
 from vllm import LLM, EngineArgs
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+import tpu_inference.envs as envs
 from tpu_inference.core import disagg_utils
 
 

tests/runner/test_tpu_runner_mesh.py

@@ -1,5 +1,4 @@
 """Unit tests for TPUModelRunner mesh initialization."""
-import os
 from unittest.mock import Mock, patch
 
 import pytest

tpu_inference/executors/ray_distributed_executor.py

@@ -74,7 +74,8 @@ def _init_executor(self) -> None:
         self.forward_dag: Optional[ray.dag.CompiledDAG] = None
 
         # Ensure Ray compiled DAG channel type is set for vLLM
-        os.environ["VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = tpu_envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
+        os.environ[
+            "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE"] = tpu_envs.VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE
 
         # Currently, this requires USE_RAY_SPMD_WORKER=True.
         self.use_ray_compiled_dag = True
@@ -91,7 +92,8 @@ def _init_executor(self) -> None:
         # Ensure Ray usage stats collection setting is propagated to Ray workers.
         # Ray workers inherit environment variables, so we explicitly set this
         # based on our configuration (defaults to "0" to disable stats).
-        os.environ["RAY_USAGE_STATS_ENABLED"] = tpu_envs.RAY_USAGE_STATS_ENABLED
+        os.environ[
+            "RAY_USAGE_STATS_ENABLED"] = tpu_envs.RAY_USAGE_STATS_ENABLED
 
         # Create the parallel GPU workers.
         self._init_workers_ray(placement_group)

tpu_inference/layers/common/sharding.py

@@ -1,6 +1,5 @@
 import json
 import math
-import os
 from dataclasses import asdict, dataclass
 from typing import TYPE_CHECKING, List, Optional
 

tpu_inference/runner/tpu_runner.py

@@ -1,6 +1,5 @@
 import copy
 import functools
-import os
 import random
 from contextlib import nullcontext
 from dataclasses import dataclass
@@ -13,8 +12,6 @@
 import torch
 import vllm.envs as vllm_envs
 from flax import nnx
-
-import tpu_inference.envs as envs
 from jax.experimental import mesh_utils
 from jax.sharding import NamedSharding, PartitionSpec
 from torchax.ops.mappings import j2t_dtype
@@ -37,6 +34,7 @@
     KVConnectorModelRunnerMixin
 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
 
+import tpu_inference.envs as envs
 from tpu_inference import utils as common_utils
 from tpu_inference.layers.common.attention_metadata import AttentionMetadata
 from tpu_inference.layers.common.sharding import (MESH_AXIS_NAMES,