Avoid downloading engine more than once (#137)

In cases when we run TRT-LLM trusses with more than one worker we
download engine files each time worker runs. I've added safeguard on
this so we download files only once.

llama/llama-2-7b-trt-llm/model/model.py

@@ -6,7 +6,7 @@
 import numpy as np
 from client import TritonClient, UserData
 from transformers import AutoTokenizer
-from utils import download_engine, prepare_grpc_tensor
+from utils import download_engine, prepare_grpc_tensor, server_loaded
 
 TRITON_MODEL_REPOSITORY_PATH = Path("/packages/inflight_batcher_llm/")
 
@@ -45,11 +45,14 @@ def load(self):
 
         # Download model from Hugging Face Hub if specified
         if is_external_engine_repo:
-            download_engine(
-                engine_repository=self._config["model_metadata"]["engine_repository"],
-                fp=self._data_dir,
-                auth_token=hf_access_token,
-            )
+            if not server_loaded():
+                download_engine(
+                    engine_repository=self._config["model_metadata"][
+                        "engine_repository"
+                    ],
+                    fp=self._data_dir,
+                    auth_token=hf_access_token,
+                )
 
         # Load Triton Server and model
         tokenizer_repository = self._config["model_metadata"]["tokenizer_repository"]

mistral/mistral-7b-instruct-chat-trt-llm-smooth-quant/model/model.py

@@ -6,7 +6,7 @@
 import numpy as np
 from client import TritonClient, UserData
 from transformers import AutoTokenizer
-from utils import download_engine, prepare_grpc_tensor
+from utils import download_engine, prepare_grpc_tensor, server_loaded
 
 TRITON_MODEL_REPOSITORY_PATH = Path("/packages/inflight_batcher_llm/")
 
@@ -45,11 +45,14 @@ def load(self):
 
         # Download model from Hugging Face Hub if specified
         if is_external_engine_repo:
-            download_engine(
-                engine_repository=self._config["model_metadata"]["engine_repository"],
-                fp=self._data_dir,
-                auth_token=hf_access_token,
-            )
+            if not server_loaded():
+                download_engine(
+                    engine_repository=self._config["model_metadata"][
+                        "engine_repository"
+                    ],
+                    fp=self._data_dir,
+                    auth_token=hf_access_token,
+                )
 
         # Load Triton Server and model
         tokenizer_repository = self._config["model_metadata"]["tokenizer_repository"]

mistral/mistral-7b-instruct-chat-trt-llm-weights-only-quant/model/model.py

@@ -6,7 +6,7 @@
 import numpy as np
 from client import TritonClient, UserData
 from transformers import AutoTokenizer
-from utils import download_engine, prepare_grpc_tensor
+from utils import download_engine, prepare_grpc_tensor, server_loaded
 
 TRITON_MODEL_REPOSITORY_PATH = Path("/packages/inflight_batcher_llm/")
 
@@ -45,11 +45,14 @@ def load(self):
 
         # Download model from Hugging Face Hub if specified
         if is_external_engine_repo:
-            download_engine(
-                engine_repository=self._config["model_metadata"]["engine_repository"],
-                fp=self._data_dir,
-                auth_token=hf_access_token,
-            )
+            if not server_loaded():
+                download_engine(
+                    engine_repository=self._config["model_metadata"][
+                        "engine_repository"
+                    ],
+                    fp=self._data_dir,
+                    auth_token=hf_access_token,
+                )
 
         # Load Triton Server and model
         tokenizer_repository = self._config["model_metadata"]["tokenizer_repository"]

mistral/mistral-7b-instruct-chat-trt-llm/model/model.py

@@ -6,7 +6,7 @@
 import numpy as np
 from client import TritonClient, UserData
 from transformers import AutoTokenizer
-from utils import download_engine, prepare_grpc_tensor
+from utils import download_engine, prepare_grpc_tensor, server_loaded
 
 TRITON_MODEL_REPOSITORY_PATH = Path("/packages/inflight_batcher_llm/")
 
@@ -45,11 +45,14 @@ def load(self):
 
         # Download model from Hugging Face Hub if specified
         if is_external_engine_repo:
-            download_engine(
-                engine_repository=self._config["model_metadata"]["engine_repository"],
-                fp=self._data_dir,
-                auth_token=hf_access_token,
-            )
+            if not server_loaded():
+                download_engine(
+                    engine_repository=self._config["model_metadata"][
+                        "engine_repository"
+                    ],
+                    fp=self._data_dir,
+                    auth_token=hf_access_token,
+                )
 
         # Load Triton Server and model
         tokenizer_repository = self._config["model_metadata"]["tokenizer_repository"]

mistral/mistral-7b-trt-llm-build-engine/model/model.py

@@ -7,7 +7,7 @@
 from build_engine_utils import BuildConfig, build_engine
 from client import TritonClient, UserData
 from transformers import AutoTokenizer
-from utils import download_engine, prepare_grpc_tensor
+from utils import download_engine, prepare_grpc_tensor, server_loaded
 
 TRITON_MODEL_REPOSITORY_PATH = Path("/packages/inflight_batcher_llm/")
 
@@ -46,11 +46,14 @@ def load(self):
 
         # Download model from Hugging Face Hub if specified
         if is_external_engine_repo:
-            download_engine(
-                engine_repository=self._config["model_metadata"]["engine_repository"],
-                fp=self._data_dir,
-                auth_token=hf_access_token,
-            )
+            if not server_loaded():
+                download_engine(
+                    engine_repository=self._config["model_metadata"][
+                        "engine_repository"
+                    ],
+                    fp=self._data_dir,
+                    auth_token=hf_access_token,
+                )
         tokenizer_repository = self._config["model_metadata"]["tokenizer_repository"]
         if "engine_build" in self._config["model_metadata"]:
             if not is_external_engine_repo:

mistral/mixtral-8x7b-instruct-trt-llm-weights-only-quant/model/model.py

@@ -6,7 +6,7 @@
 import numpy as np
 from client import TritonClient, UserData
 from transformers import AutoTokenizer
-from utils import download_engine, prepare_grpc_tensor
+from utils import download_engine, prepare_grpc_tensor, server_loaded
 
 TRITON_MODEL_REPOSITORY_PATH = Path("/packages/inflight_batcher_llm/")
 
@@ -45,11 +45,14 @@ def load(self):
 
         # Download model from Hugging Face Hub if specified
         if is_external_engine_repo:
-            download_engine(
-                engine_repository=self._config["model_metadata"]["engine_repository"],
-                fp=self._data_dir,
-                auth_token=hf_access_token,
-            )
+            if not server_loaded():
+                download_engine(
+                    engine_repository=self._config["model_metadata"][
+                        "engine_repository"
+                    ],
+                    fp=self._data_dir,
+                    auth_token=hf_access_token,
+                )
 
         # Load Triton Server and model
         tokenizer_repository = self._config["model_metadata"]["tokenizer_repository"]

mistral/mixtral-8x7b-instruct-trt-llm/model/model.py

@@ -6,7 +6,7 @@
 import numpy as np
 from client import TritonClient, UserData
 from transformers import AutoTokenizer
-from utils import download_engine, prepare_grpc_tensor
+from utils import download_engine, prepare_grpc_tensor, server_loaded
 
 TRITON_MODEL_REPOSITORY_PATH = Path("/packages/inflight_batcher_llm/")
 
@@ -45,11 +45,14 @@ def load(self):
 
         # Download model from Hugging Face Hub if specified
         if is_external_engine_repo:
-            download_engine(
-                engine_repository=self._config["model_metadata"]["engine_repository"],
-                fp=self._data_dir,
-                auth_token=hf_access_token,
-            )
+            if not server_loaded():
+                download_engine(
+                    engine_repository=self._config["model_metadata"][
+                        "engine_repository"
+                    ],
+                    fp=self._data_dir,
+                    auth_token=hf_access_token,
+                )
 
         # Load Triton Server and model
         tokenizer_repository = self._config["model_metadata"]["tokenizer_repository"]

templates/trt-llm/model/model.py

@@ -6,7 +6,7 @@
 import numpy as np
 from client import TritonClient, UserData
 from transformers import AutoTokenizer
-from utils import download_engine, prepare_grpc_tensor
+from utils import download_engine, prepare_grpc_tensor, server_loaded
 
 TRITON_MODEL_REPOSITORY_PATH = Path("/packages/inflight_batcher_llm/")
 
@@ -45,11 +45,14 @@ def load(self):
 
         # Download model from Hugging Face Hub if specified
         if is_external_engine_repo:
-            download_engine(
-                engine_repository=self._config["model_metadata"]["engine_repository"],
-                fp=self._data_dir,
-                auth_token=hf_access_token,
-            )
+            if not server_loaded():
+                download_engine(
+                    engine_repository=self._config["model_metadata"][
+                        "engine_repository"
+                    ],
+                    fp=self._data_dir,
+                    auth_token=hf_access_token,
+                )
 
         # Load Triton Server and model
         tokenizer_repository = self._config["model_metadata"]["tokenizer_repository"]