diff --git a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py index 98e5e6a4ffe0e..86b3ae017c080 100644 --- a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +++ b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py @@ -123,6 +123,8 @@ def __init__( forecast_horizon: int, data_granularity_unit: str, data_granularity_count: int, + display_name: str, + model_display_name: str | None = None, optimization_objective: str | None = None, column_specs: dict[str, str] | None = None, column_transformations: list[dict[str, dict[str, str]]] | None = None, @@ -145,7 +147,12 @@ def __init__( **kwargs, ) -> None: super().__init__( - region=region, impersonation_chain=impersonation_chain, parent_model=parent_model, **kwargs + display_name=display_name, + model_display_name=model_display_name, + region=region, + impersonation_chain=impersonation_chain, + parent_model=parent_model, + **kwargs, ) self.dataset_id = dataset_id self.target_column = target_column diff --git a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py index 81a4090407b7c..33ab07d075396 100644 --- a/providers/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +++ b/providers/src/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py @@ -509,6 +509,8 @@ def __init__( *, command: Sequence[str] = [], region: str, + display_name: str, + model_display_name: str | None = None, parent_model: str | None = None, impersonation_chain: str | Sequence[str] | None = None, dataset_id: str | None = None, @@ -517,6 +519,8 @@ def __init__( **kwargs, ) -> None: super().__init__( + display_name=display_name, + model_display_name=model_display_name, region=region, parent_model=parent_model, impersonation_chain=impersonation_chain, @@ -962,6 +966,8 @@ def __init__( python_package_gcs_uri: str, python_module_name: str, region: str, + display_name: str, + model_display_name: str | None = None, parent_model: str | None = None, impersonation_chain: str | Sequence[str] | None = None, dataset_id: str | None = None, @@ -970,6 +976,8 @@ def __init__( **kwargs, ) -> None: super().__init__( + display_name=display_name, + model_display_name=model_display_name, region=region, parent_model=parent_model, impersonation_chain=impersonation_chain, @@ -1423,6 +1431,8 @@ def __init__( script_path: str, requirements: Sequence[str] | None = None, region: str, + display_name: str, + model_display_name: str | None = None, parent_model: str | None = None, impersonation_chain: str | Sequence[str] | None = None, dataset_id: str | None = None, @@ -1431,6 +1441,8 @@ def __init__( **kwargs, ) -> None: super().__init__( + display_name=display_name, + model_display_name=model_display_name, region=region, parent_model=parent_model, impersonation_chain=impersonation_chain,