update migration logic

basetenlabs · Dec 11, 2024 · c8b6de7 · c8b6de7
1 parent 029622d
commit c8b6de7
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/truss/base/trt_llm_config.py b/truss/base/trt_llm_config.py
@@ -213,14 +213,20 @@ def migrate_runtime_fields(cls, data: Any) -> Any:
                 valid_build_fields[key] = value
             else:
                 if key in TrussTRTLLMRuntimeConfiguration.__annotations__:
-                    logger.warning(f"Setting runtime.{key}: {value}")
+                    logger.warning(f"Found runtime.{key}: {value} in build config")
                     extra_runtime_fields[key] = value
         if extra_runtime_fields:
             logger.warning(
-                f"Found extra fields {list(extra_runtime_fields.keys())} in build configuration, fields were migrated to runtime configuration."
-                " This migration of deprecated fields is scheduled for removal, please upgrade to the latest truss version and update configs according to https://docs.baseten.co/performance/engine-builder-config."
+                f"Found extra fields {list(extra_runtime_fields.keys())} in build configuration, unspecified runtime fields will be configured using these values."
+                " This configuration of deprecated fields is scheduled for removal, please upgrade to the latest truss version and update configs according to https://docs.baseten.co/performance/engine-builder-config."
+            )
+            data.get("runtime").update(
+                {
+                    k: v
+                    for k, v in extra_runtime_fields.items()
+                    if k not in data.get("runtime")
+                }
             )
-            data.update({"runtime": extra_runtime_fields})
 
         data.update({"build": valid_build_fields})
         return data

diff --git a/truss/tests/conftest.py b/truss/tests/conftest.py
@@ -808,15 +808,15 @@ def deprecated_trtllm_config(default_config) -> Dict[str, Any]:
             "enable_chunked_context": True,
             "batch_scheduler_policy": TrussTRTLLMBatchSchedulerPolicy.MAX_UTILIZATION.value,
             "request_default_max_tokens": 10,
-            "total_token_limit": 100,
+            "total_token_limit": 50,
             # end deprecated fields
             "checkpoint_repository": {
                 "source": "HF",
                 "repo": "meta/llama4-500B",
             },
             "gather_all_token_logits": False,
         },
-        "runtime": {},
+        "runtime": {"total_token_limit": 100},
     }
     return trtllm_config