wehub-resource-sync
56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
import pytest
|
|
|
|
import ray
|
|
from ray.tests.conftest import _ray_start_cluster
|
|
|
|
|
|
@pytest.fixture
|
|
def llm_config_with_mock_engine(llm_config):
|
|
# Make sure engine is mocked.
|
|
if llm_config.runtime_env is None:
|
|
llm_config.runtime_env = {}
|
|
llm_config.runtime_env.setdefault("env_vars", {})[
|
|
"RAYLLM_VLLM_ENGINE_CLS"
|
|
] = "ray.llm.tests.serve.mocks.mock_vllm_engine.MockVLLMEngine"
|
|
yield llm_config
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def ray_tpu_cluster():
|
|
"""
|
|
Simulates a Ray cluster with a multi-host TPU v6e-16 slice (4x4 topology).
|
|
"""
|
|
pod_type = "v6e-16"
|
|
topology = "4x4"
|
|
|
|
with _ray_start_cluster() as cluster:
|
|
# A 4x4 v6e slice has 16 chips. We simulate 4 hosts with 4 chips each.
|
|
for i in range(4):
|
|
env_vars = {
|
|
"TPU_NAME": "test-slice",
|
|
"TPU_WORKER_ID": str(i),
|
|
"TPU_ACCELERATOR_TYPE": pod_type,
|
|
"TPU_TOPOLOGY": topology,
|
|
}
|
|
labels = {
|
|
"ray.io/tpu-slice-name": "test-slice",
|
|
"ray.io/tpu-worker-id": str(i),
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
}
|
|
resources = {"TPU": 4, "accelerator_type:TPU-V6E": 4}
|
|
|
|
# The first node is the "head" of the slice
|
|
if i == 0:
|
|
resources[f"TPU-{pod_type}-head"] = 1
|
|
|
|
cluster.add_node(
|
|
num_cpus=8,
|
|
resources=resources,
|
|
labels=labels,
|
|
env_vars=env_vars,
|
|
)
|
|
|
|
ray.init(address=cluster.address)
|
|
yield cluster
|
|
ray.shutdown()
|