diff --git a/cloud_pipelines_backend/instrumentation/gcs_tracing.py b/cloud_pipelines_backend/instrumentation/gcs_tracing.py new file mode 100644 index 0000000..ffd90a0 --- /dev/null +++ b/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 diff --git a/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py b/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py index 57f959f..de76f17 100644 --- a/cloud_pipelines_backend/launchers/google_kubernetes_launchers.py +++ b/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, ) diff --git a/cloud_pipelines_backend/launchers/kubernetes_launchers.py b/cloud_pipelines_backend/launchers/kubernetes_launchers.py index 9b1a661..2015594 100644 --- a/cloud_pipelines_backend/launchers/kubernetes_launchers.py +++ b/cloud_pipelines_backend/launchers/kubernetes_launchers.py @@ -724,7 +724,7 @@ 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, @@ -732,9 +732,7 @@ def __init__( 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,