wehub-resource-sync
146 lines
4.6 KiB
Python
146 lines
4.6 KiB
Python
"""Session state for the ``@mlflow.test`` pytest plugin.
|
|
|
|
Tracks the single test run per pytest session and which test is currently
|
|
executing (for trace tagging).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import datetime
|
|
import logging
|
|
import threading
|
|
import uuid
|
|
|
|
import mlflow
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
TAG_TEST_NAME = "mlflow.test.name"
|
|
TAG_SESSION_ID = "mlflow.test.session_id"
|
|
TAG_CASE_ID = "mlflow.test.case_id"
|
|
|
|
_lock = threading.Lock()
|
|
_session_id: str | None = None
|
|
_run_id: str | None = None
|
|
_run_owned: bool = False
|
|
_any_test_failed: bool = False
|
|
_num_tests: int = 0
|
|
_total_duration_ms: int = 0
|
|
# True once the (opt-in) pytest plugin registered in this process; @mlflow.test
|
|
# checks it to fail loudly instead of silently doing nothing.
|
|
_plugin_active: bool = False
|
|
|
|
_current = threading.local()
|
|
|
|
|
|
def set_plugin_active() -> None:
|
|
global _plugin_active
|
|
_plugin_active = True
|
|
|
|
|
|
def is_plugin_active() -> bool:
|
|
return _plugin_active
|
|
|
|
|
|
def set_current_test(test_name: str | None, case_id: str | None = None) -> None:
|
|
_current.value = (test_name, case_id)
|
|
|
|
|
|
def current_test() -> tuple[str | None, str | None]:
|
|
return getattr(_current, "value", (None, None))
|
|
|
|
|
|
def reset(session_id: str | None = None) -> None:
|
|
global _session_id, _run_id, _run_owned, _any_test_failed, _num_tests, _total_duration_ms
|
|
if not session_id:
|
|
stamp = datetime.datetime.now().strftime("%Y%m%dT%H%M%S")
|
|
session_id = f"{stamp}-{uuid.uuid4().hex[:6]}"
|
|
with _lock:
|
|
_session_id = session_id
|
|
_run_id = None
|
|
_run_owned = False
|
|
_any_test_failed = False
|
|
_num_tests = 0
|
|
_total_duration_ms = 0
|
|
|
|
|
|
def session_id() -> str:
|
|
if _session_id is None:
|
|
reset()
|
|
return _session_id
|
|
|
|
|
|
def record_test(*, failed: bool, duration_ms: int) -> None:
|
|
"""Record the outcome and duration of one ``@mlflow.test``-marked test."""
|
|
global _any_test_failed, _num_tests, _total_duration_ms
|
|
with _lock:
|
|
_num_tests += 1
|
|
_total_duration_ms += duration_ms
|
|
if failed:
|
|
_any_test_failed = True
|
|
|
|
|
|
def ensure_run() -> str | None:
|
|
"""Open the test run, once per session. Thread-safe.
|
|
|
|
If a run is already active (e.g. opened by a user fixture), start a nested
|
|
child run rather than reusing/retagging the user's run -- we don't know what
|
|
that run is for, so we never mutate it.
|
|
"""
|
|
global _run_id, _run_owned
|
|
with _lock:
|
|
if _run_id is not None:
|
|
return _run_id
|
|
|
|
# Imported lazily: the plugin loads in every pytest run, so a
|
|
# module-level import of a PR-new symbol breaks collection against a
|
|
# version-skewed mlflow.
|
|
from mlflow.utils.mlflow_tags import MLFLOW_RUN_TYPE, MLFLOW_RUN_TYPE_TEST
|
|
|
|
tags = {MLFLOW_RUN_TYPE: MLFLOW_RUN_TYPE_TEST, TAG_SESSION_ID: session_id()}
|
|
try:
|
|
# Nest under a user-opened run rather than reusing/retagging it.
|
|
nested = mlflow.active_run() is not None
|
|
run = mlflow.start_run(nested=nested, tags=tags)
|
|
except Exception as e:
|
|
_logger.warning("mlflow.test: could not start test run: %s", e)
|
|
return None
|
|
_run_id = run.info.run_id
|
|
_run_owned = True
|
|
return _run_id
|
|
|
|
|
|
def finalize() -> None:
|
|
"""Emit the session telemetry event and end the test run.
|
|
|
|
Telemetry fires whenever a ``@mlflow.test`` ran (even if run creation
|
|
failed). The run status reflects only ``@mlflow.test`` outcomes, not other
|
|
tests that happen to run in the same pytest session. Runs once on the main
|
|
thread at ``pytest_sessionfinish``.
|
|
"""
|
|
global _run_id, _run_owned
|
|
with _lock:
|
|
if _num_tests > 0:
|
|
# Imported lazily and defensively: this plugin loads in every pytest
|
|
# run, so a module-level telemetry import would break collection if
|
|
# mlflow.telemetry is unavailable or version-skewed.
|
|
try:
|
|
from mlflow.telemetry.events import MlflowTestEvent
|
|
from mlflow.telemetry.track import _record_event
|
|
|
|
_record_event(
|
|
MlflowTestEvent,
|
|
{"num_tests": _num_tests},
|
|
duration_ms=_total_duration_ms,
|
|
)
|
|
except Exception as e:
|
|
_logger.debug("mlflow.test: skipped telemetry: %s", e)
|
|
if _run_owned and _run_id is not None:
|
|
status = "FAILED" if _any_test_failed else "FINISHED"
|
|
try:
|
|
mlflow.end_run(status=status)
|
|
except Exception as e:
|
|
_logger.warning("Failed to end run %s: %s", _run_id, e)
|
|
_run_id = None
|
|
_run_owned = False
|