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