feat: add support of google-cloud-storage v3 dependency

PiperOrigin-RevId: 825128354

Author: speedstorm1

google/cloud/aiplatform/utils/gcs_utils.py

@@ -17,14 +17,18 @@
 
 import datetime
 import glob
+# Version detection and compatibility layer for google-cloud-storage v2/v3
+from importlib.metadata import version as get_version
 import logging
 import os
 import pathlib
 import tempfile
 from typing import Optional, TYPE_CHECKING
+import warnings
 
 from google.auth import credentials as auth_credentials
 from google.cloud import storage
+from packaging.version import Version
 
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform.utils import resource_manager_utils
@@ -35,6 +39,66 @@
 _logger = logging.getLogger(__name__)
 
 
+# Detect google-cloud-storage version once at module load
+try:
+  _GCS_VERSION = Version(get_version("google-cloud-storage"))
+except Exception:
+  # Fallback if version detection fails (should not happen in normal use)
+  _GCS_VERSION = Version("2.0.0")
+
+_USE_FROM_URI = _GCS_VERSION >= Version("3.0.0")
+
+# Warn users on v2 about upcoming deprecation
+if _GCS_VERSION < Version("3.0.0"):
+  warnings.warn(
+      "Support for google-cloud-storage < 3.0.0 will be removed in a future"
+      " version of google-cloud-aiplatform. Please upgrade to"
+      " google-cloud-storage >= 3.0.0.",
+      FutureWarning,
+      stacklevel=2,
+  )
+
+
+def blob_from_uri(uri: str, client: storage.Client) -> storage.Blob:
+  """Create a Blob from a GCS URI, compatible with v2 and v3.
+
+  This function provides compatibility across google-cloud-storage versions:
+  - v3.x: Uses Blob.from_uri()
+  - v2.x: Uses Blob.from_string() (deprecated in v3)
+
+  Args:
+      uri: GCS URI (e.g., 'gs://bucket/path/to/blob')
+      client: Storage client instance
+
+  Returns:
+      storage.Blob: Blob instance
+  """
+  if _USE_FROM_URI:
+    return storage.Blob.from_uri(uri, client=client)
+  else:
+    return storage.Blob.from_string(uri, client=client)
+
+
+def bucket_from_uri(uri: str, client: storage.Client) -> storage.Bucket:
+  """Create a Bucket from a GCS URI, compatible with v2 and v3.
+
+  This function provides compatibility across google-cloud-storage versions:
+  - v3.x: Uses Bucket.from_uri()
+  - v2.x: Uses Bucket.from_string() (deprecated in v3)
+
+  Args:
+      uri: GCS bucket URI (e.g., 'gs://bucket-name')
+      client: Storage client instance
+
+  Returns:
+      storage.Bucket: Bucket instance
+  """
+  if _USE_FROM_URI:
+    return storage.Bucket.from_uri(uri, client=client)
+  else:
+    return storage.Bucket.from_string(uri, client=client)
+
+
 def upload_to_gcs(
     source_path: str,
     destination_uri: str,
@@ -79,16 +143,20 @@ def upload_to_gcs(
             destination_file_uri = (
                 destination_uri.rstrip("/") + "/" + source_file_relative_posix_path
             )
-            _logger.debug(f'Uploading "{source_file_path}" to "{destination_file_uri}"')
-            destination_blob = storage.Blob.from_string(
+            _logger.debug(
+                'Uploading "%s" to "%s"', source_file_path, destination_file_uri
+            )
+            destination_blob = blob_from_uri(
                 destination_file_uri, client=storage_client
             )
             destination_blob.upload_from_filename(filename=source_file_path)
     else:
         source_file_path = source_path
         destination_file_uri = destination_uri
-        _logger.debug(f'Uploading "{source_file_path}" to "{destination_file_uri}"')
-        destination_blob = storage.Blob.from_string(
+        _logger.debug(
+            'Uploading "%s" to "%s"', source_file_path, destination_file_uri
+        )
+        destination_blob = blob_from_uri(
             destination_file_uri, client=storage_client
         )
         destination_blob.upload_from_filename(filename=source_file_path)
@@ -234,7 +302,7 @@ def create_gcs_bucket_for_pipeline_artifacts_if_it_does_not_exist(
         credentials=credentials,
     )
 
-    pipelines_bucket = storage.Bucket.from_string(
+    pipelines_bucket = bucket_from_uri(
         uri=output_artifacts_gcs_dir,
         client=storage_client,
     )
@@ -294,9 +362,11 @@ def download_file_from_gcs(
     credentials = credentials or initializer.global_config.credentials
 
     storage_client = storage.Client(project=project, credentials=credentials)
-    source_blob = storage.Blob.from_string(source_file_uri, client=storage_client)
+    source_blob = blob_from_uri(source_file_uri, client=storage_client)
 
-    _logger.debug(f'Downloading "{source_file_uri}" to "{destination_file_path}"')
+    _logger.debug(
+        'Downloading "%s" to "%s"', source_file_uri, destination_file_path
+    )
 
     source_blob.download_to_filename(filename=destination_file_path)
 
@@ -351,36 +421,33 @@ def download_from_gcs(
 def _upload_pandas_df_to_gcs(
     df: "pandas.DataFrame", upload_gcs_path: str, file_format: str = "jsonl"
 ) -> None:
-    """Uploads the provided Pandas DataFrame to a GCS bucket.
+  """Uploads the provided Pandas DataFrame to a GCS bucket.
 
-    Args:
-        df (pandas.DataFrame):
-            Required. The Pandas DataFrame to upload.
-        upload_gcs_path (str):
-            Required. The GCS path to upload the data file.
-        file_format (str):
-            Required. The format to export the DataFrame to. Currently
-            only JSONL is supported.
+  Args:
+      df (pandas.DataFrame): Required. The Pandas DataFrame to upload.
+      upload_gcs_path (str): Required. The GCS path to upload the data file.
+      file_format (str): Required. The format to export the DataFrame to.
+        Currently only JSONL is supported.
 
-    Raises:
-        ValueError: When a file format other than JSONL is provided.
-    """
+  Raises:
+      ValueError: When a file format other than JSONL is provided.
+  """
 
-    with tempfile.TemporaryDirectory() as temp_dir:
-        local_dataset_path = os.path.join(temp_dir, "dataset.jsonl")
+  with tempfile.TemporaryDirectory() as temp_dir:
+    local_dataset_path = os.path.join(temp_dir, "dataset.jsonl")
 
-        if file_format == "jsonl":
-            df.to_json(path_or_buf=local_dataset_path, orient="records", lines=True)
-        else:
-            raise ValueError(f"Unsupported file format: {file_format}")
+    if file_format == "jsonl":
+      df.to_json(path_or_buf=local_dataset_path, orient="records", lines=True)
+    else:
+      raise ValueError(f"Unsupported file format: {file_format}")
 
-        storage_client = storage.Client(
-            project=initializer.global_config.project,
-            credentials=initializer.global_config.credentials,
-        )
-        storage.Blob.from_string(
-            uri=upload_gcs_path, client=storage_client
-        ).upload_from_filename(filename=local_dataset_path)
+    storage_client = storage.Client(
+        project=initializer.global_config.project,
+        credentials=initializer.global_config.credentials,
+    )
+    blob_from_uri(
+        uri=upload_gcs_path, client=storage_client
+    ).upload_from_filename(filename=local_dataset_path)
 
 
 def validate_gcs_path(gcs_path: str) -> None:

google/cloud/aiplatform/utils/yaml_utils.py

@@ -22,6 +22,7 @@
 from google.auth import transport
 from google.cloud import storage
 from google.cloud.aiplatform.constants import pipeline as pipeline_constants
+from google.cloud.aiplatform.utils.gcs_utils import blob_from_uri
 
 # Pattern for an Artifact Registry URL.
 _VALID_AR_URL = pipeline_constants._VALID_AR_URL
@@ -98,7 +99,7 @@ def _load_yaml_from_gs_uri(
     """
     yaml = _maybe_import_yaml()
     storage_client = storage.Client(project=project, credentials=credentials)
-    blob = storage.Blob.from_string(uri, storage_client)
+    blob = blob_from_uri(uri, storage_client)
     return yaml.safe_load(blob.download_as_bytes())
 
 

setup.py

@@ -306,8 +306,8 @@
         "proto-plus >= 1.22.3, <2.0.0",
         "protobuf>=3.20.2,<7.0.0,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5",
         "packaging >= 14.3",
-        "google-cloud-storage >= 1.32.0, < 3.0.0; python_version<'3.13'",
-        "google-cloud-storage >= 2.10.0, < 3.0.0; python_version>='3.13'",
+        "google-cloud-storage >= 1.32.0, < 4.0.0; python_version<'3.13'",
+        "google-cloud-storage >= 2.10.0, < 4.0.0; python_version>='3.13'",
         "google-cloud-bigquery >= 1.15.0, < 4.0.0, !=3.20.0",
         "google-cloud-resource-manager >= 1.3.3, < 3.0.0",
         "shapely < 3.0.0",

testing/constraints-3.10.txt

@@ -6,7 +6,7 @@ google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus==1.22.3
 protobuf
 mock==4.0.2
-google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
+google-cloud-storage==3.0.0 # Updated to v3.x, backward compatible with v2.x via wrapper
 packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 grpcio-testing==1.34.0
 mlflow==2.16.0 # Pinned to speed up installation

testing/constraints-3.11.txt

@@ -6,7 +6,7 @@ google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus
 protobuf
 mock==4.0.2
-google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
+google-cloud-storage==3.0.0 # Updated to v3.x, backward compatible with v2.x via wrapper
 packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.5.0 # Pinned until 2.9.3 is verified for Ray tests

testing/constraints-3.12.txt

@@ -5,7 +5,7 @@ google-api-core==2.21.0 # Tests google-api-core with rest async support
 google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus
 mock==4.0.2
-google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
+google-cloud-storage==3.0.0 # Updated to v3.x, backward compatible with v2.x via wrapper
 packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.5.0 # Pinned until 2.9.3 is verified for Ray tests

testing/constraints-3.13.txt

@@ -5,7 +5,7 @@ google-api-core==2.21.0 # Tests google-api-core with rest async support
 google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus
 mock==4.0.2
-google-cloud-storage==2.10.0 # Increased for kfp 2.0 compatibility
+google-cloud-storage==3.0.0 # Updated to v3.x, backward compatible with v2.x via wrapper
 packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests
 ray==2.5.0 # Pinned until 2.9.3 is verified for Ray tests

testing/constraints-3.8.txt

@@ -7,7 +7,7 @@ google-auth==2.14.1 # Tests google-auth without rest async support
 proto-plus==1.22.3
 protobuf
 mock==4.0.2
-google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
+google-cloud-storage==3.0.0 # Updated to v3.x, backward compatible with v2.x via wrapper
 packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 grpcio-testing==1.34.0
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests

testing/constraints-3.9.txt

@@ -6,7 +6,7 @@ google-auth==2.35.0 # Tests google-auth with rest async support
 proto-plus==1.22.3
 protobuf
 mock==4.0.2
-google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
+google-cloud-storage==3.0.0 # Updated to v3.x, backward compatible with v2.x via wrapper
 packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 grpcio-testing==1.34.0
 pytest-xdist==3.3.1 # Pinned to unbreak unit tests

testing/constraints-ray-2.33.0.txt

@@ -4,7 +4,7 @@ google-api-core
 proto-plus==1.22.3
 protobuf
 mock==4.0.2
-google-cloud-storage==2.2.1 # Increased for kfp 2.0 compatibility
+google-cloud-storage==3.0.0 # Updated to v3.x, backward compatible with v2.x via wrapper
 packaging==24.1 # Increased to unbreak canonicalize_version error (b/377774673)
 grpcio-testing==1.34.0
 mlflow==1.30.1 # Pinned to speed up installation