Skip to content

fix spark-kubernetes-operator compatibility #31798

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
potiuk merged 1 commit into from
Jun 27, 2023
Merged

fix spark-kubernetes-operator compatibility #31798

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 63 additions & 26 deletions airflow/providers/cncf/kubernetes/operators/spark_kubernetes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,12 @@
# under the License.
from __future__ import annotations

import datetime
from typing import TYPE_CHECKING, Sequence

from kubernetes.watch import Watch

from airflow import AirflowException
from airflow.models import BaseOperator
from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook, _load_body_to_dict

Expand All @@ -43,6 +45,7 @@ class SparkKubernetesOperator(BaseOperator):
for the to Kubernetes cluster.
:param api_group: kubernetes api group of sparkApplication
:param api_version: kubernetes api version of sparkApplication
:param watch: whether to watch the job status and logs or not
"""

template_fields: Sequence[str] = ("application_file", "namespace")
Expand All @@ -60,6 +63,7 @@ def __init__(
in_cluster: bool | None = None,
cluster_context: str | None = None,
config_file: str | None = None,
watch: bool = False,
**kwargs,
) -> None:
super().__init__(**kwargs)
Expand All @@ -72,6 +76,7 @@ def __init__(
self.in_cluster = in_cluster
self.cluster_context = cluster_context
self.config_file = config_file
self.watch = watch

self.hook = KubernetesHook(
conn_id=self.kubernetes_conn_id,
Expand All @@ -84,35 +89,67 @@ def execute(self, context: Context):
body = _load_body_to_dict(self.application_file)
name = body["metadata"]["name"]
namespace = self.namespace or self.hook.get_namespace()
namespace_event_stream = Watch().stream(
self.hook.core_v1_client.list_namespaced_pod,
namespace=namespace,
_preload_content=False,
watch=True,
label_selector=f"sparkoperator.k8s.io/app-name={name},spark-role=driver",
field_selector="status.phase=Running",
)

self.hook.create_custom_object(
group=self.api_group,
version=self.api_version,
plural=self.plural,
body=body,
namespace=namespace,
)
for event in namespace_event_stream:
if event["type"] == "ADDED":
pod_log_stream = Watch().stream(
self.hook.core_v1_client.read_namespaced_pod_log,
name=f"{name}-driver",
response = None
is_job_created = False
if self.watch:
try:
namespace_event_stream = Watch().stream(
self.hook.core_v1_client.list_namespaced_event,
namespace=namespace,
watch=True,
field_selector=f"involvedObject.kind=SparkApplication,involvedObject.name={name}",
)

response = self.hook.create_custom_object(
group=self.api_group,
version=self.api_version,
plural=self.plural,
body=body,
namespace=namespace,
_preload_content=False,
timestamps=True,
)
for line in pod_log_stream:
self.log.info(line)
else:
break
is_job_created = True
for event in namespace_event_stream:
obj = event["object"]
if event["object"].last_timestamp >= datetime.datetime.strptime(
response["metadata"]["creationTimestamp"], "%Y-%m-%dT%H:%M:%S%z"
):
self.log.info(obj.message)
if obj.reason == "SparkDriverRunning":
pod_log_stream = Watch().stream(
self.hook.core_v1_client.read_namespaced_pod_log,
name=f"{name}-driver",
namespace=namespace,
timestamps=True,
)
for line in pod_log_stream:
self.log.info(line)
elif obj.reason in [
"SparkApplicationSubmissionFailed",
"SparkApplicationFailed",
"SparkApplicationDeleted",
]:
is_job_created = False
raise AirflowException(obj.message)
elif obj.reason == "SparkApplicationCompleted":
break
else:
continue
except Exception:
if is_job_created:
self.on_kill()
raise

else:
response = self.hook.create_custom_object(
group=self.api_group,
version=self.api_version,
plural=self.plural,
body=body,
namespace=namespace,
)

return response

def on_kill(self) -> None:
body = _load_body_to_dict(self.application_file)
Expand Down
84 changes: 73 additions & 11 deletions tests/providers/cncf/kubernetes/operators/test_spark_kubernetes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,13 @@
# under the License.
from __future__ import annotations

from unittest.mock import patch
from datetime import datetime
from unittest.mock import MagicMock, patch

import pytest
from dateutil import tz

from airflow import AirflowException
from airflow.providers.cncf.kubernetes.operators.spark_kubernetes import SparkKubernetesOperator


Expand All @@ -43,13 +48,22 @@ def test_spark_kubernetes_operator(mock_kubernetes_hook):
@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.Watch.stream")
@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes._load_body_to_dict")
@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.KubernetesHook")
def test_execute(mock_kubernetes_hook, mock_load_body_to_dict, mock_stream):
def test_execute_with_watch(mock_kubernetes_hook, mock_load_body_to_dict, mock_stream):
mock_load_body_to_dict.return_value = {"metadata": {"name": "spark-app"}}

mock_kubernetes_hook.return_value.create_custom_object.return_value = {
"metadata": {"name": "spark-app", "creationTimestamp": "2022-01-01T00:00:00Z"}
}
mock_kubernetes_hook.return_value.get_namespace.return_value = "default"
mock_stream.side_effect = [[{"type": "ADDED"}], []]

op = SparkKubernetesOperator(task_id="task_id", application_file="application_file")
op.execute({})
object_mock = MagicMock()
object_mock.reason = "SparkDriverRunning"
object_mock.last_timestamp = datetime(2022, 1, 1, 23, 59, 59, tzinfo=tz.tzutc())
mock_stream.side_effect = [[{"object": object_mock}], []]

op = SparkKubernetesOperator(task_id="task_id", application_file="application_file", watch=True)
operator_output = op.execute({})

mock_kubernetes_hook.return_value.create_custom_object.assert_called_once_with(
group="sparkoperator.k8s.io",
version="v1beta2",
Expand All @@ -60,22 +74,70 @@ def test_execute(mock_kubernetes_hook, mock_load_body_to_dict, mock_stream):

assert mock_stream.call_count == 2
mock_stream.assert_any_call(
mock_kubernetes_hook.return_value.core_v1_client.list_namespaced_pod,
mock_kubernetes_hook.return_value.core_v1_client.list_namespaced_event,
namespace="default",
_preload_content=False,
watch=True,
label_selector="sparkoperator.k8s.io/app-name=spark-app,spark-role=driver",
field_selector="status.phase=Running",
field_selector="involvedObject.kind=SparkApplication,involvedObject.name=spark-app",
)

mock_stream.assert_any_call(
mock_kubernetes_hook.return_value.core_v1_client.read_namespaced_pod_log,
name="spark-app-driver",
namespace="default",
_preload_content=False,
timestamps=True,
)

assert operator_output == {"metadata": {"name": "spark-app", "creationTimestamp": "2022-01-01T00:00:00Z"}}


@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.SparkKubernetesOperator.on_kill")
@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.Watch.stream")
@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes._load_body_to_dict")
@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.KubernetesHook")
def test_raise_exception_when_job_fails(
mock_kubernetes_hook, mock_load_body_to_dict, mock_stream, mock_on_kill
):
mock_load_body_to_dict.return_value = {"metadata": {"name": "spark-app"}}

mock_kubernetes_hook.return_value.create_custom_object.return_value = {
"metadata": {"name": "spark-app", "creationTimestamp": "2022-01-01T00:00:00Z"}
}
mock_kubernetes_hook.return_value.get_namespace.return_value = "default"

object_mock = MagicMock()
object_mock.reason = "SparkApplicationFailed"
object_mock.message = "spark-app submission failed"
object_mock.last_timestamp = datetime(2022, 1, 1, 23, 59, 59, tzinfo=tz.tzutc())

mock_stream.side_effect = [[{"object": object_mock}], []]
op = SparkKubernetesOperator(task_id="task_id", application_file="application_file", watch=True)
with pytest.raises(AirflowException, match="spark-app submission failed"):
op.execute({})

assert mock_on_kill.has_called_once()


@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes._load_body_to_dict")
@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.KubernetesHook")
def test_execute_without_watch(mock_kubernetes_hook, mock_load_body_to_dict):
mock_load_body_to_dict.return_value = {"metadata": {"name": "spark-app"}}

mock_kubernetes_hook.return_value.create_custom_object.return_value = {
"metadata": {"name": "spark-app", "creationTimestamp": "2022-01-01T00:00:00Z"}
}
mock_kubernetes_hook.return_value.get_namespace.return_value = "default"

op = SparkKubernetesOperator(task_id="task_id", application_file="application_file")
operator_output = op.execute({})

mock_kubernetes_hook.return_value.create_custom_object.assert_called_once_with(
group="sparkoperator.k8s.io",
version="v1beta2",
plural="sparkapplications",
body={"metadata": {"name": "spark-app"}},
namespace="default",
)
assert operator_output == {"metadata": {"name": "spark-app", "creationTimestamp": "2022-01-01T00:00:00Z"}}


@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes._load_body_to_dict")
@patch("airflow.providers.cncf.kubernetes.operators.spark_kubernetes.KubernetesHook")
Expand Down