wehub-resource-sync
106 lines
3.3 KiB
Python
106 lines
3.3 KiB
Python
import pyarrow.fs
|
|
import pytest
|
|
|
|
from ray.train import RunConfig, ScalingConfig
|
|
|
|
|
|
def test_scaling_config_validation():
|
|
assert ScalingConfig(
|
|
num_workers=2, use_gpu=True, resources_per_worker={"CPU": 1}
|
|
).total_resources == {"CPU": 2, "GPU": 2}
|
|
|
|
with pytest.raises(ValueError, match="`use_gpu` is False but `GPU` was found in"):
|
|
ScalingConfig(num_workers=2, use_gpu=False, resources_per_worker={"GPU": 1})
|
|
|
|
with pytest.raises(ValueError, match="Cannot specify both"):
|
|
ScalingConfig(num_workers=2, use_gpu=True, use_tpu=True)
|
|
|
|
with pytest.raises(
|
|
ValueError,
|
|
match=(
|
|
"If `label_selector` is a list, it must be the same length as "
|
|
"`max_workers`"
|
|
),
|
|
):
|
|
ScalingConfig(num_workers=2, label_selector=[{"subcluster": "my_subcluster"}])
|
|
with pytest.raises(
|
|
ValueError,
|
|
match=(
|
|
"If `label_selector` is a list, it must be the same length as "
|
|
"`max_workers`"
|
|
),
|
|
):
|
|
ScalingConfig(
|
|
num_workers=(2, 3),
|
|
label_selector=[{"subcluster": "a"}, {"subcluster": "b"}],
|
|
)
|
|
|
|
|
|
def test_label_selector_per_worker():
|
|
# None -> None (no constraint; downstream consumers handle this directly).
|
|
assert ScalingConfig(num_workers=3)._label_selector_per_worker(3) is None
|
|
|
|
# Dict -> replicated per worker, decoupled from the original.
|
|
cfg = ScalingConfig(num_workers=2, label_selector={"zone": "a"})
|
|
result = cfg._label_selector_per_worker(2)
|
|
assert result == [{"zone": "a"}, {"zone": "a"}]
|
|
result[0]["zone"] = "b"
|
|
assert cfg.label_selector == {"zone": "a"}
|
|
|
|
# List -> sliced to num_workers, decoupled from the original.
|
|
cfg = ScalingConfig(
|
|
num_workers=(1, 3),
|
|
label_selector=[{"a": "1"}, {"a": "2"}, {"a": "3"}],
|
|
)
|
|
assert cfg._label_selector_per_worker(2) == [{"a": "1"}, {"a": "2"}]
|
|
|
|
|
|
def test_scaling_config_accelerator_type():
|
|
scaling_config = ScalingConfig(num_workers=2, use_gpu=True, accelerator_type="A100")
|
|
assert scaling_config.accelerator_type == "A100"
|
|
assert scaling_config._resources_per_worker_not_none == {
|
|
"GPU": 1,
|
|
"accelerator_type:A100": 0.001,
|
|
}
|
|
assert scaling_config.total_resources == {
|
|
"GPU": 2,
|
|
"accelerator_type:A100": 0.002,
|
|
}
|
|
assert scaling_config.additional_resources_per_worker == {
|
|
"accelerator_type:A100": 0.001
|
|
}
|
|
|
|
|
|
def test_scaling_config_tpu_min_workers_multiple():
|
|
with pytest.raises(ValueError, match="min_workers"):
|
|
ScalingConfig(
|
|
num_workers=(1, 2),
|
|
use_tpu=True,
|
|
topology="2x2x2",
|
|
accelerator_type="TPU-V4",
|
|
resources_per_worker={"TPU": 4},
|
|
)
|
|
|
|
|
|
def test_storage_filesystem_repr():
|
|
"""Test for https://github.com/ray-project/ray/pull/40851"""
|
|
config = RunConfig(storage_filesystem=pyarrow.fs.S3FileSystem())
|
|
repr(config)
|
|
|
|
|
|
def test_scaling_config_default_workers():
|
|
"""Test that num_workers defaults to 1 for non-TPU workloads."""
|
|
config = ScalingConfig()
|
|
assert config.num_workers == 1
|
|
assert config.total_resources == {"CPU": 1}
|
|
|
|
config_gpu = ScalingConfig(use_gpu=True)
|
|
assert config_gpu.num_workers == 1
|
|
assert config_gpu.total_resources == {"GPU": 1}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|