From aef43e2b4943c1e0489c58a0d327f1b8849d1f84 Mon Sep 17 00:00:00 2001 From: Thomas Parnell Date: Wed, 29 May 2024 11:01:40 +0000 Subject: [PATCH 1/5] Revert previous attempt at Triton patch; use CustomCacheManger instead. Co-authored-by: Chih-Chieh-Yang Signed-off-by: Thomas Parnell --- Dockerfile.ubi | 10 +++------- triton_patch/custom_cache_manager.py | 27 +++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 7 deletions(-) create mode 100644 triton_patch/custom_cache_manager.py diff --git a/Dockerfile.ubi b/Dockerfile.ubi index 0ab58963e5..36d04546b8 100644 --- a/Dockerfile.ubi +++ b/Dockerfile.ubi @@ -265,19 +265,15 @@ RUN microdnf install -y gcc \ && microdnf clean all # patch triton (fix for #720) -COPY triton_patch/cache_fix.patch . -RUN microdnf install -y patch \ - && patch /opt/vllm/lib/python3.11/site-packages/triton/runtime/cache.py cache_fix.patch \ - && microdnf remove -y patch \ - && microdnf clean all \ - && rm cache_fix.patch +COPY triton_patch/custom_cache_manager.py /opt/vllm/lib/python3.11/site-packages/triton/runtime/custom_cache_manager.py ENV HF_HUB_OFFLINE=1 \ PORT=8000 \ GRPC_PORT=8033 \ HOME=/home/vllm \ VLLM_USAGE_SOURCE=production-docker-image \ - VLLM_WORKER_MULTIPROC_METHOD=fork + VLLM_WORKER_MULTIPROC_METHOD=fork \ + TRITON_CACHE_MANAGER="triton.runtime.custom_cache_manager:CustomCacheManager" # setup non-root user for OpenShift RUN microdnf install -y shadow-utils \ diff --git a/triton_patch/custom_cache_manager.py b/triton_patch/custom_cache_manager.py new file mode 100644 index 0000000000..9c9b8d5a1b --- /dev/null +++ b/triton_patch/custom_cache_manager.py @@ -0,0 +1,27 @@ +import os +from triton.runtime.cache import FileCacheManager, default_dump_dir, default_override_dir, default_cache_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() + self.cache_dir = self.cache_dir + "_" + str(os.getpid()) + print(f"{self.cache_dir=}") + if self.cache_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) + else: + raise RuntimeError("Could not create or locate cache dir") From d5910bb47868e77070047fe454e405f356e41fc3 Mon Sep 17 00:00:00 2001 From: Thomas Parnell Date: Wed, 29 May 2024 11:06:39 +0000 Subject: [PATCH 2/5] Remove existing patch file. Signed-off-by: Thomas Parnell --- triton_patch/cache_fix.patch | 8 -------- 1 file changed, 8 deletions(-) delete mode 100644 triton_patch/cache_fix.patch diff --git a/triton_patch/cache_fix.patch b/triton_patch/cache_fix.patch deleted file mode 100644 index 97a1aa4778..0000000000 --- a/triton_patch/cache_fix.patch +++ /dev/null @@ -1,8 +0,0 @@ -4c4 -< import random ---- -> import uuid -117c117 -< rnd_id = random.randint(0, 1000000) ---- -> rnd_id = str(uuid.uuid4()) From ed54561a0bd81ba2e143044a0d11ebd169969d83 Mon Sep 17 00:00:00 2001 From: Thomas Parnell Date: Wed, 29 May 2024 11:33:27 +0000 Subject: [PATCH 3/5] Fix for ruff Signed-off-by: Thomas Parnell --- triton_patch/custom_cache_manager.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/triton_patch/custom_cache_manager.py b/triton_patch/custom_cache_manager.py index 9c9b8d5a1b..e065bafb7d 100644 --- a/triton_patch/custom_cache_manager.py +++ b/triton_patch/custom_cache_manager.py @@ -1,5 +1,11 @@ import os -from triton.runtime.cache import FileCacheManager, default_dump_dir, default_override_dir, default_cache_dir +from triton.runtime.cache import ( + FileCacheManager, + default_dump_dir, + default_override_dir, + default_cache_dir +) + class CustomCacheManager(FileCacheManager): @@ -16,7 +22,8 @@ def __init__(self, key, override=False, dump=False): 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() + self.cache_dir = os.getenv("TRITON_CACHE_DIR", + "").strip() or default_cache_dir() self.cache_dir = self.cache_dir + "_" + str(os.getpid()) print(f"{self.cache_dir=}") if self.cache_dir: From cf4eae94c1921d0d1eb17da8321c2820f43bfe78 Mon Sep 17 00:00:00 2001 From: Thomas Parnell Date: Wed, 29 May 2024 11:36:59 +0000 Subject: [PATCH 4/5] Fix for isort Signed-off-by: Thomas Parnell --- triton_patch/custom_cache_manager.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/triton_patch/custom_cache_manager.py b/triton_patch/custom_cache_manager.py index e065bafb7d..a923d9e633 100644 --- a/triton_patch/custom_cache_manager.py +++ b/triton_patch/custom_cache_manager.py @@ -1,10 +1,7 @@ import os -from triton.runtime.cache import ( - FileCacheManager, - default_dump_dir, - default_override_dir, - default_cache_dir -) + +from triton.runtime.cache import (FileCacheManager, default_cache_dir, + default_dump_dir, default_override_dir) class CustomCacheManager(FileCacheManager): From 4c5c1e75f1c99a41410e135f702e520345735e16 Mon Sep 17 00:00:00 2001 From: Joe Runde Date: Mon, 3 Jun 2024 10:20:43 -0600 Subject: [PATCH 5/5] Small updates from review Co-authored-by: Nick Hill Signed-off-by: Nick Hill Signed-off-by: Joe Runde --- triton_patch/custom_cache_manager.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/triton_patch/custom_cache_manager.py b/triton_patch/custom_cache_manager.py index a923d9e633..5c27c072e9 100644 --- a/triton_patch/custom_cache_manager.py +++ b/triton_patch/custom_cache_manager.py @@ -21,11 +21,12 @@ def __init__(self, key, override=False, dump=False): # create cache directory if it doesn't exist self.cache_dir = os.getenv("TRITON_CACHE_DIR", "").strip() or default_cache_dir() - self.cache_dir = self.cache_dir + "_" + str(os.getpid()) - print(f"{self.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=}") \ No newline at end of file