1
1
"""
2
- Tests for the simplified RayJobClusterConfig accelerator_configs behavior.
2
+ Tests for the simplified ManagedClusterConfig accelerator_configs behavior.
3
3
"""
4
4
5
5
import pytest
6
- from codeflare_sdk .ray .rayjobs .config import RayJobClusterConfig , DEFAULT_ACCELERATORS
6
+ from codeflare_sdk .ray .rayjobs .config import ManagedClusterConfig , DEFAULT_ACCELERATORS
7
7
8
8
9
9
def test_accelerator_configs_defaults_to_default_accelerators ():
10
10
"""Test that accelerator_configs defaults to DEFAULT_ACCELERATORS.copy()"""
11
- config = RayJobClusterConfig ()
11
+ config = ManagedClusterConfig ()
12
12
13
13
# Should have all the default accelerators
14
14
assert "nvidia.com/gpu" in config .accelerator_configs
@@ -27,7 +27,7 @@ def test_accelerator_configs_can_be_overridden():
27
27
"custom.com/accelerator" : "CUSTOM_ACCELERATOR" ,
28
28
}
29
29
30
- config = RayJobClusterConfig (accelerator_configs = custom_configs )
30
+ config = ManagedClusterConfig (accelerator_configs = custom_configs )
31
31
32
32
# Should have custom configs
33
33
assert config .accelerator_configs == custom_configs
@@ -46,7 +46,7 @@ def test_accelerator_configs_can_extend_defaults():
46
46
"custom.com/accelerator" : "CUSTOM_ACCEL" ,
47
47
}
48
48
49
- config = RayJobClusterConfig (accelerator_configs = extended_configs )
49
+ config = ManagedClusterConfig (accelerator_configs = extended_configs )
50
50
51
51
# Should have all defaults plus custom
52
52
assert "nvidia.com/gpu" in config .accelerator_configs
@@ -57,15 +57,15 @@ def test_accelerator_configs_can_extend_defaults():
57
57
58
58
def test_gpu_validation_works_with_defaults ():
59
59
"""Test that GPU validation works with default accelerator configs"""
60
- config = RayJobClusterConfig (head_accelerators = {"nvidia.com/gpu" : 1 })
60
+ config = ManagedClusterConfig (head_accelerators = {"nvidia.com/gpu" : 1 })
61
61
62
62
# Should not raise any errors
63
63
assert config .head_accelerators == {"nvidia.com/gpu" : 1 }
64
64
65
65
66
66
def test_gpu_validation_works_with_custom_configs ():
67
67
"""Test that GPU validation works with custom accelerator configs"""
68
- config = RayJobClusterConfig (
68
+ config = ManagedClusterConfig (
69
69
accelerator_configs = {"custom.com/accelerator" : "CUSTOM_ACCEL" },
70
70
head_accelerators = {"custom.com/accelerator" : 1 },
71
71
)
@@ -79,4 +79,55 @@ def test_gpu_validation_fails_with_unsupported_accelerator():
79
79
with pytest .raises (
80
80
ValueError , match = "GPU configuration 'unsupported.com/accelerator' not found"
81
81
):
82
- RayJobClusterConfig (head_accelerators = {"unsupported.com/accelerator" : 1 })
82
+ ManagedClusterConfig (head_accelerators = {"unsupported.com/accelerator" : 1 })
83
+
84
+
85
+ def test_ray_usage_stats_always_disabled_by_default ():
86
+ """Test that RAY_USAGE_STATS_ENABLED is always set to '0' by default"""
87
+ config = ManagedClusterConfig ()
88
+
89
+ # Should always have the environment variable set to "0"
90
+ assert "RAY_USAGE_STATS_ENABLED" in config .envs
91
+ assert config .envs ["RAY_USAGE_STATS_ENABLED" ] == "0"
92
+
93
+
94
+ def test_ray_usage_stats_overwrites_user_env ():
95
+ """Test that RAY_USAGE_STATS_ENABLED is always set to '0' even if user specifies it"""
96
+ # User tries to enable usage stats
97
+ config = ManagedClusterConfig (envs = {"RAY_USAGE_STATS_ENABLED" : "1" })
98
+
99
+ # Should still be disabled (our setting takes precedence)
100
+ assert "RAY_USAGE_STATS_ENABLED" in config .envs
101
+ assert config .envs ["RAY_USAGE_STATS_ENABLED" ] == "0"
102
+
103
+
104
+ def test_ray_usage_stats_overwrites_user_env_string ():
105
+ """Test that RAY_USAGE_STATS_ENABLED is always set to '0' even if user specifies it as string"""
106
+ # User tries to enable usage stats with string
107
+ config = ManagedClusterConfig (envs = {"RAY_USAGE_STATS_ENABLED" : "true" })
108
+
109
+ # Should still be disabled (our setting takes precedence)
110
+ assert "RAY_USAGE_STATS_ENABLED" in config .envs
111
+ assert config .envs ["RAY_USAGE_STATS_ENABLED" ] == "0"
112
+
113
+
114
+ def test_ray_usage_stats_with_other_user_envs ():
115
+ """Test that RAY_USAGE_STATS_ENABLED is set correctly while preserving other user envs"""
116
+ # User sets other environment variables
117
+ user_envs = {
118
+ "CUSTOM_VAR" : "custom_value" ,
119
+ "ANOTHER_VAR" : "another_value" ,
120
+ "RAY_USAGE_STATS_ENABLED" : "1" , # This should be overwritten
121
+ }
122
+
123
+ config = ManagedClusterConfig (envs = user_envs )
124
+
125
+ # Our setting should take precedence
126
+ assert config .envs ["RAY_USAGE_STATS_ENABLED" ] == "0"
127
+
128
+ # Other user envs should be preserved
129
+ assert config .envs ["CUSTOM_VAR" ] == "custom_value"
130
+ assert config .envs ["ANOTHER_VAR" ] == "another_value"
131
+
132
+ # Total count should be correct (3 user envs)
133
+ assert len (config .envs ) == 3
0 commit comments