wehub-resource-sync
5530 lines
197 KiB
Python
5530 lines
197 KiB
Python
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import pathlib
|
|
import posixpath
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
import urllib.parse
|
|
from dataclasses import asdict
|
|
from io import StringIO
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import flask
|
|
import pandas as pd
|
|
import pytest
|
|
import requests
|
|
from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
|
|
|
|
import mlflow.experiments
|
|
import mlflow.pyfunc
|
|
import mlflow.tracing.trace_archival_config as trace_archival_config_module
|
|
from mlflow import MlflowClient
|
|
from mlflow.artifacts import download_artifacts
|
|
from mlflow.data.pandas_dataset import from_pandas
|
|
from mlflow.entities import (
|
|
Dataset,
|
|
DatasetInput,
|
|
FallbackConfig,
|
|
FallbackStrategy,
|
|
GatewayEndpointModelConfig,
|
|
GatewayModelLinkageType,
|
|
GatewayResourceType,
|
|
InputTag,
|
|
IssueSeverity,
|
|
IssueStatus,
|
|
Metric,
|
|
Param,
|
|
RoutingStrategy,
|
|
RunInputs,
|
|
RunTag,
|
|
Span,
|
|
SpanEvent,
|
|
SpanStatusCode,
|
|
ViewType,
|
|
)
|
|
from mlflow.entities.logged_model_input import LoggedModelInput
|
|
from mlflow.entities.logged_model_output import LoggedModelOutput
|
|
from mlflow.entities.logged_model_status import LoggedModelStatus
|
|
from mlflow.entities.span import SpanAttributeKey
|
|
from mlflow.entities.trace_data import TraceData
|
|
from mlflow.entities.trace_info import TraceInfo
|
|
from mlflow.entities.trace_location import TraceLocation
|
|
from mlflow.entities.trace_metrics import (
|
|
AggregationType,
|
|
MetricAggregation,
|
|
MetricViewType,
|
|
)
|
|
from mlflow.entities.trace_state import TraceState
|
|
from mlflow.entities.trace_status import TraceStatus
|
|
from mlflow.environment_variables import (
|
|
_MLFLOW_GO_STORE_TESTING,
|
|
MLFLOW_SERVER_GRAPHQL_MAX_ALIASES,
|
|
MLFLOW_SERVER_GRAPHQL_MAX_ROOT_FIELDS,
|
|
MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT,
|
|
MLFLOW_TRACE_ARCHIVAL_CONFIG,
|
|
)
|
|
from mlflow.exceptions import MlflowException, RestException
|
|
from mlflow.genai.datasets import (
|
|
add_dataset_to_experiments,
|
|
create_dataset,
|
|
remove_dataset_from_experiments,
|
|
)
|
|
from mlflow.models import Model
|
|
from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST, ErrorCode
|
|
from mlflow.server import handlers
|
|
from mlflow.server.fastapi_app import app
|
|
from mlflow.server.handlers import initialize_backend_stores
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
from mlflow.tracing.analysis import TraceFilterCorrelationResult
|
|
from mlflow.tracing.client import TracingClient
|
|
from mlflow.tracing.constant import (
|
|
TRACE_SCHEMA_VERSION_KEY,
|
|
TraceMetricDimensionKey,
|
|
TraceMetricKey,
|
|
)
|
|
from mlflow.tracing.utils import build_otel_context
|
|
from mlflow.utils import mlflow_tags
|
|
from mlflow.utils.file_utils import TempDir, path_to_local_file_uri
|
|
from mlflow.utils.mlflow_tags import (
|
|
MLFLOW_DATASET_CONTEXT,
|
|
MLFLOW_GIT_COMMIT,
|
|
MLFLOW_PARENT_RUN_ID,
|
|
MLFLOW_PROJECT_ENTRY_POINT,
|
|
MLFLOW_SOURCE_NAME,
|
|
MLFLOW_SOURCE_TYPE,
|
|
MLFLOW_USER,
|
|
)
|
|
from mlflow.utils.os import is_windows
|
|
from mlflow.utils.proto_json_utils import message_to_json
|
|
from mlflow.utils.time import get_current_time_millis
|
|
|
|
from tests.helper_functions import get_safe_port
|
|
from tests.integration.utils import invoke_cli_runner
|
|
from tests.tracking.integration_test_utils import (
|
|
ServerThread,
|
|
_init_server,
|
|
_send_rest_tracking_post_request,
|
|
)
|
|
|
|
_logger = logging.getLogger(__name__)
|
|
|
|
|
|
@pytest.fixture(params=["file", "sqlalchemy"])
|
|
def store_type(request):
|
|
"""Provides the store type for parameterized tests."""
|
|
if request.param == "file":
|
|
pytest.skip("FileStore is no longer supported.")
|
|
return request.param
|
|
|
|
|
|
@pytest.fixture
|
|
def mlflow_client(store_type: str, tmp_path: Path, db_uri: str, monkeypatch):
|
|
"""Provides an MLflow Tracking API client pointed at the local tracking server."""
|
|
# Set passphrase for secrets management (required for encryption)
|
|
monkeypatch.setenv(
|
|
"MLFLOW_CRYPTO_KEK_PASSPHRASE", "test-passphrase-at-least-32-characters-long"
|
|
)
|
|
monkeypatch.delenv(MLFLOW_TRACE_ARCHIVAL_CONFIG.name, raising=False)
|
|
monkeypatch.setattr(trace_archival_config_module, "_TRACE_ARCHIVAL_SERVER_CONFIG_CACHE", None)
|
|
|
|
if store_type == "file":
|
|
backend_uri = tmp_path.joinpath("file").as_uri()
|
|
elif store_type == "sqlalchemy":
|
|
backend_uri = db_uri
|
|
|
|
# Force-reset backend stores before each test.
|
|
handlers._tracking_store = None
|
|
handlers._model_registry_store = None
|
|
initialize_backend_stores(backend_uri, default_artifact_root=tmp_path.as_uri())
|
|
|
|
with ServerThread(app, get_safe_port()) as url:
|
|
yield MlflowClient(url)
|
|
|
|
|
|
@pytest.fixture
|
|
def mlflow_client_with_secrets(tmp_path: Path, monkeypatch):
|
|
"""Provides an MLflow Tracking API client with fresh database for secrets management.
|
|
|
|
Creates a fresh SQLite database for each test to avoid encryption state pollution.
|
|
This is necessary because the KEK encryption state can persist across tests when
|
|
using a shared cached database.
|
|
"""
|
|
# Set passphrase for secrets management (required for encryption)
|
|
monkeypatch.setenv(
|
|
"MLFLOW_CRYPTO_KEK_PASSPHRASE", "test-passphrase-at-least-32-characters-long"
|
|
)
|
|
|
|
# Create fresh database for this test (not using cached_db)
|
|
backend_uri = f"sqlite:///{tmp_path}/mlflow.db"
|
|
artifact_uri = (tmp_path / "artifacts").as_uri()
|
|
|
|
# Initialize the store (which creates tables)
|
|
store = SqlAlchemyStore(backend_uri, artifact_uri)
|
|
store.engine.dispose()
|
|
|
|
# Force-reset backend stores before each test
|
|
handlers._tracking_store = None
|
|
handlers._model_registry_store = None
|
|
initialize_backend_stores(backend_uri, default_artifact_root=artifact_uri)
|
|
|
|
with ServerThread(app, get_safe_port()) as url:
|
|
yield MlflowClient(url)
|
|
|
|
|
|
@pytest.fixture
|
|
def cli_env(mlflow_client):
|
|
"""Provides an environment for the MLflow CLI pointed at the local tracking server."""
|
|
return {
|
|
"LC_ALL": "en_US.UTF-8",
|
|
"LANG": "en_US.UTF-8",
|
|
"MLFLOW_TRACKING_URI": mlflow_client.tracking_uri,
|
|
}
|
|
|
|
|
|
def create_experiments(client, names):
|
|
return [client.create_experiment(n) for n in names]
|
|
|
|
|
|
def test_create_get_search_experiment(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment(
|
|
"My Experiment",
|
|
artifact_location="my_location",
|
|
tags={"key1": "val1", "key2": "val2"},
|
|
)
|
|
exp = mlflow_client.get_experiment(experiment_id)
|
|
assert exp.name == "My Experiment"
|
|
if is_windows():
|
|
assert exp.artifact_location == pathlib.Path.cwd().joinpath("my_location").as_uri()
|
|
else:
|
|
assert exp.artifact_location == str(pathlib.Path.cwd().joinpath("my_location"))
|
|
assert len(exp.tags) == 2
|
|
assert exp.tags["key1"] == "val1"
|
|
assert exp.tags["key2"] == "val2"
|
|
|
|
experiments = mlflow_client.search_experiments()
|
|
assert {e.name for e in experiments} == {"My Experiment", "Default"}
|
|
mlflow_client.delete_experiment(experiment_id)
|
|
assert {e.name for e in mlflow_client.search_experiments()} == {"Default"}
|
|
assert {e.name for e in mlflow_client.search_experiments(view_type=ViewType.ACTIVE_ONLY)} == {
|
|
"Default"
|
|
}
|
|
assert {e.name for e in mlflow_client.search_experiments(view_type=ViewType.DELETED_ONLY)} == {
|
|
"My Experiment"
|
|
}
|
|
assert {e.name for e in mlflow_client.search_experiments(view_type=ViewType.ALL)} == {
|
|
"My Experiment",
|
|
"Default",
|
|
}
|
|
active_exps_paginated = mlflow_client.search_experiments(max_results=1)
|
|
assert {e.name for e in active_exps_paginated} == {"Default"}
|
|
assert active_exps_paginated.token is None
|
|
|
|
all_exps_paginated = mlflow_client.search_experiments(max_results=1, view_type=ViewType.ALL)
|
|
first_page_names = {e.name for e in all_exps_paginated}
|
|
all_exps_second_page = mlflow_client.search_experiments(
|
|
max_results=1, view_type=ViewType.ALL, page_token=all_exps_paginated.token
|
|
)
|
|
second_page_names = {e.name for e in all_exps_second_page}
|
|
assert len(first_page_names) == 1
|
|
assert len(second_page_names) == 1
|
|
assert first_page_names.union(second_page_names) == {"Default", "My Experiment"}
|
|
|
|
|
|
def test_create_experiment_validation(mlflow_client):
|
|
def assert_bad_request(payload, expected_error_message):
|
|
response = _send_rest_tracking_post_request(
|
|
mlflow_client.tracking_uri,
|
|
"/api/2.0/mlflow/experiments/create",
|
|
payload,
|
|
)
|
|
assert response.status_code == 400
|
|
assert expected_error_message in response.text
|
|
|
|
assert_bad_request(
|
|
{
|
|
"name": 123,
|
|
},
|
|
"Invalid value 123 for parameter 'name'",
|
|
)
|
|
assert_bad_request({}, "Missing value for required parameter 'name'.")
|
|
assert_bad_request(
|
|
{
|
|
"name": "experiment name",
|
|
"artifact_location": 9.0,
|
|
"tags": [{"key": "key", "value": "value"}],
|
|
},
|
|
"Invalid value 9.0 for parameter 'artifact_location'",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"name": "experiment name",
|
|
"artifact_location": "my_location",
|
|
"tags": "5",
|
|
},
|
|
"Invalid value \\\"5\\\" for parameter 'tags'",
|
|
)
|
|
|
|
|
|
def test_delete_restore_experiment(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("Deleterious")
|
|
assert mlflow_client.get_experiment(experiment_id).lifecycle_stage == "active"
|
|
mlflow_client.delete_experiment(experiment_id)
|
|
assert mlflow_client.get_experiment(experiment_id).lifecycle_stage == "deleted"
|
|
mlflow_client.restore_experiment(experiment_id)
|
|
assert mlflow_client.get_experiment(experiment_id).lifecycle_stage == "active"
|
|
|
|
|
|
def test_delete_restore_experiment_cli(mlflow_client, cli_env):
|
|
experiment_name = "DeleteriousCLI"
|
|
invoke_cli_runner(
|
|
mlflow.experiments.commands,
|
|
["create", "--experiment-name", experiment_name],
|
|
env=cli_env,
|
|
)
|
|
experiment_id = mlflow_client.get_experiment_by_name(experiment_name).experiment_id
|
|
assert mlflow_client.get_experiment(experiment_id).lifecycle_stage == "active"
|
|
invoke_cli_runner(
|
|
mlflow.experiments.commands, ["delete", "-x", str(experiment_id)], env=cli_env
|
|
)
|
|
assert mlflow_client.get_experiment(experiment_id).lifecycle_stage == "deleted"
|
|
invoke_cli_runner(
|
|
mlflow.experiments.commands, ["restore", "-x", str(experiment_id)], env=cli_env
|
|
)
|
|
assert mlflow_client.get_experiment(experiment_id).lifecycle_stage == "active"
|
|
|
|
|
|
def test_rename_experiment(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("BadName")
|
|
assert mlflow_client.get_experiment(experiment_id).name == "BadName"
|
|
mlflow_client.rename_experiment(experiment_id, "GoodName")
|
|
assert mlflow_client.get_experiment(experiment_id).name == "GoodName"
|
|
|
|
|
|
def test_rename_experiment_cli(mlflow_client, cli_env):
|
|
bad_experiment_name = "CLIBadName"
|
|
good_experiment_name = "CLIGoodName"
|
|
|
|
invoke_cli_runner(
|
|
mlflow.experiments.commands, ["create", "-n", bad_experiment_name], env=cli_env
|
|
)
|
|
experiment_id = mlflow_client.get_experiment_by_name(bad_experiment_name).experiment_id
|
|
assert mlflow_client.get_experiment(experiment_id).name == bad_experiment_name
|
|
invoke_cli_runner(
|
|
mlflow.experiments.commands,
|
|
[
|
|
"rename",
|
|
"--experiment-id",
|
|
str(experiment_id),
|
|
"--new-name",
|
|
good_experiment_name,
|
|
],
|
|
env=cli_env,
|
|
)
|
|
assert mlflow_client.get_experiment(experiment_id).name == good_experiment_name
|
|
|
|
|
|
@pytest.mark.parametrize("parent_run_id_kwarg", [None, "my-parent-id"])
|
|
def test_create_run_all_args(mlflow_client, parent_run_id_kwarg):
|
|
user = "username"
|
|
source_name = "Hello"
|
|
entry_point = "entry"
|
|
source_version = "abc"
|
|
create_run_kwargs = {
|
|
"start_time": 456,
|
|
"run_name": "my name",
|
|
"tags": {
|
|
MLFLOW_USER: user,
|
|
MLFLOW_SOURCE_TYPE: "LOCAL",
|
|
MLFLOW_SOURCE_NAME: source_name,
|
|
MLFLOW_PROJECT_ENTRY_POINT: entry_point,
|
|
MLFLOW_GIT_COMMIT: source_version,
|
|
MLFLOW_PARENT_RUN_ID: "7",
|
|
"my": "tag",
|
|
"other": "tag",
|
|
},
|
|
}
|
|
experiment_id = mlflow_client.create_experiment(
|
|
f"Run A Lot (parent_run_id={parent_run_id_kwarg})"
|
|
)
|
|
created_run = mlflow_client.create_run(experiment_id, **create_run_kwargs)
|
|
run_id = created_run.info.run_id
|
|
_logger.info(f"Run id={run_id}")
|
|
fetched_run = mlflow_client.get_run(run_id)
|
|
for run in [created_run, fetched_run]:
|
|
assert run.info.run_id == run_id
|
|
assert run.info.experiment_id == experiment_id
|
|
assert run.info.user_id == user
|
|
assert run.info.start_time == create_run_kwargs["start_time"]
|
|
assert run.info.run_name == "my name"
|
|
for tag in create_run_kwargs["tags"]:
|
|
assert tag in run.data.tags
|
|
assert run.data.tags.get(MLFLOW_USER) == user
|
|
assert run.data.tags.get(MLFLOW_PARENT_RUN_ID) == parent_run_id_kwarg or "7"
|
|
assert [run.info for run in mlflow_client.search_runs([experiment_id])] == [run.info]
|
|
|
|
|
|
def test_create_run_defaults(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("Run A Little")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
run = mlflow_client.get_run(run_id)
|
|
assert run.info.run_id == run_id
|
|
assert run.info.experiment_id == experiment_id
|
|
assert run.info.user_id == "unknown"
|
|
|
|
|
|
def test_log_metrics_params_tags(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("Oh My")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
mlflow_client.log_metric(run_id, key="metric", value=123.456, timestamp=789, step=2)
|
|
mlflow_client.log_metric(run_id, key="nan_metric", value=float("nan"))
|
|
mlflow_client.log_metric(run_id, key="inf_metric", value=float("inf"))
|
|
mlflow_client.log_metric(run_id, key="-inf_metric", value=-float("inf"))
|
|
mlflow_client.log_metric(run_id, key="stepless-metric", value=987.654, timestamp=321)
|
|
mlflow_client.log_param(run_id, "param", "value")
|
|
mlflow_client.set_tag(run_id, "taggity", "do-dah")
|
|
run = mlflow_client.get_run(run_id)
|
|
assert run.data.metrics.get("metric") == 123.456
|
|
assert math.isnan(run.data.metrics.get("nan_metric"))
|
|
assert run.data.metrics.get("inf_metric") >= 1.7976931348623157e308
|
|
assert run.data.metrics.get("-inf_metric") <= -1.7976931348623157e308
|
|
assert run.data.metrics.get("stepless-metric") == 987.654
|
|
assert run.data.params.get("param") == "value"
|
|
assert run.data.tags.get("taggity") == "do-dah"
|
|
metric_history0 = mlflow_client.get_metric_history(run_id, "metric")
|
|
assert len(metric_history0) == 1
|
|
metric0 = metric_history0[0]
|
|
assert metric0.key == "metric"
|
|
assert metric0.value == 123.456
|
|
assert metric0.timestamp == 789
|
|
assert metric0.step == 2
|
|
metric_history1 = mlflow_client.get_metric_history(run_id, "stepless-metric")
|
|
assert len(metric_history1) == 1
|
|
metric1 = metric_history1[0]
|
|
assert metric1.key == "stepless-metric"
|
|
assert metric1.value == 987.654
|
|
assert metric1.timestamp == 321
|
|
assert metric1.step == 0
|
|
|
|
metric_history = mlflow_client.get_metric_history(run_id, "a_test_accuracy")
|
|
assert metric_history == []
|
|
|
|
|
|
def test_log_metric_validation(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("metrics validation")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
def assert_bad_request(payload, expected_error_message):
|
|
response = _send_rest_tracking_post_request(
|
|
mlflow_client.tracking_uri,
|
|
"/api/2.0/mlflow/runs/log-metric",
|
|
payload,
|
|
)
|
|
assert response.status_code == 400
|
|
assert expected_error_message in response.text
|
|
|
|
assert_bad_request(
|
|
{
|
|
"run_id": 31,
|
|
"key": "metric",
|
|
"value": 41,
|
|
"timestamp": 59,
|
|
"step": 26,
|
|
},
|
|
"Invalid value 31 for parameter 'run_id' supplied",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
"key": 31,
|
|
"value": 41,
|
|
"timestamp": 59,
|
|
"step": 26,
|
|
},
|
|
"Invalid value 31 for parameter 'key' supplied",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
"key": "foo",
|
|
"value": 31,
|
|
"timestamp": 59,
|
|
"step": "foo",
|
|
},
|
|
"Invalid value \\\"foo\\\" for parameter 'step' supplied",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
"key": "foo",
|
|
"value": 31,
|
|
"timestamp": "foo",
|
|
"step": 41,
|
|
},
|
|
"Invalid value \\\"foo\\\" for parameter 'timestamp' supplied",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": None,
|
|
"key": "foo",
|
|
"value": 31,
|
|
"timestamp": 59,
|
|
"step": 41,
|
|
},
|
|
"Missing value for required parameter 'run_id'",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
# Missing key
|
|
"value": 31,
|
|
"timestamp": 59,
|
|
"step": 41,
|
|
},
|
|
"Missing value for required parameter 'key'",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
"key": None,
|
|
"value": 31,
|
|
"timestamp": 59,
|
|
"step": 41,
|
|
},
|
|
"Missing value for required parameter 'key'",
|
|
)
|
|
|
|
|
|
def test_log_metric_model(mlflow_client: MlflowClient):
|
|
experiment_id = mlflow_client.create_experiment("metrics validation")
|
|
run = mlflow_client.create_run(experiment_id)
|
|
model = mlflow_client.create_logged_model(experiment_id)
|
|
mlflow_client.log_metric(
|
|
run.info.run_id,
|
|
key="metric",
|
|
value=0.5,
|
|
timestamp=123456789,
|
|
step=1,
|
|
dataset_name="name",
|
|
dataset_digest="digest",
|
|
model_id=model.model_id,
|
|
)
|
|
|
|
model = mlflow_client.get_logged_model(model.model_id)
|
|
assert model.metrics == [
|
|
Metric(
|
|
key="metric",
|
|
value=0.5,
|
|
timestamp=123456789,
|
|
step=1,
|
|
model_id=model.model_id,
|
|
dataset_name="name",
|
|
dataset_digest="digest",
|
|
run_id=run.info.run_id,
|
|
)
|
|
]
|
|
|
|
|
|
def test_log_param_validation(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("params validation")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
def assert_bad_request(payload, expected_error_message):
|
|
response = _send_rest_tracking_post_request(
|
|
mlflow_client.tracking_uri,
|
|
"/api/2.0/mlflow/runs/log-parameter",
|
|
payload,
|
|
)
|
|
assert response.status_code == 400
|
|
assert expected_error_message in response.text
|
|
|
|
assert_bad_request(
|
|
{
|
|
"run_id": 31,
|
|
"key": "param",
|
|
"value": 41,
|
|
},
|
|
"Invalid value 31 for parameter 'run_id' supplied",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
"key": 31,
|
|
"value": 41,
|
|
},
|
|
"Invalid value 31 for parameter 'key' supplied",
|
|
)
|
|
|
|
|
|
def test_log_param_with_empty_string_as_value(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment(
|
|
test_log_param_with_empty_string_as_value.__name__
|
|
)
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
mlflow_client.log_param(run_id, "param_key", "")
|
|
assert {"param_key": ""}.items() <= mlflow_client.get_run(run_id).data.params.items()
|
|
|
|
|
|
def test_set_tag_with_empty_string_as_value(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment(
|
|
test_set_tag_with_empty_string_as_value.__name__
|
|
)
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
mlflow_client.set_tag(run_id, "tag_key", "")
|
|
assert {"tag_key": ""}.items() <= mlflow_client.get_run(run_id).data.tags.items()
|
|
|
|
|
|
def test_log_batch_containing_params_and_tags_with_empty_string_values(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment(
|
|
test_log_batch_containing_params_and_tags_with_empty_string_values.__name__
|
|
)
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
mlflow_client.log_batch(
|
|
run_id=run_id,
|
|
params=[Param("param_key", "")],
|
|
tags=[RunTag("tag_key", "")],
|
|
)
|
|
assert {"param_key": ""}.items() <= mlflow_client.get_run(run_id).data.params.items()
|
|
assert {"tag_key": ""}.items() <= mlflow_client.get_run(run_id).data.tags.items()
|
|
|
|
|
|
def test_set_tag_validation(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("tags validation")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
def assert_bad_request(payload, expected_error_message):
|
|
response = _send_rest_tracking_post_request(
|
|
mlflow_client.tracking_uri,
|
|
"/api/2.0/mlflow/runs/set-tag",
|
|
payload,
|
|
)
|
|
assert response.status_code == 400
|
|
assert expected_error_message in response.text
|
|
|
|
assert_bad_request(
|
|
{
|
|
"run_id": 31,
|
|
"key": "tag",
|
|
"value": 41,
|
|
},
|
|
"Invalid value 31 for parameter 'run_id' supplied",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
"key": "param",
|
|
"value": 41,
|
|
},
|
|
"Invalid value 41 for parameter 'value' supplied",
|
|
)
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
# Missing key
|
|
"value": "value",
|
|
},
|
|
"Missing value for required parameter 'key'",
|
|
)
|
|
|
|
response = _send_rest_tracking_post_request(
|
|
mlflow_client.tracking_uri,
|
|
"/api/2.0/mlflow/runs/set-tag",
|
|
{
|
|
"run_uuid": run_id,
|
|
"key": "key",
|
|
"value": "value",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
|
|
def test_path_validation(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("tags validation")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
invalid_path = "../path"
|
|
|
|
def assert_response(resp):
|
|
assert resp.status_code == 400
|
|
body = response.json()
|
|
assert body["error_code"] == "INVALID_PARAMETER_VALUE"
|
|
assert body["message"] == "Invalid path"
|
|
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/artifacts/list",
|
|
params={"run_id": run_id, "path": invalid_path},
|
|
)
|
|
assert_response(response)
|
|
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/get-artifact",
|
|
params={"run_id": run_id, "path": invalid_path},
|
|
)
|
|
assert_response(response)
|
|
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}//model-versions/get-artifact",
|
|
params={"name": "model", "version": 1, "path": invalid_path},
|
|
)
|
|
assert_response(response)
|
|
|
|
|
|
def test_set_experiment_tag(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("SetExperimentTagTest")
|
|
mlflow_client.set_experiment_tag(experiment_id, "dataset", "imagenet1K")
|
|
experiment = mlflow_client.get_experiment(experiment_id)
|
|
assert "dataset" in experiment.tags
|
|
assert experiment.tags["dataset"] == "imagenet1K"
|
|
# test that updating a tag works
|
|
mlflow_client.set_experiment_tag(experiment_id, "dataset", "birdbike")
|
|
experiment = mlflow_client.get_experiment(experiment_id)
|
|
assert "dataset" in experiment.tags
|
|
assert experiment.tags["dataset"] == "birdbike"
|
|
# test that setting a tag on 1 experiment does not impact another experiment.
|
|
experiment_id_2 = mlflow_client.create_experiment("SetExperimentTagTest2")
|
|
experiment2 = mlflow_client.get_experiment(experiment_id_2)
|
|
assert len(experiment2.tags) == 0
|
|
# test that setting a tag on different experiments maintain different values across experiments
|
|
mlflow_client.set_experiment_tag(experiment_id_2, "dataset", "birds200")
|
|
experiment = mlflow_client.get_experiment(experiment_id)
|
|
experiment2 = mlflow_client.get_experiment(experiment_id_2)
|
|
assert "dataset" in experiment.tags
|
|
assert experiment.tags["dataset"] == "birdbike"
|
|
assert "dataset" in experiment2.tags
|
|
assert experiment2.tags["dataset"] == "birds200"
|
|
# test can set multi-line tags
|
|
mlflow_client.set_experiment_tag(experiment_id, "multiline tag", "value2\nvalue2\nvalue2")
|
|
experiment = mlflow_client.get_experiment(experiment_id)
|
|
assert "multiline tag" in experiment.tags
|
|
assert experiment.tags["multiline tag"] == "value2\nvalue2\nvalue2"
|
|
|
|
|
|
def test_set_experiment_tag_with_empty_string_as_value(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment(
|
|
test_set_experiment_tag_with_empty_string_as_value.__name__
|
|
)
|
|
mlflow_client.set_experiment_tag(experiment_id, "tag_key", "")
|
|
assert {"tag_key": ""}.items() <= mlflow_client.get_experiment(experiment_id).tags.items()
|
|
|
|
|
|
def test_delete_experiment_tag(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("DeleteExperimentTagTest")
|
|
mlflow_client.set_experiment_tag(experiment_id, "dataset", "imagenet1K")
|
|
experiment = mlflow_client.get_experiment(experiment_id)
|
|
assert experiment.tags["dataset"] == "imagenet1K"
|
|
# test that deleting a tag works
|
|
mlflow_client.delete_experiment_tag(experiment_id, "dataset")
|
|
experiment = mlflow_client.get_experiment(experiment_id)
|
|
assert "dataset" not in experiment.tags
|
|
|
|
|
|
def test_delete_tag(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("DeleteTagExperiment")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
mlflow_client.log_metric(run_id, key="metric", value=123.456, timestamp=789, step=2)
|
|
mlflow_client.log_metric(run_id, key="stepless-metric", value=987.654, timestamp=321)
|
|
mlflow_client.log_param(run_id, "param", "value")
|
|
mlflow_client.set_tag(run_id, "taggity", "do-dah")
|
|
run = mlflow_client.get_run(run_id)
|
|
assert "taggity" in run.data.tags
|
|
assert run.data.tags["taggity"] == "do-dah"
|
|
mlflow_client.delete_tag(run_id, "taggity")
|
|
run = mlflow_client.get_run(run_id)
|
|
assert "taggity" not in run.data.tags
|
|
with pytest.raises(MlflowException, match=r"Run .+ not found"):
|
|
mlflow_client.delete_tag("fake_run_id", "taggity")
|
|
with pytest.raises(MlflowException, match="No tag with name: fakeTag"):
|
|
mlflow_client.delete_tag(run_id, "fakeTag")
|
|
mlflow_client.delete_run(run_id)
|
|
with pytest.raises(MlflowException, match=f"The run {run_id} must be in"):
|
|
mlflow_client.delete_tag(run_id, "taggity")
|
|
|
|
|
|
def test_log_batch(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("Batch em up")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
mlflow_client.log_batch(
|
|
run_id=run_id,
|
|
metrics=[Metric("metric", 123.456, 789, 3)],
|
|
params=[Param("param", "value")],
|
|
tags=[RunTag("taggity", "do-dah")],
|
|
)
|
|
run = mlflow_client.get_run(run_id)
|
|
assert run.data.metrics.get("metric") == 123.456
|
|
assert run.data.params.get("param") == "value"
|
|
assert run.data.tags.get("taggity") == "do-dah"
|
|
metric_history = mlflow_client.get_metric_history(run_id, "metric")
|
|
assert len(metric_history) == 1
|
|
metric = metric_history[0]
|
|
assert metric.key == "metric"
|
|
assert metric.value == 123.456
|
|
assert metric.timestamp == 789
|
|
assert metric.step == 3
|
|
|
|
|
|
def test_log_batch_validation(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("log_batch validation")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
def assert_bad_request(payload, expected_error_message):
|
|
response = _send_rest_tracking_post_request(
|
|
mlflow_client.tracking_uri,
|
|
"/api/2.0/mlflow/runs/log-batch",
|
|
payload,
|
|
)
|
|
assert response.status_code == 400
|
|
assert expected_error_message in response.text
|
|
|
|
for request_parameter in ["metrics", "params", "tags"]:
|
|
assert_bad_request(
|
|
{
|
|
"run_id": run_id,
|
|
request_parameter: "foo",
|
|
},
|
|
f"Invalid value \\\"foo\\\" for parameter '{request_parameter}' supplied",
|
|
)
|
|
|
|
## Should 400 if missing timestamp
|
|
assert_bad_request(
|
|
{"run_id": run_id, "metrics": [{"key": "mae", "value": 2.5}]},
|
|
"Missing value for required parameter 'metrics[0].timestamp'",
|
|
)
|
|
|
|
## Should 200 if timestamp provided but step is not
|
|
response = _send_rest_tracking_post_request(
|
|
mlflow_client.tracking_uri,
|
|
"/api/2.0/mlflow/runs/log-batch",
|
|
{
|
|
"run_id": run_id,
|
|
"metrics": [{"key": "mae", "value": 2.5, "timestamp": 123456789}],
|
|
},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
|
|
@pytest.mark.xfail(reason="Tracking server does not support logged-model endpoints yet")
|
|
@pytest.mark.allow_infer_pip_requirements_fallback
|
|
def test_log_model(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("Log models")
|
|
with TempDir(chdr=True):
|
|
model_paths = [f"model/path/{i}" for i in range(3)]
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
with mlflow.start_run(experiment_id=experiment_id) as run:
|
|
for i, m in enumerate(model_paths):
|
|
mlflow.pyfunc.log_model(name=m, loader_module="mlflow.pyfunc")
|
|
mlflow.pyfunc.save_model(
|
|
m,
|
|
mlflow_model=Model(artifact_path=m, run_id=run.info.run_id),
|
|
loader_module="mlflow.pyfunc",
|
|
)
|
|
model = Model.load(os.path.join(m, "MLmodel"))
|
|
run = mlflow.get_run(run.info.run_id)
|
|
tag = run.data.tags["mlflow.log-model.history"]
|
|
models = json.loads(tag)
|
|
model.utc_time_created = models[i]["utc_time_created"]
|
|
|
|
history_model_meta = models[i].copy()
|
|
original_model_uuid = history_model_meta.pop("model_uuid")
|
|
model_meta = model.get_tags_dict().copy()
|
|
new_model_uuid = model_meta.pop("model_uuid")
|
|
assert history_model_meta == model_meta
|
|
assert original_model_uuid != new_model_uuid
|
|
assert len(models) == i + 1
|
|
for j in range(0, i + 1):
|
|
assert models[j]["artifact_path"] == model_paths[j]
|
|
|
|
|
|
def test_set_terminated_defaults(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("Terminator 1")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
assert mlflow_client.get_run(run_id).info.status == "RUNNING"
|
|
assert mlflow_client.get_run(run_id).info.end_time is None
|
|
mlflow_client.set_terminated(run_id)
|
|
assert mlflow_client.get_run(run_id).info.status == "FINISHED"
|
|
assert mlflow_client.get_run(run_id).info.end_time <= get_current_time_millis()
|
|
|
|
|
|
def test_set_terminated_status(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("Terminator 2")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
assert mlflow_client.get_run(run_id).info.status == "RUNNING"
|
|
assert mlflow_client.get_run(run_id).info.end_time is None
|
|
mlflow_client.set_terminated(run_id, "FAILED")
|
|
assert mlflow_client.get_run(run_id).info.status == "FAILED"
|
|
assert mlflow_client.get_run(run_id).info.end_time <= get_current_time_millis()
|
|
|
|
|
|
def test_artifacts(mlflow_client, tmp_path):
|
|
experiment_id = mlflow_client.create_experiment("Art In Fact")
|
|
experiment_info = mlflow_client.get_experiment(experiment_id)
|
|
assert experiment_info.artifact_location.startswith(path_to_local_file_uri(str(tmp_path)))
|
|
artifact_path = urllib.parse.urlparse(experiment_info.artifact_location).path
|
|
assert posixpath.split(artifact_path)[-1] == experiment_id
|
|
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
assert created_run.info.artifact_uri.startswith(experiment_info.artifact_location)
|
|
run_id = created_run.info.run_id
|
|
src_dir = tmp_path.joinpath("test_artifacts_src")
|
|
src_dir.mkdir()
|
|
src_file = os.path.join(src_dir, "my.file")
|
|
with open(src_file, "w") as f:
|
|
f.write("Hello, World!")
|
|
mlflow_client.log_artifact(run_id, src_file, None)
|
|
mlflow_client.log_artifacts(run_id, src_dir, "dir")
|
|
|
|
root_artifacts_list = mlflow_client.list_artifacts(run_id)
|
|
assert {a.path for a in root_artifacts_list} == {"my.file", "dir"}
|
|
|
|
dir_artifacts_list = mlflow_client.list_artifacts(run_id, "dir")
|
|
assert {a.path for a in dir_artifacts_list} == {"dir/my.file"}
|
|
|
|
all_artifacts = download_artifacts(
|
|
run_id=run_id, artifact_path=".", tracking_uri=mlflow_client.tracking_uri
|
|
)
|
|
with open(f"{all_artifacts}/my.file") as f:
|
|
assert f.read() == "Hello, World!"
|
|
with open(f"{all_artifacts}/dir/my.file") as f:
|
|
assert f.read() == "Hello, World!"
|
|
|
|
dir_artifacts = download_artifacts(
|
|
run_id=run_id, artifact_path="dir", tracking_uri=mlflow_client.tracking_uri
|
|
)
|
|
with open(f"{dir_artifacts}/my.file") as f:
|
|
assert f.read() == "Hello, World!"
|
|
|
|
|
|
def test_search_pagination(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("search_pagination")
|
|
runs = [mlflow_client.create_run(experiment_id, start_time=1).info.run_id for _ in range(0, 10)]
|
|
runs = sorted(runs)
|
|
result = mlflow_client.search_runs([experiment_id], max_results=4, page_token=None)
|
|
assert [r.info.run_id for r in result] == runs[0:4]
|
|
assert result.token is not None
|
|
result = mlflow_client.search_runs([experiment_id], max_results=4, page_token=result.token)
|
|
assert [r.info.run_id for r in result] == runs[4:8]
|
|
assert result.token is not None
|
|
result = mlflow_client.search_runs([experiment_id], max_results=4, page_token=result.token)
|
|
assert [r.info.run_id for r in result] == runs[8:]
|
|
assert result.token is None
|
|
|
|
|
|
def test_search_validation(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("search_validation")
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=r"Invalid value 123456789 for parameter 'max_results' supplied",
|
|
):
|
|
mlflow_client.search_runs([experiment_id], max_results=123456789)
|
|
|
|
|
|
def test_get_experiment_by_name(mlflow_client):
|
|
name = "test_get_experiment_by_name"
|
|
experiment_id = mlflow_client.create_experiment(name)
|
|
res = mlflow_client.get_experiment_by_name(name)
|
|
assert res.experiment_id == experiment_id
|
|
assert res.name == name
|
|
assert mlflow_client.get_experiment_by_name("idontexist") is None
|
|
|
|
|
|
def test_get_experiment(mlflow_client):
|
|
name = "test_get_experiment"
|
|
experiment_id = mlflow_client.create_experiment(name)
|
|
res = mlflow_client.get_experiment(experiment_id)
|
|
assert res.experiment_id == experiment_id
|
|
assert res.name == name
|
|
|
|
|
|
def test_search_experiments(mlflow_client):
|
|
# To ensure the default experiment and non-default experiments have different creation_time
|
|
# for deterministic search results, send a request to the server and initialize the tracking
|
|
# store.
|
|
assert mlflow_client.search_experiments()[0].name == "Default"
|
|
|
|
experiments = [
|
|
("a", {"key": "value"}),
|
|
("ab", {"key": "vaLue"}),
|
|
("Abc", None),
|
|
]
|
|
experiment_ids = []
|
|
for name, tags in experiments:
|
|
# sleep for windows file system current_time precision in Python to enforce
|
|
# deterministic ordering based on last_update_time (creation_time due to no
|
|
# mutation of experiment state)
|
|
time.sleep(0.001)
|
|
experiment_ids.append(mlflow_client.create_experiment(name, tags=tags))
|
|
|
|
# filter_string
|
|
experiments = mlflow_client.search_experiments(filter_string="attribute.name = 'a'")
|
|
assert [e.name for e in experiments] == ["a"]
|
|
experiments = mlflow_client.search_experiments(filter_string="attribute.name != 'a'")
|
|
assert [e.name for e in experiments] == ["Abc", "ab", "Default"]
|
|
experiments = mlflow_client.search_experiments(filter_string="name LIKE 'a%'")
|
|
assert [e.name for e in experiments] == ["ab", "a"]
|
|
experiments = mlflow_client.search_experiments(filter_string="tag.key = 'value'")
|
|
assert [e.name for e in experiments] == ["a"]
|
|
experiments = mlflow_client.search_experiments(filter_string="tag.key != 'value'")
|
|
assert [e.name for e in experiments] == ["ab"]
|
|
experiments = mlflow_client.search_experiments(filter_string="tag.key ILIKE '%alu%'")
|
|
assert [e.name for e in experiments] == ["ab", "a"]
|
|
|
|
# order_by
|
|
experiments = mlflow_client.search_experiments(order_by=["name DESC"])
|
|
assert [e.name for e in experiments] == ["ab", "a", "Default", "Abc"]
|
|
|
|
# max_results
|
|
experiments = mlflow_client.search_experiments(max_results=2)
|
|
assert [e.name for e in experiments] == ["Abc", "ab"]
|
|
# page_token
|
|
experiments = mlflow_client.search_experiments(page_token=experiments.token)
|
|
assert [e.name for e in experiments] == ["a", "Default"]
|
|
|
|
# view_type
|
|
time.sleep(0.001)
|
|
mlflow_client.delete_experiment(experiment_ids[1])
|
|
experiments = mlflow_client.search_experiments(view_type=ViewType.ACTIVE_ONLY)
|
|
assert [e.name for e in experiments] == ["Abc", "a", "Default"]
|
|
experiments = mlflow_client.search_experiments(view_type=ViewType.DELETED_ONLY)
|
|
assert [e.name for e in experiments] == ["ab"]
|
|
experiments = mlflow_client.search_experiments(view_type=ViewType.ALL)
|
|
assert [e.name for e in experiments] == ["Abc", "ab", "a", "Default"]
|
|
|
|
|
|
def test_get_metric_history_bulk_rejects_invalid_requests(mlflow_client):
|
|
def assert_response(resp, message_part):
|
|
assert resp.status_code == 400
|
|
response_json = resp.json()
|
|
assert response_json.get("error_code") == "INVALID_PARAMETER_VALUE"
|
|
assert message_part in response_json.get("message", "")
|
|
|
|
response_no_run_ids_field = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={"metric_key": "key"},
|
|
)
|
|
assert_response(
|
|
response_no_run_ids_field,
|
|
"GetMetricHistoryBulk request must specify at least one run_id",
|
|
)
|
|
|
|
response_empty_run_ids = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={"run_id": [], "metric_key": "key"},
|
|
)
|
|
assert_response(
|
|
response_empty_run_ids,
|
|
"GetMetricHistoryBulk request must specify at least one run_id",
|
|
)
|
|
|
|
response_too_many_run_ids = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={"run_id": [f"id_{i}" for i in range(1000)], "metric_key": "key"},
|
|
)
|
|
assert_response(
|
|
response_too_many_run_ids,
|
|
"GetMetricHistoryBulk request cannot specify more than",
|
|
)
|
|
|
|
response_no_metric_key_field = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={"run_id": ["123"]},
|
|
)
|
|
assert_response(
|
|
response_no_metric_key_field,
|
|
"GetMetricHistoryBulk request must specify a metric_key",
|
|
)
|
|
|
|
|
|
def test_get_metric_history_bulk_returns_expected_metrics_in_expected_order(
|
|
mlflow_client,
|
|
):
|
|
experiment_id = mlflow_client.create_experiment("get metric history bulk")
|
|
created_run1 = mlflow_client.create_run(experiment_id)
|
|
run_id1 = created_run1.info.run_id
|
|
created_run2 = mlflow_client.create_run(experiment_id)
|
|
run_id2 = created_run2.info.run_id
|
|
created_run3 = mlflow_client.create_run(experiment_id)
|
|
run_id3 = created_run3.info.run_id
|
|
|
|
metricA_history = [
|
|
{"key": "metricA", "timestamp": 1, "step": 2, "value": 10.0},
|
|
{"key": "metricA", "timestamp": 1, "step": 3, "value": 11.0},
|
|
{"key": "metricA", "timestamp": 1, "step": 3, "value": 12.0},
|
|
{"key": "metricA", "timestamp": 2, "step": 3, "value": 12.0},
|
|
]
|
|
for metric in metricA_history:
|
|
mlflow_client.log_metric(run_id1, **metric)
|
|
metric_for_run2 = dict(metric)
|
|
metric_for_run2["value"] += 1.0
|
|
mlflow_client.log_metric(run_id2, **metric_for_run2)
|
|
|
|
metricB_history = [
|
|
{"key": "metricB", "timestamp": 7, "step": -2, "value": -100.0},
|
|
{"key": "metricB", "timestamp": 8, "step": 0, "value": 0.0},
|
|
{"key": "metricB", "timestamp": 8, "step": 0, "value": 1.0},
|
|
{"key": "metricB", "timestamp": 9, "step": 1, "value": 12.0},
|
|
]
|
|
for metric in metricB_history:
|
|
mlflow_client.log_metric(run_id1, **metric)
|
|
metric_for_run2 = dict(metric)
|
|
metric_for_run2["value"] += 1.0
|
|
mlflow_client.log_metric(run_id2, **metric_for_run2)
|
|
|
|
response_run1_metricA = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={"run_id": [run_id1], "metric_key": "metricA"},
|
|
)
|
|
assert response_run1_metricA.status_code == 200
|
|
assert response_run1_metricA.json().get("metrics") == [
|
|
{**metric, "run_id": run_id1} for metric in metricA_history
|
|
]
|
|
|
|
response_run2_metricB = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={"run_id": [run_id2], "metric_key": "metricB"},
|
|
)
|
|
assert response_run2_metricB.status_code == 200
|
|
assert response_run2_metricB.json().get("metrics") == [
|
|
{**metric, "run_id": run_id2, "value": metric["value"] + 1.0} for metric in metricB_history
|
|
]
|
|
|
|
response_run1_run2_metricA = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={"run_id": [run_id1, run_id2], "metric_key": "metricA"},
|
|
)
|
|
assert response_run1_run2_metricA.status_code == 200
|
|
assert response_run1_run2_metricA.json().get("metrics") == sorted(
|
|
[{**metric, "run_id": run_id1} for metric in metricA_history]
|
|
+ [
|
|
{**metric, "run_id": run_id2, "value": metric["value"] + 1.0}
|
|
for metric in metricA_history
|
|
],
|
|
key=lambda metric: metric["run_id"],
|
|
)
|
|
|
|
response_run1_run2_run_3_metricB = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={"run_id": [run_id1, run_id2, run_id3], "metric_key": "metricB"},
|
|
)
|
|
assert response_run1_run2_run_3_metricB.status_code == 200
|
|
assert response_run1_run2_run_3_metricB.json().get("metrics") == sorted(
|
|
[{**metric, "run_id": run_id1} for metric in metricB_history]
|
|
+ [
|
|
{**metric, "run_id": run_id2, "value": metric["value"] + 1.0}
|
|
for metric in metricB_history
|
|
],
|
|
key=lambda metric: metric["run_id"],
|
|
)
|
|
|
|
|
|
def test_get_metric_history_bulk_respects_max_results(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("get metric history bulk")
|
|
run_id = mlflow_client.create_run(experiment_id).info.run_id
|
|
max_results = 2
|
|
|
|
metricA_history = [
|
|
{"key": "metricA", "timestamp": 1, "step": 2, "value": 10.0},
|
|
{"key": "metricA", "timestamp": 1, "step": 3, "value": 11.0},
|
|
{"key": "metricA", "timestamp": 1, "step": 3, "value": 12.0},
|
|
{"key": "metricA", "timestamp": 2, "step": 3, "value": 12.0},
|
|
]
|
|
for metric in metricA_history:
|
|
mlflow_client.log_metric(run_id, **metric)
|
|
|
|
response_limited = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk",
|
|
params={
|
|
"run_id": [run_id],
|
|
"metric_key": "metricA",
|
|
"max_results": max_results,
|
|
},
|
|
)
|
|
assert response_limited.status_code == 200
|
|
assert response_limited.json().get("metrics") == [
|
|
{**metric, "run_id": run_id} for metric in metricA_history[:max_results]
|
|
]
|
|
|
|
|
|
def test_get_metric_history_bulk_calls_optimized_impl_when_expected(tmp_path):
|
|
from mlflow.server.handlers import get_metric_history_bulk_handler
|
|
|
|
path = path_to_local_file_uri(str(tmp_path.joinpath("sqlalchemy.db")))
|
|
uri = ("sqlite://" if sys.platform == "win32" else "sqlite:////") + path[len("file://") :]
|
|
mock_store = mock.Mock(wraps=SqlAlchemyStore(uri, str(tmp_path)))
|
|
|
|
flask_app = flask.Flask("test_flask_app")
|
|
|
|
class MockRequestArgs:
|
|
def __init__(self, args_dict):
|
|
self.args_dict = args_dict
|
|
|
|
def to_dict(
|
|
self,
|
|
flat,
|
|
):
|
|
return self.args_dict
|
|
|
|
def get(self, key, default=None):
|
|
return self.args_dict.get(key, default)
|
|
|
|
with (
|
|
mock.patch("mlflow.server.handlers._get_tracking_store", return_value=mock_store),
|
|
flask_app.test_request_context() as mock_context,
|
|
):
|
|
run_ids = [str(i) for i in range(10)]
|
|
mock_context.request.args = MockRequestArgs({
|
|
"run_id": run_ids,
|
|
"metric_key": "mock_key",
|
|
})
|
|
|
|
get_metric_history_bulk_handler()
|
|
|
|
mock_store.get_metric_history_bulk.assert_called_once_with(
|
|
run_ids=run_ids,
|
|
metric_key="mock_key",
|
|
max_results=25000,
|
|
)
|
|
|
|
|
|
def test_get_metric_history_respects_max_results(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("test max_results")
|
|
run = mlflow_client.create_run(experiment_id)
|
|
run_id = run.info.run_id
|
|
|
|
metric_history = [
|
|
{"key": "test_metric", "value": float(i), "step": i, "timestamp": 1000 + i}
|
|
for i in range(5)
|
|
]
|
|
for metric in metric_history:
|
|
mlflow_client.log_metric(run_id, **metric)
|
|
|
|
# Test without max_results - should return all metrics
|
|
all_metrics = mlflow_client.get_metric_history(run_id, "test_metric")
|
|
assert len(all_metrics) == 5
|
|
|
|
# Test with max_results=3 - should return only 3 metrics
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history",
|
|
params={"run_id": run_id, "metric_key": "test_metric", "max_results": 3},
|
|
)
|
|
assert response.status_code == 200
|
|
response_data = response.json()
|
|
assert len(response_data["metrics"]) == 3
|
|
|
|
returned_metrics = response_data["metrics"]
|
|
for i, metric in enumerate(returned_metrics):
|
|
assert metric["key"] == "test_metric"
|
|
assert metric["value"] == float(i)
|
|
if _MLFLOW_GO_STORE_TESTING.get():
|
|
assert int(metric["step"]) == i
|
|
else:
|
|
assert metric["step"] == i
|
|
|
|
|
|
def test_get_metric_history_with_page_token(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("test page_token")
|
|
run = mlflow_client.create_run(experiment_id)
|
|
run_id = run.info.run_id
|
|
|
|
metric_history = [
|
|
{"key": "test_metric", "value": float(i), "step": i, "timestamp": 1000 + i}
|
|
for i in range(10)
|
|
]
|
|
for metric in metric_history:
|
|
mlflow_client.log_metric(run_id, **metric)
|
|
|
|
page_size = 4
|
|
|
|
first_response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history",
|
|
params={
|
|
"run_id": run_id,
|
|
"metric_key": "test_metric",
|
|
"max_results": page_size,
|
|
},
|
|
)
|
|
assert first_response.status_code == 200
|
|
first_data = first_response.json()
|
|
first_metrics = first_data["metrics"]
|
|
first_token = first_data.get("next_page_token")
|
|
|
|
assert first_token is not None
|
|
assert len(first_metrics) == 4
|
|
|
|
second_response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history",
|
|
params={
|
|
"run_id": run_id,
|
|
"metric_key": "test_metric",
|
|
"max_results": page_size,
|
|
"page_token": first_token,
|
|
},
|
|
)
|
|
assert second_response.status_code == 200
|
|
second_data = second_response.json()
|
|
second_metrics = second_data["metrics"]
|
|
second_token = second_data.get("next_page_token")
|
|
|
|
assert second_token is not None
|
|
assert len(second_metrics) == 4
|
|
|
|
third_response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history",
|
|
params={
|
|
"run_id": run_id,
|
|
"metric_key": "test_metric",
|
|
"max_results": page_size,
|
|
"page_token": second_token,
|
|
},
|
|
)
|
|
assert third_response.status_code == 200
|
|
third_data = third_response.json()
|
|
third_metrics = third_data["metrics"]
|
|
third_token = third_data.get("next_page_token")
|
|
|
|
assert third_token is None
|
|
assert len(third_metrics) == 2
|
|
|
|
all_paginated_metrics = first_metrics + second_metrics + third_metrics
|
|
assert len(all_paginated_metrics) == 10
|
|
|
|
for i, metric in enumerate(all_paginated_metrics):
|
|
assert metric["key"] == "test_metric"
|
|
assert metric["value"] == float(i)
|
|
if _MLFLOW_GO_STORE_TESTING.get():
|
|
assert int(metric["step"]) == i
|
|
else:
|
|
assert metric["step"] == i
|
|
if _MLFLOW_GO_STORE_TESTING.get():
|
|
assert int(metric["timestamp"]) == 1000 + i
|
|
else:
|
|
assert metric["timestamp"] == 1000 + i
|
|
|
|
# Test with invalid page_token
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history",
|
|
params={
|
|
"run_id": run_id,
|
|
"metric_key": "test_metric",
|
|
"page_token": "invalid_token",
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
response_data = response.json()
|
|
assert "INVALID_PARAMETER_VALUE" in response_data.get("error_code", "")
|
|
|
|
|
|
def test_get_metric_history_without_max_results_returns_full_history(mlflow_client):
|
|
# Regression test: an unset proto2 `max_results` reads as 0, which previously became
|
|
# a `LIMIT 1` query that returned an empty page with a never-advancing next_page_token
|
|
experiment_id = mlflow_client.create_experiment("test no max_results")
|
|
run = mlflow_client.create_run(experiment_id)
|
|
run_id = run.info.run_id
|
|
|
|
for i in range(10):
|
|
mlflow_client.log_metric(run_id, key="test_metric", value=float(i), step=i)
|
|
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history",
|
|
params={"run_id": run_id, "metric_key": "test_metric"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert len(data["metrics"]) == 10
|
|
assert data.get("next_page_token") is None
|
|
|
|
|
|
@pytest.mark.parametrize("max_results", [0, -5])
|
|
def test_get_metric_history_rejects_non_positive_max_results(mlflow_client, max_results):
|
|
experiment_id = mlflow_client.create_experiment(f"test max_results {max_results}")
|
|
run = mlflow_client.create_run(experiment_id)
|
|
run_id = run.info.run_id
|
|
mlflow_client.log_metric(run_id, key="test_metric", value=1.0, step=0)
|
|
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history",
|
|
params={"run_id": run_id, "metric_key": "test_metric", "max_results": max_results},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "max_results" in response.text
|
|
|
|
|
|
def test_get_metric_history_bulk_interval_rejects_invalid_requests(mlflow_client):
|
|
def assert_response(resp, message_part):
|
|
assert resp.status_code == 400
|
|
response_json = resp.json()
|
|
assert response_json.get("error_code") == "INVALID_PARAMETER_VALUE"
|
|
assert message_part in response_json.get("message", "")
|
|
|
|
url = f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk-interval"
|
|
|
|
assert_response(
|
|
requests.get(url, params={"metric_key": "key"}),
|
|
"Missing value for required parameter 'run_ids'.",
|
|
)
|
|
|
|
assert_response(
|
|
requests.get(url, params={"run_ids": [], "metric_key": "key"}),
|
|
"Missing value for required parameter 'run_ids'.",
|
|
)
|
|
|
|
assert_response(
|
|
requests.get(
|
|
url,
|
|
params={"run_ids": [f"id_{i}" for i in range(1000)], "metric_key": "key"},
|
|
),
|
|
"GetMetricHistoryBulkInterval request must specify at most 100 run_ids.",
|
|
)
|
|
|
|
assert_response(
|
|
requests.get(url, params={"run_ids": ["123"], "metric_key": "key", "max_results": 0}),
|
|
"max_results must be between 1 and 2500",
|
|
)
|
|
|
|
assert_response(
|
|
requests.get(url, params={"run_ids": ["123"], "metric_key": ""}),
|
|
"Missing value for required parameter 'metric_key'",
|
|
)
|
|
|
|
assert_response(
|
|
requests.get(url, params={"run_ids": ["123"], "max_results": 5}),
|
|
"Missing value for required parameter 'metric_key'",
|
|
)
|
|
|
|
assert_response(
|
|
requests.get(
|
|
url,
|
|
params={
|
|
"run_ids": ["123"],
|
|
"metric_key": "key",
|
|
"start_step": 1,
|
|
"end_step": 0,
|
|
"max_results": 5,
|
|
},
|
|
),
|
|
"end_step must be greater than start_step. ",
|
|
)
|
|
|
|
assert_response(
|
|
requests.get(
|
|
url,
|
|
params={
|
|
"run_ids": ["123"],
|
|
"metric_key": "key",
|
|
"start_step": 1,
|
|
"max_results": 5,
|
|
},
|
|
),
|
|
"If either start step or end step are specified, both must be specified.",
|
|
)
|
|
|
|
|
|
def test_get_metric_history_bulk_interval_respects_max_results(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("get metric history bulk")
|
|
run_id1 = mlflow_client.create_run(experiment_id).info.run_id
|
|
metric_history = [
|
|
{"key": "metricA", "timestamp": 1, "step": i, "value": 10.0} for i in range(10)
|
|
]
|
|
for metric in metric_history:
|
|
mlflow_client.log_metric(run_id1, **metric)
|
|
|
|
url = f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/metrics/get-history-bulk-interval"
|
|
response_limited = requests.get(
|
|
url,
|
|
params={"run_ids": [run_id1], "metric_key": "metricA", "max_results": 5},
|
|
)
|
|
assert response_limited.status_code == 200
|
|
expected_steps = [0, 2, 4, 6, 8, 9]
|
|
expected_metrics = [
|
|
{**metric, "run_id": run_id1}
|
|
for metric in metric_history
|
|
if metric["step"] in expected_steps
|
|
]
|
|
assert response_limited.json().get("metrics") == expected_metrics
|
|
|
|
# with start_step and end_step
|
|
response_limited = requests.get(
|
|
url,
|
|
params={
|
|
"run_ids": [run_id1],
|
|
"metric_key": "metricA",
|
|
"start_step": 0,
|
|
"end_step": 4,
|
|
"max_results": 5,
|
|
},
|
|
)
|
|
assert response_limited.status_code == 200
|
|
assert response_limited.json().get("metrics") == [
|
|
{**metric, "run_id": run_id1} for metric in metric_history[:5]
|
|
]
|
|
|
|
# multiple runs
|
|
run_id2 = mlflow_client.create_run(experiment_id).info.run_id
|
|
metric_history2 = [
|
|
{"key": "metricA", "timestamp": 1, "step": i, "value": 10.0} for i in range(20)
|
|
]
|
|
for metric in metric_history2:
|
|
mlflow_client.log_metric(run_id2, **metric)
|
|
response_limited = requests.get(
|
|
url,
|
|
params={
|
|
"run_ids": [run_id1, run_id2],
|
|
"metric_key": "metricA",
|
|
"max_results": 5,
|
|
},
|
|
)
|
|
expected_steps = [0, 4, 8, 9, 12, 16, 19]
|
|
expected_metrics = []
|
|
for run_id, metric_history in [
|
|
(run_id1, metric_history),
|
|
(run_id2, metric_history2),
|
|
]:
|
|
expected_metrics.extend([
|
|
{**metric, "run_id": run_id}
|
|
for metric in metric_history
|
|
if metric["step"] in expected_steps
|
|
])
|
|
assert response_limited.json().get("metrics") == expected_metrics
|
|
|
|
# test metrics with same steps
|
|
metric_history_timestamp2 = [
|
|
{"key": "metricA", "timestamp": 2, "step": i, "value": 10.0} for i in range(10)
|
|
]
|
|
for metric in metric_history_timestamp2:
|
|
mlflow_client.log_metric(run_id1, **metric)
|
|
|
|
response_limited = requests.get(
|
|
url,
|
|
params={"run_ids": [run_id1], "metric_key": "metricA", "max_results": 5},
|
|
)
|
|
assert response_limited.status_code == 200
|
|
expected_steps = [0, 2, 4, 6, 8, 9]
|
|
expected_metrics = [
|
|
{"key": "metricA", "timestamp": j, "step": i, "value": 10.0, "run_id": run_id1}
|
|
for i in expected_steps
|
|
for j in [1, 2]
|
|
]
|
|
assert response_limited.json().get("metrics") == expected_metrics
|
|
|
|
|
|
def test_search_dataset_handler_rejects_invalid_requests(mlflow_client):
|
|
def assert_response(resp, message_part):
|
|
assert resp.status_code == 400
|
|
response_json = resp.json()
|
|
assert response_json.get("error_code") == "INVALID_PARAMETER_VALUE"
|
|
assert message_part in response_json.get("message", "")
|
|
|
|
response_no_experiment_id_field = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/experiments/search-datasets",
|
|
json={},
|
|
)
|
|
assert_response(
|
|
response_no_experiment_id_field,
|
|
"SearchDatasets request must specify at least one experiment_id.",
|
|
)
|
|
|
|
response_empty_experiment_id_field = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/experiments/search-datasets",
|
|
json={"experiment_ids": []},
|
|
)
|
|
assert_response(
|
|
response_empty_experiment_id_field,
|
|
"SearchDatasets request must specify at least one experiment_id.",
|
|
)
|
|
|
|
response_too_many_experiment_ids = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/experiments/search-datasets",
|
|
json={"experiment_ids": [f"id_{i}" for i in range(1000)]},
|
|
)
|
|
assert_response(
|
|
response_too_many_experiment_ids,
|
|
"SearchDatasets request cannot specify more than",
|
|
)
|
|
|
|
|
|
def test_search_dataset_handler_returns_expected_results(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("log inputs test")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
dataset1 = Dataset(
|
|
name="name1",
|
|
digest="digest1",
|
|
source_type="source_type1",
|
|
source="source1",
|
|
)
|
|
dataset_inputs1 = [
|
|
DatasetInput(
|
|
dataset=dataset1,
|
|
tags=[InputTag(key=MLFLOW_DATASET_CONTEXT, value="training")],
|
|
)
|
|
]
|
|
mlflow_client.log_inputs(run_id, dataset_inputs1)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/experiments/search-datasets",
|
|
json={"experiment_ids": [experiment_id]},
|
|
)
|
|
expected = {
|
|
"experiment_id": experiment_id,
|
|
"name": "name1",
|
|
"digest": "digest1",
|
|
"context": "training",
|
|
}
|
|
|
|
assert response.status_code == 200
|
|
assert response.json().get("dataset_summaries") == [expected]
|
|
|
|
|
|
def test_create_model_version_with_path_source(mlflow_client):
|
|
name = "model"
|
|
mlflow_client.create_registered_model(name)
|
|
exp_id = mlflow_client.create_experiment("test")
|
|
run = mlflow_client.create_run(experiment_id=exp_id)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": run.info.artifact_uri[len("file://") :],
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# run_id is not specified
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": run.info.artifact_uri[len("file://") :],
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "To use a local path as a model version" in response.json()["message"]
|
|
|
|
# run_id is specified but source is not in the run's artifact directory
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "/tmp",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "To use a local path as a model version" in response.json()["message"]
|
|
|
|
|
|
def test_create_model_version_with_non_local_source(mlflow_client):
|
|
name = "model"
|
|
mlflow_client.create_registered_model(name)
|
|
exp_id = mlflow_client.create_experiment("test")
|
|
run = mlflow_client.create_run(experiment_id=exp_id)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": run.info.artifact_uri[len("file://") :],
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# Test that remote uri's supplied as a source with absolute paths work fine
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts:/models",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# A single trailing slash
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts:/models/",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# Multiple trailing slashes
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts:/models///",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# Multiple slashes
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts:/models/foo///bar",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts://host:9000/models",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# Multiple dots
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts://host:9000/models/artifact/..../",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# Test that invalid remote uri's cannot be created
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts://host:9000/models/../../../",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "If supplying a source as an http, https," in response.json()["message"]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "http://host:9000/models/../../../",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "If supplying a source as an http, https," in response.json()["message"]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "https://host/api/2.0/mlflow-artifacts/artifacts/../../../",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "If supplying a source as an http, https," in response.json()["message"]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "s3a://my_bucket/api/2.0/mlflow-artifacts/artifacts/../../../",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "If supplying a source as an http, https," in response.json()["message"]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "ftp://host:8888/api/2.0/mlflow-artifacts/artifacts/../../../",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "If supplying a source as an http, https," in response.json()["message"]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts://host:9000/models/..%2f..%2fartifacts",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "If supplying a source as an http, https," in response.json()["message"]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "mlflow-artifacts://host:9000/models/artifact%00",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "If supplying a source as an http, https," in response.json()["message"]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": f"dbfs:/{run.info.run_id}/artifacts/a%3f/../../../../../../../../../../",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "Invalid model version source" in response.json()["message"]
|
|
|
|
model = mlflow_client.create_logged_model(experiment_id=exp_id)
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": model.artifact_location,
|
|
"model_id": model.model_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": model.model_uri,
|
|
"model_id": model.model_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "file:///path/to/model",
|
|
"model_id": model.model_id,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
|
|
|
|
def test_create_model_version_with_file_uri(mlflow_client):
|
|
name = "test"
|
|
mlflow_client.create_registered_model(name)
|
|
exp_id = mlflow_client.create_experiment("test")
|
|
run = mlflow_client.create_run(experiment_id=exp_id)
|
|
assert run.info.artifact_uri.startswith("file://")
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": run.info.artifact_uri,
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": f"{run.info.artifact_uri}/model",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": f"{run.info.artifact_uri}/.",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": f"{run.info.artifact_uri}/model/..",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
# run_id is not specified
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": run.info.artifact_uri,
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "To use a local path as a model version" in response.json()["message"]
|
|
|
|
# run_id is specified but source is not in the run's artifact directory
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "file:///tmp",
|
|
},
|
|
)
|
|
assert response.status_code == 400
|
|
assert "To use a local path as a model version" in response.json()["message"]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/2.0/mlflow/model-versions/create",
|
|
json={
|
|
"name": name,
|
|
"source": "file://123.456.789.123/path/to/source",
|
|
"run_id": run.info.run_id,
|
|
},
|
|
)
|
|
assert response.status_code == 500, response.json()
|
|
assert "is not a valid remote uri" in response.json()["message"]
|
|
|
|
|
|
def test_create_model_version_with_validation_regex(db_uri: str):
|
|
port = get_safe_port()
|
|
with subprocess.Popen(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"mlflow",
|
|
"server",
|
|
"--port",
|
|
str(port),
|
|
"--backend-store-uri",
|
|
db_uri,
|
|
],
|
|
env=(
|
|
os.environ.copy()
|
|
| {
|
|
"MLFLOW_CREATE_MODEL_VERSION_SOURCE_VALIDATION_REGEX": r"^mlflow-artifacts:/.*$",
|
|
"MLFLOW_SERVER_ENABLE_JOB_EXECUTION": "false",
|
|
}
|
|
),
|
|
) as proc:
|
|
try:
|
|
# Wait for the server to start
|
|
for _ in range(30):
|
|
try:
|
|
if requests.get(f"http://localhost:{port}/health").ok:
|
|
break
|
|
except requests.ConnectionError:
|
|
time.sleep(1)
|
|
else:
|
|
raise RuntimeError("Failed to connect to the MLflow server")
|
|
|
|
# Test that the validation regex works as expected
|
|
client = MlflowClient(f"http://localhost:{port}")
|
|
name = "test"
|
|
client.create_registered_model(name)
|
|
# Invalid source
|
|
with pytest.raises(MlflowException, match="Invalid model version source"):
|
|
client.create_model_version(name, source="s3://path/to/model")
|
|
# Valid source
|
|
experiment_id = client.create_experiment("test")
|
|
run = client.create_run(experiment_id=experiment_id)
|
|
assert run.info.artifact_uri.startswith("mlflow-artifacts:/")
|
|
client.create_model_version(
|
|
name, source=f"{run.info.artifact_uri}/model", run_id=run.info.run_id
|
|
)
|
|
finally:
|
|
proc.terminate()
|
|
proc.wait()
|
|
|
|
|
|
@pytest.mark.xfail(reason="Tracking server does not support logged-model endpoints yet")
|
|
def test_logging_model_with_local_artifact_uri(mlflow_client):
|
|
from sklearn.linear_model import LogisticRegression
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
with mlflow.start_run() as run:
|
|
assert run.info.artifact_uri.startswith("file://")
|
|
mlflow.sklearn.log_model(LogisticRegression(), name="model", registered_model_name="rmn")
|
|
mlflow.pyfunc.load_model("models:/rmn/1")
|
|
|
|
|
|
def test_log_input(mlflow_client, tmp_path):
|
|
df = pd.DataFrame([[1, 2, 3], [1, 2, 3]], columns=["a", "b", "c"])
|
|
path = tmp_path / "temp.csv"
|
|
df.to_csv(path)
|
|
dataset = from_pandas(df, source=path)
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
|
|
with mlflow.start_run() as run:
|
|
mlflow.log_input(dataset, "train", {"foo": "baz"})
|
|
|
|
dataset_inputs = mlflow_client.get_run(run.info.run_id).inputs.dataset_inputs
|
|
|
|
assert len(dataset_inputs) == 1
|
|
assert dataset_inputs[0].dataset.name == "dataset"
|
|
assert dataset_inputs[0].dataset.digest == "f0f3e026"
|
|
assert dataset_inputs[0].dataset.source_type == "local"
|
|
assert json.loads(dataset_inputs[0].dataset.source) == {"uri": str(path)}
|
|
assert json.loads(dataset_inputs[0].dataset.schema) == {
|
|
"mlflow_colspec": [
|
|
{"name": "a", "type": "long", "required": True},
|
|
{"name": "b", "type": "long", "required": True},
|
|
{"name": "c", "type": "long", "required": True},
|
|
]
|
|
}
|
|
assert json.loads(dataset_inputs[0].dataset.profile) == {
|
|
"num_rows": 2,
|
|
"num_elements": 6,
|
|
}
|
|
|
|
assert len(dataset_inputs[0].tags) == 2
|
|
assert dataset_inputs[0].tags[0].key == "foo"
|
|
assert dataset_inputs[0].tags[0].value == "baz"
|
|
assert dataset_inputs[0].tags[1].key == mlflow_tags.MLFLOW_DATASET_CONTEXT
|
|
assert dataset_inputs[0].tags[1].value == "train"
|
|
|
|
|
|
def test_create_model_version_model_id(mlflow_client):
|
|
name = "model"
|
|
mlflow_client.create_registered_model(name)
|
|
exp_id = mlflow_client.create_experiment("test")
|
|
model = mlflow_client.create_logged_model(experiment_id=exp_id)
|
|
mlflow_client.create_model_version(
|
|
name=name,
|
|
source=model.artifact_location,
|
|
model_id=model.model_id,
|
|
)
|
|
model = mlflow_client.get_logged_model(model.model_id)
|
|
assert model.tags["mlflow.modelVersions"] == '[{"name": "model", "version": 1}]'
|
|
mlflow_client.create_model_version(
|
|
name=name,
|
|
source=model.artifact_location,
|
|
model_id=model.model_id,
|
|
)
|
|
model = mlflow_client.get_logged_model(model.model_id)
|
|
assert (
|
|
model.tags["mlflow.modelVersions"]
|
|
== '[{"name": "model", "version": 1}, {"name": "model", "version": 2}]'
|
|
)
|
|
|
|
|
|
def test_log_inputs(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("log inputs test")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
dataset1 = Dataset(
|
|
name="name1",
|
|
digest="digest1",
|
|
source_type="source_type1",
|
|
source="source1",
|
|
)
|
|
dataset_inputs1 = [DatasetInput(dataset=dataset1, tags=[InputTag(key="tag1", value="value1")])]
|
|
|
|
mlflow_client.log_inputs(run_id, dataset_inputs1)
|
|
run = mlflow_client.get_run(run_id)
|
|
assert len(run.inputs.dataset_inputs) == 1
|
|
|
|
assert isinstance(run.inputs, RunInputs)
|
|
assert isinstance(run.inputs.dataset_inputs[0], DatasetInput)
|
|
assert isinstance(run.inputs.dataset_inputs[0].dataset, Dataset)
|
|
assert run.inputs.dataset_inputs[0].dataset.name == "name1"
|
|
assert run.inputs.dataset_inputs[0].dataset.digest == "digest1"
|
|
assert run.inputs.dataset_inputs[0].dataset.source_type == "source_type1"
|
|
assert run.inputs.dataset_inputs[0].dataset.source == "source1"
|
|
assert len(run.inputs.dataset_inputs[0].tags) == 1
|
|
assert run.inputs.dataset_inputs[0].tags[0].key == "tag1"
|
|
assert run.inputs.dataset_inputs[0].tags[0].value == "value1"
|
|
|
|
|
|
def test_log_inputs_validation(mlflow_client):
|
|
def assert_bad_request(payload, expected_error_message):
|
|
response = _send_rest_tracking_post_request(
|
|
mlflow_client.tracking_uri,
|
|
"/api/2.0/mlflow/runs/log-inputs",
|
|
payload,
|
|
)
|
|
assert response.status_code == 400
|
|
assert expected_error_message in response.text
|
|
|
|
dataset = Dataset(
|
|
name="name1",
|
|
digest="digest1",
|
|
source_type="source_type1",
|
|
source="source1",
|
|
)
|
|
tags = [InputTag(key="tag1", value="value1")]
|
|
dataset_inputs = [
|
|
json.loads(message_to_json(DatasetInput(dataset=dataset, tags=tags).to_proto()))
|
|
]
|
|
assert_bad_request(
|
|
{
|
|
"datasets": dataset_inputs,
|
|
},
|
|
"Missing value for required parameter 'run_id'",
|
|
)
|
|
|
|
|
|
def test_log_inputs_model(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("log inputs test")
|
|
run = mlflow_client.create_run(experiment_id)
|
|
model = mlflow_client.create_logged_model(experiment_id=experiment_id)
|
|
dataset = Dataset(
|
|
name="name1",
|
|
digest="digest1",
|
|
source_type="source_type1",
|
|
source="source1",
|
|
)
|
|
dataset_inputs = [
|
|
DatasetInput(
|
|
dataset=dataset,
|
|
tags=[InputTag(key=MLFLOW_DATASET_CONTEXT, value="training")],
|
|
)
|
|
]
|
|
mlflow_client.log_inputs(
|
|
run.info.run_id,
|
|
models=[LoggedModelInput(model_id=model.model_id)],
|
|
datasets=dataset_inputs,
|
|
)
|
|
run = mlflow_client.get_run(run.info.run_id)
|
|
assert len(run.inputs.model_inputs) == 1
|
|
|
|
|
|
def test_update_run_name_without_changing_status(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("update run name")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
mlflow_client.set_terminated(created_run.info.run_id, "FINISHED")
|
|
|
|
mlflow_client.update_run(created_run.info.run_id, name="name_abc")
|
|
updated_run_info = mlflow_client.get_run(created_run.info.run_id).info
|
|
assert updated_run_info.run_name == "name_abc"
|
|
assert updated_run_info.status == "FINISHED"
|
|
|
|
|
|
def test_create_promptlab_run_handler_rejects_invalid_requests(mlflow_client):
|
|
def assert_response(resp, message_part):
|
|
assert resp.status_code == 400
|
|
response_json = resp.json()
|
|
assert response_json.get("error_code") == "INVALID_PARAMETER_VALUE"
|
|
assert message_part in response_json.get("message", "")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/runs/create-promptlab-run",
|
|
json={},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"CreatePromptlabRun request must specify experiment_id.",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/runs/create-promptlab-run",
|
|
json={"experiment_id": "123"},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"CreatePromptlabRun request must specify prompt_template.",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/runs/create-promptlab-run",
|
|
json={"experiment_id": "123", "prompt_template": "my_prompt_template"},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"CreatePromptlabRun request must specify prompt_parameters.",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/runs/create-promptlab-run",
|
|
json={
|
|
"experiment_id": "123",
|
|
"prompt_template": "my_prompt_template",
|
|
"prompt_parameters": [{"key": "my_key", "value": "my_value"}],
|
|
},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"CreatePromptlabRun request must specify model_route.",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/runs/create-promptlab-run",
|
|
json={
|
|
"experiment_id": "123",
|
|
"prompt_template": "my_prompt_template",
|
|
"prompt_parameters": [{"key": "my_key", "value": "my_value"}],
|
|
"model_route": "my_route",
|
|
},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"CreatePromptlabRun request must specify model_input.",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/runs/create-promptlab-run",
|
|
json={
|
|
"experiment_id": "123",
|
|
"prompt_template": "my_prompt_template",
|
|
"prompt_parameters": [{"key": "my_key", "value": "my_value"}],
|
|
"model_route": "my_route",
|
|
"model_input": "my_input",
|
|
},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"CreatePromptlabRun request must specify mlflow_version.",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/runs/create-promptlab-run",
|
|
json={
|
|
"experiment_id": "123",
|
|
"prompt_template": "my_prompt_template",
|
|
"prompt_parameters": [{"key": "my_key", "value": "my_value"}],
|
|
"model_route": "my_route",
|
|
"model_input": "my_input",
|
|
"mlflow_version": "1.0.0",
|
|
},
|
|
)
|
|
|
|
|
|
def test_create_promptlab_run_handler_returns_expected_results(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("log inputs test")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/runs/create-promptlab-run",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"run_name": "my_run_name",
|
|
"prompt_template": "my_prompt_template",
|
|
"prompt_parameters": [{"key": "my_key", "value": "my_value"}],
|
|
"model_route": "my_route",
|
|
"model_parameters": [{"key": "temperature", "value": "0.1"}],
|
|
"model_input": "my_input",
|
|
"model_output": "my_output",
|
|
"model_output_parameters": [{"key": "latency", "value": "100"}],
|
|
"mlflow_version": "1.0.0",
|
|
"user_id": "username",
|
|
"start_time": 456,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
run_json = response.json()
|
|
assert run_json["run"]["info"]["run_name"] == "my_run_name"
|
|
assert run_json["run"]["info"]["experiment_id"] == experiment_id
|
|
assert run_json["run"]["info"]["user_id"] == "username"
|
|
assert run_json["run"]["info"]["status"] == "FINISHED"
|
|
assert run_json["run"]["info"]["start_time"] == 456
|
|
|
|
assert {"key": "model_route", "value": "my_route"} in run_json["run"]["data"]["params"]
|
|
assert {"key": "prompt_template", "value": "my_prompt_template"} in run_json["run"]["data"][
|
|
"params"
|
|
]
|
|
assert {"key": "temperature", "value": "0.1"} in run_json["run"]["data"]["params"]
|
|
|
|
assert {
|
|
"key": "mlflow.loggedArtifacts",
|
|
"value": '[{"path": "eval_results_table.json", "type": "table"}]',
|
|
} in run_json["run"]["data"]["tags"]
|
|
assert {"key": "mlflow.runSourceType", "value": "PROMPT_ENGINEERING"} in run_json["run"][
|
|
"data"
|
|
]["tags"]
|
|
|
|
|
|
def test_gateway_proxy_handler_rejects_invalid_requests(mlflow_client):
|
|
def assert_response(resp, message_part):
|
|
assert resp.status_code == 400
|
|
response_json = resp.json()
|
|
assert response_json.get("error_code") == "INVALID_PARAMETER_VALUE"
|
|
assert message_part in response_json.get("message", "")
|
|
|
|
with _init_server(
|
|
backend_uri=mlflow_client.tracking_uri,
|
|
root_artifact_uri=mlflow_client.tracking_uri,
|
|
extra_env={"MLFLOW_DEPLOYMENTS_TARGET": "http://localhost:5001"},
|
|
server_type="flask",
|
|
) as url:
|
|
patched_client = MlflowClient(url)
|
|
|
|
response = requests.post(
|
|
f"{patched_client.tracking_uri}/ajax-api/2.0/mlflow/gateway-proxy",
|
|
json={},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"Deployments proxy request must specify a gateway_path.",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{patched_client.tracking_uri}/ajax-api/2.0/mlflow/gateway-proxy",
|
|
json={"gateway_path": "foo/bar"},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"Invalid gateway_path: foo/bar for method: POST",
|
|
)
|
|
|
|
response = requests.post(
|
|
f"{patched_client.tracking_uri}/ajax-api/2.0/mlflow/gateway-proxy",
|
|
json={"gateway_path": "foo/bar/baz"},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"Invalid gateway_path: foo/bar/baz for method: POST",
|
|
)
|
|
|
|
response = requests.get(
|
|
f"{patched_client.tracking_uri}/ajax-api/2.0/mlflow/gateway-proxy",
|
|
params={"gateway_path": "hello/world"},
|
|
)
|
|
assert_response(
|
|
response,
|
|
"Invalid gateway_path: hello/world for method: GET",
|
|
)
|
|
|
|
# Unsupported method
|
|
response = requests.delete(
|
|
f"{patched_client.tracking_uri}/ajax-api/2.0/mlflow/gateway-proxy",
|
|
)
|
|
assert response.status_code == 405
|
|
|
|
|
|
def test_upload_artifact_handler_rejects_invalid_requests(mlflow_client):
|
|
def assert_response(resp, message_part):
|
|
assert resp.status_code == 400
|
|
response_json = resp.json()
|
|
assert response_json.get("error_code") == "INVALID_PARAMETER_VALUE"
|
|
assert message_part in response_json.get("message", "")
|
|
|
|
experiment_id = mlflow_client.create_experiment("upload_artifacts_test")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/upload-artifact", params={}
|
|
)
|
|
assert_response(response, "Request must specify run_uuid.")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/upload-artifact",
|
|
params={
|
|
"run_uuid": created_run.info.run_id,
|
|
},
|
|
)
|
|
assert_response(response, "Request must specify path.")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/upload-artifact",
|
|
params={"run_uuid": created_run.info.run_id, "path": ""},
|
|
)
|
|
assert_response(response, "Request must specify path.")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/upload-artifact",
|
|
params={"run_uuid": created_run.info.run_id, "path": "../test.txt"},
|
|
)
|
|
assert_response(response, "Invalid path")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/upload-artifact",
|
|
params={
|
|
"run_uuid": created_run.info.run_id,
|
|
"path": "test.txt",
|
|
},
|
|
)
|
|
assert_response(response, "Request must specify data.")
|
|
|
|
large_data = b"x" * (10 * 1024 * 1024 + 1)
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/upload-artifact",
|
|
params={
|
|
"run_uuid": created_run.info.run_id,
|
|
"path": "test.txt",
|
|
},
|
|
data=large_data,
|
|
)
|
|
assert_response(response, "Artifact size is too large")
|
|
|
|
|
|
def test_upload_artifact_handler(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("upload_artifacts_test")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/upload-artifact",
|
|
params={
|
|
"run_uuid": created_run.info.run_id,
|
|
"path": "test.txt",
|
|
},
|
|
data="hello world",
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/get-artifact",
|
|
params={
|
|
"run_uuid": created_run.info.run_id,
|
|
"path": "test.txt",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.text == "hello world"
|
|
|
|
|
|
def test_graphql_handler(mlflow_client):
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": 'query testQuery {test(inputString: "abc") { output }}',
|
|
"operationName": "testQuery",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
assert response.status_code == 200
|
|
|
|
|
|
def test_graphql_handler_batching_raise_error(mlflow_client):
|
|
# Test max root fields limit
|
|
batch_query = (
|
|
"query testQuery {"
|
|
+ " ".join([
|
|
f"key_{i}: " + 'test(inputString: "abc") { output }'
|
|
for i in range(int(MLFLOW_SERVER_GRAPHQL_MAX_ROOT_FIELDS.get()) + 2)
|
|
])
|
|
+ "}"
|
|
)
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": batch_query,
|
|
"operationName": "testQuery",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
assert response.status_code == 200
|
|
assert (
|
|
f"GraphQL queries should have at most {MLFLOW_SERVER_GRAPHQL_MAX_ROOT_FIELDS.get()}"
|
|
in response.json()["errors"][0]
|
|
)
|
|
|
|
# Test max aliases limit
|
|
batch_query = (
|
|
'query testQuery {mlflowGetExperiment(input: {experimentId: "123"}) {'
|
|
+ " ".join(
|
|
f"experiment_{i}: " + "experiment { name }"
|
|
for i in range(int(MLFLOW_SERVER_GRAPHQL_MAX_ALIASES.get()) + 2)
|
|
)
|
|
+ "}}"
|
|
)
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": batch_query,
|
|
"operationName": "testQuery",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
assert (
|
|
f"queries should have at most {MLFLOW_SERVER_GRAPHQL_MAX_ALIASES.get()} aliases"
|
|
in response.json()["errors"][0]
|
|
)
|
|
|
|
# Test max depth limit
|
|
inner = "name"
|
|
for _ in range(12):
|
|
inner = f"name {{ {inner} }}"
|
|
deep_query = (
|
|
'query testQuery { mlflowGetExperiment(input: {experimentId: "123"}) { experiment { '
|
|
+ inner
|
|
+ " } } }"
|
|
)
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": deep_query,
|
|
"operationName": "testQuery",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
assert "Query exceeds maximum depth of 10" in response.json()["errors"][0]
|
|
|
|
# Test max selections limit
|
|
# Exceed the 1000 selection limit
|
|
selections = [f"field_{i} {{ name }}" for i in range(1002)]
|
|
selections_query = (
|
|
'query testQuery { mlflowGetExperiment(input: {experimentId: "123"}) { experiment { '
|
|
+ " ".join(selections)
|
|
+ " } } }"
|
|
)
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": selections_query,
|
|
"operationName": "testQuery",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
assert "Query exceeds maximum total selections of 1000" in response.json()["errors"][0]
|
|
|
|
|
|
def test_get_experiment_graphql(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("GraphqlTest")
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": 'query testQuery {mlflowGetExperiment(input: {experimentId: "'
|
|
+ experiment_id
|
|
+ '"}) { experiment { name effectiveTraceArchivalRetention } }}',
|
|
"operationName": "testQuery",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
assert response.status_code == 200
|
|
json = response.json()
|
|
assert json["data"]["mlflowGetExperiment"]["experiment"]["name"] == "GraphqlTest"
|
|
assert json["data"]["mlflowGetExperiment"]["experiment"]["effectiveTraceArchivalRetention"] in (
|
|
None,
|
|
"",
|
|
)
|
|
|
|
|
|
def test_get_run_and_experiment_graphql(mlflow_client):
|
|
name = "GraphqlTest"
|
|
mlflow_client.create_registered_model(name)
|
|
experiment_id = mlflow_client.create_experiment(name)
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
mlflow_client.create_model_version("GraphqlTest", "runs:/graphql_test/model", run_id)
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": f"""
|
|
query testQuery @component(name: "Test") {{
|
|
mlflowGetRun(input: {{runId: "{run_id}"}}) {{
|
|
run {{
|
|
info {{
|
|
status
|
|
}}
|
|
experiment {{
|
|
name
|
|
}}
|
|
modelVersions {{
|
|
name
|
|
}}
|
|
}}
|
|
}}
|
|
}}
|
|
""",
|
|
"operationName": "testQuery",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
assert response.status_code == 200
|
|
json = response.json()
|
|
assert json["errors"] is None
|
|
assert json["data"]["mlflowGetRun"]["run"]["info"]["status"] == created_run.info.status
|
|
assert json["data"]["mlflowGetRun"]["run"]["experiment"]["name"] == name
|
|
assert json["data"]["mlflowGetRun"]["run"]["modelVersions"][0]["name"] == name
|
|
|
|
|
|
def test_legacy_start_and_end_trace_v2(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("start end trace")
|
|
|
|
# Trace CRUD APIs are not directly exposed as public API of MlflowClient,
|
|
# so we use the underlying tracking client to test them.
|
|
store = mlflow_client._tracing_client.store
|
|
|
|
# Helper function to remove auto-added system tags (mlflow.xxx) from testing
|
|
def _exclude_system_tags(tags: dict[str, str]):
|
|
return {k: v for k, v in tags.items() if not k.startswith("mlflow.")}
|
|
|
|
trace_info = store.deprecated_start_trace_v2(
|
|
experiment_id=experiment_id,
|
|
timestamp_ms=1000,
|
|
request_metadata={
|
|
"meta1": "apple",
|
|
"meta2": "grape",
|
|
},
|
|
tags={
|
|
"tag1": "football",
|
|
"tag2": "basketball",
|
|
},
|
|
)
|
|
assert trace_info.request_id is not None
|
|
assert trace_info.experiment_id == experiment_id
|
|
assert trace_info.timestamp_ms == 1000
|
|
assert trace_info.execution_time_ms == 0
|
|
assert trace_info.status == TraceStatus.IN_PROGRESS
|
|
assert trace_info.request_metadata == {
|
|
"meta1": "apple",
|
|
"meta2": "grape",
|
|
}
|
|
assert _exclude_system_tags(trace_info.tags) == {
|
|
"tag1": "football",
|
|
"tag2": "basketball",
|
|
}
|
|
|
|
trace_info = store.deprecated_end_trace_v2(
|
|
request_id=trace_info.request_id,
|
|
timestamp_ms=3000,
|
|
status=TraceStatus.OK,
|
|
request_metadata={
|
|
"meta1": "orange",
|
|
"meta3": "banana",
|
|
},
|
|
tags={
|
|
"tag1": "soccer",
|
|
"tag3": "tennis",
|
|
},
|
|
)
|
|
assert trace_info.request_id is not None
|
|
assert trace_info.experiment_id == experiment_id
|
|
assert trace_info.timestamp_ms == 1000
|
|
assert trace_info.execution_time_ms == 2000
|
|
assert trace_info.status == TraceStatus.OK
|
|
assert trace_info.request_metadata == {
|
|
"meta1": "orange",
|
|
"meta2": "grape",
|
|
"meta3": "banana",
|
|
}
|
|
assert _exclude_system_tags(trace_info.tags) == {
|
|
"tag1": "soccer",
|
|
"tag2": "basketball",
|
|
"tag3": "tennis",
|
|
}
|
|
|
|
|
|
def test_start_trace(mlflow_client):
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow.set_experiment("start end trace").experiment_id
|
|
|
|
# Helper function to remove auto-added system tags (mlflow.xxx) from testing
|
|
def _exclude_system_keys(d: dict[str, str]):
|
|
return {k: v for k, v in d.items() if not k.startswith("mlflow.")}
|
|
|
|
with mock.patch("mlflow.tracing.export.mlflow_v3._logger.warning") as mock_warning:
|
|
with mlflow.start_span(name="test") as span:
|
|
mlflow.update_current_trace(
|
|
tags={
|
|
"tag1": "football",
|
|
"tag2": "basketball",
|
|
},
|
|
metadata={
|
|
"meta1": "apple",
|
|
"meta2": "grape",
|
|
},
|
|
)
|
|
|
|
trace = mlflow_client.get_trace(span.trace_id, flush=True)
|
|
assert trace.info.trace_id == span.trace_id
|
|
assert trace.info.experiment_id == experiment_id
|
|
assert trace.info.request_time > 0
|
|
assert trace.info.execution_duration is not None
|
|
assert trace.info.state == TraceState.OK
|
|
assert _exclude_system_keys(trace.info.trace_metadata) == {
|
|
"meta1": "apple",
|
|
"meta2": "grape",
|
|
}
|
|
assert trace.info.trace_metadata[TRACE_SCHEMA_VERSION_KEY] == "3"
|
|
assert _exclude_system_keys(trace.info.tags) == {
|
|
"tag1": "football",
|
|
"tag2": "basketball",
|
|
}
|
|
|
|
# No "Failed to log span to MLflow backend" warning should be issued
|
|
for call in mock_warning.call_args_list:
|
|
assert "Failed to log span to MLflow backend" not in str(call)
|
|
|
|
|
|
def test_get_trace(mlflow_client):
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow_client.create_experiment("get trace")
|
|
span = mlflow_client.start_trace(name="test", experiment_id=experiment_id)
|
|
mlflow_client.end_trace(request_id=span.request_id, status=TraceStatus.OK)
|
|
trace = mlflow_client.get_trace(span.request_id, flush=True)
|
|
assert trace is not None
|
|
assert trace.info.request_id == span.request_id
|
|
assert trace.info.experiment_id == experiment_id
|
|
assert trace.info.state == TraceState.OK
|
|
assert len(trace.data.spans) == 1
|
|
assert trace.data.spans[0].name == "test"
|
|
assert trace.data.spans[0].status.status_code == SpanStatusCode.OK
|
|
assert trace.data.spans[0].status.description == ""
|
|
|
|
|
|
def test_search_traces(mlflow_client):
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow_client.create_experiment("search traces")
|
|
|
|
# Create test traces
|
|
def _create_trace(name, status):
|
|
span = mlflow_client.start_trace(name=name, experiment_id=experiment_id)
|
|
mlflow_client.end_trace(request_id=span.request_id, status=status)
|
|
return span.request_id
|
|
|
|
# Flush between creations to ensure distinct timestamps. Without this, all three traces
|
|
# can land in the same millisecond on a fast local server, making max_results ordering
|
|
# non-deterministic.
|
|
request_id_1 = _create_trace(name="trace1", status=TraceStatus.OK)
|
|
mlflow.flush_trace_async_logging()
|
|
request_id_2 = _create_trace(name="trace2", status=TraceStatus.OK)
|
|
mlflow.flush_trace_async_logging()
|
|
request_id_3 = _create_trace(name="trace3", status=TraceStatus.ERROR)
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
def _get_request_ids(traces):
|
|
return [t.info.request_id for t in traces]
|
|
|
|
# Validate search
|
|
traces = mlflow_client.search_traces(locations=[experiment_id])
|
|
assert set(_get_request_ids(traces)) == {request_id_3, request_id_2, request_id_1}
|
|
assert traces.token is None
|
|
|
|
traces = mlflow_client.search_traces(
|
|
locations=[experiment_id],
|
|
filter_string="status = 'OK'",
|
|
order_by=["timestamp ASC"],
|
|
)
|
|
assert set(_get_request_ids(traces)) == {request_id_1, request_id_2}
|
|
assert traces.token is None
|
|
|
|
traces = mlflow_client.search_traces(
|
|
locations=[experiment_id],
|
|
max_results=2,
|
|
)
|
|
assert set(_get_request_ids(traces)) == {request_id_3, request_id_2}
|
|
assert traces.token is not None
|
|
traces = mlflow_client.search_traces(
|
|
locations=[experiment_id],
|
|
page_token=traces.token,
|
|
)
|
|
assert _get_request_ids(traces) == [request_id_1]
|
|
assert traces.token is None
|
|
|
|
|
|
def test_search_traces_parameter_validation(mlflow_client):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="Locations must be a list of experiment IDs",
|
|
):
|
|
mlflow_client.search_traces(locations=["catalog.schema"])
|
|
|
|
|
|
def test_search_traces_match_text(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support full text search")
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow_client.create_experiment("search traces full text")
|
|
|
|
# Create test traces
|
|
def _create_trace(name, attributes):
|
|
span = mlflow_client.start_trace(name=name, experiment_id=experiment_id)
|
|
span.set_attributes(attributes)
|
|
mlflow_client.end_trace(request_id=span.trace_id, status=TraceStatus.OK)
|
|
return span.trace_id
|
|
|
|
trace_id_1 = _create_trace(name="trace1", attributes={"test": "value1"})
|
|
trace_id_2 = _create_trace(name="trace2", attributes={"test": "value2"})
|
|
trace_id_3 = _create_trace(name="trace3", attributes={"test3": "I like it"})
|
|
|
|
traces = mlflow_client.search_traces(locations=[experiment_id], flush=True)
|
|
assert len([t.info.trace_id for t in traces]) == 3
|
|
assert traces.token is None
|
|
|
|
traces = mlflow_client.search_traces(
|
|
locations=[experiment_id], filter_string="trace.text LIKE '%trace%'"
|
|
)
|
|
assert len([t.info.trace_id for t in traces]) == 3
|
|
assert traces.token is None
|
|
|
|
traces = mlflow_client.search_traces(
|
|
locations=[experiment_id], filter_string="trace.text LIKE '%value%'"
|
|
)
|
|
assert {t.info.trace_id for t in traces} == {trace_id_1, trace_id_2}
|
|
|
|
traces = mlflow_client.search_traces(
|
|
locations=[experiment_id], filter_string="trace.text LIKE '%I like it%'"
|
|
)
|
|
assert [t.info.trace_id for t in traces] == [trace_id_3]
|
|
|
|
|
|
def test_delete_traces(mlflow_client):
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow_client.create_experiment("delete traces")
|
|
|
|
def _create_trace(name, status):
|
|
span = mlflow_client.start_trace(name=name, experiment_id=experiment_id)
|
|
mlflow_client.end_trace(request_id=span.request_id, status=status)
|
|
return span.request_id
|
|
|
|
def _is_trace_exists(request_id):
|
|
try:
|
|
trace_info = mlflow_client._tracing_client.get_trace_info(request_id)
|
|
return trace_info is not None
|
|
except RestException as e:
|
|
if e.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST):
|
|
return False
|
|
raise
|
|
|
|
# Case 1: Delete all traces under experiment ID
|
|
request_id_1 = _create_trace(name="trace1", status=TraceStatus.OK)
|
|
request_id_2 = _create_trace(name="trace2", status=TraceStatus.OK)
|
|
mlflow.flush_trace_async_logging()
|
|
assert _is_trace_exists(request_id_1)
|
|
assert _is_trace_exists(request_id_2)
|
|
|
|
deleted_count = mlflow_client.delete_traces(experiment_id, max_timestamp_millis=int(1e15))
|
|
assert deleted_count == 2
|
|
assert not _is_trace_exists(request_id_1)
|
|
assert not _is_trace_exists(request_id_2)
|
|
|
|
# Case 2: Delete with max_traces limit
|
|
request_id_1 = _create_trace(name="trace1", status=TraceStatus.OK)
|
|
time.sleep(0.1) # Add some time gap to avoid timestamp collision
|
|
request_id_2 = _create_trace(name="trace2", status=TraceStatus.OK)
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
deleted_count = mlflow_client.delete_traces(
|
|
experiment_id, max_traces=1, max_timestamp_millis=int(1e15)
|
|
)
|
|
assert deleted_count == 1
|
|
# TODO: Currently the deletion order in the file store is random (based on
|
|
# the order of the trace files in the directory), so we don't validate which
|
|
# one is deleted. Uncomment the following lines once the deletion order is fixed.
|
|
# assert not _is_trace_exists(request_id_1) # Old created trace should be deleted
|
|
# assert _is_trace_exists(request_id_2)
|
|
|
|
# Case 3: Delete with explicit request ID
|
|
request_id_1 = _create_trace(name="trace1", status=TraceStatus.OK)
|
|
request_id_2 = _create_trace(name="trace2", status=TraceStatus.OK)
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
deleted_count = mlflow_client.delete_traces(experiment_id, trace_ids=[request_id_1])
|
|
assert deleted_count == 1
|
|
assert not _is_trace_exists(request_id_1)
|
|
assert _is_trace_exists(request_id_2)
|
|
|
|
|
|
def test_calculate_trace_filter_correlation(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support calculate_trace_filter_correlation")
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow_client.create_experiment("correlation test")
|
|
|
|
def _create_trace(name, tags):
|
|
span = mlflow_client.start_trace(name=name, experiment_id=experiment_id, tags=tags)
|
|
mlflow_client.end_trace(request_id=span.request_id, status=TraceStatus.OK)
|
|
return span.request_id
|
|
|
|
for i in range(6):
|
|
_create_trace(f"trace-prod-tool-{i}", {"env": "prod", "span_type": "TOOL"})
|
|
|
|
for i in range(4):
|
|
_create_trace(f"trace-dev-{i}", {"env": "dev", "span_type": "LLM" if i >= 1 else "TOOL"})
|
|
|
|
client = TracingClient(tracking_uri=mlflow_client.tracking_uri)
|
|
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
result = client.calculate_trace_filter_correlation(
|
|
experiment_ids=[experiment_id],
|
|
filter_string1="tags.env = 'prod'",
|
|
filter_string2="tags.span_type = 'TOOL'",
|
|
)
|
|
|
|
assert isinstance(result, TraceFilterCorrelationResult)
|
|
assert result.total_count == 10
|
|
assert result.filter1_count == 6
|
|
assert result.filter2_count == 7
|
|
assert result.joint_count == 6
|
|
assert 0.6 < result.npmi < 0.8
|
|
assert result.npmi_smoothed is not None
|
|
|
|
result2 = client.calculate_trace_filter_correlation(
|
|
experiment_ids=[experiment_id],
|
|
filter_string1="tags.env = 'dev'",
|
|
filter_string2="tags.span_type = 'LLM'",
|
|
)
|
|
|
|
assert result2.total_count == 10
|
|
assert result2.filter1_count == 4
|
|
assert result2.filter2_count == 3
|
|
assert result2.joint_count == 3
|
|
assert result2.npmi > 0.5
|
|
|
|
result3 = client.calculate_trace_filter_correlation(
|
|
experiment_ids=[experiment_id],
|
|
filter_string1="tags.env = 'staging'",
|
|
filter_string2="tags.span_type = 'TOOL'",
|
|
)
|
|
|
|
assert result3.total_count == 10
|
|
assert result3.filter1_count == 0
|
|
assert result3.filter2_count == 7
|
|
assert result3.joint_count == 0
|
|
assert math.isnan(result3.npmi)
|
|
|
|
with pytest.raises(MlflowException, match="Invalid"):
|
|
client.calculate_trace_filter_correlation(
|
|
experiment_ids=[experiment_id],
|
|
filter_string1="invalid.filter = 'test'",
|
|
filter_string2="tags.span_type = 'TOOL'",
|
|
)
|
|
|
|
|
|
def test_set_and_delete_trace_tag(mlflow_client):
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow_client.create_experiment("set delete tag")
|
|
|
|
# Create test trace
|
|
trace_info = mlflow_client._tracing_client.start_trace(
|
|
TraceInfo(
|
|
trace_id="tr-1234",
|
|
trace_location=TraceLocation.from_experiment_id(experiment_id),
|
|
request_time=1000,
|
|
execution_duration=2000,
|
|
state=TraceState.OK,
|
|
tags={
|
|
"tag1": "red",
|
|
"tag2": "blue",
|
|
},
|
|
)
|
|
)
|
|
|
|
# Validate set tag
|
|
mlflow_client.set_trace_tag(trace_info.request_id, "tag1", "green")
|
|
trace_info = mlflow_client._tracing_client.get_trace_info(trace_info.request_id)
|
|
assert trace_info.tags["tag1"] == "green"
|
|
|
|
# Validate delete tag
|
|
mlflow_client.delete_trace_tag(trace_info.request_id, "tag2")
|
|
trace_info = mlflow_client._tracing_client.get_trace_info(trace_info.request_id)
|
|
assert "tag2" not in trace_info.tags
|
|
|
|
|
|
def test_query_trace_metrics(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support query trace metrics")
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow_client.create_experiment("query trace metrics")
|
|
|
|
# Create test traces
|
|
def _create_trace(name, status):
|
|
span = mlflow_client.start_trace(name=name, experiment_id=experiment_id)
|
|
mlflow_client.end_trace(request_id=span.request_id, status=status)
|
|
return span.request_id
|
|
|
|
_create_trace(name="trace1", status=TraceStatus.OK)
|
|
_create_trace(name="trace2", status=TraceStatus.OK)
|
|
_create_trace(name="trace3", status=TraceStatus.ERROR)
|
|
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
metrics = mlflow_client._tracing_client.store.query_trace_metrics(
|
|
experiment_ids=[experiment_id],
|
|
view_type=MetricViewType.TRACES,
|
|
metric_name=TraceMetricKey.TRACE_COUNT,
|
|
aggregations=[MetricAggregation(aggregation_type=AggregationType.COUNT)],
|
|
dimensions=[TraceMetricDimensionKey.TRACE_STATUS],
|
|
)
|
|
assert len(metrics) == 2
|
|
assert asdict(metrics[0]) == {
|
|
"metric_name": TraceMetricKey.TRACE_COUNT,
|
|
"dimensions": {TraceMetricDimensionKey.TRACE_STATUS: "ERROR"},
|
|
"values": {"COUNT": 1},
|
|
}
|
|
|
|
assert asdict(metrics[1]) == {
|
|
"metric_name": TraceMetricKey.TRACE_COUNT,
|
|
"dimensions": {TraceMetricDimensionKey.TRACE_STATUS: "OK"},
|
|
"values": {"COUNT": 2},
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("allow_partial", [True, False])
|
|
def test_get_trace_handler(mlflow_client, allow_partial: bool, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support get trace handler")
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
|
|
with mlflow.start_span(name="test") as span:
|
|
span.set_attributes({"fruit": "apple"})
|
|
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/3.0/mlflow/traces/get",
|
|
params={"trace_id": span.trace_id, "allow_partial": allow_partial},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
|
|
trace = response.json()["trace"]
|
|
assert trace["trace_info"]["trace_id"] == span.trace_id
|
|
assert len(trace["spans"]) == 1
|
|
assert trace["spans"][0]["name"] == "test"
|
|
attributes = trace["spans"][0]["attributes"]
|
|
assert {"key": "fruit", "value": {"string_value": "apple"}} in attributes
|
|
|
|
|
|
def test_get_trace_artifact_handler(mlflow_client):
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
|
|
with mlflow.start_span(name="test") as span:
|
|
span.set_attributes({"fruit": "apple"})
|
|
span.add_event(SpanEvent("test_event", timestamp=99999, attributes={"foo": "bar"}))
|
|
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/get-trace-artifact",
|
|
params={"request_id": span.trace_id},
|
|
)
|
|
assert response.status_code == 200
|
|
assert response.headers["Content-Disposition"] == "attachment; filename=traces.json"
|
|
|
|
# Validate content
|
|
trace_data = TraceData.from_dict(json.loads(response.text))
|
|
assert trace_data.spans[0].to_dict() == span.to_dict()
|
|
|
|
|
|
def test_link_traces_to_run_and_search_traces(mlflow_client, store_type):
|
|
# Skip file store because it doesn't support linking traces to runs
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support linking traces to runs")
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow.set_experiment("link traces to run test").experiment_id
|
|
|
|
run = mlflow_client.create_run(experiment_id)
|
|
run_id = run.info.run_id
|
|
|
|
# 1. Trace created under a run
|
|
with mlflow.start_run(run_id=run_id):
|
|
with mlflow.start_span(name="trace1") as span1:
|
|
span1.set_attributes({"test": "value1"})
|
|
trace_id_1 = span1.trace_id
|
|
|
|
# 2. Trace associated with a run
|
|
with mlflow.start_span(name="trace2") as span2:
|
|
span2.set_attributes({"test": "value2"})
|
|
trace_id_2 = span2.trace_id
|
|
mlflow_client.link_traces_to_run(trace_ids=[trace_id_2], run_id=run_id)
|
|
|
|
# 3. Trace not associated with a run
|
|
with mlflow.start_span(name="trace3") as span3:
|
|
span3.set_attributes({"test": "value3"})
|
|
trace_id_3 = span3.trace_id
|
|
|
|
# Search traces without run_id filter - should return all traces in experiment
|
|
all_traces = mlflow_client.search_traces(locations=[experiment_id], flush=True)
|
|
assert {t.info.trace_id for t in all_traces} == {trace_id_1, trace_id_2, trace_id_3}
|
|
|
|
# Search traces with run_id filter - should return only linked traces
|
|
linked_traces = mlflow_client.search_traces(
|
|
locations=[experiment_id], filter_string=f"attribute.run_id = '{run_id}'"
|
|
)
|
|
linked_trace_ids = [t.info.trace_id for t in linked_traces]
|
|
assert len(linked_trace_ids) == 2
|
|
assert set(linked_trace_ids) == {trace_id_1, trace_id_2}
|
|
|
|
|
|
def test_get_metric_history_bulk_interval_graphql(mlflow_client):
|
|
name = "GraphqlTest"
|
|
mlflow_client.create_registered_model(name)
|
|
experiment_id = mlflow_client.create_experiment(name)
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
|
|
metric_name = "metric_0"
|
|
for i in range(10):
|
|
mlflow_client.log_metric(created_run.info.run_id, metric_name, i, step=i)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": f"""
|
|
query testQuery {{
|
|
mlflowGetMetricHistoryBulkInterval(input: {{
|
|
runIds: ["{created_run.info.run_id}"],
|
|
metricKey: "{metric_name}",
|
|
}}) {{
|
|
metrics {{
|
|
key
|
|
timestamp
|
|
value
|
|
}}
|
|
}}
|
|
}}
|
|
""",
|
|
"operationName": "testQuery",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
json = response.json()
|
|
expected = [{"key": metric_name, "timestamp": mock.ANY, "value": i} for i in range(10)]
|
|
assert json["data"]["mlflowGetMetricHistoryBulkInterval"]["metrics"] == expected
|
|
|
|
|
|
def test_search_runs_graphql(mlflow_client):
|
|
name = "GraphqlTest"
|
|
mlflow_client.create_registered_model(name)
|
|
experiment_id = mlflow_client.create_experiment(name)
|
|
created_run_1 = mlflow_client.create_run(experiment_id)
|
|
created_run_2 = mlflow_client.create_run(experiment_id)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": f"""
|
|
mutation testMutation {{
|
|
mlflowSearchRuns(input: {{ experimentIds: ["{experiment_id}"] }}) {{
|
|
runs {{
|
|
info {{
|
|
runId
|
|
}}
|
|
}}
|
|
}}
|
|
}}
|
|
""",
|
|
"operationName": "testMutation",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
json = response.json()
|
|
expected = [
|
|
{"info": {"runId": created_run_2.info.run_id}},
|
|
{"info": {"runId": created_run_1.info.run_id}},
|
|
]
|
|
assert json["data"]["mlflowSearchRuns"]["runs"] == expected
|
|
|
|
|
|
def test_list_artifacts_graphql(mlflow_client, tmp_path):
|
|
name = "GraphqlTest"
|
|
experiment_id = mlflow_client.create_experiment(name)
|
|
created_run_id = mlflow_client.create_run(experiment_id).info.run_id
|
|
file_path = tmp_path / "test.txt"
|
|
file_path.write_text("hello world")
|
|
mlflow_client.log_artifact(created_run_id, file_path.absolute().as_posix())
|
|
mlflow_client.log_artifact(created_run_id, file_path.absolute().as_posix(), "testDir")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": f"""
|
|
query testQuery {{
|
|
files: mlflowListArtifacts(input: {{
|
|
runId: "{created_run_id}",
|
|
}}) {{
|
|
files {{
|
|
path
|
|
isDir
|
|
fileSize
|
|
}}
|
|
}}
|
|
}}
|
|
""",
|
|
"operationName": "testQuery",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
json = response.json()
|
|
file_expected = [
|
|
{"path": "test.txt", "isDir": False, "fileSize": "11"},
|
|
{"path": "testDir", "isDir": True, "fileSize": "0"},
|
|
]
|
|
assert json["data"]["files"]["files"] == file_expected
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": f"""
|
|
query testQuery {{
|
|
subdir: mlflowListArtifacts(input: {{
|
|
runId: "{created_run_id}",
|
|
path: "testDir",
|
|
}}) {{
|
|
files {{
|
|
path
|
|
isDir
|
|
fileSize
|
|
}}
|
|
}}
|
|
}}
|
|
""",
|
|
"operationName": "testQuery",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
json = response.json()
|
|
subdir_expected = [
|
|
{"path": "testDir/test.txt", "isDir": False, "fileSize": "11"},
|
|
]
|
|
assert json["data"]["subdir"]["files"] == subdir_expected
|
|
|
|
|
|
def test_search_datasets_graphql(mlflow_client):
|
|
name = "GraphqlTest"
|
|
experiment_id = mlflow_client.create_experiment(name)
|
|
created_run_id = mlflow_client.create_run(experiment_id).info.run_id
|
|
dataset1 = Dataset(
|
|
name="test-dataset-1",
|
|
digest="12345",
|
|
source_type="script",
|
|
source="test",
|
|
)
|
|
dataset_input1 = DatasetInput(dataset=dataset1, tags=[])
|
|
dataset2 = Dataset(
|
|
name="test-dataset-2",
|
|
digest="12346",
|
|
source_type="script",
|
|
source="test",
|
|
)
|
|
dataset_input2 = DatasetInput(
|
|
dataset=dataset2, tags=[InputTag(key=MLFLOW_DATASET_CONTEXT, value="training")]
|
|
)
|
|
mlflow_client.log_inputs(created_run_id, [dataset_input1, dataset_input2])
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": f"""
|
|
mutation testMutation {{
|
|
mlflowSearchDatasets(input:{{experimentIds: ["{experiment_id}"]}}) {{
|
|
datasetSummaries {{
|
|
experimentId
|
|
name
|
|
digest
|
|
context
|
|
}}
|
|
}}
|
|
}}
|
|
""",
|
|
"operationName": "testMutation",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
json = response.json()
|
|
|
|
def sort_dataset_summaries(l1):
|
|
return sorted(l1, key=lambda x: x["digest"])
|
|
|
|
expected = sort_dataset_summaries([
|
|
{
|
|
"experimentId": experiment_id,
|
|
"name": "test-dataset-2",
|
|
"digest": "12346",
|
|
"context": "training",
|
|
},
|
|
{
|
|
"experimentId": experiment_id,
|
|
"name": "test-dataset-1",
|
|
"digest": "12345",
|
|
"context": "",
|
|
},
|
|
])
|
|
assert (
|
|
sort_dataset_summaries(json["data"]["mlflowSearchDatasets"]["datasetSummaries"]) == expected
|
|
)
|
|
|
|
|
|
def test_create_logged_model(mlflow_client: MlflowClient):
|
|
exp_id = mlflow_client.create_experiment("create_logged_model")
|
|
model = mlflow_client.create_logged_model(exp_id)
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert model.model_id == loaded_model.model_id
|
|
|
|
model = mlflow_client.create_logged_model(exp_id, name="my_model")
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert model.name == "my_model"
|
|
|
|
model = mlflow_client.create_logged_model(exp_id, model_type="LLM")
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert model.model_type == "LLM"
|
|
|
|
model = mlflow_client.create_logged_model(exp_id, source_run_id="123")
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert model.source_run_id == "123"
|
|
|
|
model = mlflow_client.create_logged_model(exp_id, params={"param": "value"})
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert model.params == {"param": "value"}
|
|
|
|
model = mlflow_client.create_logged_model(exp_id, tags={"tag": "value"})
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert model.tags == {"tag": "value"}
|
|
|
|
|
|
def test_log_logged_model_params(mlflow_client: MlflowClient):
|
|
exp_id = mlflow_client.create_experiment("create_logged_model")
|
|
model = mlflow_client.create_logged_model(exp_id)
|
|
mlflow_client.log_model_params(model.model_id, {"param": "value"})
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert loaded_model.params == {"param": "value"}
|
|
|
|
|
|
def test_finalize_logged_model(mlflow_client: MlflowClient):
|
|
exp_id = mlflow_client.create_experiment("create_logged_model")
|
|
model = mlflow_client.create_logged_model(exp_id)
|
|
finalized_model = mlflow_client.finalize_logged_model(model.model_id, LoggedModelStatus.READY)
|
|
assert finalized_model.status == LoggedModelStatus.READY
|
|
|
|
finalized_model = mlflow_client.finalize_logged_model(model.model_id, LoggedModelStatus.FAILED)
|
|
assert finalized_model.status == LoggedModelStatus.FAILED
|
|
|
|
|
|
def test_delete_logged_model(mlflow_client: MlflowClient):
|
|
exp_id = mlflow_client.create_experiment("delete_logged_model")
|
|
model = mlflow_client.create_logged_model(experiment_id=exp_id)
|
|
mlflow_client.delete_logged_model(model.model_id)
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
mlflow_client.get_logged_model(model.model_id)
|
|
|
|
models = mlflow_client.search_logged_models(experiment_ids=[exp_id])
|
|
assert len(models) == 0
|
|
|
|
|
|
def test_set_logged_model_tags(mlflow_client: MlflowClient):
|
|
exp_id = mlflow_client.create_experiment("create_logged_model")
|
|
model = mlflow_client.create_logged_model(exp_id)
|
|
mlflow_client.set_logged_model_tags(model.model_id, {"tag1": "value1", "tag2": "value2"})
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert loaded_model.tags == {"tag1": "value1", "tag2": "value2"}
|
|
|
|
mlflow_client.set_logged_model_tags(model.model_id, {"tag1": "value3"})
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert loaded_model.tags == {"tag1": "value3", "tag2": "value2"}
|
|
|
|
|
|
def test_delete_logged_model_tag(mlflow_client: MlflowClient):
|
|
exp_id = mlflow_client.create_experiment("create_logged_model")
|
|
model = mlflow_client.create_logged_model(exp_id)
|
|
mlflow_client.set_logged_model_tags(model.model_id, {"tag1": "value1", "tag2": "value2"})
|
|
mlflow_client.delete_logged_model_tag(model.model_id, "tag1")
|
|
loaded_model = mlflow_client.get_logged_model(model.model_id)
|
|
assert loaded_model.tags == {"tag2": "value2"}
|
|
|
|
with pytest.raises(MlflowException, match="No tag with key"):
|
|
mlflow_client.delete_logged_model_tag(model.model_id, "tag1")
|
|
|
|
|
|
def test_search_logged_models(mlflow_client: MlflowClient):
|
|
exp_id = mlflow_client.create_experiment("create_logged_model")
|
|
model_1 = mlflow_client.create_logged_model(exp_id)
|
|
time.sleep(0.001) # to ensure different created time
|
|
models = mlflow_client.search_logged_models(experiment_ids=[exp_id])
|
|
assert [m.name for m in models] == [model_1.name]
|
|
|
|
# max_results
|
|
model_2 = mlflow_client.create_logged_model(exp_id)
|
|
page_1 = mlflow_client.search_logged_models(experiment_ids=[exp_id], max_results=1)
|
|
assert [m.name for m in page_1] == [model_2.name]
|
|
assert page_1.token is not None
|
|
|
|
# pagination
|
|
page_2 = mlflow_client.search_logged_models(
|
|
experiment_ids=[exp_id], max_results=1, page_token=page_1.token
|
|
)
|
|
assert [m.name for m in page_2] == [model_1.name]
|
|
assert page_2.token is None
|
|
|
|
# filter_string
|
|
models = mlflow_client.search_logged_models(
|
|
experiment_ids=[exp_id], filter_string=f"name = {model_1.name!r}"
|
|
)
|
|
assert [m.name for m in models] == [model_1.name]
|
|
|
|
# datasets
|
|
run_1 = mlflow_client.create_run(exp_id)
|
|
mlflow_client.log_metric(
|
|
run_1.info.run_id,
|
|
key="metric",
|
|
value=1,
|
|
dataset_name="dataset",
|
|
dataset_digest="123",
|
|
model_id=model_1.model_id,
|
|
)
|
|
models = mlflow_client.search_logged_models(
|
|
experiment_ids=[exp_id],
|
|
datasets=[{"dataset_name": "dataset", "dataset_digest": "123"}],
|
|
)
|
|
|
|
assert [m.name for m in models] == [model_1.name]
|
|
|
|
# order_by
|
|
models = mlflow_client.search_logged_models(
|
|
experiment_ids=[exp_id],
|
|
order_by=[{"field_name": "creation_timestamp", "ascending": False}],
|
|
)
|
|
assert [m.name for m in models] == [model_2.name, model_1.name]
|
|
|
|
|
|
def test_log_outputs(mlflow_client: MlflowClient):
|
|
exp_id = mlflow_client.create_experiment("log_outputs")
|
|
run = mlflow_client.create_run(experiment_id=exp_id)
|
|
model = mlflow_client.create_logged_model(experiment_id=exp_id)
|
|
model_outputs = [LoggedModelOutput(model.model_id, 1)]
|
|
mlflow_client.log_outputs(run.info.run_id, model_outputs)
|
|
run = mlflow_client.get_run(run.info.run_id)
|
|
assert run.outputs.model_outputs == model_outputs
|
|
|
|
|
|
def test_list_logged_model_artifacts(mlflow_client: MlflowClient):
|
|
class Model(mlflow.pyfunc.PythonModel):
|
|
def predict(self, context, model_input):
|
|
return model_input
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
model_info = mlflow.pyfunc.log_model(name="model", python_model=Model())
|
|
resp = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/logged-models/{model_info.model_id}/artifacts/directories"
|
|
)
|
|
assert resp.status_code == 200
|
|
data = resp.json()
|
|
paths = [f["path"] for f in data["files"]]
|
|
assert "MLmodel" in paths
|
|
|
|
|
|
def test_get_logged_model_artifact(mlflow_client: MlflowClient):
|
|
class Model(mlflow.pyfunc.PythonModel):
|
|
def predict(self, context, model_input):
|
|
return model_input
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
model_info = mlflow.pyfunc.log_model(name="model", python_model=Model())
|
|
resp = requests.get(
|
|
f"{mlflow_client.tracking_uri}/ajax-api/2.0/mlflow/logged-models/{model_info.model_id}/artifacts/files",
|
|
params={"artifact_file_path": "MLmodel"},
|
|
)
|
|
assert resp.status_code == 200
|
|
assert model_info.model_id in resp.text
|
|
|
|
|
|
def test_suppress_url_printing(mlflow_client: MlflowClient, monkeypatch):
|
|
monkeypatch.setenv(MLFLOW_SUPPRESS_PRINTING_URL_TO_STDOUT.name, "true")
|
|
exp_id = mlflow_client.create_experiment("test_suppress_url_printing")
|
|
run = mlflow_client.create_run(experiment_id=exp_id)
|
|
captured_output = StringIO()
|
|
monkeypatch.setattr(sys, "stdout", captured_output)
|
|
mlflow_client._tracking_client._log_url(run.info.run_id)
|
|
assert captured_output.getvalue() == ""
|
|
|
|
|
|
def test_log_url_includes_workspace_when_set(mlflow_client: MlflowClient, monkeypatch):
|
|
exp_id = mlflow_client.create_experiment("test_log_url_workspace")
|
|
run = mlflow_client.create_run(experiment_id=exp_id)
|
|
captured_output = StringIO()
|
|
monkeypatch.setattr(sys, "stdout", captured_output)
|
|
monkeypatch.setattr(
|
|
"mlflow.tracking._tracking_service.client.get_workspace_url", lambda: "http://localhost"
|
|
)
|
|
monkeypatch.setattr(
|
|
"mlflow.tracking._tracking_service.client.get_request_workspace", lambda: "team-space"
|
|
)
|
|
|
|
mlflow_client._tracking_client._log_url(run.info.run_id)
|
|
|
|
out = captured_output.getvalue()
|
|
expected_fragment = f"/#/experiments/{exp_id}/runs/{run.info.run_id}?workspace=team-space"
|
|
assert expected_fragment in out
|
|
|
|
|
|
def test_assessments_end_to_end(mlflow_client):
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
|
|
# Set up experiment and trace
|
|
experiment_id = mlflow_client.create_experiment("assessment_crud_test")
|
|
trace_info = mlflow_client.start_trace(name="test_trace", experiment_id=experiment_id)
|
|
mlflow_client.end_trace(request_id=trace_info.request_id)
|
|
mlflow.flush_trace_async_logging()
|
|
|
|
# CREATE initial feedback assessment
|
|
feedback_payload = {
|
|
"assessment": {
|
|
"assessment_name": "quality_score",
|
|
"feedback": {"value": {"rating": 4, "comments": "Good response"}},
|
|
"source": {"source_type": "HUMAN", "source_id": "evaluator@company.com"},
|
|
"rationale": "Response was accurate and helpful",
|
|
"metadata": {"model": "gpt-4", "version": "1.0"},
|
|
}
|
|
}
|
|
|
|
# CREATE assessment
|
|
create_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments",
|
|
json=feedback_payload,
|
|
)
|
|
assert create_response.status_code == 200
|
|
assessment = create_response.json()["assessment"]
|
|
assessment_id = assessment["assessment_id"]
|
|
|
|
# Verify creation
|
|
assert assessment["assessment_name"] == "quality_score"
|
|
assert assessment["feedback"]["value"]["rating"] == 4
|
|
assert assessment["source"]["source_type"] == "HUMAN"
|
|
assert assessment["valid"] is True
|
|
|
|
# GET assessment
|
|
get_response = requests.get(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments/{assessment_id}"
|
|
)
|
|
assert get_response.status_code == 200
|
|
retrieved = get_response.json()["assessment"]
|
|
assert retrieved["assessment_id"] == assessment_id
|
|
assert retrieved["feedback"]["value"]["rating"] == 4
|
|
|
|
# UPDATE assessment
|
|
update_payload = {
|
|
"assessment": {
|
|
"assessment_id": assessment_id,
|
|
"trace_id": trace_info.request_id,
|
|
"assessment_name": "updated_quality_score",
|
|
"feedback": {"value": {"rating": 5, "comments": "Excellent response"}},
|
|
"rationale": "Actually, the response was excellent",
|
|
"metadata": {"model": "gpt-4", "version": "2.0"},
|
|
},
|
|
"update_mask": "assessmentName,feedback,rationale,metadata",
|
|
}
|
|
|
|
update_response = requests.patch(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments/{assessment_id}",
|
|
json=update_payload,
|
|
)
|
|
assert update_response.status_code == 200
|
|
updated = update_response.json()["assessment"]
|
|
assert updated["assessment_name"] == "updated_quality_score"
|
|
assert updated["feedback"]["value"]["rating"] == 5
|
|
assert updated["rationale"] == "Actually, the response was excellent"
|
|
|
|
# CREATE override assessment
|
|
override_payload = {
|
|
"assessment": {
|
|
"assessment_name": "corrected_quality_score",
|
|
"feedback": {"value": {"rating": 3, "comments": "Actually needs improvement"}},
|
|
"source": {"source_type": "HUMAN", "source_id": "senior_evaluator@company.com"},
|
|
"overrides": assessment_id,
|
|
}
|
|
}
|
|
|
|
override_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments",
|
|
json=override_payload,
|
|
)
|
|
assert override_response.status_code == 200
|
|
override_assessment = override_response.json()["assessment"]
|
|
override_id = override_assessment["assessment_id"]
|
|
|
|
# Verify original is now invalid
|
|
get_original = requests.get(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments/{assessment_id}"
|
|
)
|
|
assert get_original.status_code == 200
|
|
assert get_original.json()["assessment"]["valid"] is False
|
|
|
|
# Verify override is valid
|
|
get_override = requests.get(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments/{override_id}"
|
|
)
|
|
assert get_override.status_code == 200
|
|
assert get_override.json()["assessment"]["valid"] is True
|
|
assert get_override.json()["assessment"]["overrides"] == assessment_id
|
|
|
|
# DELETE override assessment (should restore original)
|
|
delete_response = requests.delete(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments/{override_id}"
|
|
)
|
|
assert delete_response.status_code == 200
|
|
|
|
# Verify override is deleted
|
|
get_deleted = requests.get(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments/{override_id}"
|
|
)
|
|
assert get_deleted.status_code == 404
|
|
|
|
# Verify original is restored to valid
|
|
get_restored = requests.get(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments/{assessment_id}"
|
|
)
|
|
assert get_restored.status_code == 200
|
|
assert get_restored.json()["assessment"]["valid"] is True
|
|
|
|
# CREATE expectation assessment to test different type
|
|
expectation_payload = {
|
|
"assessment": {
|
|
"assessment_name": "response_time_check",
|
|
"expectation": {"value": {"threshold_ms": 1000, "actual_ms": 750, "passed": True}},
|
|
"source": {"source_type": "CODE", "source_id": "automated_test"},
|
|
}
|
|
}
|
|
|
|
expectation_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments",
|
|
json=expectation_payload,
|
|
)
|
|
assert expectation_response.status_code == 200
|
|
expectation = expectation_response.json()["assessment"]
|
|
expectation_id = expectation["assessment_id"]
|
|
|
|
# Verify expectation was created correctly
|
|
expectation_value = json.loads(expectation["expectation"]["serialized_value"]["value"])
|
|
assert expectation_value["passed"] is True
|
|
assert expectation_value["threshold_ms"] == 1000
|
|
assert expectation_value["actual_ms"] == 750
|
|
assert expectation["source"]["source_type"] == "CODE"
|
|
|
|
# Clean up - delete remaining assessments
|
|
for aid in [assessment_id, expectation_id]:
|
|
delete_resp = requests.delete(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/traces/{trace_info.request_id}/assessments/{aid}"
|
|
)
|
|
assert delete_resp.status_code == 200
|
|
|
|
|
|
def test_graphql_nan_metric_handling(mlflow_client):
|
|
experiment_id = mlflow_client.create_experiment("test_graphql_nan_metrics")
|
|
created_run = mlflow_client.create_run(experiment_id)
|
|
run_id = created_run.info.run_id
|
|
|
|
# Log a normal metric and a NaN metric
|
|
mlflow_client.log_metric(run_id, key="normal_metric", value=123, timestamp=1, step=1)
|
|
mlflow_client.log_metric(run_id, key="nan_metric", value=math.nan, timestamp=2, step=2)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/graphql",
|
|
json={
|
|
"query": f"""
|
|
query testQuery {{
|
|
mlflowGetRun(input: {{runId: "{run_id}"}}) {{
|
|
run {{
|
|
data {{
|
|
metrics {{
|
|
key
|
|
value
|
|
timestamp
|
|
step
|
|
}}
|
|
}}
|
|
}}
|
|
}}
|
|
}}
|
|
""",
|
|
"operationName": "testQuery",
|
|
},
|
|
headers={"content-type": "application/json; charset=utf-8"},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
json_response = response.json()
|
|
assert json_response["errors"] is None
|
|
|
|
metrics = json_response["data"]["mlflowGetRun"]["run"]["data"]["metrics"]
|
|
|
|
# Find the normal metric and nan metric
|
|
normal_metric = None
|
|
nan_metric = None
|
|
for metric in metrics:
|
|
if metric["key"] == "normal_metric":
|
|
normal_metric = metric
|
|
elif metric["key"] == "nan_metric":
|
|
nan_metric = metric
|
|
|
|
# Verify normal metric has a numeric value
|
|
assert normal_metric is not None
|
|
assert normal_metric["key"] == "normal_metric"
|
|
assert normal_metric["value"] == 123
|
|
assert normal_metric["timestamp"] == "1"
|
|
assert normal_metric["step"] == "1"
|
|
|
|
# Verify NaN metric has null value
|
|
assert nan_metric is not None
|
|
assert nan_metric["key"] == "nan_metric"
|
|
assert nan_metric["value"] is None
|
|
assert nan_metric["timestamp"] == "2"
|
|
assert nan_metric["step"] == "2"
|
|
|
|
|
|
def test_create_and_get_evaluation_dataset(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Evaluation datasets not supported for FileStore")
|
|
|
|
experiment_id = mlflow_client.create_experiment("eval_dataset_test")
|
|
|
|
dataset = mlflow_client.create_dataset(
|
|
name="test_eval_dataset",
|
|
experiment_id=experiment_id,
|
|
tags={"environment": "test", "version": "1.0"},
|
|
)
|
|
|
|
assert dataset.name == "test_eval_dataset"
|
|
assert dataset.experiment_ids == [experiment_id]
|
|
assert dataset.tags["environment"] == "test"
|
|
assert dataset.tags["version"] == "1.0"
|
|
assert dataset.dataset_id is not None
|
|
|
|
retrieved = mlflow_client.get_dataset(dataset.dataset_id)
|
|
assert retrieved.name == dataset.name
|
|
assert retrieved.dataset_id == dataset.dataset_id
|
|
assert retrieved.tags == dataset.tags
|
|
|
|
|
|
def test_search_evaluation_datasets(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Evaluation datasets not supported for FileStore")
|
|
|
|
exp1 = mlflow_client.create_experiment("eval_search_exp1")
|
|
exp2 = mlflow_client.create_experiment("eval_search_exp2")
|
|
|
|
mlflow_client.create_dataset(
|
|
name="search_dataset_1", experiment_id=exp1, tags={"team": "ml", "status": "active"}
|
|
)
|
|
|
|
mlflow_client.create_dataset(
|
|
name="search_dataset_2",
|
|
experiment_id=[exp1, exp2],
|
|
tags={"team": "data", "status": "active"},
|
|
)
|
|
|
|
mlflow_client.create_dataset(
|
|
name="search_dataset_3", experiment_id=exp2, tags={"team": "ml", "status": "archived"}
|
|
)
|
|
|
|
all_datasets = mlflow_client.search_datasets()
|
|
assert len(all_datasets) >= 3
|
|
|
|
exp1_datasets = mlflow_client.search_datasets(experiment_ids=exp1)
|
|
dataset_names = [d.name for d in exp1_datasets]
|
|
assert "search_dataset_1" in dataset_names
|
|
assert "search_dataset_2" in dataset_names
|
|
|
|
ml_datasets = mlflow_client.search_datasets(filter_string="tags.team = 'ml'")
|
|
ml_names = [d.name for d in ml_datasets]
|
|
assert "search_dataset_1" in ml_names
|
|
assert "search_dataset_3" in ml_names
|
|
assert "search_dataset_2" not in ml_names
|
|
|
|
ordered_datasets = mlflow_client.search_datasets(order_by=["name ASC"])
|
|
names = [d.name for d in ordered_datasets]
|
|
assert names == sorted(names)
|
|
|
|
|
|
def test_evaluation_dataset_tag_operations(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Evaluation datasets not supported for FileStore")
|
|
|
|
experiment_id = mlflow_client.create_experiment("eval_tags_test")
|
|
|
|
dataset = mlflow_client.create_dataset(
|
|
name="tag_test_dataset",
|
|
experiment_id=experiment_id,
|
|
tags={"initial": "value", "env": "dev"},
|
|
)
|
|
|
|
mlflow_client.set_dataset_tags(dataset.dataset_id, {"env": "staging", "new_tag": "new_value"})
|
|
|
|
updated = mlflow_client.get_dataset(dataset.dataset_id)
|
|
assert updated.tags["initial"] == "value" # Original tag preserved
|
|
assert updated.tags["env"] == "staging" # Updated tag
|
|
assert updated.tags["new_tag"] == "new_value" # New tag added
|
|
|
|
mlflow_client.delete_dataset_tag(dataset.dataset_id, "new_tag")
|
|
|
|
final = mlflow_client.get_dataset(dataset.dataset_id)
|
|
assert "new_tag" not in final.tags
|
|
assert final.tags["env"] == "staging" # Other tags preserved
|
|
|
|
|
|
def test_evaluation_dataset_delete(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Evaluation datasets not supported for FileStore")
|
|
|
|
experiment_id = mlflow_client.create_experiment("eval_delete_test")
|
|
|
|
dataset = mlflow_client.create_dataset(
|
|
name="delete_test_dataset", experiment_id=experiment_id, tags={"to_delete": "yes"}
|
|
)
|
|
|
|
retrieved = mlflow_client.get_dataset(dataset.dataset_id)
|
|
assert retrieved.name == "delete_test_dataset"
|
|
|
|
mlflow_client.delete_dataset(dataset.dataset_id)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
mlflow_client.get_dataset(dataset.dataset_id)
|
|
|
|
|
|
def test_evaluation_dataset_upsert_records(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Evaluation datasets not supported for FileStore")
|
|
|
|
experiment_id = mlflow_client.create_experiment("upsert_records_test")
|
|
|
|
dataset = mlflow_client.create_dataset(
|
|
name="test_upsert_dataset",
|
|
experiment_id=experiment_id,
|
|
tags={"test": "upsert"},
|
|
)
|
|
|
|
initial_records = [
|
|
{
|
|
"inputs": {"question": "What is MLflow?"},
|
|
"expectations": {"answer": "MLflow is an ML platform"},
|
|
"tags": {"difficulty": "easy"},
|
|
},
|
|
{
|
|
"inputs": {"question": "What is Python?"},
|
|
"expectations": {"answer": "Python is a programming language"},
|
|
"tags": {"difficulty": "easy"},
|
|
},
|
|
]
|
|
|
|
# NB: MlflowClient doesn't have upsert_dataset_records method - merge_records() calls
|
|
# the store directly. We make HTTP requests here to test the REST API handler end-to-end.
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/datasets/{dataset.dataset_id}/records",
|
|
json={"records": json.dumps(initial_records)},
|
|
)
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
assert result["inserted_count"] == 2
|
|
assert result["updated_count"] == 0
|
|
|
|
update_records = [
|
|
{
|
|
"inputs": {"question": "What is MLflow?"},
|
|
"expectations": {"answer": "MLflow is an open-source ML platform"},
|
|
"tags": {"difficulty": "easy", "updated": "true"},
|
|
},
|
|
{
|
|
"inputs": {"question": "What is Docker?"},
|
|
"expectations": {"answer": "Docker is a containerization platform"},
|
|
"tags": {"difficulty": "medium"},
|
|
},
|
|
]
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/datasets/{dataset.dataset_id}/records",
|
|
json={"records": json.dumps(update_records)},
|
|
)
|
|
assert response.status_code == 200
|
|
result = response.json()
|
|
assert result["inserted_count"] == 1
|
|
assert result["updated_count"] == 1
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/datasets/invalid-id/records",
|
|
json={"records": json.dumps(initial_records)},
|
|
)
|
|
assert response.status_code != 200
|
|
|
|
|
|
def test_add_dataset_to_experiments_rest_tracking(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support dataset operations")
|
|
exp1 = mlflow_client.create_experiment("dataset_exp_1")
|
|
exp2 = mlflow_client.create_experiment("dataset_exp_2")
|
|
exp3 = mlflow_client.create_experiment("dataset_exp_3")
|
|
|
|
dataset = create_dataset(
|
|
name="test_multi_exp_dataset",
|
|
experiment_id=[exp1],
|
|
tags={"test": "multi_exp"},
|
|
)
|
|
|
|
assert len(dataset.experiment_ids) == 1
|
|
assert exp1 in dataset.experiment_ids
|
|
|
|
updated_dataset = add_dataset_to_experiments(
|
|
dataset_id=dataset.dataset_id,
|
|
experiment_ids=[exp2, exp3],
|
|
)
|
|
|
|
assert len(updated_dataset.experiment_ids) == 3
|
|
assert exp1 in updated_dataset.experiment_ids
|
|
assert exp2 in updated_dataset.experiment_ids
|
|
assert exp3 in updated_dataset.experiment_ids
|
|
|
|
retrieved = mlflow_client.get_dataset(dataset.dataset_id)
|
|
assert len(retrieved.experiment_ids) == 3
|
|
assert exp1 in retrieved.experiment_ids
|
|
assert exp2 in retrieved.experiment_ids
|
|
assert exp3 in retrieved.experiment_ids
|
|
|
|
|
|
def test_remove_dataset_from_experiments_rest_tracking(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support dataset operations")
|
|
exp1 = mlflow_client.create_experiment("dataset_remove_exp_1")
|
|
exp2 = mlflow_client.create_experiment("dataset_remove_exp_2")
|
|
exp3 = mlflow_client.create_experiment("dataset_remove_exp_3")
|
|
|
|
dataset = create_dataset(
|
|
name="test_remove_exp_dataset",
|
|
experiment_id=[exp1, exp2, exp3],
|
|
tags={"test": "remove_exp"},
|
|
)
|
|
|
|
assert len(dataset.experiment_ids) == 3
|
|
|
|
updated_dataset = remove_dataset_from_experiments(
|
|
dataset_id=dataset.dataset_id,
|
|
experiment_ids=[exp2],
|
|
)
|
|
|
|
assert len(updated_dataset.experiment_ids) == 2
|
|
assert exp1 in updated_dataset.experiment_ids
|
|
assert exp2 not in updated_dataset.experiment_ids
|
|
assert exp3 in updated_dataset.experiment_ids
|
|
|
|
retrieved = mlflow_client.get_dataset(dataset.dataset_id)
|
|
assert len(retrieved.experiment_ids) == 2
|
|
|
|
updated_dataset = remove_dataset_from_experiments(
|
|
dataset_id=dataset.dataset_id,
|
|
experiment_ids=[exp1, exp3],
|
|
)
|
|
|
|
assert len(updated_dataset.experiment_ids) == 0
|
|
|
|
retrieved = mlflow_client.get_dataset(dataset.dataset_id)
|
|
assert len(retrieved.experiment_ids) == 0
|
|
|
|
|
|
def test_add_multiple_experiments_at_once_rest_tracking(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support dataset operations")
|
|
exps = [mlflow_client.create_experiment(f"bulk_add_exp_{i}") for i in range(5)]
|
|
|
|
dataset = create_dataset(
|
|
name="test_bulk_add_dataset",
|
|
experiment_id=[exps[0]],
|
|
tags={"test": "bulk_add"},
|
|
)
|
|
|
|
updated_dataset = add_dataset_to_experiments(
|
|
dataset_id=dataset.dataset_id,
|
|
experiment_ids=exps[1:],
|
|
)
|
|
|
|
assert len(updated_dataset.experiment_ids) == 5
|
|
for exp in exps:
|
|
assert exp in updated_dataset.experiment_ids
|
|
|
|
|
|
def test_dataset_experiment_association_error_cases_rest_tracking(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support dataset operations")
|
|
exp1 = mlflow_client.create_experiment("error_test_exp")
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
add_dataset_to_experiments(
|
|
dataset_id="d-nonexistent1234567890abcdef1234",
|
|
experiment_ids=[exp1],
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
remove_dataset_from_experiments(
|
|
dataset_id="d-nonexistent1234567890abcdef1234",
|
|
experiment_ids=[exp1],
|
|
)
|
|
|
|
|
|
def test_idempotent_add_experiments_rest_tracking(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support dataset operations")
|
|
exp1 = mlflow_client.create_experiment("idempotent_test_exp_1")
|
|
exp2 = mlflow_client.create_experiment("idempotent_test_exp_2")
|
|
|
|
dataset = create_dataset(
|
|
name="test_idempotent_dataset",
|
|
experiment_id=[exp1, exp2],
|
|
tags={"test": "idempotent"},
|
|
)
|
|
|
|
assert len(dataset.experiment_ids) == 2
|
|
|
|
updated_dataset = add_dataset_to_experiments(
|
|
dataset_id=dataset.dataset_id,
|
|
experiment_ids=[exp1],
|
|
)
|
|
|
|
assert len(updated_dataset.experiment_ids) == 2
|
|
assert exp1 in updated_dataset.experiment_ids
|
|
assert exp2 in updated_dataset.experiment_ids
|
|
|
|
|
|
def test_idempotent_remove_experiments_rest_tracking(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support dataset operations")
|
|
exp1 = mlflow_client.create_experiment("remove_idempotent_test_exp_1")
|
|
exp2 = mlflow_client.create_experiment("remove_idempotent_test_exp_2")
|
|
|
|
dataset = create_dataset(
|
|
name="test_remove_idempotent_dataset",
|
|
experiment_id=[exp1],
|
|
tags={"test": "remove_idempotent"},
|
|
)
|
|
|
|
assert len(dataset.experiment_ids) == 1
|
|
|
|
updated_dataset = remove_dataset_from_experiments(
|
|
dataset_id=dataset.dataset_id,
|
|
experiment_ids=[exp2],
|
|
)
|
|
|
|
assert len(updated_dataset.experiment_ids) == 1
|
|
assert exp1 in updated_dataset.experiment_ids
|
|
|
|
|
|
def test_client_api_add_remove_experiments_rest_tracking(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support dataset operations")
|
|
exp1 = mlflow_client.create_experiment("client_api_exp_1")
|
|
exp2 = mlflow_client.create_experiment("client_api_exp_2")
|
|
exp3 = mlflow_client.create_experiment("client_api_exp_3")
|
|
|
|
dataset = mlflow_client.create_dataset(
|
|
name="test_client_api_dataset",
|
|
experiment_id=[exp1],
|
|
tags={"test": "client_api"},
|
|
)
|
|
|
|
updated_dataset = mlflow_client.add_dataset_to_experiments(
|
|
dataset_id=dataset.dataset_id,
|
|
experiment_ids=[exp2, exp3],
|
|
)
|
|
|
|
assert len(updated_dataset.experiment_ids) == 3
|
|
|
|
updated_dataset = mlflow_client.remove_dataset_from_experiments(
|
|
dataset_id=dataset.dataset_id,
|
|
experiment_ids=[exp2],
|
|
)
|
|
|
|
assert len(updated_dataset.experiment_ids) == 2
|
|
assert exp1 in updated_dataset.experiment_ids
|
|
assert exp2 not in updated_dataset.experiment_ids
|
|
assert exp3 in updated_dataset.experiment_ids
|
|
|
|
|
|
def test_scorer_CRUD(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("File store doesn't support scorer CRUD operations")
|
|
experiment_id = mlflow_client.create_experiment("test_scorer_api_experiment")
|
|
|
|
# Get the RestStore object directly
|
|
store = mlflow_client._tracking_client.store
|
|
|
|
# Test register scorer
|
|
scorer_data = {"name": "test_scorer", "original_func_name": "test_func"}
|
|
serialized_scorer = json.dumps(scorer_data)
|
|
|
|
version = store.register_scorer(experiment_id, "test_scorer", serialized_scorer)
|
|
assert version.scorer_version == 1
|
|
|
|
# Test list scorers
|
|
scorers = store.list_scorers(experiment_id)
|
|
assert len(scorers) == 1
|
|
assert scorers[0].scorer_name == "test_scorer"
|
|
assert scorers[0].scorer_version == 1
|
|
|
|
# Test list scorer versions
|
|
versions = store.list_scorer_versions(str(experiment_id), "test_scorer")
|
|
assert len(versions) == 1
|
|
assert versions[0].scorer_name == "test_scorer"
|
|
assert versions[0].scorer_version == 1
|
|
|
|
# Test get scorer (latest version)
|
|
scorer = store.get_scorer(str(experiment_id), "test_scorer")
|
|
assert scorer.scorer_name == "test_scorer"
|
|
assert scorer.scorer_version == 1
|
|
|
|
# Test get scorer (specific version)
|
|
scorer_v1 = store.get_scorer(str(experiment_id), "test_scorer", version=1)
|
|
assert scorer_v1.scorer_name == "test_scorer"
|
|
assert scorer_v1.scorer_version == 1
|
|
|
|
# Test register second version
|
|
scorer_data_v2 = {
|
|
"name": "test_scorer_v2",
|
|
"original_func_name": "test_func_v2",
|
|
}
|
|
serialized_scorer_v2 = json.dumps(scorer_data_v2)
|
|
|
|
version_v2 = store.register_scorer(str(experiment_id), "test_scorer", serialized_scorer_v2)
|
|
assert version_v2.scorer_version == 2
|
|
|
|
# Verify list scorers returns latest version
|
|
scorers_after_v2 = store.list_scorers(str(experiment_id))
|
|
assert len(scorers_after_v2) == 1
|
|
assert scorers_after_v2[0].scorer_version == 2
|
|
|
|
# Verify list versions returns both versions
|
|
versions_after_v2 = store.list_scorer_versions(str(experiment_id), "test_scorer")
|
|
assert len(versions_after_v2) == 2
|
|
|
|
# Test delete specific version
|
|
store.delete_scorer(str(experiment_id), "test_scorer", version=1)
|
|
|
|
# Verify version 1 is deleted
|
|
versions_after_delete = store.list_scorer_versions(str(experiment_id), "test_scorer")
|
|
assert len(versions_after_delete) == 1
|
|
assert versions_after_delete[0].scorer_version == 2
|
|
|
|
# Test delete all versions
|
|
store.delete_scorer(str(experiment_id), "test_scorer")
|
|
|
|
# Verify all versions are deleted
|
|
scorers_after_delete_all = store.list_scorers(str(experiment_id))
|
|
assert len(scorers_after_delete_all) == 0
|
|
|
|
# Clean up
|
|
mlflow_client.delete_experiment(experiment_id)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"filter_string",
|
|
[
|
|
"status = 'OK'",
|
|
None,
|
|
],
|
|
)
|
|
def test_online_scoring_config(mlflow_client_with_secrets, filter_string):
|
|
"""
|
|
Smoke test for online scoring configuration REST APIs.
|
|
Tests upsert_online_scoring_config and get_online_scoring_configs with both
|
|
string and None filter values (None is sent by UI when filter field is blank).
|
|
"""
|
|
experiment_id = mlflow_client_with_secrets.create_experiment("test_online_scoring")
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-secret", secret_value={"api_key": "sk-test"}, provider="openai"
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="test-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
)
|
|
],
|
|
)
|
|
|
|
scorer_data = {"instructions_judge_pydantic_data": {"model": f"gateway:/{endpoint.name}"}}
|
|
serialized_scorer = json.dumps(scorer_data)
|
|
scorer_version = store.register_scorer(experiment_id, "my_scorer", serialized_scorer)
|
|
scorer_id = scorer_version.scorer_id
|
|
|
|
config = store.upsert_online_scoring_config(
|
|
experiment_id=experiment_id,
|
|
scorer_name="my_scorer",
|
|
sample_rate=0.5,
|
|
filter_string=filter_string,
|
|
)
|
|
assert config.scorer_id == scorer_id
|
|
assert config.sample_rate == 0.5
|
|
assert config.filter_string == filter_string
|
|
assert config.experiment_id == experiment_id
|
|
|
|
configs = store.get_online_scoring_configs([scorer_id])
|
|
assert len(configs) == 1
|
|
assert configs[0].scorer_id == scorer_id
|
|
assert configs[0].sample_rate == 0.5
|
|
assert configs[0].filter_string == filter_string
|
|
|
|
# Update with different filter string to test update functionality
|
|
updated_filter = "status = 'COMPLETED'"
|
|
updated_config = store.upsert_online_scoring_config(
|
|
experiment_id=experiment_id,
|
|
scorer_name="my_scorer",
|
|
sample_rate=0.8,
|
|
filter_string=updated_filter,
|
|
)
|
|
assert updated_config.scorer_id == scorer_id
|
|
assert updated_config.sample_rate == 0.8
|
|
assert updated_config.filter_string == updated_filter
|
|
|
|
configs_after_update = store.get_online_scoring_configs([scorer_id])
|
|
assert len(configs_after_update) == 1
|
|
assert configs_after_update[0].sample_rate == 0.8
|
|
assert configs_after_update[0].filter_string == updated_filter
|
|
|
|
|
|
@pytest.mark.parametrize("use_async", [False, True])
|
|
@pytest.mark.asyncio
|
|
async def test_rest_store_logs_spans_via_otel_endpoint(mlflow_client, store_type, use_async):
|
|
"""
|
|
End-to-end test that verifies RestStore can log spans to a running server via OTLP endpoint.
|
|
|
|
This test:
|
|
1. Creates spans using MLflow's span entities
|
|
2. Uses RestStore.log_spans or log_spans_async to send them via OTLP protocol
|
|
3. Verifies the spans were stored and can be retrieved
|
|
"""
|
|
if store_type == "file":
|
|
pytest.skip("FileStore does not support OTLP span logging")
|
|
|
|
experiment_id = mlflow_client.create_experiment(f"rest_store_otel_test_{use_async}")
|
|
root_span = mlflow_client.start_trace(
|
|
f"rest_store_otel_trace_{use_async}", experiment_id=experiment_id
|
|
)
|
|
otel_span = OTelReadableSpan(
|
|
name=f"test-rest-store-span-{use_async}",
|
|
context=build_otel_context(
|
|
trace_id=int(root_span.trace_id[3:], 16), # Remove 'tr-' prefix and convert to int
|
|
span_id=0x1234567890ABCDEF,
|
|
),
|
|
parent=None,
|
|
start_time=1000000000,
|
|
end_time=2000000000,
|
|
attributes={
|
|
SpanAttributeKey.REQUEST_ID: root_span.trace_id,
|
|
"test.attribute": json.dumps(f"test-value-{use_async}"), # JSON-encoded string value
|
|
},
|
|
resource=None,
|
|
)
|
|
mlflow_span_to_log = Span(otel_span)
|
|
# Call either sync or async version based on parametrization
|
|
if use_async:
|
|
# Use await to execute the async method
|
|
result_spans = await mlflow_client._tracking_client.store.log_spans_async(
|
|
location=experiment_id, spans=[mlflow_span_to_log]
|
|
)
|
|
else:
|
|
result_spans = mlflow_client._tracking_client.store.log_spans(
|
|
location=experiment_id, spans=[mlflow_span_to_log]
|
|
)
|
|
|
|
# Verify the spans were returned (indicates successful logging)
|
|
assert len(result_spans) == 1
|
|
assert result_spans[0].name == f"test-rest-store-span-{use_async}"
|
|
|
|
|
|
# =============================================================================
|
|
# Secrets and Endpoints E2E Tests
|
|
# =============================================================================
|
|
|
|
|
|
def test_create_and_get_secret(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-api-key",
|
|
secret_value={"api_key": "sk-test-12345"},
|
|
provider="openai",
|
|
)
|
|
|
|
assert secret.secret_name == "test-api-key"
|
|
assert secret.provider == "openai"
|
|
assert secret.secret_id is not None
|
|
|
|
fetched = store.get_secret_info(secret.secret_id)
|
|
assert fetched.secret_name == "test-api-key"
|
|
assert fetched.provider == "openai"
|
|
assert fetched.secret_id == secret.secret_id
|
|
|
|
|
|
def test_update_secret(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-key",
|
|
secret_value={"api_key": "initial-value"},
|
|
provider="anthropic",
|
|
)
|
|
|
|
updated = store.update_gateway_secret(
|
|
secret_id=secret.secret_id,
|
|
secret_value={"api_key": "updated-value"},
|
|
)
|
|
|
|
assert updated.secret_id == secret.secret_id
|
|
assert updated.secret_name == "test-key"
|
|
|
|
|
|
def test_list_secret_infos(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret1 = store.create_gateway_secret(
|
|
secret_name="openai-key",
|
|
secret_value={"api_key": "sk-openai"},
|
|
provider="openai",
|
|
)
|
|
store.create_gateway_secret(
|
|
secret_name="anthropic-key",
|
|
secret_value={"api_key": "sk-ant"},
|
|
provider="anthropic",
|
|
)
|
|
|
|
all_secrets = store.list_secret_infos()
|
|
assert len(all_secrets) >= 2
|
|
|
|
openai_secrets = store.list_secret_infos(provider="openai")
|
|
assert len(openai_secrets) >= 1
|
|
assert any(s.secret_id == secret1.secret_id for s in openai_secrets)
|
|
|
|
|
|
def test_delete_secret(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="temp-key",
|
|
secret_value={"api_key": "temp-value"},
|
|
)
|
|
|
|
store.delete_gateway_secret(secret.secret_id)
|
|
|
|
all_secrets = store.list_secret_infos()
|
|
assert not any(s.secret_id == secret.secret_id for s in all_secrets)
|
|
|
|
|
|
def test_create_secret_with_dict_value(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="aws-creds",
|
|
secret_value={"aws_access_key_id": "AKIATEST1234", "aws_secret_access_key": "secret123abc"},
|
|
provider="bedrock",
|
|
)
|
|
|
|
assert secret.secret_name == "aws-creds"
|
|
assert secret.provider == "bedrock"
|
|
assert secret.secret_id is not None
|
|
assert isinstance(secret.masked_values, dict)
|
|
assert secret.masked_values == {
|
|
"aws_access_key_id": "AKI...1234",
|
|
"aws_secret_access_key": "sec...3abc",
|
|
}
|
|
|
|
|
|
def test_update_secret_with_dict_value(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="aws-creds-update",
|
|
secret_value={"api_key": "initial-value-1234"},
|
|
provider="bedrock",
|
|
)
|
|
|
|
assert isinstance(secret.masked_values, dict)
|
|
assert secret.masked_values == {"api_key": "ini...1234"}
|
|
|
|
updated = store.update_gateway_secret(
|
|
secret_id=secret.secret_id,
|
|
secret_value={
|
|
"aws_access_key_id": "NEWKEY123456",
|
|
"aws_secret_access_key": "newsecret1234",
|
|
},
|
|
)
|
|
|
|
assert updated.secret_id == secret.secret_id
|
|
assert updated.secret_name == "aws-creds-update"
|
|
assert isinstance(updated.masked_values, dict)
|
|
assert updated.masked_values == {
|
|
"aws_access_key_id": "NEW...3456",
|
|
"aws_secret_access_key": "new...1234",
|
|
}
|
|
|
|
|
|
def test_create_and_update_compound_secret_via_rest(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="bedrock-aws-creds",
|
|
secret_value={
|
|
"aws_access_key_id": "AKIAORIGINAL1234",
|
|
"aws_secret_access_key": "original-secret-key-1234",
|
|
},
|
|
provider="bedrock",
|
|
auth_config={"auth_mode": "access_keys", "aws_region_name": "us-east-1"},
|
|
)
|
|
|
|
assert secret.secret_name == "bedrock-aws-creds"
|
|
assert secret.provider == "bedrock"
|
|
assert isinstance(secret.masked_values, dict)
|
|
assert secret.masked_values == {
|
|
"aws_access_key_id": "AKI...1234",
|
|
"aws_secret_access_key": "ori...1234",
|
|
}
|
|
|
|
fetched = store.get_secret_info(secret_id=secret.secret_id)
|
|
assert fetched.secret_id == secret.secret_id
|
|
assert isinstance(fetched.masked_values, dict)
|
|
assert fetched.masked_values == secret.masked_values
|
|
|
|
updated = store.update_gateway_secret(
|
|
secret_id=secret.secret_id,
|
|
secret_value={
|
|
"aws_access_key_id": "AKIAROTATED5678",
|
|
"aws_secret_access_key": "rotated-secret-key-5678",
|
|
},
|
|
)
|
|
|
|
assert updated.secret_id == secret.secret_id
|
|
assert updated.last_updated_at > secret.created_at
|
|
assert isinstance(updated.masked_values, dict)
|
|
assert updated.masked_values == {
|
|
"aws_access_key_id": "AKI...5678",
|
|
"aws_secret_access_key": "rot...5678",
|
|
}
|
|
|
|
|
|
def test_create_and_get_endpoint(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-api-key",
|
|
secret_value={"api_key": "sk-test-12345"},
|
|
provider="openai",
|
|
)
|
|
secret2 = store.create_gateway_secret(
|
|
secret_name="test-api-key-fallback",
|
|
secret_value={"api_key": "sk-test-67890"},
|
|
provider="anthropic",
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="test-model-def",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
model_def_fallback = store.create_gateway_model_definition(
|
|
name="test-model-def-fallback",
|
|
secret_id=secret2.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-5-sonnet",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="test-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def_fallback.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
weight=1.0,
|
|
fallback_order=0,
|
|
),
|
|
],
|
|
routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
|
|
fallback_config=FallbackConfig(
|
|
strategy=FallbackStrategy.SEQUENTIAL,
|
|
max_attempts=2,
|
|
),
|
|
)
|
|
|
|
assert endpoint.name == "test-endpoint"
|
|
assert endpoint.endpoint_id is not None
|
|
assert len(endpoint.model_mappings) == 2
|
|
assert endpoint.model_mappings[0].model_definition.model_name == "gpt-4"
|
|
assert endpoint.routing_strategy == RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT
|
|
assert endpoint.fallback_config is not None
|
|
assert endpoint.fallback_config.strategy == FallbackStrategy.SEQUENTIAL
|
|
assert endpoint.fallback_config.max_attempts == 2
|
|
|
|
fetched = store.get_gateway_endpoint(endpoint.endpoint_id)
|
|
assert fetched.name == "test-endpoint"
|
|
assert fetched.endpoint_id == endpoint.endpoint_id
|
|
assert len(fetched.model_mappings) == 2
|
|
assert fetched.routing_strategy == RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT
|
|
assert fetched.fallback_config is not None
|
|
assert fetched.fallback_config.strategy == FallbackStrategy.SEQUENTIAL
|
|
assert fetched.fallback_config.max_attempts == 2
|
|
|
|
|
|
def test_create_endpoint_with_usage_tracking(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="usage-tracking-test-key",
|
|
secret_value={"api_key": "sk-usage-tracking-test"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="usage-tracking-model-def",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="usage-tracking-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
)
|
|
],
|
|
usage_tracking=True,
|
|
)
|
|
|
|
assert endpoint.usage_tracking is True
|
|
experiment_id = endpoint.experiment_id
|
|
|
|
# Experiment is automatically created with usage tracking enabled
|
|
experiment = mlflow_client_with_secrets.get_experiment(experiment_id)
|
|
assert experiment.name == "gateway/usage-tracking-endpoint"
|
|
|
|
|
|
def test_update_endpoint(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-api-key-2",
|
|
secret_value={"api_key": "sk-test-67890"},
|
|
provider="anthropic",
|
|
)
|
|
secret2 = store.create_gateway_secret(
|
|
secret_name="test-api-key-2-fallback",
|
|
secret_value={"api_key": "sk-test-99999"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="test-model-def-2",
|
|
secret_id=secret.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-5-sonnet",
|
|
)
|
|
model_def_fallback = store.create_gateway_model_definition(
|
|
name="test-model-def-2-fallback",
|
|
secret_id=secret2.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="initial-name",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
updated = store.update_gateway_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
name="updated-name",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def_fallback.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
weight=1.0,
|
|
fallback_order=0,
|
|
),
|
|
],
|
|
routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
|
|
fallback_config=FallbackConfig(
|
|
strategy=FallbackStrategy.SEQUENTIAL,
|
|
max_attempts=3,
|
|
),
|
|
)
|
|
|
|
assert updated.endpoint_id == endpoint.endpoint_id
|
|
assert updated.name == "updated-name"
|
|
assert updated.routing_strategy == RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT
|
|
assert updated.fallback_config is not None
|
|
assert updated.fallback_config.strategy == FallbackStrategy.SEQUENTIAL
|
|
assert updated.fallback_config.max_attempts == 3
|
|
assert len(updated.model_mappings) == 2
|
|
|
|
|
|
def test_list_endpoints(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret1 = store.create_gateway_secret(
|
|
secret_name="test-api-key-3",
|
|
secret_value={"api_key": "sk-test-11111"},
|
|
provider="openai",
|
|
)
|
|
secret2 = store.create_gateway_secret(
|
|
secret_name="test-api-key-4",
|
|
secret_value={"api_key": "sk-test-22222"},
|
|
provider="openai",
|
|
)
|
|
secret3 = store.create_gateway_secret(
|
|
secret_name="test-api-key-fallback-3",
|
|
secret_value={"api_key": "sk-test-44444"},
|
|
provider="anthropic",
|
|
)
|
|
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="test-model-def-3",
|
|
secret_id=secret1.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="test-model-def-4",
|
|
secret_id=secret2.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-3.5-turbo",
|
|
)
|
|
model_def3 = store.create_gateway_model_definition(
|
|
name="test-model-def-fallback-3",
|
|
secret_id=secret3.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-5-sonnet",
|
|
)
|
|
|
|
# Create endpoint without fallback
|
|
endpoint1 = store.create_gateway_endpoint(
|
|
name="endpoint-1",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
# Create endpoint with fallback
|
|
endpoint2 = store.create_gateway_endpoint(
|
|
name="endpoint-2",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def3.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
weight=1.0,
|
|
fallback_order=0,
|
|
),
|
|
],
|
|
routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
|
|
fallback_config=FallbackConfig(
|
|
strategy=FallbackStrategy.SEQUENTIAL,
|
|
max_attempts=2,
|
|
),
|
|
)
|
|
|
|
all_endpoints = store.list_gateway_endpoints()
|
|
assert len(all_endpoints) >= 2
|
|
endpoint_ids = {e.endpoint_id for e in all_endpoints}
|
|
assert endpoint1.endpoint_id in endpoint_ids
|
|
assert endpoint2.endpoint_id in endpoint_ids
|
|
|
|
# Find and verify endpoints
|
|
found_ep1 = next(e for e in all_endpoints if e.endpoint_id == endpoint1.endpoint_id)
|
|
found_ep2 = next(e for e in all_endpoints if e.endpoint_id == endpoint2.endpoint_id)
|
|
|
|
assert found_ep1.routing_strategy is None
|
|
assert found_ep1.fallback_config is None
|
|
|
|
assert found_ep2.routing_strategy == RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT
|
|
assert found_ep2.fallback_config is not None
|
|
assert found_ep2.fallback_config.strategy == FallbackStrategy.SEQUENTIAL
|
|
assert found_ep2.fallback_config.max_attempts == 2
|
|
|
|
|
|
def test_delete_endpoint(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-api-key-5",
|
|
secret_value={"api_key": "sk-test-33333"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="test-model-def-5",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="temp-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
store.delete_gateway_endpoint(endpoint.endpoint_id)
|
|
|
|
all_endpoints = store.list_gateway_endpoints()
|
|
assert not any(e.endpoint_id == endpoint.endpoint_id for e in all_endpoints)
|
|
|
|
|
|
def test_model_definitions(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="model-secret",
|
|
secret_value={"api_key": "sk-test"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="test-model-def",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
assert model_def.name == "test-model-def"
|
|
assert model_def.secret_id == secret.secret_id
|
|
assert model_def.provider == "openai"
|
|
assert model_def.model_name == "gpt-4"
|
|
assert model_def.model_definition_id is not None
|
|
|
|
fetched = store.get_gateway_model_definition(model_def.model_definition_id)
|
|
assert fetched.model_definition_id == model_def.model_definition_id
|
|
assert fetched.name == "test-model-def"
|
|
|
|
updated = store.update_gateway_model_definition(
|
|
model_definition_id=model_def.model_definition_id,
|
|
model_name="gpt-4-turbo",
|
|
)
|
|
assert updated.model_definition_id == model_def.model_definition_id
|
|
assert updated.model_name == "gpt-4-turbo"
|
|
|
|
all_defs = store.list_gateway_model_definitions()
|
|
assert any(d.model_definition_id == model_def.model_definition_id for d in all_defs)
|
|
|
|
store.delete_gateway_model_definition(model_def.model_definition_id)
|
|
|
|
all_defs_after = store.list_gateway_model_definitions()
|
|
assert not any(d.model_definition_id == model_def.model_definition_id for d in all_defs_after)
|
|
|
|
|
|
def test_attach_detach_model_to_endpoint(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="attach-detach-secret",
|
|
secret_value={"api_key": "sk-test-attach"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="attach-model-def-1",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="attach-model-def-2",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-3.5-turbo",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="attach-test-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
assert len(endpoint.model_mappings) == 1
|
|
assert endpoint.model_mappings[0].model_definition.model_name == "gpt-4"
|
|
|
|
mapping = store.attach_model_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
model_config=GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
)
|
|
|
|
assert mapping.endpoint_id == endpoint.endpoint_id
|
|
assert mapping.model_definition_id == model_def2.model_definition_id
|
|
|
|
fetched_endpoint = store.get_gateway_endpoint(endpoint.endpoint_id)
|
|
assert len(fetched_endpoint.model_mappings) == 2
|
|
|
|
store.detach_model_from_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
model_definition_id=model_def2.model_definition_id,
|
|
)
|
|
|
|
fetched_endpoint_after = store.get_gateway_endpoint(endpoint.endpoint_id)
|
|
assert len(fetched_endpoint_after.model_mappings) == 1
|
|
|
|
|
|
def test_endpoint_bindings(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="binding-secret",
|
|
secret_value={"api_key": "sk-test-44444"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="binding-model-def-1",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="binding-model-def-2",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-3.5-turbo",
|
|
)
|
|
|
|
endpoint1 = store.create_gateway_endpoint(
|
|
name="binding-test-endpoint-1",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
endpoint2 = store.create_gateway_endpoint(
|
|
name="binding-test-endpoint-2",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
binding1 = store.create_endpoint_binding(
|
|
endpoint_id=endpoint1.endpoint_id,
|
|
resource_type=GatewayResourceType.SCORER,
|
|
resource_id="job-123",
|
|
)
|
|
|
|
binding2 = store.create_endpoint_binding(
|
|
endpoint_id=endpoint1.endpoint_id,
|
|
resource_type=GatewayResourceType.SCORER,
|
|
resource_id="job-456",
|
|
)
|
|
|
|
binding3 = store.create_endpoint_binding(
|
|
endpoint_id=endpoint2.endpoint_id,
|
|
resource_type=GatewayResourceType.SCORER,
|
|
resource_id="job-789",
|
|
)
|
|
|
|
assert binding1.endpoint_id == endpoint1.endpoint_id
|
|
assert binding1.resource_type == GatewayResourceType.SCORER
|
|
assert binding1.resource_id == "job-123"
|
|
|
|
bindings_endpoint1 = store.list_endpoint_bindings(endpoint_id=endpoint1.endpoint_id)
|
|
assert len(bindings_endpoint1) == 2
|
|
resource_ids = {b.resource_id for b in bindings_endpoint1}
|
|
assert binding1.resource_id in resource_ids
|
|
assert binding2.resource_id in resource_ids
|
|
assert binding3.resource_id not in resource_ids
|
|
|
|
bindings_by_type = store.list_endpoint_bindings(resource_type=GatewayResourceType.SCORER)
|
|
assert len(bindings_by_type) >= 3
|
|
|
|
bindings_by_resource = store.list_endpoint_bindings(resource_id="job-123")
|
|
assert len(bindings_by_resource) == 1
|
|
assert bindings_by_resource[0].resource_id == binding1.resource_id
|
|
|
|
bindings_multi = store.list_endpoint_bindings(
|
|
endpoint_id=endpoint1.endpoint_id,
|
|
resource_type=GatewayResourceType.SCORER,
|
|
)
|
|
assert len(bindings_multi) == 2
|
|
|
|
store.delete_endpoint_binding(
|
|
endpoint_id=binding1.endpoint_id,
|
|
resource_type=binding1.resource_type.value,
|
|
resource_id=binding1.resource_id,
|
|
)
|
|
|
|
bindings_after = store.list_endpoint_bindings(endpoint_id=endpoint1.endpoint_id)
|
|
assert len(bindings_after) == 1
|
|
assert not any(b.resource_id == binding1.resource_id for b in bindings_after)
|
|
|
|
|
|
def test_secrets_and_endpoints_integration(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="integration-test-key",
|
|
secret_value={"api_key": "sk-integration-test"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="integration-model-def-1",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-3.5-turbo",
|
|
)
|
|
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="integration-model-def-2",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="integration-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
mapping = store.attach_model_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
model_config=GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
)
|
|
|
|
binding = store.create_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type=GatewayResourceType.SCORER,
|
|
resource_id="integration-job",
|
|
)
|
|
|
|
fetched_endpoint = store.get_gateway_endpoint(endpoint.endpoint_id)
|
|
assert len(fetched_endpoint.model_mappings) == 2
|
|
mapping_ids = {m.mapping_id for m in fetched_endpoint.model_mappings}
|
|
assert mapping.mapping_id in mapping_ids
|
|
|
|
bindings = store.list_endpoint_bindings(resource_id="integration-job")
|
|
assert len(bindings) == 1
|
|
assert bindings[0].resource_id == binding.resource_id
|
|
|
|
store.delete_endpoint_binding(
|
|
endpoint_id=binding.endpoint_id,
|
|
resource_type=binding.resource_type.value,
|
|
resource_id=binding.resource_id,
|
|
)
|
|
store.detach_model_from_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
model_definition_id=model_def2.model_definition_id,
|
|
)
|
|
store.delete_gateway_endpoint(endpoint.endpoint_id)
|
|
store.delete_gateway_model_definition(model_def1.model_definition_id)
|
|
store.delete_gateway_model_definition(model_def2.model_definition_id)
|
|
store.delete_gateway_secret(secret.secret_id)
|
|
|
|
|
|
def test_list_providers(mlflow_client_with_secrets):
|
|
import requests
|
|
|
|
base_url = mlflow_client_with_secrets._tracking_client.tracking_uri
|
|
response = requests.get(f"{base_url}/ajax-api/3.0/mlflow/gateway/supported-providers")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "providers" in data
|
|
assert isinstance(data["providers"], list)
|
|
assert len(data["providers"]) > 0
|
|
assert "openai" in data["providers"]
|
|
|
|
|
|
def test_list_models(mlflow_client_with_secrets):
|
|
import requests
|
|
|
|
base_url = mlflow_client_with_secrets._tracking_client.tracking_uri
|
|
response = requests.get(f"{base_url}/ajax-api/3.0/mlflow/gateway/supported-models")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "models" in data
|
|
assert isinstance(data["models"], list)
|
|
assert len(data["models"]) > 0
|
|
|
|
model = data["models"][0]
|
|
assert "model" in model
|
|
assert "provider" in model
|
|
assert "mode" in model
|
|
assert all(not m["model"].startswith("ft:") for m in data["models"])
|
|
|
|
response = requests.get(
|
|
f"{base_url}/ajax-api/3.0/mlflow/gateway/supported-models", params={"provider": "openai"}
|
|
)
|
|
assert response.status_code == 200
|
|
filtered_data = response.json()
|
|
assert all(m["provider"] == "openai" for m in filtered_data["models"])
|
|
|
|
|
|
def test_get_provider_config(mlflow_client_with_secrets):
|
|
import requests
|
|
|
|
base_url = mlflow_client_with_secrets._tracking_client.tracking_uri
|
|
|
|
# Test simple provider (openai) - should have single api_key auth mode
|
|
response = requests.get(
|
|
f"{base_url}/ajax-api/3.0/mlflow/gateway/provider-config",
|
|
params={"provider": "openai"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "auth_modes" in data
|
|
assert "default_mode" in data
|
|
assert data["default_mode"] == "api_key"
|
|
assert len(data["auth_modes"]) >= 1
|
|
api_key_mode = data["auth_modes"][0]
|
|
assert api_key_mode["mode"] == "api_key"
|
|
|
|
# Test multi-mode provider (bedrock) - should have multiple auth modes
|
|
response = requests.get(
|
|
f"{base_url}/ajax-api/3.0/mlflow/gateway/provider-config",
|
|
params={"provider": "bedrock"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "auth_modes" in data
|
|
assert data["default_mode"] == "api_key"
|
|
assert len(data["auth_modes"]) >= 2 # api_key, access_keys, iam_role
|
|
|
|
# Check access_keys mode structure
|
|
access_keys_mode = next(m for m in data["auth_modes"] if m["mode"] == "access_keys")
|
|
assert len(access_keys_mode["secret_fields"]) == 2 # access_key_id, secret_access_key
|
|
assert any(f["name"] == "aws_secret_access_key" for f in access_keys_mode["secret_fields"])
|
|
assert any(f["name"] == "aws_region_name" for f in access_keys_mode["config_fields"])
|
|
|
|
# Check iam_role mode exists
|
|
iam_role_mode = next(m for m in data["auth_modes"] if m["mode"] == "iam_role")
|
|
assert any(f["name"] == "aws_role_name" for f in iam_role_mode["config_fields"])
|
|
|
|
# Unknown providers get a generic fallback
|
|
response = requests.get(
|
|
f"{base_url}/ajax-api/3.0/mlflow/gateway/provider-config",
|
|
params={"provider": "unknown_provider"},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["default_mode"] == "api_key"
|
|
assert data["auth_modes"][0]["mode"] == "api_key"
|
|
assert data["auth_modes"][0]["config_fields"][0]["name"] == "api_base"
|
|
|
|
# Missing provider parameter returns 400
|
|
response = requests.get(f"{base_url}/ajax-api/3.0/mlflow/gateway/provider-config")
|
|
assert response.status_code == 400
|
|
|
|
|
|
def test_get_secrets_config_with_custom_passphrase(mlflow_client_with_secrets):
|
|
base_url = mlflow_client_with_secrets._tracking_client.tracking_uri
|
|
|
|
response = requests.get(f"{base_url}/ajax-api/3.0/mlflow/gateway/secrets/config")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["secrets_available"] is True
|
|
assert data["using_default_passphrase"] is False
|
|
|
|
|
|
def test_get_secrets_config_with_default_passphrase(tmp_path: Path, monkeypatch):
|
|
from tests.tracking.integration_test_utils import ServerThread, get_safe_port
|
|
|
|
monkeypatch.delenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", raising=False)
|
|
|
|
backend_uri = f"sqlite:///{tmp_path}/mlflow.db"
|
|
artifact_uri = (tmp_path / "artifacts").as_uri()
|
|
|
|
store = SqlAlchemyStore(backend_uri, artifact_uri)
|
|
store.engine.dispose()
|
|
|
|
handlers._tracking_store = None
|
|
handlers._model_registry_store = None
|
|
initialize_backend_stores(backend_uri, default_artifact_root=artifact_uri)
|
|
|
|
with ServerThread(app, get_safe_port()) as url:
|
|
response = requests.get(f"{url}/ajax-api/3.0/mlflow/gateway/secrets/config")
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert data["secrets_available"] is True
|
|
assert data["using_default_passphrase"] is True
|
|
|
|
|
|
def test_endpoint_with_orphaned_model_definition(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="orphan-test-key",
|
|
secret_value={"api_key": "sk-orphan-test"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="orphan-model-def",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="orphan-test-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
assert len(endpoint.model_mappings) == 1
|
|
assert endpoint.model_mappings[0].model_definition.secret_id == secret.secret_id
|
|
assert endpoint.model_mappings[0].model_definition.secret_name == "orphan-test-key"
|
|
|
|
store.delete_gateway_secret(secret.secret_id)
|
|
|
|
fetched_endpoint = store.get_gateway_endpoint(endpoint.endpoint_id)
|
|
assert len(fetched_endpoint.model_mappings) == 1
|
|
assert fetched_endpoint.model_mappings[0].model_definition.secret_id is None
|
|
assert fetched_endpoint.model_mappings[0].model_definition.secret_name is None
|
|
|
|
|
|
def test_update_model_definition_provider(mlflow_client_with_secrets):
|
|
store = mlflow_client_with_secrets._tracking_client.store
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="provider-update-secret",
|
|
secret_value={"api_key": "sk-provider-test"},
|
|
provider="openai",
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="provider-update-model-def",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
assert model_def.provider == "openai"
|
|
assert model_def.model_name == "gpt-4"
|
|
|
|
updated = store.update_gateway_model_definition(
|
|
model_definition_id=model_def.model_definition_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-5-haiku-latest",
|
|
)
|
|
|
|
assert updated.provider == "anthropic"
|
|
assert updated.model_name == "claude-3-5-haiku-latest"
|
|
|
|
fetched = store.get_gateway_model_definition(model_def.model_definition_id)
|
|
assert fetched.provider == "anthropic"
|
|
assert fetched.model_name == "claude-3-5-haiku-latest"
|
|
|
|
store.delete_gateway_model_definition(model_def.model_definition_id)
|
|
store.delete_gateway_secret(secret.secret_id)
|
|
|
|
|
|
def test_create_issue_with_all_fields(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
|
|
mlflow.set_tracking_uri(mlflow_client.tracking_uri)
|
|
experiment_id = mlflow_client.create_experiment("Issue Test")
|
|
run = mlflow_client.create_run(experiment_id)
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": "High latency issue",
|
|
"description": "API calls are taking too long",
|
|
"status": IssueStatus.PENDING.value,
|
|
"source_run_id": run.info.run_id,
|
|
"root_causes": ["Database query inefficiency", "Network latency"],
|
|
"severity": IssueSeverity.HIGH.value,
|
|
"created_by": "test-user",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "issue" in data
|
|
issue = data["issue"]
|
|
assert issue["experiment_id"] == experiment_id
|
|
assert issue["name"] == "High latency issue"
|
|
assert issue["description"] == "API calls are taking too long"
|
|
assert issue["status"] == IssueStatus.PENDING.value
|
|
assert issue["source_run_id"] == run.info.run_id
|
|
assert issue["root_causes"] == ["Database query inefficiency", "Network latency"]
|
|
assert issue["severity"] == IssueSeverity.HIGH.value
|
|
assert issue["created_by"] == "test-user"
|
|
assert "issue_id" in issue
|
|
assert "created_timestamp" in issue
|
|
assert "last_updated_timestamp" in issue
|
|
|
|
|
|
def test_create_issue_minimal_fields(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
experiment_id = mlflow_client.create_experiment("Issue Test Minimal")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": "Test issue",
|
|
"description": "Test description",
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
issue = data["issue"]
|
|
assert issue["experiment_id"] == experiment_id
|
|
assert issue["name"] == "Test issue"
|
|
assert issue["description"] == "Test description"
|
|
assert issue["status"] == IssueStatus.PENDING.value
|
|
assert "issue_id" in issue
|
|
|
|
|
|
def test_create_issue_with_required_fields(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
experiment_id = mlflow_client.create_experiment("Issue Test Required Fields")
|
|
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": "Issue with required fields only",
|
|
"description": "Testing issue creation with required fields",
|
|
"status": IssueStatus.RESOLVED.value,
|
|
},
|
|
)
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
issue = data["issue"]
|
|
assert issue["status"] == IssueStatus.RESOLVED.value
|
|
assert "issue_id" in issue
|
|
assert "created_timestamp" in issue
|
|
assert "last_updated_timestamp" in issue
|
|
|
|
|
|
def test_create_issue_invalid_experiment(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": "999999",
|
|
"name": "Test issue",
|
|
"description": "Test description",
|
|
},
|
|
)
|
|
assert response.status_code == 404
|
|
data = response.json()
|
|
assert data["error_code"] == "RESOURCE_DOES_NOT_EXIST"
|
|
|
|
|
|
def test_get_issue(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
experiment_id = mlflow_client.create_experiment("Issue Test Get")
|
|
|
|
create_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": "Test issue",
|
|
"description": "Test description",
|
|
"severity": IssueSeverity.MEDIUM.value,
|
|
},
|
|
)
|
|
issue_id = create_response.json()["issue"]["issue_id"]
|
|
|
|
get_response = requests.get(f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/{issue_id}")
|
|
assert get_response.status_code == 200
|
|
data = get_response.json()
|
|
issue = data["issue"]
|
|
assert issue["issue_id"] == issue_id
|
|
assert issue["name"] == "Test issue"
|
|
assert issue["severity"] == IssueSeverity.MEDIUM.value
|
|
|
|
|
|
def test_get_issue_not_found(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
response = requests.get(f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/nonexistent-issue")
|
|
assert response.status_code == 404
|
|
data = response.json()
|
|
assert data["error_code"] == "RESOURCE_DOES_NOT_EXIST"
|
|
|
|
|
|
def test_update_issue(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
experiment_id = mlflow_client.create_experiment("Issue Test Update")
|
|
|
|
create_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": "Original name",
|
|
"description": "Original description",
|
|
"status": IssueStatus.PENDING.value,
|
|
},
|
|
)
|
|
issue_id = create_response.json()["issue"]["issue_id"]
|
|
|
|
update_response = requests.patch(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/{issue_id}",
|
|
json={
|
|
"issue_id": issue_id,
|
|
"name": "Updated name",
|
|
"description": "Updated description",
|
|
"status": IssueStatus.RESOLVED.value,
|
|
"severity": IssueSeverity.HIGH.value,
|
|
},
|
|
)
|
|
assert update_response.status_code == 200
|
|
data = update_response.json()
|
|
issue = data["issue"]
|
|
assert issue["issue_id"] == issue_id
|
|
assert issue["name"] == "Updated name"
|
|
assert issue["description"] == "Updated description"
|
|
assert issue["status"] == IssueStatus.RESOLVED.value
|
|
assert issue["severity"] == IssueSeverity.HIGH.value
|
|
|
|
|
|
def test_search_issues_no_filters(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
experiment_id = mlflow_client.create_experiment("Issue Test Search")
|
|
|
|
for i in range(3):
|
|
requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": f"Issue {i}",
|
|
"description": f"Description {i}",
|
|
},
|
|
)
|
|
|
|
search_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/search", json={}
|
|
)
|
|
assert search_response.status_code == 200
|
|
data = search_response.json()
|
|
assert "issues" in data
|
|
assert len(data["issues"]) == 3
|
|
assert {issue["name"] for issue in data["issues"]} == {"Issue 0", "Issue 1", "Issue 2"}
|
|
assert {issue["status"] for issue in data["issues"]} == {IssueStatus.PENDING.value}
|
|
|
|
|
|
def test_search_issues_by_experiment(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
exp1 = mlflow_client.create_experiment("Issue Test Search Exp1")
|
|
exp2 = mlflow_client.create_experiment("Issue Test Search Exp2")
|
|
|
|
requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": exp1,
|
|
"name": "Issue in exp1",
|
|
"description": "Description",
|
|
},
|
|
)
|
|
|
|
requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": exp2,
|
|
"name": "Issue in exp2",
|
|
"description": "Description",
|
|
},
|
|
)
|
|
|
|
search_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/search",
|
|
json={"experiment_id": exp1},
|
|
)
|
|
assert search_response.status_code == 200
|
|
data = search_response.json()
|
|
issues = data["issues"]
|
|
assert len(issues) == 1
|
|
assert issues[0]["experiment_id"] == exp1
|
|
assert issues[0]["name"] == "Issue in exp1"
|
|
|
|
|
|
def test_search_issues_by_status(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
experiment_id = mlflow_client.create_experiment("Issue Test Search Status")
|
|
|
|
requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": "Draft issue",
|
|
"description": "Description",
|
|
"status": IssueStatus.PENDING.value,
|
|
},
|
|
)
|
|
|
|
requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": "Confirmed issue",
|
|
"description": "Description",
|
|
"status": IssueStatus.RESOLVED.value,
|
|
},
|
|
)
|
|
|
|
search_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/search",
|
|
json={"experiment_id": experiment_id, "filter_string": "status = 'resolved'"},
|
|
)
|
|
assert search_response.status_code == 200
|
|
data = search_response.json()
|
|
issues = data["issues"]
|
|
assert all(issue["status"] == IssueStatus.RESOLVED.value for issue in issues)
|
|
assert any(issue["name"] == "Confirmed issue" for issue in issues)
|
|
|
|
|
|
def test_search_issues_with_pagination(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
experiment_id = mlflow_client.create_experiment("Issue Test Pagination")
|
|
|
|
for i in range(15):
|
|
requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": f"Issue {i}",
|
|
"description": f"Description {i}",
|
|
},
|
|
)
|
|
|
|
first_page = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/search",
|
|
json={"experiment_id": experiment_id, "max_results": 10},
|
|
)
|
|
assert first_page.status_code == 200
|
|
first_data = first_page.json()
|
|
assert len(first_data["issues"]) == 10
|
|
assert "next_page_token" in first_data
|
|
assert first_data["next_page_token"] != ""
|
|
|
|
second_page = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/search",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"max_results": 10,
|
|
"page_token": first_data["next_page_token"],
|
|
},
|
|
)
|
|
assert second_page.status_code == 200
|
|
second_data = second_page.json()
|
|
assert len(second_data["issues"]) == 5
|
|
assert second_data["next_page_token"] == ""
|
|
|
|
|
|
def test_search_issues_sorted_by_timestamp(mlflow_client, store_type):
|
|
if store_type == "file":
|
|
pytest.skip("Issues are only supported in SqlAlchemyStore")
|
|
experiment_id = mlflow_client.create_experiment("Issue Test Sort")
|
|
|
|
# Create issues with slight delays to ensure different timestamps
|
|
issue_ids = []
|
|
for i in range(3):
|
|
response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues",
|
|
json={
|
|
"experiment_id": experiment_id,
|
|
"name": f"Issue {i}",
|
|
"description": f"Description {i}",
|
|
},
|
|
)
|
|
issue_ids.append(response.json()["issue"]["issue_id"])
|
|
time.sleep(0.01) # Small delay to ensure different timestamps
|
|
|
|
search_response = requests.post(
|
|
f"{mlflow_client.tracking_uri}/api/3.0/mlflow/issues/search",
|
|
json={"experiment_id": experiment_id},
|
|
)
|
|
assert search_response.status_code == 200
|
|
data = search_response.json()
|
|
issues = data["issues"]
|
|
assert len(issues) == 3
|
|
# Issues should be returned (default order is by created_timestamp descending)
|
|
assert {issue["issue_id"] for issue in issues} == set(issue_ids)
|