[Bugfix][TPU] Set readonly=True for non-root devices (#6980)

vllm-project · Jul 31, 2024 · 533d193 · 533d193
1 parent 9f0e69b
commit 533d193
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/vllm/worker/tpu_worker.py b/vllm/worker/tpu_worker.py
@@ -104,7 +104,10 @@ def init_device(self) -> None:
         # Use persistent cache to avoid XLA recompilation.
         # NOTE(woosuk): This does not completely eliminate the recompilation
         # overhead because dynamo does not cache the compiled results.
-        xr.initialize_cache(envs.VLLM_XLA_CACHE_PATH, readonly=False)
+        # NOTE(woosuk): Set readonly=False only for the rank 0 process to avoid
+        # race conditions.
+        xr.initialize_cache(envs.VLLM_XLA_CACHE_PATH,
+                            readonly=not self.is_driver_worker)
 
     def load_model(self):
         self.model_runner.load_model()