def test_adopt_launched_task(self, mock_kube_client):
executor = self.kubernetes_executor
executor.scheduler_job_id = "modified"
annotations = {
'dag_id': 'dag',
'execution_date': datetime.utcnow().isoformat(),
'task_id': 'task',
'try_number': '1',
}
ti_key = annotations_to_key(annotations)
pod = k8s.V1Pod(
metadata=k8s.V1ObjectMeta(name="foo",
labels={"airflow-worker": "bar"},
annotations=annotations))
pod_ids = {ti_key: {}}
executor.adopt_launched_task(mock_kube_client,
pod=pod,
pod_ids=pod_ids)
assert mock_kube_client.patch_namespaced_pod.call_args[1] == {
'body': {
'metadata': {
'labels': {
'airflow-worker': 'modified'
},
'annotations': annotations,
'name': 'foo',
}
},
'name': 'foo',
'namespace': None,
}
assert pod_ids == {}
assert executor.running == {ti_key}
def adopt_launched_task(self, kube_client: client.CoreV1Api,
pod: k8s.V1Pod, pod_ids: Dict[TaskInstanceKey,
k8s.V1Pod]) -> None:
"""
Patch existing pod so that the current KubernetesJobWatcher can monitor it via label selectors
:param kube_client: kubernetes client for speaking to kube API
:param pod: V1Pod spec that we will patch with new label
:param pod_ids: pod_ids we expect to patch.
"""
self.log.info("attempting to adopt pod %s", pod.metadata.name)
pod.metadata.labels[
'airflow-worker'] = pod_generator.make_safe_label_value(
str(self.scheduler_job_id))
pod_id = annotations_to_key(pod.metadata.annotations)
if pod_id not in pod_ids:
self.log.error(
"attempting to adopt taskinstance which was not specified by database: %s",
pod_id)
return
try:
kube_client.patch_namespaced_pod(
name=pod.metadata.name,
namespace=pod.metadata.namespace,
body=PodGenerator.serialize_pod(pod),
)
pod_ids.pop(pod_id)
self.running.add(pod_id)
except ApiException as e:
self.log.info("Failed to adopt pod %s. Reason: %s",
pod.metadata.name, e)
def process_watcher_task(self, task: KubernetesWatchType) -> None:
"""Process the task by watcher."""
pod_id, namespace, state, annotations, resource_version = task
self.log.info(
'Attempting to finish pod; pod_id: %s; state: %s; annotations: %s', pod_id, state, annotations
)
key = annotations_to_key(annotations=annotations)
if key:
self.log.debug('finishing job %s - %s (%s)', key, state, pod_id)
self.result_queue.put((key, state, pod_id, namespace, resource_version))
def test_try_adopt_task_instances(self, mock_adopt_completed_pods,
mock_adopt_launched_task):
executor = self.kubernetes_executor
executor.scheduler_job_id = "10"
ti_key = annotations_to_key({
'dag_id':
'dag',
'execution_date':
datetime.utcnow().isoformat(),
'task_id':
'task',
'try_number':
'1',
})
mock_ti = mock.MagicMock(queued_by_job_id="1",
external_executor_id="1",
key=ti_key)
pod = k8s.V1Pod(metadata=k8s.V1ObjectMeta(name="foo"))
mock_kube_client = mock.MagicMock()
mock_kube_client.list_namespaced_pod.return_value.items = [pod]
executor.kube_client = mock_kube_client
# First adoption
reset_tis = executor.try_adopt_task_instances([mock_ti])
mock_kube_client.list_namespaced_pod.assert_called_once_with(
namespace='default', label_selector='airflow-worker=1')
mock_adopt_launched_task.assert_called_once_with(
mock_kube_client, pod, {ti_key: mock_ti})
mock_adopt_completed_pods.assert_called_once()
assert reset_tis == [mock_ti
] # assume failure adopting when checking return
# Second adoption (queued_by_job_id and external_executor_id no longer match)
mock_kube_client.reset_mock()
mock_adopt_launched_task.reset_mock()
mock_adopt_completed_pods.reset_mock()
mock_ti.queued_by_job_id = "10" # scheduler_job would have updated this after the first adoption
executor.scheduler_job_id = "20"
# assume success adopting when checking return, `adopt_launched_task` pops `ti_key` from `pod_ids`
mock_adopt_launched_task.side_effect = lambda client, pod, pod_ids: pod_ids.pop(
ti_key)
reset_tis = executor.try_adopt_task_instances([mock_ti])
mock_kube_client.list_namespaced_pod.assert_called_once_with(
namespace='default', label_selector='airflow-worker=10')
mock_adopt_launched_task.assert_called_once(
) # Won't check args this time around as they get mutated
mock_adopt_completed_pods.assert_called_once()
assert reset_tis == [
] # This time our return is empty - no TIs to reset
