def test_k8s_agent_manage_jobs_delete_jobs(monkeypatch, cloud_api): job_mock = MagicMock() job_mock.metadata.labels = { "prefect.io/identifier": "id", "prefect.io/flow_run_id": "fr", } job_mock.metadata.name = "my_job" job_mock.status.failed = True job_mock.status.succeeded = True batch_client = MagicMock() list_job = MagicMock() list_job.metadata._continue = 0 list_job.items = [job_mock] batch_client.list_namespaced_job.return_value = list_job batch_client.delete_namespaced_job.return_value = None monkeypatch.setattr( "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client) ) pod = MagicMock() pod.metadata.name = "pod_name" pod.status.phase = "Success" core_client = MagicMock() list_pods = MagicMock() list_pods.items = [pod] core_client.list_namespaced_pod.return_value = list_pods monkeypatch.setattr( "kubernetes.client.CoreV1Api", MagicMock(return_value=core_client) ) agent = KubernetesAgent() agent.manage_jobs() assert batch_client.delete_namespaced_job.called
def test_k8s_agent_manage_jobs_client_call(monkeypatch, cloud_api): gql_return = MagicMock(return_value=MagicMock(data=MagicMock( set_flow_run_state=None))) client = MagicMock() client.return_value.graphql = gql_return monkeypatch.setattr("prefect.agent.agent.Client", client) job_mock = MagicMock() job_mock.metadata.labels = { "prefect.io/identifier": "id", "prefect.io/flow_run_id": "fr", } job_mock.metadata.name = "my_job" job_mock.status.failed = False job_mock.status.succeeded = False list_job = MagicMock() list_job.metadata._continue = 0 list_job.items = [job_mock] pod = MagicMock() pod.metadata.name = "pod_name" c_status = MagicMock() c_status.state.waiting.reason = "ErrImagePull" pod.status.container_statuses = [c_status] list_pods = MagicMock() list_pods.items = [pod] agent = KubernetesAgent() agent.batch_client.list_namespaced_job.return_value = list_job agent.core_client.list_namespaced_pod.return_value = list_pods agent.manage_jobs()
def test_k8s_agent_manage_jobs_reports_failed_pods(monkeypatch, cloud_api): gql_return = MagicMock( return_value=MagicMock( data=MagicMock( set_flow_run_state=None, write_run_logs=None, get_flow_run_state=prefect.engine.state.Success(), ) ) ) client = MagicMock() client.return_value.graphql = gql_return monkeypatch.setattr("prefect.agent.agent.Client", client) job_mock = MagicMock() job_mock.metadata.labels = { "prefect.io/identifier": "id", "prefect.io/flow_run_id": "fr", } job_mock.metadata.name = "my_job" job_mock.status.failed = True job_mock.status.succeeded = False batch_client = MagicMock() list_job = MagicMock() list_job.metadata._continue = 0 list_job.items = [job_mock] batch_client.list_namespaced_job.return_value = list_job monkeypatch.setattr( "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client) ) pod = MagicMock() pod.metadata.name = "pod_name" pod.status.phase = "Failed" terminated = MagicMock() terminated.exit_code = "code" terminated.message = "message" terminated.reason = "reason" terminated.signal = "signal" c_status = MagicMock() c_status.state.terminated = terminated pod.status.container_statuses = [c_status] pod2 = MagicMock() pod2.metadata.name = "pod_name" pod2.status.phase = "Success" core_client = MagicMock() list_pods = MagicMock() list_pods.items = [pod, pod2] core_client.list_namespaced_pod.return_value = list_pods monkeypatch.setattr( "kubernetes.client.CoreV1Api", MagicMock(return_value=core_client) ) agent = KubernetesAgent() agent.manage_jobs() assert core_client.list_namespaced_pod.called
def test_k8s_agent_manage_pending_pods(monkeypatch, cloud_api): gql_return = MagicMock( return_value=MagicMock( data=MagicMock(set_flow_run_state=None, write_run_logs=None) ) ) client = MagicMock() client.return_value.graphql = gql_return monkeypatch.setattr("prefect.agent.agent.Client", client) job_mock = MagicMock() job_mock.metadata.labels = { "prefect.io/identifier": "id", "prefect.io/flow_run_id": "fr", } job_mock.metadata.name = "my_job" job_mock.status.failed = False job_mock.status.succeeded = False batch_client = MagicMock() list_job = MagicMock() list_job.metadata._continue = 0 list_job.items = [job_mock] batch_client.list_namespaced_job.return_value = list_job monkeypatch.setattr( "kubernetes.client.BatchV1Api", MagicMock(return_value=batch_client) ) dt = pendulum.now() pod = MagicMock() pod.metadata.name = "pod_name" pod.status.phase = "Pending" event = MagicMock() event.last_timestamp = dt event.reason = "reason" event.message = "message" core_client = MagicMock() list_pods = MagicMock() list_pods.items = [pod] list_events = MagicMock() list_events.items = [event] core_client.list_namespaced_pod.return_value = list_pods core_client.list_namespaced_event.return_value = list_events monkeypatch.setattr( "kubernetes.client.CoreV1Api", MagicMock(return_value=core_client) ) agent = KubernetesAgent() agent.manage_jobs() assert agent.job_pod_event_timestamps["my_job"]["pod_name"] == dt
def test_k8s_agent_manage_jobs_event_logging_handles_bad_timestamps( monkeypatch, cloud_api, caplog): agent = KubernetesAgent() agent.client = MagicMock() # Only run once agent.batch_client.list_namespaced_job().metadata._continue = 0 # List a job job = MagicMock() # Only incomplete job logs are displayed job.status.failed = job.status.succeeded = False agent.batch_client.list_namespaced_job().items = [job] # List a pod pod = MagicMock() pod.status.phase = "Pending" # Logs are displayed for 'Pending' pods agent.core_client.list_namespaced_pod().items = [pod, MagicMock()] @dataclass class Event: last_timestamp: datetime.datetime = None reason: str = "reason" message: str = "message" event_without_timestamp = Event(message="no timestamp attr") event_without_timestamp.__dict__.pop("last_timestamp") agent.core_client.list_namespaced_event.return_value.items = [ Event(last_timestamp=pendulum.now(), message="working log"), Event(message="null timestamp"), event_without_timestamp, ] agent.logger = MagicMock() agent.manage_jobs() # One event log was streamed agent.client.write_run_logs.assert_called_once() assert "working log" in agent.client.write_run_logs.call_args[0][0][0][ "message"] debug_logs = [call[0][0] for call in agent.logger.debug.call_args_list] assert any( re.match( "Encountered K8s event .* with no timestamp: .*null timestamp", l) for l in debug_logs) assert any( re.match( "Encountered K8s event .* with no timestamp: .*no timestamp attr", l) for l in debug_logs)
def test_k8s_agent_manage_jobs_reports_empty_status(monkeypatch, cloud_api): gql_return = MagicMock(return_value=MagicMock(data=MagicMock( set_flow_run_state=None, write_run_logs=None, get_flow_run_state=prefect.engine.state.Success(), ))) client = MagicMock() client.return_value.graphql = gql_return monkeypatch.setattr("prefect.agent.agent.Client", client) job_mock = MagicMock() job_mock.metadata.labels = { "prefect.io/identifier": "id", "prefect.io/flow_run_id": "fr", } job_mock.metadata.name = "my_job" job_mock.status.failed = True job_mock.status.succeeded = False list_job = MagicMock() list_job.metadata._continue = 0 list_job.items = [job_mock] pod = MagicMock() pod.metadata.name = "pod_name" pod.status.phase = "Failed" pod.status.container_statuses = None pod2 = MagicMock() pod2.metadata.name = "pod_name" pod2.status.phase = "Success" list_pods = MagicMock() list_pods.items = [pod, pod2] agent = KubernetesAgent() agent.batch_client.list_namespaced_job.return_value = list_job agent.core_client.list_namespaced_pod.return_value = list_pods agent.manage_jobs() assert agent.core_client.list_namespaced_pod.called