wehub-resource-sync
534 lines
18 KiB
Python
534 lines
18 KiB
Python
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
|