Remove logs and add helper func

Signed-off-by: Indrajit Bhosale <[email protected]>

components/src/dynamo/trtllm/constants.py

@@ -8,4 +8,4 @@ class DisaggregationMode(Enum):
     AGGREGATED = "prefill_and_decode"
     PREFILL = "prefill"
     DECODE = "decode"
-    ENCODE = "encode"
+    ENCODE = "encode"

components/src/dynamo/trtllm/engine.py

@@ -51,7 +51,7 @@ async def initialize(self):
                 self._llm = MultimodalEncoder(
                     model=model,
                     max_batch_size=max_batch_size,
-                )                
+                )
             else:
                 self._llm = self._llm_cls(**self.engine_args)
 
@@ -106,7 +106,9 @@ def _warn_about_unsupported_field(field_name: str) -> None:
 
 
 @asynccontextmanager
-async def get_llm_engine(engine_args, disaggregation_mode: DisaggregationMode) -> AsyncGenerator[TensorRTLLMEngine, None]:
+async def get_llm_engine(
+    engine_args, disaggregation_mode: DisaggregationMode
+) -> AsyncGenerator[TensorRTLLMEngine, None]:
     engine = TensorRTLLMEngine(engine_args, disaggregation_mode)
     try:
         await engine.initialize()

components/src/dynamo/trtllm/main.py

@@ -33,7 +33,7 @@
 from tensorrt_llm.llmapi.tokenizer import tokenizer_factory
 from tensorrt_llm.metrics import MetricsCollector
 from torch.cuda import device_count
-from transformers import AutoConfig, GenerationConfig
+from transformers import AutoConfig
 
 import dynamo.nixl_connect as nixl_connect
 from dynamo.common.config_dump import dump_config

components/src/dynamo/trtllm/utils/disagg_utils.py

@@ -46,7 +46,7 @@ def decode(
             multimodal_embedding_handles=getattr(
                 disaggregated_params, "multimodal_embedding_handles", None
             ),
-            multimodal_hashes=getattr(disaggregated_params, "multimodal_hashes", None),            
+            multimodal_hashes=getattr(disaggregated_params, "multimodal_hashes", None),
         )
 
     @staticmethod
@@ -73,5 +73,5 @@ def encode(
             draft_tokens=disaggregated_params.draft_tokens,
             # E-P Disaggregated Params (for full EPD flow)
             multimodal_embedding_handles=mm_handles,
-            multimodal_hashes=mm_hashes,            
+            multimodal_hashes=mm_hashes,
         )