Added main logic

Author: oandreeva-nv

python/openai/openai_frontend/engine/triton_engine.py

@@ -1,4 +1,4 @@
-# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -78,11 +78,22 @@ class TritonModelMetadata:
 
 class TritonLLMEngine(LLMEngine):
     def __init__(
-        self, server: tritonserver.Server, tokenizer: str, backend: Optional[str] = None
+        self,
+        server: tritonserver.Server,
+        tokenizer_map: Dict[str, str] = None,
+        backend: Optional[str] = None,
     ):
         # Assume an already configured and started server
         self.server = server
-        self.tokenizer = self._get_tokenizer(tokenizer)
+        self.tokenizer_map = {}
+        if tokenizer_map:
+            for model_name, tokenizer_path in tokenizer_map.items():
+                try:
+                    self.tokenizer_map[model_name] = get_tokenizer(tokenizer_path)
+                except Exception as e:
+                    print(
+                        f"Warning: Failed to load tokenizer for {model_name} from {tokenizer_path}: {e}"
+                    )
         # TODO: Reconsider name of "backend" vs. something like "request_format"
         self.backend = backend
 
@@ -253,12 +264,12 @@ def _get_model_metadata(self) -> Dict[str, TritonModelMetadata]:
             if not backend and model.config()["platform"] == "ensemble":
                 backend = "ensemble"
             print(f"Found model: {name=}, {backend=}")
-
+            default_tokenizer = self.tokenizer_map.get("default", None)
             metadata = TritonModelMetadata(
                 name=name,
                 backend=backend,
                 model=model,
-                tokenizer=self.tokenizer,
+                tokenizer=self.tokenizer_map.get(name, default_tokenizer),
                 create_time=self.create_time,
                 request_converter=self._determine_request_converter(backend),
             )

python/openai/openai_frontend/main.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -95,12 +95,20 @@ def parse_args():
         required=True,
         help="Path to the Triton model repository holding the models to be served",
     )
+    # TODO: determine what to do with single tokenizer flag
     triton_group.add_argument(
         "--tokenizer",
         type=str,
         default=None,
         help="HuggingFace ID or local folder path of the Tokenizer to use for chat templates",
     )
+    triton_group.add_argument(
+        "--tokenizers",
+        type=str,
+        nargs="+",  # Accept multiple arguments
+        default=[],
+        help="List of HuggingFace IDs or local folder paths of Tokenizers to use. Format: model_name:tokenizer_path",
+    )
     triton_group.add_argument(
         "--backend",
         type=str,
@@ -160,8 +168,22 @@ def parse_args():
 def main():
     args = parse_args()
 
-    # Initialize a Triton Inference Server pointing at LLM models
-    server: tritonserver.Server = tritonserver.Server(
+    # Parse tokenizer mappings
+    tokenizer_map = {}
+    for tokenizer_spec in args.tokenizers:
+        try:
+            model_name, tokenizer_path = tokenizer_spec.split(":")
+            tokenizer_map[model_name] = tokenizer_path
+        except ValueError:
+            print(
+                f"Warning: Skipping invalid tokenizer specification: {tokenizer_spec}. Format should be 'model_name:tokenizer_path'"
+            )
+
+    if args.tokenizer:
+        tokenizer_map["default"] = args.tokenizer
+
+    # Initialize Triton server
+    server = tritonserver.Server(
         model_repository=args.model_repository,
         log_verbose=args.tritonserver_log_verbose_level,
         log_info=True,
@@ -170,8 +192,8 @@ def main():
     ).start(wait_until_ready=True)
 
     # Wrap Triton Inference Server in an interface-conforming "LLMEngine"
-    engine: TritonLLMEngine = TritonLLMEngine(
-        server=server, tokenizer=args.tokenizer, backend=args.backend
+    engine = TritonLLMEngine(
+        server=server, tokenizer_map=tokenizer_map, backend=args.backend
     )
 
     # Attach TritonLLMEngine as the backbone for inference and model management