apache-airflow-providers-google 14.0.0__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

airflow/providers/google/cloud/operators/translate_speech.py

@@ -22,18 +22,20 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from google.protobuf.json_format import MessageToDict
+
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.speech_to_text import CloudSpeechToTextHook
 from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.providers.google.common.links.storage import FileDetailsLink
-from google.protobuf.json_format import MessageToDict
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class CloudTranslateSpeechOperator(GoogleCloudBaseOperator):
     """
@@ -171,7 +173,6 @@ class CloudTranslateSpeechOperator(GoogleCloudBaseOperator):
         if self.audio.uri:
             FileDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 # Slice from: "gs://{BUCKET_NAME}/{FILE_NAME}" to: "{BUCKET_NAME}/{FILE_NAME}"
                 uri=self.audio.uri[5:],
                 project_id=self.project_id or translate_hook.project_id,

airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py

@@ -21,7 +21,13 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform import datasets
+from google.cloud.aiplatform.models import Model
+from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.vertex_ai.auto_ml import AutoMLHook
@@ -32,16 +38,12 @@ from airflow.providers.google.cloud.links.vertex_ai import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.deprecated import deprecated
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform import datasets
-from google.cloud.aiplatform.models import Model
-from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator):
     """The base class for operators that launch AutoML jobs on VertexAI."""
@@ -91,6 +93,13 @@ class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator):
         self.impersonation_chain = impersonation_chain
         self.hook: AutoMLHook | None = None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def on_kill(self) -> None:
         """Act as a callback called when the operator is killed; cancel any running job."""
         if self.hook:
@@ -242,12 +251,12 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
 
@@ -334,12 +343,12 @@ class CreateAutoMLImageTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
 
@@ -457,15 +466,20 @@ class CreateAutoMLTabularTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
 
+@deprecated(
+    planned_removal_date="March 24, 2026",
+    use_instead="airflow.providers.google.cloud.operators.vertex_ai.generative_model.SupervisedFineTuningTrainOperator",
+    category=AirflowProviderDeprecationWarning,
+)
 class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator):
     """Create Auto ML Video Training job."""
 
@@ -531,12 +545,12 @@ class CreateAutoMLVideoTrainingJobOperator(AutoMLTrainingJobBaseOperator):
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
 
@@ -573,16 +587,6 @@ class DeleteAutoMLTrainingJobOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
-    @property
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="training_pipeline_id",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def training_pipeline(self):
-        """Alias for ``training_pipeline_id``, used for compatibility (deprecated)."""
-        return self.training_pipeline_id
-
     def execute(self, context: Context):
         hook = AutoMLHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -650,6 +654,12 @@ class ListAutoMLTrainingJobOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = AutoMLHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -666,5 +676,5 @@ class ListAutoMLTrainingJobOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAITrainingPipelinesLink.persist(context=context, task_instance=self)
+        VertexAITrainingPipelinesLink.persist(context=context)
         return [TrainingPipeline.to_dict(result) for result in results]

airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py

@@ -24,6 +24,10 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform_v1.types import BatchPredictionJob
+
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.batch_prediction_job import BatchPredictionJobHook
@@ -33,15 +37,13 @@ from airflow.providers.google.cloud.links.vertex_ai import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.vertex_ai import CreateBatchPredictionJobTrigger
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform_v1.types import BatchPredictionJob
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.aiplatform import BatchPredictionJob as BatchPredictionJobObject, Model, explain
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
     """
@@ -229,6 +231,13 @@ class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         self.log.info("Creating Batch prediction job")
         batch_prediction_job: BatchPredictionJobObject = self.hook.submit_batch_prediction_job(
@@ -260,9 +269,10 @@ class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
         batch_prediction_job_id = batch_prediction_job.name
         self.log.info("Batch prediction job was created. Job id: %s", batch_prediction_job_id)
 
-        self.xcom_push(context, key="batch_prediction_job_id", value=batch_prediction_job_id)
+        context["ti"].xcom_push(key="batch_prediction_job_id", value=batch_prediction_job_id)
         VertexAIBatchPredictionJobLink.persist(
-            context=context, task_instance=self, batch_prediction_job_id=batch_prediction_job_id
+            context=context,
+            batch_prediction_job_id=batch_prediction_job_id,
         )
 
         if self.deferrable:
@@ -293,13 +303,11 @@ class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
         job: dict[str, Any] = event["job"]
         self.log.info("Batch prediction job %s created and completed successfully.", job["name"])
         job_id = self.hook.extract_batch_prediction_job_id(job)
-        self.xcom_push(
-            context,
+        context["ti"].xcom_push(
             key="batch_prediction_job_id",
             value=job_id,
         )
-        self.xcom_push(
-            context,
+        context["ti"].xcom_push(
             key="training_conf",
             value={
                 "training_conf_id": job_id,
@@ -425,6 +433,13 @@ class GetBatchPredictionJobOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = BatchPredictionJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -443,7 +458,8 @@ class GetBatchPredictionJobOperator(GoogleCloudBaseOperator):
             )
             self.log.info("Batch prediction job was gotten.")
             VertexAIBatchPredictionJobLink.persist(
-                context=context, task_instance=self, batch_prediction_job_id=self.batch_prediction_job
+                context=context,
+                batch_prediction_job_id=self.batch_prediction_job,
             )
             return BatchPredictionJob.to_dict(result)
         except NotFound:
@@ -515,6 +531,12 @@ class ListBatchPredictionJobsOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = BatchPredictionJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -531,5 +553,5 @@ class ListBatchPredictionJobsOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAIBatchPredictionJobListLink.persist(context=context, task_instance=self)
+        VertexAIBatchPredictionJobListLink.persist(context=context)
         return [BatchPredictionJob.to_dict(result) for result in results]

airflow/providers/google/cloud/operators/vertex_ai/custom_job.py

@@ -23,8 +23,14 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.aiplatform.models import Model
+from google.cloud.aiplatform_v1.types.dataset import Dataset
+from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
+
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobHook
 from airflow.providers.google.cloud.links.vertex_ai import (
     VertexAIModelLink,
@@ -37,21 +43,17 @@ from airflow.providers.google.cloud.triggers.vertex_ai import (
     CustomPythonPackageTrainingJobTrigger,
     CustomTrainingJobTrigger,
 )
-from airflow.providers.google.common.deprecated import deprecated
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.aiplatform.models import Model
-from google.cloud.aiplatform_v1.types.dataset import Dataset
-from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.cloud.aiplatform import (
         CustomContainerTrainingJob,
         CustomPythonPackageTrainingJob,
         CustomTrainingJob,
     )
+    from google.cloud.aiplatform_v1.types import PscInterfaceConfig
+
+    from airflow.providers.common.compat.sdk import Context
 
 
 class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
@@ -109,6 +111,7 @@ class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
         predefined_split_column_name: str | None = None,
         timestamp_split_column_name: str | None = None,
         tensorboard: str | None = None,
+        psc_interface_config: PscInterfaceConfig | None = None,
         gcp_conn_id: str = "google_cloud_default",
         impersonation_chain: str | Sequence[str] | None = None,
         **kwargs,
@@ -165,21 +168,29 @@ class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
         self.predefined_split_column_name = predefined_split_column_name
         self.timestamp_split_column_name = timestamp_split_column_name
         self.tensorboard = tensorboard
+        self.psc_interface_config = psc_interface_config
         # END Run param
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "region": self.region,
+            "project_id": self.project_id,
+        }
+
     def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any] | None:
         if event["status"] == "error":
             raise AirflowException(event["message"])
         training_pipeline = event["job"]
         custom_job_id = self.hook.extract_custom_job_id_from_training_pipeline(training_pipeline)
-        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
+        context["ti"].xcom_push(key="custom_job_id", value=custom_job_id)
         try:
             model = training_pipeline["model_to_upload"]
             model_id = self.hook.extract_model_id(model)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
             return model
         except KeyError:
             self.log.warning(
@@ -465,6 +476,8 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
             For more information on configuring your service account please visit:
             https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+    :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+        training.
     :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
@@ -578,18 +591,19 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
             sync=True,
+            psc_interface_config=self.psc_interface_config,
         )
 
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        context["ti"].xcom_push(key="custom_job_id", value=custom_job_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
     def invoke_defer(self, context: Context) -> None:
@@ -644,11 +658,12 @@ class CreateCustomContainerTrainingJobOperator(CustomTrainingJobBaseOperator):
             predefined_split_column_name=self.predefined_split_column_name,
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
+            psc_interface_config=self.psc_interface_config,
         )
         custom_container_training_job_obj.wait_for_resource_creation()
         training_pipeline_id: str = custom_container_training_job_obj.name
-        self.xcom_push(context, key="training_id", value=training_pipeline_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        context["ti"].xcom_push(key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_pipeline_id)
         self.defer(
             trigger=CustomContainerTrainingJobTrigger(
                 conn_id=self.gcp_conn_id,
@@ -923,6 +938,8 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
             For more information on configuring your service account please visit:
             https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+    :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+        training.
     :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
@@ -1035,18 +1052,19 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
             sync=True,
+            psc_interface_config=self.psc_interface_config,
         )
 
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        context["ti"].xcom_push(key="custom_job_id", value=custom_job_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
     def invoke_defer(self, context: Context) -> None:
@@ -1102,11 +1120,12 @@ class CreateCustomPythonPackageTrainingJobOperator(CustomTrainingJobBaseOperator
             predefined_split_column_name=self.predefined_split_column_name,
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
+            psc_interface_config=self.psc_interface_config,
         )
         custom_python_training_job_obj.wait_for_resource_creation()
         training_pipeline_id: str = custom_python_training_job_obj.name
-        self.xcom_push(context, key="training_id", value=training_pipeline_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        context["ti"].xcom_push(key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_pipeline_id)
         self.defer(
             trigger=CustomPythonPackageTrainingJobTrigger(
                 conn_id=self.gcp_conn_id,
@@ -1381,6 +1400,8 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             ``projects/{project}/locations/{location}/tensorboards/{tensorboard}``
             For more information on configuring your service account please visit:
             https://cloud.google.com/vertex-ai/docs/experiments/tensorboard-training
+    :param psc_interface_config: Optional. Configuration for Private Service Connect interface used for
+        training.
     :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
@@ -1498,18 +1519,19 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
             sync=True,
+            psc_interface_config=None,
         )
 
         if model:
             result = Model.to_dict(model)
             model_id = self.hook.extract_model_id(result)
-            self.xcom_push(context, key="model_id", value=model_id)
-            VertexAIModelLink.persist(context=context, task_instance=self, model_id=model_id)
+            context["ti"].xcom_push(key="model_id", value=model_id)
+            VertexAIModelLink.persist(context=context, model_id=model_id)
         else:
             result = model  # type: ignore
-        self.xcom_push(context, key="training_id", value=training_id)
-        self.xcom_push(context, key="custom_job_id", value=custom_job_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_id)
+        context["ti"].xcom_push(key="training_id", value=training_id)
+        context["ti"].xcom_push(key="custom_job_id", value=custom_job_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_id)
         return result
 
     def invoke_defer(self, context: Context) -> None:
@@ -1565,11 +1587,12 @@ class CreateCustomTrainingJobOperator(CustomTrainingJobBaseOperator):
             predefined_split_column_name=self.predefined_split_column_name,
             timestamp_split_column_name=self.timestamp_split_column_name,
             tensorboard=self.tensorboard,
+            psc_interface_config=self.psc_interface_config,
         )
         custom_training_job_obj.wait_for_resource_creation()
         training_pipeline_id: str = custom_training_job_obj.name
-        self.xcom_push(context, key="training_id", value=training_pipeline_id)
-        VertexAITrainingLink.persist(context=context, task_instance=self, training_id=training_pipeline_id)
+        context["ti"].xcom_push(key="training_id", value=training_pipeline_id)
+        VertexAITrainingLink.persist(context=context, training_id=training_pipeline_id)
         self.defer(
             trigger=CustomTrainingJobTrigger(
                 conn_id=self.gcp_conn_id,
@@ -1632,26 +1655,6 @@ class DeleteCustomTrainingJobOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
-    @property
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="training_pipeline_id",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def training_pipeline(self):
-        """Alias for ``training_pipeline_id``, used for compatibility (deprecated)."""
-        return self.training_pipeline_id
-
-    @property
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="custom_job_id",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def custom_job(self):
-        """Alias for ``custom_job_id``, used for compatibility (deprecated)."""
-        return self.custom_job_id
-
     def execute(self, context: Context):
         hook = CustomJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -1767,6 +1770,12 @@ class ListCustomTrainingJobOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context):
         hook = CustomJobHook(
             gcp_conn_id=self.gcp_conn_id,
@@ -1783,5 +1792,5 @@ class ListCustomTrainingJobOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        VertexAITrainingPipelinesLink.persist(context=context, task_instance=self)
+        VertexAITrainingPipelinesLink.persist(context=context)
         return [TrainingPipeline.to_dict(result) for result in results]