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", }, )