diff --git a/src/ploomber_core/telemetry/system_info.py b/src/ploomber_core/telemetry/system_info.py index bcca963..5b76115 100644 --- a/src/ploomber_core/telemetry/system_info.py +++ b/src/ploomber_core/telemetry/system_info.py @@ -106,3 +106,24 @@ 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_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..60ecd94 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 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""" @@ -443,35 +436,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 +479,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", + }