import logging import os from pathlib import Path from typing import Any import requests from packaging.version import Version from mlflow.environment_variables import ( _MLFLOW_TELEMETRY_LOGGING, _MLFLOW_TESTING_TELEMETRY, MLFLOW_DISABLE_TELEMETRY, ) from mlflow.telemetry.constant import ( CONFIG_STAGING_URL, CONFIG_URL, FALLBACK_UI_CONFIG, UI_CONFIG_STAGING_URL, UI_CONFIG_URL, ) from mlflow.telemetry.schemas import ENV_VAR_TO_ENVIRONMENT_MAP, Environment from mlflow.version import VERSION _logger = logging.getLogger(__name__) def _is_ci_env_or_testing() -> bool: """ Check if the current environment is a CI environment. If so, we should not track telemetry. """ env_vars = { "PYTEST_CURRENT_TEST", # https://docs.pytest.org/en/stable/example/simple.html#pytest-current-test-environment-variable "GITHUB_ACTIONS", # https://docs.github.com/en/actions/reference/variables-reference?utm_source=chatgpt.com#default-environment-variables "CI", # set by many CI providers "CIRCLECI", # https://circleci.com/docs/variables/#built-in-environment-variables "GITLAB_CI", # https://docs.gitlab.com/ci/variables/predefined_variables/#predefined-variables "JENKINS_URL", # https://www.jenkins.io/doc/book/pipeline/jenkinsfile/#using-environment-variables "TRAVIS", # https://docs.travis-ci.com/user/environment-variables/#default-environment-variables "TF_BUILD", # https://learn.microsoft.com/en-us/azure/devops/pipelines/build/variables?view=azure-devops&tabs=yaml#system-variables "BITBUCKET_BUILD_NUMBER", # https://support.atlassian.com/bitbucket-cloud/docs/variables-and-secrets/ "CODEBUILD_BUILD_ARN", # https://docs.aws.amazon.com/codebuild/latest/userguide/build-env-ref-env-vars.html "BUILDKITE", # https://buildkite.com/docs/pipelines/configure/environment-variables "TEAMCITY_VERSION", # https://www.jetbrains.com/help/teamcity/predefined-build-parameters.html#Predefined+Server+Build+Parameters "CLOUD_RUN_EXECUTION", # https://cloud.google.com/run/docs/reference/container-contract#env-vars # runbots "RUNBOT_HOST_URL", "RUNBOT_BUILD_NAME", "RUNBOT_WORKER_ID", } # For most of the cases, the env var existing means we are in CI for var in env_vars: if var in os.environ: return True return False # NB: implement the function here to avoid unnecessary imports inside databricks_utils def _is_in_databricks() -> bool: # check if in databricks runtime if "DATABRICKS_RUNTIME_VERSION" in os.environ: return True if os.path.exists("/databricks/DBR_VERSION"): return True # check if in databricks model serving environment if os.environ.get("IS_IN_DB_MODEL_SERVING_ENV", "false").lower() in ("true", "1"): return True return False def _detect_environment() -> str | None: # Check for MLflow demo deployment (e.g. demo.mlflow.org) before generic docker detection if os.environ.get("MLFLOW_DEPLOYMENT_ENV") == "demo": return Environment.DEMO.value for env_var, environment in ENV_VAR_TO_ENVIRONMENT_MAP.items(): if env_var in os.environ: return environment.value # https://docs.aws.amazon.com/sagemaker/latest/dg/nbi-metadata.html if Path("/opt/ml/metadata/resource-metadata.json").exists(): return Environment.SAGEMAKER_NOTEBOOK.value # unofficial heuristic to detect docker environment if Path("/.dockerenv").exists(): return Environment.DOCKER.value return None _IS_MLFLOW_DEV_VERSION = Version(VERSION).is_devrelease _IS_IN_CI_ENV_OR_TESTING = _is_ci_env_or_testing() _IS_IN_DATABRICKS = _is_in_databricks() _IS_MLFLOW_TESTING_TELEMETRY = _MLFLOW_TESTING_TELEMETRY.get() def is_telemetry_disabled() -> bool: try: if _IS_MLFLOW_TESTING_TELEMETRY: return False # NB: _IS_IN_DATABRICKS is intentionally NOT a disable signal here. When the # tracking URI is databricks:// or databricks-uc://, telemetry is forwarded # to the workspace's own ingestion endpoint (see TelemetryClient._forward_to_databricks). # The non-Databricks (OSS) ingestion path is separately guarded in # TelemetryClient._process_records so it is never hit from inside DBR. return ( MLFLOW_DISABLE_TELEMETRY.get() or os.environ.get("DO_NOT_TRACK", "false").lower() == "true" or _IS_IN_CI_ENV_OR_TESTING or _IS_MLFLOW_DEV_VERSION ) except Exception as e: _log_error(f"Failed to check telemetry disabled status: {e}") return True def _get_config_url(version: str, is_ui: bool = False) -> str | None: """ Get the config URL for the given MLflow version. """ version_obj = Version(version) if version_obj.is_devrelease or _IS_MLFLOW_TESTING_TELEMETRY: base_url = UI_CONFIG_STAGING_URL if is_ui else CONFIG_STAGING_URL return f"{base_url}/{version}.json" if version_obj.base_version == version or ( version_obj.is_prerelease and version_obj.pre[0] == "rc" ): base_url = UI_CONFIG_URL if is_ui else CONFIG_URL return f"{base_url}/{version}.json" return None def _log_error(message: str) -> None: if _MLFLOW_TELEMETRY_LOGGING.get(): _logger.error(message, exc_info=True) def fetch_ui_telemetry_config() -> dict[str, Any]: # Check if telemetry is disabled if is_telemetry_disabled(): return FALLBACK_UI_CONFIG # Get config URL config_url = _get_config_url(VERSION, is_ui=True) if not config_url: return FALLBACK_UI_CONFIG # Fetch config from remote URL try: response = requests.get(config_url, timeout=1) if response.status_code != 200: return FALLBACK_UI_CONFIG return response.json() except Exception as e: _log_error(f"Failed to fetch UI telemetry config: {e}") return FALLBACK_UI_CONFIG