[Bugfix] Add custom Triton cache manager to resolve MoE MP issue

Dockerfile

@@ -166,6 +166,13 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
 RUN --mount=type=bind,from=mamba-builder,src=/usr/src/mamba,target=/usr/src/mamba \
     --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install /usr/src/mamba/*.whl --no-cache-dir
+
+# workaround for https://github.com/vllm-project/vllm/issues/6103
+# until fixed in Triton upstream
+RUN --mount=type=bind,source=triton_patch,target=/context \
+    export TRITON_PATH=$(python3 -c "import triton; print(triton.__file__.strip(\"__init__.py\"))") \
+    && cp /context/custom_cache_manager.py ${TRITON_PATH}/runtime/custom_cache_manager.py
+
 #################### vLLM installation IMAGE ####################
 
 
@@ -197,7 +204,8 @@ FROM vllm-base AS vllm-openai
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install accelerate hf_transfer 'modelscope!=1.15.0'
 
-ENV VLLM_USAGE_SOURCE production-docker-image
+ENV VLLM_USAGE_SOURCE production-docker-image \
+    TRITON_CACHE_MANAGER "triton.runtime.custom_cache_manager:CustomCacheManager"
 
 ENTRYPOINT ["python3", "-m", "vllm.entrypoints.openai.api_server"]
 #################### OPENAI API SERVER ####################

triton_patch/custom_cache_manager.py

@@ -0,0 +1,32 @@
+import os
+
+from triton.runtime.cache import (FileCacheManager, default_cache_dir,
+                                  default_dump_dir, default_override_dir)
+
+
+class CustomCacheManager(FileCacheManager):
+
+    def __init__(self, key, override=False, dump=False):
+        self.key = key
+        self.lock_path = None
+        if dump:
+            self.cache_dir = default_dump_dir()
+            self.cache_dir = os.path.join(self.cache_dir, self.key)
+            self.lock_path = os.path.join(self.cache_dir, "lock")
+            os.makedirs(self.cache_dir, exist_ok=True)
+        elif override:
+            self.cache_dir = default_override_dir()
+            self.cache_dir = os.path.join(self.cache_dir, self.key)
+        else:
+            # create cache directory if it doesn't exist
+            self.cache_dir = os.getenv("TRITON_CACHE_DIR",
+                                       "").strip() or default_cache_dir()
+            if self.cache_dir:
+                self.cache_dir = f"{self.cache_dir}_{os.getpid()}"
+                self.cache_dir = os.path.join(self.cache_dir, self.key)
+                self.lock_path = os.path.join(self.cache_dir, "lock")
+                os.makedirs(self.cache_dir, exist_ok=True)
+            else:
+                raise RuntimeError("Could not create or locate cache dir")
+
+        print(f"Triton cache dir: {self.cache_dir=}")