Merge pull request #56 from PriorLabs/ENG-377

Keep track of execution environment

Author: safaricd

src/tabpfn_common_utils/telemetry/core/events.py

@@ -6,7 +6,7 @@
 from datetime import datetime, timezone
 from functools import lru_cache
 from typing import Any, Literal, Optional
-from .runtime import get_runtime
+from .runtime import get_execution_context
 from .state import get_property, set_property
 
 
@@ -179,8 +179,8 @@ def _get_runtime_kernel() -> Optional[str]:
     Returns:
         str: Runtime environment of the platform.
     """
-    runtime = get_runtime()
-    return runtime.kernel
+    exec_context = get_execution_context()
+    return exec_context.kernel
 
 
 @lru_cache(maxsize=1)

src/tabpfn_common_utils/telemetry/core/runtime.py

@@ -5,65 +5,159 @@
 import os
 import sys
 from dataclasses import dataclass
-from typing import Literal
+from typing import Callable, Dict, Literal, Mapping, Optional, Sequence, Tuple
+
+
+# The type of kernel the code is running in
+KernelType = Literal["ipython", "jupyter", "tty"]
+
+# The type of environment the code is running in
+EnvironmentType = Literal[
+    "kaggle",
+    "colab",
+    "gcp",
+    "aws",
+    "azure",
+    "databricks",
+]
+
+# Static list of environment hints, purely heuristic based on env variables.
+# This information is used purely for detecting purposes, and the values are
+# not propagated or stored in the telemetry state.
+ENV_TYPE_HINTS: Mapping[EnvironmentType, Sequence[str]] = {
+    # Notebook providers
+    "kaggle": [
+        # Kaggle kernels
+        "KAGGLE_KERNEL_RUN_TYPE",
+        "KAGGLE_URL_BASE",
+        "KAGGLE_KERNEL_INTEGRATIONS",
+        "KAGGLE_USER_SECRETS_TOKEN",
+        "KAGGLE_GCP_PROJECT",
+        "KAGGLE_GCP_ZONE",
+    ],
+    "colab": [
+        # Google Colab
+        "COLAB_GPU",
+        "COLAB_TPU_ADDR",
+        "COLAB_JUPYTER_TRANSPORT",
+        "COLAB_BACKEND_VERSION",
+    ],
+    "databricks": [
+        # Databricks clusters and runtime
+        "DATABRICKS_RUNTIME_VERSION",
+        "DATABRICKS_CLUSTER_ID",
+        "DATABRICKS_HOST",
+        "DATABRICKS_WORKSPACE_URL",
+        "DB_IS_DRIVER",
+    ],
+    # Cloud providers
+    "aws": [
+        # Generic AWS environment hints
+        "AWS_EXECUTION_ENV",
+        "AWS_REGION",
+        "AWS_DEFAULT_REGION",
+        # SageMaker
+        "SM_MODEL_DIR",
+        "SM_NUM_GPUS",
+        "SM_HOSTS",
+        "SM_CURRENT_HOST",
+        "TRAINING_JOB_NAME",
+    ],
+    "gcp": [
+        # Project hints
+        "GOOGLE_CLOUD_PROJECT",
+        "GCP_PROJECT",
+        "GCLOUD_PROJECT",
+        "CLOUDSDK_CORE_PROJECT",
+        # Cloud Run and Cloud Functions
+        "K_SERVICE",
+        "K_REVISION",
+        "K_CONFIGURATION",
+        "CLOUD_RUN_JOB",
+        # Vertex AI
+        "AIP_MODEL_DIR",
+        "AIP_DATA_FORMAT",
+        "AIP_TRAINING_DATA_URI",
+        "CLOUD_ML_JOB_ID",
+        # Cloud Shell
+        "CLOUD_SHELL",
+    ],
+    "azure": [
+        # Azure ML
+        "AZUREML_RUN_ID",
+        "AZUREML_ARM_SUBSCRIPTION",
+        "AZUREML_ARM_RESOURCEGROUP",
+        "AZUREML_ARM_WORKSPACE_NAME",
+    ],
+}
 
 
 @dataclass
-class Runtime:
-    """Runtime environment."""
+class ExecutionContext:
+    """The execution context of the current environment."""
 
     interactive: bool
-    kernel: Literal["ipython", "jupyter", "tty", "kaggle"] | None = None
+    """Whether the code is running in an interactive environment."""
+
+    kernel: Optional[KernelType] = None
+    """Low-level Python frontend or shell (e.g. IPython, Jupyter, TTY)"""
+
+    environment: Optional[EnvironmentType] = None
+    """Higher-level hosted environment or notebook platform."""
+
     ci: bool = False
+    """Whether the code is running in a CI environment."""
 
 
-def get_runtime() -> Runtime:
-    """Get the runtime environment.
+def get_execution_context() -> ExecutionContext:
+    """Get the execution context of the current environment.
 
     Returns:
-        The runtime environment.
+        The execution context of the current environment.
     """
-    # First check for Kaggle
-    if _is_kaggle():
-        return Runtime(interactive=True, kernel="kaggle")
+    # First check for environment
+    environment = _get_environment()
+
+    # Next, get kernel information
+    interactive, kernel = _get_kernel()
 
-    # Next check for CI
-    if _is_ci():
-        return Runtime(interactive=False, kernel=None, ci=True)
+    context = ExecutionContext(
+        interactive=interactive, kernel=kernel, environment=environment, ci=_is_ci()
+    )
+    return context
 
-    # Check for IPython
-    if _is_ipy():
-        return Runtime(interactive=True, kernel="ipython")
 
-    # Jupyter kernel
-    if _is_jupyter_kernel():
-        return Runtime(interactive=True, kernel="jupyter")
+def _get_kernel() -> Tuple[bool, Optional[KernelType]]:
+    """Get the kernel the code is running in.
 
-    # TTY
-    if _is_tty():
-        return Runtime(interactive=True, kernel="tty")
+    Returns:
+        A tuple of (whether the kernel is interactive, the kernel type).
+    """
+    mapping: Dict[KernelType, Callable[[], bool]] = {
+        "ipython": _is_ipy,
+        "jupyter": _is_jupyter_kernel,
+        "tty": _is_tty,
+    }
+    for kernel, func in mapping.items():
+        if func():
+            return True, kernel
+    return False, None
 
-    # Default to non-interactive
-    return Runtime(interactive=False, kernel=None)
 
+def _get_environment() -> Optional[EnvironmentType]:
+    """Get the environment the code is running in.
 
-def _is_kaggle() -> bool:
-    """Check if the current environment is running in a Kaggle kernel.
+    An environment is a higher-level hosted environment or notebook platform.
+    This is about where the code is running (Kaggle, Colab, AWS, GCP, ...).
 
     Returns:
-        bool: True if the current environment is running in a Kaggle kernel.
+        The environment the code is running in.
     """
-    # Kaggle-specific and preset env vars
-    kaggle_env_vars = [
-        "KAGGLE_KERNEL_RUN_TYPE",
-        "KAGGLE_URL_BASE",
-        "KAGGLE_KERNEL_INTEGRATIONS",
-        "KAGGLE_USER_SECRETS_TOKEN",
-    ]
-    if any(v in os.environ for v in kaggle_env_vars):
-        return True
+    for env_type, hints in ENV_TYPE_HINTS.items():
+        if any(k in os.environ for k in hints):
+            return env_type
 
-    return False
+    return None
 
 
 def _is_ipy() -> bool:

src/tabpfn_common_utils/telemetry/core/service.py

@@ -5,7 +5,7 @@
 from posthog import Posthog
 from .config import download_config
 from .events import BaseTelemetryEvent
-from .runtime import get_runtime
+from .runtime import get_execution_context
 from ...utils import singleton
 from typing import Any, Dict, Optional
 
@@ -71,8 +71,8 @@ def telemetry_enabled(cls) -> bool:
             bool: True if telemetry is enabled, False otherwise.
         """
         # Disable telemetry by default in CI environments, but allow override
-        runtime = get_runtime()
-        default_disable = "1" if runtime.ci else "0"
+        exec_context = get_execution_context()
+        default_disable = "1" if exec_context.ci else "0"
 
         disable_telemetry = os.getenv(
             "TABPFN_DISABLE_TELEMETRY", default_disable

src/tabpfn_common_utils/telemetry/interactive/flows.py

@@ -11,7 +11,7 @@
 from .prompts.base import PromptSpec
 from .prompts.newsletter import NewsletterPrompt
 from .prompts.identity import IdentityPrompt
-from ..core.runtime import get_runtime
+from ..core.runtime import get_execution_context
 
 
 def capture_session(enabled: bool = True) -> None:
@@ -142,8 +142,8 @@ def opt_in(enabled: bool = True, delta_days: int = 30, max_prompts: int = 2) ->
         return
 
     # Only show prompts in Jupyter/IPython
-    runtime = get_runtime()
-    if runtime.kernel not in {"jupyter", "ipython"}:
+    exec_context = get_execution_context()
+    if exec_context.kernel not in {"jupyter", "ipython"}:
         return
 
     # Check if prompts should be shown

tests/telemetry/core/test_runtime.py

@@ -10,7 +10,7 @@
     _is_ipy,
     _is_jupyter_kernel,
     _is_tty,
-    get_runtime,
+    get_execution_context,
 )
 
 
@@ -24,7 +24,7 @@ def setup(self) -> None:
     def test_get_runtime_ci_environment(self) -> None:
         """Test that CI environments are detected correctly."""
         with patch(f"{self.module}._is_ci", return_value=True):
-            runtime = get_runtime()
+            runtime = get_execution_context()
             assert runtime.ci is True
             assert runtime.interactive is False
             assert runtime.kernel is None
@@ -35,7 +35,7 @@ def test_get_runtime_interactive_ipython(self) -> None:
             patch(f"{self.module}._is_ci", return_value=False),
             patch(f"{self.module}._is_ipy", return_value=True),
         ):
-            runtime = get_runtime()
+            runtime = get_execution_context()
             assert runtime.interactive is True
             assert runtime.ci is False
 
@@ -46,7 +46,7 @@ def test_get_runtime_interactive_jupyter(self) -> None:
             patch(f"{self.module}._is_ipy", return_value=False),
             patch(f"{self.module}._is_jupyter_kernel", return_value=True),
         ):
-            runtime = get_runtime()
+            runtime = get_execution_context()
             assert runtime.interactive is True
             assert runtime.ci is False
 
@@ -57,7 +57,7 @@ def test_get_runtime_default_noninteractive(self) -> None:
             patch(f"{self.module}._is_ipy", return_value=False),
             patch(f"{self.module}._is_jupyter_kernel", return_value=False),
         ):
-            runtime = get_runtime()
+            runtime = get_execution_context()
             assert runtime.interactive is False
             assert runtime.ci is False