wehub-resource-sync
359 lines
13 KiB
Python
359 lines
13 KiB
Python
from unittest import mock
|
|
|
|
import kubernetes
|
|
import pytest
|
|
import yaml
|
|
from kubernetes.config.config_exception import ConfigException
|
|
|
|
from mlflow.entities import RunStatus
|
|
from mlflow.exceptions import ExecutionException
|
|
from mlflow.projects import kubernetes as kb
|
|
|
|
|
|
def test_run_command_creation():
|
|
command = [
|
|
"python train.py --alpha 0.5 --l1-ratio 0.1",
|
|
"--comment 'foo bar'",
|
|
'--comment-bis "bar foo"',
|
|
]
|
|
command = kb._get_run_command(command)
|
|
assert command == [
|
|
"python",
|
|
"train.py",
|
|
"--alpha",
|
|
"0.5",
|
|
"--l1-ratio",
|
|
"0.1",
|
|
"--comment",
|
|
"'foo bar'",
|
|
"--comment-bis",
|
|
"'bar foo'",
|
|
]
|
|
|
|
|
|
def test_valid_kubernetes_job_spec():
|
|
"""
|
|
Tests job specification for Kubernetes.
|
|
"""
|
|
custom_template = yaml.safe_load(
|
|
"apiVersion: batch/v1\n"
|
|
"kind: Job\n"
|
|
"metadata:\n"
|
|
" name: pi-with-ttl\n"
|
|
"spec:\n"
|
|
" ttlSecondsAfterFinished: 100\n"
|
|
" template:\n"
|
|
" spec:\n"
|
|
" containers:\n"
|
|
" - name: pi\n"
|
|
" image: perl\n"
|
|
" command: ['perl', '-Mbignum=bpi', '-wle']\n"
|
|
" env: \n"
|
|
" - name: DUMMY\n"
|
|
' value: "test_var"\n'
|
|
" restartPolicy: Never\n"
|
|
)
|
|
project_name = "mlflow-docker-example"
|
|
image_tag = "image_tag"
|
|
image_digest = "5e74a5a"
|
|
command = ["mlflow", "run", ".", "--env-manager", "local", "-P", "alpha=0.5"]
|
|
env_vars = {"RUN_ID": "1"}
|
|
job_definition = kb._get_kubernetes_job_definition(
|
|
project_name=project_name,
|
|
image_tag=image_tag,
|
|
image_digest=image_digest,
|
|
command=command,
|
|
env_vars=env_vars,
|
|
job_template=custom_template,
|
|
)
|
|
container_spec = job_definition["spec"]["template"]["spec"]["containers"][0]
|
|
assert container_spec["name"] == project_name
|
|
assert container_spec["image"] == image_tag + "@" + image_digest
|
|
assert container_spec["command"] == command
|
|
assert len(container_spec["env"]) == 2
|
|
assert container_spec["env"][0]["name"] == "DUMMY"
|
|
assert container_spec["env"][0]["value"] == "test_var"
|
|
assert container_spec["env"][1]["name"] == "RUN_ID"
|
|
assert container_spec["env"][1]["value"] == "1"
|
|
|
|
|
|
def test_run_kubernetes_job():
|
|
active_run = mock.Mock()
|
|
project_name = "mlflow-docker-example"
|
|
image_tag = "image_tag"
|
|
image_digest = "5e74a5a"
|
|
command = ["python train.py --alpha 0.5 --l1-ratio 0.1"]
|
|
env_vars = {"RUN_ID": "1"}
|
|
kube_context = "docker-for-desktop"
|
|
job_template = yaml.safe_load(
|
|
"apiVersion: batch/v1\n"
|
|
"kind: Job\n"
|
|
"metadata:\n"
|
|
" name: pi-with-ttl\n"
|
|
" namespace: mlflow\n"
|
|
"spec:\n"
|
|
" ttlSecondsAfterFinished: 100\n"
|
|
" template:\n"
|
|
" spec:\n"
|
|
" containers:\n"
|
|
" - name: pi\n"
|
|
" image: perl\n"
|
|
" command: ['perl', '-Mbignum=bpi', '-wle']\n"
|
|
" restartPolicy: Never\n"
|
|
)
|
|
with (
|
|
mock.patch("kubernetes.config.load_kube_config") as kube_config_mock,
|
|
mock.patch("kubernetes.client.BatchV1Api.create_namespaced_job") as kube_api_mock,
|
|
):
|
|
submitted_run_obj = kb.run_kubernetes_job(
|
|
project_name=project_name,
|
|
active_run=active_run,
|
|
image_tag=image_tag,
|
|
image_digest=image_digest,
|
|
command=command,
|
|
env_vars=env_vars,
|
|
job_template=job_template,
|
|
kube_context=kube_context,
|
|
)
|
|
|
|
assert submitted_run_obj._mlflow_run_id == active_run.info.run_id
|
|
assert submitted_run_obj._job_name.startswith(project_name)
|
|
assert submitted_run_obj._job_namespace == "mlflow"
|
|
assert kube_api_mock.call_count == 1
|
|
args = kube_config_mock.call_args_list
|
|
assert args[0][1]["context"] == kube_context
|
|
|
|
|
|
def test_run_kubernetes_job_current_kubecontext():
|
|
active_run = mock.Mock()
|
|
project_name = "mlflow-docker-example"
|
|
image_tag = "image_tag"
|
|
image_digest = "5e74a5a"
|
|
command = ["python train.py --alpha 0.5 --l1-ratio 0.1"]
|
|
env_vars = {"RUN_ID": "1"}
|
|
kube_context = None
|
|
|
|
job_template = yaml.safe_load(
|
|
"apiVersion: batch/v1\n"
|
|
"kind: Job\n"
|
|
"metadata:\n"
|
|
" name: pi-with-ttl\n"
|
|
" namespace: mlflow\n"
|
|
"spec:\n"
|
|
" ttlSecondsAfterFinished: 100\n"
|
|
" template:\n"
|
|
" spec:\n"
|
|
" containers:\n"
|
|
" - name: pi\n"
|
|
" image: perl\n"
|
|
" command: ['perl', '-Mbignum=bpi', '-wle']\n"
|
|
" restartPolicy: Never\n"
|
|
)
|
|
with (
|
|
mock.patch("kubernetes.config.load_kube_config") as kube_config_mock,
|
|
mock.patch("kubernetes.config.load_incluster_config") as incluster_kube_config_mock,
|
|
mock.patch("kubernetes.client.BatchV1Api.create_namespaced_job") as kube_api_mock,
|
|
):
|
|
submitted_run_obj = kb.run_kubernetes_job(
|
|
project_name=project_name,
|
|
active_run=active_run,
|
|
image_tag=image_tag,
|
|
image_digest=image_digest,
|
|
command=command,
|
|
env_vars=env_vars,
|
|
job_template=job_template,
|
|
kube_context=kube_context,
|
|
)
|
|
|
|
assert submitted_run_obj._mlflow_run_id == active_run.info.run_id
|
|
assert submitted_run_obj._job_name.startswith(project_name)
|
|
assert submitted_run_obj._job_namespace == "mlflow"
|
|
assert kube_api_mock.call_count == 1
|
|
assert kube_config_mock.call_count == 1
|
|
assert incluster_kube_config_mock.call_count == 0
|
|
|
|
|
|
def test_run_kubernetes_job_in_cluster():
|
|
active_run = mock.Mock()
|
|
project_name = "mlflow-docker-example"
|
|
image_tag = "image_tag"
|
|
image_digest = "5e74a5a"
|
|
command = ["python train.py --alpha 0.5 --l1-ratio 0.1"]
|
|
env_vars = {"RUN_ID": "1"}
|
|
kube_context = None
|
|
job_template = yaml.safe_load(
|
|
"apiVersion: batch/v1\n"
|
|
"kind: Job\n"
|
|
"metadata:\n"
|
|
" name: pi-with-ttl\n"
|
|
" namespace: mlflow\n"
|
|
"spec:\n"
|
|
" ttlSecondsAfterFinished: 100\n"
|
|
" template:\n"
|
|
" spec:\n"
|
|
" containers:\n"
|
|
" - name: pi\n"
|
|
" image: perl\n"
|
|
" command: ['perl', '-Mbignum=bpi', '-wle']\n"
|
|
" restartPolicy: Never\n"
|
|
)
|
|
with mock.patch("kubernetes.config.load_kube_config") as kube_config_mock:
|
|
kube_config_mock.side_effect = ConfigException()
|
|
with (
|
|
mock.patch("kubernetes.config.load_incluster_config") as incluster_kube_config_mock,
|
|
mock.patch("kubernetes.client.BatchV1Api.create_namespaced_job") as kube_api_mock,
|
|
):
|
|
submitted_run_obj = kb.run_kubernetes_job(
|
|
project_name=project_name,
|
|
active_run=active_run,
|
|
image_tag=image_tag,
|
|
image_digest=image_digest,
|
|
command=command,
|
|
env_vars=env_vars,
|
|
job_template=job_template,
|
|
kube_context=kube_context,
|
|
)
|
|
|
|
assert submitted_run_obj._mlflow_run_id == active_run.info.run_id
|
|
assert submitted_run_obj._job_name.startswith(project_name)
|
|
assert submitted_run_obj._job_namespace == "mlflow"
|
|
assert kube_api_mock.call_count == 1
|
|
assert kube_config_mock.call_count == 1
|
|
assert incluster_kube_config_mock.call_count == 1
|
|
|
|
|
|
def test_push_image_to_registry():
|
|
image_uri = "dockerhub_account/mlflow-kubernetes-example"
|
|
with mock.patch("docker.from_env") as docker_mock:
|
|
client = mock.MagicMock()
|
|
docker_mock.return_value = client
|
|
kb.push_image_to_registry(image_uri)
|
|
assert client.images.push.call_count == 1
|
|
args = client.images.push.call_args_list
|
|
assert args[0][1]["repository"] == image_uri
|
|
|
|
|
|
def test_push_image_to_registry_handling_errors():
|
|
image_uri = "dockerhub_account/mlflow-kubernetes-example"
|
|
with pytest.raises(
|
|
ExecutionException,
|
|
match="Error while pushing to docker registry: An image does not exist locally",
|
|
):
|
|
kb.push_image_to_registry(image_uri)
|
|
|
|
|
|
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_failed():
|
|
mlflow_run_id = 1
|
|
job_name = "job-name"
|
|
job_namespace = "job-namespace"
|
|
condition = kubernetes.client.models.V1JobCondition(type="Failed", status="True")
|
|
job_status = kubernetes.client.models.V1JobStatus(
|
|
active=1,
|
|
completion_time=None,
|
|
conditions=[condition],
|
|
failed=1,
|
|
start_time=1,
|
|
succeeded=None,
|
|
)
|
|
job = kubernetes.client.models.V1Job(status=job_status)
|
|
with mock.patch(
|
|
"kubernetes.client.BatchV1Api.read_namespaced_job_status", return_value=job
|
|
) as kube_api_mock:
|
|
submitted_run = kb.KubernetesSubmittedRun(mlflow_run_id, job_name, job_namespace)
|
|
assert RunStatus.FAILED == 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_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", return_value=job
|
|
) as kube_api_mock:
|
|
submitted_run = kb.KubernetesSubmittedRun(mlflow_run_id, job_name, job_namespace)
|
|
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_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", return_value=job
|
|
) as kube_api_mock:
|
|
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
|
|
|
|
|
|
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()
|