refactor: move worker init logic from main.py to workers/ module (#6063)

Signed-off-by: Yuewei Na <nv-yna@users.noreply.github.com>
Co-authored-by: Yuewei Na <nv-yna@users.noreply.github.com>

components/src/dynamo/trtllm/workers/__init__.py

@@ -4,9 +4,52 @@
 """Worker initialization modules for TensorRT-LLM backend.
 
 This package contains worker initialization functions for different modalities:
+- llm_worker: Text and multimodal LLM inference
 - video_diffusion_worker: Video generation using diffusion models
+
+The init_worker() function dispatches to the appropriate worker based on modality.
 """
 
-from dynamo.trtllm.workers.video_diffusion_worker import init_video_diffusion_worker
+import asyncio
+import logging
+
+from dynamo.runtime import DistributedRuntime
+from dynamo.trtllm.constants import Modality
+from dynamo.trtllm.utils.trtllm_utils import Config
+
+
+async def init_worker(
+    runtime: DistributedRuntime, config: Config, shutdown_event: asyncio.Event
+) -> None:
+    """Initialize the appropriate worker based on modality.
+
+    Dispatches to the correct worker initialization function based on the
+    configured modality (text, multimodal, video_diffusion, etc.).
+
+    Args:
+        runtime: The Dynamo distributed runtime.
+        config: Configuration parsed from command line.
+        shutdown_event: Event to signal shutdown.
+    """
+    logging.info(f"Initializing worker with modality={config.modality}")
+
+    modality = Modality(config.modality)
+
+    if Modality.is_diffusion(modality):
+        if modality == Modality.VIDEO_DIFFUSION:
+            from dynamo.trtllm.workers.video_diffusion_worker import (
+                init_video_diffusion_worker,
+            )
+
+            await init_video_diffusion_worker(runtime, config, shutdown_event)
+            return
+        # TODO: Add IMAGE_DIFFUSION support in follow-up PR
+        raise ValueError(f"Unsupported diffusion modality: {modality}")
+
+    # LLM modalities (text, multimodal) use the LLM worker
+    from dynamo.trtllm.workers.llm_worker import init_llm_worker
+
+    await init_llm_worker(runtime, config, shutdown_event)
+
 
-__all__ = ["init_video_diffusion_worker"]
+__all__ = ["init_worker"]