NeMo-Export-Deploy 0.1.0__py3-none-any.whl

nemo_deploy/__init__.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 nemo_deploy.deploy_base import DeployBase
+from nemo_deploy.deploy_pytriton import DeployPyTriton
+from nemo_deploy.triton_deployable import ITritonDeployable
+from nemo_export_deploy_common.package_info import __package_name__, __version__
+
+__all__ = [
+    "DeployBase",
+    "DeployPyTriton",
+    "ITritonDeployable",
+    "__version__",
+    "__package_name__",
+]

nemo_deploy/deploy_base.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import logging
+from abc import ABC, abstractmethod
+
+from nemo_deploy.triton_deployable import ITritonDeployable
+
+LOGGER = logging.getLogger("NeMo")
+
+
+class DeployBase(ABC):
+    def __init__(
+        self,
+        triton_model_name: str,
+        triton_model_version: int = 1,
+        model=None,
+        max_batch_size: int = 128,
+        http_port: int = 8000,
+        grpc_port: int = 8001,
+        address="0.0.0.0",
+        allow_grpc=True,
+        allow_http=True,
+        streaming=False,
+    ):
+        self.triton_model_name = triton_model_name
+        self.triton_model_version = triton_model_version
+        self.max_batch_size = max_batch_size
+        self.model = model
+        self.http_port = http_port
+        self.grpc_port = grpc_port
+        self.address = address
+        self.triton = None
+        self.allow_grpc = allow_grpc
+        self.allow_http = allow_http
+        self.streaming = streaming
+
+    @abstractmethod
+    def deploy(self):
+        pass
+
+    @abstractmethod
+    def serve(self):
+        pass
+
+    @abstractmethod
+    def run(self):
+        pass
+
+    @abstractmethod
+    def stop(self):
+        pass
+
+    def _is_model_deployable(self):
+        if not issubclass(type(self.model), ITritonDeployable):
+            raise Exception(
+                "This model is not deployable to Triton.nemo_deploy.ITritonDeployable class should be inherited"
+            )
+        else:
+            return True

nemo_deploy/deploy_pytriton.py

@@ -0,0 +1,173 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+
+import logging
+
+from nemo_deploy.deploy_base import DeployBase
+from nemo_export_deploy_common.import_utils import MISSING_TRITON_MSG, UnavailableError
+
+LOGGER = logging.getLogger("NeMo")
+
+try:
+    from pytriton.model_config import ModelConfig
+    from pytriton.triton import Triton, TritonConfig
+
+    HAVE_TRITON = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_TRITON = False
+
+
+class DeployPyTriton(DeployBase):
+    """Deploys any models to Triton Inference Server that implements ITritonDeployable interface in nemo_deploy.
+
+    Example:
+        from nemo_deploy import DeployPyTriton, NemoQueryLLM
+        from nemo_export.tensorrt_llm import TensorRTLLM
+
+        trt_llm_exporter = TensorRTLLM(model_dir="/path/for/model/files")
+        trt_llm_exporter.export(
+            nemo_checkpoint_path="/path/for/nemo/checkpoint",
+            model_type="llama",
+            tensor_parallelism_size=1,
+        )
+
+        nm = DeployPyTriton(model=trt_llm_exporter, triton_model_name="model_name", http_port=8000)
+        nm.deploy()
+        nm.run()
+        nq = NemoQueryLLM(url="localhost", model_name="model_name")
+
+        prompts = ["hello, testing GPT inference", "another GPT inference test?"]
+        output = nq.query_llm(prompts=prompts, max_output_len=100)
+        print("prompts: ", prompts)
+        print("")
+        print("output: ", output)
+        print("")
+
+        prompts = ["Give me some info about Paris", "Do you think Londan is a good city to visit?", "What do you think about Rome?"]
+        output = nq.query_llm(prompts=prompts, max_output_len=250)
+        print("prompts: ", prompts)
+        print("")
+        print("output: ", output)
+        print("")
+
+    """
+
+    def __init__(
+        self,
+        triton_model_name: str,
+        triton_model_version: int = 1,
+        model=None,
+        max_batch_size: int = 128,
+        http_port: int = 8000,
+        grpc_port: int = 8001,
+        address="0.0.0.0",
+        allow_grpc=True,
+        allow_http=True,
+        streaming=False,
+        pytriton_log_verbose=0,
+    ):
+        """A nemo checkpoint or model is expected for serving on Triton Inference Server.
+
+        Args:
+            triton_model_name (str): Name for the service
+            triton_model_version(int): Version for the service
+            checkpoint_path (str): path of the nemo file
+            model (ITritonDeployable): A model that implements the ITritonDeployable from nemo_deploy import ITritonDeployable
+            max_batch_size (int): max batch size
+            port (int) : port for the Triton server
+            address (str): http address for Triton server to bind.
+        """
+        super().__init__(
+            triton_model_name=triton_model_name,
+            triton_model_version=triton_model_version,
+            model=model,
+            max_batch_size=max_batch_size,
+            http_port=http_port,
+            grpc_port=grpc_port,
+            address=address,
+            allow_grpc=allow_grpc,
+            allow_http=allow_http,
+            streaming=streaming,
+        )
+        self.pytriton_log_verbose = pytriton_log_verbose
+
+    def deploy(self):
+        """Deploys any models to Triton Inference Server."""
+        if not HAVE_TRITON:
+            raise UnavailableError(MISSING_TRITON_MSG)
+
+        try:
+            if self.streaming:
+                triton_config = TritonConfig(
+                    log_verbose=self.pytriton_log_verbose,
+                    allow_grpc=self.allow_grpc,
+                    allow_http=self.allow_http,
+                    grpc_address=self.address,
+                )
+                self.triton = Triton(config=triton_config)
+                self.triton.bind(
+                    model_name=self.triton_model_name,
+                    model_version=self.triton_model_version,
+                    infer_func=self.model.triton_infer_fn_streaming,
+                    inputs=self.model.get_triton_input,
+                    outputs=self.model.get_triton_output,
+                    config=ModelConfig(decoupled=True),
+                )
+            else:
+                triton_config = TritonConfig(
+                    http_address=self.address,
+                    http_port=self.http_port,
+                    grpc_address=self.address,
+                    grpc_port=self.grpc_port,
+                    allow_grpc=self.allow_grpc,
+                    allow_http=self.allow_http,
+                )
+                self.triton = Triton(config=triton_config)
+                self.triton.bind(
+                    model_name=self.triton_model_name,
+                    model_version=self.triton_model_version,
+                    infer_func=self.model.triton_infer_fn,
+                    inputs=self.model.get_triton_input,
+                    outputs=self.model.get_triton_output,
+                    config=ModelConfig(max_batch_size=self.max_batch_size),
+                )
+        except Exception as e:
+            self.triton = None
+            LOGGER.error(e)
+
+    def serve(self):
+        """Starts serving the model and waits for the requests."""
+        if self.triton is None:
+            raise Exception("deploy should be called first.")
+
+        try:
+            self.triton.serve()
+        except Exception as e:
+            self.triton = None
+            LOGGER.error(e)
+
+    def run(self):
+        """Starts serving the model asynchronously."""
+        if self.triton is None:
+            raise Exception("deploy should be called first.")
+
+        self.triton.run()
+
+    def stop(self):
+        """Stops serving the model."""
+        if self.triton is None:
+            raise Exception("deploy should be called first.")
+
+        self.triton.stop()

nemo_deploy/deploy_ray.py

@@ -0,0 +1,140 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+
+import logging
+
+from nemo_deploy.ray_utils import find_available_port
+from nemo_export_deploy_common.import_utils import MISSING_RAY_MSG, UnavailableError
+
+try:
+    import ray
+    from ray import serve
+    from ray.serve import Application
+
+    HAVE_RAY = True
+except (ImportError, ModuleNotFoundError):
+    from unittest.mock import MagicMock
+
+    ray = MagicMock()
+    serve = MagicMock()
+    Application = MagicMock()
+    HAVE_RAY = False
+
+LOGGER = logging.getLogger("NeMo")
+
+
+class DeployRay:
+    """A class for managing Ray deployment and serving of models.
+
+    This class provides functionality to initialize Ray, start Ray Serve,
+    deploy models, and manage the lifecycle of the Ray cluster.
+
+    Attributes:
+        address (str): The address of the Ray cluster to connect to.
+        num_cpus (int): Number of CPUs to allocate for the Ray cluster.
+        num_gpus (int): Number of GPUs to allocate for the Ray cluster.
+        include_dashboard (bool): Whether to include the Ray dashboard.
+        ignore_reinit_error (bool): Whether to ignore errors when reinitializing Ray.
+        runtime_env (dict): Runtime environment configuration for Ray.
+    """
+
+    def __init__(
+        self,
+        address: str = "auto",
+        num_cpus: int = 1,
+        num_gpus: int = 1,
+        include_dashboard: bool = False,
+        ignore_reinit_error: bool = True,
+        runtime_env: dict = None,
+    ):
+        """Initialize the DeployRay instance and set up the Ray cluster.
+
+        Args:
+            address (str, optional): Address of the Ray cluster. Defaults to "auto".
+            num_cpus (int, optional): Number of CPUs to allocate. Defaults to 1.
+            num_gpus (int, optional): Number of GPUs to allocate. Defaults to 1.
+            include_dashboard (bool, optional): Whether to include the dashboard. Defaults to False.
+            ignore_reinit_error (bool, optional): Whether to ignore reinit errors. Defaults to True.
+            runtime_env (dict, optional): Runtime environment configuration. Defaults to None.
+
+        Raises:
+            Exception: If Ray is not installed.
+        """
+        if not HAVE_RAY:
+            raise UnavailableError(MISSING_RAY_MSG)
+
+        # Initialize Ray with proper configuration
+
+        try:
+            # Try to connect to existing Ray cluster
+            ray.init(
+                address=address,
+                ignore_reinit_error=ignore_reinit_error,
+                runtime_env=runtime_env,
+            )
+        except ConnectionError:
+            # If no cluster exists, start a local one
+            LOGGER.info("No existing Ray cluster found. Starting a local Ray cluster...")
+            ray.init(
+                num_cpus=num_cpus,
+                num_gpus=num_gpus,
+                include_dashboard=include_dashboard,
+                ignore_reinit_error=ignore_reinit_error,
+                runtime_env=runtime_env,
+            )
+
+    def start(self, host: str = "0.0.0.0", port: int = None):
+        """Start Ray Serve with the specified host and port.
+
+        Args:
+            host (str, optional): Host address to bind to. Defaults to "0.0.0.0".
+            port (int, optional): Port number to use. If None, an available port will be found.
+        """
+        if not port:
+            port = find_available_port(8000, host)
+        serve.start(
+            http_options={
+                "host": host,
+                "port": port,
+            }
+        )
+
+    def run(self, app: Application, model_name: str):
+        """Deploy and start serving a model using Ray Serve.
+
+        Args:
+            app (Application): The Ray Serve application to deploy.
+            model_name (str): Name to give to the deployed model.
+        """
+        serve.run(app, name=model_name)
+
+    def stop(self):
+        """Stop the Ray Serve deployment and shutdown the Ray cluster.
+
+        This method attempts to gracefully shutdown both Ray Serve and the Ray cluster.
+        If any errors occur during shutdown, they are logged as warnings.
+        """
+        try:
+            # First try to gracefully shutdown Ray Serve
+            LOGGER.info("Shutting down Ray Serve...")
+            serve.shutdown()
+        except Exception as e:
+            LOGGER.warning(f"Error during serve.shutdown(): {str(e)}")
+        try:
+            # Then try to gracefully shutdown Ray
+            LOGGER.info("Shutting down Ray...")
+            ray.shutdown()
+        except Exception as e:
+            LOGGER.warning(f"Error during ray.shutdown(): {str(e)}")

nemo_deploy/multimodal/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 nemo_deploy.multimodal.query_multimodal import NemoQueryMultimodal

nemo_deploy/multimodal/query_multimodal.py

@@ -0,0 +1,184 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 io import BytesIO
+
+import numpy as np
+import requests
+
+from nemo_deploy.utils import str_list2numpy
+from nemo_export_deploy_common.import_utils import (
+    MISSING_DECORD_MSG,
+    MISSING_PIL_MSG,
+    MISSING_TRITON_MSG,
+    UnavailableError,
+)
+
+try:
+    from PIL import Image
+
+    HAVE_PIL = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_PIL = False
+
+try:
+    from decord import VideoReader
+
+    HAVE_DECORD = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_DECORD = False
+
+try:
+    from pytriton.client import ModelClient
+
+    HAVE_TRITON = True
+except (ImportError, ModuleNotFoundError):
+    from unittest.mock import MagicMock
+
+    ModelClient = MagicMock()
+    HAVE_TRITON = False
+
+
+class NemoQueryMultimodal:
+    """Sends a query to Triton for Multimodal inference.
+
+    Example:
+        from nemo_deploy.multimodal import NemoQueryMultimodal
+
+        nq = NemoQueryMultimodal(url="localhost", model_name="neva", model_type="neva")
+
+        input_text = "Hi! What is in this image?"
+        output = nq.query(
+            input_text=input_text,
+            input_media="/path/to/image.jpg",
+            max_output_len=30,
+            top_k=1,
+            top_p=0.0,
+            temperature=1.0,
+        )
+        print("prompts: ", prompts)
+    """
+
+    def __init__(self, url, model_name, model_type):
+        self.url = url
+        self.model_name = model_name
+        self.model_type = model_type
+
+    def setup_media(self, input_media):
+        """Setup input media."""
+        if self.model_type == "video-neva":
+            if not HAVE_DECORD:
+                raise UnavailableError(MISSING_DECORD_MSG)
+
+            vr = VideoReader(input_media)
+            frames = [f.asnumpy() for f in vr]
+            return np.array(frames)
+        elif self.model_type == "lita" or self.model_type == "vita":
+            if not HAVE_DECORD:
+                raise UnavailableError(MISSING_DECORD_MSG)
+
+            vr = VideoReader(input_media)
+            frames = [f.asnumpy() for f in vr]
+            subsample_len = self.frame_len(frames)
+            sub_frames = self.get_subsampled_frames(frames, subsample_len)
+            return np.array(sub_frames)
+        elif self.model_type in ["neva", "vila", "mllama"]:
+            if not HAVE_PIL:
+                raise UnavailableError(MISSING_PIL_MSG)
+
+            if input_media.startswith("http") or input_media.startswith("https"):
+                response = requests.get(input_media, timeout=5)
+                media = Image.open(BytesIO(response.content)).convert("RGB")
+            else:
+                media = Image.open(input_media).convert("RGB")
+            return np.expand_dims(np.array(media), axis=0)
+        else:
+            raise RuntimeError(f"Invalid model type {self.model_type}")
+
+    def frame_len(self, frames):
+        """Get frame len."""
+        max_frames = 256
+        if len(frames) <= max_frames:
+            return len(frames)
+        else:
+            subsample = int(np.ceil(float(len(frames)) / max_frames))
+            return int(np.round(float(len(frames)) / subsample))
+
+    def get_subsampled_frames(self, frames, subsample_len):
+        """Get subsampled frames."""
+        idx = np.round(np.linspace(0, len(frames) - 1, subsample_len)).astype(int)
+        sub_frames = [frames[i] for i in idx]
+        return sub_frames
+
+    def query(
+        self,
+        input_text,
+        input_media,
+        batch_size=1,
+        max_output_len=30,
+        top_k=1,
+        top_p=0.0,
+        temperature=1.0,
+        repetition_penalty=1.0,
+        num_beams=1,
+        init_timeout=60.0,
+        lora_uids=None,
+    ):
+        """Run query."""
+        if not HAVE_TRITON:
+            raise UnavailableError(MISSING_TRITON_MSG)
+
+        prompts = str_list2numpy([input_text])
+        inputs = {"input_text": prompts}
+
+        media = self.setup_media(input_media)
+        if isinstance(media, dict):
+            inputs.update(media)
+        else:
+            inputs["input_media"] = np.repeat(media[np.newaxis, :, :, :, :], prompts.shape[0], axis=0)
+
+        if batch_size is not None:
+            inputs["batch_size"] = np.full(prompts.shape, batch_size, dtype=np.int_)
+
+        if max_output_len is not None:
+            inputs["max_output_len"] = np.full(prompts.shape, max_output_len, dtype=np.int_)
+
+        if top_k is not None:
+            inputs["top_k"] = np.full(prompts.shape, top_k, dtype=np.int_)
+
+        if top_p is not None:
+            inputs["top_p"] = np.full(prompts.shape, top_p, dtype=np.single)
+
+        if temperature is not None:
+            inputs["temperature"] = np.full(prompts.shape, temperature, dtype=np.single)
+
+        if repetition_penalty is not None:
+            inputs["repetition_penalty"] = np.full(prompts.shape, repetition_penalty, dtype=np.single)
+
+        if num_beams is not None:
+            inputs["num_beams"] = np.full(prompts.shape, num_beams, dtype=np.int_)
+
+        if lora_uids is not None:
+            lora_uids = np.char.encode(lora_uids, "utf-8")
+            inputs["lora_uids"] = np.full((prompts.shape[0], len(lora_uids)), lora_uids)
+
+        with ModelClient(self.url, self.model_name, init_timeout_s=init_timeout) as client:
+            result_dict = client.infer_batch(**inputs)
+            output_type = client.model_config.outputs[0].dtype
+
+            if output_type == np.bytes_:
+                sentences = np.char.decode(result_dict["outputs"].astype("bytes"), "utf-8")
+                return sentences
+            else:
+                return result_dict["outputs"]

nemo_deploy/nlp/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 nemo_deploy.nlp.query_llm import NemoQueryLLM, NemoQueryLLMHF, NemoQueryLLMPyTorch, NemoQueryTRTLLMAPI
+
+__all__ = [
+    "NemoQueryLLM",
+    "NemoQueryLLMHF",
+    "NemoQueryLLMPyTorch",
+    "NemoQueryTRTLLMAPI",
+]