import atexit import random import sys import threading import time import urllib.parse import uuid import warnings from dataclasses import asdict from functools import lru_cache from queue import Empty, Full, Queue from typing import Any, Callable, Literal import requests from mlflow.environment_variables import _MLFLOW_TELEMETRY_SESSION_ID, MLFLOW_WORKSPACE from mlflow.telemetry.constant import ( BATCH_SIZE, BATCH_TIME_INTERVAL_SECONDS, MAX_QUEUE_SIZE, MAX_WORKERS, RETRYABLE_ERRORS, UNRECOVERABLE_ERRORS, ) from mlflow.telemetry.installation_id import get_or_create_installation_id from mlflow.telemetry.schemas import Record, TelemetryConfig, TelemetryInfo, get_source_sdk from mlflow.telemetry.utils import ( _IS_IN_DATABRICKS, _IS_MLFLOW_DEV_VERSION, _detect_environment, _get_config_url, _log_error, is_telemetry_disabled, ) from mlflow.utils.credentials import get_default_host_creds from mlflow.utils.logging_utils import should_suppress_logs_in_thread, suppress_logs_in_thread from mlflow.utils.rest_utils import http_request _DATABRICKS_SCHEMES = ("databricks", "databricks-uc", "uc") # Cache per tracking URI; 16 is more than enough for any realistic number of # distinct tracking URIs within a single process. @lru_cache(maxsize=16) def _fetch_server_info(tracking_uri: str) -> dict[str, Any] | None: try: response = http_request( host_creds=get_default_host_creds(tracking_uri), endpoint="/api/3.0/mlflow/server-info", method="GET", timeout=3, max_retries=0, raise_on_status=False, ) if response.status_code == 200: return response.json() except Exception: pass return None def _enrich_http_scheme(scheme: Literal["http", "https"], store_type: str | None) -> str: store_type_to_suffix = {"FileStore": "file", "SqlStore": "sql"} if suffix := store_type_to_suffix.get(store_type): return f"{scheme}-{suffix}" return scheme def _is_localhost_uri(uri: str) -> bool | None: """ Check if the given URI points to localhost. Returns: True if the URI points to localhost, False if it points to a remote host, or None if the URI cannot be parsed or has no hostname. """ try: parsed = urllib.parse.urlparse(uri) hostname = parsed.hostname if not hostname: return None return ( hostname in (".", "::1") or hostname.startswith("localhost") or hostname.startswith("127.0.0.1") ) except Exception: return None def _get_tracking_uri_info() -> tuple[str | None, bool | None]: """ Get tracking URI information including scheme and localhost status. Returns: A tuple of (scheme, is_localhost). is_localhost is only set for http/https schemes. """ # import here to avoid circular import from mlflow.tracking._tracking_service.utils import ( _get_tracking_scheme_with_resolved_uri, get_tracking_uri, ) try: tracking_uri = get_tracking_uri() scheme = _get_tracking_scheme_with_resolved_uri(tracking_uri) # Check if http/https points to localhost is_localhost = _is_localhost_uri(tracking_uri) if scheme in ("http", "https") else None return scheme, is_localhost except Exception: return None, None class TelemetryClient: def __init__(self): self.info = asdict( TelemetryInfo( session_id=_MLFLOW_TELEMETRY_SESSION_ID.get() or uuid.uuid4().hex, environment=_detect_environment(), installation_id=get_or_create_installation_id(), ) ) self._queue: Queue[list[Record]] = Queue(maxsize=MAX_QUEUE_SIZE) self._lock = threading.RLock() self._max_workers = MAX_WORKERS self._is_stopped = False self._is_active = False self._atexit_callback_registered = False self._batch_size = BATCH_SIZE self._batch_time_interval = BATCH_TIME_INTERVAL_SECONDS self._pending_records: list[Record] = [] self._last_batch_time = time.time() self._batch_lock = threading.Lock() # consumer threads for sending records self._consumer_threads = [] self._is_config_fetched = False self.config = None def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): self._clean_up() def _fetch_config(self): def _fetch(): try: self._get_config() if self.config is None: self._is_stopped = True _set_telemetry_client(None) self._is_config_fetched = True except Exception: self._is_stopped = True self._is_config_fetched = True _set_telemetry_client(None) self._config_thread = threading.Thread( target=_fetch, name="GetTelemetryConfig", daemon=True, ) self._config_thread.start() def _get_config(self): """ Get the config for the given MLflow version. """ mlflow_version = self.info["mlflow_version"] if config_url := _get_config_url(mlflow_version): try: response = requests.get(config_url, timeout=1) if response.status_code != 200: return config = response.json() if ( config.get("mlflow_version") != mlflow_version or config.get("disable_telemetry") is True or config.get("ingestion_url") is None ): return if get_source_sdk().value in config.get("disable_sdks", []): return if sys.platform in config.get("disable_os", []): return rollout_percentage = config.get("rollout_percentage", 100) if random.randint(0, 100) > rollout_percentage: return self.config = TelemetryConfig( ingestion_url=config["ingestion_url"], disable_events=set(config.get("disable_events", [])), ) except Exception as e: _log_error(f"Failed to get telemetry config: {e}") return def add_record(self, record: Record): """ Add a record to be batched and sent to the telemetry server. """ if not self.is_active: self.activate() if self._is_stopped: return with self._batch_lock: self._pending_records.append(record) # Only send if we've reached the batch size; # time-based sending is handled by the consumer thread. if len(self._pending_records) >= self._batch_size: self._send_batch() def add_records(self, records: list[Record]): if not self.is_active: self.activate() if self._is_stopped: return with self._batch_lock: # Add records in chunks to ensure we never exceed batch_size offset = 0 while offset < len(records): # Calculate how many records we can add to reach batch_size space_left = self._batch_size - len(self._pending_records) chunk_size = min(space_left, len(records) - offset) # Add only enough records to reach batch_size self._pending_records.extend(records[offset : offset + chunk_size]) offset += chunk_size # Send batch if we've reached the limit if len(self._pending_records) >= self._batch_size: self._send_batch() def _send_batch(self): """Send the current batch of records.""" if not self._pending_records: return self._last_batch_time = time.time() try: self._queue.put(self._pending_records, block=False) self._pending_records = [] except Full: _log_error("Failed to add record to the queue, queue is full") def _process_records(self, records: list[Record], request_timeout: float = 1): """Process a batch of telemetry records.""" try: self._update_backend_store() if self.info.get("tracking_uri_scheme") in _DATABRICKS_SCHEMES: if not _IS_MLFLOW_DEV_VERSION: self._forward_to_databricks(records, request_timeout) return # Never POST to the OSS ingestion endpoint from inside a Databricks # workload. The Databricks-forwarding branch above is the only path # allowed in DBR, model serving, etc. if _IS_IN_DATABRICKS: return records = [ { "data": self.info | record.to_dict(), # use random uuid as partition key to make sure records are # distributed evenly across shards "partition-key": uuid.uuid4().hex, } for record in records ] self._send_with_retries( lambda timeout: requests.post( self.config.ingestion_url, json={"records": records}, headers={"Content-Type": "application/json"}, timeout=timeout, ), request_timeout, ) except Exception as e: _log_error(f"Failed to send telemetry records: {e}") def _forward_to_databricks(self, records: list[Record], request_timeout: float = 1): from mlflow.tracking._tracking_service.utils import get_tracking_uri from mlflow.utils.databricks_utils import get_databricks_host_creds try: creds = get_databricks_host_creds(get_tracking_uri()) except Exception as e: _log_error(f"Failed to get Databricks credentials for telemetry: {e}") return events = [] for record in records: event = dict(self.info) event.update(record.to_dict()) params_value = event.pop("params", None) if params_value is not None: event["params_json"] = params_value events.append(event) self._send_with_retries( lambda timeout: http_request( host_creds=creds, endpoint="/api/2.0/mlflow/client-telemetry/ingest", method="POST", timeout=timeout, max_retries=0, raise_on_status=False, json={"events": events}, ), request_timeout, ) def _send_with_retries( self, send_fn: Callable[[float], requests.Response], request_timeout: float = 1 ) -> None: max_attempts = 3 sleep_time = 1 for i in range(max_attempts): response = None should_retry = False try: response = send_fn(request_timeout) should_retry = response.status_code in RETRYABLE_ERRORS except (ConnectionError, TimeoutError, requests.ConnectionError, requests.Timeout): should_retry = True # NB: DO NOT retry when terminating # otherwise this increases shutdown overhead significantly if self._is_stopped: return if i < max_attempts - 1 and should_retry: time.sleep(sleep_time) elif response and response.status_code in UNRECOVERABLE_ERRORS: self._is_stopped = True self.is_active = False return else: return def _consumer(self) -> None: """Individual consumer that processes records from the queue.""" # suppress logs in the consumer thread to avoid emitting any irrelevant # logs during telemetry collection. should_suppress_logs_in_thread.set(True) while not self._is_config_fetched: time.sleep(0.1) while self.config and not self._is_stopped: try: records = self._queue.get(timeout=1) except Empty: # check if batch time interval has passed and send data if needed if time.time() - self._last_batch_time >= self._batch_time_interval: self._last_batch_time = time.time() with self._batch_lock: if self._pending_records: self._send_batch() continue self._process_records(records) self._queue.task_done() # clear the queue if config is None while self.config is None and not self._queue.empty(): try: self._queue.get_nowait() self._queue.task_done() except Empty: break # drop remaining records when terminating to avoid # causing any overhead def activate(self) -> None: """Activate the async queue to accept and handle incoming tasks.""" with self._lock: if self.is_active: return self._set_up_threads() # only fetch config when activating to avoid fetching when # no records are added self._fetch_config() # Callback to ensure remaining tasks are processed before program exit if not self._atexit_callback_registered: # This works in jupyter notebook atexit.register(self._at_exit_callback) self._atexit_callback_registered = True self.is_active = True @property def is_active(self) -> bool: return self._is_active @is_active.setter def is_active(self, value: bool) -> None: self._is_active = value def _set_up_threads(self) -> None: """Set up multiple consumer threads.""" with self._lock: # Start multiple consumer threads for i in range(self._max_workers): consumer_thread = threading.Thread( target=self._consumer, name=f"MLflowTelemetryConsumer-{i}", daemon=True, ) consumer_thread.start() self._consumer_threads.append(consumer_thread) def _at_exit_callback(self) -> None: """Callback function executed when the program is exiting.""" try: # Suppress logs/warnings during shutdown # NB: this doesn't suppress log not emitted by mlflow with suppress_logs_in_thread(), warnings.catch_warnings(): warnings.simplefilter("ignore") self.flush(terminate=True) except Exception as e: _log_error(f"Failed to flush telemetry during termination: {e}") def flush(self, terminate=False) -> None: """ Flush the async telemetry queue. Args: terminate: If True, shut down the telemetry threads after flushing. """ if not self.is_active: return if terminate: # Full shutdown for termination - signal stop and exit immediately self._is_stopped = True self.is_active = False # non-terminating flush is only used in tests else: self._config_thread.join(timeout=1) # Send any pending records before flushing with self._batch_lock: if self._pending_records and self.config and not self._is_stopped: self._send_batch() # For non-terminating flush, just wait for queue to empty try: self._queue.join() except Exception as e: _log_error(f"Failed to flush telemetry: {e}") def _resolve_tracking_scheme(self, scheme: str) -> str: if scheme not in ("http", "https"): return scheme # import here to avoid circular import from mlflow.tracking._tracking_service.utils import get_tracking_uri server_info = _fetch_server_info(get_tracking_uri()) store_type = server_info.get("store_type") if server_info else None return _enrich_http_scheme(scheme, store_type) def _update_backend_store(self): """ Backend store might be changed after mlflow is imported, we should use this method to update the backend store info at sending telemetry step. """ try: scheme, is_localhost = _get_tracking_uri_info() if scheme is not None: self.info["tracking_uri_scheme"] = self._resolve_tracking_scheme(scheme) if is_localhost is not None: self.info["is_localhost"] = is_localhost self.info["ws_enabled"] = bool(MLFLOW_WORKSPACE.get()) except Exception as e: _log_error(f"Failed to update backend store: {e}") def _clean_up(self): """Join all threads""" self.flush(terminate=True) for thread in self._consumer_threads: if thread.is_alive(): thread.join(timeout=1) _MLFLOW_TELEMETRY_CLIENT = None _client_lock = threading.Lock() def set_telemetry_client(): if is_telemetry_disabled(): # set to None again so this function can be used to # re-initialize the telemetry client _set_telemetry_client(None) else: try: _set_telemetry_client(TelemetryClient()) except Exception as e: _log_error(f"Failed to set telemetry client: {e}") _set_telemetry_client(None) def _set_telemetry_client(value: TelemetryClient | None): global _MLFLOW_TELEMETRY_CLIENT with _client_lock: _MLFLOW_TELEMETRY_CLIENT = value if value: _MLFLOW_TELEMETRY_SESSION_ID.set(value.info["session_id"]) else: _MLFLOW_TELEMETRY_SESSION_ID.unset() def get_telemetry_client() -> TelemetryClient | None: with _client_lock: return _MLFLOW_TELEMETRY_CLIENT