apache-airflow-providers-cncf-kubernetes 3.1.0__py3-none-any.whl → 10.10.0rc1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (72) hide show
  1. airflow/providers/cncf/kubernetes/__init__.py +18 -23
  2. airflow/providers/cncf/kubernetes/backcompat/__init__.py +17 -0
  3. airflow/providers/cncf/kubernetes/backcompat/backwards_compat_converters.py +31 -49
  4. airflow/providers/cncf/kubernetes/callbacks.py +200 -0
  5. airflow/providers/cncf/kubernetes/cli/__init__.py +16 -0
  6. airflow/providers/cncf/kubernetes/cli/kubernetes_command.py +195 -0
  7. airflow/providers/cncf/kubernetes/decorators/kubernetes.py +163 -0
  8. airflow/providers/cncf/kubernetes/decorators/kubernetes_cmd.py +118 -0
  9. airflow/providers/cncf/kubernetes/exceptions.py +37 -0
  10. airflow/providers/cncf/kubernetes/executors/__init__.py +17 -0
  11. airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +831 -0
  12. airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py +91 -0
  13. airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py +736 -0
  14. airflow/providers/cncf/kubernetes/executors/local_kubernetes_executor.py +306 -0
  15. airflow/providers/cncf/kubernetes/get_provider_info.py +249 -50
  16. airflow/providers/cncf/kubernetes/hooks/kubernetes.py +846 -112
  17. airflow/providers/cncf/kubernetes/k8s_model.py +62 -0
  18. airflow/providers/cncf/kubernetes/kube_client.py +156 -0
  19. airflow/providers/cncf/kubernetes/kube_config.py +125 -0
  20. airflow/providers/cncf/kubernetes/kubernetes_executor_templates/__init__.py +16 -0
  21. airflow/providers/cncf/kubernetes/kubernetes_executor_templates/basic_template.yaml +79 -0
  22. airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py +165 -0
  23. airflow/providers/cncf/kubernetes/operators/custom_object_launcher.py +368 -0
  24. airflow/providers/cncf/kubernetes/operators/job.py +646 -0
  25. airflow/providers/cncf/kubernetes/operators/kueue.py +132 -0
  26. airflow/providers/cncf/kubernetes/operators/pod.py +1417 -0
  27. airflow/providers/cncf/kubernetes/operators/resource.py +191 -0
  28. airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py +336 -35
  29. airflow/providers/cncf/kubernetes/pod_generator.py +592 -0
  30. airflow/providers/cncf/kubernetes/pod_template_file_examples/__init__.py +16 -0
  31. airflow/providers/cncf/kubernetes/pod_template_file_examples/dags_in_image_template.yaml +68 -0
  32. airflow/providers/cncf/kubernetes/pod_template_file_examples/dags_in_volume_template.yaml +74 -0
  33. airflow/providers/cncf/kubernetes/pod_template_file_examples/git_sync_template.yaml +95 -0
  34. airflow/providers/cncf/kubernetes/python_kubernetes_script.jinja2 +51 -0
  35. airflow/providers/cncf/kubernetes/python_kubernetes_script.py +92 -0
  36. airflow/providers/cncf/kubernetes/resource_convert/__init__.py +16 -0
  37. airflow/providers/cncf/kubernetes/resource_convert/configmap.py +52 -0
  38. airflow/providers/cncf/kubernetes/resource_convert/env_variable.py +39 -0
  39. airflow/providers/cncf/kubernetes/resource_convert/secret.py +40 -0
  40. airflow/providers/cncf/kubernetes/secret.py +128 -0
  41. airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py +30 -14
  42. airflow/providers/cncf/kubernetes/template_rendering.py +81 -0
  43. airflow/providers/cncf/kubernetes/triggers/__init__.py +16 -0
  44. airflow/providers/cncf/kubernetes/triggers/job.py +176 -0
  45. airflow/providers/cncf/kubernetes/triggers/pod.py +344 -0
  46. airflow/providers/cncf/kubernetes/utils/__init__.py +3 -0
  47. airflow/providers/cncf/kubernetes/utils/container.py +118 -0
  48. airflow/providers/cncf/kubernetes/utils/delete_from.py +154 -0
  49. airflow/providers/cncf/kubernetes/utils/k8s_resource_iterator.py +46 -0
  50. airflow/providers/cncf/kubernetes/utils/pod_manager.py +887 -152
  51. airflow/providers/cncf/kubernetes/utils/xcom_sidecar.py +25 -16
  52. airflow/providers/cncf/kubernetes/version_compat.py +38 -0
  53. apache_airflow_providers_cncf_kubernetes-10.10.0rc1.dist-info/METADATA +125 -0
  54. apache_airflow_providers_cncf_kubernetes-10.10.0rc1.dist-info/RECORD +62 -0
  55. {apache_airflow_providers_cncf_kubernetes-3.1.0.dist-info → apache_airflow_providers_cncf_kubernetes-10.10.0rc1.dist-info}/WHEEL +1 -2
  56. apache_airflow_providers_cncf_kubernetes-10.10.0rc1.dist-info/entry_points.txt +3 -0
  57. apache_airflow_providers_cncf_kubernetes-10.10.0rc1.dist-info/licenses/NOTICE +5 -0
  58. airflow/providers/cncf/kubernetes/backcompat/pod.py +0 -119
  59. airflow/providers/cncf/kubernetes/backcompat/pod_runtime_info_env.py +0 -56
  60. airflow/providers/cncf/kubernetes/backcompat/volume.py +0 -62
  61. airflow/providers/cncf/kubernetes/backcompat/volume_mount.py +0 -58
  62. airflow/providers/cncf/kubernetes/example_dags/example_kubernetes.py +0 -163
  63. airflow/providers/cncf/kubernetes/example_dags/example_spark_kubernetes.py +0 -66
  64. airflow/providers/cncf/kubernetes/example_dags/example_spark_kubernetes_spark_pi.yaml +0 -57
  65. airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py +0 -622
  66. apache_airflow_providers_cncf_kubernetes-3.1.0.dist-info/METADATA +0 -452
  67. apache_airflow_providers_cncf_kubernetes-3.1.0.dist-info/NOTICE +0 -6
  68. apache_airflow_providers_cncf_kubernetes-3.1.0.dist-info/RECORD +0 -29
  69. apache_airflow_providers_cncf_kubernetes-3.1.0.dist-info/entry_points.txt +0 -3
  70. apache_airflow_providers_cncf_kubernetes-3.1.0.dist-info/top_level.txt +0 -1
  71. /airflow/providers/cncf/kubernetes/{example_dags → decorators}/__init__.py +0 -0
  72. {apache_airflow_providers_cncf_kubernetes-3.1.0.dist-info → apache_airflow_providers_cncf_kubernetes-10.10.0rc1.dist-info/licenses}/LICENSE +0 -0
@@ -0,0 +1,646 @@
1
+ # Licensed to the Apache Software Foundation (ASF) under one
2
+ # or more contributor license agreements. See the NOTICE file
3
+ # distributed with this work for additional information
4
+ # regarding copyright ownership. The ASF licenses this file
5
+ # to you under the Apache License, Version 2.0 (the
6
+ # "License"); you may not use this file except in compliance
7
+ # with the License. You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing,
12
+ # software distributed under the License is distributed on an
13
+ # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
+ # KIND, either express or implied. See the License for the
15
+ # specific language governing permissions and limitations
16
+ # under the License.
17
+ """Executes a Kubernetes Job."""
18
+
19
+ from __future__ import annotations
20
+
21
+ import copy
22
+ import json
23
+ import logging
24
+ import os
25
+ import warnings
26
+ from collections.abc import Sequence
27
+ from functools import cached_property
28
+ from typing import TYPE_CHECKING, Any, Literal
29
+
30
+ from kubernetes.client import BatchV1Api, models as k8s
31
+ from kubernetes.client.api_client import ApiClient
32
+ from kubernetes.client.rest import ApiException
33
+
34
+ from airflow.configuration import conf
35
+ from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
36
+ from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
37
+ from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
38
+ add_unique_suffix,
39
+ create_unique_id,
40
+ )
41
+ from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
42
+ from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator, merge_objects
43
+ from airflow.providers.cncf.kubernetes.triggers.job import KubernetesJobTrigger
44
+ from airflow.providers.cncf.kubernetes.utils.pod_manager import EMPTY_XCOM_RESULT, PodNotFoundException
45
+ from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_1_PLUS
46
+
47
+ if AIRFLOW_V_3_1_PLUS:
48
+ from airflow.sdk import BaseOperator
49
+ else:
50
+ from airflow.models import BaseOperator
51
+ from airflow.utils import yaml
52
+ from airflow.utils.context import Context
53
+
54
+ if TYPE_CHECKING:
55
+ from airflow.utils.context import Context
56
+
57
+ log = logging.getLogger(__name__)
58
+
59
+
60
+ class KubernetesJobOperator(KubernetesPodOperator):
61
+ """
62
+ Executes a Kubernetes Job.
63
+
64
+ .. seealso::
65
+ For more information on how to use this operator, take a look at the guide:
66
+ :ref:`howto/operator:KubernetesJobOperator`
67
+
68
+ .. note::
69
+ If you use `Google Kubernetes Engine <https://cloud.google.com/kubernetes-engine/>`__
70
+ and Airflow is not running in the same cluster, consider using
71
+ :class:`~airflow.providers.google.cloud.operators.kubernetes_engine.GKEStartJobOperator`, which
72
+ simplifies the authorization process.
73
+
74
+ :param job_template_file: path to job template file (templated)
75
+ :param full_job_spec: The complete JodSpec
76
+ :param backoff_limit: Specifies the number of retries before marking this job failed. Defaults to 6
77
+ :param completion_mode: CompletionMode specifies how Pod completions are tracked. It can be `NonIndexed` (default) or `Indexed`.
78
+ :param completions: Specifies the desired number of successfully finished pods the job should be run with.
79
+ :param manual_selector: manualSelector controls generation of pod labels and pod selectors.
80
+ :param parallelism: Specifies the maximum desired number of pods the job should run at any given time.
81
+ :param selector: The selector of this V1JobSpec.
82
+ :param suspend: Suspend specifies whether the Job controller should create Pods or not.
83
+ :param ttl_seconds_after_finished: ttlSecondsAfterFinished limits the lifetime of a Job that has finished execution (either Complete or Failed).
84
+ :param wait_until_job_complete: Whether to wait until started job finished execution (either Complete or
85
+ Failed). Default is False.
86
+ :param job_poll_interval: Interval in seconds between polling the job status. Default is 10.
87
+ Used if the parameter `wait_until_job_complete` set True.
88
+ :param deferrable: Run operator in the deferrable mode. Note that the parameter
89
+ `wait_until_job_complete` must be set True.
90
+ :param on_kill_propagation_policy: Whether and how garbage collection will be performed. Default is 'Foreground'.
91
+ Acceptable values are:
92
+ 'Orphan' - orphan the dependents;
93
+ 'Background' - allow the garbage collector to delete the dependents in the background;
94
+ 'Foreground' - a cascading policy that deletes all dependents in the foreground.
95
+ Default value is 'Foreground'.
96
+ :param discover_pods_retry_number: Number of time list_namespaced_pod will be performed to discover
97
+ already running pods.
98
+ :param unwrap_single: Unwrap single result from the pod. For example, when set to `True` - if the XCom
99
+ result should be `['res']`, the final result would be `'res'`. Default is True to support backward
100
+ compatibility.
101
+ """
102
+
103
+ template_fields: Sequence[str] = tuple({"job_template_file"} | set(KubernetesPodOperator.template_fields))
104
+
105
+ def __init__(
106
+ self,
107
+ *,
108
+ job_template_file: str | None = None,
109
+ full_job_spec: k8s.V1Job | None = None,
110
+ backoff_limit: int | None = None,
111
+ completion_mode: str | None = None,
112
+ completions: int | None = None,
113
+ manual_selector: bool | None = None,
114
+ parallelism: int | None = None,
115
+ selector: k8s.V1LabelSelector | None = None,
116
+ suspend: bool | None = None,
117
+ ttl_seconds_after_finished: int | None = None,
118
+ wait_until_job_complete: bool = False,
119
+ job_poll_interval: float = 10,
120
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
121
+ on_kill_propagation_policy: Literal["Foreground", "Background", "Orphan"] = "Foreground",
122
+ discover_pods_retry_number: int = 3,
123
+ unwrap_single: bool = True,
124
+ **kwargs,
125
+ ) -> None:
126
+ self._pod = None
127
+ super().__init__(**kwargs)
128
+ self.job_template_file = job_template_file
129
+ self.full_job_spec = full_job_spec
130
+ self.job_request_obj: k8s.V1Job | None = None
131
+ self.job: k8s.V1Job | None = None
132
+ self.backoff_limit = backoff_limit
133
+ self.completion_mode = completion_mode
134
+ self.completions = completions
135
+ self.manual_selector = manual_selector
136
+ self.parallelism = parallelism
137
+ self.selector = selector
138
+ self.suspend = suspend
139
+ self.ttl_seconds_after_finished = ttl_seconds_after_finished
140
+ self.wait_until_job_complete = wait_until_job_complete
141
+ self.job_poll_interval = job_poll_interval
142
+ self.deferrable = deferrable
143
+ self.on_kill_propagation_policy = on_kill_propagation_policy
144
+ self.discover_pods_retry_number = discover_pods_retry_number
145
+ self.unwrap_single = unwrap_single
146
+
147
+ @property
148
+ def pod(self):
149
+ warnings.warn(
150
+ "`pod` parameter is deprecated, please use `pods`",
151
+ AirflowProviderDeprecationWarning,
152
+ stacklevel=2,
153
+ )
154
+ return self.pods[0] if self.pods else None
155
+
156
+ @pod.setter
157
+ def pod(self, value):
158
+ self._pod = value
159
+
160
+ @cached_property
161
+ def _incluster_namespace(self):
162
+ from pathlib import Path
163
+
164
+ path = Path("/var/run/secrets/kubernetes.io/serviceaccount/namespace")
165
+ return path.exists() and path.read_text() or None
166
+
167
+ @cached_property
168
+ def hook(self) -> KubernetesHook:
169
+ hook = KubernetesHook(
170
+ conn_id=self.kubernetes_conn_id,
171
+ in_cluster=self.in_cluster,
172
+ config_file=self.config_file,
173
+ cluster_context=self.cluster_context,
174
+ )
175
+ return hook
176
+
177
+ @cached_property
178
+ def job_client(self) -> BatchV1Api:
179
+ return self.hook.batch_v1_client
180
+
181
+ def create_job(self, job_request_obj: k8s.V1Job) -> k8s.V1Job:
182
+ self.log.debug("Starting job:\n%s", yaml.safe_dump(job_request_obj.to_dict()))
183
+ self.hook.create_job(job=job_request_obj)
184
+
185
+ return job_request_obj
186
+
187
+ def execute(self, context: Context):
188
+ if self.deferrable and not self.wait_until_job_complete:
189
+ self.log.warning(
190
+ "Deferrable mode is available only with parameter `wait_until_job_complete=True`. "
191
+ "Please, set it up."
192
+ )
193
+ if (self.get_logs or self.do_xcom_push) and not self.wait_until_job_complete:
194
+ self.log.warning(
195
+ "Getting Logs and pushing to XCom are available only with parameter `wait_until_job_complete=True`. "
196
+ "Please, set it up."
197
+ )
198
+ self.job_request_obj = self.build_job_request_obj(context)
199
+ self.job = self.create_job( # must set `self.job` for `on_kill`
200
+ job_request_obj=self.job_request_obj
201
+ )
202
+
203
+ ti = context["ti"]
204
+ ti.xcom_push(key="job_name", value=self.job.metadata.name)
205
+ ti.xcom_push(key="job_namespace", value=self.job.metadata.namespace)
206
+
207
+ self.pods: Sequence[k8s.V1Pod] | None = None
208
+ if self.parallelism is None and self.pod is None:
209
+ self.pods = [
210
+ self.get_or_create_pod(
211
+ pod_request_obj=self.pod_request_obj,
212
+ context=context,
213
+ )
214
+ ]
215
+ else:
216
+ self.pods = self.get_pods(pod_request_obj=self.pod_request_obj, context=context)
217
+
218
+ if self.wait_until_job_complete and self.deferrable:
219
+ self.execute_deferrable()
220
+ return
221
+
222
+ if self.wait_until_job_complete:
223
+ if self.do_xcom_push:
224
+ xcom_result = []
225
+ for pod in self.pods:
226
+ self.pod_manager.await_container_completion(
227
+ pod=pod, container_name=self.base_container_name
228
+ )
229
+ self.pod_manager.await_xcom_sidecar_container_start(pod=pod)
230
+ xcom_result.append(self.extract_xcom(pod=pod))
231
+ self.job = self.hook.wait_until_job_complete(
232
+ job_name=self.job.metadata.name,
233
+ namespace=self.job.metadata.namespace,
234
+ job_poll_interval=self.job_poll_interval,
235
+ )
236
+ if self.get_logs:
237
+ for pod in self.pods:
238
+ self.pod_manager.fetch_requested_container_logs(
239
+ pod=pod,
240
+ containers=self.container_logs,
241
+ follow_logs=True,
242
+ )
243
+
244
+ ti.xcom_push(key="job", value=self.job.to_dict())
245
+ if self.wait_until_job_complete:
246
+ if error_message := self.hook.is_job_failed(job=self.job):
247
+ raise AirflowException(
248
+ f"Kubernetes job '{self.job.metadata.name}' is failed with error '{error_message}'"
249
+ )
250
+ if self.do_xcom_push:
251
+ return xcom_result
252
+
253
+ def execute_deferrable(self):
254
+ self.defer(
255
+ trigger=KubernetesJobTrigger(
256
+ job_name=self.job.metadata.name,
257
+ job_namespace=self.job.metadata.namespace,
258
+ pod_names=[pod.metadata.name for pod in self.pods],
259
+ pod_namespace=self.pods[0].metadata.namespace,
260
+ base_container_name=self.base_container_name,
261
+ kubernetes_conn_id=self.kubernetes_conn_id,
262
+ cluster_context=self.cluster_context,
263
+ config_file=self.config_file,
264
+ in_cluster=self.in_cluster,
265
+ poll_interval=self.job_poll_interval,
266
+ get_logs=self.get_logs,
267
+ do_xcom_push=self.do_xcom_push,
268
+ ),
269
+ method_name="execute_complete",
270
+ )
271
+
272
+ def execute_complete(self, context: Context, event: dict, **kwargs):
273
+ ti = context["ti"]
274
+ ti.xcom_push(key="job", value=event["job"])
275
+ if event["status"] == "error":
276
+ raise AirflowException(event["message"])
277
+
278
+ if self.get_logs:
279
+ for pod_name in event["pod_names"]:
280
+ pod_namespace = event["pod_namespace"]
281
+ pod = self.hook.get_pod(pod_name, pod_namespace)
282
+ if not pod:
283
+ raise PodNotFoundException("Could not find pod after resuming from deferral")
284
+ self._write_logs(pod)
285
+
286
+ if self.do_xcom_push:
287
+ xcom_results: list[Any | None] = []
288
+ for xcom_result in event["xcom_result"]:
289
+ if isinstance(xcom_result, str) and xcom_result.rstrip() == EMPTY_XCOM_RESULT:
290
+ self.log.info("xcom result file is empty.")
291
+ xcom_results.append(None)
292
+ continue
293
+ self.log.info("xcom result: \n%s", xcom_result)
294
+ xcom_results.append(json.loads(xcom_result))
295
+ return xcom_results[0] if self.unwrap_single and len(xcom_results) == 1 else xcom_results
296
+
297
+ @staticmethod
298
+ def deserialize_job_template_file(path: str) -> k8s.V1Job:
299
+ """
300
+ Generate a Job from a file.
301
+
302
+ Unfortunately we need access to the private method
303
+ ``_ApiClient__deserialize_model`` from the kubernetes client.
304
+ This issue is tracked here: https://github.com/kubernetes-client/python/issues/977.
305
+
306
+ :param path: Path to the file
307
+ :return: a kubernetes.client.models.V1Job
308
+ """
309
+ if os.path.exists(path):
310
+ with open(path) as stream:
311
+ job = yaml.safe_load(stream)
312
+ else:
313
+ job = None
314
+ log.warning("Template file %s does not exist", path)
315
+
316
+ api_client = ApiClient()
317
+ return api_client._ApiClient__deserialize_model(job, k8s.V1Job)
318
+
319
+ def on_kill(self) -> None:
320
+ if self.job:
321
+ job = self.job
322
+ kwargs = {
323
+ "name": job.metadata.name,
324
+ "namespace": job.metadata.namespace,
325
+ "propagation_policy": self.on_kill_propagation_policy,
326
+ }
327
+ if self.termination_grace_period is not None:
328
+ kwargs.update(grace_period_seconds=self.termination_grace_period)
329
+ self.job_client.delete_namespaced_job(**kwargs)
330
+
331
+ def build_job_request_obj(self, context: Context | None = None) -> k8s.V1Job:
332
+ """
333
+ Return V1Job object based on job template file, full job spec, and other operator parameters.
334
+
335
+ The V1Job attributes are derived (in order of precedence) from operator params, full job spec, job
336
+ template file.
337
+ """
338
+ self.log.debug("Creating job for KubernetesJobOperator task %s", self.task_id)
339
+ if self.job_template_file:
340
+ self.log.debug("Job template file found, will parse for base job")
341
+ job_template = self.deserialize_job_template_file(self.job_template_file)
342
+ if self.full_job_spec:
343
+ job_template = self.reconcile_jobs(job_template, self.full_job_spec)
344
+ elif self.full_job_spec:
345
+ job_template = self.full_job_spec
346
+ else:
347
+ job_template = k8s.V1Job(metadata=k8s.V1ObjectMeta())
348
+
349
+ pod_template = super().build_pod_request_obj(context)
350
+ pod_template_spec = k8s.V1PodTemplateSpec(
351
+ metadata=pod_template.metadata,
352
+ spec=pod_template.spec,
353
+ )
354
+ self.pod_request_obj = pod_template
355
+
356
+ job = k8s.V1Job(
357
+ api_version="batch/v1",
358
+ kind="Job",
359
+ metadata=k8s.V1ObjectMeta(
360
+ namespace=self.namespace,
361
+ labels=self.labels,
362
+ name=self.name,
363
+ annotations=self.annotations,
364
+ ),
365
+ spec=k8s.V1JobSpec(
366
+ active_deadline_seconds=self.active_deadline_seconds,
367
+ backoff_limit=self.backoff_limit,
368
+ completion_mode=self.completion_mode,
369
+ completions=self.completions,
370
+ manual_selector=self.manual_selector,
371
+ parallelism=self.parallelism,
372
+ selector=self.selector,
373
+ suspend=self.suspend,
374
+ template=pod_template_spec,
375
+ ttl_seconds_after_finished=self.ttl_seconds_after_finished,
376
+ ),
377
+ )
378
+
379
+ job = self.reconcile_jobs(job_template, job)
380
+
381
+ if not job.metadata.name:
382
+ job.metadata.name = create_unique_id(
383
+ task_id=self.task_id, unique=self.random_name_suffix, max_length=80
384
+ )
385
+ elif self.random_name_suffix:
386
+ # user has supplied job name, we're just adding suffix
387
+ job.metadata.name = add_unique_suffix(name=job.metadata.name)
388
+
389
+ job.metadata.name = f"job-{job.metadata.name}"
390
+
391
+ if not job.metadata.namespace:
392
+ hook_namespace = self.hook.get_namespace()
393
+ job_namespace = self.namespace or hook_namespace or self._incluster_namespace or "default"
394
+ job.metadata.namespace = job_namespace
395
+
396
+ self.log.info("Building job %s ", job.metadata.name)
397
+
398
+ return job
399
+
400
+ @staticmethod
401
+ def reconcile_jobs(base_job: k8s.V1Job, client_job: k8s.V1Job | None) -> k8s.V1Job:
402
+ """
403
+ Merge Kubernetes Job objects.
404
+
405
+ :param base_job: has the base attributes which are overwritten if they exist
406
+ in the client job and remain if they do not exist in the client_job
407
+ :param client_job: the job that the client wants to create.
408
+ :return: the merged jobs
409
+
410
+ This can't be done recursively as certain fields are overwritten and some are concatenated.
411
+ """
412
+ if client_job is None:
413
+ return base_job
414
+
415
+ client_job_cp = copy.deepcopy(client_job)
416
+ client_job_cp.spec = KubernetesJobOperator.reconcile_job_specs(base_job.spec, client_job_cp.spec)
417
+ client_job_cp.metadata = PodGenerator.reconcile_metadata(base_job.metadata, client_job_cp.metadata)
418
+ client_job_cp = merge_objects(base_job, client_job_cp)
419
+
420
+ return client_job_cp
421
+
422
+ @staticmethod
423
+ def reconcile_job_specs(
424
+ base_spec: k8s.V1JobSpec | None, client_spec: k8s.V1JobSpec | None
425
+ ) -> k8s.V1JobSpec | None:
426
+ """
427
+ Merge Kubernetes JobSpec objects.
428
+
429
+ :param base_spec: has the base attributes which are overwritten if they exist
430
+ in the client_spec and remain if they do not exist in the client_spec
431
+ :param client_spec: the spec that the client wants to create.
432
+ :return: the merged specs
433
+ """
434
+ if base_spec and not client_spec:
435
+ return base_spec
436
+ if not base_spec and client_spec:
437
+ return client_spec
438
+ if client_spec and base_spec:
439
+ client_spec.template.spec = PodGenerator.reconcile_specs(
440
+ base_spec.template.spec, client_spec.template.spec
441
+ )
442
+ client_spec.template.metadata = PodGenerator.reconcile_metadata(
443
+ base_spec.template.metadata, client_spec.template.metadata
444
+ )
445
+ return merge_objects(base_spec, client_spec)
446
+
447
+ return None
448
+
449
+ def get_pods(
450
+ self, pod_request_obj: k8s.V1Pod, context: Context, *, exclude_checked: bool = True
451
+ ) -> Sequence[k8s.V1Pod]:
452
+ """Return an already-running pods if exists."""
453
+ label_selector = self._build_find_pod_label_selector(context, exclude_checked=exclude_checked)
454
+ pod_list: Sequence[k8s.V1Pod] = []
455
+ retry_number: int = 0
456
+
457
+ while len(pod_list) != self.parallelism or retry_number <= self.discover_pods_retry_number:
458
+ pod_list = self.client.list_namespaced_pod(
459
+ namespace=pod_request_obj.metadata.namespace,
460
+ label_selector=label_selector,
461
+ ).items
462
+ retry_number += 1
463
+
464
+ if len(pod_list) == 0:
465
+ raise AirflowException(f"No pods running with labels {label_selector}")
466
+
467
+ for pod_instance in pod_list:
468
+ self.log_matching_pod(pod=pod_instance, context=context)
469
+
470
+ return pod_list
471
+
472
+
473
+ class KubernetesDeleteJobOperator(BaseOperator):
474
+ """
475
+ Delete a Kubernetes Job.
476
+
477
+ .. seealso::
478
+ For more information on how to use this operator, take a look at the guide:
479
+ :ref:`howto/operator:KubernetesDeleteJobOperator`
480
+
481
+ :param name: name of the Job.
482
+ :param namespace: the namespace to run within kubernetes.
483
+ :param kubernetes_conn_id: The :ref:`kubernetes connection id <howto/connection:kubernetes>`
484
+ for the Kubernetes cluster.
485
+ :param config_file: The path to the Kubernetes config file. (templated)
486
+ If not specified, default value is ``~/.kube/config``
487
+ :param in_cluster: run kubernetes client with in_cluster configuration.
488
+ :param cluster_context: context that points to kubernetes cluster.
489
+ Ignored when in_cluster is True. If None, current-context is used. (templated)
490
+ :param delete_on_status: Condition for performing delete operation depending on the job status. Values:
491
+ ``None`` - delete the job regardless of its status, "Complete" - delete only successfully completed
492
+ jobs, "Failed" - delete only failed jobs. (default: ``None``)
493
+ :param wait_for_completion: Whether to wait for the job to complete. (default: ``False``)
494
+ :param poll_interval: Interval in seconds between polling the job status. Used when the `delete_on_status`
495
+ parameter is set. (default: 10.0)
496
+ """
497
+
498
+ template_fields: Sequence[str] = (
499
+ "config_file",
500
+ "name",
501
+ "namespace",
502
+ "cluster_context",
503
+ )
504
+
505
+ def __init__(
506
+ self,
507
+ *,
508
+ name: str,
509
+ namespace: str,
510
+ kubernetes_conn_id: str | None = KubernetesHook.default_conn_name,
511
+ config_file: str | None = None,
512
+ in_cluster: bool | None = None,
513
+ cluster_context: str | None = None,
514
+ delete_on_status: str | None = None,
515
+ wait_for_completion: bool = False,
516
+ poll_interval: float = 10.0,
517
+ **kwargs,
518
+ ) -> None:
519
+ super().__init__(**kwargs)
520
+ self.name = name
521
+ self.namespace = namespace
522
+ self.kubernetes_conn_id = kubernetes_conn_id
523
+ self.config_file = config_file
524
+ self.in_cluster = in_cluster
525
+ self.cluster_context = cluster_context
526
+ self.delete_on_status = delete_on_status
527
+ self.wait_for_completion = wait_for_completion
528
+ self.poll_interval = poll_interval
529
+
530
+ @cached_property
531
+ def hook(self) -> KubernetesHook:
532
+ return KubernetesHook(
533
+ conn_id=self.kubernetes_conn_id,
534
+ in_cluster=self.in_cluster,
535
+ config_file=self.config_file,
536
+ cluster_context=self.cluster_context,
537
+ )
538
+
539
+ @cached_property
540
+ def client(self) -> BatchV1Api:
541
+ return self.hook.batch_v1_client
542
+
543
+ def execute(self, context: Context):
544
+ try:
545
+ if self.delete_on_status not in ("Complete", "Failed", None):
546
+ raise AirflowException(
547
+ "The `delete_on_status` parameter must be one of 'Complete', 'Failed' or None. "
548
+ "The current value is %s",
549
+ str(self.delete_on_status),
550
+ )
551
+
552
+ if self.wait_for_completion:
553
+ job = self.hook.wait_until_job_complete(
554
+ job_name=self.name, namespace=self.namespace, job_poll_interval=self.poll_interval
555
+ )
556
+ else:
557
+ job = self.hook.get_job_status(job_name=self.name, namespace=self.namespace)
558
+
559
+ if (
560
+ self.delete_on_status is None
561
+ or (self.delete_on_status == "Complete" and self.hook.is_job_successful(job=job))
562
+ or (self.delete_on_status == "Failed" and self.hook.is_job_failed(job=job))
563
+ ):
564
+ self.log.info("Deleting kubernetes Job: %s", self.name)
565
+ self.client.delete_namespaced_job(name=self.name, namespace=self.namespace)
566
+ self.log.info("Kubernetes job was deleted.")
567
+ else:
568
+ self.log.info(
569
+ "Deletion of the job %s was skipped due to settings of on_status=%s",
570
+ self.name,
571
+ self.delete_on_status,
572
+ )
573
+ except ApiException as e:
574
+ if e.status == 404:
575
+ self.log.info("The Kubernetes job %s does not exist.", self.name)
576
+ else:
577
+ raise e
578
+
579
+
580
+ class KubernetesPatchJobOperator(BaseOperator):
581
+ """
582
+ Update a Kubernetes Job.
583
+
584
+ .. seealso::
585
+ For more information on how to use this operator, take a look at the guide:
586
+ :ref:`howto/operator:KubernetesPatchJobOperator`
587
+
588
+ :param name: name of the Job
589
+ :param namespace: the namespace to run within kubernetes
590
+ :param body: Job json object with parameters for update
591
+ https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.25/#job-v1-batch
592
+ e.g. ``{"spec": {"suspend": True}}``
593
+ :param kubernetes_conn_id: The :ref:`kubernetes connection id <howto/connection:kubernetes>`
594
+ for the Kubernetes cluster.
595
+ :param config_file: The path to the Kubernetes config file. (templated)
596
+ If not specified, default value is ``~/.kube/config``
597
+ :param in_cluster: run kubernetes client with in_cluster configuration.
598
+ :param cluster_context: context that points to kubernetes cluster.
599
+ Ignored when in_cluster is True. If None, current-context is used. (templated)
600
+ """
601
+
602
+ template_fields: Sequence[str] = (
603
+ "config_file",
604
+ "name",
605
+ "namespace",
606
+ "body",
607
+ "cluster_context",
608
+ )
609
+
610
+ def __init__(
611
+ self,
612
+ *,
613
+ name: str,
614
+ namespace: str,
615
+ body: object,
616
+ kubernetes_conn_id: str | None = KubernetesHook.default_conn_name,
617
+ config_file: str | None = None,
618
+ in_cluster: bool | None = None,
619
+ cluster_context: str | None = None,
620
+ **kwargs,
621
+ ) -> None:
622
+ super().__init__(**kwargs)
623
+ self.name = name
624
+ self.namespace = namespace
625
+ self.body = body
626
+ self.kubernetes_conn_id = kubernetes_conn_id
627
+ self.config_file = config_file
628
+ self.in_cluster = in_cluster
629
+ self.cluster_context = cluster_context
630
+
631
+ @cached_property
632
+ def hook(self) -> KubernetesHook:
633
+ return KubernetesHook(
634
+ conn_id=self.kubernetes_conn_id,
635
+ in_cluster=self.in_cluster,
636
+ config_file=self.config_file,
637
+ cluster_context=self.cluster_context,
638
+ )
639
+
640
+ def execute(self, context: Context) -> dict:
641
+ self.log.info("Updating existing Job: %s", self.name)
642
+ job_object = self.hook.patch_namespaced_job(
643
+ job_name=self.name, namespace=self.namespace, body=self.body
644
+ )
645
+ self.log.info("Job was updated.")
646
+ return k8s.V1Job.to_dict(job_object)