diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8be555de3a6b2..f1697038d4015 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -331,7 +331,6 @@ repos: ^airflow\/providers\/google\/cloud\/transfers\/bigquery_to_mysql\.py$| ^airflow\/providers\/amazon\/aws\/transfers\/redshift_to_s3\.py$| ^airflow\/providers\/google\/cloud\/operators\/compute\.py$| - ^airflow\/providers\/google\/cloud\/operators\/vertex_ai\/custom_job\.py$| ^airflow\/providers\/amazon\/aws\/operators\/emr\.py$| ^airflow\/providers\/amazon\/aws\/operators\/eks\.py$ )$ diff --git a/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py b/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py index 0b3d221675e07..e94c39efea1d3 100644 --- a/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +++ b/airflow/providers/google/cloud/operators/vertex_ai/custom_job.py @@ -445,9 +445,19 @@ def __init__( self, *, command: Sequence[str] = [], + region: str, + parent_model: str | None = None, + impersonation_chain: str | Sequence[str] | None = None, + dataset_id: str | None = None, **kwargs, ) -> None: - super().__init__(**kwargs) + super().__init__( + region=region, + parent_model=parent_model, + impersonation_chain=impersonation_chain, + dataset_id=dataset_id, + **kwargs, + ) self.command = command def execute(self, context: Context): @@ -816,9 +826,19 @@ def __init__( *, python_package_gcs_uri: str, python_module_name: str, + region: str, + parent_model: str | None = None, + impersonation_chain: str | Sequence[str] | None = None, + dataset_id: str | None = None, **kwargs, ) -> None: - super().__init__(**kwargs) + super().__init__( + region=region, + parent_model=parent_model, + impersonation_chain=impersonation_chain, + dataset_id=dataset_id, + **kwargs, + ) self.python_package_gcs_uri = python_package_gcs_uri self.python_module_name = python_module_name @@ -1190,9 +1210,19 @@ def __init__( *, script_path: str, requirements: Sequence[str] | None = None, + region: str, + parent_model: str | None = None, + impersonation_chain: str | Sequence[str] | None = None, + dataset_id: str | None = None, **kwargs, ) -> None: - super().__init__(**kwargs) + super().__init__( + region=region, + parent_model=parent_model, + impersonation_chain=impersonation_chain, + dataset_id=dataset_id, + **kwargs, + ) self.requirements = requirements self.script_path = script_path