Resolve conflicts

Signed-off-by: Xing Liu <xingliu14@gmail.com>

Author: xingliu14

tests/worker/tpu_worker_test.py

@@ -63,22 +63,22 @@ def test_init_success(self, mock_vllm_config):
         assert worker.profile_dir is None
         assert worker.devices == ['tpu:0']
 
-    @patch('tpu_inference.worker.tpu_worker_jax.vllm_envs')
+    @patch('tpu_inference.worker.tpu_worker.vllm_envs')
     def test_init_with_profiler_on_rank_zero(self, mock_envs,
                                              mock_vllm_config):
         """Tests that the profiler directory is set correctly on rank 0."""
-        mock_vllm_envs.VLLM_TORCH_PROFILER_DIR = "/tmp/profiles"
+        mock_envs.VLLM_TORCH_PROFILER_DIR = "/tmp/profiles"
         worker = TPUWorker(vllm_config=mock_vllm_config,
                            local_rank=0,
                            rank=0,
                            distributed_init_method="test_method")
         assert worker.profile_dir == "/tmp/profiles"
 
-    @patch('tpu_inference.worker.tpu_worker_jax.vllm_envs')
+    @patch('tpu_inference.worker.tpu_worker.vllm_envs')
     def test_init_with_profiler_on_other_ranks(self, mock_envs,
                                                mock_vllm_config):
         """Tests that the profiler directory is NOT set on non-rank 0 workers."""
-        mock_vllm_envs.VLLM_TORCH_PROFILER_DIR = "/tmp/profiles"
+        mock_envs.VLLM_TORCH_PROFILER_DIR = "/tmp/profiles"
         worker = TPUWorker(vllm_config=mock_vllm_config,
                            local_rank=1,
                            rank=1,