NeMo-Export-Deploy 0.7.0__tar.gz → 0.8.0__tar.gz

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/NeMo_Export_Deploy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NeMo-Export-Deploy
-Version: 0.7.0
+Version: 0.8.0
 Summary: NeMo Export and Deploy - a library to export and deploy LLMs and MMs
 License: Apache 2.0
 Requires-Python: <3.13,>=3.12
@@ -31,7 +31,7 @@ Requires-Dist: torchmetrics>=0.11.0
 Requires-Dist: wandb==0.27.2
 Requires-Dist: webdataset>=0.2.86
 Requires-Dist: nvidia-pytriton; platform_system != "Darwin"
-Requires-Dist: flashinfer-python==0.6.8.post1; platform_system != "Darwin"
+Requires-Dist: flashinfer-python; platform_system != "Darwin"
 Requires-Dist: Pillow; platform_system != "Darwin" and platform_machine != "aarch64"
 Requires-Dist: decord; platform_system != "Darwin" and platform_machine != "aarch64"
 Requires-Dist: pyparsing>2.0.2
@@ -42,7 +42,7 @@ Requires-Dist: ijson
 Requires-Dist: pyarrow<21.0.0,>=20.0.0
 Provides-Extra: inframework
 Provides-Extra: vllm
-Requires-Dist: vllm~=0.24.0; extra == "vllm"
+Requires-Dist: vllm~=0.29.0; extra == "vllm"
 Requires-Dist: pandas; extra == "vllm"
 Requires-Dist: timm; extra == "vllm"
 Provides-Extra: trt-onnx

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/NeMo_Export_Deploy.egg-info/SOURCES.txt

@@ -26,6 +26,7 @@ nemo_deploy/llm/inference/nemo_io.py
 nemo_deploy/llm/inference/nemo_utils.py
 nemo_deploy/llm/inference/tron_utils.py
 nemo_deploy/multimodal/__init__.py
+nemo_deploy/multimodal/image_url_validator.py
 nemo_deploy/multimodal/megatron_multimodal_deployable.py
 nemo_deploy/multimodal/megatron_multimodal_deployable_ray.py
 nemo_deploy/multimodal/query_multimodal.py

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/NeMo_Export_Deploy.egg-info/requires.txt

@@ -28,7 +28,7 @@ pyarrow<21.0.0,>=20.0.0
 
 [:platform_system != "Darwin"]
 nvidia-pytriton
-flashinfer-python==0.6.8.post1
+flashinfer-python
 
 [:platform_system != "Darwin" and platform_machine != "aarch64"]
 Pillow
@@ -48,6 +48,6 @@ onnxscript>=0.6.0
 transformers==4.51.3
 
 [vllm]
-vllm~=0.24.0
+vllm~=0.29.0
 pandas
 timm

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NeMo-Export-Deploy
-Version: 0.7.0
+Version: 0.8.0
 Summary: NeMo Export and Deploy - a library to export and deploy LLMs and MMs
 License: Apache 2.0
 Requires-Python: <3.13,>=3.12
@@ -31,7 +31,7 @@ Requires-Dist: torchmetrics>=0.11.0
 Requires-Dist: wandb==0.27.2
 Requires-Dist: webdataset>=0.2.86
 Requires-Dist: nvidia-pytriton; platform_system != "Darwin"
-Requires-Dist: flashinfer-python==0.6.8.post1; platform_system != "Darwin"
+Requires-Dist: flashinfer-python; platform_system != "Darwin"
 Requires-Dist: Pillow; platform_system != "Darwin" and platform_machine != "aarch64"
 Requires-Dist: decord; platform_system != "Darwin" and platform_machine != "aarch64"
 Requires-Dist: pyparsing>2.0.2
@@ -42,7 +42,7 @@ Requires-Dist: ijson
 Requires-Dist: pyarrow<21.0.0,>=20.0.0
 Provides-Extra: inframework
 Provides-Extra: vllm
-Requires-Dist: vllm~=0.24.0; extra == "vllm"
+Requires-Dist: vllm~=0.29.0; extra == "vllm"
 Requires-Dist: pandas; extra == "vllm"
 Requires-Dist: timm; extra == "vllm"
 Provides-Extra: trt-onnx

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_deploy/llm/inference/inference_base.py

@@ -16,6 +16,7 @@
 import atexit
 import logging
 import os
+import socket
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
 
@@ -414,6 +415,22 @@ class MCoreEngineWithCleanup:
         return getattr(self.mcore_engine, name)
 
 
+def _default_coordinator_host() -> str:
+    """Return a bindable address for the data-parallel inference coordinator.
+
+    MCore defaults the coordinator host to ``socket.gethostname()``, which is a
+    *name*, not an address. Inside Docker it resolves to the container ID and ZMQ
+    cannot bind it (``ZMQError: No such device``). Resolving that name to its IP
+    yields an address that is both bindable and reachable from other containers,
+    which is what the Ray path already gets by exporting the node IP as
+    ``MASTER_ADDR``. Fall back to loopback only if resolution fails.
+    """
+    try:
+        return socket.gethostbyname(socket.gethostname())
+    except OSError:
+        return "127.0.0.1"
+
+
 def create_mcore_engine(
     path: Path,
     num_devices: Optional[int] = None,
@@ -519,14 +536,10 @@ def create_mcore_engine(
         buffer_size_gb=int(buffer_size_gb),
         max_requests=max_batch_size,
         block_size_tokens=block_size_tokens,
-        # Coordinator mode constructs one DynamicInferenceEngine per rank independently, so a
-        # per-request runtime toggle (as _infer_fn used to do) only ever reaches the primary
-        # rank's engine. Prompt log probs require this to be False on every rank, so it's set
-        # here unconditionally at construction time instead.
-        materialize_only_last_token_logits=False,
+        materialize_only_last_token_logits=True,
     )
 
-    coordinator_host = os.environ.get("MASTER_ADDR")
+    coordinator_host = os.environ.get("MASTER_ADDR") or _default_coordinator_host()
 
     llm = MegatronLLM(
         model=model,

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_deploy/llm/inference/tron_utils.py

@@ -261,7 +261,7 @@ def _set_random_seed(
     te_rng_tracker: bool = False,
     inference_rng_tracker: bool = False,
 ) -> None:
-    """Set random seed for reproducability.
+    """Set random seed for reproducibility.
 
     Args:
         seed_ (int): Base random seed to use
@@ -332,7 +332,7 @@ def _initialize_tp_communicators(model_config: Union[GPTConfig, T5Config], micro
             bootstrap_backend=bootstrap_backend,
         )
     except TypeError:
-        # Fallback for older TE versions
+        # Fallback for older Transformer Engine versions
         if bootstrap_backend != "mpi":
             LOGGER.info("Warning: Transformer Engine may only support MPI bootstrap backend")
 

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_deploy/llm/megatronllm_deployable.py

@@ -69,7 +69,7 @@ class MegatronLLMDeployable(ITritonDeployable):
         context_parallel_size (int): context parallelism.
         expert_model_parallel_size (int): expert parallelism.
         params_dtype (torch.dtype): max input length.
-        inference_batch_times_seqlen_threshold (int): squence threshold.
+        inference_batch_times_seqlen_threshold (int): sequence threshold.
         inference_max_seq_length (int): max_seq_length for inference. Required by MCoreEngine (>=0.12). Defaults to
         4096.
         max_batch_size (int): max batch size for inference. Defaults to 32.
@@ -205,6 +205,12 @@ class MegatronLLMDeployable(ITritonDeployable):
                     stop_words=stop_words,
                 )
 
+                if log_probs:
+                    dynamic_engine = getattr(self.mcore_engine, "engine", None)
+                    if dynamic_engine is not None:
+                        dynamic_engine.materialize_only_last_token_logits = False
+                        dynamic_engine.context.config.materialize_only_last_token_logits = False
+
                 self.generate(prompts, inference_params)
             else:
                 return
@@ -394,6 +400,22 @@ class MegatronLLMDeployable(ITritonDeployable):
         if apply_chat_template:
             prompts = [self.apply_chat_template(prompt) for prompt in prompts]
 
+        if torch.distributed.is_initialized():
+            if torch.distributed.get_world_size() > 1:
+                torch.distributed.broadcast(torch.tensor([0], dtype=torch.long, device="cuda"), src=0)
+                broadcast_list(prompts, src=0)
+                broadcast_list(
+                    data=[
+                        temperature,
+                        top_k,
+                        top_p,
+                        num_tokens_to_generate,
+                        log_probs,
+                        stop_words,
+                    ],
+                    src=0,
+                )
+
         # cast top_k,top_p to native int, float since typecheck assert statements added in MCore0.13 error otherwise
         # skip_prompt_log_probs=False (default) includes prompt tokens in top-N logprobs when top_logprobs>0.
         inference_params = SamplingParams(
@@ -407,6 +429,17 @@ class MegatronLLMDeployable(ITritonDeployable):
             stop_words=stop_words,
         )
 
+        # Mcore's dynamic inference engine defaults materialize_only_last_token_logits=True for
+        # performance, but prompt log probs require all token logits to be materialized
+        # (prompt log probs are required for logprob eval benchmarks).
+        # Toggle it on both the engine and the context config (controls the
+        # model forward pass and log prob calculations).
+        dynamic_engine = getattr(self.mcore_engine, "engine", None)
+        needs_all_logits = log_probs or bool(top_logprobs)
+        if dynamic_engine is not None and needs_all_logits:
+            dynamic_engine.materialize_only_last_token_logits = False
+            dynamic_engine.context.config.materialize_only_last_token_logits = False
+
         results = self.generate(prompts, inference_params)
         # Handle DynamicInferenceRequestRecord objects by merging them into a single request
         results = [
@@ -422,7 +455,7 @@ class MegatronLLMDeployable(ITritonDeployable):
         if log_probs:
             output_log_probs = []
             for r in results:
-                # Convert to torch tensor and then move to cpu as generated_log_probs is a list and cant be moved
+                # Convert to torch tensor and then move to cpu as generated_log_probs is a list and cannot be moved
                 # to cpu otherwise
                 if echo:
                     lp = torch.tensor(r.prompt_log_probs + r.generated_log_probs).cpu().detach().numpy()
@@ -448,7 +481,7 @@ class MegatronLLMDeployable(ITritonDeployable):
         if top_logprobs:
             output_top_n_log_probs = []
             for r in results:
-                # Convert to torch tensor and then move to cpu as generated_log_probs is a list and cant be moved
+                # Convert to torch tensor and then move to cpu as generated_log_probs is a list and cannot be moved
                 # to cpu otherwise.
                 # top_logprobs for input tokens is supported with MCore 0.13 and above.
                 if echo:

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_deploy/llm/megatronllm_deployable_ray.py

@@ -104,6 +104,8 @@ class ModelWorker:
                 legacy_model_format=legacy_model_format,
                 **model_config_kwargs,
             )
+            if rank != 0:
+                self.model.generate_other_ranks()
         except Exception as e:
             LOGGER.error(f"Replica {replica_id} - Failed to initialize model for rank {rank}: {str(e)}")
             raise

nemo_export_deploy-0.8.0/nemo_deploy/multimodal/image_url_validator.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+# Copyright (c) 2026, 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 base64
+import ipaddress
+import socket
+from io import BytesIO
+from urllib.parse import urljoin, urlparse
+
+import urllib3
+
+# Ranges that must never be reachable via a request-controlled image URL.
+# 169.254.0.0/16 is the cloud IMDS range (AWS/GCP/Azure 169.254.169.254) —
+# the primary SSRF target in cloud deployments.
+_BLOCKED_NETWORKS = [
+    ipaddress.ip_network("127.0.0.0/8"),  # loopback — server's own local services
+    ipaddress.ip_network("10.0.0.0/8"),  # RFC 1918 private
+    ipaddress.ip_network("172.16.0.0/12"),  # RFC 1918 private
+    ipaddress.ip_network("192.168.0.0/16"),  # RFC 1918 private
+    ipaddress.ip_network("169.254.0.0/16"),  # link-local / cloud IMDS
+    ipaddress.ip_network("::1/128"),  # IPv6 loopback
+    ipaddress.ip_network("fc00::/7"),  # IPv6 unique-local
+    ipaddress.ip_network("fe80::/10"),  # IPv6 link-local (IPv6 cloud IMDS equivalent)
+]
+
+
+def _is_blocked_ip(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
+    """Return True if ip must not be reachable via an image URL fetch."""
+    if any(ip in net for net in _BLOCKED_NETWORKS):
+        return True
+    return not ip.is_global or ip.is_multicast or ip.is_reserved or ip.is_unspecified
+
+
+def _resolve_and_validate(hostname: str) -> list[tuple[int, str]]:
+    """Resolve hostname to every A/AAAA address and reject if any is blocked.
+
+    Returns a list of (address_family, ip_string) tuples for all resolved
+    addresses. Validating every returned address (rather than just the
+    first) guards against DNS answers that mix a public address with a
+    private/link-local one (DNS rebinding / multi-answer SSRF).
+    """
+    try:
+        infos = socket.getaddrinfo(hostname, None)
+    except socket.gaierror as exc:
+        raise ValueError(f"Cannot resolve image URL hostname '{hostname}': {exc}") from exc
+    resolved: list[tuple[int, str]] = []
+    seen: set[str] = set()
+    for family, _, _, _, sockaddr in infos:
+        ip_str = sockaddr[0]
+        if ip_str in seen:
+            continue
+        seen.add(ip_str)
+        ip = ipaddress.ip_address(ip_str)
+        if _is_blocked_ip(ip):
+            raise ValueError(
+                f"Image URL hostname '{hostname}' resolves to a blocked address ({ip}). "
+                "Private, loopback, link-local, and other non-global addresses are not allowed."
+            )
+        resolved.append((family, ip_str))
+    if not resolved:
+        raise ValueError(f"Cannot resolve image URL hostname '{hostname}': no addresses returned.")
+    return resolved
+
+
+def validate_image_url(url: str) -> None:
+    """Raise ValueError if url is not a safe http/https URL.
+
+    Rejects file://, non-http(s) schemes, and URLs that resolve to
+    private/link-local/loopback ranges (SSRF guard). This is a point-in-time
+    check; use fetch_image_bytes_safely()/fetch_image_data_uri_safely() to
+    actually retrieve the image, since those pin validation to the
+    connection itself and guard redirects.
+    """
+    parsed = urlparse(url)
+    if parsed.scheme not in ("http", "https"):
+        raise ValueError(f"Unsupported image URL scheme '{parsed.scheme}'. Only http and https are allowed.")
+    hostname = parsed.hostname
+    if not hostname:
+        raise ValueError("Image URL has no hostname.")
+    _resolve_and_validate(hostname)
+
+
+def _pinned_request(url: str, hostname: str, ip: str, family: int, timeout: float) -> urllib3.HTTPResponse:
+    """Issue a GET request connected directly to ip, without following redirects.
+
+    Connecting to the pre-validated IP directly (instead of letting the HTTP
+    client re-resolve hostname) closes the TOCTOU window between validation
+    and connection that DNS rebinding exploits. The Host header and TLS SNI
+    still use the original hostname so virtual hosting and certificate
+    validation behave normally.
+    """
+    parsed = urlparse(url)
+    port = parsed.port or (443 if parsed.scheme == "https" else 80)
+    path = parsed.path or "/"
+    if parsed.query:
+        path = f"{path}?{parsed.query}"
+    pool_host = f"[{ip}]" if family == socket.AF_INET6 else ip
+    headers = {"Host": hostname}
+    if parsed.scheme == "https":
+        pool = urllib3.HTTPSConnectionPool(
+            pool_host,
+            port=port,
+            timeout=timeout,
+            assert_hostname=hostname,
+            server_hostname=hostname,
+            cert_reqs="CERT_REQUIRED",
+        )
+    else:
+        pool = urllib3.HTTPConnectionPool(pool_host, port=port, timeout=timeout)
+    try:
+        return pool.request("GET", path, headers=headers, redirect=False, preload_content=True)
+    finally:
+        pool.close()
+
+
+def fetch_image_bytes_safely(url: str, timeout: float = 5, max_redirects: int = 5) -> bytes:
+    """Safely fetch image bytes from url, guarding against SSRF.
+
+    Unlike validate_image_url() followed by a separate request, this
+    resolves and validates every hostname (including redirect targets)
+    immediately before connecting, connects directly to the validated IP
+    address, and never auto-follows redirects — each hop is re-validated
+    from scratch.
+    """
+    current_url = url
+    for _ in range(max_redirects + 1):
+        parsed = urlparse(current_url)
+        if parsed.scheme not in ("http", "https"):
+            raise ValueError(f"Unsupported image URL scheme '{parsed.scheme}'. Only http and https are allowed.")
+        hostname = parsed.hostname
+        if not hostname:
+            raise ValueError("Image URL has no hostname.")
+        resolved = _resolve_and_validate(hostname)
+        family, ip = resolved[0]
+        response = _pinned_request(current_url, hostname, ip, family, timeout)
+        if response.status in (301, 302, 303, 307, 308):
+            location = response.headers.get("Location")
+            if not location:
+                raise ValueError(f"Redirect response from '{current_url}' is missing a Location header.")
+            current_url = urljoin(current_url, location)
+            continue
+        if response.status >= 400:
+            raise ValueError(f"Failed to fetch image URL '{current_url}': HTTP {response.status}")
+        return response.data
+    raise ValueError(f"Exceeded max redirects ({max_redirects}) while fetching image URL '{url}'.")
+
+
+def fetch_image_data_uri_safely(url: str, timeout: float = 5, max_redirects: int = 5) -> str:
+    """Safely fetch an image and return it as a base64 data URI.
+
+    Useful for handing an image to code (e.g. third-party libraries) that
+    would otherwise perform its own, unguarded network fetch given a raw
+    URL — the network I/O happens here, under the SSRF guard, instead.
+    """
+    from PIL import Image
+
+    data = fetch_image_bytes_safely(url, timeout=timeout, max_redirects=max_redirects)
+    image_format = (Image.open(BytesIO(data)).format or "JPEG").lower()
+    encoded = base64.b64encode(data).decode("ascii")
+    return f"data:image/{image_format};base64,{encoded}"

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_deploy/multimodal/megatron_multimodal_deployable.py

@@ -173,6 +173,15 @@ class MegatronMultimodalDeployable(ITritonDeployable):
         if isinstance(self.inference_wrapped_model, QwenVLInferenceWrapper):
             from qwen_vl_utils import process_vision_info
 
+            from nemo_deploy.multimodal.image_url_validator import fetch_image_data_uri_safely
+
+            # data: URIs are inline base64 and never trigger a network request.
+            # All other values are treated as URLs: fetch them ourselves under the SSRF
+            # guard and hand qwen_vl_utils a data URI, since it would otherwise perform
+            # its own unguarded network fetch for a raw URL.
+            if not image_source.startswith("data:"):
+                image_source = fetch_image_data_uri_safely(image_source)
+
             messages = [
                 {
                     "role": "user",

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_deploy/multimodal/query_multimodal.py

@@ -17,7 +17,6 @@ from io import BytesIO
 from typing import List, Optional
 
 import numpy as np
-import requests
 
 from nemo_deploy.utils import str_list2numpy
 from nemo_export_deploy_common.import_utils import (
@@ -100,8 +99,10 @@ class NemoQueryMultimodal:
                 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")
+                from nemo_deploy.multimodal.image_url_validator import fetch_image_bytes_safely
+
+                image_bytes = fetch_image_bytes_safely(input_media, timeout=5)
+                media = Image.open(BytesIO(image_bytes)).convert("RGB")
             else:
                 media = Image.open(input_media).convert("RGB")
             return np.expand_dims(np.array(media), axis=0)

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_deploy/service/fastapi_interface_to_pytriton.py

@@ -276,7 +276,7 @@ async def chat_completions_v1(request: ChatCompletionRequest):
     prompts = request.messages
     if not isinstance(request.messages, list):
         prompts = [request.messages]
-    # Serialize the dictionary to a JSON string represnetation to be able to convert to numpy array
+    # Serialize the dictionary to a JSON string representation to be able to convert to numpy array
     # (str_list2numpy) and back to list (str_ndarray2list) as required by PyTriton. Using the dictionaries directly
     # with these methods is not possible as they expect string type.
     json_prompts = [dict_to_str(prompts)]

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_deploy/service/fastapi_interface_to_pytriton_multimodal.py

@@ -298,7 +298,7 @@ async def chat_completions_v1(request: MultimodalChatCompletionRequest):
                 # Remove image_url field
                 content.pop("image_url", None)
 
-    # Serialize the dictionary to a JSON string represnetation to be able to convert to numpy array
+    # Serialize the dictionary to a JSON string representation to be able to convert to numpy array
     # (str_list2numpy) and back to list (str_ndarray2list) as required by PyTriton. Using the dictionaries directly
     # with these methods is not possible as they expect string type.
     json_prompts = [dict_to_str(prompts)]

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_export/model_adapters/embedding/embedding_adapter.py

@@ -255,7 +255,7 @@ def get_llama_bidirectional_hf_model(
     # configure pooling
     pooling_module = Pooling(pooling_mode=pooling_mode)
 
-    # NV-Embed-v1 model has seperate embedding model and a built-in pooling module
+    # NV-Embed-v1 model has separate embedding model and a built-in pooling module
     if (
         model.__class__.__name__ == "NVEmbedModel"
         and hasattr(model, "latent_attention_model")

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_export/utils/lora_converter.py

@@ -201,7 +201,7 @@ def convert_lora_nemo_to_canonical(lora_nemo, save_path, hf_format=False, donor_
                         new_key = replace_number_add_offset(key, layer_offset)
                         lora_state_dict[tp][new_key] = value
 
-        # TODO: currently suport tp=1
+        # TODO: currently support tp=1
         lora_state_dict = lora_state_dict[0]
         if lora_config["peft"]["lora_tuning"].get("variant", "nemo") == "nemo":
             lora_config["peft"]["lora_tuning"]["variant"] = "canonical"

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_export/utils/utils.py

@@ -120,10 +120,10 @@ def get_example_inputs(
 
 
 def validate_fp8_network(network) -> None:
-    """Checks the network to ensure it's compatible with fp8 precison.
+    """Checks the network to ensure it's compatible with fp8 precision.
 
     Raises:
-        ValueError if netowrk doesn't container Q/DQ FP8 layers
+        ValueError if network doesn't container Q/DQ FP8 layers
     """
     import tensorrt as trt
 

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_export/vllm_exporter.py

@@ -112,7 +112,6 @@ class vLLMExporter(ITritonDeployable):
         quantization: str = None,
         seed: int = 0,
         gpu_memory_utilization: float = 0.9,
-        swap_space: float = 4,
         cpu_offload_gb: float = 0,
         enforce_eager: bool = False,
         task: Literal["auto", "generate", "embedding"] = "auto",
@@ -132,7 +131,6 @@ class vLLMExporter(ITritonDeployable):
             quantization (str, optional): Quantization type. Defaults to None.
             seed (int, optional): Random seed. Defaults to 0.
             gpu_memory_utilization (float, optional): Fraction of GPU memory to use. Defaults to 0.9.
-            swap_space (float, optional): Amount of swap space (in GB) to use. Defaults to 4.
             cpu_offload_gb (float, optional): Amount of CPU offload memory (in GB). Defaults to 0.
             enforce_eager (bool, optional): Whether to enforce eager execution. Defaults to False.
             task (Literal["auto", "generate", "embedding"], optional): Task type for vLLM. Defaults to "auto".
@@ -203,7 +201,6 @@ class vLLMExporter(ITritonDeployable):
                     quantization=quantization,
                     seed=seed,
                     gpu_memory_utilization=gpu_memory_utilization,
-                    swap_space=swap_space,
                     cpu_offload_gb=cpu_offload_gb,
                     enforce_eager=enforce_eager,
                     runner=task,
@@ -219,7 +216,6 @@ class vLLMExporter(ITritonDeployable):
                 quantization=quantization,
                 seed=seed,
                 gpu_memory_utilization=gpu_memory_utilization,
-                swap_space=swap_space,
                 cpu_offload_gb=cpu_offload_gb,
                 enforce_eager=enforce_eager,
                 runner=task,

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_export_deploy_common/import_utils.py

@@ -394,7 +394,7 @@ def gpu_only_import(module, *, alt=None) -> Tuple[object, bool]:
 
     return safe_import(
         module,
-        msg=f"{module} is not enabled in non GPU-enabled installations or environemnts. {GPU_INSTALL_STRING}",
+        msg=f"{module} is not enabled in non GPU-enabled installations or environments. {GPU_INSTALL_STRING}",
         alt=alt,
     )
 

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/nemo_export_deploy_common/package_info.py

@@ -14,7 +14,7 @@
 
 
 MAJOR = 0
-MINOR = 7
+MINOR = 8
 PATCH = 0
 PRE_RELEASE = ""
 

{nemo_export_deploy-0.7.0 → nemo_export_deploy-0.8.0}/pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
     "wandb==0.27.2",
     "webdataset>=0.2.86",
     "nvidia-pytriton ; platform_system != 'Darwin' ",
-    "flashinfer-python==0.6.8.post1 ; platform_system != 'Darwin'",
+    "flashinfer-python; platform_system != 'Darwin'",
     "Pillow ; platform_system != 'Darwin' and platform_machine != 'aarch64'",
     "decord ; platform_system != 'Darwin' and platform_machine != 'aarch64'",
     "pyparsing>2.0.2",
@@ -80,7 +80,7 @@ dependencies = [
 
 [project.optional-dependencies]
 inframework = []
-vllm = ["vllm~=0.24.0", "pandas", "timm"]
+vllm = ["vllm~=0.29.0", "pandas", "timm"]
 trt-onnx = [
     "tensorrt==10.16.1.11",
     "onnx==1.21.0",
@@ -106,11 +106,15 @@ nemo-run = ["nemo-run"]
 [tool.uv.sources]
 xformers = [{ index = "pytorch-cu130" }]
 torch = [{ index = "pytorch-cu130" }]
+# flashinfer-cubin isn't published to PyPI since 0.6.14; both packages are pulled from
+# flashinfer's own wheel index, same as vLLM's requirements/cuda.txt.
+flashinfer-python = [{ index = "flashinfer" }]
+flashinfer-cubin = [{ index = "flashinfer" }]
 vllm = [
     { index = "pytorch-cu130", marker = "python_version < '3.9' and platform_machine == 'x86_64'" },
     { index = "pypi", marker = "platform_machine == 'aarch64'" },
 ]
-megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "97e8dba4060ae74a110b9da6255eb0a66b016b2d" }
+megatron-bridge = { git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge.git", rev = "02fcce5a103ca0c88a945319075348729e0b666e" }
 # Use the released nvidia-resiliency-ext (>=0.6.0) rather than a git dev build:
 # megatron-core's runtime is_nvrx_min_version() asserts the installed package
 # version is >=0.6.0, and a git build reports 0.6.0.devN (< 0.6.0), which fails
@@ -150,7 +154,6 @@ override-dependencies = [
     "transformer-engine[pytorch]>=2.16.0a0,<2.19.0",
     "open-clip-torch>=3.2.0",
     "megatron-energon[av-decode]>=6.0,<7.dev0",
-    "flashinfer-python==0.6.8.post1",
     "datasets>=3.3.0",
     "flash-linear-attention>=0.3.0,<0.4.dev0",
     "patchelf; sys_platform=='never'",
@@ -162,10 +165,16 @@ override-dependencies = [
     "transformers==5.8.1",
     "protobuf~=6.33.5",
     "opencv-python-headless; sys_platform == 'never'",
-    "cryptography>=43.0.0,<47",
+    "cryptography>=50.0.0,<51",
     "onnx<1.20.0",
-    "flashinfer-python==0.6.8.post1",
-    "flashinfer-cubin==0.6.8.post1",
+    # flashinfer-python and flashinfer-cubin must be kept in lockstep (flashinfer-python does a
+    # runtime version-match check); flashinfer-cubin isn't published to PyPI since 0.6.14, so both
+    # are sourced from flashinfer's own wheel index below. Pin matches vLLM's own requirements/cuda.txt.
+    "flashinfer-python==0.6.18.post1",
+    "flashinfer-cubin==0.6.18.post1",
+    "fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git@v1.1.0",  # pin 1.1.0; PyPI sdist is incomplete, install from git tag (see DSv4/GLM-5 READMEs)
+    "nvidia-cudnn-frontend==1.26.0",
+    "nvidia-cutlass-dsl[cu13]; sys_platform == 'never'",    # Rely on system site packages install
 ]
 prerelease = "allow"
 
@@ -173,6 +182,16 @@ prerelease = "allow"
 name = "nvidia-resiliency-ext"
 version = "0.6.0"
 
+[[tool.uv.dependency-metadata]]
+name = "fast-hadamard-transform"
+version = "1.1.0"
+requires-dist = ["torch", "packaging", "ninja"]
+
+# torch-memory-saver's build backend reads TMS_CUDA_MAJOR at build time. Keep this in
+# sync with docker/Dockerfile's BASE_IMAGE CUDA major version (currently cuda13.x).
+[tool.uv.extra-build-variables]
+torch-memory-saver = { TMS_CUDA_MAJOR = "13" }
+
 [[tool.uv.index]]
 name = "pypi"
 url = "https://pypi.org/simple"
@@ -183,6 +202,19 @@ name = "pytorch-cu130"
 url = "https://download.pytorch.org/whl/cu130"
 explicit = true
 
+[[tool.uv.index]]
+name = "flashinfer"
+url = "https://flashinfer.ai/whl/"
+explicit = true
+
+# `[tool.uv.sources]`/`[[tool.uv.index]]` above are only honored by uv's project interface
+# (uv sync/lock/add). The release CI builds the wheel via the pip-compatible interface
+# (`uv pip install -e .`), which does NOT read that config and only reads `[tool.uv.pip]`.
+# flashinfer-cubin/flashinfer-python aren't published on PyPI since 0.6.14, so without this
+# extra index `uv pip install -e .` can't resolve megatron-bridge's flashinfer-cubin pin.
+[tool.uv.pip]
+extra-index-url = ["https://flashinfer.ai/whl/"]
+
 [tool.pytest.ini_options]
 addopts = "--durations=15 -s -rA -x"
 testpaths = ["tests"]
@@ -191,18 +223,14 @@ python_files = "test_*.py"
 
 [tool.coverage.run]
 concurrency = ["thread", "multiprocessing"]
+relative_files = true
 omit = [
-    "/tmp/*",
-    "/workspace/tests/*",
-    "/workspace/nemo_export_deploy_common/*",
-    "/workspace/*.py",
+    "*/tests/*",
+    "*/nemo_export_deploy_common/*",
     "*_ray.py",
     "*package_info.py",
 ]
 
-[tool.coverage.paths]
-source = [".", "/workspace", "/home/runner/work/Export-Deploy/Export-Deploy"]
-
 [tool.ruff]
 line-length = 120