issue on truss config

Author: michaelfeil

truss/base/trt_llm_config.py

@@ -19,14 +19,6 @@
     os.environ.get("ENGINE_BUILDER_TRUSS_RUNTIME_MIGRATION", "False") == "True"
 )
 
-# Workaround for briton import without truss installed.
-try:
-    from truss.base.truss_config import Resources
-except ImportError:
-
-    class Resources(BaseModel):  # type: ignore
-        pass
-
 
 class TrussTRTLLMModel(str, Enum):
     ENCODER = "encoder"
@@ -156,19 +148,33 @@ class TrussTRTLLMBuildConfiguration(BaseModel):
         TrussTRTLLMPluginConfiguration()
     )
     num_builder_gpus: Optional[int] = None
-    build_resources: Optional[Resources] = None
+    build_resources: Optional[Any] = None
     speculator: Optional[TrussSpeculatorConfiguration] = None
 
     class Config:
         extra = "forbid"
 
     def __init__(self, **data):
+        data = self.parse_build_resources(data)
         super().__init__(**data)
         self._validate_kv_cache_flags()
         self._validate_speculator_config()
         self._bei_specfic_migration()
         self._depreacate_num_builder_gpus()
 
+    @staticmethod
+    def parse_build_resources(data):
+        build_resources = data.get("build_resources")
+        if build_resources:
+            try:
+                from truss.base.truss_config import Resources
+
+                print("build_resources", build_resources)
+                data["build_resources"] = Resources.from_dict(build_resources)
+            except ImportError:
+                pass
+        return data
+
     def _depreacate_num_builder_gpus(self):
         if self.num_builder_gpus:
             logger.warning(

truss/tests/trt_llm/test_trt_llm_config.py

@@ -74,11 +74,13 @@ def test_trt_llm_chunked_prefill_fix(trtllm_config):
 
 
 def test_trt_llm_builder_gpus(trtllm_config):
+    db_dict = {"cpu": "1", "accelerator": "A10G:2", "memory": "1G", "node_count": None}
+    trtllm_config["trt_llm"]["build"]["build_resources"] = db_dict
+
     trt_llm_config = TRTLLMConfiguration(**trtllm_config["trt_llm"])
-    trt_llm_config.build.build_resources = Resources(
-        cpu=1, accelerator="A10G:2", memory="1G"
-    )
-    assert trt_llm_config.build.build_resources.accelerator == "A10G:2"
+    build_resources = Resources.from_dict(db_dict)
+    trt_llm_config.build.build_resources = build_resources
+    assert trt_llm_config.build.build_resources == build_resources
 
 
 def test_trt_llm_lookahead_decoding(trtllm_config):