# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Kubernetes sub-commands."""
from __future__ import annotations
import sys
from datetime import datetime, timedelta
from pathlib import Path
from kubernetes import client
from kubernetes.client.api_client import ApiClient
from kubernetes.client.rest import ApiException
from airflow.exceptions import UnknownExecutorException
from airflow.executors.executor_loader import ExecutorLoader
from airflow.models import DagModel, DagRun, TaskInstance
from airflow.providers.cncf.kubernetes import pod_generator
from airflow.providers.cncf.kubernetes.executors.kubernetes_executor import KubeConfig, KubernetesExecutor
from airflow.providers.cncf.kubernetes.kube_client import (
_get_executor_kube_client,
_team_kwargs,
get_kube_client,
)
from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import create_unique_id
from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator, generate_pod_command_args
from airflow.providers.cncf.kubernetes.version_compat import (
AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_1_PLUS,
AIRFLOW_V_3_2_PLUS,
)
from airflow.providers.common.compat.sdk import conf
from airflow.utils import cli as cli_utils, yaml
from airflow.utils.providers_configuration_loader import providers_configuration_loaded
from airflow.utils.types import DagRunType
try:
from airflow.serialization.serialized_objects import create_scheduler_operator
except ImportError:
[docs]
create_scheduler_operator = lambda t: t
if AIRFLOW_V_3_1_PLUS:
from airflow.utils.cli import get_bagged_dag
else:
from airflow.utils.cli import get_dag as get_bagged_dag # type: ignore[attr-defined,no-redef]
@cli_utils.action_cli
@providers_configuration_loaded
[docs]
def generate_pod_yaml(args):
"""Generate yaml files for each task in the DAG. Used for testing output of KubernetesExecutor."""
logical_date = args.logical_date if AIRFLOW_V_3_0_PLUS else args.execution_date
if AIRFLOW_V_3_0_PLUS:
dag = get_bagged_dag(bundle_names=args.bundle_name, dag_id=args.dag_id)
else:
dag = get_bagged_dag(subdir=args.subdir, dag_id=args.dag_id)
yaml_output_path = Path(args.output_path) / "airflow_yaml_output"
dm = DagModel(dag_id=dag.dag_id)
if AIRFLOW_V_3_0_PLUS:
dr = DagRun(dag.dag_id, logical_date=logical_date)
dr.run_id = DagRun.generate_run_id(
run_type=DagRunType.MANUAL, logical_date=logical_date, run_after=logical_date
)
dm.bundle_name = args.bundle_name if args.bundle_name else "default"
dm.relative_fileloc = dag.relative_fileloc
else:
dr = DagRun(dag.dag_id, execution_date=logical_date)
dr.run_id = DagRun.generate_run_id(run_type=DagRunType.MANUAL, execution_date=logical_date)
executor_conf = None
if AIRFLOW_V_3_2_PLUS and args.team:
from airflow.executors.base_executor import ExecutorConf
executor_conf = ExecutorConf(team_name=args.team)
kube_config = KubeConfig(executor_conf=executor_conf)
for task in dag.tasks:
if AIRFLOW_V_3_0_PLUS:
from uuid6 import uuid7
serialized_task = create_scheduler_operator(task)
ti = TaskInstance(serialized_task, run_id=dr.run_id, dag_version_id=uuid7())
else:
ti = TaskInstance(task, run_id=dr.run_id)
ti.dag_run = dr
ti.dag_model = dm
command_args = generate_pod_command_args(ti)
pod = PodGenerator.construct_pod(
dag_id=args.dag_id,
task_id=ti.task_id,
pod_id=create_unique_id(args.dag_id, ti.task_id),
try_number=ti.try_number,
kube_image=kube_config.kube_image,
date=ti.logical_date if AIRFLOW_V_3_0_PLUS else ti.execution_date,
args=command_args,
pod_override_object=PodGenerator.from_obj(ti.executor_config),
scheduler_job_id="worker-config",
namespace=kube_config.executor_namespace,
base_worker_pod=PodGenerator.deserialize_model_file(kube_config.pod_template_file),
with_mutation_hook=True,
)
api_client = ApiClient()
date_string = pod_generator.datetime_to_label_safe_datestring(logical_date)
yaml_file_name = f"{args.dag_id}_{ti.task_id}_{date_string}.yml"
yaml_output_path.mkdir(parents=True, exist_ok=True)
with open(yaml_output_path / yaml_file_name, "w") as output:
sanitized_pod = api_client.sanitize_for_serialization(pod)
output.write(yaml.dump(sanitized_pod))
print(f"YAML output can be found at {yaml_output_path}")
def _get_pod_completion_time(pod):
"""
Return the time the pod's last container finished.
Scans both ``container_statuses`` and ``init_container_statuses`` and returns the
latest ``finished_at`` timestamp. Falls back to the latest condition
``last_transition_time`` (which is updated at the terminal transition), and finally
to ``creation_timestamp`` as a last resort.
"""
statuses = [*(pod.status.container_statuses or []), *(pod.status.init_container_statuses or [])]
times = [
s.state.terminated.finished_at
for s in statuses
if s.state and s.state.terminated and s.state.terminated.finished_at
]
if times:
return max(times)
condition_times = [
c.last_transition_time for c in (pod.status.conditions or []) if c.last_transition_time
]
if condition_times:
return max(condition_times)
return pod.metadata.creation_timestamp
@cli_utils.action_cli(check_db=False)
@providers_configuration_loaded
[docs]
def cleanup_pods(args):
"""Clean up k8s pods in evicted/failed/succeeded/pending states."""
team_name = args.team
if team_name and not AIRFLOW_V_3_2_PLUS:
raise SystemExit("cleanup-pods --team requires Airflow 3.2+.")
if args.executor and not team_name:
raise SystemExit("cleanup-pods --executor requires --team.")
namespace = args.namespace or conf.get("kubernetes_executor", "namespace", **_team_kwargs(team_name))
if team_name:
executor_cls = _get_team_executor_cls(team_name, args.executor)
if not args.allow_shared_namespace:
_check_namespace_not_shared(team_name, namespace)
min_pending_minutes = args.min_pending_minutes
# protect newly created pods from deletion
if min_pending_minutes < 5:
min_pending_minutes = 5
min_completed_minutes = args.min_completed_minutes
# https://kubernetes.io/docs/concepts/workloads/pods/pod-lifecycle/
# All Containers in the Pod have terminated in success, and will not be restarted.
pod_succeeded = "succeeded"
# The Pod has been accepted by the Kubernetes cluster,
# but one or more of the containers has not been set up and made ready to run.
pod_pending = "pending"
# All Containers in the Pod have terminated, and at least one Container has terminated in failure.
# That is, the Container either exited with non-zero status or was terminated by the system.
pod_failed = "failed"
# https://kubernetes.io/docs/tasks/administer-cluster/out-of-resource/
pod_reason_evicted = "evicted"
# If pod is failed and restartPolicy is:
# * Always: Restart Container; Pod phase stays Running.
# * OnFailure: Restart Container; Pod phase stays Running.
# * Never: Pod phase becomes Failed.
pod_restart_policy_never = "never"
if args.verbose:
print("Loading Kubernetes configuration")
if team_name:
kube_client = _get_executor_kube_client(
team_name=team_name, executor_client_factory=executor_cls.client_factory
)
else:
kube_client = get_kube_client()
if args.verbose:
print(f"Listing pods in namespace {namespace}")
airflow_pod_labels = [
"dag_id",
"task_id",
"try_number",
"airflow_version",
]
list_kwargs = {"namespace": namespace, "limit": 500, "label_selector": ",".join(airflow_pod_labels)}
while True:
pod_list = kube_client.list_namespaced_pod(**list_kwargs)
for pod in pod_list.items:
pod_name = pod.metadata.name
if args.verbose:
print(f"Inspecting pod {pod_name}")
pod_phase = pod.status.phase.lower()
pod_reason = pod.status.reason.lower() if pod.status.reason else ""
pod_restart_policy = pod.spec.restart_policy.lower()
current_time = datetime.now(pod.metadata.creation_timestamp.tzinfo)
is_terminal = (
pod_phase == pod_succeeded
or (pod_phase == pod_failed and pod_restart_policy == pod_restart_policy_never)
or (pod_reason == pod_reason_evicted)
)
is_terminal_old_enough = is_terminal and (
min_completed_minutes == 0
or current_time - _get_pod_completion_time(pod) > timedelta(minutes=min_completed_minutes)
)
is_pending_too_long = (
pod_phase == pod_pending
and current_time - pod.metadata.creation_timestamp > timedelta(minutes=min_pending_minutes)
)
if is_terminal_old_enough or is_pending_too_long:
print(
f'Deleting pod "{pod_name}" phase "{pod_phase}" and reason "{pod_reason}", '
f'restart policy "{pod_restart_policy}"'
)
try:
_delete_pod(pod.metadata.name, namespace, kube_client)
except ApiException as e:
print(f"Can't remove POD: {e}", file=sys.stderr)
else:
if args.verbose:
print(f"No action taken on pod {pod_name}")
continue_token = pod_list.metadata._continue
if not continue_token:
break
list_kwargs["_continue"] = continue_token
def _get_team_executor_cls(team_name: str, executor: str | None) -> type[KubernetesExecutor]:
"""Return the class of the team's executor named ``executor``, or of its default one."""
# Only the team's own executors: a global one doesn't connect with the team's config.
names = [n for n in ExecutorLoader.get_executor_names(validate_teams=False) if n.team_name == team_name]
if not names:
raise SystemExit(f"Team {team_name!r} has no executors configured.")
try:
if hasattr(ExecutorLoader, "resolve_executor_name"):
name = ExecutorLoader.resolve_executor_name(executor, team_name, executor_names=names)
elif executor: # Airflow < 3.4
name = ExecutorLoader.lookup_executor_name_by_str(
executor, team_name=team_name, validate_teams=False
)
else:
name = names[0]
except UnknownExecutorException as e:
raise SystemExit(str(e))
executor_cls, _ = ExecutorLoader.import_executor_cls(name)
if not issubclass(executor_cls, KubernetesExecutor):
raise SystemExit(
f"Executor {name.module_path} of team {team_name!r} isn't a KubernetesExecutor. "
"Pass --executor to pick one that is."
)
return executor_cls
def _check_namespace_not_shared(team_name: str, namespace: str) -> None:
"""Refuse a namespace another scope's KubernetesExecutor also uses, since their pods can't be told apart."""
# Only scopes that run a KubernetesExecutor make pods; the others still have a default namespace.
scopes = {
n.team_name
for n in ExecutorLoader.get_executor_names(validate_teams=False)
if n.team_name != team_name
and issubclass(ExecutorLoader.import_executor_cls(n)[0], KubernetesExecutor)
}
for other in sorted(scopes, key=lambda t: t or ""):
if conf.get("kubernetes_executor", "namespace", **_team_kwargs(other)) == namespace:
owner = f"a KubernetesExecutor of team {other!r}" if other else "a global KubernetesExecutor"
raise SystemExit(
f"Namespace {namespace!r} of team {team_name!r} is also used by {owner}, so cleanup-pods "
"would delete its pods too. Pass --allow-shared-namespace if that executor runs on "
"another cluster."
)
def _delete_pod(name, namespace, kube_client):
"""
Delete a namespaced pod.
Helper Function for cleanup_pods.
"""
delete_options = client.V1DeleteOptions()
print(f'Deleting POD "{name}" from "{namespace}" namespace')
api_response = kube_client.delete_namespaced_pod(name=name, namespace=namespace, body=delete_options)
print(api_response)