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 9246fc51e32 Add reserved IP ranges to Vertex AI pipeline jobs (#72560)
9246fc51e32 is described below
commit 9246fc51e32f822adf0da0a73274c53abd91e60e
Author: Zhibo Lin <[email protected]>
AuthorDate: Wed Sep 9 19:14:13 2026 +0800
Add reserved IP ranges to Vertex AI pipeline jobs (#72560)
* Add reserved IP ranges to Vertex AI pipeline jobs
* Split Vertex AI pipeline job hook tests
---
.../google/cloud/hooks/vertex_ai/pipeline_job.py | 10 ++++++
.../cloud/operators/vertex_ai/pipeline_job.py | 6 ++++
.../cloud/hooks/vertex_ai/test_pipeline_job.py | 38 ++++++++++++++++++++++
.../unit/google/cloud/operators/test_vertex_ai.py | 5 ++-
4 files changed, 58 insertions(+), 1 deletion(-)
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py
b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py
index 17770cc14c9..efd597beeb8 100644
---
a/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py
+++
b/providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py
@@ -185,6 +185,7 @@ class PipelineJobHook(GoogleBaseHook, OperationHelper):
network: str | None = None,
create_request_timeout: float | None = None,
experiment: str | experiment_resources.Experiment | None = None,
+ reserved_ip_ranges: list[str] | None = None,
# END: run param
) -> PipelineJob:
"""
@@ -234,6 +235,9 @@ class PipelineJobHook(GoogleBaseHook, OperationHelper):
PipelineJob. Metrics produced by the PipelineJob as system.Metric
Artifacts will be associated as
metrics to the current Experiment Run. Pipeline parameters will be
associated as parameters to
the current Experiment Run.
+ :param reserved_ip_ranges: Optional. A list of names for the reserved
IP ranges under the VPC
+ network that can be used for this PipelineJob. If set, the
PipelineJob will only use IP
+ addresses from these ranges.
"""
self._pipeline_job = self.get_pipeline_job_object(
display_name=display_name,
@@ -252,6 +256,7 @@ class PipelineJobHook(GoogleBaseHook, OperationHelper):
self._pipeline_job.submit(
service_account=service_account,
network=network,
+ reserved_ip_ranges=reserved_ip_ranges,
create_request_timeout=create_request_timeout,
experiment=experiment,
)
@@ -279,6 +284,7 @@ class PipelineJobHook(GoogleBaseHook, OperationHelper):
network: str | None = None,
create_request_timeout: float | None = None,
experiment: str | experiment_resources.Experiment | None = None,
+ reserved_ip_ranges: list[str] | None = None,
# END: run param
) -> PipelineJob:
"""
@@ -331,6 +337,9 @@ class PipelineJobHook(GoogleBaseHook, OperationHelper):
Metrics produced by the PipelineJob as system.Metric Artifacts
will be associated as metrics
to the current Experiment Run. Pipeline parameters will be
associated as parameters to
the current Experiment Run.
+ :param reserved_ip_ranges: Optional. A list of names for the reserved
IP ranges under the VPC
+ network that can be used for this PipelineJob. If set, the
PipelineJob will only use IP
+ addresses from these ranges.
"""
self._pipeline_job = self.get_pipeline_job_object(
display_name=display_name,
@@ -349,6 +358,7 @@ class PipelineJobHook(GoogleBaseHook, OperationHelper):
self._pipeline_job.submit(
service_account=service_account,
network=network,
+ reserved_ip_ranges=reserved_ip_ranges,
create_request_timeout=create_request_timeout,
experiment=experiment,
)
diff --git
a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py
b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py
index 29359ac2d23..bb6a75546a4 100644
---
a/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py
+++
b/providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py
@@ -92,6 +92,9 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
Metrics produced by the PipelineJob as system.Metric Artifacts will be
associated as metrics
to the current Experiment Run. Pipeline parameters will be associated
as parameters to
the current Experiment Run.
+ :param reserved_ip_ranges: Optional. A list of names for the reserved IP
ranges under the VPC
+ network that can be used for this PipelineJob. If set, the PipelineJob
will only use IP
+ addresses from these ranges.
:param gcp_conn_id: The connection ID to use connecting to Google Cloud.
:param impersonation_chain: Optional service account to impersonate using
short-term
credentials, or chained list of accounts required to get the
access_token
@@ -137,6 +140,7 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
network: str | None = None,
create_request_timeout: float | None = None,
experiment: str | experiment_resources.Experiment | None = None,
+ reserved_ip_ranges: list[str] | None = None,
gcp_conn_id: str = "google_cloud_default",
impersonation_chain: str | Sequence[str] | None = None,
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
@@ -160,6 +164,7 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
self.network = network
self.create_request_timeout = create_request_timeout
self.experiment = experiment
+ self.reserved_ip_ranges = reserved_ip_ranges
self.gcp_conn_id = gcp_conn_id
self.impersonation_chain = impersonation_chain
self.deferrable = deferrable
@@ -189,6 +194,7 @@ class RunPipelineJobOperator(GoogleCloudBaseOperator):
failure_policy=self.failure_policy,
service_account=self.service_account,
network=self.network,
+ reserved_ip_ranges=self.reserved_ip_ranges,
create_request_timeout=self.create_request_timeout,
experiment=self.experiment,
)
diff --git
a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_pipeline_job.py
b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_pipeline_job.py
index fc1df7b8d15..f5a9ec016b6 100644
---
a/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_pipeline_job.py
+++
b/providers/google/tests/unit/google/cloud/hooks/vertex_ai/test_pipeline_job.py
@@ -246,6 +246,44 @@ class TestPipelineJobWithoutDefaultProjectIdHook:
)
mock_client.return_value.common_location_path.assert_called_once_with(TEST_PROJECT_ID,
TEST_REGION)
+ @pytest.mark.parametrize(
+ "reserved_ip_ranges", [None, [], ["range-1", "range-2"]], ids=["none",
"empty", "multiple"]
+ )
+
@mock.patch(PIPELINE_JOB_STRING.format("PipelineJobHook.get_pipeline_job_object"))
+ def test_run_pipeline_job_forwards_reserved_ip_ranges(
+ self, mock_get_pipeline_job_object, reserved_ip_ranges
+ ) -> None:
+ self.hook.run_pipeline_job(
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ display_name="display-name",
+ template_path="gs://bucket/template.json",
+ reserved_ip_ranges=reserved_ip_ranges,
+ )
+
+ assert
mock_get_pipeline_job_object.return_value.submit.call_args.kwargs["reserved_ip_ranges"]
== (
+ reserved_ip_ranges
+ )
+
+ @pytest.mark.parametrize(
+ "reserved_ip_ranges", [None, [], ["range-1", "range-2"]], ids=["none",
"empty", "multiple"]
+ )
+
@mock.patch(PIPELINE_JOB_STRING.format("PipelineJobHook.get_pipeline_job_object"))
+ def test_submit_pipeline_job_forwards_reserved_ip_ranges(
+ self, mock_get_pipeline_job_object, reserved_ip_ranges
+ ) -> None:
+ self.hook.submit_pipeline_job(
+ project_id=TEST_PROJECT_ID,
+ region=TEST_REGION,
+ display_name="display-name",
+ template_path="gs://bucket/template.json",
+ reserved_ip_ranges=reserved_ip_ranges,
+ )
+
+ assert
mock_get_pipeline_job_object.return_value.submit.call_args.kwargs["reserved_ip_ranges"]
== (
+ reserved_ip_ranges
+ )
+
class TestPipelineJobAsyncHook:
@pytest.mark.asyncio
diff --git
a/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py
b/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py
index 7fee751c47d..d333ca331ed 100644
--- a/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py
+++ b/providers/google/tests/unit/google/cloud/operators/test_vertex_ai.py
@@ -2881,9 +2881,10 @@ class TestVertexAIDeleteModelVersionOperator:
class TestVertexAIRunPipelineJobOperator:
+ @pytest.mark.parametrize("reserved_ip_ranges", [None, [], ["range-1",
"range-2"]])
@mock.patch(VERTEX_AI_PATH.format("pipeline_job.PipelineJobHook"))
@mock.patch("google.cloud.aiplatform_v1.types.PipelineJob.to_dict")
- def test_execute(self, to_dict_mock, mock_hook):
+ def test_execute(self, to_dict_mock, mock_hook, reserved_ip_ranges):
op = RunPipelineJobOperator(
task_id=TASK_ID,
gcp_conn_id=GCP_CONN_ID,
@@ -2902,6 +2903,7 @@ class TestVertexAIRunPipelineJobOperator:
failure_policy="",
service_account="",
network="",
+ reserved_ip_ranges=reserved_ip_ranges,
create_request_timeout=None,
experiment=None,
)
@@ -2922,6 +2924,7 @@ class TestVertexAIRunPipelineJobOperator:
failure_policy="",
service_account="",
network="",
+ reserved_ip_ranges=reserved_ip_ranges,
create_request_timeout=None,
experiment=None,
)
