pass through backend args

bindings/python/pyproject.toml

@@ -30,6 +30,7 @@ dependencies = [
     "setproctitle",
     "grpcio",
     "grpcio-health-checking",
+    "PyYAML",
 ]
 
 [project.optional-dependencies]

bindings/python/src/smg/serve.py

@@ -31,7 +31,9 @@ class WorkerLauncher(ABC):
     """Abstract base class for backend worker launchers."""
 
     @abstractmethod
-    def build_command(self, args: argparse.Namespace, host: str, port: int) -> list[str]:
+    def build_command(
+        self, args: argparse.Namespace, backend_args: list[str], host: str, port: int
+    ) -> list[str]:
         """Build the CLI command list to launch a worker."""
         ...
 
@@ -58,18 +60,38 @@ def gpu_env(self, args: argparse.Namespace, dp_rank: int, env: dict | None = Non
         base_gpu = dp_rank * tp_size
         gpu_ids = range(base_gpu, base_gpu + tp_size)
         env["CUDA_VISIBLE_DEVICES"] = ",".join(str(g) for g in gpu_ids)
+        env["PYTHONUNBUFFERED"] = "1"
         return env
 
+    def _filter_backend_args(self, backend_args: list[str], filter_args: list[str]) -> list[str]:
+        """Filter backend_args to only include those relevant for vLLM's grpc_server."""
+        filtered_backend_args = []
+        skip_next = False
+        for arg in backend_args:
+            if skip_next:
+                skip_next = False
+                continue
+            if arg in filter_args:
+                skip_next = True  # Assume all valid args take a value
+                continue
+            filtered_backend_args.append(arg)
+
+        return filtered_backend_args
+
     def launch(
-        self,
-        args: argparse.Namespace,
-        host: str,
-        port: int,
-        env: dict | None = None,
+        self, args: argparse.Namespace, backend_args: list[str], host: str, port: int, env: dict
     ) -> subprocess.Popen:
         """Launch the worker subprocess."""
-        cmd = self.build_command(args, host, port)
-        return subprocess.Popen(cmd, start_new_session=True, env=env or os.environ.copy())
+        cmd = self.build_command(args, backend_args, host, port)
+        logger.info("Launching worker with command: %s", " ".join(cmd))
+
+        return subprocess.Popen(
+            cmd,
+            start_new_session=True,
+            env=env,
+            stdout=sys.stdout,
+            stderr=sys.stderr,
+        )
 
 
 class SglangWorkerLauncher(WorkerLauncher):
@@ -78,7 +100,9 @@ class SglangWorkerLauncher(WorkerLauncher):
     def _get_tp_size(self, args: argparse.Namespace) -> int:
         return getattr(args, "tp_size", 1)
 
-    def build_command(self, args: argparse.Namespace, host: str, port: int) -> list[str]:
+    def build_command(
+        self, args: argparse.Namespace, backend_args: list[str], host: str, port: int
+    ) -> list[str]:
         cmd = [
             sys.executable,
             "-m",
@@ -92,6 +116,8 @@ def build_command(self, args: argparse.Namespace, host: str, port: int) -> list[
         ]
         if getattr(args, "connection_mode", "grpc") == "grpc":
             cmd.append("--grpc-mode")
+
+        cmd.extend(self._filter_backend_args(backend_args, ["--model-path", "--host", "--port"]))
         return cmd
 
 
@@ -101,20 +127,29 @@ class VllmWorkerLauncher(WorkerLauncher):
     def _get_tp_size(self, args: argparse.Namespace) -> int:
         return getattr(args, "tensor_parallel_size", 1)
 
-    def build_command(self, args: argparse.Namespace, host: str, port: int) -> list[str]:
-        if getattr(args, "connection_mode", "grpc") != "grpc":
-            raise ValueError("vLLM backend only supports grpc connection mode")
-        return [
+    def build_command(
+        self, args: argparse.Namespace, backend_args: list[str], host: str, port: int
+    ) -> list[str]:
+        vllm_entry_points = (
+            "vllm.entrypoints.grpc_server"
+            if getattr(args, "connection_mode", "grpc") == "grpc"
+            else "vllm.entrypoints.openai.api_server"
+        )
+
+        cmd = [
             sys.executable,
             "-m",
-            "vllm.entrypoints.grpc_server",
+            vllm_entry_points,
             "--model",
             getattr(args, "model", ""),
             "--host",
             host,
             "--port",
             str(port),
         ]
+        cmd.extend(self._filter_backend_args(backend_args, ["--model", "--host", "--port"]))
+
+        return cmd
 
 
 class TrtllmWorkerLauncher(WorkerLauncher):
@@ -125,12 +160,46 @@ class TrtllmWorkerLauncher(WorkerLauncher):
     """
 
     def _get_tp_size(self, args: argparse.Namespace) -> int:
-        return getattr(args, "tp_size", 1)
+        """Get tensor parallel size from args or config file.
+
+        Priority: args.tp_size > args.tensor_parallel_size > config file > default(1)
+        """
+        # Try --tp-size argument first
+        tp_size = getattr(args, "tp_size", None)
+        if tp_size is not None:
+            return tp_size
+
+        # Try --tensor-parallel-size (vLLM-style naming)
+        tp_size = getattr(args, "tensor_parallel_size", None)
+        if tp_size is not None:
+            return tp_size
+
+        # Try reading from config YAML file
+        config_path = getattr(args, "config", None)
+        if config_path:
+            try:
+                import yaml
+
+                with open(config_path) as f:
+                    config = yaml.safe_load(f)
+                    if config and "tensor_parallel_size" in config:
+                        return config["tensor_parallel_size"]
+                    if config and "tp_size" in config:
+                        return config["tp_size"]
+            except Exception as e:
+                logger.warning(
+                    "Failed to read tensor_parallel_size from config %s: %s", config_path, e
+                )
 
-    def build_command(self, args: argparse.Namespace, host: str, port: int) -> list[str]:
+        return 1
+
+    def build_command(
+        self, args: argparse.Namespace, backend_args: list[str], host: str, port: int
+    ) -> list[str]:
         if getattr(args, "connection_mode", "grpc") != "grpc":
             raise ValueError("TensorRT-LLM backend only supports grpc connection mode")
-        return [
+
+        cmd = [
             sys.executable,
             "-m",
             "tensorrt_llm.commands.serve",
@@ -142,6 +211,11 @@ def build_command(self, args: argparse.Namespace, host: str, port: int) -> list[
             str(port),
         ]
 
+        # Add optional config file
+        cmd.extend(self._filter_backend_args(backend_args, ["--model", "--host", "--port"]))
+
+        return cmd
+
 
 BACKEND_LAUNCHERS: dict[str, type[WorkerLauncher]] = {
     "sglang": SglangWorkerLauncher,
@@ -271,9 +345,21 @@ def _add_vllm_args(parser: argparse.ArgumentParser) -> None:
 
 
 def _add_trtllm_stub_args(parser: argparse.ArgumentParser) -> None:
-    """Stub for TRT-LLM args until full integration."""
-    group = parser.add_argument_group("TRT-LLM Options (stub)")
-    group.add_argument("--model", type=str, help="Model path")
+    """Add TensorRT-LLM specific arguments.
+
+    Note: TensorRT-LLM doesn't provide a centralized argument manager like
+    vLLM's EngineArgs. We manually add the most commonly used arguments.
+    TP size is read from the config file, not passed as CLI argument.
+    """
+    group = parser.add_argument_group("TensorRT-LLM Options")
+    group.add_argument("--model", type=str, help="Model path (HuggingFace ID or local path)")
+    group.add_argument("--tp-size", type=str, help="Tensor parallel size (overrides config file)")
+    group.add_argument(
+        "--config",
+        type=str,
+        required=False,
+        help="Config file path (YAML, optional - must contain tensor_parallel_size if provided)",
+    )
 
 
 BACKEND_ARG_ADDERS = {
@@ -283,7 +369,7 @@ def _add_trtllm_stub_args(parser: argparse.ArgumentParser) -> None:
 }
 
 BACKEND_CHOICES = list(BACKEND_ARG_ADDERS.keys())
-DEFAULT_BACKEND = "sglang"
+DEFAULT_BACKEND = os.getenv("SMG_DEFAULT_BACKEND", "sglang")
 
 
 # ---------------------------------------------------------------------------
@@ -304,7 +390,7 @@ def add_serve_args(parser: argparse.ArgumentParser) -> None:
         "--connection-mode",
         default="grpc",
         choices=["grpc", "http"],
-        help="Connection mode for workers (default: grpc). Note: vllm and trtllm only support grpc",
+        help="Connection mode for workers (default: grpc). Note: trtllm only support grpc",
     )
     # Router host/port - may be overridden by backend (e.g. sglang)
     group.add_argument(
@@ -353,26 +439,31 @@ def _import_backend_args(backend: str, parser: argparse.ArgumentParser) -> None:
 
 def parse_serve_args(
     argv: list[str] | None = None,
-) -> tuple[str, argparse.Namespace]:
+) -> tuple[str, argparse.Namespace, list[str]]:
     """Two-pass argument parsing for serve command.
 
-    Pass 1: Extract --backend with parse_known_args (no backend imports).
-    Pass 2: Build full parser with backend-specific + router args.
+    Pass 1: Parse with only serve + router args (parse_known_args). We get
+    --backend and known flags; any remaining argv tokens are left as
+    backend_args (no backend-specific parser loaded yet).
+    Pass 2: Build full parser with serve + backend-specific + router args,
+    then parse_args(argv) to get the full namespace. backend_args from pass 1
+    are returned unchanged for the launcher to pass through to worker commands.
 
     Returns:
-        Tuple of (backend_name, parsed_namespace).
+        Tuple of (backend_name, parsed_namespace, backend_args).
+        backend_args: argv tokens not recognized in pass 1 (e.g. --config path).
     """
     if argv is None:
         argv = []
 
-    # Pass 1: extract --backend (lightweight, no backend imports)
+    # Pass 1: serve + router args only; unknown tokens become backend_args
     pre_parser = argparse.ArgumentParser(add_help=False)
-    pre_parser.add_argument("--backend", default=DEFAULT_BACKEND, choices=BACKEND_CHOICES)
-    pre_args, _ = pre_parser.parse_known_args(argv)
-    backend = pre_args.backend
+    add_serve_args(pre_parser)
+    RouterArgs.add_cli_args(pre_parser, use_router_prefix=True, exclude_host_port=True)
+    serve_router_args, backend_args = pre_parser.parse_known_args(argv)
+    backend = serve_router_args.backend
 
-    # Pass 2: build full parser with backend-specific args
-    # Use conflict_handler='resolve' to let backend args override serve defaults
+    # Pass 2: full parser with backend-specific args; resolve so backend can override
     parser = argparse.ArgumentParser(
         description=f"Launch {backend} worker(s) + gateway router",
         conflict_handler="resolve",
@@ -382,7 +473,7 @@ def parse_serve_args(
     RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True)
 
     args = parser.parse_args(argv)
-    return backend, args
+    return backend, args, backend_args
 
 
 # ---------------------------------------------------------------------------
@@ -395,9 +486,10 @@ def parse_serve_args(
 class ServeOrchestrator:
     """Coordinate worker launch, health checking, router startup, and shutdown."""
 
-    def __init__(self, backend: str, args: argparse.Namespace):
+    def __init__(self, backend: str, args: argparse.Namespace, backend_args: list[str]):
         self.backend = backend
         self.args = args
+        self.backend_args = backend_args
         self.launcher: WorkerLauncher = BACKEND_LAUNCHERS[backend]()
         self.workers: list[tuple[subprocess.Popen, int]] = []
         self._shutting_down = False
@@ -424,7 +516,7 @@ def _launch_workers(self) -> None:
         host = self.args.worker_host
         for dp_rank, port in enumerate(ports):
             env = self.launcher.gpu_env(self.args, dp_rank)
-            proc = self.launcher.launch(self.args, host, port, env)
+            proc = self.launcher.launch(self.args, self.backend_args, host, port, env)
             self.workers.append((proc, port))
             logger.info(
                 "Launched %s worker on %s:%d (pid %d, GPUs: %s)",
@@ -500,6 +592,10 @@ def _cleanup_workers(self) -> None:
 
 def serve_main(argv: list[str] | None = None) -> None:
     """Parse serve args, create orchestrator, and run."""
-    backend, args = parse_serve_args(argv)
-    orchestrator = ServeOrchestrator(backend, args)
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+    backend, args, backend_args = parse_serve_args(argv)
+    orchestrator = ServeOrchestrator(backend, args, backend_args)
     orchestrator.run()