wehub-resource-sync
992 lines
33 KiB
Python
992 lines
33 KiB
Python
import os
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.tests.conftest import _ray_start_cluster
|
|
from ray.train import RunConfig, ScalingConfig, UserCallback
|
|
from ray.train.v2._internal.constants import (
|
|
HEALTH_CHECK_INTERVAL_S_ENV_VAR,
|
|
is_v2_enabled,
|
|
)
|
|
from ray.train.v2.jax import JaxTrainer
|
|
|
|
assert is_v2_enabled()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_tpu_single_host():
|
|
"""Start a mock single-host TPU Ray cluster with 2x4 v6e (8 chips per host)."""
|
|
with _ray_start_cluster() as cluster:
|
|
# Simulate one node with 8 TPU chips.
|
|
cluster.add_node(
|
|
num_cpus=4,
|
|
resources={"TPU": 8, "accelerator_type:TPU-V6E": 1},
|
|
env_vars={"TPU_ACCELERATOR_TYPE": "v6e-8"},
|
|
)
|
|
|
|
ray.init(address=cluster.address)
|
|
|
|
yield cluster
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_tpu_multi_host():
|
|
"""
|
|
Simulates a Ray cluster with two multi-host TPU v4-16 slices.
|
|
"""
|
|
|
|
pod_type = "v4-16"
|
|
topology = "2x2x2"
|
|
|
|
with _ray_start_cluster() as cluster:
|
|
# First TPU slice - v4-16 with 2 hosts and 4 chips/host.
|
|
slice_a_head_env = {
|
|
"TPU_NAME": "slice-A",
|
|
"TPU_WORKER_ID": "0",
|
|
"TPU_ACCELERATOR_TYPE": pod_type,
|
|
"TPU_TOPOLOGY": topology,
|
|
}
|
|
slice_a_head_labels = {
|
|
"ray.io/tpu-slice-name": "slice-A",
|
|
"ray.io/tpu-worker-id": "0",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
}
|
|
slice_a_worker_env = {
|
|
"TPU_NAME": "slice-A",
|
|
"TPU_WORKER_ID": "1",
|
|
"TPU_ACCELERATOR_TYPE": pod_type,
|
|
"TPU_TOPOLOGY": topology,
|
|
}
|
|
slice_a_worker_labels = {
|
|
"ray.io/tpu-slice-name": "slice-A",
|
|
"ray.io/tpu-worker-id": "1",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
}
|
|
cluster.add_node(
|
|
num_cpus=8,
|
|
resources={
|
|
"TPU": 4,
|
|
f"TPU-{pod_type}-head": 1,
|
|
"accelerator_type:TPU-V4": 1,
|
|
},
|
|
env_vars=slice_a_head_env,
|
|
labels=slice_a_head_labels,
|
|
)
|
|
cluster.add_node(
|
|
num_cpus=8,
|
|
resources={"TPU": 4, "accelerator_type:TPU-V4": 1},
|
|
env_vars=slice_a_worker_env,
|
|
labels=slice_a_worker_labels,
|
|
)
|
|
# Second TPU slice - v4-16 with 2 hosts and 4 chips/host.
|
|
slice_b_head_env = {
|
|
"TPU_NAME": "slice-B",
|
|
"TPU_WORKER_ID": "0",
|
|
"TPU_ACCELERATOR_TYPE": pod_type,
|
|
"TPU_TOPOLOGY": topology,
|
|
}
|
|
slice_b_head_labels = {
|
|
"ray.io/tpu-slice-name": "slice-B",
|
|
"ray.io/tpu-worker-id": "0",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
}
|
|
slice_b_worker_env = {
|
|
"TPU_NAME": "slice-B",
|
|
"TPU_WORKER_ID": "1",
|
|
"TPU_ACCELERATOR_TYPE": pod_type,
|
|
"TPU_TOPOLOGY": topology,
|
|
}
|
|
slice_b_worker_labels = {
|
|
"ray.io/tpu-slice-name": "slice-B",
|
|
"ray.io/tpu-worker-id": "1",
|
|
"ray.io/tpu-pod-type": pod_type,
|
|
}
|
|
cluster.add_node(
|
|
num_cpus=8,
|
|
resources={
|
|
"TPU": 4,
|
|
f"TPU-{pod_type}-head": 1,
|
|
"accelerator_type:TPU-V4": 1,
|
|
},
|
|
env_vars=slice_b_head_env,
|
|
labels=slice_b_head_labels,
|
|
)
|
|
cluster.add_node(
|
|
num_cpus=8,
|
|
resources={"TPU": 4, "accelerator_type:TPU-V4": 1},
|
|
env_vars=slice_b_worker_env,
|
|
labels=slice_b_worker_labels,
|
|
)
|
|
|
|
ray.init(address=cluster.address)
|
|
yield cluster
|
|
ray.shutdown()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reduce_health_check_interval(monkeypatch):
|
|
monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, "0.2")
|
|
yield
|
|
|
|
|
|
def train_func():
|
|
import jax
|
|
|
|
import ray
|
|
from ray import train
|
|
|
|
train_ctx = train.get_context()
|
|
rank = train_ctx.get_world_rank()
|
|
|
|
devices = jax.devices()
|
|
node_labels = ray.get_runtime_context().get_node_labels()
|
|
slice_name = node_labels.get("ray.io/tpu-slice-name")
|
|
current_ip = ray.util.get_node_ip_address()
|
|
|
|
megascale_vars = {
|
|
"MEGASCALE_SLICE_ID": os.environ.get("MEGASCALE_SLICE_ID"),
|
|
"MEGASCALE_NUM_SLICES": os.environ.get("MEGASCALE_NUM_SLICES"),
|
|
"MEGASCALE_COORDINATOR_ADDRESS": os.environ.get(
|
|
"MEGASCALE_COORDINATOR_ADDRESS"
|
|
),
|
|
"MEGASCALE_PORT": os.environ.get("MEGASCALE_PORT"),
|
|
}
|
|
|
|
train.report(
|
|
{
|
|
"worker_id": rank,
|
|
"slice_name": slice_name,
|
|
"node_ip": current_ip,
|
|
"devices": [str(d) for d in devices],
|
|
**megascale_vars,
|
|
}
|
|
)
|
|
|
|
|
|
class CustomMetricsCallback(UserCallback):
|
|
"""
|
|
In Ray Train V2, reporting metrics from all workers is a no-op, so we
|
|
utilize this callback to access the results in our tests.
|
|
"""
|
|
|
|
def __init__(self, actor_name: str):
|
|
self.actor_name = actor_name
|
|
|
|
def after_report(self, run_context, metrics, checkpoint):
|
|
# Connect to the specific verify actor for this test.
|
|
sink = ray.get_actor(self.actor_name)
|
|
sink.log.remote(metrics)
|
|
|
|
|
|
@ray.remote
|
|
class VerificationActor:
|
|
"""
|
|
This Actor is called from the custom metrics callback and saves
|
|
the reported metrics from each test.
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.reports = []
|
|
|
|
def log(self, metrics):
|
|
self.reports.extend(metrics)
|
|
|
|
def get_reports(self):
|
|
return self.reports
|
|
|
|
|
|
def test_tpu_single_host(ray_tpu_single_host, tmp_path):
|
|
"""
|
|
Tests single-host scheduling with no topology value. In this case, the
|
|
JaxTrainer should skip the multi-host slice scheduling logic and setup
|
|
with a single TPU worker.
|
|
"""
|
|
actor_name = "test_tpu_single_host"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
num_workers=1,
|
|
resources_per_worker={"TPU": 8},
|
|
accelerator_type="TPU-V6E",
|
|
),
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={"env_vars": {"JAX_PLATFORMS": "cpu"}},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
# Fetch metrics result using the verification actor.
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
|
|
# The train func should have ran on one single-host TPU.
|
|
assert len(reports) == 1, f"Expected 1 report, got {len(reports)}"
|
|
report = reports[0]
|
|
assert report["worker_id"] == 0
|
|
|
|
# Validate we do not automatically set megascale vars for single-slice.
|
|
for r in reports:
|
|
assert r.get("MEGASCALE_SLICE_ID") is None
|
|
assert r.get("MEGASCALE_NUM_SLICES") is None
|
|
assert r.get("MEGASCALE_COORDINATOR_ADDRESS") is None
|
|
|
|
|
|
def test_tpu_single_slice_multi_host(ray_tpu_multi_host, tmp_path):
|
|
"""
|
|
Tests scheduling on a single multi-host slice. The number of workers
|
|
is set by the user to match the number of hosts in the slice, with each
|
|
worker consuming the full resources on that host.
|
|
"""
|
|
actor_name = "test_tpu_single_slice_multi_host"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
accelerator_type="TPU-V4",
|
|
topology="2x2x2",
|
|
num_workers=2,
|
|
),
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={"env_vars": {"JAX_PLATFORMS": "cpu"}},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
# Fetch metrics result from each worker using the verification actor.
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
|
|
# Verify two TPU workers on the same slice ran the training func.
|
|
assert (
|
|
len(reports) == 2
|
|
), f"Expected 2 workers to report metrics, got {len(reports)}"
|
|
worker_ids = {r["worker_id"] for r in reports}
|
|
assert worker_ids == {0, 1}, "Expected unique worker IDs from 0 to N-1."
|
|
slices_used = {r["slice_name"] for r in reports}
|
|
assert len(slices_used) == 1, "Expected workers to be scheduled to 1 slice."
|
|
assert next(iter(slices_used)) in ("slice-A", "slice-B")
|
|
|
|
# Validate we do not automatically set megascale vars for single-slice.
|
|
for r in reports:
|
|
assert r.get("MEGASCALE_SLICE_ID") is None
|
|
assert r.get("MEGASCALE_NUM_SLICES") is None
|
|
assert r.get("MEGASCALE_COORDINATOR_ADDRESS") is None
|
|
|
|
|
|
def test_tpu_multi_slice_multi_host(ray_tpu_multi_host, tmp_path):
|
|
"""
|
|
Tests execution of TPU workers across multiple multi-host slices. The
|
|
user specifies num_workers equal to the total number of hosts across all
|
|
slices.
|
|
"""
|
|
actor_name = "test_tpu_multi_slice_multi_host"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
accelerator_type="TPU-V4",
|
|
topology="2x2x2",
|
|
num_workers=4,
|
|
),
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={"env_vars": {"JAX_PLATFORMS": "cpu"}},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
# Fetch metrics result from each worker using the verification actor.
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
|
|
# Verify execution of all 4 TPU workers across both slices.
|
|
assert (
|
|
len(reports) == 4
|
|
), f"Expected 4 workers to report metrics, got {len(reports)}"
|
|
worker_ids = {r["worker_id"] for r in reports}
|
|
assert worker_ids == {0, 1, 2, 3}, "Expected unique worker IDs from 0 to N-1."
|
|
slices_used = {r["slice_name"] for r in reports}
|
|
assert len(slices_used) == 2, "Expected workers to schedule across 2 slices."
|
|
assert "slice-A" in slices_used
|
|
assert "slice-B" in slices_used
|
|
|
|
# Verify megascale coordinator address set to IP of worker 0.
|
|
worker_0_report = next(r for r in reports if r["worker_id"] == 0)
|
|
expected_coordinator_ip = worker_0_report["node_ip"]
|
|
|
|
for r in reports:
|
|
assert r["MEGASCALE_COORDINATOR_ADDRESS"] == expected_coordinator_ip
|
|
assert r["MEGASCALE_NUM_SLICES"] == "2"
|
|
# When the user does not set MEGASCALE_PORT in their pod spec, Ray
|
|
# Train falls back to the default port from
|
|
# ``get_tpu_coordinator_env_vars``.
|
|
assert r["MEGASCALE_PORT"] == "8081"
|
|
|
|
# Validate MEGASCALE_SLICE_ID set based on indexed TPU Pod name.
|
|
slice_a_reports = [r for r in reports if r["slice_name"] == "slice-A"]
|
|
slice_b_reports = [r for r in reports if r["slice_name"] == "slice-B"]
|
|
|
|
assert list({r["MEGASCALE_SLICE_ID"] for r in slice_a_reports}) == ["0"]
|
|
assert list({r["MEGASCALE_SLICE_ID"] for r in slice_b_reports}) == ["1"]
|
|
|
|
|
|
def test_tpu_multi_slice_overrides_stale_megascale_env(ray_tpu_multi_host, tmp_path):
|
|
"""
|
|
Tests JaxTrainer's MEGASCALE_* env var handling when the underlying TPU
|
|
node provider has already baked values into the pod environment.
|
|
|
|
This simulates the multi-slice fault tolerance scenario where one slice was
|
|
preempted and the replacement pods were provisioned by the TPU node
|
|
provider with env vars reflecting a different slice id / total slice count
|
|
than the worker group Ray Train is currently scheduling. Without an
|
|
override, libtpu's megascale topology coordinator would wait for slices
|
|
that no longer exist and hang TPU initialization.
|
|
"""
|
|
actor_name = "test_tpu_multi_slice_overrides_stale_megascale_env"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
user_port_override = "9999"
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
accelerator_type="TPU-V4",
|
|
topology="2x2x2",
|
|
num_workers=4,
|
|
),
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={
|
|
"env_vars": {
|
|
"JAX_PLATFORMS": "cpu",
|
|
# Stale values, as the TPU node provider would inject on
|
|
# a freshly-provisioned replacement slice. These should
|
|
# all be overridden by Ray Train.
|
|
"MEGASCALE_COORDINATOR_ADDRESS": "stale-coordinator:9999",
|
|
"MEGASCALE_NUM_SLICES": "3",
|
|
"MEGASCALE_SLICE_ID": "2",
|
|
# User-customized port (e.g. set in pod spec to avoid a
|
|
# conflict with another process). This should be
|
|
# preserved by Ray Train, NOT overridden.
|
|
"MEGASCALE_PORT": user_port_override,
|
|
}
|
|
},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
assert (
|
|
len(reports) == 4
|
|
), f"Expected 4 workers to report metrics, got {len(reports)}"
|
|
|
|
# All workers should see the controller-computed coordinator address (the
|
|
# IP of worker 0), NOT the stale "stale-coordinator:9999" value.
|
|
worker_0_report = next(r for r in reports if r["worker_id"] == 0)
|
|
expected_coordinator_ip = worker_0_report["node_ip"]
|
|
|
|
for r in reports:
|
|
assert r["MEGASCALE_COORDINATOR_ADDRESS"] == expected_coordinator_ip, (
|
|
"Expected MEGASCALE_COORDINATOR_ADDRESS to be overridden by Ray "
|
|
f"Train, but worker {r['worker_id']} still has stale value "
|
|
f"{r['MEGASCALE_COORDINATOR_ADDRESS']}"
|
|
)
|
|
assert r["MEGASCALE_NUM_SLICES"] == "2", (
|
|
"Expected MEGASCALE_NUM_SLICES to be overridden to the actual "
|
|
f"slice count (2), but worker {r['worker_id']} still has stale "
|
|
f"value {r['MEGASCALE_NUM_SLICES']}"
|
|
)
|
|
assert r["MEGASCALE_PORT"] == user_port_override, (
|
|
"Expected MEGASCALE_PORT to preserve the user-provided pod-spec "
|
|
f"value '{user_port_override}', but worker {r['worker_id']} has "
|
|
f"value '{r['MEGASCALE_PORT']}'. Ray Train should not override "
|
|
"user-customized ports."
|
|
)
|
|
|
|
# Validate MEGASCALE_SLICE_ID was overridden per slice (0 / 1) rather than
|
|
# the stale "2" the provider injected on every worker.
|
|
slice_a_reports = [r for r in reports if r["slice_name"] == "slice-A"]
|
|
slice_b_reports = [r for r in reports if r["slice_name"] == "slice-B"]
|
|
assert list({r["MEGASCALE_SLICE_ID"] for r in slice_a_reports}) == ["0"]
|
|
assert list({r["MEGASCALE_SLICE_ID"] for r in slice_b_reports}) == ["1"]
|
|
|
|
|
|
def test_multi_slice_manual_resources(ray_tpu_multi_host, tmp_path):
|
|
"""
|
|
Tests execution of TPU workers across multiple multi-host slices when
|
|
`resources_per_worker` is specified. The JaxTrainer should execute across
|
|
both slices with num_workers workers of the specified resources.
|
|
"""
|
|
actor_name = "test_multi_slice_manual_resources"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
accelerator_type="TPU-V4",
|
|
topology="2x2x2",
|
|
resources_per_worker={"TPU": 1}, # 1 CPU added by default per-bundle.
|
|
num_workers=16,
|
|
),
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={"env_vars": {"JAX_PLATFORMS": "cpu"}},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
# Fetch metrics result from each worker using the verification actor.
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
|
|
# Verify execution of all 16 TPU workers across both v4-16 slices.
|
|
assert (
|
|
len(reports) == 16
|
|
), f"Expected 16 workers to report metrics, got {len(reports)}"
|
|
worker_ids = {r["worker_id"] for r in reports}
|
|
assert worker_ids == set(range(16)), "Expected unique worker IDs from 0 to N-1."
|
|
slices_used = {r["slice_name"] for r in reports}
|
|
assert len(slices_used) == 2, "Expected workers to span 2 slices."
|
|
assert "slice-A" in slices_used
|
|
assert "slice-B" in slices_used
|
|
|
|
# Verify megascale coordinator address set to IP of worker 0.
|
|
worker_0_report = next(r for r in reports if r["worker_id"] == 0)
|
|
expected_coordinator_ip = worker_0_report["node_ip"]
|
|
|
|
for r in reports:
|
|
assert r["MEGASCALE_COORDINATOR_ADDRESS"] == expected_coordinator_ip
|
|
assert r["MEGASCALE_NUM_SLICES"] == "2"
|
|
|
|
# Validate MEGASCALE_SLICE_ID set based on indexed TPU Pod name.
|
|
slice_a_reports = [r for r in reports if r["slice_name"] == "slice-A"]
|
|
slice_b_reports = [r for r in reports if r["slice_name"] == "slice-B"]
|
|
|
|
assert list({r["MEGASCALE_SLICE_ID"] for r in slice_a_reports}) == ["0"]
|
|
assert list({r["MEGASCALE_SLICE_ID"] for r in slice_b_reports}) == ["1"]
|
|
|
|
|
|
def test_tpu_multi_slice_uneven_workers(ray_tpu_multi_host, tmp_path):
|
|
"""
|
|
Tests that ScalingConfig raises a ValueError if the requested num_workers
|
|
does not divide evenly across TPU slices of the requested topology.
|
|
"""
|
|
# Default resources (1 worker per host).
|
|
with pytest.raises(ValueError, match="must be a multiple of"):
|
|
ScalingConfig(
|
|
use_tpu=True,
|
|
accelerator_type="TPU-V4",
|
|
topology="2x2x2",
|
|
num_workers=3, # Expect a multiple of 2.
|
|
)
|
|
# Explicit resources (1 TPU chip per worker).
|
|
with pytest.raises(ValueError, match="must be a multiple of"):
|
|
ScalingConfig(
|
|
use_tpu=True,
|
|
accelerator_type="TPU-V4",
|
|
topology="2x2x1",
|
|
resources_per_worker={"TPU": 1},
|
|
num_workers=6, # Expect a multiple of 4.
|
|
)
|
|
|
|
|
|
def train_func_with_data_multi_host():
|
|
import unittest.mock
|
|
|
|
import jax
|
|
|
|
import ray
|
|
from ray import train
|
|
|
|
train_ctx = train.get_context()
|
|
rank = train_ctx.get_world_rank()
|
|
|
|
devices = jax.devices()
|
|
local_devices = jax.local_devices()
|
|
node_labels = ray.get_runtime_context().get_node_labels()
|
|
slice_name = node_labels.get("ray.io/tpu-slice-name")
|
|
|
|
ds_shard = train.get_dataset_shard("train")
|
|
|
|
batches = []
|
|
# Mock process_allgather because the JAX CPU backend does not support multiprocess communication.
|
|
# Since our mock test distributes exact same data lengths uniformly, all workers report the identical
|
|
# array `[has_batch, local_batch_size]`. We simulate the gather by stacking the local array identically.
|
|
def mock_process_allgather(arr):
|
|
import jax.numpy as jnp
|
|
|
|
return jnp.stack([arr] * jax.process_count())
|
|
|
|
with unittest.mock.patch(
|
|
"jax.experimental.multihost_utils.process_allgather",
|
|
side_effect=mock_process_allgather,
|
|
):
|
|
# Each batch has 4 rows. iter_jax_batches should combine the batches from all workers and return a batch of size 4 * num_workers.
|
|
for batch in ds_shard.iter_jax_batches(batch_size=4, synchronize_batches=True):
|
|
batches.append(batch["id"].shape)
|
|
|
|
train.report(
|
|
{
|
|
"worker_id": rank,
|
|
"slice_name": slice_name,
|
|
"devices": len(devices),
|
|
"local_devices": len(local_devices),
|
|
"batches": batches,
|
|
}
|
|
)
|
|
|
|
|
|
def test_multi_host_data_iterator(ray_tpu_multi_host, tmp_path):
|
|
actor_name = "test_multi_host_data_iterator"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
import ray
|
|
|
|
# Create 32 rows, batch size 4.
|
|
# We have 4 workers. Each custom shard will have 8 rows (2 batches).
|
|
ds = ray.data.range(32)
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func_with_data_multi_host,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
accelerator_type="TPU-V4",
|
|
topology="2x2x2",
|
|
num_workers=4,
|
|
),
|
|
datasets={"train": ds},
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={
|
|
"env_vars": {
|
|
"JAX_PLATFORMS": "cpu",
|
|
"XLA_FLAGS": "--xla_force_host_platform_device_count=4",
|
|
}
|
|
},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
# Fetch metrics result from each worker using the verification actor.
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
|
|
# Verify execution of all 4 TPU workers across both slices.
|
|
assert (
|
|
len(reports) == 4
|
|
), f"Expected 4 workers to report metrics, got {len(reports)}"
|
|
|
|
worker_ids = {r["worker_id"] for r in reports}
|
|
assert worker_ids == {0, 1, 2, 3}, "Expected unique worker IDs from 0 to N-1."
|
|
|
|
for r in reports:
|
|
# With batch_size=4 and 8 rows per shard, each worker sees 2 batches.
|
|
# jax.devices() will return 16 global devices (CPU).
|
|
assert r["devices"] == 16
|
|
assert r["local_devices"] == 4
|
|
assert len(r["batches"]) == 2
|
|
|
|
# Each batch should have global batch size of 16.
|
|
for batch_shape in r["batches"]:
|
|
assert batch_shape == (16,)
|
|
|
|
|
|
def test_single_host_data_iterator_padding(ray_tpu_single_host, tmp_path):
|
|
actor_name = "test_single_host_data_iterator_padding"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
import numpy as np
|
|
|
|
import ray
|
|
|
|
# Create 44 rows. Batch size 16.
|
|
# 44 / 16 = 2 batches of 16, and 1 batch of 12.
|
|
# With paddings, we expect 3 batches of 16.
|
|
ds = ray.data.from_items([{"features": np.ones((8,))} for _ in range(44)])
|
|
|
|
def train_func():
|
|
from ray import train
|
|
|
|
ds_shard = train.get_dataset_shard("train")
|
|
batches = []
|
|
for batch in ds_shard.iter_jax_batches(
|
|
batch_size=16,
|
|
paddings=-1,
|
|
):
|
|
batches.append(batch["features"].shape)
|
|
train.report({"batches": batches})
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
num_workers=1,
|
|
resources_per_worker={"TPU": 8},
|
|
accelerator_type="TPU-V6E",
|
|
),
|
|
datasets={"train": ds},
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={
|
|
"env_vars": {
|
|
"JAX_PLATFORMS": "cpu",
|
|
"XLA_FLAGS": "--xla_force_host_platform_device_count=8",
|
|
}
|
|
},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
assert len(reports) == 1
|
|
assert len(reports[0]["batches"]) == 3
|
|
for shape in reports[0]["batches"]:
|
|
assert shape == (16, 8)
|
|
|
|
|
|
def test_single_host_data_iterator_dtypes(ray_tpu_single_host, tmp_path):
|
|
actor_name = "test_single_host_data_iterator_dtypes"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
import numpy as np
|
|
|
|
import ray
|
|
|
|
ds = ray.data.from_items([{"features": np.ones((8,))} for _ in range(32)])
|
|
|
|
def train_func():
|
|
import jax.numpy as jnp
|
|
|
|
from ray import train
|
|
|
|
ds_shard = train.get_dataset_shard("train")
|
|
batches = []
|
|
for batch in ds_shard.iter_jax_batches(
|
|
batch_size=16,
|
|
dtypes=jnp.float16,
|
|
):
|
|
assert batch["features"].dtype == jnp.float16
|
|
batches.append(batch["features"].shape)
|
|
train.report({"batches": batches})
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
num_workers=1,
|
|
resources_per_worker={"TPU": 8},
|
|
accelerator_type="TPU-V6E",
|
|
),
|
|
datasets={"train": ds},
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={
|
|
"env_vars": {
|
|
"JAX_PLATFORMS": "cpu",
|
|
"XLA_FLAGS": "--xla_force_host_platform_device_count=8",
|
|
}
|
|
},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
assert len(reports) == 1
|
|
assert len(reports[0]["batches"]) == 2
|
|
for shape in reports[0]["batches"]:
|
|
assert shape == (16, 8)
|
|
|
|
|
|
def train_func_with_data_single_host(config):
|
|
import jax
|
|
import numpy as np
|
|
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
|
|
|
|
from ray import train
|
|
|
|
devices = jax.devices()
|
|
local_devices = jax.local_devices()
|
|
|
|
# We requested 8 TPU resources per worker, so XLA_FLAGS should expose 8 mock local devices
|
|
assert len(devices) == 8
|
|
assert len(local_devices) == 8
|
|
|
|
# 8 local devices across a 2x4 mesh
|
|
mesh = Mesh(np.array(local_devices).reshape(2, 4), ("x", "y"))
|
|
named_sharding = NamedSharding(mesh, P("x", "y"))
|
|
|
|
ds_shard = train.get_dataset_shard("train")
|
|
|
|
batches = []
|
|
drop_last = config.get("drop_last", False) if config else False
|
|
# Local batch size must be evenly divisible by 8 (num_local_devices)
|
|
for batch in ds_shard.iter_jax_batches(
|
|
batch_size=16,
|
|
drop_last=drop_last,
|
|
):
|
|
arr = jax.device_put(batch["features"], named_sharding)
|
|
assert arr.sharding == named_sharding
|
|
batches.append(arr.shape)
|
|
|
|
train.report(
|
|
{
|
|
"devices": len(devices),
|
|
"local_devices": len(local_devices),
|
|
"batches": batches,
|
|
}
|
|
)
|
|
|
|
|
|
def test_single_host_data_iterator_2d(ray_tpu_single_host, tmp_path):
|
|
actor_name = "test_single_host_data_iterator_2d"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
import numpy as np
|
|
|
|
import ray
|
|
|
|
# Create 48 rows. Each row is {"features": np.ones((8,))}
|
|
# We have 1 worker. The worker processes 48 rows in batches of 16.
|
|
# The first 3 batches are (16, 8) in shape.
|
|
ds = ray.data.from_items([{"features": np.ones((8,))} for _ in range(48)])
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func_with_data_single_host,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
num_workers=1,
|
|
resources_per_worker={"TPU": 8},
|
|
accelerator_type="TPU-V6E",
|
|
),
|
|
datasets={"train": ds},
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
# Mock 8 local CPU devices exclusively for this worker process.
|
|
# This enables jax.device_put to successfully execute local 2D resharding
|
|
# completely in-memory because jax.process_count() defaults to 1.
|
|
worker_runtime_env={
|
|
"env_vars": {
|
|
"JAX_PLATFORMS": "cpu",
|
|
"XLA_FLAGS": "--xla_force_host_platform_device_count=8",
|
|
}
|
|
},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
assert len(reports) == 1
|
|
|
|
for r in reports:
|
|
assert r["devices"] == 8
|
|
assert r["local_devices"] == 8
|
|
assert len(r["batches"]) == 3
|
|
|
|
for batch_shape in r["batches"]:
|
|
assert batch_shape == (16, 8)
|
|
|
|
|
|
def test_single_host_data_iterator_2d_truncation(ray_tpu_single_host, tmp_path):
|
|
actor_name = "test_single_host_data_iterator_2d_truncation"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
import numpy as np
|
|
|
|
import ray
|
|
|
|
# Create 60 rows. Each row is {"features": np.ones((8,))}
|
|
# 60 is not divisible by 16, so the last 12 rows should be truncated.
|
|
# Thus, each local batch is (16, 8) in shape and we expect exactly 3 batches.
|
|
ds = ray.data.from_items([{"features": np.ones((8,))} for _ in range(60)])
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func_with_data_single_host,
|
|
train_loop_config={"drop_last": True},
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
num_workers=1,
|
|
resources_per_worker={"TPU": 8},
|
|
accelerator_type="TPU-V6E",
|
|
),
|
|
datasets={"train": ds},
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
# Mock 8 local CPU devices exclusively for this worker process.
|
|
# This enables jax.device_put to successfully execute local 2D resharding
|
|
# completely in-memory because jax.process_count() defaults to 1.
|
|
worker_runtime_env={
|
|
"env_vars": {
|
|
"JAX_PLATFORMS": "cpu",
|
|
"XLA_FLAGS": "--xla_force_host_platform_device_count=8",
|
|
}
|
|
},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
assert len(reports) == 1
|
|
|
|
for r in reports:
|
|
assert r["devices"] == 8
|
|
assert r["local_devices"] == 8
|
|
assert len(r["batches"]) == 3
|
|
|
|
for batch_shape in r["batches"]:
|
|
assert batch_shape == (16, 8)
|
|
|
|
|
|
def test_single_host_data_iterator_2d_truncation_failure(ray_tpu_single_host, tmp_path):
|
|
import numpy as np
|
|
|
|
import ray
|
|
|
|
# Create 60 rows. Each row is {"features": np.ones((8,))}
|
|
# 60 is not divisible by 16. So if drop_last=False (the default), it should raise a ValueError.
|
|
ds = ray.data.from_items([{"features": np.ones((8,))} for _ in range(60)])
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func_with_data_single_host,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
num_workers=1,
|
|
resources_per_worker={"TPU": 8},
|
|
accelerator_type="TPU-V6E",
|
|
),
|
|
datasets={"train": ds},
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
worker_runtime_env={
|
|
"env_vars": {
|
|
"JAX_PLATFORMS": "cpu",
|
|
"XLA_FLAGS": "--xla_force_host_platform_device_count=8",
|
|
}
|
|
},
|
|
),
|
|
)
|
|
with pytest.raises(
|
|
Exception, match="evenly divisible by the number of local JAX devices"
|
|
):
|
|
trainer.fit()
|
|
|
|
|
|
def test_scaling_config_validation():
|
|
with pytest.raises(
|
|
ValueError, match="Cannot set `label_selector` when `use_tpu=True`"
|
|
):
|
|
ScalingConfig(
|
|
num_workers=2,
|
|
use_tpu=True,
|
|
topology="2x2x2",
|
|
accelerator_type="TPU-V4",
|
|
label_selector={"subcluster": "my_subcluster"},
|
|
)
|
|
|
|
|
|
def train_func_with_collate_fn(config):
|
|
import numpy as np
|
|
|
|
from ray import train
|
|
from ray.data.collate_fn import NumpyBatchCollateFn
|
|
|
|
class CustomCollateFn(NumpyBatchCollateFn):
|
|
def __call__(self, batch):
|
|
# Combine "col1" and "col2" columns into a single "features" tensor
|
|
return np.stack((batch["col1"], batch["col2"]), axis=1)
|
|
|
|
ds_shard = train.get_dataset_shard("train")
|
|
|
|
batches = []
|
|
for batch in ds_shard.iter_jax_batches(
|
|
batch_size=16,
|
|
collate_fn=CustomCollateFn(),
|
|
):
|
|
# The output of collate_fn is now a single jax.Array (global sharding)
|
|
# Check a specific value: it should be [1.0, 2.0] at the first position
|
|
assert np.array_equal(batch[0], np.array([1.0, 2.0]))
|
|
batches.append(batch.shape)
|
|
|
|
train.report(
|
|
{
|
|
"batches": batches,
|
|
}
|
|
)
|
|
|
|
|
|
def test_single_host_data_iterator_collate_fn(ray_tpu_single_host, tmp_path):
|
|
actor_name = "test_single_host_data_iterator_collate_fn"
|
|
verify_actor = VerificationActor.options(name=actor_name).remote()
|
|
|
|
import ray
|
|
|
|
# Create 32 rows. Each row has "col1" (ones) and "col2" (twos).
|
|
ds = ray.data.from_items([{"col1": 1.0, "col2": 2.0} for _ in range(32)])
|
|
|
|
trainer = JaxTrainer(
|
|
train_loop_per_worker=train_func_with_collate_fn,
|
|
scaling_config=ScalingConfig(
|
|
use_tpu=True,
|
|
num_workers=1,
|
|
resources_per_worker={"TPU": 8},
|
|
accelerator_type="TPU-V6E",
|
|
),
|
|
datasets={"train": ds},
|
|
run_config=RunConfig(
|
|
storage_path=str(tmp_path),
|
|
callbacks=[CustomMetricsCallback(actor_name)],
|
|
worker_runtime_env={
|
|
"env_vars": {
|
|
"JAX_PLATFORMS": "cpu",
|
|
"XLA_FLAGS": "--xla_force_host_platform_device_count=8",
|
|
}
|
|
},
|
|
),
|
|
)
|
|
result = trainer.fit()
|
|
assert result.error is None
|
|
|
|
reports = ray.get(verify_actor.get_reports.remote())
|
|
assert len(reports) == 1
|
|
|
|
for r in reports:
|
|
# We expect 2 batches (32 rows / 16 batch_size)
|
|
# Each batch should have shape (16, 2)
|
|
assert len(r["batches"]) == 2
|
|
for batch_shape in r["batches"]:
|
|
assert batch_shape == (16, 2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|