def test_submitted_run_get_status_succeeded(): mlflow_run_id = 1 job_name = "job-name" job_namespace = "job-namespace" condition = kubernetes.client.models.V1JobCondition(type="Complete", status="True") job_status = kubernetes.client.models.V1JobStatus( active=None, completion_time=None, conditions=[condition], failed=None, start_time=None, succeeded=1, ) job = kubernetes.client.models.V1Job(status=job_status) with mock.patch("kubernetes.client.BatchV1Api.read_namespaced_job_status" ) as kube_api_mock: kube_api_mock.return_value = job submitted_run = kb.KubernetesSubmittedRun(mlflow_run_id, job_name, job_namespace) print("status", submitted_run.get_status()) assert RunStatus.FINISHED == submitted_run.get_status() assert kube_api_mock.call_count == 1 args = kube_api_mock.call_args_list assert args[0][1]["name"] == job_name assert args[0][1]["namespace"] == job_namespace
def test_state_transitions(): mlflow_run_id = 1 job_name = "job-name" job_namespace = "job-namespace" submitted_run = kb.KubernetesSubmittedRun(mlflow_run_id, job_name, job_namespace) with mock.patch("kubernetes.client.BatchV1Api.read_namespaced_job_status") as kube_api_mock: def set_return_value(**kwargs): job_status = kubernetes.client.models.V1JobStatus(**kwargs) kube_api_mock.return_value = kubernetes.client.models.V1Job(status=job_status) set_return_value() assert RunStatus.SCHEDULED == submitted_run.get_status() set_return_value(start_time=1) assert RunStatus.RUNNING == submitted_run.get_status() set_return_value(start_time=1, failed=1) assert RunStatus.RUNNING == submitted_run.get_status() set_return_value(start_time=1, failed=1) assert RunStatus.RUNNING == submitted_run.get_status() set_return_value(start_time=1, failed=1, active=1) assert RunStatus.RUNNING == submitted_run.get_status() set_return_value(start_time=1, failed=1, succeeded=1) assert RunStatus.RUNNING == submitted_run.get_status() set_return_value(start_time=1, failed=1, succeeded=1, completion_time=2) assert RunStatus.RUNNING == submitted_run.get_status() condition = kubernetes.client.models.V1JobCondition(type="Complete", status="True") set_return_value( conditions=[condition], failed=1, start_time=1, completion_time=2, succeeded=1 ) assert RunStatus.FINISHED == submitted_run.get_status()
def test_submitted_run_get_status_killed(): mlflow_run_id = 1 job_name = "job-name" job_namespace = "job-namespace" with mock.patch("kubernetes.client.BatchV1Api.delete_namespaced_job") as kube_api_mock: submitted_run = kb.KubernetesSubmittedRun(mlflow_run_id, job_name, job_namespace) submitted_run.cancel() assert RunStatus.KILLED == submitted_run.get_status() assert kube_api_mock.call_count == 1 args = kube_api_mock.call_args_list assert args[0][1]["name"] == job_name assert args[0][1]["namespace"] == job_namespace
def test_submitted_run_get_status_running(): mlflow_run_id = 1 job_name = "job-name" job_namespace = "job-namespace" job_status = kubernetes.client.models.V1JobStatus( active=1, completion_time=None, conditions=None, failed=1, start_time=1, succeeded=1 ) job = kubernetes.client.models.V1Job(status=job_status) with mock.patch("kubernetes.client.BatchV1Api.read_namespaced_job_status") as kube_api_mock: kube_api_mock.return_value = job submitted_run = kb.KubernetesSubmittedRun(mlflow_run_id, job_name, job_namespace) assert RunStatus.RUNNING == submitted_run.get_status() assert kube_api_mock.call_count == 1 args = kube_api_mock.call_args_list assert args[0][1]["name"] == job_name assert args[0][1]["namespace"] == job_namespace