Skip to content

Commit

Permalink
faster telemetry + simpler initialization (#82)
Browse files Browse the repository at this point in the history
* more efficient system info

* adds classmethod to initialize telemetry with default settings

* lint

* changelog and docs

* update doc

* ci
  • Loading branch information
edublancas authored Jan 12, 2024
1 parent 8a8cfe9 commit bf83ec4
Show file tree
Hide file tree
Showing 7 changed files with 101 additions and 31 deletions.
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
19 changes: 10 additions & 9 deletions doc/telemetry.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
"click",
"pyyaml",
"posthog",
'importlib-metadata;python_version<"3.8"',
]

DEV = [
Expand Down
34 changes: 34 additions & 0 deletions src/ploomber_core/telemetry/system_info.py
Original file line number Diff line number Diff line change
Expand Up @@ -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}"
Expand Down Expand Up @@ -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),
}
50 changes: 28 additions & 22 deletions src/ploomber_core/telemetry/telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -69,6 +59,9 @@
logging.getLogger("posthog").disabled = True


SYSTEM_INFO = get_system_info()


class UserSettings(Config):
"""User-customizable settings"""

Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand All @@ -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,
Expand Down
15 changes: 15 additions & 0 deletions tests/telemetry/test_system_info.py
Original file line number Diff line number Diff line change
@@ -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",
}
9 changes: 9 additions & 0 deletions tests/telemetry/test_telemetry.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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"

0 comments on commit bf83ec4

Please sign in to comment.