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
