diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py index 87d838333baea..ad691581e4af2 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/hooks/spark_submit.py @@ -818,7 +818,7 @@ def submit(self, application: str = "", **kwargs: Any) -> str | None: f"{self._submit_log_tail}" ) finally: - # K8s-API tracking defers post-submit commands to _poll_k8s_driver_via_api's finally + # K8s-API tracking defers post-submit commands to poll_k8s_driver_via_api's finally # block so they run once after the driver reaches a terminal state. Spark cluster-mode # driver tracking defers them to poll_until_complete for the same reason. All other # modes run them here, immediately after spark-submit exits. @@ -1061,7 +1061,7 @@ def _query_yarn_application_status(self, application_id: str) -> tuple[str, str, f"{application_id}: {resp.text[:200]}" ) from exc - def _kill_yarn_application(self, application_id: str) -> None: + def kill_yarn_application(self, application_id: str) -> None: """PUT ``/ws/v1/cluster/apps/{id}/state`` to kill the application (best-effort).""" try: url = f"{self._get_yarn_rm_base_url()}/ws/v1/cluster/apps/{application_id}/state" @@ -1175,7 +1175,7 @@ def _start_driver_status_tracking(self) -> None: f"returncode = {returncode}" ) - def _poll_k8s_driver_via_api(self) -> str | None: + def poll_k8s_driver_via_api(self) -> str | None: """ Poll the K8s driver pod phase until it reaches a terminal state. @@ -1385,10 +1385,28 @@ def on_kill(self) -> None: # state because `yarn_track_via_rm_api=True` deliberately terminates # `_submit_sp` right after submission to free the JVM. if self._yarn_application_id and self._yarn_track_via_rm_api: - self._kill_yarn_application(self._yarn_application_id) + self.kill_yarn_application(self._yarn_application_id) self._run_post_submit_commands() + @property + def conf(self) -> dict[str, Any]: + """Return the live Spark conf dict (not a copy); callers may mutate keys in place.""" + return self._conf + + @property + def kubernetes_driver_pod(self) -> str | None: + return self._kubernetes_driver_pod + + @kubernetes_driver_pod.setter + def kubernetes_driver_pod(self, value: str | None) -> None: + self._kubernetes_driver_pod = value + + @property + def yarn_application_id(self) -> str | None: + """YARN application id captured from spark-submit logs. Get-only.""" + return self._yarn_application_id + def query_yarn_application_status(self, application_id: str) -> str: """ Return a normalized single string status for the ResumableJobMixin interface. diff --git a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py index 48d7015532643..848f2896e66a9 100644 --- a/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py +++ b/providers/apache/spark/src/airflow/providers/apache/spark/operators/spark_submit.py @@ -96,9 +96,9 @@ class _KubernetesSparkSubmitBackend(_SparkSubmitDeploymentBackend): """Logic for tracking Spark driver pods in Kubernetes.""" def submit_job(self, context: Context) -> str | None: - self.hook._conf[_K8S_WAIT_APP_COMPLETION_CONF] = "false" + self.hook.conf[_K8S_WAIT_APP_COMPLETION_CONF] = "false" self.hook.submit(self.operator.application) - pod_name = self.hook._kubernetes_driver_pod + pod_name = self.hook.kubernetes_driver_pod namespace = self.hook._connection["namespace"] if not pod_name: raise RuntimeError("spark-submit did not capture a K8s driver pod name") @@ -133,8 +133,8 @@ def is_job_succeeded(self, status: str) -> bool: def poll_until_complete(self, external_id: str, context: Context) -> None: if external_id is not None: _, pod_name = self.operator._parse_k8s_external_id(external_id) - self.hook._kubernetes_driver_pod = pod_name - terminal_phase = self.hook._poll_k8s_driver_via_api() + self.hook.kubernetes_driver_pod = pod_name + terminal_phase = self.hook.poll_k8s_driver_via_api() # Cache only when the pod actually reached Succeeded, the 404/vanished path # returns None for cases like: pod deleted by on_kill or garbage collected after failure) # and must not be cached, otherwise a retry would see "Succeeded" and skip resubmission. @@ -150,15 +150,15 @@ class _YarnSparkSubmitBackend(_SparkSubmitDeploymentBackend): """Logic for tracking Spark applications in YARN cluster mode.""" def submit_job(self, context: Context) -> str | None: - if self.hook._conf.get("spark.yarn.submit.waitAppCompletion", "").strip().lower() == "true": + if self.hook.conf.get("spark.yarn.submit.waitAppCompletion", "").strip().lower() == "true": raise ValueError( "spark.yarn.submit.waitAppCompletion=true cannot be set for cluster mode as it conflicts " "with the need to exit spark-submit immediately to persist the application ID for tracking. " "Either remove the explicit conf or set durable=False." ) - self.hook._conf["spark.yarn.submit.waitAppCompletion"] = "false" + self.hook.conf["spark.yarn.submit.waitAppCompletion"] = "false" self.hook.submit(self.operator.application) - app_id = self.hook._yarn_application_id + app_id = self.hook.yarn_application_id if not app_id: raise RuntimeError("spark-submit did not produce a YARN application ID") self.operator.log.info("YARN application submitted: %s", app_id) @@ -181,10 +181,10 @@ def poll_until_complete(self, external_id: str, context: Context) -> None: self.hook._run_post_submit_commands() def on_kill(self) -> None: - if self.hook._yarn_application_id: + if self.hook.yarn_application_id: # spark-submit has already exited (waitAppCompletion=false), so the hook's # CLI-based kill has nothing to terminate. Kill the YARN app via REST API instead. - self.hook._kill_yarn_application(self.hook._yarn_application_id) + self.hook.kill_yarn_application(self.hook.yarn_application_id) else: self.hook.on_kill() diff --git a/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py b/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py index d25ddbf7f76ee..1a130af742a60 100644 --- a/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py +++ b/providers/apache/spark/tests/unit/apache/spark/hooks/test_spark_submit.py @@ -1632,7 +1632,7 @@ def test_poll_k8s_driver_succeeds(self, mock_get_client): mock_client.read_namespaced_pod.side_effect = [running_pod, succeeded_pod] with patch.object(hook, "_run_post_submit_commands"): - hook._poll_k8s_driver_via_api() + hook.poll_k8s_driver_via_api() assert mock_client.delete_namespaced_pod.call_args.args[:2] == ("spark-app-abc-driver", "mynamespace") @@ -1647,7 +1647,7 @@ def test_poll_k8s_driver_raises_on_failed(self, mock_get_client): mock_client.read_namespaced_pod.return_value = failed_pod with pytest.raises(RuntimeError, match="phase=Failed"): - hook._poll_k8s_driver_via_api() + hook.poll_k8s_driver_via_api() @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_poll_k8s_driver_raises_after_consecutive_unknown(self, mock_get_client): @@ -1659,7 +1659,7 @@ def test_poll_k8s_driver_raises_after_consecutive_unknown(self, mock_get_client) mock_client.read_namespaced_pod.return_value = V1Pod(status=V1PodStatus(phase="Unknown")) with patch("time.sleep"), pytest.raises(RuntimeError, match="Unknown phase"): - hook._poll_k8s_driver_via_api() + hook.poll_k8s_driver_via_api() # assert that it was polled minimum 3 times to confirm the Unknown status before raising assert mock_client.read_namespaced_pod.call_count == 3 @@ -1677,13 +1677,13 @@ def test_poll_k8s_driver_tolerates_transient_api_errors(self, mock_get_client, _ mock_client.read_namespaced_pod.side_effect = [api_error, api_error, succeeded_pod] with patch.object(hook, "_run_post_submit_commands"): - hook._poll_k8s_driver_via_api() + hook.poll_k8s_driver_via_api() assert mock_client.read_namespaced_pod.call_count == 3 @patch("airflow.providers.cncf.kubernetes.kube_client.get_kube_client") def test_post_submit_commands_run_exactly_once_on_k8s_path(self, mock_get_client): - """_run_post_submit_commands must fire exactly once: in _poll_k8s_driver_via_api finally.""" + """_run_post_submit_commands must fire exactly once: in poll_k8s_driver_via_api finally.""" hook = SparkSubmitHook(conn_id="spark_k8s_cluster", track_driver_via_k8s_api=True) hook._kubernetes_driver_pod = "spark-app-abc-driver" hook._kubernetes_application_id = "spark-abc" @@ -1692,7 +1692,7 @@ def test_post_submit_commands_run_exactly_once_on_k8s_path(self, mock_get_client mock_client.read_namespaced_pod.return_value = V1Pod(status=V1PodStatus(phase="Succeeded")) with patch.object(hook, "_run_post_submit_commands") as mock_cmd: - hook._poll_k8s_driver_via_api() + hook.poll_k8s_driver_via_api() mock_cmd.assert_called_once() @@ -1708,7 +1708,7 @@ def test_poll_k8s_driver_raises_after_consecutive_api_errors(self, mock_get_clie mock_client.read_namespaced_pod.side_effect = api_error with pytest.raises(RuntimeError, match="K8s API unreachable"): - hook._poll_k8s_driver_via_api() + hook.poll_k8s_driver_via_api() assert mock_client.read_namespaced_pod.call_count == 3 @@ -1722,7 +1722,7 @@ def test_poll_k8s_driver_exits_cleanly_on_404(self, mock_get_client): mock_client = mock_get_client.return_value mock_client.read_namespaced_pod.side_effect = kube_client.ApiException(status=404, reason="Not Found") - hook._poll_k8s_driver_via_api() + hook.poll_k8s_driver_via_api() mock_client.delete_namespaced_pod.assert_not_called() @@ -2214,3 +2214,21 @@ def test_on_kill_tolerates_rm_failure(self, mock_put): hook.on_kill() mock_put.assert_called_once() + + def test_conf_mutation_writes_through(self): + hook = SparkSubmitHook(conn_id="") + hook.conf["spark.foo"] = "bar" + assert hook._conf["spark.foo"] == "bar" + + def test_kubernetes_driver_pod_writes_through(self): + hook = SparkSubmitHook(conn_id="") + hook.kubernetes_driver_pod = "spark-app-abc-driver" + assert hook._kubernetes_driver_pod == "spark-app-abc-driver" + assert hook.kubernetes_driver_pod == "spark-app-abc-driver" + + def test_yarn_application_id_is_get_only(self): + hook = SparkSubmitHook(conn_id="") + hook._yarn_application_id = "application_1486558679801_1820" + assert hook.yarn_application_id == "application_1486558679801_1820" + with pytest.raises(AttributeError): + hook.yarn_application_id = "application_1486558679801_1821" diff --git a/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py b/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py index 3b525582ed7c2..6b3cd610b4899 100644 --- a/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py +++ b/providers/apache/spark/tests/unit/apache/spark/operators/test_spark_submit.py @@ -748,8 +748,8 @@ def capture(url, timeout): def test_yarn_first_run_persists_app_id_before_polling(self): operator = self._make_operator() operator._hook = self._make_hook(is_yarn_cluster=True) - operator._hook._conf = {} - operator._hook._yarn_application_id = "application_1234_0001" + operator._hook.conf = {} + operator._hook.yarn_application_id = "application_1234_0001" operator._hook.submit.return_value = None task_store = FakeTaskStateStore() @@ -791,8 +791,8 @@ def test_yarn_retry_skips_already_succeeded_app(self): def test_yarn_retry_resubmits_after_failed_app(self): operator = self._make_operator() operator._hook = self._make_hook(is_yarn_cluster=True) - operator._hook._conf = {} - operator._hook._yarn_application_id = "application_1234_0002" + operator._hook.conf = {} + operator._hook.yarn_application_id = "application_1234_0002" operator._hook.submit.return_value = None task_store = FakeTaskStateStore({"spark_job_id": "application_1234_0001"}) @@ -808,19 +808,19 @@ def test_yarn_retry_resubmits_after_failed_app(self): def test_yarn_injects_wait_app_completion_false(self): operator = self._make_operator() hook = self._make_hook(is_yarn_cluster=True) - hook._conf = {} - hook._yarn_application_id = "application_1234_0001" + hook.conf = {} + hook.yarn_application_id = "application_1234_0001" hook.submit.return_value = None operator._hook = hook operator.submit_job(context={}) - assert hook._conf.get("spark.yarn.submit.waitAppCompletion") == "false" + assert hook.conf.get("spark.yarn.submit.waitAppCompletion") == "false" def test_yarn_raises_if_wait_app_completion_true(self): operator = self._make_operator() hook = self._make_hook(is_yarn_cluster=True) - hook._conf = {"spark.yarn.submit.waitAppCompletion": "true"} + hook.conf = {"spark.yarn.submit.waitAppCompletion": "true"} operator._hook = hook with pytest.raises( @@ -869,16 +869,16 @@ def simulate_failed_tracking(): operator.poll_until_complete("driver-001", {}) def test_on_kill_sends_authenticated_kill_to_yarn_rm(self): - """operator.on_kill() must call _kill_yarn_application so Kerberos auth is applied.""" + """operator.on_kill() must call kill_yarn_application so Kerberos auth is applied.""" operator = self._make_operator() hook = self._make_hook(is_yarn_cluster=True) hook._is_yarn_cluster_mode = True - hook._yarn_application_id = "application_1234_0001" + hook.yarn_application_id = "application_1234_0001" operator._hook = hook operator.on_kill() - hook._kill_yarn_application.assert_called_once_with("application_1234_0001") + hook.kill_yarn_application.assert_called_once_with("application_1234_0001") def test_yarn_cluster_reconnect_without_rm_api_raises(self): """durable=True + yarn_track_via_rm_api=False must raise - RM API is required for resume.""" @@ -917,7 +917,7 @@ def _make_k8s_hook(self): hook._is_kubernetes = True hook._is_yarn = False hook._is_yarn_cluster_mode = False - hook._conf = {} + hook.conf = {} return hook def test_execute_calls_submit_then_poll_when_flag_set(self): @@ -926,12 +926,12 @@ def test_execute_calls_submit_then_poll_when_flag_set(self): operator._hook = hook call_order = [] hook.submit.side_effect = lambda *a, **kw: call_order.append("submit") - hook._poll_k8s_driver_via_api.side_effect = lambda: call_order.append("poll") + hook.poll_k8s_driver_via_api.side_effect = lambda: call_order.append("poll") operator.execute(context={}) hook.submit.assert_called_once_with("test.jar") - hook._poll_k8s_driver_via_api.assert_called_once() + hook.poll_k8s_driver_via_api.assert_called_once() assert call_order == ["submit", "poll"] def test_execute_falls_through_to_plain_submit_when_flag_off(self): @@ -944,25 +944,25 @@ def test_execute_falls_through_to_plain_submit_when_flag_off(self): operator.execute(context={}) hook.submit.assert_called_once_with("test.jar") - hook._poll_k8s_driver_via_api.assert_not_called() + hook.poll_k8s_driver_via_api.assert_not_called() def test_k8s_submit_job_returns_encoded_external_id(self): operator = self._make_operator(track_driver_via_k8s_api=True) hook = self._make_k8s_hook() - hook._kubernetes_driver_pod = "spark-abc-driver" + hook.kubernetes_driver_pod = "spark-abc-driver" hook._connection = {"namespace": "mynamespace"} operator._hook = hook result = operator.submit_job(context={}) assert result == "mynamespace:spark-abc-driver" - assert hook._conf.get("spark.kubernetes.submission.waitAppCompletion") == "false" + assert hook.conf.get("spark.kubernetes.submission.waitAppCompletion") == "false" hook.submit.assert_called_once_with("test.jar") def test_k8s_submit_job_raises_when_pod_name_missing(self): operator = self._make_operator(track_driver_via_k8s_api=True) hook = self._make_k8s_hook() - hook._kubernetes_driver_pod = None + hook.kubernetes_driver_pod = None hook._connection = {"namespace": "mynamespace"} operator._hook = hook @@ -1043,14 +1043,14 @@ def test_k8s_poll_until_complete_sets_pod_name_and_calls_poll_api(self): operator.poll_until_complete("mynamespace:spark-abc-driver", {}) - assert hook._kubernetes_driver_pod == "spark-abc-driver" - hook._poll_k8s_driver_via_api.assert_called_once() + assert hook.kubernetes_driver_pod == "spark-abc-driver" + hook.poll_k8s_driver_via_api.assert_called_once() @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="task_state_store requires Airflow 3.3+") def test_k8s_poll_until_complete_writes_succeeded_to_task_store(self): operator = self._make_operator(track_driver_via_k8s_api=True) hook = self._make_k8s_hook() - hook._poll_k8s_driver_via_api.return_value = "Succeeded" + hook.poll_k8s_driver_via_api.return_value = "Succeeded" operator._hook = hook task_store = FakeTaskStateStore() @@ -1062,7 +1062,7 @@ def test_k8s_poll_until_complete_writes_succeeded_to_task_store(self): def test_k8s_polling_does_not_write_task_store_when_reconnect_disabled(self): operator = self._make_operator(track_driver_via_k8s_api=True, durable=False) hook = self._make_k8s_hook() - hook._poll_k8s_driver_via_api.return_value = "Succeeded" + hook.poll_k8s_driver_via_api.return_value = "Succeeded" operator._hook = hook task_store = FakeTaskStateStore() @@ -1073,7 +1073,7 @@ def test_k8s_polling_does_not_write_task_store_when_reconnect_disabled(self): def test_k8s_poll_until_complete_does_not_cache_and_reraises_on_failure(self): operator = self._make_operator(track_driver_via_k8s_api=True) hook = self._make_k8s_hook() - hook._poll_k8s_driver_via_api.side_effect = RuntimeError("Spark application failed (phase=Failed)") + hook.poll_k8s_driver_via_api.side_effect = RuntimeError("Spark application failed (phase=Failed)") operator._hook = hook task_store = FakeTaskStateStore() @@ -1096,7 +1096,7 @@ def test_k8s_execute_persists_pod_id_when_durable(self): """execute() with durable=True stores the pod ID in task_store before polling.""" operator = self._make_operator(track_driver_via_k8s_api=True, durable=True) hook = self._make_k8s_hook() - hook._kubernetes_driver_pod = "spark-abc-driver" + hook.kubernetes_driver_pod = "spark-abc-driver" hook._connection = {"namespace": "mynamespace"} operator._hook = hook task_store = FakeTaskStateStore() @@ -1119,7 +1119,7 @@ def test_k8s_execute_durable_false_does_not_persist_pod_id(self): """execute() with durable=False does not write spark_job_id to task_store.""" operator = self._make_operator(track_driver_via_k8s_api=True, durable=False) hook = self._make_k8s_hook() - hook._kubernetes_driver_pod = "spark-abc-driver" + hook.kubernetes_driver_pod = "spark-abc-driver" hook._connection = {"namespace": "mynamespace"} operator._hook = hook task_store = FakeTaskStateStore()