error loudly for development pushes for TRT-LLM built trusses (#1050)

* add live_reload config override for TRT-LLM built trusses

* add unary default for max_beam_width

* fail loudly instead

* add field validator

* fix tests

* add for missing case

* remove redundant helper

* bump

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truss"
-version = "0.9.25rc9"
+version = "0.9.25rc10"
 description = "A seamless bridge from model development to model delivery"
 license = "MIT"
 readme = "README.md"

truss/cli/cli.py

@@ -35,6 +35,7 @@
 from truss.util.config_checks import (
     check_and_update_memory_for_trt_llm_builder,
     check_secrets_for_trt_llm_builder,
+    uses_trt_llm_builder,
 )
 from truss.util.errors import RemoteNetworkError
 
@@ -860,6 +861,10 @@ def push(
         console.print(
             f"Automatically increasing memory for trt-llm builder to {TRTLLM_MIN_MEMORY_REQUEST_GI}Gi."
         )
+    if uses_trt_llm_builder(tr) and not publish:
+        live_reload_disabled_text = "Development mode is currently not supported for trusses using TRT-LLM build flow, push as a published model using --publish"
+        console.print(live_reload_disabled_text, style="red")
+        sys.exit(1)
 
     # TODO(Abu): This needs to be refactored to be more generic
     service = remote_provider.push(

truss/config/trt_llm.py

@@ -3,7 +3,7 @@
 from enum import Enum
 from typing import Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 from rich.console import Console
 
 logging.basicConfig(level=logging.INFO)
@@ -53,7 +53,7 @@ class TrussTRTLLMBuildConfiguration(BaseModel):
     max_input_len: int
     max_output_len: int
     max_batch_size: int
-    max_beam_width: int
+    max_beam_width: Optional[int] = 1
     max_prompt_embedding_table_size: int = 0
     checkpoint_repository: CheckpointRepository
     gather_all_token_logits: bool = False
@@ -70,6 +70,12 @@ class TrussTRTLLMBuildConfiguration(BaseModel):
     kv_cache_free_gpu_mem_fraction: float = 0.9
     num_builder_gpus: Optional[int] = None
 
+    @field_validator("max_beam_width", mode="after")
+    @classmethod
+    def ensure_unary_max_beam_width(cls, value):
+        if value and value != 1:
+            raise ValueError("Non-unary max_beam_width not supported")
+
 
 class TrussTRTLLMServingConfiguration(BaseModel):
     engine_repository: str

truss/contexts/image_builder/serving_image_builder.py

@@ -339,9 +339,6 @@ def copy_into_build_dir(from_path: Path, path_in_build_dir: str):
         # Copy over truss
         copy_tree_path(truss_dir, build_dir, ignore_patterns=truss_ignore_patterns)
 
-        # Copy over template truss for TRT-LLM (we overwrite the model and packages dir)
-        # Most of the code is pulled from upstream triton-inference-server tensorrtllm_backend
-        # https://github.com/triton-inference-server/tensorrtllm_backend/tree/v0.9.0/all_models/inflight_batcher_llm
         if config.trt_llm is not None:
             is_audio_model = (
                 config.trt_llm.build.base_model == TrussTRTLLMModel.WHISPER

truss/tests/test_config.py

@@ -55,7 +55,7 @@ def trtllm_config(default_config) -> Dict[str, Any]:
             "max_input_len": 1024,
             "max_output_len": 1024,
             "max_batch_size": 512,
-            "max_beam_width": 1,
+            "max_beam_width": None,
             "checkpoint_repository": {
                 "source": "HF",
                 "repo": "meta/llama4-500B",

truss/util/config_checks.py

@@ -8,6 +8,12 @@
 from truss.truss_handle import TrussHandle
 
 
+def uses_trt_llm_builder(tr: TrussHandle) -> bool:
+    return (
+        tr.spec.config.trt_llm is not None and tr.spec.config.trt_llm.build is not None
+    )
+
+
 def check_secrets_for_trt_llm_builder(tr: TrussHandle) -> bool:
     if tr.spec.config.trt_llm and tr.spec.config.trt_llm.build:
         source = tr.spec.config.trt_llm.build.checkpoint_repository.source