wehub-resource-sync
323 lines
9.8 KiB
Python
323 lines
9.8 KiB
Python
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from ray.data.util.jax_util import jax_sync_generator
|
|
|
|
pytest.importorskip("jax")
|
|
|
|
|
|
def test_jax_sync_generator_empty_batch(ray_start_regular_shared):
|
|
def empty_batch_iterable():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
yield {} # Empty dict batch
|
|
yield {"a": np.array([4, 5, 6])}
|
|
|
|
gen = jax_sync_generator(empty_batch_iterable(), drop_last=True, batch_size=3)
|
|
results = list(gen)
|
|
|
|
assert len(results) == 2
|
|
assert np.array_equal(results[0]["a"], np.array([1, 2, 3]))
|
|
assert np.array_equal(results[1]["a"], np.array([4, 5, 6]))
|
|
|
|
|
|
def test_jax_sync_generator_empty_column(ray_start_regular_shared):
|
|
def empty_column_iterable():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
yield {"a": np.array([])} # Dict with empty column
|
|
yield {"a": np.array([4, 5, 6])}
|
|
|
|
gen = jax_sync_generator(empty_column_iterable(), drop_last=True, batch_size=3)
|
|
results = list(gen)
|
|
|
|
assert len(results) == 2
|
|
assert np.array_equal(results[0]["a"], np.array([1, 2, 3]))
|
|
assert np.array_equal(results[1]["a"], np.array([4, 5, 6]))
|
|
|
|
|
|
def test_jax_sync_generator_no_sync(ray_start_regular_shared):
|
|
def batches():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
yield {"a": np.array([4, 5, 6])}
|
|
|
|
# Should work fine in single process even with sync=False
|
|
gen = jax_sync_generator(
|
|
batches(), drop_last=True, batch_size=3, synchronize_batches=False
|
|
)
|
|
results = list(gen)
|
|
assert len(results) == 2
|
|
|
|
|
|
def test_jax_sync_generator_padding(ray_start_regular_shared):
|
|
def batches():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
yield {"a": np.array([4, 5])}
|
|
|
|
# Should pad the second batch to size 3 with value -1
|
|
gen = jax_sync_generator(
|
|
batches(),
|
|
drop_last=False,
|
|
batch_size=3,
|
|
paddings=-1,
|
|
synchronize_batches=False,
|
|
)
|
|
results = list(gen)
|
|
|
|
assert len(results) == 2
|
|
assert len(results[0]["a"]) == 3
|
|
assert len(results[1]["a"]) == 3
|
|
assert np.array_equal(results[1]["a"], np.array([4, 5, -1]))
|
|
|
|
|
|
def test_jax_sync_generator_drop_last(ray_start_regular_shared):
|
|
def batches():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
yield {"a": np.array([4, 5])}
|
|
|
|
# Should drop the second batch because it's not size 3 and padding=None
|
|
# Note: in single host, it doesn't drop unless we use _iter_batches(drop_last=True)
|
|
# But jax_sync_generator with drop_last=True will raise error if sizes don't match min.
|
|
# Actually, jax_sync_generator logic for single host just passes through if not sync.
|
|
|
|
# Let's test single host with divisibility check failure
|
|
gen = jax_sync_generator(
|
|
batches(),
|
|
drop_last=False,
|
|
batch_size=3,
|
|
paddings=None,
|
|
synchronize_batches=False,
|
|
)
|
|
results = list(gen)
|
|
assert len(results) == 2 # Both yielded because 2 is divisible by 1 local device
|
|
|
|
# Let's force num_local_devices = 4 for testing error handling
|
|
from unittest.mock import patch
|
|
|
|
with patch("jax.local_device_count", return_value=4):
|
|
gen = jax_sync_generator(
|
|
batches(),
|
|
drop_last=True,
|
|
batch_size=3,
|
|
paddings=None,
|
|
synchronize_batches=False,
|
|
)
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="evenly divisible by the number of local JAX devices",
|
|
):
|
|
list(gen)
|
|
|
|
|
|
def test_jax_sync_generator_multi_host_uneven_batches_with_padding(
|
|
ray_start_regular_shared,
|
|
):
|
|
|
|
from unittest.mock import patch
|
|
|
|
def batches():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
# Host 0 ends here, Host 1 has more
|
|
|
|
# Mock jax environment: 2 hosts, 1 device per host
|
|
with patch("jax.process_count", return_value=2), patch(
|
|
"jax.local_device_count", return_value=1
|
|
), patch(
|
|
"ray.data.util.jax_util._convert_batch",
|
|
side_effect=lambda x, sharding, **kwargs: x,
|
|
):
|
|
|
|
def mock_process_allgather(arr):
|
|
# Simulate Host 1 having more data
|
|
# local_infos for host 0: [1, 3, 0, 0, ...] (from batches())
|
|
# arr is a JAX array because jax_sync_generator does jnp.array(local_infos)
|
|
# Convert to numpy for easy manipulation
|
|
h0_infos = np.array(arr)
|
|
|
|
h1_infos = np.zeros_like(h0_infos)
|
|
h1_infos[0] = 1 # batch 1 exists
|
|
h1_infos[1] = 3 # batch 1 size
|
|
h1_infos[2] = 1 # batch 2 exists
|
|
h1_infos[3] = 3 # batch 2 size
|
|
|
|
import jax.numpy as jnp
|
|
|
|
return jnp.array([h0_infos, h1_infos])
|
|
|
|
with patch(
|
|
"jax.experimental.multihost_utils.process_allgather",
|
|
side_effect=mock_process_allgather,
|
|
):
|
|
# Host 0 uses jax_sync_generator
|
|
gen = jax_sync_generator(
|
|
batches(),
|
|
drop_last=False,
|
|
batch_size=3,
|
|
paddings=-1,
|
|
synchronize_batches=True,
|
|
)
|
|
# Should yield 2 batches: one real, one dummy
|
|
results = list(gen)
|
|
assert len(results) == 2
|
|
assert np.array_equal(results[0]["a"], np.array([1, 2, 3]))
|
|
assert np.array_equal(results[1]["a"], np.array([-1, -1, -1]))
|
|
|
|
|
|
def test_jax_sync_generator_multi_host_uneven_batches_drop_last(
|
|
ray_start_regular_shared,
|
|
):
|
|
|
|
from unittest.mock import patch
|
|
|
|
def batches():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
yield {"a": np.array([4, 5, 6])}
|
|
|
|
with patch("jax.process_count", return_value=2), patch(
|
|
"jax.local_device_count", return_value=1
|
|
), patch(
|
|
"ray.data.util.jax_util._convert_batch",
|
|
side_effect=lambda x, sharding, **kwargs: x,
|
|
):
|
|
|
|
def mock_process_allgather(arr):
|
|
# Host 0 has 2 batches, Host 1 has 1 batch
|
|
h0_infos = np.array(arr)
|
|
h1_infos = np.zeros_like(h0_infos)
|
|
h1_infos[0] = 1
|
|
h1_infos[1] = 3
|
|
import jax.numpy as jnp
|
|
|
|
return jnp.array([h0_infos, h1_infos])
|
|
|
|
with patch(
|
|
"jax.experimental.multihost_utils.process_allgather",
|
|
side_effect=mock_process_allgather,
|
|
):
|
|
gen = jax_sync_generator(
|
|
batches(),
|
|
drop_last=True,
|
|
batch_size=3,
|
|
synchronize_batches=True,
|
|
)
|
|
# Should yield only 1 batch and stop
|
|
results = list(gen)
|
|
assert len(results) == 1
|
|
assert np.array_equal(results[0]["a"], np.array([1, 2, 3]))
|
|
|
|
|
|
def test_jax_sync_generator_multi_host_uneven_batch_sizes_fail(
|
|
ray_start_regular_shared,
|
|
):
|
|
|
|
from unittest.mock import patch
|
|
|
|
def batches():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
|
|
with patch("jax.process_count", return_value=2), patch(
|
|
"jax.local_device_count", return_value=1
|
|
), patch(
|
|
"ray.data.util.jax_util._convert_batch",
|
|
side_effect=lambda x, sharding, **kwargs: x,
|
|
):
|
|
|
|
def mock_process_allgather(arr):
|
|
# Host 0 batch size 3, Host 1 batch size 2
|
|
h0_infos = np.array(arr)
|
|
h1_infos = h0_infos.copy()
|
|
h1_infos[1] = 2
|
|
import jax.numpy as jnp
|
|
|
|
return jnp.array([h0_infos, h1_infos])
|
|
|
|
with patch(
|
|
"jax.experimental.multihost_utils.process_allgather",
|
|
side_effect=mock_process_allgather,
|
|
):
|
|
gen = jax_sync_generator(
|
|
batches(),
|
|
drop_last=False,
|
|
batch_size=3,
|
|
synchronize_batches=True,
|
|
)
|
|
with pytest.raises(
|
|
ValueError, match="Uneven batch sizes detected across JAX workers"
|
|
):
|
|
list(gen)
|
|
|
|
|
|
def test_jax_sync_generator_multi_host_uneven_num_batches_fail(
|
|
ray_start_regular_shared,
|
|
):
|
|
|
|
from unittest.mock import patch
|
|
|
|
def batches():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
yield {"a": np.array([4, 5, 6])}
|
|
|
|
with patch("jax.process_count", return_value=2), patch(
|
|
"jax.local_device_count", return_value=1
|
|
), patch(
|
|
"ray.data.util.jax_util._convert_batch",
|
|
side_effect=lambda x, sharding, **kwargs: x,
|
|
):
|
|
|
|
def mock_process_allgather(arr):
|
|
# Host 0 has 2 batches, Host 1 has 1 batch, no padding
|
|
h0_infos = np.array(arr)
|
|
h1_infos = np.zeros_like(h0_infos)
|
|
h1_infos[0] = 1
|
|
h1_infos[1] = 3
|
|
import jax.numpy as jnp
|
|
|
|
return jnp.array([h0_infos, h1_infos])
|
|
|
|
with patch(
|
|
"jax.experimental.multihost_utils.process_allgather",
|
|
side_effect=mock_process_allgather,
|
|
):
|
|
gen = jax_sync_generator(
|
|
batches(),
|
|
drop_last=False,
|
|
batch_size=3,
|
|
synchronize_batches=True,
|
|
)
|
|
with pytest.raises(
|
|
ValueError,
|
|
match="Uneven number of batches detected across JAX workers",
|
|
):
|
|
list(gen)
|
|
|
|
|
|
def test_jax_sync_generator_with_dtypes(ray_start_regular_shared):
|
|
def batches():
|
|
yield {"a": np.array([1, 2, 3])}
|
|
|
|
import jax.numpy as jnp
|
|
|
|
dtypes = {"a": jnp.float16}
|
|
|
|
# Mock _convert_batch to capture dtypes
|
|
mock_convert = MagicMock(side_effect=lambda x, sharding, dtypes=None: x)
|
|
|
|
with patch("ray.data.util.jax_util._convert_batch", mock_convert):
|
|
gen = jax_sync_generator(
|
|
batches(),
|
|
drop_last=False,
|
|
batch_size=3,
|
|
dtypes=dtypes,
|
|
synchronize_batches=False,
|
|
)
|
|
list(gen)
|
|
|
|
# Verify that dtypes was passed to _convert_batch
|
|
mock_convert.assert_called_once()
|
|
assert mock_convert.call_args[1]["dtypes"] == dtypes
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|