31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123 | class SparkOnK8STrigger(BaseTrigger):
"""Watch a Spark application on Kubernetes.
startup_deadline is an absolute Unix timestamp, preserved when a trigger restarts.
It only limits the Pending phase; None disables the startup timeout.
"""
def __init__(
self,
*,
driver_pod_name: str,
namespace: str = "default",
kubernetes_conn_id: str = "kubernetes_default",
poll_interval: int = 10,
startup_deadline: float | None = None,
):
super().__init__()
self.driver_pod_name = driver_pod_name
self.namespace = namespace
self.kubernetes_conn_id = kubernetes_conn_id
self.poll_interval = poll_interval
self.startup_deadline = startup_deadline
def serialize(self) -> tuple[str, dict[str, Any]]:
return (
"spark_on_k8s.airflow.triggers.SparkOnK8STrigger",
{
"driver_pod_name": self.driver_pod_name,
"namespace": self.namespace,
"kubernetes_conn_id": self.kubernetes_conn_id,
"poll_interval": self.poll_interval,
"startup_deadline": self.startup_deadline,
},
)
async def run(self) -> AsyncIterator[TriggerEvent]:
from spark_on_k8s.utils.async_app_manager import AsyncSparkAppManager, SparkAppStartupTimeoutError
from spark_on_k8s.utils.spark_app_status import SparkAppDeletedError, SparkAppStatus
try:
k8s_client_manager = _AirflowKubernetesAsyncClientManager(
kubernetes_conn_id=self.kubernetes_conn_id,
)
async_spark_app_manager = AsyncSparkAppManager(
k8s_client_manager=k8s_client_manager,
)
status = await async_spark_app_manager.wait_for_app(
namespace=self.namespace,
pod_name=self.driver_pod_name,
poll_interval=self.poll_interval,
startup_deadline=self.startup_deadline,
)
if not isinstance(status, SparkAppStatus):
# Only a manager that reports no final status needs this read; another read could only fail.
status = await async_spark_app_manager.app_status(
namespace=self.namespace,
pod_name=self.driver_pod_name,
)
yield TriggerEvent(
{
"namespace": self.namespace,
"pod_name": self.driver_pod_name,
"status": status,
}
)
except SparkAppDeletedError as e:
yield TriggerEvent(
{
"namespace": self.namespace,
"pod_name": self.driver_pod_name,
"status": "deleted",
"error": str(e),
}
)
except SparkAppStartupTimeoutError as e:
yield TriggerEvent(
{
"namespace": self.namespace,
"pod_name": self.driver_pod_name,
"status": "startup_timeout",
"error": str(e),
}
)
except Exception as e:
yield TriggerEvent(
{
"namespace": self.namespace,
"pod_name": self.driver_pod_name,
"status": "error",
"error": str(e),
"stacktrace": traceback.format_exc(),
}
)
|