Skip to content

Commit ef9c6c2

Browse files
committed
Trying to bump job file coverage again
Signed-off-by: Pat O'Connor <[email protected]>
1 parent 6f6442f commit ef9c6c2

File tree

1 file changed

+57
-53
lines changed

1 file changed

+57
-53
lines changed

src/codeflare_sdk/ray/job/test_job.py

Lines changed: 57 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,9 @@ def test_ray_job_spec_creation():
1414
entrypoint_memory=1024,
1515
entrypoint_resources={"custom.com/resource": 1.0},
1616
cluster_name="test-cluster",
17-
cluster_namespace="test-ns"
17+
cluster_namespace="test-ns",
1818
)
19-
19+
2020
assert spec.entrypoint == "python script.py"
2121
assert spec.submission_id == "test-123"
2222
assert spec.runtime_env == {"pip": ["numpy", "pandas"]}
@@ -35,47 +35,56 @@ def test_ray_job_yaml_generation_full():
3535
spec = RayJobSpec(
3636
entrypoint="python -c 'import ray; print(ray.cluster_resources())'",
3737
submission_id="test-submission-456",
38-
runtime_env={"pip": ["numpy==1.24.0", "pandas"], "env_vars": {"RAY_LOG_LEVEL": "DEBUG"}},
38+
runtime_env={
39+
"pip": ["numpy==1.24.0", "pandas"],
40+
"env_vars": {"RAY_LOG_LEVEL": "DEBUG"},
41+
},
3942
metadata={"job_timeout_s": "1800", "author": "test-user"},
4043
entrypoint_num_cpus=2.5,
4144
entrypoint_num_gpus=1.0,
4245
entrypoint_memory=2048,
4346
entrypoint_resources={"custom.com/special": 1.0},
4447
cluster_name="ml-cluster",
45-
cluster_namespace="ml-namespace"
48+
cluster_namespace="ml-namespace",
4649
)
47-
50+
4851
job = RayJob(
4952
metadata={
5053
"name": "comprehensive-test-job",
5154
"namespace": "test-namespace",
52-
"labels": {"app": "ml-training"}
55+
"labels": {"app": "ml-training"},
5356
},
54-
spec=spec
57+
spec=spec,
5558
)
56-
59+
5760
yaml_dict = job.to_dict()
58-
61+
5962
# Verify top-level structure
6063
assert yaml_dict["apiVersion"] == "ray.io/v1"
6164
assert yaml_dict["kind"] == "RayJob"
6265
assert yaml_dict["metadata"]["name"] == "comprehensive-test-job"
6366
assert yaml_dict["metadata"]["namespace"] == "test-namespace"
6467
assert yaml_dict["metadata"]["labels"] == {"app": "ml-training"}
65-
68+
6669
# Verify spec section
6770
spec_dict = yaml_dict["spec"]
68-
assert spec_dict["entrypoint"] == "python -c 'import ray; print(ray.cluster_resources())'"
71+
assert (
72+
spec_dict["entrypoint"]
73+
== "python -c 'import ray; print(ray.cluster_resources())'"
74+
)
6975
assert spec_dict["submission_id"] == "test-submission-456"
70-
assert spec_dict["runtime_env"] == {"pip": ["numpy==1.24.0", "pandas"], "env_vars": {"RAY_LOG_LEVEL": "DEBUG"}}
76+
assert spec_dict["runtime_env"] == {
77+
"pip": ["numpy==1.24.0", "pandas"],
78+
"env_vars": {"RAY_LOG_LEVEL": "DEBUG"},
79+
}
7180
assert spec_dict["metadata"] == {"job_timeout_s": "1800", "author": "test-user"}
7281
assert spec_dict["entrypoint_num_cpus"] == 2.5
7382
assert spec_dict["entrypoint_num_gpus"] == 1.0
7483
assert spec_dict["entrypoint_memory"] == 2048
7584
assert spec_dict["entrypoint_resources"] == {"custom.com/special": 1.0}
7685
assert spec_dict["cluster_name"] == "ml-cluster"
7786
assert spec_dict["cluster_namespace"] == "ml-namespace"
78-
87+
7988
# Verify status section (should use spec status when job.status is None)
8089
status_dict = yaml_dict["status"]
8190
assert status_dict["status"] == RayJobStatus.PENDING
@@ -88,19 +97,16 @@ def test_ray_job_yaml_generation_full():
8897
def test_ray_job_yaml_generation_minimal():
8998
"""Test RayJob YAML generation with minimal required fields."""
9099
spec = RayJobSpec(entrypoint="python minimal_job.py")
91-
92-
job = RayJob(
93-
metadata={"name": "minimal-job"},
94-
spec=spec
95-
)
96-
100+
101+
job = RayJob(metadata={"name": "minimal-job"}, spec=spec)
102+
97103
yaml_dict = job.to_dict()
98-
104+
99105
# Verify structure
100106
assert yaml_dict["apiVersion"] == "ray.io/v1"
101107
assert yaml_dict["kind"] == "RayJob"
102108
assert yaml_dict["metadata"]["name"] == "minimal-job"
103-
109+
104110
# Verify spec has only required field and defaults to None for others
105111
spec_dict = yaml_dict["spec"]
106112
assert spec_dict["entrypoint"] == "python minimal_job.py"
@@ -113,7 +119,7 @@ def test_ray_job_yaml_generation_minimal():
113119
assert spec_dict["entrypoint_resources"] is None
114120
assert spec_dict["cluster_name"] is None
115121
assert spec_dict["cluster_namespace"] is None
116-
122+
117123
# Verify default status
118124
status_dict = yaml_dict["status"]
119125
assert status_dict["status"] == RayJobStatus.PENDING
@@ -123,9 +129,9 @@ def test_ray_job_yaml_with_existing_status():
123129
"""Test RayJob YAML generation when status is pre-populated."""
124130
spec = RayJobSpec(
125131
entrypoint="python running_job.py",
126-
status=RayJobStatus.PENDING # This should be overridden by job.status
132+
status=RayJobStatus.PENDING, # This should be overridden by job.status
127133
)
128-
134+
129135
# Simulate status from Kubernetes controller
130136
existing_status = {
131137
"status": "RUNNING",
@@ -135,18 +141,18 @@ def test_ray_job_yaml_with_existing_status():
135141
"driver_info": {
136142
"id": "driver-abc123",
137143
"node_ip_address": "10.244.1.5",
138-
"pid": "12345"
139-
}
144+
"pid": "12345",
145+
},
140146
}
141-
147+
142148
job = RayJob(
143149
metadata={"name": "status-test-job", "namespace": "test-ns"},
144150
spec=spec,
145-
status=existing_status
151+
status=existing_status,
146152
)
147-
153+
148154
yaml_dict = job.to_dict()
149-
155+
150156
# Should use existing status, not spec status
151157
assert yaml_dict["status"] == existing_status
152158
assert yaml_dict["status"]["status"] == "RUNNING"
@@ -163,36 +169,38 @@ def test_ray_job_yaml_with_complex_runtime_env():
163169
"env_vars": {
164170
"CUDA_VISIBLE_DEVICES": "0,1",
165171
"PYTHONPATH": "/opt/ml/code",
166-
"HF_HOME": "/tmp/huggingface"
172+
"HF_HOME": "/tmp/huggingface",
167173
},
168174
"working_dir": "./training_code",
169-
"py_modules": ["utils", "models"]
175+
"py_modules": ["utils", "models"],
170176
}
171-
177+
172178
spec = RayJobSpec(
173179
entrypoint="python train_model.py --epochs 100",
174180
runtime_env=complex_runtime_env,
175181
entrypoint_num_gpus=2.0,
176-
entrypoint_memory=8192
182+
entrypoint_memory=8192,
177183
)
178-
184+
179185
job = RayJob(
180-
metadata={"name": "complex-env-job", "namespace": "ml-training"},
181-
spec=spec
186+
metadata={"name": "complex-env-job", "namespace": "ml-training"}, spec=spec
182187
)
183-
188+
184189
yaml_dict = job.to_dict()
185-
190+
186191
# Verify complex runtime_env is preserved exactly
187192
spec_runtime_env = yaml_dict["spec"]["runtime_env"]
188193
assert spec_runtime_env["pip"] == ["torch==1.13.0", "transformers", "datasets"]
189-
assert spec_runtime_env["conda"]["dependencies"] == ["python=3.9", "cudatoolkit=11.8"]
194+
assert spec_runtime_env["conda"]["dependencies"] == [
195+
"python=3.9",
196+
"cudatoolkit=11.8",
197+
]
190198
assert spec_runtime_env["env_vars"]["CUDA_VISIBLE_DEVICES"] == "0,1"
191199
assert spec_runtime_env["env_vars"]["PYTHONPATH"] == "/opt/ml/code"
192200
assert spec_runtime_env["env_vars"]["HF_HOME"] == "/tmp/huggingface"
193201
assert spec_runtime_env["working_dir"] == "./training_code"
194202
assert spec_runtime_env["py_modules"] == ["utils", "models"]
195-
203+
196204
# Verify other fields
197205
assert yaml_dict["spec"]["entrypoint"] == "python train_model.py --epochs 100"
198206
assert yaml_dict["spec"]["entrypoint_num_gpus"] == 2.0
@@ -203,23 +211,19 @@ def test_ray_job_yaml_different_statuses():
203211
"""Test RayJob YAML generation with different status values."""
204212
statuses_to_test = [
205213
RayJobStatus.PENDING,
206-
RayJobStatus.RUNNING,
214+
RayJobStatus.RUNNING,
207215
RayJobStatus.SUCCEEDED,
208216
RayJobStatus.FAILED,
209-
RayJobStatus.STOPPED
217+
RayJobStatus.STOPPED,
210218
]
211-
219+
212220
for status in statuses_to_test:
213-
spec = RayJobSpec(
214-
entrypoint=f"python job_{status.lower()}.py",
215-
status=status
216-
)
217-
221+
spec = RayJobSpec(entrypoint=f"python job_{status.lower()}.py", status=status)
222+
218223
job = RayJob(
219-
metadata={"name": f"job-{status.lower()}", "namespace": "test"},
220-
spec=spec
224+
metadata={"name": f"job-{status.lower()}", "namespace": "test"}, spec=spec
221225
)
222-
226+
223227
yaml_dict = job.to_dict()
224228
assert yaml_dict["status"]["status"] == status
225-
assert yaml_dict["spec"]["entrypoint"] == f"python job_{status.lower()}.py"
229+
assert yaml_dict["spec"]["entrypoint"] == f"python job_{status.lower()}.py"

0 commit comments

Comments
 (0)