f

Signed-off-by: oliver könig <okoenig@nvidia.com>

Author: ko3n1g

nemo_export/trt_llm/converter/model_converter.py

@@ -16,7 +16,6 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
-import tensorrt_llm
 import torch
 
 from nemo_export.trt_llm.converter.model_to_trt_llm_ckpt import (
@@ -27,6 +26,7 @@
 from nemo_export_deploy_common.import_utils import UnavailableError
 
 try:
+    import tensorrt_llm
     from tensorrt_llm._utils import pad_vocab_size
     from tensorrt_llm.functional import non_gated_version
     from tensorrt_llm.layers import MoeConfig

nemo_export/trt_llm/converter/utils.py

@@ -18,7 +18,13 @@
 import numpy as np
 import tensorrt_llm
 import torch
-from tensorrt_llm._utils import mpi_comm, torch_to_numpy
+
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from tensorrt_llm._utils import mpi_comm, torch_to_numpy
+except (ImportError, ModuleNotFoundError) as e:
+    raise UnavailableError("tensorrt_llm is not available. Please install it with `pip install tensorrt-llm`.") from e
 
 # A global dicts to store exported weights.
 # This is set to be a global variable to avoid extra code modification from tensorrt_llm.

nemo_export/trt_llm/tensorrt_llm_run.py

@@ -24,19 +24,23 @@
 
 import numpy as np
 import tensorrt as trt
-import tensorrt_llm
 import torch
-from mpi4py.futures import MPIPoolExecutor
-from tensorrt_llm.builder import Engine
-from tensorrt_llm.lora_manager import LoraManager
-from tensorrt_llm.quantization import QuantMode
 from transformers import PreTrainedTokenizer
 
 from nemo_export_deploy_common.import_utils import UnavailableError
 
 LOGGER = logging.getLogger("NeMo")
 
 try:
+    from mpi4py.futures import MPIPoolExecutor
+except (ImportError, ModuleNotFoundError) as e:
+    raise UnavailableError("mpi4py is not available. Please install it with `pip install mpi4py`.") from e
+
+try:
+    import tensorrt_llm
+    from tensorrt_llm.builder import Engine
+    from tensorrt_llm.lora_manager import LoraManager
+    from tensorrt_llm.quantization import QuantMode
     from tensorrt_llm.runtime import (
         ModelConfig,
         ModelRunner,

nemo_export/vllm/model_loader.py

@@ -18,12 +18,17 @@
 
 import safetensors.torch
 import torch
-from vllm.config import ModelConfig
-from vllm.model_executor.model_loader import BaseModelLoader, get_model
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
 
 from nemo_export.utils import load_model_weights
 from nemo_export.vllm.model_config import NemoModelConfig
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from vllm.config import ModelConfig
+    from vllm.model_executor.model_loader import BaseModelLoader, get_model
+    from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("vllm is not installed. Please install it with `pip install vllm`.")
 
 LOGGER = logging.getLogger("NeMo")
 

nemo_export/vllm_hf_exporter.py

@@ -16,13 +16,22 @@
 from typing import List
 
 import numpy as np
-from pytriton.decorators import batch, first_value
-from pytriton.model_config import Tensor
-from vllm import LLM, SamplingParams
-from vllm.lora.request import LoRARequest
 
 from nemo_deploy import ITritonDeployable
 from nemo_deploy.utils import cast_output, str_ndarray2list
+from nemo_export_deploy_common.import_utils import UnavailableError
+
+try:
+    from pytriton.decorators import batch, first_value
+    from pytriton.model_config import Tensor
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("pytriton is not installed. Please install it with `pip install nvidia-pytriton`.")
+
+try:
+    from vllm import LLM, SamplingParams
+    from vllm.lora.request import LoRARequest
+except (ImportError, ModuleNotFoundError):
+    raise UnavailableError("vllm is not installed. Please install it with `pip install vllm`.")
 
 
 class vLLMHFExporter(ITritonDeployable):