This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 9a28876d514 Validate Gemini batch results_folder in the worker that
writes results (#70393)
9a28876d514 is described below
commit 9a28876d51496f5119b825515934694018c25c76
Author: Shahar Epstein <[email protected]>
AuthorDate: Sat Aug 1 07:42:20 2026 +0300
Validate Gemini batch results_folder in the worker that writes results
(#70393)
In deferrable mode the results file is written by execute_complete in a
resume worker where execute() -- and its results_folder pre-flight check
-- never ran. A folder missing there (heterogeneous workers, relative
path against a different CWD) surfaced as a raw FileNotFoundError after
the billable batch job already succeeded, instead of the operator's
clear validation error.
---
.../providers/google/cloud/operators/gen_ai.py | 18 ++++++--
.../unit/google/cloud/operators/test_gen_ai.py | 50 ++++++++++++++++++++++
2 files changed, 64 insertions(+), 4 deletions(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
index 0801bb24a75..65d8a943374 100644
--- a/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
+++ b/providers/google/src/airflow/providers/google/cloud/operators/gen_ai.py
@@ -499,6 +499,11 @@ class
GenAIGeminiCreateBatchJobOperator(GoogleCloudBaseOperator):
raise AirflowException("Something went wrong during waiting of the
batch job.")
return job
+ def _validate_results_folder(self):
+ # Re-checked where the results file is written: deferral resume skips
execute().
+ if self.results_folder and not
os.path.exists(os.path.abspath(self.results_folder)):
+ raise AirflowException("path to results_folder does not exist,
please provide correct path")
+
def _prepare_results_for_xcom(self, job):
results = []
if job.dest and job.dest.inlined_responses:
@@ -514,6 +519,7 @@ class
GenAIGeminiCreateBatchJobOperator(GoogleCloudBaseOperator):
self.log.warning("Error found in the inline result")
results.append(inline_response.error)
elif job.dest and job.dest.file_name:
+ self._validate_results_folder()
file_content_bytes =
self.hook.download_file(file_name=job.dest.file_name)
file_content = file_content_bytes.decode("utf-8")
file_name = job.display_name or job.name.replace("/", "-")
@@ -541,8 +547,7 @@ class
GenAIGeminiCreateBatchJobOperator(GoogleCloudBaseOperator):
if self.results_folder and not isinstance(self.input_source, str):
raise AirflowException("results_folder works only when
input_source is file name")
- if self.results_folder and not
os.path.exists(os.path.abspath(self.results_folder)):
- raise AirflowException("path to results_folder does not exist,
please provide correct path")
+ self._validate_results_folder()
if self.deferrable:
self.defer(
@@ -944,6 +949,11 @@ class
GenAIGeminiCreateEmbeddingsBatchJobOperator(GoogleCloudBaseOperator):
raise AirflowException("Something went wrong during waiting of the
batch job: %s", e)
return job
+ def _validate_results_folder(self):
+ # Re-checked where the results file is written: deferral resume skips
execute().
+ if self.results_folder and not
os.path.exists(os.path.abspath(self.results_folder)):
+ raise AirflowException("path to results_folder does not exist,
please provide correct path")
+
def _prepare_results_for_xcom(self, job):
results = []
if job.dest and job.dest.inlined_embed_content_responses:
@@ -959,6 +969,7 @@ class
GenAIGeminiCreateEmbeddingsBatchJobOperator(GoogleCloudBaseOperator):
self.log.warning("Error found in the inline result")
results.append(inline_embed_response.error)
elif job.dest and job.dest.file_name:
+ self._validate_results_folder()
file_content_bytes =
self.hook.download_file(file_name=job.dest.file_name)
file_content = file_content_bytes.decode("utf-8")
file_name = job.display_name or job.name.replace("/", "-")
@@ -986,8 +997,7 @@ class
GenAIGeminiCreateEmbeddingsBatchJobOperator(GoogleCloudBaseOperator):
if self.results_folder and not isinstance(self.input_source, str):
raise AirflowException("results_folder works only when
input_source is file name")
- if self.results_folder and not
os.path.exists(os.path.abspath(self.results_folder)):
- raise AirflowException("path to results_folder does not exist,
please provide correct path")
+ self._validate_results_folder()
if self.deferrable:
self.defer(
trigger=GenAIGeminiCreateEmbeddingsBatchJobTrigger(
diff --git a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
index 3e814f9bd9a..1071f2de829 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_gen_ai.py
@@ -457,6 +457,31 @@ class TestGenAIGeminiCreateBatchJobOperator:
):
op.execute(context={"ti": mock.MagicMock()})
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def
test_prepare_results_for_xcom_results_folder_not_exists_raises_airflow_exception(self,
mock_hook):
+ op = GenAIGeminiCreateBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ model=TEST_GEMINI_MODEL,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ input_source=TEST_FILE_NAME,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ results_folder=TEST_FILE_PATH,
+ )
+ mock_job = mock.MagicMock()
+ mock_job.dest.inlined_responses = None
+ mock_job.dest.file_name = "results-file"
+
+ with pytest.raises(
+ AirflowException,
+ match="path to results_folder does not exist, please provide
correct path",
+ ):
+ op._prepare_results_for_xcom(mock_job)
+
+ mock_hook.return_value.download_file.assert_not_called()
+
@mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
def test__wait_until_complete_exception_raises_airflow_exception(self,
mock_hook):
op = GenAIGeminiCreateBatchJobOperator(
@@ -832,6 +857,31 @@ class TestGenAIGeminiCreateEmbeddingsBatchJobOperator:
):
op.execute(context={"ti": mock.MagicMock()})
+ @mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
+ def
test_prepare_results_for_xcom_results_folder_not_exists_raises_airflow_exception(self,
mock_hook):
+ op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
+ task_id=TASK_ID,
+ project_id=GCP_PROJECT,
+ location=GCP_LOCATION,
+ input_source=TEST_FILE_NAME,
+ model=EMBEDDING_MODEL,
+ gemini_api_key=TEST_GEMINI_API_KEY,
+ gcp_conn_id=GCP_CONN_ID,
+ impersonation_chain=IMPERSONATION_CHAIN,
+ results_folder=TEST_FILE_PATH,
+ )
+ mock_job = mock.MagicMock()
+ mock_job.dest.inlined_embed_content_responses = None
+ mock_job.dest.file_name = "results-file"
+
+ with pytest.raises(
+ AirflowException,
+ match="path to results_folder does not exist, please provide
correct path",
+ ):
+ op._prepare_results_for_xcom(mock_job)
+
+ mock_hook.return_value.download_file.assert_not_called()
+
@mock.patch(GEN_AI_PATH.format("GenAIGeminiAPIHook"))
def test__wait_until_complete_exception_raises_airflow_exception(self,
mock_hook):
op = GenAIGeminiCreateEmbeddingsBatchJobOperator(
