fix(specdec): harden multimodal media delivery

Signed-off-by: Slawomir Kierat <skierat@nvidia.com>

Author: skierat

examples/speculative_decoding/distributed_generate/launch_multimodal.sh

@@ -63,7 +63,9 @@ fi
 
 mkdir -p "$OUTPUT_PATH"
 
-CONTAINER_IMAGE=${CONTAINER_IMAGE:-/lustre/fs1/portfolios/coreai/users/skierat/lmsysorg+sglang+v0.5.3-cu129.sqsh}
+# Set CONTAINER_IMAGE to a local .sqsh image to avoid pulling from the registry.
+DEFAULT_CONTAINER_IMAGE="lmsysorg/sglang:v0.5.3-cu129"
+CONTAINER_IMAGE="${CONTAINER_IMAGE:-$DEFAULT_CONTAINER_IMAGE}"
 
 counter=$START_SHARD
 worker_pids=()

examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py

@@ -31,6 +31,7 @@
 import tqdm
 
 QWEN_IMAGE_TOKEN = "<|vision_start|><|image_pad|><|vision_end|>"
+_UNRESOLVED_MEDIA_PATHS: set[str] = set()
 
 
 def _load_json_or_jsonl(path: str) -> list[dict[str, Any]]:
@@ -101,6 +102,8 @@ def _resolve_media_path(
 ) -> str | None:
     if not path:
         return None
+    if path.startswith(("http://", "https://", "data:")):
+        return path
     candidate = Path(path)
     if candidate.is_absolute() and candidate.exists():
         return str(candidate)
@@ -121,7 +124,10 @@ def _resolve_media_path(
                 return str(rooted)
     if candidate.exists():
         return str(candidate)
-    return path
+    if path not in _UNRESOLVED_MEDIA_PATHS:
+        print(f"WARNING: could not resolve media path: {path}")
+        _UNRESOLVED_MEDIA_PATHS.add(path)
+    return None
 
 
 def _as_openai_media_value(
@@ -132,9 +138,16 @@ def _as_openai_media_value(
     if media_url_base:
         candidate = Path(path)
         if candidate.is_absolute():
-            return f"{media_url_base.rstrip('/')}{quote(str(candidate), safe='/')}"
-        if not candidate.is_absolute():
-            return f"{media_url_base.rstrip('/')}/{quote(path, safe='/')}"
+            if media_root:
+                try:
+                    path = str(candidate.relative_to(media_root))
+                except ValueError:
+                    # The local HTTP server deliberately exposes only media_root.
+                    # Leave paths outside it local for SGLang to resolve directly.
+                    return path
+            else:
+                return f"{media_url_base.rstrip('/')}{quote(path, safe='/')}"
+        return f"{media_url_base.rstrip('/')}/{quote(path, safe='/')}"
     # Do not convert local paths to file://. This SGLang build falls through to
     # the base64 loader for file:// videos and raises "Incorrect padding".
     return path

examples/speculative_decoding/distributed_generate/worker_multimodal.sh

@@ -91,19 +91,19 @@ trap 'exit 130' INT
 trap 'exit 143' TERM
 
 if [ "$API_MODE" = "openai" ]; then
-    python3 -m http.server "$MEDIA_HTTP_PORT" --bind 127.0.0.1 --directory / \
+    python3 -m http.server "$MEDIA_HTTP_PORT" --bind 127.0.0.1 --directory /media_data \
         >/tmp/multimodal_media_http_${MEDIA_HTTP_PORT}.log 2>&1 &
     MEDIA_HTTP_PID=$!
 
     echo "Waiting for media HTTP server at $MEDIA_URL_BASE..."
     for _ in $(seq 1 30); do
-        if curl -fsS "$MEDIA_URL_BASE/media_data/" >/dev/null 2>&1; then
+        if curl -fsS "$MEDIA_URL_BASE/" >/dev/null 2>&1; then
             echo "Media HTTP server is up."
             break
         fi
         sleep 1
     done
-    if ! curl -fsS "$MEDIA_URL_BASE/media_data/" >/dev/null 2>&1; then
+    if ! curl -fsS "$MEDIA_URL_BASE/" >/dev/null 2>&1; then
         echo "ERROR: media HTTP server did not start. See /tmp/multimodal_media_http_${MEDIA_HTTP_PORT}.log" >&2
         exit 1
     fi

tests/examples/speculative_decoding/test_merge_dflash_datasets.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Tests for the DFlash JSONL merger recipe."""
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).parents[3]
+MERGER = REPO_ROOT / "examples/speculative_decoding/recipes/merge_dflash_datasets.py"
+
+
+def write_record(path: Path, record_id: str, prompt: str, answer: str) -> None:
+    """Write a one-turn DFlash-training conversation."""
+
+    path.write_text(
+        json.dumps(
+            {
+                "id": record_id,
+                "messages": [
+                    {"role": "user", "content": prompt},
+                    {"role": "assistant", "content": answer},
+                ],
+            }
+        )
+        + "\n",
+        encoding="utf-8",
+    )
+
+
+def test_merge_keeps_synthetic_variants_together(tmp_path):
+    """Single and parallel merges deduplicate matching variants from one synthetic shard."""
+
+    synthetic_output = tmp_path / "synthetic-output"
+    synthetic_output.mkdir()
+    write_record(
+        synthetic_output / "output-00000-00000-temp-0.0.jsonl",
+        "duplicate-first",
+        "Describe the image.",
+        "The image shows a blue square.",
+    )
+    write_record(
+        synthetic_output / "output-00000-00000-temp-0.1.jsonl",
+        "duplicate-second",
+        "Describe the image.",
+        "The image shows a blue square.",
+    )
+    write_record(
+        synthetic_output / "output-00001-00001-temp-0.0.jsonl",
+        "synthetic-unique",
+        "Describe the video.",
+        "The video shows a red circle.",
+    )
+    curated_text = tmp_path / "curated-text.jsonl"
+    write_record(curated_text, "curated-unique", "What is two plus two?", "Four.")
+    expected_ids = {"duplicate-first", "synthetic-unique", "curated-unique"}
+    for jobs in (1, 2):
+        output = tmp_path / f"merged-{jobs}.jsonl"
+        subprocess.run(
+            [
+                sys.executable,
+                str(MERGER),
+                "--source",
+                f"synthetic={synthetic_output}",
+                "--source",
+                f"curated={curated_text}",
+                "--output",
+                str(output),
+                "--jobs",
+                str(jobs),
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+        records = [json.loads(line) for line in output.read_text(encoding="utf-8").splitlines()]
+        assert {record["id"] for record in records} == expected_ids
+        assert not list(tmp_path.glob(f".{output.name}.parallel-*"))

tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Tests for the multimodal SGLang generation client."""
+
+import importlib.util
+from pathlib import Path
+
+_SCRIPT_PATH = (
+    Path(__file__).parents[3]
+    / "examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py"
+)
+_spec = importlib.util.spec_from_file_location("server_generate_vlm_sglang", _SCRIPT_PATH)
+assert _spec is not None and _spec.loader is not None
+server_generate_vlm_sglang = importlib.util.module_from_spec(_spec)
+_spec.loader.exec_module(server_generate_vlm_sglang)
+
+
+def test_resolve_media_path_supports_local_and_remote_media(tmp_path):
+    """Existing local media and supported remote media remain usable."""
+
+    local_media = tmp_path / "image.jpg"
+    local_media.touch()
+    remote_media = "https://example.com/image.jpg"
+
+    assert server_generate_vlm_sglang._resolve_media_path("image.jpg", str(tmp_path), None) == str(
+        local_media
+    )
+    assert (
+        server_generate_vlm_sglang._resolve_media_path(remote_media, str(tmp_path), None)
+        == remote_media
+    )
+
+
+def test_resolve_media_path_returns_none_and_warns_once_for_missing_media(monkeypatch, capsys):
+    """Missing local media can fall back to video or be skipped by the caller."""
+
+    monkeypatch.setattr(server_generate_vlm_sglang, "_UNRESOLVED_MEDIA_PATHS", set())
+
+    assert server_generate_vlm_sglang._resolve_media_path("missing.mp4", None, None) is None
+    assert server_generate_vlm_sglang._resolve_media_path("missing.mp4", None, None) is None
+
+    assert capsys.readouterr().out == "WARNING: could not resolve media path: missing.mp4\n"
+
+
+def test_openai_media_value_is_relative_to_the_media_root(tmp_path):
+    """The local HTTP server exposes media files, not the container root."""
+
+    media_root = tmp_path / "media"
+    media_path = media_root / "videos" / "clip.mp4"
+    media_path.parent.mkdir(parents=True)
+    media_path.touch()
+    input_path = tmp_path / "input" / "private.json"
+    input_path.parent.mkdir()
+    input_path.touch()
+
+    assert (
+        server_generate_vlm_sglang._as_openai_media_value(
+            str(media_path), "http://127.0.0.1:18080", str(media_root), None
+        )
+        == "http://127.0.0.1:18080/videos/clip.mp4"
+    )
+    assert server_generate_vlm_sglang._as_openai_media_value(
+        str(input_path), "http://127.0.0.1:18080", str(media_root), None
+    ) == str(input_path)