From bf83ec4d825090db85aa509bd2c4ac36875abb5c Mon Sep 17 00:00:00 2001 From: Eduardo Blancas Date: Fri, 12 Jan 2024 10:04:33 -0600 Subject: [PATCH] faster telemetry + simpler initialization (#82) * more efficient system info * adds classmethod to initialize telemetry with default settings * lint * changelog and docs * update doc * ci --- CHANGELOG.md | 4 ++ doc/telemetry.md | 19 ++++---- setup.py | 1 + src/ploomber_core/telemetry/system_info.py | 34 +++++++++++++++ src/ploomber_core/telemetry/telemetry.py | 50 ++++++++++++---------- tests/telemetry/test_system_info.py | 15 +++++++ tests/telemetry/test_telemetry.py | 9 ++++ 7 files changed, 101 insertions(+), 31 deletions(-) create mode 100644 tests/telemetry/test_system_info.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 093a1bf..5967335 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## 0.2.20dev +* [Feature] Obtain system info once for faster initialization +* [Feature] Add `Telemetry.from_package` for simpler initialization + + ## 0.2.19 (2023-12-12) * [Fix] `@modify_exceptions` decorator now compatible with `ClickException` ([#81](https://github.com/ploomber/core/issues/81)) diff --git a/doc/telemetry.md b/doc/telemetry.md index a955b7e..5559d02 100644 --- a/doc/telemetry.md +++ b/doc/telemetry.md @@ -27,16 +27,17 @@ from ploomber_core.telemetry import Telemetry Initialize it with the API key, name of the package, and version. Here's an example: +```{versionadded} 0.2.20 +`Telemetry.from_package` +``` + ```{code-cell} ipython3 -from ploomber_core import __version__ - -telemetry = Telemetry( - api_key="SOMEKEY", - # change for the actual name - package_name="ploomber-core", - # change for the actual version - version=__version__, -) +telemetry = Telemetry.from_package(package_name="ploomber-core") +``` + +```{note} +`Telemetry.from_package` is the simplest way to initialize the telemetry object, you might use the constructor directly if you want to customize +the configuration. ``` To log call functions: diff --git a/setup.py b/setup.py index 96aa91e..e82d350 100644 --- a/setup.py +++ b/setup.py @@ -17,6 +17,7 @@ "click", "pyyaml", "posthog", + 'importlib-metadata;python_version<"3.8"', ] DEV = [ diff --git a/src/ploomber_core/telemetry/system_info.py b/src/ploomber_core/telemetry/system_info.py index bcca963..75a5b94 100644 --- a/src/ploomber_core/telemetry/system_info.py +++ b/src/ploomber_core/telemetry/system_info.py @@ -4,6 +4,12 @@ import sys +try: + import importlib.metadata as importlib_metadata +except ModuleNotFoundError: + import importlib_metadata + + def python_version(): py_version = sys.version_info return f"{py_version.major}.{py_version.minor}.{py_version.micro}" @@ -106,3 +112,31 @@ def is_airflow(): def is_argo(): """Returns: True for Argo env""" return "ARGO_AGENT_TASK_WORKERS" in os.environ or "ARGO_KUBELET_PORT" in os.environ + + +def safe_call(function): + try: + return function() + except Exception: + return None + + +def get_package_version(package_name): + try: + return importlib_metadata.version(package_name) + except importlib_metadata.PackageNotFoundError: + return None + + +def get_system_info(): + return { + "os": safe_call(get_os), + "python_version": safe_call(python_version), + "env": safe_call(get_env), + "is_docker": safe_call(is_docker), + "is_colab": safe_call(is_colab), + "is_paperspace": safe_call(is_paperspace), + "is_slurm": safe_call(is_slurm), + "is_airflow": safe_call(is_airflow), + "is_argo": safe_call(is_argo), + } diff --git a/src/ploomber_core/telemetry/telemetry.py b/src/ploomber_core/telemetry/telemetry.py index 18d08ed..ac9816a 100644 --- a/src/ploomber_core/telemetry/telemetry.py +++ b/src/ploomber_core/telemetry/telemetry.py @@ -46,17 +46,7 @@ from ploomber_core.telemetry import validate_inputs from ploomber_core.config import Config -from ploomber_core.telemetry.system_info import ( - get_env, - get_os, - is_airflow, - is_argo, - is_colab, - is_docker, - is_paperspace, - is_slurm, - python_version, -) +from ploomber_core.telemetry.system_info import get_system_info, get_package_version TELEMETRY_VERSION = "0.5" DEFAULT_HOME_DIR = str(Path.home() / ".ploomber") @@ -69,6 +59,9 @@ logging.getLogger("posthog").disabled = True +SYSTEM_INFO = get_system_info() + + class UserSettings(Config): """User-customizable settings""" @@ -420,6 +413,16 @@ def __init__(self, api_key, package_name, version): self.package_name = package_name self.version = version + @classmethod + def from_package(cls, package_name): + """ + Initialize a Telemetry client with the default configuration for + a package with the given name + """ + default_api_key = "phc_P9SpSeypyPwxrMdFn2edOOEooQioF2axppyEeDwtMSP" + version = get_package_version(package_name) + return cls(api_key=default_api_key, package_name=package_name, version=version) + def log_api(self, action, client_time=None, total_runtime=None, metadata=None): """ This function logs through an API call, assigns parameters @@ -443,35 +446,38 @@ def log_api(self, action, client_time=None, total_runtime=None, metadata=None): metadata["uid_issue"] = uid uid = None - py_version = python_version() - docker_container = is_docker() cloud = is_cloud_user() email = email_registered() - colab = is_colab() + colab = SYSTEM_INFO["is_colab"] + if colab: metadata["colab"] = colab - paperspace = is_paperspace() + paperspace = SYSTEM_INFO["is_paperspace"] + if paperspace: metadata["paperspace"] = paperspace - slurm = is_slurm() + slurm = SYSTEM_INFO["is_slurm"] + if slurm: metadata["slurm"] = slurm - airflow = is_airflow() + airflow = SYSTEM_INFO["is_airflow"] + if airflow: metadata["airflow"] = airflow - argo = is_argo() + argo = SYSTEM_INFO["is_argo"] + if argo: metadata["argo"] = argo if "dag" in metadata: metadata["dag"] = parse_dag(metadata["dag"]) - os = get_os() - environment = get_env() + os = SYSTEM_INFO["os"] + environment = SYSTEM_INFO["env"] if telemetry_enabled: (event_id, uid, action, client_time, elapsed_time) = validate_entries( @@ -483,10 +489,10 @@ def log_api(self, action, client_time=None, total_runtime=None, metadata=None): "action": action, "client_time": str(client_time), "total_runtime": total_runtime, - "python_version": py_version, + "python_version": SYSTEM_INFO["python_version"], "version": self.version, "package_name": self.package_name, - "docker_container": docker_container, + "docker_container": SYSTEM_INFO["is_docker"], "cloud": cloud, "email": email, "os": os, diff --git a/tests/telemetry/test_system_info.py b/tests/telemetry/test_system_info.py new file mode 100644 index 0000000..b5017a1 --- /dev/null +++ b/tests/telemetry/test_system_info.py @@ -0,0 +1,15 @@ +from ploomber_core.telemetry.system_info import get_system_info + + +def test_get_system_info(): + assert set(get_system_info()) == { + "env", + "is_airflow", + "is_argo", + "is_colab", + "is_docker", + "is_paperspace", + "is_slurm", + "os", + "python_version", + } diff --git a/tests/telemetry/test_telemetry.py b/tests/telemetry/test_telemetry.py index da2090d..4eaebcf 100644 --- a/tests/telemetry/test_telemetry.py +++ b/tests/telemetry/test_telemetry.py @@ -16,6 +16,7 @@ from ploomber_core.telemetry import system_info from ploomber_core.telemetry.validate_inputs import str_param, opt_str_param + from ploomber_core.exceptions import BaseException MOCK_API_KEY = "phc_P1dsjk20bijsabdaib2eu" @@ -1074,3 +1075,11 @@ def test_check_cloud(tmp_directory, monkeypatch, capsys, last_cloud_check): assert expected in captured.out assert config["last_cloud_check"] == now + + +def test_from_package(): + _telemetry = telemetry.Telemetry.from_package("ploomber-core") + + assert _telemetry.api_key + assert _telemetry.version + assert _telemetry.package_name == "ploomber-core"