chore: import upstream snapshot with attribution

release/llm_tests/batch/test_batch_multi_node_vllm.py

@@ -0,0 +1,60 @@
+import pytest
+
+import ray
+from ray.data.llm import build_processor, vLLMEngineProcessorConfig
+
+
+@pytest.fixture(autouse=True)
+def cleanup_ray_resources():
+    """Automatically cleanup Ray resources between tests to prevent conflicts."""
+    yield
+    ray.shutdown()
+
+
+@pytest.mark.parametrize(
+    "tp_size,pp_size",
+    [
+        (2, 4),
+        (4, 2),
+    ],
+)
+def test_vllm_multi_node(tp_size, pp_size):
+    config = vLLMEngineProcessorConfig(
+        model_source="facebook/opt-1.3b",
+        engine_kwargs=dict(
+            enable_prefix_caching=True,
+            enable_chunked_prefill=True,
+            max_num_batched_tokens=4096,
+            pipeline_parallel_size=pp_size,
+            tensor_parallel_size=tp_size,
+            distributed_executor_backend="ray",
+        ),
+        tokenize_stage=False,
+        detokenize_stage=False,
+        concurrency=1,
+        batch_size=64,
+        chat_template_stage=False,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            prompt=f"You are a calculator. {row['id']} ** 3 = ?",
+            sampling_params=dict(
+                temperature=0.3,
+                max_tokens=20,
+                detokenize=True,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            resp=row["generated_text"],
+        ),
+    )
+
+    ds = ray.data.range(60)
+    ds = processor(ds)
+    ds = ds.materialize()
+
+    outs = ds.take_all()
+    assert len(outs) == 60
+    assert all("resp" in out for out in outs)