From 8df0cb9296bc0ee3e93b19416f71c1e5e32e798d Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 28 Oct 2026 15:05:50 -0300 Subject: [PATCH 1/2] extend get_models for custom async clients Closes #1048 Signed-off-by: Andre Manoel --- .../data_designer/interface/data_designer.py | 42 +++- .../tests/interface/test_data_designer.py | 15 -- .../interface/test_data_designer_models.py | 220 ++++++++++++++++++ 3 files changed, 255 insertions(+), 22 deletions(-) create mode 100644 packages/data-designer/tests/interface/test_data_designer_models.py diff --git a/packages/data-designer/src/data_designer/interface/data_designer.py b/packages/data-designer/src/data_designer/interface/data_designer.py index f773686d0..2078293eb 100644 --- a/packages/data-designer/src/data_designer/interface/data_designer.py +++ b/packages/data-designer/src/data_designer/interface/data_designer.py @@ -27,6 +27,7 @@ from data_designer.config.models import ( ModelConfig, ModelProvider, + load_model_configs, ) from data_designer.config.preview_results import PreviewResults from data_designer.config.run_config import JinjaRenderingEngine, RunConfig @@ -44,6 +45,7 @@ from data_designer.engine.model_provider import ModelProviderRegistry, resolve_model_provider_registry from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode from data_designer.engine.models.errors import RETRYABLE_MODEL_ERRORS +from data_designer.engine.models.factory import create_model_registry from data_designer.engine.readiness import run_readiness_check from data_designer.engine.resources.person_reader import ( PersonReader, @@ -647,7 +649,13 @@ def set_run_config(self, run_config: RunConfig) -> None: self._run_config = run_config self._request_admission = self._create_request_admission_controller() - def get_models(self, model_aliases: list[str]) -> dict[str, ModelFacade]: + def get_models( + self, + model_aliases: list[str], + *, + model_configs: list[ModelConfig] | str | Path | None = None, + client_concurrency_mode: ClientConcurrencyMode = ClientConcurrencyMode.SYNC, + ) -> dict[str, ModelFacade]: """Get a dict of ModelFacade instances for custom column development. Use this to experiment with custom column generator functions outside of @@ -656,17 +664,37 @@ def get_models(self, model_aliases: list[str]) -> dict[str, ModelFacade]: Args: model_aliases: List of model aliases to include in the dict. + model_configs: Model configurations, a YAML string, or a config file path. + If None, uses the default model configurations. When provided, aliases + are resolved only from these configurations. + client_concurrency_mode: Sync clients by default. Use + `ClientConcurrencyMode.ASYNC` for async facade methods. Returns: Dict mapping alias to ModelFacade instance. + + Notes: + Creates no dataset artifacts and requires no managed person assets. + Call `close()` on each sync facade or await `aclose()` on each async + facade when finished. """ - config_builder = DataDesignerConfigBuilder() - resource_provider = self._create_resource_provider( - "dev", - config_builder, - client_concurrency_mode=ClientConcurrencyMode.SYNC, + client_concurrency_mode = ClientConcurrencyMode(client_concurrency_mode) + resolved_configs = get_default_model_configs() if model_configs is None else load_model_configs(model_configs) + if len({config.alias for config in resolved_configs}) != len(resolved_configs): + raise ValueError("Model config aliases must be unique.") + + model_registry = create_model_registry( + model_configs=resolved_configs, + secret_resolver=self._secret_resolver, + model_provider_registry=self._model_provider_registry, + client_concurrency_mode=client_concurrency_mode, + run_config=self._run_config, + request_admission=self._request_admission, + request_event_sink=self._open_telemetry if self._run_config.otel_metrics_port is not None else None, ) - return {alias: resource_provider.model_registry.get_model(model_alias=alias) for alias in model_aliases} + for alias in model_aliases: + model_registry.get_model_config(model_alias=alias) + return {alias: model_registry.get_model(model_alias=alias) for alias in model_aliases} def _resolve_model_providers(self, model_providers: list[ModelProvider] | None) -> list[ModelProvider]: if model_providers is None: diff --git a/packages/data-designer/tests/interface/test_data_designer.py b/packages/data-designer/tests/interface/test_data_designer.py index 8348a3333..b0af59e1c 100644 --- a/packages/data-designer/tests/interface/test_data_designer.py +++ b/packages/data-designer/tests/interface/test_data_designer.py @@ -44,7 +44,6 @@ HuggingFaceSeedSource, ) from data_designer.config.seed_source_dataframe import DataFrameSeedSource -from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode from data_designer.engine.models.errors import ( RETRYABLE_MODEL_ERRORS, ModelAPIConnectionError, @@ -594,20 +593,6 @@ def test_run_config_keeps_error_rate_when_disabled(stub_artifact_path, stub_mode assert data_designer.run_config.effective_shutdown_error_rate == 1.0 -def test_get_models_uses_sync_clients(stub_artifact_path, stub_model_providers): - data_designer = DataDesigner(artifact_path=stub_artifact_path, model_providers=stub_model_providers) - - with patch.object(data_designer, "_create_resource_provider") as mock_resource_provider_method: - mock_resource_provider = MagicMock() - mock_resource_provider.model_registry.get_model.return_value = MagicMock() - mock_resource_provider_method.return_value = mock_resource_provider - - data_designer.get_models(["stub-model"]) - - _, kwargs = mock_resource_provider_method.call_args - assert kwargs["client_concurrency_mode"] == ClientConcurrencyMode.SYNC - - def test_create_forwards_on_batch_complete_callback( stub_artifact_path: Path, stub_model_providers: list[ModelProvider], diff --git a/packages/data-designer/tests/interface/test_data_designer_models.py b/packages/data-designer/tests/interface/test_data_designer_models.py new file mode 100644 index 000000000..ad97b7cec --- /dev/null +++ b/packages/data-designer/tests/interface/test_data_designer_models.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +import json +from collections.abc import Iterator +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import httpx +import pytest +from pydantic import ValidationError +from pytest_httpx import HTTPXMock + +import data_designer.interface.data_designer as dd_mod +from data_designer.config.errors import InvalidConfigError +from data_designer.config.models import ModelConfig, ModelProvider +from data_designer.config.run_config import RequestAdmissionTuningConfig, RunConfig +from data_designer.engine.models.clients.adapters.http_model_client import ClientConcurrencyMode +from data_designer.engine.models.utils import ChatMessage +from data_designer.engine.secret_resolver import EnvironmentResolver +from data_designer.interface.data_designer import DataDesigner + + +@pytest.fixture +def model_designer( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + stub_model_providers: list[ModelProvider], +) -> Iterator[DataDesigner]: + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("PROVIDER_1_API_KEY", "resolved-test-key") + with patch.object(dd_mod, "create_person_reader", side_effect=AssertionError("Person reader is not needed")): + yield DataDesigner( + model_providers=stub_model_providers, + secret_resolver=EnvironmentResolver(), + managed_assets_path=tmp_path / "missing-assets", + ) + assert not (tmp_path / "artifacts").exists() + assert not (tmp_path / "missing-assets").exists() + + +@pytest.fixture +def custom_model_configs(stub_model_configs: list[ModelConfig]) -> list[ModelConfig]: + return [stub_model_configs[0].model_copy(update={"alias": alias}) for alias in ("judge", "patient")] + + +@pytest.fixture +def completion_payload() -> dict[str, Any]: + return { + "id": "test-completion", + "model": "stub-model", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + } + + +@pytest.mark.parametrize("config_source", ["defaults", "list", "yaml", "path"]) +def test_get_models_sync( + model_designer: DataDesigner, + custom_model_configs: list[ModelConfig], + stub_model_configs: list[ModelConfig], + completion_payload: dict[str, Any], + httpx_mock: HTTPXMock, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + config_source: str, +) -> None: + monkeypatch.setattr(dd_mod, "get_default_model_configs", lambda: stub_model_configs) + configs: list[ModelConfig] | str | Path | None = custom_model_configs + if config_source == "defaults": + configs = None + elif config_source in {"yaml", "path"}: + configs = json.dumps({"model_configs": [config.model_dump(mode="json") for config in custom_model_configs]}) + if config_source == "path": + config_path = tmp_path / "models.yaml" + config_path.write_text(configs) + configs = config_path + + aliases = ["stub-model"] if configs is None else ["patient", "judge", "patient"] + models = model_designer.get_models(aliases, model_configs=configs) + assert list(models) == list(dict.fromkeys(aliases)) + httpx_mock.add_response(json=completion_payload, is_reusable=True) + try: + for alias, model in models.items(): + assert model.model_alias == alias + assert model.model_name == "stub-model" + assert model.completion([ChatMessage.as_user("test")]).message.content == "hello" + assert model.usage_stats.token_usage.total_tokens == 5 + assert model.usage_stats.request_usage.successful_requests == 1 + finally: + for model in models.values(): + model.close() + + for request in httpx_mock.get_requests(): + assert str(request.url) == "https://api.provider-1.com/v1/chat/completions" + assert request.headers["Authorization"] == "Bearer resolved-test-key" + assert json.loads(request.content)["model"] == "stub-model" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_defaults", [False, True]) +async def test_get_models_async_shares_admission_and_preserves_lifetime( + model_designer: DataDesigner, + custom_model_configs: list[ModelConfig], + completion_payload: dict[str, Any], + httpx_mock: HTTPXMock, + monkeypatch: pytest.MonkeyPatch, + use_defaults: bool, +) -> None: + for config in custom_model_configs: + config.inference_parameters.max_parallel_requests = 1 + model_designer.set_run_config(RunConfig(request_admission=RequestAdmissionTuningConfig(startup_ramp_seconds=0))) + monkeypatch.setattr(dd_mod, "get_default_model_configs", lambda: custom_model_configs) + configs = None if use_defaults else custom_model_configs + judge = model_designer.get_models( + ["judge"], model_configs=configs, client_concurrency_mode=ClientConcurrencyMode.ASYNC + )["judge"] + patient = model_designer.get_models( + ["patient"], model_configs=configs, client_concurrency_mode=ClientConcurrencyMode.ASYNC + )["patient"] + in_flight = peak = 0 + + async def respond(request: httpx.Request) -> httpx.Response: + nonlocal in_flight, peak + assert request.headers["Authorization"] == "Bearer resolved-test-key" + in_flight += 1 + peak = max(peak, in_flight) + await asyncio.sleep(0.01) + in_flight -= 1 + return httpx.Response(200, json=completion_payload) + + httpx_mock.add_callback(respond, is_reusable=True) + messages = [ChatMessage.as_user("test")] + try: + responses = await asyncio.gather(judge.acompletion(messages), patient.acompletion(messages)) + assert [response.message.content for response in responses] == ["hello", "hello"] + assert peak == 1 + await judge.aclose() + assert (await patient.acompletion(messages)).message.content == "hello" + assert patient.usage_stats.token_usage.total_tokens == 10 + assert patient.usage_stats.request_usage.successful_requests == 2 + finally: + await judge.aclose() + await patient.aclose() + + +def test_get_models_retries( + model_designer: DataDesigner, + custom_model_configs: list[ModelConfig], + completion_payload: dict[str, Any], + httpx_mock: HTTPXMock, +) -> None: + httpx_mock.add_response(status_code=503, headers={"Retry-After": "0"}, json={"error": {"message": "retry"}}) + httpx_mock.add_response(json=completion_payload) + model = model_designer.get_models(["judge"], model_configs=custom_model_configs)["judge"] + try: + assert model.completion([ChatMessage.as_user("test")]).message.content == "hello" + assert len(httpx_mock.get_requests()) == 2 + assert model.usage_stats.request_usage.successful_requests == 1 + finally: + model.close() + + +@pytest.mark.parametrize( + "config_source,aliases", + [ + ("defaults", ["stub-model", "missing"]), + ("custom", ["judge", "missing"]), + ("custom", ["judge", "stub-model"]), + ("empty", ["stub-model"]), + ], +) +def test_get_models_validates_aliases_before_creating_clients( + model_designer: DataDesigner, + custom_model_configs: list[ModelConfig], + stub_model_configs: list[ModelConfig], + monkeypatch: pytest.MonkeyPatch, + config_source: str, + aliases: list[str], +) -> None: + monkeypatch.setattr(dd_mod, "get_default_model_configs", lambda: stub_model_configs) + configs = {"defaults": None, "custom": custom_model_configs, "empty": []}[config_source] + with patch("data_designer.engine.models.clients.factory.create_model_client") as create_client: + with pytest.raises(ValueError, match="No model config with alias"): + model_designer.get_models(aliases, model_configs=configs) + create_client.assert_not_called() + + +@pytest.mark.parametrize( + "configs,error,match", + [ + ("{}", InvalidConfigError, "model_configs"), + ('{"model_configs": [{}]}', ValidationError, "alias"), + ], +) +def test_get_models_validates_config( + model_designer: DataDesigner, configs: str, error: type[Exception], match: str +) -> None: + with pytest.raises(error, match=match): + model_designer.get_models(["judge"], model_configs=configs) + + +def test_get_models_rejects_duplicate_configs( + model_designer: DataDesigner, custom_model_configs: list[ModelConfig] +) -> None: + with pytest.raises(ValueError, match="aliases must be unique"): + model_designer.get_models(["judge"], model_configs=custom_model_configs * 2) + + +def test_get_models_rejects_invalid_concurrency_mode(model_designer: DataDesigner) -> None: + with pytest.raises(ValueError, match="not a valid ClientConcurrencyMode"): + model_designer.get_models([], client_concurrency_mode="invalid") + + +def test_get_models_empty_aliases(model_designer: DataDesigner) -> None: + assert model_designer.get_models([], model_configs=[]) == {} From 071805f88cf33958f451515b041d6f09a1b39103 Mon Sep 17 00:00:00 2001 From: Andre Manoel Date: Wed, 28 Oct 2026 16:53:12 -0300 Subject: [PATCH 2/2] protect standalone model admission limits Signed-off-by: Andre Manoel --- .../data_designer/engine/models/factory.py | 3 + .../data_designer/engine/models/registry.py | 16 ++++- .../models/request_admission/controller.py | 39 ++++++++++ .../request_admission/test_controller.py | 71 ++++++++++++++++++- .../data_designer/interface/data_designer.py | 6 ++ .../interface/test_data_designer_models.py | 65 +++++++++++++++++ 6 files changed, 196 insertions(+), 4 deletions(-) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/factory.py b/packages/data-designer-engine/src/data_designer/engine/models/factory.py index 6104ca00d..1d0e6285f 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/factory.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/factory.py @@ -29,6 +29,7 @@ def create_model_registry( run_config: RunConfig | None = None, request_admission: AdaptiveRequestAdmissionController | None = None, request_event_sink: RequestAdmissionEventSink | None = None, + freeze_request_limits: bool = False, ) -> ModelRegistry: """Factory function for creating a ModelRegistry instance. @@ -50,6 +51,7 @@ def create_model_registry( request_admission: Optional shared request-admission controller. When omitted, a new controller is created from ``run_config``. request_event_sink: Optional direct sink for request-admission and model-request events. + freeze_request_limits: Reject changes to registered alias limits on the shared controller. Returns: A configured ModelRegistry instance. @@ -91,6 +93,7 @@ def model_facade_factory( model_facade_factory=model_facade_factory, request_admission=request_admission, retry_config=RetryConfig(), + freeze_request_limits=freeze_request_limits, ) diff --git a/packages/data-designer-engine/src/data_designer/engine/models/registry.py b/packages/data-designer-engine/src/data_designer/engine/models/registry.py index a732a8c43..235f45993 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/registry.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/registry.py @@ -66,12 +66,14 @@ def __init__( model_facade_factory: ModelFacadeFactory | None = None, request_admission: AdaptiveRequestAdmissionController | None = None, retry_config: RetryConfig | None = None, + freeze_request_limits: bool = False, ) -> None: self._secret_resolver = secret_resolver self._model_provider_registry = model_provider_registry self._model_facade_factory = model_facade_factory self._request_admission = request_admission self._retry_config = retry_config + self._freeze_request_limits = freeze_request_limits self._model_configs: dict[str, ModelConfig] = {} self._models: dict[str, ModelFacade] = {} self._set_model_configs(model_configs) @@ -109,14 +111,23 @@ def get_model(self, *, model_alias: str) -> ModelFacade: # Lazy initialization: only create model facade when first requested if model_alias not in self._models: - self._models[model_alias] = self._get_model(self._model_configs[model_alias]) + self._models[model_alias] = self._get_model(self.get_model_config(model_alias=model_alias)) return self._models[model_alias] def get_model_config(self, *, model_alias: str) -> ModelConfig: if model_alias not in self._model_configs: raise ValueError(f"No model config with alias {model_alias!r} found!") - return self._model_configs[model_alias] + model_config = self._model_configs[model_alias] + if self._request_admission is not None: + self._request_admission.validate_registration( + provider_name=model_config.provider, + model_id=model_config.model, + alias=model_config.alias, + max_parallel_requests=model_config.inference_parameters.max_parallel_requests, + freeze=self._freeze_request_limits, + ) + return model_config def get_model_usage_stats(self, total_time_elapsed: float) -> dict[str, dict[str, Any]]: return { @@ -375,5 +386,6 @@ def _get_model(self, model_config: ModelConfig) -> ModelFacade: model_id=model_config.model, alias=model_config.alias, max_parallel_requests=model_config.inference_parameters.max_parallel_requests, + freeze=self._freeze_request_limits, ) return facade diff --git a/packages/data-designer-engine/src/data_designer/engine/models/request_admission/controller.py b/packages/data-designer-engine/src/data_designer/engine/models/request_admission/controller.py index a2047dad3..b7fb3841f 100644 --- a/packages/data-designer-engine/src/data_designer/engine/models/request_admission/controller.py +++ b/packages/data-designer-engine/src/data_designer/engine/models/request_admission/controller.py @@ -100,6 +100,7 @@ def pressure(self) -> RequestPressureSnapshotProvider: ... @dataclass class _GlobalCapState: limits_by_alias: dict[str, int] = field(default_factory=dict) + frozen_aliases: set[str] = field(default_factory=set) effective_max: int = 0 def register_alias(self, alias: str, max_parallel: int) -> None: @@ -140,6 +141,21 @@ def pressure(self) -> RequestPressureSnapshotProvider: def config(self) -> RequestAdmissionConfig: return self._config + def validate_registration( + self, + *, + provider_name: str, + model_id: str, + alias: str, + max_parallel_requests: int, + freeze: bool = False, + ) -> None: + """Check alias-limit compatibility without changing shared admission state.""" + with self._lock: + self._validate_registration_locked( + ProviderModelKey(provider_name, model_id), alias, max_parallel_requests, freeze + ) + def register( self, *, @@ -147,13 +163,18 @@ def register( model_id: str, alias: str, max_parallel_requests: int, + freeze: bool = False, ) -> None: + """Register an alias's limit, optionally preventing later changes within this controller.""" events: list[RequestAdmissionEvent] = [] with self._lock: key = ProviderModelKey(provider_name, model_id) + self._validate_registration_locked(key, alias, max_parallel_requests, freeze) cap = self._global_caps.setdefault(key, _GlobalCapState()) before = cap.effective_max cap.register_alias(alias, max_parallel_requests) + if freeze: + cap.frozen_aliases.add(alias) self._sequence += 1 for resource, state in self._domains.items(): if resource.provider_model_key == key: @@ -429,6 +450,24 @@ def global_snapshots(self) -> Mapping[ProviderModelKey, ProviderModelPressureSna now = time.monotonic() return {key: self._global_snapshot_locked(key, now) for key in self._global_caps} + def _validate_registration_locked( + self, key: ProviderModelKey, alias: str, max_parallel_requests: int, freeze: bool + ) -> None: + cap = self._global_caps.get(key) + if cap is None: + return + registered_limit = cap.limits_by_alias.get(alias) + if ( + registered_limit is not None + and registered_limit != max(1, max_parallel_requests) + and (freeze or alias in cap.frozen_aliases) + ): + raise ValueError( + f"Model alias {alias!r} for provider {key.provider_name!r}, model {key.model_id!r} is already registered " + f"with max_parallel_requests={registered_limit}; cannot use {max_parallel_requests}. " + "Use matching limits, reset via DataDesigner.set_run_config(), or use a new DataDesigner instance." + ) + def _queued_waiter_ahead_locked(self, item: RequestAdmissionItem, now: float) -> bool: if not self._queue.has_waiters: return False diff --git a/packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py b/packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py index fb4dbef88..d5f369390 100644 --- a/packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py +++ b/packages/data-designer-engine/tests/engine/models/request_admission/test_controller.py @@ -37,9 +37,13 @@ def _item(domain: RequestDomain = RequestDomain.CHAT, timeout: float | None = No ) -def _controller(cap: int = 2, config: RequestAdmissionConfig | None = None) -> AdaptiveRequestAdmissionController: +def _controller( + cap: int = 2, config: RequestAdmissionConfig | None = None, *, freeze: bool = False +) -> AdaptiveRequestAdmissionController: controller = AdaptiveRequestAdmissionController(config) - controller.register(provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=cap) + controller.register( + provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=cap, freeze=freeze + ) return controller @@ -79,6 +83,69 @@ def test_request_admission_enforces_provider_model_aggregate_cap() -> None: assert denied.reason == "no_capacity" +@pytest.mark.parametrize("initial_cap,new_cap", [(1, 10), (10, 1)]) +@pytest.mark.parametrize("method", ["register", "validate_registration"]) +def test_request_admission_rejects_conflicting_alias_cap(initial_cap: int, new_cap: int, method: str) -> None: + controller = _controller(cap=initial_cap, freeze=True) + + with pytest.raises(ValueError, match="already registered with max_parallel_requests"): + getattr(controller, method)( + provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=new_cap + ) + + snapshot = controller.global_snapshot("nvidia", "nemotron") + assert snapshot.static_cap == initial_cap + assert snapshot.raw_caps == {"default": initial_cap} + + +def test_request_admission_accepts_identical_alias_cap() -> None: + controller = _controller(cap=2, freeze=True) + controller.register(provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=2) + + lease = controller.try_acquire(_item()) + assert isinstance(lease, RequestAdmissionLease) + assert controller.global_snapshot("nvidia", "nemotron").static_cap == 2 + controller.release(lease, RequestReleaseOutcome(kind="success")) + + +@pytest.mark.parametrize("provider,model", [("other-provider", "nemotron"), ("nvidia", "other-model")]) +def test_request_admission_alias_caps_are_independent(provider: str, model: str) -> None: + controller = _controller(cap=1, freeze=True) + controller.register(provider_name=provider, model_id=model, alias="default", max_parallel_requests=10, freeze=True) + + assert controller.global_snapshot("nvidia", "nemotron").static_cap == 1 + assert controller.global_snapshot(provider, model).static_cap == 10 + + +@pytest.mark.parametrize("initial_cap,new_cap", [(1, 10), (10, 1)]) +def test_request_admission_unfrozen_alias_can_be_reconfigured(initial_cap: int, new_cap: int) -> None: + controller = _controller(cap=initial_cap) + controller.register(provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=new_cap) + + assert controller.global_snapshot("nvidia", "nemotron").static_cap == new_cap + + +def test_request_admission_rejects_freezing_conflicting_alias_cap() -> None: + controller = _controller(cap=1) + + with pytest.raises(ValueError, match="already registered with max_parallel_requests"): + controller.register( + provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=10, freeze=True + ) + + controller.register(provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=10) + assert controller.global_snapshot("nvidia", "nemotron").static_cap == 10 + + +def test_request_admission_validation_does_not_register_alias() -> None: + controller = AdaptiveRequestAdmissionController() + controller.validate_registration( + provider_name="nvidia", model_id="nemotron", alias="default", max_parallel_requests=1, freeze=True + ) + + assert controller.global_snapshots() == {} + + def test_request_admission_duplicate_release_does_not_corrupt_counts() -> None: controller = _controller(cap=1) item = _item() diff --git a/packages/data-designer/src/data_designer/interface/data_designer.py b/packages/data-designer/src/data_designer/interface/data_designer.py index 2078293eb..bc012afd1 100644 --- a/packages/data-designer/src/data_designer/interface/data_designer.py +++ b/packages/data-designer/src/data_designer/interface/data_designer.py @@ -677,6 +677,11 @@ def get_models( Creates no dataset artifacts and requires no managed person assets. Call `close()` on each sync facade or await `aclose()` on each async facade when finished. + + Request limits are shared with `create()`, `preview()`, and other calls on this instance. + Reusing an alias for the same provider/model with a different + `max_parallel_requests` raises ValueError, even after closing its facade. + Call `set_run_config()` or use a new instance to reset shared limits. """ client_concurrency_mode = ClientConcurrencyMode(client_concurrency_mode) resolved_configs = get_default_model_configs() if model_configs is None else load_model_configs(model_configs) @@ -691,6 +696,7 @@ def get_models( run_config=self._run_config, request_admission=self._request_admission, request_event_sink=self._open_telemetry if self._run_config.otel_metrics_port is not None else None, + freeze_request_limits=True, ) for alias in model_aliases: model_registry.get_model_config(model_alias=alias) diff --git a/packages/data-designer/tests/interface/test_data_designer_models.py b/packages/data-designer/tests/interface/test_data_designer_models.py index ad97b7cec..3cbfeedd8 100644 --- a/packages/data-designer/tests/interface/test_data_designer_models.py +++ b/packages/data-designer/tests/interface/test_data_designer_models.py @@ -148,6 +148,71 @@ async def respond(request: httpx.Request) -> httpx.Response: await patient.aclose() +@pytest.mark.asyncio +@pytest.mark.parametrize("initial_cap,new_cap", [(1, 10), (10, 1), (1, 1)]) +async def test_get_models_preserves_registered_alias_cap( + model_designer: DataDesigner, + custom_model_configs: list[ModelConfig], + completion_payload: dict[str, Any], + httpx_mock: HTTPXMock, + initial_cap: int, + new_cap: int, +) -> None: + model_designer.set_run_config( + RunConfig(request_admission=RequestAdmissionTuningConfig(startup_ramp_seconds=0, successes_until_increase=1)) + ) + config = custom_model_configs[0].model_copy(deep=True) + config.inference_parameters.max_parallel_requests = initial_cap + first = model_designer.get_models( + ["judge"], model_configs=[config], client_concurrency_mode=ClientConcurrencyMode.ASYNC + )["judge"] + config = config.model_copy(deep=True) + config.inference_parameters.max_parallel_requests = new_cap + second = None + in_flight = peak = 0 + + async def respond(request: httpx.Request) -> httpx.Response: + nonlocal in_flight, peak + in_flight += 1 + peak = max(peak, in_flight) + await asyncio.sleep(0.01) + in_flight -= 1 + return httpx.Response(200, json=completion_payload) + + httpx_mock.add_callback(respond, is_reusable=True) + messages = [ChatMessage.as_user("test")] + try: + if initial_cap != new_cap: + patient = config.model_copy(update={"alias": "patient"}, deep=True) + patient.inference_parameters.max_parallel_requests = 1 + with pytest.raises(ValueError, match="already registered with max_parallel_requests"): + model_designer.get_models( + ["patient", "judge"], + model_configs=[patient, config], + client_concurrency_mode=ClientConcurrencyMode.ASYNC, + ) + else: + second = model_designer.get_models( + ["judge"], model_configs=[config], client_concurrency_mode=ClientConcurrencyMode.ASYNC + )["judge"] + await first.acompletion(messages) + await asyncio.gather(first.acompletion(messages), (second or first).acompletion(messages)) + assert peak == min(initial_cap, 2) + await first.aclose() + if second is None: + with pytest.raises(ValueError, match="already registered with max_parallel_requests"): + model_designer.get_models(["judge"], model_configs=[config]) + model_designer.set_run_config(RunConfig()) + second = model_designer.get_models( + ["judge"], model_configs=[config], client_concurrency_mode=ClientConcurrencyMode.ASYNC + )["judge"] + assert (await second.acompletion(messages)).message.content == "hello" + finally: + await first.aclose() + if second is not None: + await second.aclose() + + def test_get_models_retries( model_designer: DataDesigner, custom_model_configs: list[ModelConfig],