wehub-resource-sync
3103 lines
108 KiB
Python
3103 lines
108 KiB
Python
import json
|
|
import uuid
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
from mlflow.entities import (
|
|
BudgetAction,
|
|
BudgetDuration,
|
|
BudgetDurationUnit,
|
|
BudgetTargetScope,
|
|
BudgetUnit,
|
|
FallbackConfig,
|
|
FallbackStrategy,
|
|
GatewayBudgetPolicy,
|
|
GatewayEndpoint,
|
|
GatewayEndpointBinding,
|
|
GatewayEndpointModelConfig,
|
|
GatewayEndpointModelMapping,
|
|
GatewayEndpointTag,
|
|
GatewayGuardrail,
|
|
GatewayGuardrailConfig,
|
|
GatewayModelDefinition,
|
|
GatewayModelLinkageType,
|
|
GatewaySecretInfo,
|
|
GuardrailAction,
|
|
GuardrailStage,
|
|
RoutingStrategy,
|
|
)
|
|
from mlflow.environment_variables import MLFLOW_ENABLE_WORKSPACES, MLFLOW_TRACKING_URI
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import (
|
|
INVALID_PARAMETER_VALUE,
|
|
INVALID_STATE,
|
|
RESOURCE_ALREADY_EXISTS,
|
|
RESOURCE_DOES_NOT_EXIST,
|
|
ErrorCode,
|
|
)
|
|
from mlflow.store.tracking.dbmodels.models import (
|
|
SqlExperiment,
|
|
SqlExperimentTag,
|
|
SqlGatewayBudgetPolicy,
|
|
SqlGatewayEndpoint,
|
|
SqlGatewayEndpointBinding,
|
|
SqlGatewayEndpointModelMapping,
|
|
SqlGatewayEndpointTag,
|
|
SqlGatewayGuardrail,
|
|
SqlGatewayGuardrailConfig,
|
|
SqlGatewayModelDefinition,
|
|
SqlGatewaySecret,
|
|
SqlOnlineScoringConfig,
|
|
SqlScorer,
|
|
SqlScorerVersion,
|
|
SqlSpan,
|
|
SqlSpanMetrics,
|
|
SqlTraceInfo,
|
|
SqlTraceMetadata,
|
|
)
|
|
from mlflow.store.tracking.gateway.config_resolver import (
|
|
get_endpoint_config,
|
|
get_resource_endpoint_configs,
|
|
)
|
|
from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig, GatewayModelConfig
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
from mlflow.store.tracking.sqlalchemy_workspace_store import WorkspaceAwareSqlAlchemyStore
|
|
from mlflow.tracing.constant import SpanMetricKey, TraceMetadataKey
|
|
from mlflow.utils.workspace_context import WorkspaceContext
|
|
from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME
|
|
|
|
pytestmark = pytest.mark.notrackingurimock
|
|
|
|
TEST_PASSPHRASE = "test-passphrase-for-gateway-tests"
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def set_kek_passphrase(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", TEST_PASSPHRASE)
|
|
|
|
|
|
def _cleanup_database(store: SqlAlchemyStore):
|
|
"""Clean up gateway-specific tables after each test."""
|
|
with store.ManagedSessionMaker(read_only=False) as session:
|
|
# Delete all rows in gateway tables in dependency order
|
|
for model in (
|
|
SqlGatewayGuardrailConfig,
|
|
SqlGatewayGuardrail,
|
|
SqlGatewayBudgetPolicy,
|
|
SqlGatewayEndpointTag,
|
|
SqlGatewayEndpointBinding,
|
|
SqlGatewayEndpointModelMapping,
|
|
SqlGatewayEndpoint,
|
|
SqlGatewayModelDefinition,
|
|
SqlGatewaySecret,
|
|
SqlOnlineScoringConfig,
|
|
SqlScorerVersion,
|
|
SqlScorer,
|
|
SqlSpanMetrics,
|
|
SqlSpan,
|
|
SqlTraceMetadata,
|
|
SqlTraceInfo,
|
|
SqlExperimentTag,
|
|
SqlExperiment,
|
|
):
|
|
session.query(model).delete()
|
|
|
|
# Ensure the default experiment exists in the default workspace (ID 0).
|
|
with WorkspaceContext(DEFAULT_WORKSPACE_NAME):
|
|
store._create_default_experiment(session)
|
|
|
|
|
|
@pytest.fixture(autouse=True, params=[False, True], ids=["workspace-disabled", "workspace-enabled"])
|
|
def workspaces_enabled(request, monkeypatch):
|
|
enabled = request.param
|
|
monkeypatch.setenv(MLFLOW_ENABLE_WORKSPACES.name, "true" if enabled else "false")
|
|
if enabled:
|
|
# Use a unique workspace per test to avoid name collisions on shared DBs.
|
|
workspace_name = f"gateway-test-{uuid.uuid4().hex}"
|
|
with WorkspaceContext(workspace_name):
|
|
yield enabled
|
|
else:
|
|
yield enabled
|
|
|
|
|
|
@pytest.fixture
|
|
def store(tmp_path: Path, db_uri: str, workspaces_enabled):
|
|
artifact_uri = tmp_path / "artifacts"
|
|
artifact_uri.mkdir(exist_ok=True)
|
|
store_cls = WorkspaceAwareSqlAlchemyStore if workspaces_enabled else SqlAlchemyStore
|
|
if db_uri_env := MLFLOW_TRACKING_URI.get():
|
|
s = store_cls(db_uri_env, artifact_uri.as_uri())
|
|
yield s
|
|
_cleanup_database(s)
|
|
else:
|
|
s = store_cls(db_uri, artifact_uri.as_uri())
|
|
yield s
|
|
_cleanup_database(s)
|
|
|
|
|
|
# =============================================================================
|
|
# Secret Operations
|
|
# =============================================================================
|
|
|
|
|
|
def test_create_gateway_secret(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="my-api-key",
|
|
secret_value={"api_key": "sk-test-123456"},
|
|
provider="openai",
|
|
created_by="test-user",
|
|
)
|
|
|
|
assert isinstance(secret, GatewaySecretInfo)
|
|
assert secret.secret_id.startswith("s-")
|
|
assert secret.secret_name == "my-api-key"
|
|
assert secret.provider == "openai"
|
|
assert secret.created_by == "test-user"
|
|
assert isinstance(secret.masked_values, dict)
|
|
assert "api_key" in secret.masked_values
|
|
assert "sk-test-123456" not in secret.masked_values["api_key"]
|
|
|
|
|
|
def test_create_gateway_secret_with_auth_config(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="bedrock-creds",
|
|
secret_value={"api_key": "aws-secret-key"},
|
|
provider="bedrock",
|
|
auth_config={"region": "us-east-1", "project_id": "my-project"},
|
|
)
|
|
|
|
assert secret.secret_name == "bedrock-creds"
|
|
assert secret.provider == "bedrock"
|
|
|
|
|
|
def test_create_gateway_secret_with_dict_value(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="multi-secret",
|
|
secret_value={
|
|
"aws_access_key_id": "AKIA1234567890",
|
|
"aws_secret_access_key": "secret-key-here",
|
|
},
|
|
provider="bedrock",
|
|
auth_config={"auth_mode": "access_keys", "aws_region_name": "us-west-2"},
|
|
)
|
|
|
|
assert secret.secret_name == "multi-secret"
|
|
assert secret.provider == "bedrock"
|
|
assert isinstance(secret.masked_values, dict)
|
|
assert "aws_access_key_id" in secret.masked_values
|
|
assert "aws_secret_access_key" in secret.masked_values
|
|
assert "AKIA1234567890" not in secret.masked_values["aws_access_key_id"]
|
|
assert "secret-key-here" not in secret.masked_values["aws_secret_access_key"]
|
|
|
|
|
|
def test_create_gateway_secret_duplicate_name_raises(store: SqlAlchemyStore):
|
|
store.create_gateway_secret(secret_name="duplicate-name", secret_value={"api_key": "value1"})
|
|
|
|
with pytest.raises(MlflowException, match="already exists") as exc:
|
|
store.create_gateway_secret(
|
|
secret_name="duplicate-name", secret_value={"api_key": "value2"}
|
|
)
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS)
|
|
|
|
|
|
def test_get_gateway_secret_info_by_id(store: SqlAlchemyStore):
|
|
created = store.create_gateway_secret(
|
|
secret_name="test-secret",
|
|
secret_value={"api_key": "secret-value"},
|
|
provider="anthropic",
|
|
)
|
|
|
|
retrieved = store.get_secret_info(secret_id=created.secret_id)
|
|
|
|
assert retrieved.secret_id == created.secret_id
|
|
assert retrieved.secret_name == "test-secret"
|
|
assert retrieved.provider == "anthropic"
|
|
|
|
|
|
def test_get_gateway_secret_info_by_name(store: SqlAlchemyStore):
|
|
created = store.create_gateway_secret(
|
|
secret_name="named-secret",
|
|
secret_value={"api_key": "secret-value"},
|
|
)
|
|
|
|
retrieved = store.get_secret_info(secret_name="named-secret")
|
|
|
|
assert retrieved.secret_id == created.secret_id
|
|
assert retrieved.secret_name == "named-secret"
|
|
|
|
|
|
def test_get_gateway_secret_info_requires_one_of_id_or_name(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="Exactly one of") as exc:
|
|
store.get_secret_info()
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
with pytest.raises(MlflowException, match="Exactly one of") as exc:
|
|
store.get_secret_info(secret_id="id", secret_name="name")
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
|
|
def test_get_gateway_secret_info_not_found(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="not found") as exc:
|
|
store.get_secret_info(secret_id="nonexistent")
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
def test_update_gateway_secret(store: SqlAlchemyStore):
|
|
created = store.create_gateway_secret(
|
|
secret_name="rotate-me",
|
|
secret_value={"api_key": "old-value"},
|
|
)
|
|
original_updated_at = created.last_updated_at
|
|
|
|
updated = store.update_gateway_secret(
|
|
secret_id=created.secret_id,
|
|
secret_value={"api_key": "new-value"},
|
|
updated_by="rotator-user",
|
|
)
|
|
|
|
assert updated.secret_id == created.secret_id
|
|
assert updated.last_updated_by == "rotator-user"
|
|
assert updated.last_updated_at > original_updated_at
|
|
|
|
|
|
def test_update_gateway_secret_with_auth_config(store: SqlAlchemyStore):
|
|
created = store.create_gateway_secret(
|
|
secret_name="auth-update",
|
|
secret_value={"api_key": "value"},
|
|
auth_config={"region": "us-east-1"},
|
|
)
|
|
|
|
store.update_gateway_secret(
|
|
secret_id=created.secret_id,
|
|
secret_value={"api_key": "new-value"},
|
|
auth_config={"region": "eu-west-1", "new_key": "new_value"},
|
|
)
|
|
|
|
|
|
def test_update_gateway_secret_clear_auth_config(store: SqlAlchemyStore):
|
|
created = store.create_gateway_secret(
|
|
secret_name="clear-auth",
|
|
secret_value={"api_key": "value"},
|
|
auth_config={"region": "us-east-1"},
|
|
)
|
|
|
|
store.update_gateway_secret(
|
|
secret_id=created.secret_id,
|
|
secret_value={"api_key": "new-value"},
|
|
auth_config={},
|
|
)
|
|
|
|
|
|
def test_delete_gateway_secret(store: SqlAlchemyStore):
|
|
created = store.create_gateway_secret(
|
|
secret_name="to-delete", secret_value={"api_key": "value"}
|
|
)
|
|
|
|
store.delete_gateway_secret(created.secret_id)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.get_secret_info(secret_id=created.secret_id)
|
|
|
|
|
|
def test_list_gateway_secret_infos(store: SqlAlchemyStore):
|
|
s1 = store.create_gateway_secret(
|
|
secret_name="openai-1", secret_value={"api_key": "v1"}, provider="openai"
|
|
)
|
|
s2 = store.create_gateway_secret(
|
|
secret_name="openai-2", secret_value={"api_key": "v2"}, provider="openai"
|
|
)
|
|
s3 = store.create_gateway_secret(
|
|
secret_name="anthropic-1", secret_value={"api_key": "v3"}, provider="anthropic"
|
|
)
|
|
created_ids = {s1.secret_id, s2.secret_id, s3.secret_id}
|
|
|
|
all_secrets = store.list_secret_infos()
|
|
all_ids = {s.secret_id for s in all_secrets}
|
|
assert created_ids.issubset(all_ids)
|
|
|
|
openai_secrets = store.list_secret_infos(provider="openai")
|
|
openai_ids = {s.secret_id for s in openai_secrets}
|
|
assert {s1.secret_id, s2.secret_id}.issubset(openai_ids)
|
|
assert s3.secret_id not in openai_ids
|
|
assert all(s.provider == "openai" for s in openai_secrets)
|
|
|
|
|
|
def test_secret_id_and_name_are_immutable_at_database_level(store: SqlAlchemyStore):
|
|
"""
|
|
Verify that secret_id and secret_name cannot be modified at the database level.
|
|
|
|
These fields are used as AAD (Additional Authenticated Data) in AES-GCM encryption.
|
|
If they are modified, decryption will fail. A database trigger enforces this immutability
|
|
to prevent any code path from accidentally allowing mutation.
|
|
"""
|
|
from sqlalchemy import text
|
|
from sqlalchemy.exc import DatabaseError, IntegrityError, OperationalError
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="immutable-test",
|
|
secret_value={"api_key": "test-value"},
|
|
provider="openai",
|
|
)
|
|
|
|
def attempt_mutation(session):
|
|
session.execute(
|
|
text("UPDATE secrets SET secret_name = :new_name WHERE secret_id = :id"),
|
|
{"new_name": "modified-name", "id": secret.secret_id},
|
|
)
|
|
session.flush()
|
|
|
|
with store.ManagedSessionMaker(read_only=False) as session:
|
|
with pytest.raises((DatabaseError, IntegrityError, OperationalError)):
|
|
attempt_mutation(session)
|
|
|
|
retrieved = store.get_secret_info(secret_id=secret.secret_id)
|
|
assert retrieved.secret_name == "immutable-test"
|
|
|
|
|
|
# =============================================================================
|
|
# Model Definition Operations
|
|
# =============================================================================
|
|
|
|
|
|
def test_create_gateway_model_definition(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="test-key", secret_value={"api_key": "value"})
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="gpt-4-turbo",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4-turbo-preview",
|
|
created_by="test-user",
|
|
)
|
|
|
|
assert isinstance(model_def, GatewayModelDefinition)
|
|
assert model_def.model_definition_id.startswith("d-")
|
|
assert model_def.name == "gpt-4-turbo"
|
|
assert model_def.secret_id == secret.secret_id
|
|
assert model_def.secret_name == "test-key"
|
|
assert model_def.provider == "openai"
|
|
assert model_def.model_name == "gpt-4-turbo-preview"
|
|
|
|
|
|
def test_create_gateway_model_definition_duplicate_name_raises(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="dup-key", secret_value={"api_key": "value"})
|
|
|
|
store.create_gateway_model_definition(
|
|
name="duplicate-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="already exists") as exc:
|
|
store.create_gateway_model_definition(
|
|
name="duplicate-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS)
|
|
|
|
|
|
def test_create_gateway_model_definition_nonexistent_secret_raises(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="not found") as exc:
|
|
store.create_gateway_model_definition(
|
|
name="orphan-model",
|
|
secret_id="nonexistent",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
def test_get_gateway_model_definition_by_id(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="get-key", secret_value={"api_key": "value"})
|
|
created = store.create_gateway_model_definition(
|
|
name="model-by-id",
|
|
secret_id=secret.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-sonnet",
|
|
)
|
|
|
|
retrieved = store.get_gateway_model_definition(model_definition_id=created.model_definition_id)
|
|
|
|
assert retrieved.model_definition_id == created.model_definition_id
|
|
assert retrieved.name == "model-by-id"
|
|
|
|
|
|
def test_get_gateway_model_definition_by_name(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="name-key", secret_value={"api_key": "value"})
|
|
created = store.create_gateway_model_definition(
|
|
name="model-by-name",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
retrieved = store.get_gateway_model_definition(name="model-by-name")
|
|
|
|
assert retrieved.model_definition_id == created.model_definition_id
|
|
|
|
|
|
def test_get_gateway_model_definition_requires_one_of_id_or_name(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="Exactly one of") as exc:
|
|
store.get_gateway_model_definition()
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
|
|
def test_list_gateway_model_definitions(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="list-key", secret_value={"api_key": "value"})
|
|
|
|
store.create_gateway_model_definition(
|
|
name="list-model-1", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
store.create_gateway_model_definition(
|
|
name="list-model-2",
|
|
secret_id=secret.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3",
|
|
)
|
|
|
|
all_defs = store.list_gateway_model_definitions()
|
|
assert len(all_defs) >= 2
|
|
|
|
openai_defs = store.list_gateway_model_definitions(provider="openai")
|
|
assert all(d.provider == "openai" for d in openai_defs)
|
|
|
|
|
|
def test_update_gateway_model_definition(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="update-key", secret_value={"api_key": "value"}
|
|
)
|
|
created = store.create_gateway_model_definition(
|
|
name="update-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
updated = store.update_gateway_model_definition(
|
|
model_definition_id=created.model_definition_id,
|
|
model_name="gpt-4-turbo",
|
|
updated_by="updater",
|
|
)
|
|
|
|
assert updated.model_name == "gpt-4-turbo"
|
|
assert updated.last_updated_by == "updater"
|
|
|
|
|
|
def test_delete_gateway_model_definition(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="delete-key", secret_value={"api_key": "value"}
|
|
)
|
|
created = store.create_gateway_model_definition(
|
|
name="delete-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
|
|
store.delete_gateway_model_definition(created.model_definition_id)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.get_gateway_model_definition(model_definition_id=created.model_definition_id)
|
|
|
|
|
|
def test_delete_gateway_model_definition_in_use_raises(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="in-use-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="in-use-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
store.create_gateway_endpoint(
|
|
name="uses-model",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="currently in use") as exc:
|
|
store.delete_gateway_model_definition(model_def.model_definition_id)
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_STATE)
|
|
|
|
|
|
# =============================================================================
|
|
# Endpoint Operations
|
|
# =============================================================================
|
|
|
|
|
|
def test_create_gateway_endpoint(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="ep-key", secret_value={"api_key": "value"})
|
|
model_def = store.create_gateway_model_definition(
|
|
name="ep-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
|
|
# Create an experiment to link with the endpoint
|
|
experiment_id = store.create_experiment("test-experiment")
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="my-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
created_by="test-user",
|
|
usage_tracking=True,
|
|
experiment_id=experiment_id,
|
|
)
|
|
|
|
assert isinstance(endpoint, GatewayEndpoint)
|
|
assert endpoint.endpoint_id.startswith("e-")
|
|
assert endpoint.name == "my-endpoint"
|
|
assert len(endpoint.model_mappings) == 1
|
|
assert endpoint.model_mappings[0].model_definition_id == model_def.model_definition_id
|
|
assert endpoint.usage_tracking is True
|
|
assert endpoint.experiment_id == experiment_id
|
|
|
|
|
|
def test_create_gateway_endpoint_auto_creates_experiment(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="auto-exp-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="auto-exp-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="auto-exp-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
|
|
assert endpoint.experiment_id is not None
|
|
|
|
experiment = store.get_experiment(endpoint.experiment_id)
|
|
assert experiment.name == "gateway/auto-exp-endpoint"
|
|
assert experiment.tags.get("mlflow.experiment.sourceType") == "GATEWAY"
|
|
assert experiment.tags.get("mlflow.experiment.sourceId") == endpoint.endpoint_id
|
|
assert experiment.tags.get("mlflow.experiment.isGateway") == "true"
|
|
|
|
|
|
def test_create_gateway_endpoint_usage_tracking_defaults_to_true(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="default-ut-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="default-ut-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
|
|
# Create endpoint without specifying usage_tracking
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="default-ut-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
assert endpoint.usage_tracking is True
|
|
# An experiment should be auto-created since usage_tracking defaults to True
|
|
assert endpoint.experiment_id is not None
|
|
|
|
|
|
def test_create_gateway_endpoint_empty_models_raises(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="at least one") as exc:
|
|
store.create_gateway_endpoint(name="empty-endpoint", model_configs=[])
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
|
|
def test_create_gateway_endpoint_nonexistent_model_raises(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="not found") as exc:
|
|
store.create_gateway_endpoint(
|
|
name="orphan-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id="nonexistent",
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
def test_get_gateway_endpoint_by_id(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="get-ep-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="get-ep-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
created = store.create_gateway_endpoint(
|
|
name="get-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
retrieved = store.get_gateway_endpoint(endpoint_id=created.endpoint_id)
|
|
|
|
assert retrieved.endpoint_id == created.endpoint_id
|
|
assert retrieved.name == "get-endpoint"
|
|
|
|
|
|
def test_get_gateway_endpoint_by_name(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="name-ep-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="name-ep-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
created = store.create_gateway_endpoint(
|
|
name="named-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
retrieved = store.get_gateway_endpoint(name="named-endpoint")
|
|
|
|
assert retrieved.endpoint_id == created.endpoint_id
|
|
|
|
|
|
def test_get_gateway_endpoint_requires_one_of_id_or_name(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="Exactly one of") as exc:
|
|
store.get_gateway_endpoint()
|
|
assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
|
|
|
|
|
|
def test_update_gateway_endpoint(store: SqlAlchemyStore):
|
|
secret1 = store.create_gateway_secret(
|
|
secret_name="upd-ep-key1", secret_value={"api_key": "value1"}
|
|
)
|
|
secret2 = store.create_gateway_secret(
|
|
secret_name="upd-ep-key2", secret_value={"api_key": "value2"}
|
|
)
|
|
secret3 = store.create_gateway_secret(
|
|
secret_name="upd-ep-key3", secret_value={"api_key": "value3"}
|
|
)
|
|
secret4 = store.create_gateway_secret(
|
|
secret_name="upd-ep-key4", secret_value={"api_key": "value4"}
|
|
)
|
|
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="upd-ep-model1",
|
|
secret_id=secret1.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="upd-ep-model2",
|
|
secret_id=secret2.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-5-sonnet-20241022",
|
|
)
|
|
model_def3 = store.create_gateway_model_definition(
|
|
name="upd-ep-model3",
|
|
secret_id=secret3.secret_id,
|
|
provider="cohere",
|
|
model_name="command-r-plus",
|
|
)
|
|
model_def4 = store.create_gateway_model_definition(
|
|
name="upd-ep-model4",
|
|
secret_id=secret4.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
# Create endpoint with model1 as PRIMARY
|
|
created = store.create_gateway_endpoint(
|
|
name="update-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
# Verify initial state
|
|
assert len(created.model_mappings) == 1
|
|
assert created.model_mappings[0].model_definition_id == model_def1.model_definition_id
|
|
assert created.model_mappings[0].linkage_type == GatewayModelLinkageType.PRIMARY
|
|
assert created.routing_strategy is None
|
|
assert created.fallback_config is None
|
|
|
|
# Test 1: Basic update - rename endpoint
|
|
renamed = store.update_gateway_endpoint(
|
|
endpoint_id=created.endpoint_id,
|
|
name="renamed-endpoint",
|
|
updated_by="updater",
|
|
)
|
|
|
|
assert renamed.name == "renamed-endpoint"
|
|
assert renamed.last_updated_by == "updater"
|
|
|
|
# Test 2: Update with routing strategy and fallback config
|
|
with_fallback = store.update_gateway_endpoint(
|
|
endpoint_id=created.endpoint_id,
|
|
routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
|
|
fallback_config=FallbackConfig(
|
|
strategy=FallbackStrategy.SEQUENTIAL,
|
|
max_attempts=2,
|
|
),
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
weight=1.0,
|
|
fallback_order=0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def3.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
weight=1.0,
|
|
fallback_order=1,
|
|
),
|
|
],
|
|
updated_by="updater2",
|
|
)
|
|
|
|
assert with_fallback.routing_strategy == RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT
|
|
assert with_fallback.fallback_config is not None
|
|
assert with_fallback.fallback_config.strategy == FallbackStrategy.SEQUENTIAL
|
|
assert with_fallback.fallback_config.max_attempts == 2
|
|
assert len(with_fallback.model_mappings) == 3
|
|
fallback_mappings = [
|
|
m
|
|
for m in with_fallback.model_mappings
|
|
if m.linkage_type == GatewayModelLinkageType.FALLBACK
|
|
]
|
|
assert len(fallback_mappings) == 2
|
|
assert with_fallback.last_updated_by == "updater2"
|
|
|
|
# Test 3: Update PRIMARY models and FALLBACK models
|
|
with_new_models = store.update_gateway_endpoint(
|
|
endpoint_id=created.endpoint_id,
|
|
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.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def4.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
weight=1.0,
|
|
fallback_order=0,
|
|
),
|
|
],
|
|
fallback_config=FallbackConfig(
|
|
strategy=FallbackStrategy.SEQUENTIAL,
|
|
max_attempts=1,
|
|
),
|
|
updated_by="updater3",
|
|
)
|
|
|
|
# Verify PRIMARY models were replaced (in model_mappings)
|
|
assert len(with_new_models.model_mappings) == 3
|
|
primary_mappings = [
|
|
m
|
|
for m in with_new_models.model_mappings
|
|
if m.linkage_type == GatewayModelLinkageType.PRIMARY
|
|
]
|
|
assert len(primary_mappings) == 2
|
|
|
|
primary_model_ids = {m.model_definition_id for m in primary_mappings}
|
|
assert primary_model_ids == {model_def2.model_definition_id, model_def3.model_definition_id}
|
|
|
|
|
|
def test_delete_gateway_endpoint(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="del-ep-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="del-ep-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
created = store.create_gateway_endpoint(
|
|
name="delete-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
store.delete_gateway_endpoint(created.endpoint_id)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.get_gateway_endpoint(endpoint_id=created.endpoint_id)
|
|
|
|
|
|
def test_list_gateway_endpoints(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="list-ep-key", secret_value={"api_key": "value"}
|
|
)
|
|
secret_fallback = store.create_gateway_secret(
|
|
secret_name="list-ep-fallback-key", secret_value={"api_key": "fallback-value"}
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="list-ep-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
model_def_fallback = store.create_gateway_model_definition(
|
|
name="list-ep-fallback-model",
|
|
secret_id=secret_fallback.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-5-sonnet-20241022",
|
|
)
|
|
|
|
# Create endpoint without fallback
|
|
ep1 = store.create_gateway_endpoint(
|
|
name="list-endpoint-1",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
# Create endpoint with fallback config and routing strategy
|
|
ep2 = store.create_gateway_endpoint(
|
|
name="list-endpoint-2",
|
|
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,
|
|
),
|
|
)
|
|
|
|
endpoints = store.list_gateway_endpoints()
|
|
assert len(endpoints) >= 2
|
|
|
|
# Find our test endpoints
|
|
found_ep1 = next((e for e in endpoints if e.endpoint_id == ep1.endpoint_id), None)
|
|
found_ep2 = next((e for e in endpoints if e.endpoint_id == ep2.endpoint_id), None)
|
|
|
|
assert found_ep1 is not None
|
|
assert found_ep2 is not None
|
|
|
|
|
|
# =============================================================================
|
|
# Model Mapping Operations
|
|
# =============================================================================
|
|
|
|
|
|
def test_attach_model_to_gateway_endpoint(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="attach-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="attach-model-1", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="attach-model-2",
|
|
secret_id=secret.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="attach-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=2.0,
|
|
),
|
|
created_by="attacher",
|
|
)
|
|
|
|
assert isinstance(mapping, GatewayEndpointModelMapping)
|
|
assert mapping.mapping_id.startswith("m-")
|
|
assert mapping.endpoint_id == endpoint.endpoint_id
|
|
assert mapping.model_definition_id == model_def2.model_definition_id
|
|
assert mapping.weight == 2.0
|
|
|
|
|
|
def test_attach_duplicate_model_raises(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="dup-attach-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="dup-attach-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="dup-attach-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="already attached") as exc:
|
|
store.attach_model_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
model_config=GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
)
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_ALREADY_EXISTS)
|
|
|
|
|
|
def test_detach_model_from_gateway_endpoint(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="detach-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="detach-model-1",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="detach-model-2",
|
|
secret_id=secret.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="detach-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
store.detach_model_from_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
model_definition_id=model_def1.model_definition_id,
|
|
)
|
|
|
|
updated_endpoint = store.get_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
assert len(updated_endpoint.model_mappings) == 1
|
|
model_def_id = updated_endpoint.model_mappings[0].model_definition_id
|
|
assert model_def_id == model_def2.model_definition_id
|
|
|
|
|
|
def test_detach_nonexistent_mapping_raises(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="no-map-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="no-map-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="no-map-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="not attached") as exc:
|
|
store.detach_model_from_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
model_definition_id="nonexistent-model",
|
|
)
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
# =============================================================================
|
|
# Binding Operations
|
|
# =============================================================================
|
|
|
|
|
|
def test_create_gateway_endpoint_binding(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="bind-key", secret_value={"api_key": "value"})
|
|
model_def = store.create_gateway_model_definition(
|
|
name="bind-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="bind-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
binding = store.create_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="job-123",
|
|
created_by="binder",
|
|
)
|
|
|
|
assert isinstance(binding, GatewayEndpointBinding)
|
|
assert binding.endpoint_id == endpoint.endpoint_id
|
|
assert binding.resource_type == "scorer"
|
|
assert binding.resource_id == "job-123"
|
|
assert binding.created_by == "binder"
|
|
|
|
|
|
def test_delete_gateway_endpoint_binding(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="del-bind-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="del-bind-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="del-bind-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
store.create_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="job-456",
|
|
)
|
|
|
|
store.delete_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="job-456",
|
|
)
|
|
|
|
bindings = store.list_endpoint_bindings(endpoint_id=endpoint.endpoint_id)
|
|
assert len(bindings) == 0
|
|
|
|
|
|
def test_list_gateway_endpoint_bindings(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="list-bind-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="list-bind-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="list-bind-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
store.create_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="job-1",
|
|
)
|
|
store.create_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="job-2",
|
|
)
|
|
|
|
bindings = store.list_endpoint_bindings(endpoint_id=endpoint.endpoint_id)
|
|
assert len(bindings) == 2
|
|
|
|
filtered = store.list_endpoint_bindings(resource_type="scorer", resource_id="job-1")
|
|
assert len(filtered) == 1
|
|
assert filtered[0].resource_id == "job-1"
|
|
|
|
|
|
# =============================================================================
|
|
# Config Resolver Operations
|
|
# =============================================================================
|
|
|
|
|
|
def test_get_resource_gateway_endpoint_configs(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="resolver-key",
|
|
secret_value={"api_key": "sk-secret-value-123"},
|
|
provider="openai",
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="resolver-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4-turbo",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="resolver-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
store.create_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="resolver-job-123",
|
|
)
|
|
|
|
configs = get_resource_endpoint_configs(
|
|
resource_type="scorer",
|
|
resource_id="resolver-job-123",
|
|
store=store,
|
|
)
|
|
|
|
assert len(configs) == 1
|
|
config = configs[0]
|
|
assert config.endpoint_id == endpoint.endpoint_id
|
|
assert config.endpoint_name == "resolver-endpoint"
|
|
assert len(config.models) == 1
|
|
|
|
model_config = config.models[0]
|
|
assert model_config.model_definition_id == model_def.model_definition_id
|
|
assert model_config.provider == "openai"
|
|
assert model_config.model_name == "gpt-4-turbo"
|
|
assert model_config.secret_value == {"api_key": "sk-secret-value-123"}
|
|
|
|
|
|
def test_get_resource_endpoint_configs_with_auth_config(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="auth-resolver-key",
|
|
secret_value={"api_key": "aws-secret"},
|
|
provider="bedrock",
|
|
auth_config={"region": "us-east-1", "profile": "default"},
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="auth-resolver-model",
|
|
secret_id=secret.secret_id,
|
|
provider="bedrock",
|
|
model_name="anthropic.claude-3",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="auth-resolver-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
store.create_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="auth-job",
|
|
)
|
|
|
|
configs = get_resource_endpoint_configs(
|
|
resource_type="scorer",
|
|
resource_id="auth-job",
|
|
store=store,
|
|
)
|
|
|
|
model_config = configs[0].models[0]
|
|
assert model_config.auth_config == {"region": "us-east-1", "profile": "default"}
|
|
|
|
|
|
def test_get_resource_endpoint_configs_with_dict_secret(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="aws-creds",
|
|
secret_value={
|
|
"aws_access_key_id": "AKIA1234567890",
|
|
"aws_secret_access_key": "secret-key-value",
|
|
},
|
|
provider="bedrock",
|
|
auth_config={"auth_mode": "access_keys", "aws_region_name": "us-west-2"},
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="aws-model",
|
|
secret_id=secret.secret_id,
|
|
provider="bedrock",
|
|
model_name="anthropic.claude-3",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="aws-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
store.create_endpoint_binding(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="aws-job",
|
|
)
|
|
|
|
configs = get_resource_endpoint_configs(
|
|
resource_type="scorer",
|
|
resource_id="aws-job",
|
|
store=store,
|
|
)
|
|
|
|
model_config = configs[0].models[0]
|
|
assert model_config.secret_value == {
|
|
"aws_access_key_id": "AKIA1234567890",
|
|
"aws_secret_access_key": "secret-key-value",
|
|
}
|
|
assert model_config.auth_config == {
|
|
"auth_mode": "access_keys",
|
|
"aws_region_name": "us-west-2",
|
|
}
|
|
|
|
|
|
def test_get_resource_endpoint_configs_no_bindings(store: SqlAlchemyStore):
|
|
configs = get_resource_endpoint_configs(
|
|
resource_type="scorer",
|
|
resource_id="nonexistent-resource",
|
|
store=store,
|
|
)
|
|
|
|
assert configs == []
|
|
|
|
|
|
def test_get_resource_endpoint_configs_multiple_endpoints(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="multi-key", secret_value={"api_key": "value"})
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="multi-model-1",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="multi-model-2",
|
|
secret_id=secret.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3",
|
|
)
|
|
endpoint1 = store.create_gateway_endpoint(
|
|
name="multi-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="multi-endpoint-2",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
store.create_endpoint_binding(
|
|
endpoint_id=endpoint1.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="multi-resource",
|
|
)
|
|
store.create_endpoint_binding(
|
|
endpoint_id=endpoint2.endpoint_id,
|
|
resource_type="scorer",
|
|
resource_id="multi-resource",
|
|
)
|
|
|
|
configs = get_resource_endpoint_configs(
|
|
resource_type="scorer",
|
|
resource_id="multi-resource",
|
|
store=store,
|
|
)
|
|
|
|
assert len(configs) == 2
|
|
endpoint_names = {c.endpoint_name for c in configs}
|
|
assert endpoint_names == {"multi-endpoint-1", "multi-endpoint-2"}
|
|
|
|
|
|
def test_get_gateway_endpoint_config(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="ep-config-key",
|
|
secret_value={"api_key": "sk-endpoint-secret-789"},
|
|
provider="openai",
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="ep-config-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4o",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="ep-config-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
config = get_endpoint_config(
|
|
endpoint_name=endpoint.name,
|
|
store=store,
|
|
)
|
|
|
|
assert config.endpoint_id == endpoint.endpoint_id
|
|
assert config.endpoint_name == "ep-config-endpoint"
|
|
assert len(config.models) == 1
|
|
|
|
model_config = config.models[0]
|
|
assert model_config.model_definition_id == model_def.model_definition_id
|
|
assert model_config.provider == "openai"
|
|
assert model_config.model_name == "gpt-4o"
|
|
assert model_config.secret_value == {"api_key": "sk-endpoint-secret-789"}
|
|
|
|
|
|
def test_get_gateway_endpoint_config_with_auth_config(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="ep-auth-key",
|
|
secret_value={"api_key": "bedrock-secret"},
|
|
provider="bedrock",
|
|
auth_config={"region": "eu-west-1", "project_id": "test-project"},
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="ep-auth-model",
|
|
secret_id=secret.secret_id,
|
|
provider="bedrock",
|
|
model_name="anthropic.claude-3-sonnet",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="ep-auth-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
config = get_endpoint_config(
|
|
endpoint_name=endpoint.name,
|
|
store=store,
|
|
)
|
|
|
|
model_config = config.models[0]
|
|
assert model_config.auth_config == {"region": "eu-west-1", "project_id": "test-project"}
|
|
|
|
|
|
def test_get_gateway_endpoint_config_multiple_models(store: SqlAlchemyStore):
|
|
secret1 = store.create_gateway_secret(
|
|
secret_name="ep-multi-key-1", secret_value={"api_key": "secret-1"}
|
|
)
|
|
secret2 = store.create_gateway_secret(
|
|
secret_name="ep-multi-key-2", secret_value={"api_key": "secret-2"}
|
|
)
|
|
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="ep-multi-model-1",
|
|
secret_id=secret1.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="ep-multi-model-2",
|
|
secret_id=secret2.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-opus",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="ep-multi-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
config = get_endpoint_config(
|
|
endpoint_name=endpoint.name,
|
|
store=store,
|
|
)
|
|
|
|
assert len(config.models) == 2
|
|
providers = {m.provider for m in config.models}
|
|
assert providers == {"openai", "anthropic"}
|
|
model_names = {m.model_name for m in config.models}
|
|
assert model_names == {"gpt-4", "claude-3-opus"}
|
|
|
|
|
|
def test_get_gateway_endpoint_config_nonexistent_endpoint_raises(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="not found") as exc:
|
|
get_endpoint_config(
|
|
endpoint_name="nonexistent-endpoint",
|
|
store=store,
|
|
)
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
def test_get_gateway_endpoint_config_experiment_id_is_string(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="exp-id-test-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="exp-id-test-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="exp-id-test-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
usage_tracking=True,
|
|
)
|
|
|
|
config = get_endpoint_config(
|
|
endpoint_name=endpoint.name,
|
|
store=store,
|
|
)
|
|
|
|
# Verify experiment_id is a string (not an integer from SQLAlchemy)
|
|
assert config.experiment_id is not None
|
|
assert isinstance(config.experiment_id, str)
|
|
|
|
|
|
# =============================================================================
|
|
# Serialization & Caching Tests
|
|
# =============================================================================
|
|
|
|
|
|
def test_gateway_model_config_roundtrip():
|
|
original = GatewayModelConfig(
|
|
model_definition_id="md-123",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
secret_value={"api_key": "sk-test"},
|
|
auth_config={"region": "us-east-1"},
|
|
weight=0.7,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
fallback_order=2,
|
|
)
|
|
serialized = json.loads(json.dumps(original.to_dict()))
|
|
restored = GatewayModelConfig.from_dict(serialized)
|
|
assert restored == original
|
|
assert isinstance(restored.linkage_type, GatewayModelLinkageType)
|
|
|
|
|
|
def test_gateway_endpoint_config_roundtrip_with_fallback():
|
|
original = GatewayEndpointConfig(
|
|
endpoint_id="e-abc",
|
|
endpoint_name="my-ep",
|
|
models=[
|
|
GatewayModelConfig(
|
|
model_definition_id="md-1",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
secret_value={"api_key": "sk-1"},
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
),
|
|
GatewayModelConfig(
|
|
model_definition_id="md-2",
|
|
provider="anthropic",
|
|
model_name="claude-3",
|
|
secret_value={"api_key": "sk-2"},
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
fallback_order=0,
|
|
),
|
|
],
|
|
routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
|
|
fallback_config=FallbackConfig(
|
|
strategy=FallbackStrategy.SEQUENTIAL,
|
|
max_attempts=3,
|
|
),
|
|
experiment_id="42",
|
|
)
|
|
serialized = json.loads(json.dumps(original.to_dict()))
|
|
restored = GatewayEndpointConfig.from_dict(serialized)
|
|
assert restored == original
|
|
assert isinstance(restored.routing_strategy, RoutingStrategy)
|
|
assert isinstance(restored.models[0].linkage_type, GatewayModelLinkageType)
|
|
assert isinstance(restored.fallback_config.strategy, FallbackStrategy)
|
|
|
|
|
|
def test_gateway_endpoint_config_roundtrip_no_fallback():
|
|
original = GatewayEndpointConfig(
|
|
endpoint_id="e-xyz",
|
|
endpoint_name="simple-ep",
|
|
models=[
|
|
GatewayModelConfig(
|
|
model_definition_id="md-1",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
secret_value={"api_key": "sk-1"},
|
|
),
|
|
],
|
|
)
|
|
serialized = json.loads(json.dumps(original.to_dict()))
|
|
restored = GatewayEndpointConfig.from_dict(serialized)
|
|
assert restored == original
|
|
|
|
|
|
def test_endpoint_config_cache_hit_skips_db(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="cache-hit-key", secret_value={"api_key": "sk-cache"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="cache-hit-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="cache-hit-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
# First call populates the cache
|
|
config1 = get_endpoint_config(endpoint_name=endpoint.name, store=store)
|
|
assert config1.endpoint_name == endpoint.name
|
|
|
|
# Second call should hit the cache and not open a new DB session
|
|
original_session_maker = store.ManagedSessionMaker
|
|
with mock.patch.object(
|
|
store, "ManagedSessionMaker", wraps=original_session_maker
|
|
) as mock_session:
|
|
config2 = get_endpoint_config(endpoint_name=endpoint.name, store=store)
|
|
mock_session.assert_not_called()
|
|
|
|
assert config2.endpoint_name == config1.endpoint_name
|
|
assert config2.models[0].secret_value == config1.models[0].secret_value
|
|
|
|
|
|
def test_endpoint_config_cache_invalidated_on_secret_update(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="cache-inv-key", secret_value={"api_key": "sk-old"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="cache-inv-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="cache-inv-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
# Populate cache
|
|
config1 = get_endpoint_config(endpoint_name=endpoint.name, store=store)
|
|
assert config1.models[0].secret_value == {"api_key": "sk-old"}
|
|
|
|
# Update secret — should invalidate cache
|
|
store.update_gateway_secret(
|
|
secret_id=secret.secret_id,
|
|
secret_value={"api_key": "sk-new"},
|
|
)
|
|
|
|
# Next call should fetch fresh data from DB
|
|
config2 = get_endpoint_config(endpoint_name=endpoint.name, store=store)
|
|
assert config2.models[0].secret_value == {"api_key": "sk-new"}
|
|
|
|
|
|
def test_endpoint_config_cache_invalidated_on_endpoint_delete(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="cache-del-key", secret_value={"api_key": "sk-del"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="cache-del-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="cache-del-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
# Populate cache
|
|
config = get_endpoint_config(endpoint_name=endpoint.name, store=store)
|
|
assert config.endpoint_id == endpoint.endpoint_id
|
|
|
|
# Delete endpoint — should invalidate cache
|
|
store.delete_gateway_endpoint(endpoint.endpoint_id)
|
|
|
|
# Next call should hit DB and raise since endpoint is gone
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
get_endpoint_config(endpoint_name=endpoint.name, store=store)
|
|
|
|
|
|
def test_gateway_endpoint_config_roundtrip_fallback_with_none_fields():
|
|
original = GatewayEndpointConfig(
|
|
endpoint_id="e-none",
|
|
endpoint_name="none-fallback-ep",
|
|
models=[
|
|
GatewayModelConfig(
|
|
model_definition_id="md-1",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
secret_value={"api_key": "sk-1"},
|
|
),
|
|
],
|
|
fallback_config=FallbackConfig(strategy=None, max_attempts=None),
|
|
)
|
|
serialized = json.loads(json.dumps(original.to_dict()))
|
|
restored = GatewayEndpointConfig.from_dict(serialized)
|
|
assert restored == original
|
|
|
|
|
|
def test_endpoint_config_cache_workspace_isolation(store: SqlAlchemyStore, workspaces_enabled):
|
|
if not workspaces_enabled:
|
|
pytest.skip("workspace isolation only relevant when workspaces enabled")
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="ws-iso-key", secret_value={"api_key": "sk-ws"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="ws-iso-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="ws-iso-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
# Populate cache in current workspace
|
|
config = get_endpoint_config(endpoint_name=endpoint.name, store=store)
|
|
assert config.endpoint_id == endpoint.endpoint_id
|
|
|
|
# Switch to a different workspace — same endpoint name should NOT be served from cache
|
|
with WorkspaceContext(f"other-workspace-{uuid.uuid4().hex}"):
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
get_endpoint_config(endpoint_name=endpoint.name, store=store)
|
|
|
|
|
|
# =============================================================================
|
|
# Endpoint Tag Operations
|
|
# =============================================================================
|
|
|
|
|
|
def test_set_gateway_endpoint_tag(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(secret_name="tag-key", secret_value={"api_key": "value"})
|
|
model_def = store.create_gateway_model_definition(
|
|
name="tag-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="tag-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
tag = GatewayEndpointTag(key="env", value="production")
|
|
store.set_gateway_endpoint_tag(endpoint.endpoint_id, tag)
|
|
|
|
retrieved = store.get_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
assert len(retrieved.tags) == 1
|
|
assert retrieved.tags[0].key == "env"
|
|
assert retrieved.tags[0].value == "production"
|
|
|
|
|
|
def test_set_gateway_endpoint_tag_update_existing(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="tag-upd-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="tag-upd-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="tag-upd-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
store.set_gateway_endpoint_tag(endpoint.endpoint_id, GatewayEndpointTag(key="env", value="dev"))
|
|
store.set_gateway_endpoint_tag(
|
|
endpoint.endpoint_id, GatewayEndpointTag(key="env", value="production")
|
|
)
|
|
|
|
retrieved = store.get_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
assert len(retrieved.tags) == 1
|
|
assert retrieved.tags[0].key == "env"
|
|
assert retrieved.tags[0].value == "production"
|
|
|
|
|
|
def test_set_multiple_endpoint_tags(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="multi-tag-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="multi-tag-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="multi-tag-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
store.set_gateway_endpoint_tag(
|
|
endpoint.endpoint_id, GatewayEndpointTag(key="env", value="production")
|
|
)
|
|
store.set_gateway_endpoint_tag(endpoint.endpoint_id, GatewayEndpointTag(key="team", value="ml"))
|
|
store.set_gateway_endpoint_tag(
|
|
endpoint.endpoint_id, GatewayEndpointTag(key="version", value="v1")
|
|
)
|
|
|
|
retrieved = store.get_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
assert len(retrieved.tags) == 3
|
|
tag_dict = {t.key: t.value for t in retrieved.tags}
|
|
assert tag_dict == {"env": "production", "team": "ml", "version": "v1"}
|
|
|
|
|
|
def test_set_gateway_endpoint_tag_nonexistent_endpoint_raises(store: SqlAlchemyStore):
|
|
tag = GatewayEndpointTag(key="env", value="production")
|
|
with pytest.raises(MlflowException, match="not found") as exc:
|
|
store.set_gateway_endpoint_tag("nonexistent-endpoint", tag)
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
def test_delete_gateway_endpoint_tag(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="del-tag-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="del-tag-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="del-tag-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
store.set_gateway_endpoint_tag(
|
|
endpoint.endpoint_id, GatewayEndpointTag(key="env", value="production")
|
|
)
|
|
store.set_gateway_endpoint_tag(endpoint.endpoint_id, GatewayEndpointTag(key="team", value="ml"))
|
|
|
|
store.delete_gateway_endpoint_tag(endpoint.endpoint_id, "env")
|
|
|
|
retrieved = store.get_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
assert len(retrieved.tags) == 1
|
|
assert retrieved.tags[0].key == "team"
|
|
|
|
|
|
def test_delete_gateway_endpoint_tag_nonexistent_endpoint_raises(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="not found") as exc:
|
|
store.delete_gateway_endpoint_tag("nonexistent-endpoint", "env")
|
|
assert exc.value.error_code == ErrorCode.Name(RESOURCE_DOES_NOT_EXIST)
|
|
|
|
|
|
def test_delete_gateway_endpoint_tag_nonexistent_key_no_op(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="del-noop-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="del-noop-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="del-noop-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
# Should not raise even if tag doesn't exist
|
|
store.delete_gateway_endpoint_tag(endpoint.endpoint_id, "nonexistent-key")
|
|
|
|
retrieved = store.get_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
assert len(retrieved.tags) == 0
|
|
|
|
|
|
def test_endpoint_tags_deleted_with_endpoint(store: SqlAlchemyStore):
|
|
secret = store.create_gateway_secret(
|
|
secret_name="cascade-tag-key", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="cascade-tag-model",
|
|
secret_id=secret.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="cascade-tag-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
store.set_gateway_endpoint_tag(
|
|
endpoint.endpoint_id, GatewayEndpointTag(key="env", value="production")
|
|
)
|
|
|
|
store.delete_gateway_endpoint(endpoint.endpoint_id)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.get_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
|
|
|
|
# =============================================================================
|
|
# Scorer-Endpoint Integration Tests
|
|
# =============================================================================
|
|
|
|
|
|
def _create_gateway_endpoint(store: SqlAlchemyStore, name: str) -> GatewayEndpoint:
|
|
"""Helper to create a gateway endpoint for scorer tests."""
|
|
secret = store.create_gateway_secret(
|
|
secret_name=f"{name}-secret", secret_value={"api_key": "value"}
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name=f"{name}-model", secret_id=secret.secret_id, provider="openai", model_name="gpt-4"
|
|
)
|
|
return store.create_gateway_endpoint(
|
|
name=name,
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
)
|
|
|
|
|
|
def test_register_scorer_resolves_endpoint_name_to_id(store: SqlAlchemyStore):
|
|
experiment_id = store.create_experiment(f"scorer-endpoint-test-{uuid.uuid4().hex}")
|
|
endpoint = _create_gateway_endpoint(store, "test-endpoint")
|
|
|
|
serialized_scorer = json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": f"gateway:/{endpoint.name}",
|
|
"instructions": "Rate the response",
|
|
}
|
|
})
|
|
|
|
scorer = store.register_scorer(experiment_id, "my-scorer", serialized_scorer)
|
|
|
|
# The returned scorer should have the endpoint name (resolved from ID)
|
|
stored_data = json.loads(scorer._serialized_scorer)
|
|
stored_model = stored_data["instructions_judge_pydantic_data"]["model"]
|
|
assert stored_model == f"gateway:/{endpoint.name}"
|
|
|
|
|
|
def test_register_scorer_with_nonexistent_endpoint_raises(store: SqlAlchemyStore):
|
|
experiment_id = store.create_experiment(f"scorer-nonexistent-endpoint-test-{uuid.uuid4().hex}")
|
|
|
|
serialized_scorer = json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": "gateway:/nonexistent-endpoint",
|
|
"instructions": "Rate the response",
|
|
}
|
|
})
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.register_scorer(experiment_id, "my-scorer", serialized_scorer)
|
|
|
|
|
|
def test_get_scorer_resolves_endpoint_id_to_name(store: SqlAlchemyStore):
|
|
experiment_id = store.create_experiment(f"get-scorer-endpoint-test-{uuid.uuid4().hex}")
|
|
endpoint = _create_gateway_endpoint(store, "get-test-endpoint")
|
|
|
|
serialized_scorer = json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": f"gateway:/{endpoint.name}",
|
|
"instructions": "Rate the response",
|
|
}
|
|
})
|
|
|
|
store.register_scorer(experiment_id, "my-scorer", serialized_scorer)
|
|
|
|
# When retrieving, the endpoint ID should be resolved back to name
|
|
retrieved = store.get_scorer(experiment_id, "my-scorer")
|
|
retrieved_data = json.loads(retrieved._serialized_scorer)
|
|
retrieved_model = retrieved_data["instructions_judge_pydantic_data"]["model"]
|
|
assert retrieved_model == f"gateway:/{endpoint.name}"
|
|
|
|
|
|
def test_get_scorer_with_deleted_endpoint_sets_model_to_null(store: SqlAlchemyStore):
|
|
experiment_id = store.create_experiment(f"deleted-endpoint-scorer-test-{uuid.uuid4().hex}")
|
|
endpoint = _create_gateway_endpoint(store, "to-delete-endpoint")
|
|
|
|
serialized_scorer = json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": f"gateway:/{endpoint.name}",
|
|
"instructions": "Rate the response",
|
|
}
|
|
})
|
|
|
|
store.register_scorer(experiment_id, "my-scorer", serialized_scorer)
|
|
|
|
# Delete the endpoint
|
|
store.delete_gateway_endpoint(endpoint.endpoint_id)
|
|
|
|
# Retrieving should set model to null
|
|
retrieved = store.get_scorer(experiment_id, "my-scorer")
|
|
retrieved_data = json.loads(retrieved._serialized_scorer)
|
|
assert retrieved_data["instructions_judge_pydantic_data"]["model"] is None
|
|
|
|
|
|
def test_list_scorers_batch_resolves_endpoint_ids(store: SqlAlchemyStore):
|
|
experiment_id = store.create_experiment(f"list-scorers-endpoint-test-{uuid.uuid4().hex}")
|
|
endpoint1 = _create_gateway_endpoint(store, "list-endpoint-1")
|
|
endpoint2 = _create_gateway_endpoint(store, "list-endpoint-2")
|
|
|
|
# Create scorers with different endpoints
|
|
store.register_scorer(
|
|
experiment_id,
|
|
"scorer-1",
|
|
json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": f"gateway:/{endpoint1.name}",
|
|
"instructions": "Rate 1",
|
|
}
|
|
}),
|
|
)
|
|
store.register_scorer(
|
|
experiment_id,
|
|
"scorer-2",
|
|
json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": f"gateway:/{endpoint2.name}",
|
|
"instructions": "Rate 2",
|
|
}
|
|
}),
|
|
)
|
|
store.register_scorer(
|
|
experiment_id,
|
|
"scorer-3",
|
|
json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": "openai:/gpt-4",
|
|
"instructions": "Rate 3",
|
|
}
|
|
}),
|
|
)
|
|
|
|
scorers = store.list_scorers(experiment_id)
|
|
assert len(scorers) == 3
|
|
|
|
# Find each scorer and verify model resolution
|
|
scorer_map = {s.scorer_name: s for s in scorers}
|
|
|
|
data1 = json.loads(scorer_map["scorer-1"]._serialized_scorer)
|
|
assert data1["instructions_judge_pydantic_data"]["model"] == f"gateway:/{endpoint1.name}"
|
|
|
|
data2 = json.loads(scorer_map["scorer-2"]._serialized_scorer)
|
|
assert data2["instructions_judge_pydantic_data"]["model"] == f"gateway:/{endpoint2.name}"
|
|
|
|
data3 = json.loads(scorer_map["scorer-3"]._serialized_scorer)
|
|
assert data3["instructions_judge_pydantic_data"]["model"] == "openai:/gpt-4"
|
|
|
|
|
|
# =============================================================================
|
|
# Fallback Routing Tests
|
|
# =============================================================================
|
|
|
|
|
|
def test_create_gateway_endpoint_with_fallback_routing(store: SqlAlchemyStore):
|
|
# Create secrets and model definitions
|
|
secret1 = store.create_gateway_secret(
|
|
secret_name="fallback-key-1", secret_value={"api_key": "sk-model1"}
|
|
)
|
|
secret2 = store.create_gateway_secret(
|
|
secret_name="fallback-key-2", secret_value={"api_key": "sk-model2"}
|
|
)
|
|
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="fallback-model-1",
|
|
secret_id=secret1.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="fallback-model-2",
|
|
secret_id=secret2.secret_id,
|
|
provider="anthropic",
|
|
model_name="claude-3-5-sonnet-20241022",
|
|
)
|
|
|
|
# Create a third model for PRIMARY
|
|
secret3 = store.create_gateway_secret(
|
|
secret_name="primary-key", secret_value={"api_key": "sk-primary"}
|
|
)
|
|
model_def3 = store.create_gateway_model_definition(
|
|
name="primary-model",
|
|
secret_id=secret3.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
# Create endpoint with fallback configuration
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="fallback-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def3.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
weight=1.0,
|
|
fallback_order=0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.FALLBACK,
|
|
weight=1.0,
|
|
fallback_order=1,
|
|
),
|
|
],
|
|
created_by="test-user",
|
|
routing_strategy=None, # Fallback is independent of routing strategy
|
|
fallback_config=FallbackConfig(
|
|
strategy=FallbackStrategy.SEQUENTIAL,
|
|
max_attempts=2,
|
|
),
|
|
)
|
|
|
|
# Verify endpoint was created with fallback config
|
|
assert isinstance(endpoint, GatewayEndpoint)
|
|
assert endpoint.endpoint_id.startswith("e-")
|
|
assert endpoint.name == "fallback-endpoint"
|
|
assert endpoint.routing_strategy is None # No routing strategy needed for fallback
|
|
assert endpoint.fallback_config is not None
|
|
assert isinstance(endpoint.fallback_config, FallbackConfig)
|
|
assert endpoint.fallback_config.strategy == FallbackStrategy.SEQUENTIAL
|
|
assert endpoint.fallback_config.max_attempts == 2
|
|
assert len(endpoint.model_mappings) == 3
|
|
fallback_ids = [
|
|
m.model_definition_id
|
|
for m in endpoint.model_mappings
|
|
if m.linkage_type == GatewayModelLinkageType.FALLBACK
|
|
]
|
|
assert model_def1.model_definition_id in fallback_ids
|
|
assert model_def2.model_definition_id in fallback_ids
|
|
|
|
# Verify linkage types and fallback_order
|
|
# PRIMARY models are in endpoint.model_mappings
|
|
assert len(endpoint.model_mappings) == 3
|
|
primary_mapping = next(
|
|
m for m in endpoint.model_mappings if m.linkage_type == GatewayModelLinkageType.PRIMARY
|
|
)
|
|
assert primary_mapping.model_definition_id == model_def3.model_definition_id
|
|
|
|
|
|
def test_create_gateway_endpoint_with_traffic_split(store: SqlAlchemyStore):
|
|
secret1 = store.create_gateway_secret(
|
|
secret_name="traffic-split-key-1", secret_value={"api_key": "sk-test1"}
|
|
)
|
|
secret2 = store.create_gateway_secret(
|
|
secret_name="traffic-split-key-2", secret_value={"api_key": "sk-test2"}
|
|
)
|
|
|
|
model_def1 = store.create_gateway_model_definition(
|
|
name="traffic-split-model-1",
|
|
secret_id=secret1.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
)
|
|
model_def2 = store.create_gateway_model_definition(
|
|
name="traffic-split-model-2",
|
|
secret_id=secret2.secret_id,
|
|
provider="openai",
|
|
model_name="gpt-3.5-turbo",
|
|
)
|
|
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="traffic-split-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def1.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def2.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=1.0,
|
|
),
|
|
],
|
|
routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
|
|
)
|
|
|
|
assert endpoint.routing_strategy == RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT
|
|
assert len(endpoint.model_mappings) == 2
|
|
# All should be PRIMARY linkages for traffic split
|
|
for mapping in endpoint.model_mappings:
|
|
assert mapping.linkage_type == GatewayModelLinkageType.PRIMARY
|
|
|
|
|
|
# =============================================================================
|
|
# Budget Policy Operations
|
|
# =============================================================================
|
|
|
|
|
|
def test_create_budget_policy(store: SqlAlchemyStore):
|
|
policy = store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=100.0,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.MONTHS, value=1),
|
|
target_scope=BudgetTargetScope.GLOBAL,
|
|
budget_action=BudgetAction.ALERT,
|
|
created_by="admin",
|
|
)
|
|
assert isinstance(policy, GatewayBudgetPolicy)
|
|
assert policy.budget_policy_id.startswith("bp-")
|
|
assert policy.budget_unit == BudgetUnit.USD
|
|
assert policy.budget_amount == 100.0
|
|
assert policy.duration.unit == BudgetDurationUnit.MONTHS
|
|
assert policy.duration.value == 1
|
|
assert policy.target_scope == BudgetTargetScope.GLOBAL
|
|
assert policy.budget_action == BudgetAction.ALERT
|
|
assert policy.created_by == "admin"
|
|
assert policy.created_at > 0
|
|
assert policy.last_updated_at > 0
|
|
|
|
|
|
def test_get_budget_policy_by_id(store: SqlAlchemyStore):
|
|
created = store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=75.0,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.HOURS, value=24),
|
|
target_scope=BudgetTargetScope.WORKSPACE,
|
|
budget_action=BudgetAction.REJECT,
|
|
)
|
|
fetched = store.get_budget_policy(budget_policy_id=created.budget_policy_id)
|
|
assert fetched.budget_policy_id == created.budget_policy_id
|
|
assert fetched.budget_unit == BudgetUnit.USD
|
|
assert fetched.budget_amount == 75.0
|
|
assert fetched.duration.unit == BudgetDurationUnit.HOURS
|
|
assert fetched.duration.value == 24
|
|
assert fetched.target_scope == BudgetTargetScope.WORKSPACE
|
|
assert fetched.budget_action == BudgetAction.REJECT
|
|
|
|
|
|
def test_get_budget_policy_not_found_raises(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="BudgetPolicy"):
|
|
store.get_budget_policy(budget_policy_id="bp-nonexistent")
|
|
|
|
|
|
def test_update_budget_policy(store: SqlAlchemyStore):
|
|
created = store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=100.0,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.MONTHS, value=1),
|
|
target_scope=BudgetTargetScope.GLOBAL,
|
|
budget_action=BudgetAction.ALERT,
|
|
)
|
|
updated = store.update_budget_policy(
|
|
budget_policy_id=created.budget_policy_id,
|
|
budget_amount=200.0,
|
|
budget_action=BudgetAction.REJECT,
|
|
updated_by="editor",
|
|
)
|
|
assert updated.budget_amount == 200.0
|
|
assert updated.budget_action == BudgetAction.REJECT
|
|
assert updated.last_updated_by == "editor"
|
|
assert updated.last_updated_at >= created.last_updated_at
|
|
# Unchanged fields should remain
|
|
assert updated.duration.unit == BudgetDurationUnit.MONTHS
|
|
assert updated.duration.value == 1
|
|
assert updated.target_scope == BudgetTargetScope.GLOBAL
|
|
|
|
|
|
def test_update_budget_policy_not_found_raises(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="BudgetPolicy"):
|
|
store.update_budget_policy(
|
|
budget_policy_id="bp-nonexistent",
|
|
budget_amount=999.0,
|
|
)
|
|
|
|
|
|
def test_delete_budget_policy(store: SqlAlchemyStore):
|
|
created = store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=10.0,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.DAYS, value=1),
|
|
target_scope=BudgetTargetScope.GLOBAL,
|
|
budget_action=BudgetAction.ALERT,
|
|
)
|
|
store.delete_budget_policy(created.budget_policy_id)
|
|
|
|
with pytest.raises(MlflowException, match="BudgetPolicy"):
|
|
store.get_budget_policy(budget_policy_id=created.budget_policy_id)
|
|
|
|
|
|
def test_delete_budget_policy_not_found_raises(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="BudgetPolicy"):
|
|
store.delete_budget_policy("bp-nonexistent")
|
|
|
|
|
|
def test_list_budget_policies(store: SqlAlchemyStore):
|
|
store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=100.0,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.MONTHS, value=1),
|
|
target_scope=BudgetTargetScope.GLOBAL,
|
|
budget_action=BudgetAction.ALERT,
|
|
)
|
|
store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=50.0,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.DAYS, value=7),
|
|
target_scope=BudgetTargetScope.WORKSPACE,
|
|
budget_action=BudgetAction.REJECT,
|
|
)
|
|
|
|
all_policies = store.list_budget_policies()
|
|
assert len(all_policies) == 2
|
|
assert all_policies.token is None
|
|
|
|
|
|
def test_list_budget_policies_empty(store: SqlAlchemyStore):
|
|
policies = store.list_budget_policies()
|
|
assert policies == []
|
|
assert policies.token is None
|
|
|
|
|
|
def test_list_budget_policies_pagination(store: SqlAlchemyStore):
|
|
for i in range(3):
|
|
store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=100.0 + i,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.MONTHS, value=1),
|
|
target_scope=BudgetTargetScope.GLOBAL,
|
|
budget_action=BudgetAction.ALERT,
|
|
)
|
|
|
|
page1 = store.list_budget_policies(max_results=2)
|
|
assert len(page1) == 2
|
|
assert page1.token is not None
|
|
|
|
page2 = store.list_budget_policies(max_results=2, page_token=page1.token)
|
|
assert len(page2) == 1
|
|
assert page2.token is None
|
|
|
|
all_ids = [p.budget_policy_id for p in page1] + [p.budget_policy_id for p in page2]
|
|
assert len(set(all_ids)) == 3
|
|
|
|
|
|
def test_list_budget_policies_invalid_max_results(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="max_results"):
|
|
store.list_budget_policies(max_results=0)
|
|
|
|
|
|
def test_create_budget_policy_all_duration_units(store: SqlAlchemyStore):
|
|
for du in BudgetDurationUnit:
|
|
policy = store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=100.0,
|
|
duration=BudgetDuration(unit=du, value=1),
|
|
target_scope=BudgetTargetScope.GLOBAL,
|
|
budget_action=BudgetAction.ALERT,
|
|
)
|
|
assert policy.duration.unit == du
|
|
|
|
|
|
def test_create_budget_policy_all_budget_actions(store: SqlAlchemyStore):
|
|
for action in BudgetAction:
|
|
policy = store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=100.0,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.MONTHS, value=1),
|
|
target_scope=BudgetTargetScope.GLOBAL,
|
|
budget_action=action,
|
|
)
|
|
assert policy.budget_action == action
|
|
|
|
|
|
# =============================================================================
|
|
# Guardrail Tests
|
|
# =============================================================================
|
|
|
|
|
|
def _create_scorer(store: SqlAlchemyStore, endpoint_name: str | None = None):
|
|
"""Helper to create a scorer for guardrail tests."""
|
|
name_suffix = uuid.uuid4().hex[:8]
|
|
ep_name = endpoint_name or f"guardrail-ep-{name_suffix}"
|
|
endpoint = _create_gateway_endpoint(store, ep_name)
|
|
experiment_id = store.create_experiment(f"guardrail-scorer-exp-{name_suffix}")
|
|
|
|
serialized_scorer = json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": f"gateway:/{endpoint.name}",
|
|
"instructions": "Is this input safe?",
|
|
}
|
|
})
|
|
|
|
return store.register_scorer(experiment_id, f"safety-judge-{name_suffix}", serialized_scorer)
|
|
|
|
|
|
def _create_scorer_versions(store: SqlAlchemyStore):
|
|
"""Helper to create multiple versions of a scorer for guardrail tests."""
|
|
name_suffix = uuid.uuid4().hex[:8]
|
|
endpoint = _create_gateway_endpoint(store, f"guardrail-ep-{name_suffix}")
|
|
experiment_id = store.create_experiment(f"guardrail-scorer-exp-{name_suffix}")
|
|
scorer_name = f"safety-judge-{name_suffix}"
|
|
|
|
serialized_scorer_v1 = json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": f"gateway:/{endpoint.name}",
|
|
"instructions": "Is this input safe?",
|
|
}
|
|
})
|
|
serialized_scorer_v2 = json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": f"gateway:/{endpoint.name}",
|
|
"instructions": "Is this input still safe?",
|
|
}
|
|
})
|
|
|
|
scorer_v1 = store.register_scorer(experiment_id, scorer_name, serialized_scorer_v1)
|
|
scorer_v2 = store.register_scorer(experiment_id, scorer_name, serialized_scorer_v2)
|
|
return scorer_v1, scorer_v2
|
|
|
|
|
|
def test_create_gateway_guardrail(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
created_by="test-user",
|
|
)
|
|
|
|
assert isinstance(guardrail, GatewayGuardrail)
|
|
assert guardrail.guardrail_id.startswith("gr-")
|
|
assert guardrail.name == "test-guardrail"
|
|
assert guardrail.scorer.scorer_id == scorer.scorer_id
|
|
assert guardrail.scorer.scorer_version == scorer.scorer_version
|
|
assert guardrail.stage == GuardrailStage.BEFORE
|
|
assert guardrail.action == GuardrailAction.VALIDATION
|
|
assert guardrail.created_by == "test-user"
|
|
assert guardrail.last_updated_by == "test-user"
|
|
assert guardrail.created_at > 0
|
|
assert guardrail.last_updated_at > 0
|
|
|
|
|
|
def test_create_gateway_guardrail_after_sanitization(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.AFTER,
|
|
action=GuardrailAction.SANITIZATION,
|
|
)
|
|
|
|
assert guardrail.stage == GuardrailStage.AFTER
|
|
assert guardrail.action == GuardrailAction.SANITIZATION
|
|
|
|
|
|
def test_create_gateway_guardrail_rejects_unknown_scorer(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="Scorer with ID 'missing-scorer' not found"):
|
|
store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id="missing-scorer",
|
|
scorer_version=1,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
|
|
def test_create_gateway_guardrail_rejects_unknown_scorer_version(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=rf"Scorer with ID '{scorer.scorer_id}' and version 999 not found",
|
|
):
|
|
store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=999,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
|
|
def test_get_gateway_guardrail(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
|
|
created = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
fetched = store.get_gateway_guardrail(created.guardrail_id)
|
|
|
|
assert fetched.guardrail_id == created.guardrail_id
|
|
assert fetched.scorer.scorer_id == scorer.scorer_id
|
|
assert fetched.scorer.scorer_version == scorer.scorer_version
|
|
assert fetched.stage == GuardrailStage.BEFORE
|
|
assert fetched.action == GuardrailAction.VALIDATION
|
|
|
|
|
|
def test_get_gateway_guardrail_not_found(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.get_gateway_guardrail("gr-nonexistent")
|
|
|
|
|
|
def test_delete_gateway_guardrail(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
store.delete_gateway_guardrail(guardrail.guardrail_id)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.get_gateway_guardrail(guardrail.guardrail_id)
|
|
|
|
|
|
def test_delete_gateway_guardrail_not_found(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.delete_gateway_guardrail("gr-nonexistent")
|
|
|
|
|
|
def test_list_gateway_guardrails(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
|
|
for stage in GuardrailStage:
|
|
store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=stage,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
result = store.list_gateway_guardrails()
|
|
assert len(result) == 2
|
|
assert result.token is None
|
|
|
|
|
|
def test_list_gateway_guardrails_empty(store: SqlAlchemyStore):
|
|
result = store.list_gateway_guardrails()
|
|
assert result == []
|
|
assert result.token is None
|
|
|
|
|
|
def test_list_gateway_guardrails_pagination(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
|
|
for _ in range(3):
|
|
store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
page1 = store.list_gateway_guardrails(max_results=2)
|
|
assert len(page1) == 2
|
|
assert page1.token is not None
|
|
|
|
page2 = store.list_gateway_guardrails(max_results=2, page_token=page1.token)
|
|
assert len(page2) == 1
|
|
assert page2.token is None
|
|
|
|
all_ids = sorted([g.guardrail_id for g in page1] + [g.guardrail_id for g in page2])
|
|
assert len(all_ids) == 3
|
|
assert len(set(all_ids)) == 3
|
|
|
|
|
|
def test_add_guardrail_to_endpoint(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-endpoint-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
config = store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
execution_order=5,
|
|
created_by="test-user",
|
|
)
|
|
|
|
assert isinstance(config, GatewayGuardrailConfig)
|
|
assert config.endpoint_id == endpoint.endpoint_id
|
|
assert config.guardrail_id == guardrail.guardrail_id
|
|
assert config.execution_order == 5
|
|
assert config.created_by == "test-user"
|
|
|
|
|
|
def test_add_guardrail_to_endpoint_null_execution_order(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-null-order-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
config = store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
)
|
|
|
|
assert config.execution_order is None
|
|
|
|
|
|
def test_add_guardrail_to_endpoint_duplicate_raises(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-dup-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="already added"):
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
)
|
|
|
|
|
|
def test_update_endpoint_guardrail_config(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-update-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
execution_order=5,
|
|
)
|
|
|
|
updated = store.update_endpoint_guardrail_config(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
execution_order=10,
|
|
)
|
|
|
|
assert updated.execution_order == 10
|
|
assert updated.guardrail_id == guardrail.guardrail_id
|
|
|
|
|
|
def test_update_endpoint_guardrail_config_not_found(store: SqlAlchemyStore):
|
|
endpoint = _create_gateway_endpoint(store, f"gr-upd-nf-{uuid.uuid4().hex[:8]}")
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.update_endpoint_guardrail_config(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id="gr-nonexistent",
|
|
execution_order=1,
|
|
)
|
|
|
|
|
|
def test_add_guardrail_to_nonexistent_endpoint_raises(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id="e-nonexistent",
|
|
guardrail_id=guardrail.guardrail_id,
|
|
)
|
|
|
|
|
|
def test_add_nonexistent_guardrail_to_endpoint_raises(store: SqlAlchemyStore):
|
|
endpoint = _create_gateway_endpoint(store, f"gr-noguard-{uuid.uuid4().hex[:8]}")
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id="gr-nonexistent",
|
|
)
|
|
|
|
|
|
def test_remove_guardrail_from_endpoint(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-remove-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
)
|
|
|
|
store.remove_guardrail_from_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
)
|
|
|
|
configs = store.list_endpoint_guardrail_configs(endpoint.endpoint_id)
|
|
assert configs == []
|
|
|
|
|
|
def test_remove_guardrail_from_endpoint_not_found(store: SqlAlchemyStore):
|
|
endpoint = _create_gateway_endpoint(store, f"gr-removenf-{uuid.uuid4().hex[:8]}")
|
|
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.remove_guardrail_from_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id="gr-nonexistent",
|
|
)
|
|
|
|
|
|
def test_list_endpoint_guardrail_configs(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-list-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrails = [
|
|
store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
for _ in range(3)
|
|
]
|
|
|
|
configs = [
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=g.guardrail_id,
|
|
execution_order=(i + 1) * 10,
|
|
)
|
|
for i, g in enumerate(guardrails)
|
|
]
|
|
|
|
configs = store.list_endpoint_guardrail_configs(endpoint.endpoint_id)
|
|
assert len(configs) == 3
|
|
assert [c.execution_order for c in configs] == [10, 20, 30]
|
|
assert [c.guardrail_id for c in configs] == [g.guardrail_id for g in guardrails]
|
|
for config, guardrail in zip(configs, guardrails):
|
|
assert config.guardrail is not None
|
|
assert config.guardrail.guardrail_id == guardrail.guardrail_id
|
|
assert config.guardrail.stage == GuardrailStage.BEFORE
|
|
|
|
|
|
def test_list_endpoint_guardrail_configs_null_order_sorts_last(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-null-last-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrail_with_order = store.create_gateway_guardrail(
|
|
name="test-guardrail-with-order",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
guardrail_null_order = store.create_gateway_guardrail(
|
|
name="test-guardrail-null-order",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail_null_order.guardrail_id,
|
|
)
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail_with_order.guardrail_id,
|
|
execution_order=5,
|
|
)
|
|
|
|
configs = store.list_endpoint_guardrail_configs(endpoint.endpoint_id)
|
|
assert len(configs) == 2
|
|
assert configs[0].execution_order == 5
|
|
assert configs[0].guardrail_id == guardrail_with_order.guardrail_id
|
|
assert configs[1].execution_order is None
|
|
assert configs[1].guardrail_id == guardrail_null_order.guardrail_id
|
|
|
|
|
|
def test_list_endpoint_guardrail_configs_empty(store: SqlAlchemyStore):
|
|
endpoint = _create_gateway_endpoint(store, f"gr-empty-{uuid.uuid4().hex[:8]}")
|
|
configs = store.list_endpoint_guardrail_configs(endpoint.endpoint_id)
|
|
assert configs == []
|
|
|
|
|
|
def test_list_endpoint_guardrail_configs_nonexistent_endpoint(store: SqlAlchemyStore):
|
|
with pytest.raises(MlflowException, match="not found"):
|
|
store.list_endpoint_guardrail_configs("e-nonexistent")
|
|
|
|
|
|
def test_delete_guardrail_cascades_to_configs(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-cascade-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
)
|
|
|
|
store.delete_gateway_guardrail(guardrail.guardrail_id)
|
|
|
|
configs = store.list_endpoint_guardrail_configs(endpoint.endpoint_id)
|
|
assert configs == []
|
|
|
|
|
|
def test_delete_endpoint_cascades_to_guardrail_configs(store: SqlAlchemyStore):
|
|
scorer = _create_scorer(store)
|
|
endpoint = _create_gateway_endpoint(store, f"gr-ep-cascade-{uuid.uuid4().hex[:8]}")
|
|
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="test-guardrail",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
)
|
|
|
|
store.add_guardrail_to_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
)
|
|
|
|
store.delete_gateway_endpoint(endpoint.endpoint_id)
|
|
|
|
# Guardrail itself should still exist
|
|
fetched = store.get_gateway_guardrail(guardrail.guardrail_id)
|
|
assert fetched.guardrail_id == guardrail.guardrail_id
|
|
|
|
|
|
# =============================================================================
|
|
# sum_gateway_trace_cost Tests
|
|
# =============================================================================
|
|
|
|
|
|
def _insert_trace_with_cost(
|
|
session,
|
|
experiment_id,
|
|
trace_id,
|
|
timestamp_ms,
|
|
span_costs,
|
|
is_gateway=True,
|
|
endpoint_id="ep-test",
|
|
):
|
|
trace = SqlTraceInfo(
|
|
request_id=trace_id,
|
|
experiment_id=experiment_id,
|
|
timestamp_ms=timestamp_ms,
|
|
execution_time_ms=100,
|
|
status="OK",
|
|
)
|
|
session.add(trace)
|
|
session.flush()
|
|
|
|
if is_gateway:
|
|
metadata = SqlTraceMetadata(
|
|
request_id=trace_id,
|
|
key=TraceMetadataKey.GATEWAY_ENDPOINT_ID,
|
|
value=endpoint_id,
|
|
)
|
|
session.add(metadata)
|
|
|
|
for span_id, cost in span_costs:
|
|
span = SqlSpan(
|
|
trace_id=trace_id,
|
|
experiment_id=experiment_id,
|
|
span_id=span_id,
|
|
status="OK",
|
|
start_time_unix_nano=timestamp_ms * 1_000_000,
|
|
end_time_unix_nano=(timestamp_ms + 100) * 1_000_000,
|
|
content="{}",
|
|
)
|
|
session.add(span)
|
|
session.flush()
|
|
|
|
metric = SqlSpanMetrics(
|
|
trace_id=trace_id,
|
|
span_id=span_id,
|
|
key=SpanMetricKey.TOTAL_COST,
|
|
value=cost,
|
|
)
|
|
session.add(metric)
|
|
|
|
session.flush()
|
|
|
|
|
|
def test_sum_gateway_trace_cost_basic(store: SqlAlchemyStore):
|
|
exp = store.create_experiment("cost-test-basic")
|
|
exp_id = int(exp)
|
|
|
|
with store.ManagedSessionMaker(read_only=False) as session:
|
|
_insert_trace_with_cost(session, exp_id, "t1", 1000, [("s1", 0.05), ("s2", 0.03)])
|
|
_insert_trace_with_cost(session, exp_id, "t2", 2000, [("s1", 0.10)])
|
|
|
|
total = store.sum_gateway_trace_cost(start_time_ms=0, end_time_ms=5000)
|
|
assert abs(total - 0.18) < 1e-9
|
|
|
|
|
|
def test_sum_gateway_trace_cost_excludes_non_gateway(store: SqlAlchemyStore):
|
|
exp = store.create_experiment("cost-test-non-gw")
|
|
exp_id = int(exp)
|
|
|
|
with store.ManagedSessionMaker(read_only=False) as session:
|
|
_insert_trace_with_cost(session, exp_id, "gw1", 1000, [("s1", 0.10)], is_gateway=True)
|
|
_insert_trace_with_cost(session, exp_id, "nongw1", 1000, [("s1", 0.50)], is_gateway=False)
|
|
|
|
total = store.sum_gateway_trace_cost(start_time_ms=0, end_time_ms=5000)
|
|
assert abs(total - 0.10) < 1e-9
|
|
|
|
|
|
def test_sum_gateway_trace_cost_time_window(store: SqlAlchemyStore):
|
|
exp = store.create_experiment("cost-test-window")
|
|
exp_id = int(exp)
|
|
|
|
with store.ManagedSessionMaker(read_only=False) as session:
|
|
_insert_trace_with_cost(session, exp_id, "early", 500, [("s1", 0.01)])
|
|
_insert_trace_with_cost(session, exp_id, "in-window", 1500, [("s1", 0.05)])
|
|
_insert_trace_with_cost(session, exp_id, "late", 3000, [("s1", 0.99)])
|
|
|
|
# Only in-window trace should be included
|
|
total = store.sum_gateway_trace_cost(start_time_ms=1000, end_time_ms=2000)
|
|
assert abs(total - 0.05) < 1e-9
|
|
|
|
|
|
def test_sum_gateway_trace_cost_workspace_filter(store: SqlAlchemyStore):
|
|
with store.ManagedSessionMaker(read_only=False) as session:
|
|
# Create two experiments in different workspaces
|
|
exp_ws_a = SqlExperiment(
|
|
name=f"cost-ws-a-{uuid.uuid4().hex}",
|
|
artifact_location="/tmp/a",
|
|
lifecycle_stage="active",
|
|
workspace="workspace-a",
|
|
)
|
|
exp_ws_b = SqlExperiment(
|
|
name=f"cost-ws-b-{uuid.uuid4().hex}",
|
|
artifact_location="/tmp/b",
|
|
lifecycle_stage="active",
|
|
workspace="workspace-b",
|
|
)
|
|
session.add_all([exp_ws_a, exp_ws_b])
|
|
session.flush()
|
|
|
|
_insert_trace_with_cost(session, exp_ws_a.experiment_id, "t-a", 1000, [("s1", 0.10)])
|
|
_insert_trace_with_cost(session, exp_ws_b.experiment_id, "t-b", 1000, [("s1", 0.25)])
|
|
|
|
# Filtering by workspace-a should only include 0.10
|
|
total_a = store.sum_gateway_trace_cost(
|
|
start_time_ms=0, end_time_ms=5000, workspace="workspace-a"
|
|
)
|
|
assert abs(total_a - 0.10) < 1e-9
|
|
|
|
# Filtering by workspace-b should only include 0.25
|
|
total_b = store.sum_gateway_trace_cost(
|
|
start_time_ms=0, end_time_ms=5000, workspace="workspace-b"
|
|
)
|
|
assert abs(total_b - 0.25) < 1e-9
|
|
|
|
# No workspace filter should include both
|
|
total_all = store.sum_gateway_trace_cost(start_time_ms=0, end_time_ms=5000)
|
|
assert abs(total_all - 0.35) < 1e-9
|
|
|
|
|
|
def test_sum_gateway_trace_cost_empty(store: SqlAlchemyStore):
|
|
total = store.sum_gateway_trace_cost(start_time_ms=0, end_time_ms=5000)
|
|
assert total == 0.0
|