diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 7375dd4a2..8f942b041 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -29,7 +29,6 @@ steps: depends_on: image-build commands: - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py - - pytest -s -v tests/e2e/offline_inference/test_ovis_image.py agents: queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU plugins: @@ -93,11 +92,28 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" - - label: "Diffusion Parallelism Test" - timeout_in_minutes: 25 + - label: "Diffusion Sequence Parallelism Test" + timeout_in_minutes: 20 depends_on: image-build commands: - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py + agents: + queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + + - label: "Diffusion Tensor Parallelism Test" + timeout_in_minutes: 20 + depends_on: image-build + commands: - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py agents: queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU diff --git a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py index cefda8915..cd22b34e0 100644 --- a/tests/e2e/offline_inference/test_diffusion_cpu_offload.py +++ b/tests/e2e/offline_inference/test_diffusion_cpu_offload.py @@ -38,8 +38,10 @@ def test_cpu_offload_diffusion_model(model_name: str): guidance_scale=0.0, generator=torch.Generator("cuda").manual_seed(42), ) + peak = monitor.peak_used_mb + monitor.stop() - return monitor.peak_used_mb + return peak offload_peak_memory = inference(offload=True) no_offload_peak_memory = inference(offload=False) diff --git a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py index f4fb4371e..a9fdec4dc 100644 --- a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py +++ b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py @@ -119,7 +119,7 @@ def _run_zimage_generate( return _extract_single_image([last_output]), median_time_s, peak_memory_mb finally: - m.close() + monitor.stop() cleanup_dist_env_and_memory() @@ -160,7 +160,7 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): mean_abs_diff, max_abs_diff = _diff_metrics(tp1_img, tp2_img) mean_threshold = 3e-2 - max_threshold = 3.5e-1 + max_threshold = 5e-1 print( "Z-Image TP image diff stats (TP=1 vs TP=2): " f"mean_abs_diff={mean_abs_diff:.6e}, max_abs_diff={max_abs_diff:.6e}; " diff --git a/tests/utils.py b/tests/utils.py index 8e5593d65..ead286760 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -500,7 +500,7 @@ class GPUMemoryMonitor: pass time.sleep(self.interval) - self._thread = threading.Thread(target=monitor_loop, daemon=True) + self._thread = threading.Thread(target=monitor_loop, daemon=False) self._thread.start() def stop(self) -> None: