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
