Trace GCS upload/download duration and bytes

cloud_pipelines_backend/instrumentation/gcs_tracing.py

@@ -0,0 +1,64 @@
+"""GCS storage provider that emits one OTel span per upload/download so transfer duration is measurable."""
+
+from __future__ import annotations
+
+import collections.abc
+import contextlib
+
+from opentelemetry import trace
+from opentelemetry.trace import StatusCode
+
+from cloud_pipelines.orchestration.storage_providers import google_cloud_storage
+
+_tracer = trace.get_tracer("tangle.storage")
+
+
+@contextlib.contextmanager
+def _gcs_operation_span(
+    operation: str, uri: str
+) -> collections.abc.Iterator[trace.Span]:
+    with _tracer.start_as_current_span(
+        f"gcs.{operation}",
+        attributes={"gcs.operation": operation, "gcs.uri": uri},
+    ) as span:
+        try:
+            yield span
+        except Exception as exception:
+            span.set_status(StatusCode.ERROR)
+            span.record_exception(exception)
+            raise
+
+
+class TracingGoogleCloudStorageProvider(
+    google_cloud_storage.GoogleCloudStorageProvider
+):
+    def upload(
+        self,
+        source_path: str,
+        destination_uri: google_cloud_storage.GoogleCloudStorageUri,
+    ) -> None:
+        with _gcs_operation_span("upload", destination_uri.uri):
+            super().upload(source_path, destination_uri)
+
+    def upload_bytes(
+        self, data: bytes, destination_uri: google_cloud_storage.GoogleCloudStorageUri
+    ) -> None:
+        with _gcs_operation_span("upload_bytes", destination_uri.uri) as span:
+            span.set_attribute("gcs.bytes", len(data))
+            super().upload_bytes(data, destination_uri)
+
+    def download(
+        self,
+        source_uri: google_cloud_storage.GoogleCloudStorageUri,
+        destination_path: str,
+    ) -> None:
+        with _gcs_operation_span("download", source_uri.uri):
+            super().download(source_uri, destination_path)
+
+    def download_bytes(
+        self, source_uri: google_cloud_storage.GoogleCloudStorageUri
+    ) -> bytes:
+        with _gcs_operation_span("download_bytes", source_uri.uri) as span:
+            data = super().download_bytes(source_uri)
+            span.set_attribute("gcs.bytes", len(data))
+            return data

cloud_pipelines_backend/launchers/google_kubernetes_launchers.py

@@ -2,9 +2,8 @@
 
 from kubernetes import client as k8s_client_lib
 
-from cloud_pipelines.orchestration.storage_providers import google_cloud_storage
-
 from . import kubernetes_launchers
+from ..instrumentation import gcs_tracing
 
 if typing.TYPE_CHECKING:
     from google.cloud import storage
@@ -44,9 +43,7 @@ def __init__(
             pod_labels=pod_labels,
             pod_annotations={"gke-gcsfuse/volumes": "true"} | (pod_annotations or {}),
             pod_postprocessor=final_pod_postporocessor,
-            _storage_provider=google_cloud_storage.GoogleCloudStorageProvider(
-                gcs_client
-            ),
+            _storage_provider=gcs_tracing.TracingGoogleCloudStorageProvider(gcs_client),
             _create_volume_and_volume_mount=kubernetes_launchers._create_volume_and_volume_mount_google_cloud_storage,
         )
 
@@ -87,8 +84,6 @@ def __init__(
             pod_annotations={"gke-gcsfuse/volumes": "true"} | (pod_annotations or {}),
             pod_postprocessor=final_pod_postporocessor,
             always_launch_jobs=always_launch_jobs,
-            _storage_provider=google_cloud_storage.GoogleCloudStorageProvider(
-                gcs_client
-            ),
+            _storage_provider=gcs_tracing.TracingGoogleCloudStorageProvider(gcs_client),
             _create_volume_and_volume_mount=kubernetes_launchers._create_volume_and_volume_mount_google_cloud_storage,
         )

cloud_pipelines_backend/launchers/kubernetes_launchers.py

@@ -724,17 +724,15 @@ def __init__(
             pod_postprocessors.append(pod_postprocessor)
         final_pod_postporocessor = _create_pod_postprocessor_stack(pod_postprocessors)
 
-        from cloud_pipelines.orchestration.storage_providers import google_cloud_storage
+        from ..instrumentation import gcs_tracing
 
         super().__init__(
             namespace=namespace,
             service_account_name=service_account_name,
             api_client=api_client,
             request_timeout=request_timeout,
             pod_name_prefix=pod_name_prefix,
-            _storage_provider=google_cloud_storage.GoogleCloudStorageProvider(
-                gcs_client
-            ),
+            _storage_provider=gcs_tracing.TracingGoogleCloudStorageProvider(gcs_client),
             pod_labels=pod_labels,
             pod_annotations={"gke-gcsfuse/volumes": "true"} | (pod_annotations or {}),
             pod_postprocessor=final_pod_postporocessor,