feat: onboard dynamo-sdk basic and kv-router examples

container/Dockerfile.vllm

@@ -118,6 +118,10 @@ ENV VLLM_KV_CAPI_PATH="/opt/triton/bindings/lib/libtriton_distributed_llm_capi.s
 # FIXME: Copy more specific folders in for dev/debug after directory restructure
 COPY . /workspace
 
+RUN source /opt/triton/venv/bin/activate && \
+    cd /workspace/deploy/compoundai/sdk && \
+    uv pip install -e .
+
 # FIXME: May want a modification with triton-distributed banner on entry
 ENTRYPOINT ["/opt/nvidia/nvidia_entrypoint.sh"]
 

deploy/compoundai/sdk/pyproject.toml

@@ -19,11 +19,13 @@ version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
 authors = [
-    { name = "mabdulwahhab", email = "mabdulwahhab@nvidia.com" }
+    { name = "mabdulwahhab", email = "mabdulwahhab@nvidia.com" },
+    { name = "biswa.panda", email = "biswaranjanp@nvidia.com" }
 ]
-requires-python = "==3.10.*"
 dependencies = [
-    "bentoml>=1.4.1",
+    # TODO: currently compatible with bentoml==1.4.1, 
+    # v1.4.2 has removed normalize_identifier which is used in compoundai
+    "bentoml==1.4.1",
     "types-psutil==7.0.0.20250218",
 ]
 

deploy/compoundai/sdk/src/compoundai/__init__.py

@@ -12,3 +12,13 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+
+from compoundai.sdk.decorators import nova_endpoint, nova_api, async_onstart
+from compoundai.sdk.service import service
+from compoundai.sdk.dependency import depends
+from bentoml import api
+from bentoml._internal.context import server_context
+from compoundai.sdk.image import NOVA_IMAGE
+
+
+tdist_context = {}

deploy/compoundai/sdk/src/compoundai/cli/serve_nova.py

@@ -26,9 +26,12 @@
 from typing import Any
 
 import click
-from triton_distributed_rs import DistributedRuntime, triton_endpoint, triton_worker
+from triton_distributed.runtime import DistributedRuntime, triton_endpoint, triton_worker
+
+from compoundai import tdist_context
 
 logger = logging.getLogger("compoundai.serve.nova")
+logger.setLevel(logging.INFO)
 
 
 def generate_run_id():
@@ -69,9 +72,9 @@ def main(
     from bentoml._internal.log import configure_server_logging
 
     run_id = generate_run_id()
-
-    # print the contents of the environment variable BENTOML_RUNNER_MAP
-    print(f"BENTOML_RUNNER_MAP: {os.environ['BENTOML_RUNNER_MAP']}")
+    tdist_context["service_name"] = service_name
+    tdist_context["runner_map"] = runner_map
+    tdist_context["worker_id"] = worker_id
 
     # Import service first to check configuration
     service = import_service(bento_identifier)
@@ -100,10 +103,17 @@ def main(
     if service.is_nova_component():
         if worker_id is not None:
             server_context.worker_index = worker_id
-        class_instance = service.inner()
 
         @triton_worker()
         async def worker(runtime: DistributedRuntime):
+            global tdist_context
+            tdist_context["runtime"] = runtime
+            # init service instance, calls __init__
+            #TODO:bis: delete this print
+            print(
+                f"SERVE NOVA triton_worker sn: {service_name} tdist_context: {tdist_context}, runtime: {runtime}"
+            )
+            # logger.info(f"setting tdist_context: {tdist_context}, runtime: {runtime}")
             if service_name and service_name != service.name:
                 server_context.service_type = "service"
             else:
@@ -123,21 +133,6 @@ async def worker(runtime: DistributedRuntime):
                 await component.create_service()
                 logger.info(f"[{run_id}] Created {service.name} component")
 
-                # Run startup hooks before setting up endpoints
-                for name, member in vars(class_instance.__class__).items():
-                    if callable(member) and getattr(
-                        member, "__bentoml_startup_hook__", False
-                    ):
-                        logger.info(f"[{run_id}] Running startup hook: {name}")
-                        result = getattr(class_instance, name)()
-                        if inspect.isawaitable(result):
-                            await result
-                            logger.info(
-                                f"[{run_id}] Completed async startup hook: {name}"
-                            )
-                        else:
-                            logger.info(f"[{run_id}] Completed startup hook: {name}")
-
                 # Set runtime on all dependencies
                 for dep in service.dependencies.values():
                     dep.set_runtime(runtime)
@@ -151,28 +146,51 @@ async def worker(runtime: DistributedRuntime):
                     raise ValueError(error_msg)
 
                 print(f"[{run_id}] Nova endpoints: {nova_endpoints}")
+                endpoints = []
                 for name, endpoint in nova_endpoints.items():
                     td_endpoint = component.endpoint(name)
                     logger.info(f"[{run_id}] Registering endpoint '{name}'")
+                    endpoints.append(td_endpoint)
                     # Bind an instance of inner to the endpoint
+                    # bound_method = endpoint.func.__get__(service.inner())
+                tdist_context["component"] = component
+                tdist_context["endpoints"] = endpoints
+                class_instance = service.inner()
+                twm = []
+                for name, endpoint in nova_endpoints.items():
                     bound_method = endpoint.func.__get__(class_instance)
                     # Only pass request type for now, use Any for response
                     # TODO: Handle a triton_endpoint not having types
                     # TODO: Handle multiple endpoints in a single component
                     triton_wrapped_method = triton_endpoint(endpoint.request_type, Any)(
                         bound_method
                     )
-                    result = await td_endpoint.serve_endpoint(triton_wrapped_method)
-                    # WARNING: unreachable code :( because serve blocks
-                    logger.info(f"[{run_id}] Result: {result}")
-                    logger.info(f"[{run_id}] Registered endpoint '{name}'")
-
-                logger.info(
-                    f"[{run_id}] Started {service.name} instance with all endpoints registered"
-                )
+                    twm.append(triton_wrapped_method)
+                #TODO:bis: delete this print
+                print(f"SETTING sn:{service_name} tdist_context: {tdist_context}")
+                # Run startup hooks before setting up endpoints
+                for name, member in vars(class_instance.__class__).items():
+                    if callable(member) and getattr(
+                        member, "__bentoml_startup_hook__", False
+                    ):
+                        logger.info(f"[{run_id}] Running startup hook: {name}")
+                        result = getattr(class_instance, name)()
+                        if inspect.isawaitable(result):
+                            # await on startup hook async_onstart
+                            await result
+                            logger.info(
+                                f"[{run_id}] Completed async startup hook: {name}"
+                            )
+                        else:
+                            logger.info(f"[{run_id}] Completed startup hook: {name}")
+                #TODO:bis: delete this print
+                print(f"serve nova STARTING sn:{service_name}, endpoints:{endpoints}")
                 logger.info(
-                    f"[{run_id}] Available endpoints: {service.list_nova_endpoints()}"
+                    f"[{run_id}] Starting {service.name} instance with all registered endpoints"
                 )
+                #TODO:bis: convert to list
+                result = await endpoints[0].serve_endpoint(twm[0])
+
             except Exception as e:
                 logger.error(f"[{run_id}] Error in Nova component setup: {str(e)}")
                 raise

deploy/compoundai/sdk/src/compoundai/sdk/dependency.py

@@ -133,6 +133,36 @@ def __init__(
         self._nova_client: Optional[NovaClient] = None
         self._runtime = None
 
+    # offers an escape hatch to get the endpoint directly
+    async def get_endpoint(self, name: str) -> Any:
+        """
+        usage:
+        dep = depends(Worker)
+
+        ...
+        await dep.get_endpoint("generate") # equivalent to the following
+        router_client = (
+            await runtime.namespace("triton-init")
+            .component("router")
+            .endpoint("generate")
+            .client()
+        )
+
+        """
+        # TODO: Read the runtime from the tdist since it is not stored in global
+        if self._runtime is None:
+            print("Get Endpoint: Runtime not set for NovaDependency. Cannot get endpoint.")
+            raise ValueError("Runtime not set for NovaDependency")
+
+        address = self.on.nova_address()
+        comp_ns, comp_name = address
+        print("Get Endpoint: Nova ADDRESS: ", address)
+        return await self._runtime.namespace(comp_ns)\
+            .component(comp_name)\
+            .endpoint(name)\
+            .client()
+
+
     def set_runtime(self, runtime: Any) -> None:
         """Set the Nova runtime for this dependency"""
         self._runtime = runtime

deploy/compoundai/sdk/src/compoundai/sdk/image.py

@@ -16,5 +16,10 @@
 # wrapper over bento images to handle TritonDistributed base image
 
 import bentoml
+import os
 
-NOVA_IMAGE = bentoml.images.PythonImage(base_image="triton-distributed:latest-vllm")
+
+#TODO: "triton-distributed:latest-vllm" image will not be available to image builder
+# so we need to use the nova-base image. We'd consider publishing the image to a public registry.
+image_name = os.getenv("NOVA_IMAGE", "nvcr.io/nvidian/nim-llm-dev/yatai-bentos:nova-base-0886e19")
+NOVA_IMAGE = bentoml.images.PythonImage(base_image=image_name)

deploy/compoundai/sdk/src/compoundai/sdk/service.py

@@ -91,12 +91,8 @@ def nova_address(self) -> Tuple[str, str]:
             except (json.JSONDecodeError, ValueError) as e:
                 raise ValueError(f"Failed to parse BENTOML_RUNNER_MAP: {str(e)}") from e
 
-        # Ensure namespace and name are not None
-        namespace = self._nova_config.namespace or "default"
-        name = self._nova_config.name or self.inner.__name__
-
-        print(f"Using default Nova address: {namespace}/{name}")
-        return (namespace, name)
+        print(f"Using default Nova address: {self._nova_config.namespace}/{self._nova_config.name}")
+        return (self._nova_config.namespace, self._nova_config.name)
 
     def get_nova_endpoints(self) -> Dict[str, NovaEndpoint]:
         """Get all registered Nova endpoints"""