From 9ba4820c6f8c7b367e3afd227d9ba0f3664eeb37 Mon Sep 17 00:00:00 2001 From: Qiu Qin Date: Mon, 6 Jan 2025 13:35:47 -0500 Subject: [PATCH 1/4] Use object storage client for telemetry. --- ads/telemetry/base.py | 27 +++++++++++++++++---------- ads/telemetry/client.py | 25 +++++++++++++++++-------- 2 files changed, 34 insertions(+), 18 deletions(-) diff --git a/ads/telemetry/base.py b/ads/telemetry/base.py index 03fb2b735..953bc3bd5 100644 --- a/ads/telemetry/base.py +++ b/ads/telemetry/base.py @@ -1,17 +1,18 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- # Copyright (c) 2024 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ import logging -from ads import set_auth +import oci + from ads.common import oci_client as oc -from ads.common.auth import default_signer +from ads.common.auth import default_signer, resource_principal from ads.config import OCI_RESOURCE_PRINCIPAL_VERSION - logger = logging.getLogger(__name__) + + class TelemetryBase: """Base class for Telemetry Client.""" @@ -25,15 +26,21 @@ def __init__(self, bucket: str, namespace: str = None) -> None: namespace : str, optional Namespace of the OCI object storage bucket, by default None. """ + # Use resource principal as authentication method if available, + # however, do not change the ADS authentication if user configured it by set_auth. if OCI_RESOURCE_PRINCIPAL_VERSION: - set_auth("resource_principal") - self._auth = default_signer() - self.os_client = oc.OCIClientFactory(**self._auth).object_storage + self._auth = resource_principal() + else: + self._auth = default_signer() + self.os_client: oci.object_storage.ObjectStorageClient = oc.OCIClientFactory( + **self._auth + ).object_storage self.bucket = bucket self._namespace = namespace self._service_endpoint = None - logger.debug(f"Initialized Telemetry. Namespace: {self.namespace}, Bucket: {self.bucket}") - + logger.debug( + f"Initialized Telemetry. Namespace: {self.namespace}, Bucket: {self.bucket}" + ) @property def namespace(self) -> str: @@ -58,5 +65,5 @@ def service_endpoint(self): Tenancy-specific endpoint. """ if not self._service_endpoint: - self._service_endpoint = self.os_client.base_client.endpoint + self._service_endpoint = str(self.os_client.base_client.endpoint) return self._service_endpoint diff --git a/ads/telemetry/client.py b/ads/telemetry/client.py index 2b3f0fec3..700e1f167 100644 --- a/ads/telemetry/client.py +++ b/ads/telemetry/client.py @@ -1,17 +1,19 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- # Copyright (c) 2024 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ import logging import threading +import traceback import urllib.parse -import requests + +import oci from requests import Response -from .base import TelemetryBase + from ads.config import DEBUG_TELEMETRY +from .base import TelemetryBase logger = logging.getLogger(__name__) @@ -32,7 +34,7 @@ class TelemetryClient(TelemetryBase): >>> import traceback >>> from ads.telemetry.client import TelemetryClient >>> AQUA_BUCKET = os.environ.get("AQUA_BUCKET", "service-managed-models") - >>> AQUA_BUCKET_NS = os.environ.get("AQUA_BUCKET_NS", "ociodscdev") + >>> AQUA_BUCKET_NS = os.environ.get("AQUA_BUCKET_NS", "namespace") >>> telemetry = TelemetryClient(bucket=AQUA_BUCKET, namespace=AQUA_BUCKET_NS) >>> telemetry.record_event_async(category="aqua/service/model", action="create") # records create action >>> telemetry.record_event_async(category="aqua/service/model/create", action="shape", detail="VM.GPU.A10.1") @@ -69,16 +71,23 @@ def record_event( raise ValueError("Please specify the category and the action.") if detail: category, action = f"{category}/{action}", detail + # Here `endpoint`` is for debugging purpose + # For some federated/domain users, the `endpoint` may not be a valid URL endpoint = f"{self.service_endpoint}/n/{self.namespace}/b/{self.bucket}/o/telemetry/{category}/{action}" - headers = {"User-Agent": self._encode_user_agent(**kwargs)} logger.debug(f"Sending telemetry to endpoint: {endpoint}") - signer = self._auth["signer"] - response = requests.head(endpoint, auth=signer, headers=headers) - logger.debug(f"Telemetry status code: {response.status_code}") + + self.os_client.base_client.user_agent = self._encode_user_agent(**kwargs) + response: oci.response.Response = self.os_client.head_object( + namespace_name=self.namespace, + bucket_name=self.bucket, + object_name=f"telemetry/{category}/{action}", + ) + logger.debug(f"Telemetry status: {response.status}") return response except Exception as e: if DEBUG_TELEMETRY: logger.error(f"There is an error recording telemetry: {e}") + traceback.print_exc() def record_event_async( self, category: str = None, action: str = None, detail: str = None, **kwargs From b0600c8ba93556925400cddf196909dd048616c2 Mon Sep 17 00:00:00 2001 From: Qiu Qin Date: Mon, 6 Jan 2025 15:22:33 -0500 Subject: [PATCH 2/4] Handle telemetry request when object does not exist. --- ads/telemetry/client.py | 31 +++++++++++++++++++++---------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/ads/telemetry/client.py b/ads/telemetry/client.py index 700e1f167..7e4f5b697 100644 --- a/ads/telemetry/client.py +++ b/ads/telemetry/client.py @@ -7,9 +7,9 @@ import threading import traceback import urllib.parse +from typing import Optional import oci -from requests import Response from ads.config import DEBUG_TELEMETRY @@ -47,7 +47,7 @@ def _encode_user_agent(**kwargs): def record_event( self, category: str = None, action: str = None, detail: str = None, **kwargs - ) -> Response: + ) -> Optional[int]: """Send a head request to generate an event record. Parameters @@ -64,7 +64,12 @@ def record_event( Returns ------- - Response + int + The status code for the telemetry request. + 200: The the object exists for the telemetry request + 404: The the object does not exist for the telemetry request. + Note that for telemetry purpose, the object does not need to be exist. + `None` will be returned if the telemetry request failed. """ try: if not category or not action: @@ -77,17 +82,23 @@ def record_event( logger.debug(f"Sending telemetry to endpoint: {endpoint}") self.os_client.base_client.user_agent = self._encode_user_agent(**kwargs) - response: oci.response.Response = self.os_client.head_object( - namespace_name=self.namespace, - bucket_name=self.bucket, - object_name=f"telemetry/{category}/{action}", - ) - logger.debug(f"Telemetry status: {response.status}") - return response + try: + response: oci.response.Response = self.os_client.head_object( + namespace_name=self.namespace, + bucket_name=self.bucket, + object_name=f"telemetry/{category}/{action}", + ) + logger.debug(f"Telemetry status: {response.status}") + return response.status + except oci.exceptions.ServiceError as ex: + if ex.status == 404: + return ex.status + raise ex except Exception as e: if DEBUG_TELEMETRY: logger.error(f"There is an error recording telemetry: {e}") traceback.print_exc() + return None def record_event_async( self, category: str = None, action: str = None, detail: str = None, **kwargs From 7450aa924d7baefcac22860e30182a1779b71690 Mon Sep 17 00:00:00 2001 From: Qiu Qin Date: Mon, 6 Jan 2025 15:22:46 -0500 Subject: [PATCH 3/4] Update telemetry unit test. --- .../telemetry/test_telemetry_client.py | 72 +++++++++++-------- 1 file changed, 41 insertions(+), 31 deletions(-) diff --git a/tests/unitary/default_setup/telemetry/test_telemetry_client.py b/tests/unitary/default_setup/telemetry/test_telemetry_client.py index a0efd8cb1..fa5c733c1 100644 --- a/tests/unitary/default_setup/telemetry/test_telemetry_client.py +++ b/tests/unitary/default_setup/telemetry/test_telemetry_client.py @@ -1,59 +1,69 @@ #!/usr/bin/env python -# -*- coding: utf-8 -*- # Copyright (c) 2024 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ +from unittest.mock import patch - -from unittest.mock import patch, PropertyMock +import oci from ads.telemetry.client import TelemetryClient -class TestTelemetryClient: - """Contains unittests for TelemetryClient.""" +TEST_CONFIG = { + "tenancy": "ocid1.tenancy.oc1..unique_ocid", + "user": "ocid1.user.oc1..unique_ocid", + "fingerprint": "00:00:00:00:00:00:00:00:00:00:00:00:00:00:00:00", + "key_file": "//", + "region": "test_region", +} - endpoint = "https://objectstorage.us-ashburn-1.oraclecloud.com" +EXPECTED_ENDPOINT = "https://objectstorage.test_region.oraclecloud.com" - def mocked_requests_head(*args, **kwargs): - class MockResponse: - def __init__(self, status_code): - self.status_code = status_code - return MockResponse(200) +class TestTelemetryClient: + """Contains unittests for TelemetryClient.""" - @patch('requests.head', side_effect=mocked_requests_head) - @patch('ads.telemetry.client.TelemetryClient.service_endpoint', new_callable=PropertyMock, - return_value=endpoint) - def test_telemetry_client_record_event(self, mock_endpoint, mock_head): - """Tests TelemetryClient.record_event() with category/action and path, respectively. - """ + @patch("oci.base_client.BaseClient.request") + @patch("oci.signer.Signer") + def test_telemetry_client_record_event(self, signer, request_call): + """Tests TelemetryClient.record_event() with category/action and path, respectively.""" data = { "cmd": "ads aqua model list", "category": "aqua/service/model", "action": "list", "bucket": "test_bucket", "namespace": "test_namespace", - "value": { - "keyword": "test_service_model_name_or_id" - } + "value": {"keyword": "test_service_model_name_or_id"}, } category = data["category"] action = data["action"] bucket = data["bucket"] namespace = data["namespace"] value = data["value"] - expected_endpoint = f"{self.endpoint}/n/{namespace}/b/{bucket}/o/telemetry/{category}/{action}" - telemetry = TelemetryClient(bucket=bucket, namespace=namespace) + with patch("oci.config.from_file", return_value=TEST_CONFIG): + telemetry = TelemetryClient(bucket=bucket, namespace=namespace) telemetry.record_event(category=category, action=action) telemetry.record_event(category=category, action=action, **value) - expected_headers = [ - {'User-Agent': ''}, - {'User-Agent': 'keyword=test_service_model_name_or_id'} + expected_agent_headers = [ + "", + "keyword=test_service_model_name_or_id", ] - i = 0 - for call_args in mock_head.call_args_list: - args, kwargs = call_args - assert all(endpoint == expected_endpoint for endpoint in args) - assert kwargs['headers'] == expected_headers[i] - i += 1 + + assert len(request_call.call_args_list) == 2 + expected_url = f"{EXPECTED_ENDPOINT}/n/{namespace}/b/{bucket}/o/telemetry/{category}/{action}" + + # Event #1, no user-agent + args, _ = request_call.call_args_list[0] + request: oci.request.Request = args[0] + operation = args[2] + assert request.url == expected_url + assert operation == "head_object" + assert request.header_params["user-agent"] == expected_agent_headers[0] + + # Event #2, with user-agent + args, _ = request_call.call_args_list[1] + request: oci.request.Request = args[0] + operation = args[2] + assert request.url == expected_url + assert operation == "head_object" + assert request.header_params["user-agent"] == expected_agent_headers[1] From 3c07c7b5d30f7ea3917d0a475d4b2abba097af0e Mon Sep 17 00:00:00 2001 From: Qiu Qin Date: Mon, 6 Jan 2025 15:41:15 -0500 Subject: [PATCH 4/4] Update copyright year. --- ads/telemetry/base.py | 2 +- ads/telemetry/client.py | 2 +- tests/unitary/default_setup/telemetry/test_telemetry_client.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/ads/telemetry/base.py b/ads/telemetry/base.py index 953bc3bd5..a104d3e54 100644 --- a/ads/telemetry/base.py +++ b/ads/telemetry/base.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -# Copyright (c) 2024 Oracle and/or its affiliates. +# Copyright (c) 2024, 2025 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ import logging diff --git a/ads/telemetry/client.py b/ads/telemetry/client.py index 7e4f5b697..614925996 100644 --- a/ads/telemetry/client.py +++ b/ads/telemetry/client.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -# Copyright (c) 2024 Oracle and/or its affiliates. +# Copyright (c) 2024, 2025 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ diff --git a/tests/unitary/default_setup/telemetry/test_telemetry_client.py b/tests/unitary/default_setup/telemetry/test_telemetry_client.py index fa5c733c1..d98fcd352 100644 --- a/tests/unitary/default_setup/telemetry/test_telemetry_client.py +++ b/tests/unitary/default_setup/telemetry/test_telemetry_client.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -# Copyright (c) 2024 Oracle and/or its affiliates. +# Copyright (c) 2024, 2025 Oracle and/or its affiliates. # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/ from unittest.mock import patch