fix(pytorch_xdit): fail fast on missing image and on incomplete WAN caches

cvs/lib/unittests/test_utils_lib.py

@@ -1,4 +1,6 @@
 # cvs/lib/unittests/test_utils_lib.py
+import os
+import shlex
 import unittest
 from unittest.mock import patch
 
@@ -38,6 +40,17 @@ def test_get_model_from_rocm_smi_output_defaults_to_mi300x_when_unrecognized(sel
         smi_output = 'Device Name:        AMD Radeon Graphics\nDevice ID:          0x1234\n'
         self.assertEqual(utils_lib.get_model_from_rocm_smi_output(smi_output), 'mi300x')
 
+    def test_wan_hf_snapshot_offline_check_commands_paths_quoted(self):
+        snap_root = '/data/my hf cache/snapshots/abc123'
+        cmds = utils_lib.wan_hf_snapshot_offline_check_commands(snap_root)
+        self.assertIn('configuration.json', cmds)
+        self.assertIn('low_noise diffusion shards (6 x >500MiB)', cmds)
+        quoted_cfg = shlex.quote(os.path.join(snap_root, 'configuration.json'))
+        self.assertIn(quoted_cfg, cmds['configuration.json'])
+        for label, cmd in cmds.items():
+            self.assertIn('OK', cmd, msg=label)
+            self.assertIn('MISSING', cmd, msg=label)
+
 
 class TestResolveTestConfigPlaceholdersAorta(unittest.TestCase):
     """Aorta benchmark YAML uses the same resolver as other CVS test suites (see tests/benchmark/test_aorta.py)."""

cvs/lib/utils_lib.py

@@ -9,6 +9,7 @@
 import os
 import sys
 import json
+import shlex
 
 import pytest
 from cvs.lib import globals
@@ -448,6 +449,50 @@ def cluster_target_output_label(cluster_node_key: str) -> str:
     return str(cluster_node_key).strip().replace("/", "_")
 
 
+def wan_hf_snapshot_offline_check_commands(snapshot_dir_host: str) -> dict:
+    """
+    Remote-shell checks for a Wan2.2 I2V-A14B-style Hugging Face snapshot (host path).
+
+    Each command prints OK on success and MISSING on failure. This catches the common
+    false positive where the snapshot directory exists while `hf download` is still
+    running or LFS pointer files were not fully materialized.
+
+    Intended for snapshots of `Wan-AI/Wan2.2-I2V-A14B` (and forks with the same tree).
+    Callers should only run these when `model_repo` is known to use this layout.
+    """
+    root = snapshot_dir_host.rstrip('/')
+    checks = {}
+
+    def ok_file(rel: str) -> str:
+        return f"test -f {shlex.quote(os.path.join(root, *rel.split('/')))} && echo OK || echo MISSING"
+
+    checks['configuration.json'] = ok_file('configuration.json')
+    checks['low_noise_model/config.json'] = ok_file('low_noise_model/config.json')
+    checks['high_noise_model/config.json'] = ok_file('high_noise_model/config.json')
+
+    low_dir = shlex.quote(os.path.join(root, 'low_noise_model'))
+    high_dir = shlex.quote(os.path.join(root, 'high_noise_model'))
+    checks['low_noise diffusion shards (6 x >500MiB)'] = (
+        f's=$(find -L {low_dir} -maxdepth 1 -type f -name "diffusion_pytorch_model-*.safetensors" '
+        f'-size +500M 2>/dev/null | wc -l); test "$s" -eq 6 && echo OK || echo MISSING'
+    )
+    checks['high_noise diffusion shards (6 x >500MiB)'] = (
+        f's=$(find -L {high_dir} -maxdepth 1 -type f -name "diffusion_pytorch_model-*.safetensors" '
+        f'-size +500M 2>/dev/null | wc -l); test "$s" -eq 6 && echo OK || echo MISSING'
+    )
+
+    vae = shlex.quote(os.path.join(root, 'Wan2.1_VAE.pth'))
+    t5 = shlex.quote(os.path.join(root, 'models_t5_umt5-xxl-enc-bf16.pth'))
+    checks['Wan2.1_VAE.pth (>100MiB, not pointer-only)'] = (
+        f"test -f {vae} && test -n \"$(find -L {vae} -size +100M 2>/dev/null)\" && echo OK || echo MISSING"
+    )
+    checks['models_t5_umt5-xxl-enc-bf16.pth (>1GiB, not pointer-only)'] = (
+        f"test -f {t5} && test -n \"$(find -L {t5} -size +1G 2>/dev/null)\" && echo OK || echo MISSING"
+    )
+
+    return checks
+
+
 def collect_system_metadata(phdl, cluster_dict, config_dict, test_command=None, env_vars=None):
     """
     Collect comprehensive system metadata from compute nodes for test reporting.

cvs/tests/inference/pytorch_xdit/pytorch_xdit_flux1_dev_single.py

@@ -280,8 +280,8 @@ def s_phdl(cluster_dict):
     # Single-node mode: execute locally ONLY when the target actually refers to this machine.
     #
     # Rationale: users often specify a remote node IP/hostname in cluster.json even for a
-    # single-node run. Always forcing local execution will run benchmarks on the login node
-    # (no GPUs/ROCm) and fail in confusing ways.
+    # single-node run. Without this check, that target would run on this host instead of the
+    # target node's GPUs/ROCm, and fail in confusing ways.
     if len(node_list) == 1:
         target = node_list[0]
         if _is_local_target(target):
@@ -516,8 +516,8 @@ def test_run_flux1_benchmark(s_phdl, inference_dict, benchmark_params_dict, hf_t
     """
     globals.error_list = []
 
-    # Preflight: ensure all nodes have GPU-capable hardware. Running on a login node (no /dev/kfd)
-    # will cause ROCm + container init to fail and produce no timing.json.
+    # Preflight: ensure all nodes expose /dev/kfd (ROCm). Missing device nodes usually means
+    # the target is not suitable for this GPU container workload.
     log.info(f"Checking /dev/kfd on {len(s_phdl.host_list)} node(s)")
     kfd_check = s_phdl.exec("test -e /dev/kfd && echo KFD_OK || echo KFD_MISSING", print_console=False)
     missing_kfd_nodes = []
@@ -531,12 +531,22 @@ def test_run_flux1_benchmark(s_phdl, inference_dict, benchmark_params_dict, hf_t
     if missing_kfd_nodes:
         fail_test(
             f"ROCm device node /dev/kfd not found on {len(missing_kfd_nodes)} node(s): {', '.join(missing_kfd_nodes)}. "
-            f"This test must be run on GPU compute nodes (e.g., via an interactive SLURM allocation)."
+            f"This test requires ROCm GPU nodes with /dev/kfd on each target."
         )
         update_test_result()
         return
 
     container_image = inference_dict['container_image']
+    missing_img_nodes = docker_lib.nodes_missing_docker_image(s_phdl, container_image)
+    if missing_img_nodes:
+        fail_test(
+            f"Container image not found locally on {len(missing_img_nodes)} node(s): {', '.join(missing_img_nodes)}. "
+            f"Configured image: {container_image}. Pull it on each target node before running this benchmark "
+            f"(for example: docker pull {container_image})."
+        )
+        update_test_result()
+        return
+
     container_name = inference_dict['container_name']
     hf_home = inference_dict['hf_home']
     output_base_dir = inference_dict['output_base_dir']

cvs/tests/inference/pytorch_xdit/pytorch_xdit_wan22_14b_single.py

@@ -23,6 +23,7 @@
     get_model_from_rocm_smi_output,
     resolve_cluster_config_placeholders,
     resolve_test_config_placeholders,
+    wan_hf_snapshot_offline_check_commands,
 )
 from cvs.lib import docker_lib
 from cvs.lib import globals
@@ -257,8 +258,8 @@ def s_phdl(cluster_dict):
     # Single-node mode: execute locally ONLY when the target actually refers to this machine.
     #
     # Rationale: users often specify a remote node IP/hostname in cluster.json even for a
-    # single-node run. Always forcing local execution will run benchmarks on the login node
-    # (no GPUs/ROCm) and silently "pass" until parsing fails.
+    # single-node run. Without this check, that target would run on this host instead of the
+    # target node's GPUs/ROCm, and fail in confusing ways.
     if len(node_list) == 1:
         target = node_list[0]
         if _is_local_target(target):
@@ -361,6 +362,20 @@ def test_verify_hf_cache_or_download(s_phdl, inference_dict, hf_token):
             update_test_result()
             return
 
+        # Align with HF snapshot checks for Wan2.2 I2V-A14B when staging via absolute path.
+        if "Wan2.2-I2V-A14B" in model_repo or "Wan2.2-I2V-A14B" in host_model_path:
+            for label, cmd in wan_hf_snapshot_offline_check_commands(host_model_path).items():
+                res = s_phdl.exec(cmd, print_console=False)
+                bad = [n for n, out in (res or {}).items() if "OK" not in (out or "")]
+                if bad:
+                    fail_test(
+                        "WAN local model directory looks incomplete (Wan2.2-I2V-A14B layout). "
+                        f"Check '{label}' failed on {len(bad)} node(s): {', '.join(bad)}. "
+                        f"Model path: {host_model_path}."
+                    )
+                    update_test_result()
+                    return
+
         inference_dict["_resolved_model_mount_host"] = host_model_path
         inference_dict["_resolved_ckpt_dir_container"] = "/model"
         log.info(f"Using local model path: {host_model_path} (mounted to /model in container) on all nodes")
@@ -393,6 +408,20 @@ def test_verify_hf_cache_or_download(s_phdl, inference_dict, hf_token):
     inference_dict["_resolved_ckpt_dir_container"] = f"/hf_home/hub/models--{model_path_safe}/snapshots/{model_rev}"
     log.info(f"Using pre-cached snapshot: {inference_dict['_resolved_ckpt_dir_container']} on all nodes")
 
+    # Stronger offline checks for Wan2.2 I2V-A14B HF layout (avoids passing while download is partial).
+    if "Wan2.2-I2V-A14B" in model_repo:
+        for label, cmd in wan_hf_snapshot_offline_check_commands(snapshot_dir_host).items():
+            res = s_phdl.exec(cmd, print_console=False)
+            bad = [n for n, out in (res or {}).items() if "OK" not in (out or "")]
+            if bad:
+                fail_test(
+                    "WAN Hugging Face snapshot looks incomplete (or not a Wan2.2-I2V-A14B-style tree). "
+                    f"Check '{label}' failed on {len(bad)} node(s): {', '.join(bad)}. "
+                    f"Snapshot path: {snapshot_dir_host}. Wait for downloads to finish or fix hf_home / model_rev."
+                )
+                update_test_result()
+                return
+
     update_test_result()
 
 
@@ -406,7 +435,36 @@ def test_run_wan22_benchmark(s_phdl, inference_dict, benchmark_params_dict, hf_t
     """
     globals.error_list = []
 
+    # Preflight: same ROCm device check as FLUX — avoids messy container failures on non-GPU nodes.
+    log.info(f"Checking /dev/kfd on {len(s_phdl.host_list)} node(s)")
+    kfd_check = s_phdl.exec("test -e /dev/kfd && echo KFD_OK || echo KFD_MISSING", print_console=False)
+    missing_kfd_nodes = []
+    for node, output in kfd_check.items():
+        if "KFD_OK" not in (output or ""):
+            missing_kfd_nodes.append(node)
+            log.error(f"ROCm device node /dev/kfd not found on {node}")
+        else:
+            log.info(f"/dev/kfd found on {node}")
+
+    if missing_kfd_nodes:
+        fail_test(
+            f"ROCm device node /dev/kfd not found on {len(missing_kfd_nodes)} node(s): {', '.join(missing_kfd_nodes)}. "
+            f"This test requires ROCm GPU nodes with /dev/kfd on each target."
+        )
+        update_test_result()
+        return
+
     container_image = inference_dict['container_image']
+    missing_img_nodes = docker_lib.nodes_missing_docker_image(s_phdl, container_image)
+    if missing_img_nodes:
+        fail_test(
+            f"Container image not found locally on {len(missing_img_nodes)} node(s): {', '.join(missing_img_nodes)}. "
+            f"Configured image: {container_image}. Pull it on each target node before running this benchmark "
+            f"(for example: docker pull {container_image})."
+        )
+        update_test_result()
+        return
+
     container_name = inference_dict['container_name']
     hf_home = inference_dict['hf_home']
     output_base_dir = inference_dict['output_base_dir']