wehub-resource-sync
62 lines
1.6 KiB
Python
62 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import textwrap
|
|
from pathlib import Path
|
|
|
|
from mlflow.pytest.session import TAG_TEST_NAME
|
|
from mlflow.tracking import MlflowClient
|
|
|
|
_GENERATED_TEST = """
|
|
import mlflow
|
|
from mlflow.genai.scorers import scorer
|
|
|
|
|
|
@scorer
|
|
def always_pass(*, outputs):
|
|
return True
|
|
|
|
|
|
@mlflow.test
|
|
def test_marked():
|
|
result = mlflow.genai.evaluate(
|
|
data=[{"inputs": {"text": "hi"}, "outputs": "hi"}],
|
|
scorers=[always_pass],
|
|
)
|
|
assert result.passed, result.reason
|
|
"""
|
|
|
|
|
|
def test_evaluate_traces_tagged_with_test_identity(tmp_path: Path):
|
|
test_file = tmp_path / "test_generated.py"
|
|
test_file.write_text(textwrap.dedent(_GENERATED_TEST))
|
|
|
|
tracking_uri = f"sqlite:///{tmp_path / 'mlflow.db'}"
|
|
result = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"pytest",
|
|
test_file.name,
|
|
"-p",
|
|
"no:cacheprovider",
|
|
# The plugin is opt-in (no pytest11 entry point).
|
|
"-p",
|
|
"mlflow.pytest.plugin",
|
|
"-q",
|
|
],
|
|
cwd=tmp_path,
|
|
capture_output=True,
|
|
text=True,
|
|
env={**os.environ, "MLFLOW_TRACKING_URI": tracking_uri},
|
|
)
|
|
assert result.returncode == 0, result.stdout + result.stderr
|
|
|
|
client = MlflowClient(tracking_uri=tracking_uri)
|
|
experiment_ids = [e.experiment_id for e in client.search_experiments()]
|
|
traces = client.search_traces(experiment_ids=experiment_ids) if experiment_ids else []
|
|
assert traces
|
|
assert all(t.info.tags.get(TAG_TEST_NAME, "").endswith("::test_marked") for t in traces)
|