chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_config.py

@@ -0,0 +1,105 @@
+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__]))