[Bugfix] Fix issue 5205, add trust-remote-code (#5213)

Signed-off-by: zhumingjue <zhumingjue@huawei.com>

.buildkite/cuda/test-nightly.yml

@@ -1604,6 +1604,47 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: ":full_moon: Diffusion · Wan VAE Spatial Shard Correctness Test"
+        timeout_in_minutes: 60
+        commands:
+          - >-
+            pytest -sv tests/diffusion/distributed/test_wan_spatial_shard.py
+            -m "full_model and diffusion and H100 and distributed_cuda"
+            --run-level "full_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
   - label: ":bar_chart: Testcase Statistics"
     key: nightly-testcase-statistics
     timeout_in_minutes: 120

tests/dfx/reliability/test_reliability_voxcpm2.py

@@ -43,6 +43,9 @@ RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
         "server_params": {
             "model": "openbmb/VoxCPM2",
             "stage_config_name": "voxcpm2.yaml",
+            # VoxCPM2 uses a custom HF tokenizer; CLI default trust_remote_code=False
+            # would otherwise override deploy YAML settings during server startup.
+            "server_args": ["--trust-remote-code", "--disable-log-stats"],
         },
     },
 ]

tests/diffusion/distributed/test_wan_spatial_shard.py

@@ -6,18 +6,19 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
+from tests.helpers.mark import hardware_test
 from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
     DistributedAutoencoderKLWan,
 )
 from vllm_omni.platforms import current_omni_platform
 
-# Module-level markers apply to the CPU unit tests below. The multi-GPU
-# correctness test at the end of the file adds the nightly/distributed markers
-# and a skip guard so it only runs where >= 2 CUDA devices are available.
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+# CPU unit tests are marked core_model + cpu. The multi-GPU correctness test at
+# the end uses hardware_test (H100 x2) plus full_model / diffusion / parallel.
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_split_for_parallel_decode_pads_uneven_height():
     x = torch.arange(1 * 1 * 1 * 5 * 2, dtype=torch.float32).reshape(1, 1, 1, 5, 2)
 
@@ -34,6 +35,8 @@ def test_split_for_parallel_decode_pads_uneven_height():
     assert torch.equal(local[..., 1, :], torch.zeros_like(local[..., 1, :]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_split_for_parallel_decode_pads_uneven_width():
     x = torch.arange(1 * 1 * 1 * 2 * 5, dtype=torch.float32).reshape(1, 1, 1, 2, 5)
 
@@ -51,6 +54,8 @@ def test_split_for_parallel_decode_pads_uneven_width():
     assert torch.equal(local[..., :, 1], torch.zeros_like(local[..., :, 1]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_split_for_parallel_decode_rejects_invalid_split_dim():
     x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32)
 
@@ -64,6 +69,8 @@ def test_split_for_parallel_decode_rejects_invalid_split_dim():
         )
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_split_for_parallel_decode_rejects_zero_world_size():
     x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32)
 
@@ -76,6 +83,8 @@ def test_split_for_parallel_decode_rejects_zero_world_size():
         )
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_split_for_parallel_decode_rejects_rank_out_of_range():
     x = torch.zeros((1, 1, 1, 4, 4), dtype=torch.float32)
 
@@ -88,6 +97,8 @@ def test_split_for_parallel_decode_rejects_rank_out_of_range():
         )
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_gather_and_trim_height(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3))
 
@@ -104,6 +115,8 @@ def test_gather_and_trim_height(monkeypatch: pytest.MonkeyPatch):
     assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_gather_and_trim_width(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3))
 
@@ -120,6 +133,8 @@ def test_gather_and_trim_width(monkeypatch: pytest.MonkeyPatch):
     assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_gather_and_trim_rank0_only_assembles_on_rank0(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 3))
 
@@ -136,6 +151,8 @@ def test_gather_and_trim_rank0_only_assembles_on_rank0(monkeypatch: pytest.Monke
     assert torch.equal(out.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0, 2.0]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_gather_and_trim_rank0_only_returns_empty_on_non_zero_rank(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (1, 3))
 
@@ -156,6 +173,8 @@ def test_gather_and_trim_rank0_only_returns_empty_on_non_zero_rank(monkeypatch:
     assert out.numel() == 0
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_reshard_from_trimmed_height_pads_invalid_rows(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
 
@@ -183,6 +202,8 @@ def test_reshard_from_trimmed_height_pads_invalid_rows(monkeypatch: pytest.Monke
     assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (2, 3))
 
@@ -210,6 +231,8 @@ def test_reshard_from_trimmed_width_pads_invalid_columns(monkeypatch: pytest.Mon
     assert torch.equal(out.flatten(), torch.tensor([4.0, 0.0]))
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (0, 1))
 
@@ -225,6 +248,8 @@ def test_halo_exchange_single_rank_noop(monkeypatch: pytest.MonkeyPatch):
     assert recv_bottom is None
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.MonkeyPatch):
     x = torch.ones((1, 1, 2, 2))
 
@@ -239,6 +264,8 @@ def test_dist_zero_pad_only_applies_global_height_edges(monkeypatch: pytest.Monk
     assert last.shape == (1, 1, 3, 3)
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.MonkeyPatch):
     x = torch.ones((1, 1, 2, 2))
 
@@ -253,6 +280,8 @@ def test_dist_zero_pad_only_applies_global_width_edges(monkeypatch: pytest.Monke
     assert last.shape == (1, 1, 2, 3)
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(
         "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
@@ -270,6 +299,8 @@ def test_spatial_shard_height_gate_falls_back_for_partial_group(monkeypatch: pyt
     assert vae._spatial_shard_decode_enabled(z) is False
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch):
     monkeypatch.setattr(
         "vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan.dist.get_world_size",
@@ -286,6 +317,8 @@ def test_spatial_shard_width_gate_selects_width(monkeypatch: pytest.MonkeyPatch)
     assert vae._spatial_shard_decode_enabled(z) is True
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
 def test_tile_mode_disables_spatial_shard_decode():
     vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
     vae.distributed_executor = SimpleNamespace(group=object(), parallel_size=2, parallel_mode="tile")
@@ -298,7 +331,7 @@ def test_tile_mode_disables_spatial_shard_decode():
 
 
 # =============================================================================
-# Multi-GPU numerical-correctness test (nightly / full_model, distributed CUDA)
+# Multi-GPU numerical-correctness test (nightly Diffusion Test group, H100 x2)
 #
 # Spawns a small process group and verifies that spatial_shard_height/spatial_shard_width decode match
 # a single-process (non-distributed) reference decode of the same latent within
@@ -383,11 +416,7 @@ def _spatial_shard_decode_worker(rank: int, split_dim: str, return_dict, master_
 @pytest.mark.full_model
 @pytest.mark.diffusion
 @pytest.mark.parallel
-@pytest.mark.distributed_cuda
-@pytest.mark.skipif(
-    current_omni_platform.get_device_count() < _SPATIAL_SHARD_WORLD_SIZE,
-    reason="Requires >= 2 accelerator devices",
-)
+@hardware_test(res={"cuda": "H100"}, num_cards=_SPATIAL_SHARD_WORLD_SIZE)
 @pytest.mark.parametrize("split_dim", ["height", "width"])
 def test_spatial_shard_decode_matches_reference(split_dim: str):
     manager = mp.get_context("spawn").Manager()

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -368,6 +368,7 @@ def _run_online(stage_configs_path: str, output_path: Path) -> tuple[Image.Image
         "300",
         "--init-timeout",
         "900",
+        "--trust-remote-code",
     ]
     try:
         with OmniServer(MODEL_PATH, server_args, use_omni=True) as omni_server:

tests/e2e/online_serving/test_dynin_omni_expansion.py

@@ -26,7 +26,14 @@ T2I_PROMPT = "A high quality detailed living room interior photo."
 T2S_PROMPT = "Please read this sentence naturally: Hello from online serving."
 I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors."
 
-TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, stage_init_timeout=600)]
+TEST_PARAMS = [
+    OmniServerParams(
+        model=MODEL,
+        stage_config_path=STAGE_CONFIG,
+        stage_init_timeout=600,
+        server_args=["--trust-remote-code"],
+    )
+]
 _STAGE_COUNT = 3
 _I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False}