async def test_job_copy_minimal(
self,
job,
mock_trainml,
):
requested_config = dict(name="job_copy", )
expected_kwargs = {
"type": "notebook",
"gpu_type": "1060-id",
"gpu_count": 1,
"disk_size": 10,
"worker_commands": None,
"workers": None,
"environment": None,
"data": None,
"source_job_uuid": "job-id-1",
}
api_response = specimen.Job(
mock_trainml,
**{
"customer_uuid": "cus-id-1",
"job_uuid": "job-id-2",
"name": "job_copy",
"type": "notebook",
"status": "new",
},
)
mock_trainml.jobs.create = AsyncMock(return_value=api_response)
new_job = await job.copy(**requested_config)
mock_trainml.jobs.create.assert_called_once_with(
"job_copy", **expected_kwargs)
assert new_job.id == "job-id-2"
async def test_job_wait_for_current_status(self, mock_trainml):
job = specimen.Job(
mock_trainml,
**{
"customer_uuid": "cus-id-1",
"job_uuid": "job-id-1",
"name": "test notebook",
"type": "notebook",
"status": "running",
},
)
api_response = None
mock_trainml._query = AsyncMock(return_value=api_response)
await job.wait_for("running")
mock_trainml._query.assert_not_called()
def test_job_default_ws_msg_handler_multiple_workers(
self, mock_trainml, capsys):
job = specimen.Job(
mock_trainml,
**{
"customer_uuid":
"cus-id-1",
"job_uuid":
"job-id-1",
"name":
"test notebook",
"type":
"notebook",
"status":
"running",
"workers": [
{
"rig_uuid": "rig-id-1",
"job_worker_uuid": "worker-id-1",
"command": "jupyter lab",
"status": "running",
},
{
"rig_uuid": "rig-id-1",
"job_worker_uuid": "worker-id-2",
"command": "jupyter lab",
"status": "running",
},
],
},
)
data = {
"msg": "Epoch (1/1000)\n",
"time": 1613079345318,
"type": "subscription",
"stream": "worker-id-1",
"job_worker_uuid": "worker-id-1",
}
handler = job._get_msg_handler(None)
handler(data)
captured = capsys.readouterr()
assert (captured.out ==
"02/11/2021, 15:35:45: Worker 1 - Epoch (1/1000)\n")
async def test_job_copy_training_job_failure(
self,
mock_trainml,
):
job = specimen.Job(
mock_trainml,
**{
"customer_uuid": "cus-id-1",
"job_uuid": "job-id-1",
"name": "test training",
"type": "training",
"status": "running",
},
)
mock_trainml.jobs.create = AsyncMock()
with raises(SpecificationError):
await job.copy("new job")
mock_trainml.jobs.create.assert_not_called()
async def test_job_attach_immediate_return(self, mock_trainml):
job_spec = {
"customer_uuid": "cus-id-1",
"job_uuid": "job-id-1",
"name": "test training",
"type": "training",
"status": "finished",
"workers": [],
}
job = specimen.Job(
mock_trainml,
**job_spec,
)
api_response = None
mock_trainml._ws_subscribe = AsyncMock(return_value=api_response)
job.refresh = AsyncMock(return_value=job_spec)
await job.attach()
mock_trainml._ws_subscribe.assert_not_called()
def test_job_get_connection_details_local_model_and_input_data(
self, mock_trainml):
job = specimen.Job(
mock_trainml,
**{
"customer_uuid": "cus-id-1",
"job_uuid": "job-id-1",
"name": "test notebook",
"type": "notebook",
"status": "new",
"model": {
"source_type": "local",
"source_uri": "~/model_dir"
},
"data": {
"datasets": [],
"input_type": "local",
"input_uri": "~/data_dir",
"status": "ready",
},
"vpn": {
"status": "new",
"cidr": "10.106.171.0/24",
"client": {
"port": "36017",
"id": "cus-id-1",
"address": "10.106.171.253",
"ssh_port": 46600,
},
},
},
)
details = job.get_connection_details()
expected_details = dict(
cidr="10.106.171.0/24",
ssh_port=46600,
model_path="~/model_dir",
input_path="~/data_dir",
output_path=None,
)
assert details == expected_details
async def test_job_attach(self, mock_trainml):
job_spec = {
"customer_uuid":
"cus-id-1",
"job_uuid":
"job-id-1",
"name":
"test training",
"type":
"training",
"status":
"running",
"workers": [
{
"rig_uuid": "rig-id-1",
"job_worker_uuid": "worker-id-1",
"command": "jupyter lab",
"status": "running",
},
{
"rig_uuid": "rig-id-1",
"job_worker_uuid": "worker-id-2",
"command": "jupyter lab",
"status": "running",
},
],
}
job = specimen.Job(
mock_trainml,
**job_spec,
)
api_response = None
mock_trainml._ws_subscribe = AsyncMock(return_value=api_response)
job.refresh = AsyncMock(return_value=job_spec)
await job.attach()
mock_trainml._ws_subscribe.assert_called_once()
def test_job_bool(self, job, mock_trainml):
empty_job = specimen.Job(mock_trainml)
assert bool(job)
assert not bool(empty_job)
def job(mock_trainml):
yield specimen.Job(
mock_trainml,
**{
"customer_uuid":
"cus-id-1",
"job_uuid":
"job-id-1",
"name":
"test notebook",
"start":
"2021-02-11T15:46:22.455Z",
"type":
"notebook",
"status":
"new",
"credits_per_hour":
0.1,
"credits":
0.1007,
"workers": [{
"rig_uuid": "rig-id-1",
"job_worker_uuid": "worker-id-1",
"command": "jupyter lab",
"status": "new",
}],
"worker_status":
"new",
"resources": {
"gpu_count": 1,
"gpu_type_id": "1060-id",
"disk_size": 10,
},
"model": {
"size": 7176192,
"git_uri": "[email protected]:trainML/test-private.git",
"status": "new",
},
"data": {
"datasets": [
{
"dataset_uuid": "data-id-1",
"name": "first one",
"type": "public",
"size": 184549376,
},
{
"dataset_uuid": "data-id-2",
"name": "second one",
"type": "public",
"size": 5068061409,
},
],
"status":
"ready",
},
"environment": {
"type":
"DEEPLEARNING_PY38",
"image_size":
44966989795,
"env": [
{
"value": "env1val",
"key": "env1"
},
{
"value": "env2val",
"key": "env2"
},
],
"worker_key_types": ["aws", "gcp"],
"status":
"new",
},
"vpn": {
"status": "new",
"cidr": "10.106.171.0/24",
"client": {
"port": "36017",
"id": "cus-id-1",
"address": "10.106.171.253",
},
"net_prefix_type_id": 1,
},
"nb_token":
"token",
},
)
