ODSC-41193: Adds API section to the telemetry. (#140)

Author: mrDzurb

ads/config.py

@@ -14,7 +14,7 @@
 OCI_ODSC_SERVICE_ENDPOINT = os.environ.get("OCI_ODSC_SERVICE_ENDPOINT")
 OCI_IDENTITY_SERVICE_ENDPOINT = os.environ.get("OCI_IDENTITY_SERVICE_ENDPOINT")
 NB_SESSION_COMPARTMENT_OCID = os.environ.get("NB_SESSION_COMPARTMENT_OCID")
-PROJECT_OCID = os.environ.get("PROJECT_OCID")
+PROJECT_OCID = os.environ.get("PROJECT_OCID") or os.environ.get("PIPELINE_PROJECT_OCID")
 NB_SESSION_OCID = os.environ.get("NB_SESSION_OCID")
 USER_OCID = os.environ.get("USER_OCID")
 OCI_RESOURCE_PRINCIPAL_VERSION = os.environ.get("OCI_RESOURCE_PRINCIPAL_VERSION")
@@ -24,17 +24,26 @@
 OCI_REGION_METADATA = os.environ.get("OCI_REGION_METADATA")
 JOB_RUN_OCID = os.environ.get("JOB_RUN_OCID")
 JOB_RUN_COMPARTMENT_OCID = os.environ.get("JOB_RUN_COMPARTMENT_OCID")
+PIPELINE_RUN_OCID = os.environ.get("PIPELINE_RUN_OCID")
+PIPELINE_RUN_COMPARTMENT_OCID = os.environ.get("PIPELINE_RUN_COMPARTMENT_OCID")
+PIPELINE_COMPARTMENT_OCID = os.environ.get("PIPELINE_COMPARTMENT_OCID")
+
 CONDA_BUCKET_NAME = os.environ.get("CONDA_BUCKET_NAME", "service-conda-packs")
 CONDA_BUCKET_NS = os.environ.get("CONDA_BUCKET_NS", "id19sfcrra6z")
 OCI_RESOURCE_PRINCIPAL_RPT_ENDPOINT = os.environ.get(
     "OCI_RESOURCE_PRINCIPAL_RPT_ENDPOINT"
 )
-COMPARTMENT_OCID = NB_SESSION_COMPARTMENT_OCID or JOB_RUN_COMPARTMENT_OCID
+COMPARTMENT_OCID = (
+    NB_SESSION_COMPARTMENT_OCID
+    or JOB_RUN_COMPARTMENT_OCID
+    or PIPELINE_RUN_COMPARTMENT_OCID
+)
 MD_OCID = os.environ.get("MD_OCID")
 DATAFLOW_RUN_OCID = os.environ.get("DATAFLOW_RUN_ID")
+
 RESOURCE_OCID = (
-    NB_SESSION_OCID or JOB_RUN_OCID or MD_OCID
-)  # We can add DATAFLOW_RUN_OCID here. Needs impact analysis
+    NB_SESSION_OCID or JOB_RUN_OCID or MD_OCID or PIPELINE_RUN_OCID or DATAFLOW_RUN_OCID
+)
 NO_CONTAINER = os.environ.get("NO_CONTAINER")
 
 

ads/telemetry/telemetry.py

@@ -1,22 +1,26 @@
 #!/usr/bin/env python
-# -*- coding: utf-8; -*-
+# -*- coding: utf-8 -*--
 
-# Copyright (c) 2022 Oracle and/or its affiliates.
+# Copyright (c) 2022, 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
-from ads import __version__
+
+import os
 from enum import Enum, auto
+from typing import Any, Dict, Optional
 
 import ads.config
+from ads import __version__
+from ads.common import logger
 
+LIBRARY = "Oracle-ads"
 EXTRA_USER_AGENT_INFO = "EXTRA_USER_AGENT_INFO"
 USER_AGENT_KEY = "additional_user_agent"
-ENV_MD_OCID = "MD_OCID"
 UNKNOWN = "UNKNOWN"
 
 
 class Surface(Enum):
     """
-    An Enum class for labeling the surface where ADS is being used
+    An Enum class for labeling the surface where ADS is being used.
     """
 
     WORKSTATION = auto()
@@ -25,6 +29,7 @@ class Surface(Enum):
     DATASCIENCE_MODEL_DEPLOYMENT = auto()
     DATAFLOW = auto()
     OCI_SERVICE = auto()
+    DATASCIENCE_PIPELINE = auto()
 
     @classmethod
     def surface(cls):
@@ -43,12 +48,33 @@ def surface(cls):
                 surface = cls.DATASCIENCE_MODEL_DEPLOYMENT
             elif ads.config.DATAFLOW_RUN_OCID:
                 surface = cls.DATAFLOW
+            elif ads.config.PIPELINE_RUN_OCID:
+                surface = cls.DATASCIENCE_PIPELINE
         return surface
 
 
-def update_oci_client_config(config={}):
-    if not config.get(USER_AGENT_KEY):
-        config[
-            USER_AGENT_KEY
-        ] = f"Oracle-ads/version={__version__}/surface={Surface.surface().name}"  # To be enabled in future - /api={os.environ.get(EXTRA_USER_AGENT_INFO,UNKNOWN)}"
+def update_oci_client_config(config: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    """Adds user agent information to the config if it is not setup yet.
+
+    Returns
+    -------
+    Dict
+        The updated configuration.
+    """
+
+    try:
+        config = config or {}
+        if not config.get(USER_AGENT_KEY):
+            config.update(
+                {
+                    USER_AGENT_KEY: (
+                        f"{LIBRARY}/version={__version__}/"
+                        f"surface={Surface.surface().name}/"
+                        f"api={os.environ.get(EXTRA_USER_AGENT_INFO,UNKNOWN) or UNKNOWN}"
+                    )
+                }
+            )
+    except Exception as ex:
+        logger.debug(ex)
+
     return config

tests/unitary/default_setup/telemetry/test_agent.py

@@ -3,12 +3,19 @@
 # Copyright (c) 2022, 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import importlib
 from unittest.mock import patch
+
 import pytest
-from pytest import MonkeyPatch
-import importlib
+# from pytest import MonkeyPatch
 
 import ads
+from ads.telemetry.telemetry import (
+    EXTRA_USER_AGENT_INFO,
+    LIBRARY,
+    UNKNOWN,
+    USER_AGENT_KEY,
+)
 
 
 class TestUserAgent:
@@ -31,47 +38,77 @@ def test_user_agent_api_keys_using_test_profile(self, mock_signer):
             auth_info = ads.auth.api_keys("test_path", "TEST_PROFILE")
             assert (
                 auth_info["config"].get("additional_user_agent")
-                == f"Oracle-ads/version={ads.__version__}/surface=WORKSTATION"
+                == f"{LIBRARY}/version={ads.__version__}/surface=WORKSTATION/api={UNKNOWN}"
             )
 
     @patch("oci.auth.signers.get_resource_principals_signer")
-    def test_user_agent_rp(self, mock_signer):
+    def test_user_agent_rp(self, mock_signer, monkeypatch):
+        monkeypatch.delenv("OCI_RESOURCE_PRINCIPAL_VERSION", raising=False)
         importlib.reload(ads.config)
         importlib.reload(ads.telemetry)
         auth_info = ads.auth.resource_principal()
         assert (
             auth_info["config"].get("additional_user_agent")
-            == f"Oracle-ads/version={ads.__version__}/surface=WORKSTATION"
+            == f"{LIBRARY}/version={ads.__version__}/surface=WORKSTATION/api={UNKNOWN}"
         )
 
     @patch("oci.signer.load_private_key_from_file")
-    def test_user_agent_default_signer(self, mock_load_key_file):
-        monkeypatch = MonkeyPatch()
+    def test_user_agent_default_signer(self, mock_load_key_file, monkeypatch):
+        # monkeypatch = MonkeyPatch()
         monkeypatch.delenv("OCI_RESOURCE_PRINCIPAL_VERSION", raising=False)
         importlib.reload(ads.config)
         with patch("oci.config.from_file", return_value=self.test_config):
             auth_info = ads.auth.default_signer()
             assert (
                 auth_info["config"].get("additional_user_agent")
-                == f"Oracle-ads/version={ads.__version__}/surface=WORKSTATION"
+                == f"{LIBRARY}/version={ads.__version__}/surface=WORKSTATION/api={UNKNOWN}"
             )
 
     @pytest.mark.parametrize(
-        "RESOURCE_KEY, USER_AGENT_VALUE",
+        "INPUT_DATA, EXPECTED_RESULT",
         [
-            ("MD_OCID", "DATASCIENCE_MODEL_DEPLOYMENT"),
-            ("JOB_RUN_OCID", "DATASCIENCE_JOB"),
-            ("NB_SESSION_OCID", "DATASCIENCE_NOTEBOOK"),
-            ("DATAFLOW_RUN_ID", "DATAFLOW"),
+            (
+                {"RESOURCE_KEY": "MD_OCID", EXTRA_USER_AGENT_INFO: ""},
+                {
+                    "USER_AGENT_VALUE": "DATASCIENCE_MODEL_DEPLOYMENT",
+                    EXTRA_USER_AGENT_INFO: UNKNOWN,
+                },
+            ),
+            (
+                {"RESOURCE_KEY": "JOB_RUN_OCID", EXTRA_USER_AGENT_INFO: "test_api"},
+                {
+                    "USER_AGENT_VALUE": "DATASCIENCE_JOB",
+                    EXTRA_USER_AGENT_INFO: "test_api",
+                },
+            ),
+            (
+                {"RESOURCE_KEY": "NB_SESSION_OCID", EXTRA_USER_AGENT_INFO: None},
+                {
+                    "USER_AGENT_VALUE": "DATASCIENCE_NOTEBOOK",
+                    EXTRA_USER_AGENT_INFO: UNKNOWN,
+                },
+            ),
+            (
+                {
+                    "RESOURCE_KEY": "DATAFLOW_RUN_ID",
+                    EXTRA_USER_AGENT_INFO: "some_class&some_method",
+                },
+                {
+                    "USER_AGENT_VALUE": "DATAFLOW",
+                    EXTRA_USER_AGENT_INFO: "some_class&some_method",
+                },
+            ),
         ],
     )
     @patch("oci.signer.load_private_key_from_file")
     def test_user_agent_default_signer_known_resources(
-        self, mock_load_key_file, RESOURCE_KEY, USER_AGENT_VALUE
+        self,mock_load_key_file, monkeypatch, INPUT_DATA, EXPECTED_RESULT
     ):
-        monkeypatch = MonkeyPatch()
+        # monkeypatch = MonkeyPatch()
         monkeypatch.setenv("OCI_RESOURCE_PRINCIPAL_VERSION", "1.1")
-        monkeypatch.setenv(RESOURCE_KEY, "1234")
+        monkeypatch.setenv(INPUT_DATA["RESOURCE_KEY"], "1234")
+        if INPUT_DATA[EXTRA_USER_AGENT_INFO] is not None:
+            monkeypatch.setenv(EXTRA_USER_AGENT_INFO, INPUT_DATA[EXTRA_USER_AGENT_INFO])
 
         importlib.reload(ads.config)
         importlib.reload(ads)
@@ -82,17 +119,18 @@ def test_user_agent_default_signer_known_resources(
             auth_info = ads.auth.default_signer()
             assert (
                 auth_info["config"].get("additional_user_agent")
-                == f"Oracle-ads/version={ads.__version__}/surface={USER_AGENT_VALUE}"
+                == f"{LIBRARY}/version={ads.__version__}/surface={EXPECTED_RESULT['USER_AGENT_VALUE']}/api={EXPECTED_RESULT[EXTRA_USER_AGENT_INFO]}"
             )
         monkeypatch.delenv("OCI_RESOURCE_PRINCIPAL_VERSION", raising=False)
-        monkeypatch.delenv(RESOURCE_KEY, raising=False)
+        monkeypatch.delenv(INPUT_DATA["RESOURCE_KEY"], raising=False)
+        monkeypatch.delenv(EXTRA_USER_AGENT_INFO, raising=False)
 
     @patch("oci.signer.Signer")
     def test_user_agent_default_singer_ociservice(
         self,
         mock_signer,
+        monkeypatch
     ):
-        monkeypatch = MonkeyPatch()
         monkeypatch.setenv("OCI_RESOURCE_PRINCIPAL_VERSION", "1.1")
 
         importlib.reload(ads.config)
@@ -102,6 +140,6 @@ def test_user_agent_default_singer_ociservice(
             auth_info = ads.auth.default_signer()
             assert (
                 auth_info["config"].get("additional_user_agent")
-                == f"Oracle-ads/version={ads.__version__}/surface=OCI_SERVICE"
+                == f"{LIBRARY}/version={ads.__version__}/surface=OCI_SERVICE/api={UNKNOWN}"
             )
         monkeypatch.delenv("OCI_RESOURCE_PRINCIPAL_VERSION", raising=False)