wehub-resource-sync
85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
"""
|
|
Helpers for making sure a Databricks SQL warehouse is running before MLflow tracing API calls
|
|
that require it.
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
from datetime import timedelta
|
|
|
|
from mlflow.environment_variables import (
|
|
MLFLOW_SQL_WAREHOUSE_AUTO_START,
|
|
MLFLOW_SQL_WAREHOUSE_AUTO_START_TIMEOUT_SECONDS,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
_CACHE_TTL_SECONDS = 60.0
|
|
|
|
# warehouse_id -> monotonic deadline by which the "RUNNING" verification expires.
|
|
#
|
|
# Concurrent callers may race and hit the SDK more than once on a cold cache; that's fine.
|
|
# `warehouses.get` is cheap and `start_and_wait` is idempotent on the server. The cache's
|
|
# purpose is to eliminate SDK hops in the steady state, not to single-flight the cold path.
|
|
_verified_running: dict[str, float] = {}
|
|
|
|
|
|
def _get_workspace_client():
|
|
from databricks.sdk import WorkspaceClient
|
|
|
|
return WorkspaceClient()
|
|
|
|
|
|
def ensure_sql_warehouse_running(warehouse_id: str) -> None:
|
|
"""
|
|
Verify the SQL warehouse is in ``RUNNING`` state, starting it and waiting if necessary.
|
|
|
|
No-op when ``MLFLOW_SQL_WAREHOUSE_AUTO_START`` is false. Results are cached per-process
|
|
for ``_CACHE_TTL_SECONDS`` to avoid hammering the SDK across closely-spaced calls.
|
|
The ``start_and_wait`` timeout is taken from
|
|
``MLFLOW_SQL_WAREHOUSE_AUTO_START_TIMEOUT_SECONDS``.
|
|
|
|
Args:
|
|
warehouse_id: The Databricks SQL warehouse ID to check.
|
|
|
|
Raises:
|
|
MlflowException: When the warehouse fails to reach ``RUNNING`` (timeout or other
|
|
SDK error).
|
|
"""
|
|
if not MLFLOW_SQL_WAREHOUSE_AUTO_START.get():
|
|
return
|
|
|
|
deadline = _verified_running.get(warehouse_id)
|
|
if deadline is not None and time.monotonic() < deadline:
|
|
return
|
|
|
|
from databricks.sdk.service.sql import State
|
|
|
|
client = _get_workspace_client()
|
|
info = client.warehouses.get(warehouse_id)
|
|
|
|
if info.state != State.RUNNING:
|
|
timeout = MLFLOW_SQL_WAREHOUSE_AUTO_START_TIMEOUT_SECONDS.get()
|
|
_logger.info(
|
|
f"SQL warehouse '{warehouse_id}' is {info.state.value}; starting it and "
|
|
f"waiting up to {timeout}s for RUNNING."
|
|
)
|
|
try:
|
|
client.warehouses.start_and_wait(warehouse_id, timeout=timedelta(seconds=timeout))
|
|
except TimeoutError as e:
|
|
raise MlflowException(
|
|
f"Timed out after {timeout}s waiting for SQL warehouse '{warehouse_id}' to "
|
|
f"reach RUNNING state. Increase the timeout via the "
|
|
f"`{MLFLOW_SQL_WAREHOUSE_AUTO_START_TIMEOUT_SECONDS.name}` environment "
|
|
f"variable, or start the warehouse explicitly and retry."
|
|
) from e
|
|
except Exception as e:
|
|
raise MlflowException(
|
|
f"Failed to start SQL warehouse '{warehouse_id}': {e}. Start the warehouse "
|
|
f"explicitly and retry, or set `MLFLOW_SQL_WAREHOUSE_AUTO_START=false` to "
|
|
f"disable this preflight."
|
|
) from e
|
|
|
|
_verified_running[warehouse_id] = time.monotonic() + _CACHE_TTL_SECONDS
|