wehub-resource-sync
303 lines
11 KiB
Python
303 lines
11 KiB
Python
import asyncio
|
|
import json
|
|
from dataclasses import asdict
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
from google.protobuf.json_format import Parse
|
|
|
|
from ray.core.generated.gcs_pb2 import JobsAPIInfo
|
|
from ray.dashboard.modules.job.common import (
|
|
JobErrorType,
|
|
JobInfo,
|
|
JobInfoStorageClient,
|
|
JobStatus,
|
|
JobSubmitRequest,
|
|
http_uri_components_to_uri,
|
|
uri_to_http_components,
|
|
validate_request_type,
|
|
)
|
|
|
|
|
|
class TestJobSubmitRequestValidation:
|
|
def test_validate_entrypoint(self):
|
|
r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
|
|
assert r.entrypoint == "abc"
|
|
|
|
with pytest.raises(TypeError, match="required positional argument"):
|
|
validate_request_type({}, JobSubmitRequest)
|
|
|
|
with pytest.raises(TypeError, match="must be a string"):
|
|
validate_request_type({"entrypoint": 123}, JobSubmitRequest)
|
|
|
|
def test_validate_submission_id(self):
|
|
r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
|
|
assert r.entrypoint == "abc"
|
|
assert r.submission_id is None
|
|
|
|
r = validate_request_type(
|
|
{"entrypoint": "abc", "submission_id": "123"}, JobSubmitRequest
|
|
)
|
|
assert r.entrypoint == "abc"
|
|
assert r.submission_id == "123"
|
|
|
|
with pytest.raises(TypeError, match="must be a string"):
|
|
validate_request_type(
|
|
{"entrypoint": 123, "submission_id": 1}, JobSubmitRequest
|
|
)
|
|
|
|
def test_validate_runtime_env(self):
|
|
r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
|
|
assert r.entrypoint == "abc"
|
|
assert r.runtime_env is None
|
|
|
|
r = validate_request_type(
|
|
{"entrypoint": "abc", "runtime_env": {"hi": "hi2"}}, JobSubmitRequest
|
|
)
|
|
assert r.entrypoint == "abc"
|
|
assert r.runtime_env == {"hi": "hi2"}
|
|
|
|
with pytest.raises(TypeError, match="must be a dict"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "runtime_env": 123}, JobSubmitRequest
|
|
)
|
|
|
|
with pytest.raises(TypeError, match="keys must be strings"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "runtime_env": {1: "hi"}}, JobSubmitRequest
|
|
)
|
|
|
|
def test_validate_metadata(self):
|
|
r = validate_request_type({"entrypoint": "abc"}, JobSubmitRequest)
|
|
assert r.entrypoint == "abc"
|
|
assert r.metadata is None
|
|
|
|
r = validate_request_type(
|
|
{"entrypoint": "abc", "metadata": {"hi": "hi2"}}, JobSubmitRequest
|
|
)
|
|
assert r.entrypoint == "abc"
|
|
assert r.metadata == {"hi": "hi2"}
|
|
|
|
with pytest.raises(TypeError, match="must be a dict"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "metadata": 123}, JobSubmitRequest
|
|
)
|
|
|
|
with pytest.raises(TypeError, match="keys must be strings"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "metadata": {1: "hi"}}, JobSubmitRequest
|
|
)
|
|
|
|
with pytest.raises(TypeError, match="values must be strings"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "metadata": {"hi": 1}}, JobSubmitRequest
|
|
)
|
|
|
|
def test_validate_entrypoint_label_selector(self):
|
|
r = validate_request_type(
|
|
{
|
|
"entrypoint": "abc",
|
|
"entrypoint_label_selector": {"fragile_node": "!1"},
|
|
},
|
|
JobSubmitRequest,
|
|
)
|
|
assert r.entrypoint_label_selector == {"fragile_node": "!1"}
|
|
|
|
with pytest.raises(TypeError, match="must be a dict"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "entrypoint_label_selector": "bad"},
|
|
JobSubmitRequest,
|
|
)
|
|
|
|
with pytest.raises(TypeError, match="keys must be strings"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "entrypoint_label_selector": {1: "bad"}},
|
|
JobSubmitRequest,
|
|
)
|
|
|
|
with pytest.raises(TypeError, match="values must be strings"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "entrypoint_label_selector": {"k": 1}},
|
|
JobSubmitRequest,
|
|
)
|
|
|
|
def test_entrypoint_resources_disallow_strings(self):
|
|
with pytest.raises(TypeError, match="values must be numbers"):
|
|
validate_request_type(
|
|
{"entrypoint": "abc", "entrypoint_resources": {"Custom": "1"}},
|
|
JobSubmitRequest,
|
|
)
|
|
|
|
|
|
def test_uri_to_http_and_back():
|
|
assert uri_to_http_components("gcs://hello.zip") == ("gcs", "hello.zip")
|
|
assert uri_to_http_components("gcs://hello.whl") == ("gcs", "hello.whl")
|
|
|
|
with pytest.raises(ValueError, match="'blah' is not a valid Protocol"):
|
|
uri_to_http_components("blah://halb.zip")
|
|
|
|
with pytest.raises(ValueError, match="does not end in .zip or .whl"):
|
|
assert uri_to_http_components("gcs://hello.not_zip")
|
|
|
|
with pytest.raises(ValueError, match="does not end in .zip or .whl"):
|
|
assert uri_to_http_components("gcs://hello")
|
|
|
|
assert http_uri_components_to_uri("gcs", "hello.zip") == "gcs://hello.zip"
|
|
assert http_uri_components_to_uri("blah", "halb.zip") == "blah://halb.zip"
|
|
assert http_uri_components_to_uri("blah", "halb.whl") == "blah://halb.whl"
|
|
|
|
for original_uri in ["gcs://hello.zip", "gcs://fasdf.whl"]:
|
|
new_uri = http_uri_components_to_uri(*uri_to_http_components(original_uri))
|
|
assert new_uri == original_uri
|
|
|
|
|
|
def test_dynamic_status_message():
|
|
info = JobInfo(
|
|
status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_num_cpus=1
|
|
)
|
|
assert "may be waiting for resources" in info.message
|
|
|
|
info = JobInfo(
|
|
status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_num_gpus=1
|
|
)
|
|
assert "may be waiting for resources" in info.message
|
|
|
|
info = JobInfo(status=JobStatus.PENDING, entrypoint="echo hi", entrypoint_memory=4)
|
|
assert "may be waiting for resources" in info.message
|
|
|
|
info = JobInfo(
|
|
status=JobStatus.PENDING,
|
|
entrypoint="echo hi",
|
|
entrypoint_resources={"Custom": 1},
|
|
)
|
|
assert "may be waiting for resources" in info.message
|
|
|
|
info = JobInfo(
|
|
status=JobStatus.PENDING, entrypoint="echo hi", runtime_env={"conda": "env"}
|
|
)
|
|
assert "may be waiting for the runtime environment" in info.message
|
|
|
|
|
|
def test_job_info_to_json():
|
|
info = JobInfo(
|
|
status=JobStatus.PENDING,
|
|
entrypoint="echo hi",
|
|
entrypoint_num_cpus=1,
|
|
entrypoint_num_gpus=1,
|
|
entrypoint_memory=4,
|
|
entrypoint_resources={"Custom": 1},
|
|
runtime_env={"pip": ["pkg"]},
|
|
)
|
|
expected_items = {
|
|
"status": "PENDING",
|
|
"message": (
|
|
"Job has not started yet. It may be waiting for resources "
|
|
"(CPUs, GPUs, memory, custom resources) to become available. "
|
|
"It may be waiting for the runtime environment to be set up."
|
|
),
|
|
"entrypoint": "echo hi",
|
|
"entrypoint_num_cpus": 1,
|
|
"entrypoint_num_gpus": 1,
|
|
"entrypoint_memory": 4,
|
|
"entrypoint_resources": {"Custom": 1},
|
|
"runtime_env_json": '{"pip": ["pkg"]}',
|
|
}
|
|
|
|
# Check that the expected items are in the JSON.
|
|
assert expected_items.items() <= info.to_json().items()
|
|
|
|
new_job_info = JobInfo.from_json(info.to_json())
|
|
assert new_job_info == info
|
|
|
|
# If `status` is just a string, then operations like status.is_terminal()
|
|
# would fail, so we should make sure that it's a JobStatus.
|
|
assert isinstance(new_job_info.status, JobStatus)
|
|
|
|
|
|
def test_job_info_json_to_proto():
|
|
"""Test that JobInfo JSON can be converted to JobsAPIInfo protobuf."""
|
|
info = JobInfo(
|
|
status=JobStatus.PENDING,
|
|
entrypoint="echo hi",
|
|
error_type=JobErrorType.JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE,
|
|
start_time=123,
|
|
end_time=456,
|
|
metadata={"hi": "hi2"},
|
|
entrypoint_num_cpus=1,
|
|
entrypoint_num_gpus=1,
|
|
entrypoint_memory=4,
|
|
entrypoint_resources={"Custom": 1},
|
|
runtime_env={"pip": ["pkg"]},
|
|
driver_agent_http_address="http://localhost:1234",
|
|
driver_node_id="node_id",
|
|
)
|
|
info_json = json.dumps(info.to_json())
|
|
info_proto = Parse(info_json, JobsAPIInfo())
|
|
assert info_proto.status == "PENDING"
|
|
assert info_proto.entrypoint == "echo hi"
|
|
assert info_proto.start_time == 123
|
|
assert info_proto.end_time == 456
|
|
assert info_proto.metadata == {"hi": "hi2"}
|
|
assert info_proto.entrypoint_num_cpus == 1
|
|
assert info_proto.entrypoint_num_gpus == 1
|
|
assert info_proto.entrypoint_memory == 4
|
|
assert info_proto.entrypoint_resources == {"Custom": 1}
|
|
assert info_proto.runtime_env_json == '{"pip": ["pkg"]}'
|
|
assert info_proto.message == (
|
|
"Job has not started yet. It may be waiting for resources "
|
|
"(CPUs, GPUs, memory, custom resources) to become available. "
|
|
"It may be waiting for the runtime environment to be set up."
|
|
)
|
|
assert info_proto.error_type == "JOB_SUPERVISOR_ACTOR_UNSCHEDULABLE"
|
|
assert info_proto.driver_agent_http_address == "http://localhost:1234"
|
|
assert info_proto.driver_node_id == "node_id"
|
|
|
|
minimal_info = JobInfo(status=JobStatus.PENDING, entrypoint="echo hi")
|
|
minimal_info_json = json.dumps(minimal_info.to_json())
|
|
minimal_info_proto = Parse(minimal_info_json, JobsAPIInfo())
|
|
assert minimal_info_proto.status == "PENDING"
|
|
assert minimal_info_proto.entrypoint == "echo hi"
|
|
for unset_optional_field in [
|
|
"entrypoint_num_cpus",
|
|
"entrypoint_num_gpus",
|
|
"entrypoint_memory",
|
|
"runtime_env_json",
|
|
"error_type",
|
|
"driver_agent_http_address",
|
|
"driver_node_id",
|
|
]:
|
|
assert not minimal_info_proto.HasField(unset_optional_field)
|
|
|
|
|
|
def test_get_all_jobs_filters_out_none_job_info():
|
|
prefix = JobInfoStorageClient.JOB_DATA_KEY_PREFIX
|
|
mock_gcs = MagicMock()
|
|
mock_gcs.async_internal_kv_keys = AsyncMock(
|
|
return_value=[
|
|
(prefix + "job1").encode(),
|
|
(prefix + "job2").encode(),
|
|
]
|
|
)
|
|
|
|
storage = JobInfoStorageClient(mock_gcs)
|
|
job_info_1 = JobInfo(status=JobStatus.RUNNING, entrypoint="echo 1")
|
|
|
|
async def mock_get_info(job_id, timeout=30):
|
|
if job_id == "job1":
|
|
return job_info_1
|
|
return None
|
|
|
|
storage.get_info = mock_get_info
|
|
|
|
result = asyncio.run(storage.get_all_jobs())
|
|
|
|
assert result == {"job1": job_info_1}
|
|
for job_id, job_info in result.items():
|
|
asdict(job_info) # This should not raise an exception
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|