diff --git a/.env.sample b/.env.sample index d72da9a..395cb55 100644 --- a/.env.sample +++ b/.env.sample @@ -12,6 +12,10 @@ JOINTFM_SCHEMA_VERSION=v1 # 1. Deployment ID: SDK builds the hosted predictionsUnstructured URL. # JOINTFM_DEPLOYMENT_ID= +# 1b. Comma-separated deployment IDs for load-balanced hosted calls (same checkpoint). +# Requires at least two unique IDs. Mutually exclusive with JOINTFM_DEPLOYMENT_ID and other selectors. +# JOINTFM_DEPLOYMENT_IDS=chevron-id,research-id + # 2. Deployment URL: SDK appends predictionsUnstructured. # JOINTFM_DEPLOYMENT_URL=https://app.datarobot.com/api/v2/deployments/ diff --git a/README.md b/README.md index cc0248a..f8e2a52 100644 --- a/README.md +++ b/README.md @@ -38,7 +38,7 @@ The direct local service exposes `GET /healthz` and `POST /predict`. Structured SDK defaults live in `jointfm_client.configuration.JointFMConfig` and are mirrored in the checked-in `config.sample.yaml`. Copy `config.sample.yaml` to `config.yaml` and change only the fields needed for your deployment or transport defaults. `JointFMClient.from_env()` and `load_settings()` read `config.yaml` by default, then layer `.env` values over it, then layer process environment variables or the supplied `env` mapping over both. Explicit Python arguments such as `timeout=` and `retry_config=` still override YAML transport defaults. -`JointFMClient.from_env()` and `load_settings()` resolve `JOINTFM_SCHEMA_VERSION` and exactly one service selector from that layered configuration. `JOINTFM_MODEL_VERSION` is optional: when unset the SDK discovers the model version from `/healthz` on first use, and when set the SDK validates it against `/healthz` as a drift-detection guard. Hosted selectors also require `DATAROBOT_ENDPOINT` and `DATAROBOT_API_TOKEN`; the direct local selector does not use DataRobot credentials. Missing credentials, missing schema version, malformed credentials, unsupported schema versions, missing selectors, and multiple selectors raise `JointFMConfigurationError`. +`JointFMClient.from_env()` and `load_settings()` resolve `JOINTFM_SCHEMA_VERSION` and exactly one service selector from that layered configuration. Hosted options include `JOINTFM_DEPLOYMENT_ID` or load-balanced `JOINTFM_DEPLOYMENT_IDS` (comma-separated same-checkpoint peers; mutually exclusive with other selectors). `JOINTFM_MODEL_VERSION` is optional: when unset the SDK discovers the model version from `/healthz` on first use, and when set the SDK validates it against `/healthz` as a drift-detection guard. Hosted selectors also require `DATAROBOT_ENDPOINT` and `DATAROBOT_API_TOKEN`; the direct local selector does not use DataRobot credentials. Missing credentials, missing schema version, malformed credentials, unsupported schema versions, missing selectors, and multiple selectors raise `JointFMConfigurationError`. `DATAROBOT_ENDPOINT` must be a normalized HTTPS DataRobot API v2 URL ending in `/api/v2`; the SDK stores it without a trailing slash. `DATAROBOT_API_TOKEN` must be non-empty and whitespace-free. The token is excluded from `JointFMSettings` repr output. @@ -86,6 +86,7 @@ JOINTFM_SCHEMA_VERSION=v1 Choose exactly one service selector: - `JOINTFM_DEPLOYMENT_ID`: builds `DATAROBOT_ENDPOINT.rstrip("/") + "/"` plus `deployments/{deployment_id}/predictionsUnstructured` +- `JOINTFM_DEPLOYMENT_IDS`: comma-separated hosted deployment IDs (≥2 unique, same checkpoint) for round-robin load balancing; mutually exclusive with other selectors - `JOINTFM_DEPLOYMENT_URL`: appends `/predictionsUnstructured` to a hosted deployment URL - `JOINTFM_PREDICT_URL`: uses a full hosted prediction URL ending in `/predictionsUnstructured` - `JOINTFM_DEPLOYMENT_TARGET` with `JOINTFM_PULUMI_OUTPUTS_PATH`: resolves a named target from saved Pulumi outputs JSON, preferring `deployment_id`, then `deployment_url`, then `predict_url` @@ -245,6 +246,7 @@ DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2 DATAROBOT_API_TOKEN= JOINTFM_SCHEMA_VERSION=v1 JOINTFM_DEPLOYMENT_ID= +# Or: JOINTFM_DEPLOYMENT_IDS=chevron-id,research-id # Optional drift-detection pin; the SDK discovers the model version from /healthz when unset: # JOINTFM_MODEL_VERSION=jointfm-inference:0.2.0+ckpt.fin-2026-05-22 ``` diff --git a/config.sample.yaml b/config.sample.yaml index ecf4bea..924e567 100644 --- a/config.sample.yaml +++ b/config.sample.yaml @@ -6,6 +6,7 @@ environment: datarobot_endpoint: DATAROBOT_ENDPOINT datarobot_api_token: DATAROBOT_API_TOKEN deployment_id: JOINTFM_DEPLOYMENT_ID + deployment_ids: JOINTFM_DEPLOYMENT_IDS deployment_url: JOINTFM_DEPLOYMENT_URL predict_url: JOINTFM_PREDICT_URL deployment_target: JOINTFM_DEPLOYMENT_TARGET @@ -17,6 +18,7 @@ deployment: datarobot_endpoint: null datarobot_api_token: null deployment_id: null + deployment_ids: null deployment_url: null predict_url: null deployment_target: null diff --git a/docs/api-reference.md b/docs/api-reference.md index 6f10e74..380c9dc 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -34,7 +34,7 @@ This reference covers the supported public Python surface exported by `jointfm_c | Name | Purpose | | --- | --- | -| `JointFMSettings` | Validated hosted or local service settings: optional normalized DataRobot endpoint, optional secret token, health and prediction URLs, service selector, schema pin, model pin, and optional selector details. The API token is excluded from `repr`. | +| `JointFMSettings` | Validated hosted or local service settings: optional normalized DataRobot endpoint, optional secret token, health and prediction URLs, service selector, schema pin, optional load-balanced `instances` pool, model pin, and optional selector details. The API token is excluded from `repr`. | | `JointFMConfig` | Top-level structured configuration loaded from defaults, YAML, and explicit overrides. | | `PathConfig` | Default local file names for `config.yaml`, `config.sample.yaml`, and `.env`. | | `EnvironmentVariableConfig` | Environment variable names consumed by settings loading. | @@ -119,6 +119,7 @@ All SDK-specific exceptions inherit from `JointFMError`. | `JOINTFM_SCHEMA_VERSION` | Hosted calls | Request schema pin. The SDK supports only `v1`. | | `JOINTFM_MODEL_VERSION` | Hosted calls | Exact JointFM deployment model version expected from the service-health payload and prediction responses. | | `JOINTFM_DEPLOYMENT_ID` | One selector | Deployment ID used to build hosted health and prediction URLs. | +| `JOINTFM_DEPLOYMENT_IDS` | One selector | Comma-separated hosted deployment IDs for round-robin load balancing (at least two unique IDs). Mutually exclusive with other selectors. Peers must share `model_version` and `checkpoint_version`; the SDK uses the minimum `max_sample_count`. | | `JOINTFM_DEPLOYMENT_URL` | One selector | Hosted deployment URL; the SDK derives the `/predictionsUnstructured` route from it and reuses that route for health probes. | | `JOINTFM_PREDICT_URL` | One selector | Full hosted prediction URL ending in `/predictionsUnstructured`; the SDK derives the owning deployment URL. | | `JOINTFM_DEPLOYMENT_TARGET` | One selector with outputs path | Key in a saved Pulumi outputs JSON file. | @@ -126,7 +127,7 @@ All SDK-specific exceptions inherit from `JointFMError`. | `JOINTFM_LOCAL_BASE_URL` | One selector | Direct local JointFM REST service base URL. The SDK calls `GET /healthz` and `POST /predict` without DataRobot authorization. | | `DATAROBOT_DEPLOYMENT_ID` | Optional live tests | Hosted deployment ID used only by the optional live smoke test so normal CI does not call DataRobot accidentally. | -Set exactly one selector among `JOINTFM_DEPLOYMENT_ID`, `JOINTFM_DEPLOYMENT_URL`, `JOINTFM_PREDICT_URL`, `JOINTFM_DEPLOYMENT_TARGET`, and `JOINTFM_LOCAL_BASE_URL`. +Set exactly one selector among `JOINTFM_DEPLOYMENT_ID`, `JOINTFM_DEPLOYMENT_IDS`, `JOINTFM_DEPLOYMENT_URL`, `JOINTFM_PREDICT_URL`, `JOINTFM_DEPLOYMENT_TARGET`, and `JOINTFM_LOCAL_BASE_URL`. ## V1 Payload Fields diff --git a/src/jointfm_client/__init__.py b/src/jointfm_client/__init__.py index a5810b5..791338a 100644 --- a/src/jointfm_client/__init__.py +++ b/src/jointfm_client/__init__.py @@ -99,6 +99,7 @@ UnsupportedSchemaVersionError, UnsupportedServiceContractError, ) +from jointfm_client.pool import JointFMInstancePool from jointfm_client.notebooks import ( WORKSPACE_ROOT_MARKERS, bootstrap_notebook, @@ -108,6 +109,7 @@ DATAROBOT_API_TOKEN_ENV, DATAROBOT_ENDPOINT_ENV, JOINTFM_DEPLOYMENT_ID_ENV, + JOINTFM_DEPLOYMENT_IDS_ENV, JOINTFM_DEPLOYMENT_TARGET_ENV, JOINTFM_DEPLOYMENT_URL_ENV, JOINTFM_LOCAL_BASE_URL_ENV, @@ -115,6 +117,7 @@ JOINTFM_PREDICT_URL_ENV, JOINTFM_PULUMI_OUTPUTS_PATH_ENV, JOINTFM_SCHEMA_VERSION_ENV, + JointFMInstanceSettings, JointFMSettings, build_datarobot_prediction_headers, build_hosted_deployment_url, @@ -170,6 +173,7 @@ "IMPORT_NAMESPACE", "MeanForecastResult", "JOINTFM_DEPLOYMENT_ID_ENV", + "JOINTFM_DEPLOYMENT_IDS_ENV", "JOINTFM_DEPLOYMENT_TARGET_ENV", "JOINTFM_DEPLOYMENT_URL_ENV", "JOINTFM_LOCAL_BASE_URL_ENV", @@ -189,6 +193,8 @@ "JointFMResponseDecodeError", "JointFMResponseError", "JointFMServiceError", + "JointFMInstancePool", + "JointFMInstanceSettings", "JointFMRetryConfig", "JointFMSettings", "JointFMTimeoutConfig", diff --git a/src/jointfm_client/client.py b/src/jointfm_client/client.py index 05a1a3a..c7c80fa 100644 --- a/src/jointfm_client/client.py +++ b/src/jointfm_client/client.py @@ -17,6 +17,7 @@ from __future__ import annotations from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import re from typing import Any, Self, cast @@ -57,6 +58,7 @@ QuantileForecastResult, SampleForecastResult, ) +from jointfm_client.pool import JointFMInstancePool from jointfm_client.settings import ( JointFMSettings, load_settings, @@ -108,6 +110,7 @@ def __init__( self._datarobot_request_id_headers = datarobot_request_id_headers self._health_metadata: HealthMetadata | None = None self._sample_batch_cap: int | None = None + self._pool: JointFMInstancePool | None = None @classmethod def from_env( @@ -152,10 +155,21 @@ def health(self, *, cache: bool = False, refresh: bool = False) -> HealthMetadat deployment gateway only proxies the unstructured prediction route; the container short-circuits that body before any schema or model version validation and returns the same typed health payload. + + When ``JOINTFM_DEPLOYMENT_IDS`` is set, reachable peers are probed and must + share ``model_version`` and ``checkpoint_version``; the sample-batch cap is + the minimum ``max_sample_count`` across those peers. """ if cache and not refresh and self._health_metadata is not None: return self._health_metadata + if self._uses_pool(): + metadata = self._require_pool().probe_all_health() + self._sample_batch_cap = metadata.max_sample_count + if cache: + self._health_metadata = metadata + return metadata + if self._uses_predict_route_for_health(): payload = self._fetch_hosted_health_payload() else: @@ -195,14 +209,14 @@ def _fetch_hosted_health_payload(self) -> Mapping[str, Any]: def predict(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: """Submit one V1 JSON prediction payload to the configured endpoint.""" - predict_url = self._require_predict_url("predict") + self._require_predict_url("predict") model_version = payload.get("model_version") if not isinstance(model_version, str): raise JointFMConfigurationError( "JointFMClient.predict() requires payload['model_version']" ) self._resolve_model_version(model_version=model_version) - response_payload = self._transport_for_request().post_json(predict_url, payload) + response_payload = self._post_predict_json(payload) ForecastResponse.raise_for_errors(response_payload) return response_payload @@ -243,7 +257,7 @@ def forecast( | None = None, ) -> ForecastResponse: """Build and submit a forecast request from tabular history inputs.""" - predict_url = self._require_predict_url("forecast") + self._require_predict_url("forecast") resolved_model_version = self._resolve_model_version( model_version=model_version, ) @@ -303,18 +317,16 @@ def forecast( ) sample_cap = self._resolve_sample_batch_cap(payload) if sample_cap is not None: - return self._forecast_sample_batches(predict_url, payload, sample_cap) + return self._forecast_sample_batches(payload, sample_cap) try: - response_payload = self._transport_for_request().post_json( - predict_url, payload - ) + response_payload = self._post_predict_json(payload) except JointFMHTTPStatusError as error: sample_cap = _sample_batch_cap_from_error(error, payload) if sample_cap is None: raise self._sample_batch_cap = sample_cap - return self._forecast_sample_batches(predict_url, payload, sample_cap) + return self._forecast_sample_batches(payload, sample_cap) return _forecast_response_from_payload(response_payload, payload) @@ -484,37 +496,31 @@ def _resolve_sample_batch_cap(self, payload: Mapping[str, Any]) -> int | None: def _forecast_sample_batches( self, - predict_url: str, payload: Mapping[str, Any], sample_cap: int, ) -> SampleForecastResult: requested_samples = cast(int, payload["n_samples"]) remaining_samples = requested_samples + batch_payloads: list[dict[str, Any]] = [] batch_index = 0 - batch_results: list[SampleForecastResult] = [] while remaining_samples > 0: batch_samples = min(sample_cap, remaining_samples) batch_payload = dict(payload) batch_payload["n_samples"] = batch_samples _set_batch_seed(batch_payload, batch_index) - response_payload = self._transport_for_request().post_json( - predict_url, - batch_payload, - ) - batch_result = _forecast_response_from_payload( - response_payload, - batch_payload, - ) - if not isinstance(batch_result, SampleForecastResult): - raise JointFMServiceError( - "JointFM forecast response violated the V1 contract: " - "sample batching requires sample forecast responses" - ) - batch_results.append(batch_result) + batch_payloads.append(batch_payload) remaining_samples -= batch_samples batch_index += 1 + if self._uses_pool() and len(batch_payloads) > 1: + batch_results = self._forecast_sample_batches_parallel(batch_payloads) + else: + batch_results = [ + self._sample_forecast_from_batch_payload(batch_payload) + for batch_payload in batch_payloads + ] + try: return _merge_sample_forecast_results(batch_results, payload) except ValueError as error: @@ -522,6 +528,87 @@ def _forecast_sample_batches( f"JointFM forecast response violated the V1 contract: {error}" ) from error + def _forecast_sample_batches_parallel( + self, batch_payloads: Sequence[Mapping[str, Any]] + ) -> list[SampleForecastResult]: + pool = self._require_pool() + max_workers = min(len(batch_payloads), pool.instance_count) + + def _run_batch(item: tuple[int, Mapping[str, Any]]) -> SampleForecastResult: + batch_index, batch_payload = item + response_payload = pool.post_json_to( + pool.instance_at(batch_index), batch_payload + ) + return self._sample_forecast_from_response(response_payload, batch_payload) + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + return list( + executor.map( + _run_batch, + enumerate(batch_payloads), + ) + ) + + def _sample_forecast_from_batch_payload( + self, batch_payload: Mapping[str, Any] + ) -> SampleForecastResult: + return self._sample_forecast_from_response( + self._post_predict_json(batch_payload), batch_payload + ) + + def _sample_forecast_from_response( + self, response_payload: Mapping[str, Any], batch_payload: Mapping[str, Any] + ) -> SampleForecastResult: + batch_result = _forecast_response_from_payload(response_payload, batch_payload) + if not isinstance(batch_result, SampleForecastResult): + raise JointFMServiceError( + "JointFM forecast response violated the V1 contract: " + "sample batching requires sample forecast responses" + ) + return batch_result + + def _post_predict_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + if self._uses_pool(): + return self._require_pool().post_json(payload) + predict_url = self._require_predict_url("predict") + return self._transport_for_request().post_json(predict_url, payload) + + def _uses_pool(self) -> bool: + return self.settings is not None and len(self.settings.instances) > 1 + + def _require_pool(self) -> JointFMInstancePool: + if not self._uses_pool(): + raise JointFMConfigurationError( + "JointFMClient pool routing requires multiple deployment instances" + ) + if self._pool is None: + assert self.settings is not None + # One Session per peer for thread-safe parallel sample batches. + # An injected transport is reused across peers (tests/mocks only); + # production from_env builds a distinct fail-fast transport each. + if self._transport is not None: + transports = tuple(self._transport for _ in self.settings.instances) + else: + transports = tuple( + self._new_pool_peer_transport() for _ in self.settings.instances + ) + self._pool = JointFMInstancePool( + instances=self.settings.instances, + transports=transports, + expected_model_version=self.settings.model_version, + ) + return self._pool + + def _new_pool_peer_transport(self) -> JSONTransport: + assert self.settings is not None + return JointFMHTTPTransport.from_settings( + self.settings, + timeout=self._timeout, + retry_config=JointFMRetryConfig(max_attempts=1), + response_body_excerpt_characters=(self._response_body_excerpt_characters), + datarobot_request_id_headers=self._datarobot_request_id_headers, + ) + def _require_settings(self, method_name: str) -> JointFMSettings: if self.settings is None: raise JointFMConfigurationError( @@ -576,6 +663,10 @@ def _resolve_model_version( assert self._health_metadata is not None return self._health_metadata.model_version + # Pool peers must share checkpoint identity before any traffic. + if self._uses_pool() and self._health_metadata is None: + self.health(cache=True) + normalized_model_version = validate_jointfm_model_version( configured_model_version ) diff --git a/src/jointfm_client/configuration.py b/src/jointfm_client/configuration.py index c2205f1..a82b7bd 100644 --- a/src/jointfm_client/configuration.py +++ b/src/jointfm_client/configuration.py @@ -54,6 +54,7 @@ class EnvironmentVariableConfig(_ConfigModel): datarobot_endpoint: str = "DATAROBOT_ENDPOINT" datarobot_api_token: str = "DATAROBOT_API_TOKEN" deployment_id: str = "JOINTFM_DEPLOYMENT_ID" + deployment_ids: str = "JOINTFM_DEPLOYMENT_IDS" deployment_url: str = "JOINTFM_DEPLOYMENT_URL" predict_url: str = "JOINTFM_PREDICT_URL" deployment_target: str = "JOINTFM_DEPLOYMENT_TARGET" @@ -92,6 +93,7 @@ class HostedDeploymentConfig(_ConfigModel): datarobot_endpoint: str | None = None datarobot_api_token: str | None = Field(default=None, repr=False) deployment_id: str | None = None + deployment_ids: str | None = None deployment_url: str | None = None predict_url: str | None = None deployment_target: str | None = None @@ -113,6 +115,7 @@ def to_environment_values( values, environment.datarobot_api_token, self.datarobot_api_token ) _set_if_configured(values, environment.deployment_id, self.deployment_id) + _set_if_configured(values, environment.deployment_ids, self.deployment_ids) _set_if_configured(values, environment.deployment_url, self.deployment_url) _set_if_configured(values, environment.predict_url, self.predict_url) _set_if_configured( @@ -355,6 +358,7 @@ def _set_if_configured(values: dict[str, str], name: str, value: str | None) -> DATAROBOT_ENDPOINT_ENV: Final = DEFAULT_ENVIRONMENT_CONFIG.datarobot_endpoint DATAROBOT_API_TOKEN_ENV: Final = DEFAULT_ENVIRONMENT_CONFIG.datarobot_api_token JOINTFM_DEPLOYMENT_ID_ENV: Final = DEFAULT_ENVIRONMENT_CONFIG.deployment_id +JOINTFM_DEPLOYMENT_IDS_ENV: Final = DEFAULT_ENVIRONMENT_CONFIG.deployment_ids JOINTFM_DEPLOYMENT_URL_ENV: Final = DEFAULT_ENVIRONMENT_CONFIG.deployment_url JOINTFM_PREDICT_URL_ENV: Final = DEFAULT_ENVIRONMENT_CONFIG.predict_url JOINTFM_DEPLOYMENT_TARGET_ENV: Final = DEFAULT_ENVIRONMENT_CONFIG.deployment_target @@ -433,6 +437,7 @@ def _set_if_configured(values: dict[str, str], name: str, value: str | None) -> "ForecastCsvConfig", "HostedDeploymentConfig", "JOINTFM_DEPLOYMENT_ID_ENV", + "JOINTFM_DEPLOYMENT_IDS_ENV", "JOINTFM_DEPLOYMENT_TARGET_ENV", "JOINTFM_DEPLOYMENT_URL_ENV", "JOINTFM_LOCAL_BASE_URL_ENV", diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py new file mode 100644 index 0000000..a17443b --- /dev/null +++ b/src/jointfm_client/pool.py @@ -0,0 +1,330 @@ +# Copyright 2026 DataRobot, Inc. and its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Round-robin load balancing across multiple hosted JointFM deployments.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, replace +import logging +import threading +import time +from typing import Any + +from jointfm_client.configuration import DEFAULT_RETRY_STATUS_CODES +from jointfm_client.contract import ( + HEALTH_REQUEST_TYPE, + HealthMetadata, + validate_service_metadata, +) +from jointfm_client.exceptions import ( + JointFMHTTPStatusError, + JointFMRequestError, + UnsupportedModelVersionError, + UnsupportedServiceContractError, +) +from jointfm_client.settings import JointFMInstanceSettings +from jointfm_client.transport import JSONTransport + +logger = logging.getLogger(__name__) + +_POOL_RETRYABLE_HTTP_STATUS_CODES = frozenset(DEFAULT_RETRY_STATUS_CODES) +_DEFAULT_PEER_COOLDOWN_SECONDS = 30.0 + + +class PoolPeer: + """One hosted deployment endpoint with a Session-safe POST lock.""" + + def __init__( + self, + settings: JointFMInstanceSettings, + transport: JSONTransport, + post_lock: threading.Lock, + ) -> None: + self.settings = settings + self.transport = transport + self._post_lock = post_lock + + @property + def deployment_id(self) -> str: + """Deployment identifier used for routing and health accounting.""" + return self.settings.deployment_id + + def post_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """POST ``payload`` to this peer's predict URL under the peer lock.""" + with self._post_lock: + return self.transport.post_json(self.settings.predict_url, payload) + + +@dataclass(frozen=True, slots=True) +class HealthProbeResult: + """Merged health metadata and the peer IDs that passed the gate.""" + + metadata: HealthMetadata + healthy_ids: frozenset[str] + + +class PeerRoutingState: + """Active-set, cooldown, and round-robin cursor for pool peers.""" + + def __init__( + self, + *, + deployment_ids: Sequence[str], + peer_cooldown_seconds: float, + ) -> None: + self._peer_cooldown_seconds = peer_cooldown_seconds + self._lock = threading.Lock() + self._index = 0 + self._active_ids = set(deployment_ids) + self._cooldown_until: dict[str, float] = {} + + def eligible(self, peers: Sequence[PoolPeer]) -> tuple[PoolPeer, ...]: + """Return peers currently preferred for routing.""" + with self._lock: + return self._eligible_unlocked(peers) + + def next(self, peers: Sequence[PoolPeer]) -> PoolPeer: + """Return the next eligible peer using round-robin selection.""" + with self._lock: + active = self._eligible_unlocked(peers) + peer = active[self._index % len(active)] + self._index = (self._index + 1) % len(active) + return peer + + def at(self, peers: Sequence[PoolPeer], index: int) -> PoolPeer: + """Return the eligible peer pinned for ``index`` (sticky batch mapping).""" + active = self.eligible(peers) + return active[index % len(active)] + + def set_healthy(self, healthy_ids: Sequence[str]) -> None: + """Replace the active set with health-passing peers and clear their cooldowns.""" + with self._lock: + self._active_ids = set(healthy_ids) + self._index = 0 + for deployment_id in healthy_ids: + self._cooldown_until.pop(deployment_id, None) + + def cool_down(self, deployment_id: str) -> None: + """Temporarily exclude ``deployment_id`` from preferred routing.""" + with self._lock: + self._cooldown_until[deployment_id] = ( + time.monotonic() + self._peer_cooldown_seconds + ) + + def reactivate(self, deployment_id: str) -> None: + """Mark ``deployment_id`` active and clear any cooldown.""" + with self._lock: + self._active_ids.add(deployment_id) + self._cooldown_until.pop(deployment_id, None) + + def _eligible_unlocked(self, peers: Sequence[PoolPeer]) -> tuple[PoolPeer, ...]: + now = time.monotonic() + active = tuple(peer for peer in peers if peer.deployment_id in self._active_ids) + not_cooling = tuple( + peer + for peer in active + if self._cooldown_until.get(peer.deployment_id, 0.0) <= now + ) + # All cooling: still try health-active peers rather than stall. + return not_cooling or active or tuple(peers) + + +class PoolHealthGate: + """Probe peers and require matching model/checkpoint across healthy ones.""" + + def __init__(self, *, expected_model_version: str | None = None) -> None: + self._expected_model_version = expected_model_version + + def probe(self, peers: Sequence[PoolPeer]) -> HealthProbeResult: + """Probe peers; skip per-peer failures; fail only when none are usable.""" + healthy: list[tuple[PoolPeer, HealthMetadata]] = [] + last_error: BaseException | None = None + for peer in peers: + try: + payload = peer.post_json({"request_type": HEALTH_REQUEST_TYPE}) + validate_service_metadata( + payload, expected_model_version=self._expected_model_version + ) + healthy.append((peer, HealthMetadata.from_payload(payload))) + except Exception as error: + # Skip unreachable or incompatible peers; a bad backup must not + # take down a healthy primary. + last_error = error + _log_unavailable(peer.deployment_id, error) + + if not healthy: + assert last_error is not None + raise last_error + + reference = healthy[0][1] + max_samples = reference.max_sample_count + for peer, metadata in healthy[1:]: + _reject_metadata_mismatch(peer.deployment_id, metadata, reference) + max_samples = min(max_samples, metadata.max_sample_count) + + metadata = ( + reference + if max_samples == reference.max_sample_count + else replace(reference, max_sample_count=max_samples) + ) + return HealthProbeResult( + metadata=metadata, + healthy_ids=frozenset(peer.deployment_id for peer, _ in healthy), + ) + + +class JointFMInstancePool: + """Round-robin JointFM requests across hosted deployments. + + Use one fail-fast transport per peer (``max_attempts=1``); this pool retries + peers. Each transport is locked independently so parallel batches stay + concurrent across peers while failover onto a busy peer stays Session-safe. + """ + + def __init__( + self, + *, + instances: Sequence[JointFMInstanceSettings], + transports: Sequence[JSONTransport], + expected_model_version: str | None = None, + peer_cooldown_seconds: float = _DEFAULT_PEER_COOLDOWN_SECONDS, + ) -> None: + if len(instances) < 2: + raise ValueError("JointFMInstancePool requires at least two instances") + if len(transports) != len(instances): + raise ValueError( + "JointFMInstancePool requires one transport per instance: " + f"got {len(transports)} transports for {len(instances)} instances" + ) + if peer_cooldown_seconds < 0: + raise ValueError("peer_cooldown_seconds must be >= 0") + # Key locks by transport identity so a shared injected Session serializes, + # while distinct per-peer Sessions stay concurrent. + locks_by_transport: dict[int, threading.Lock] = {} + peers: list[PoolPeer] = [] + for instance, transport in zip(instances, transports, strict=True): + post_lock = locks_by_transport.setdefault(id(transport), threading.Lock()) + peers.append(PoolPeer(instance, transport, post_lock)) + self._peers = tuple(peers) + self._peers_by_id = {peer.deployment_id: peer for peer in self._peers} + self._routing = PeerRoutingState( + deployment_ids=tuple(peer.deployment_id for peer in self._peers), + peer_cooldown_seconds=peer_cooldown_seconds, + ) + self._health_gate = PoolHealthGate( + expected_model_version=expected_model_version + ) + + @property + def instance_count(self) -> int: + """Number of peers currently eligible for routing.""" + return len(self._routing.eligible(self._peers)) + + def next_instance(self) -> JointFMInstanceSettings: + """Return the next eligible instance using round-robin selection.""" + return self._routing.next(self._peers).settings + + def instance_at(self, index: int) -> JointFMInstanceSettings: + """Return the eligible instance pinned for ``index`` (sticky batch mapping).""" + return self._routing.at(self._peers, index).settings + + def probe_all_health(self) -> HealthMetadata: + """Probe peers; require matching model/checkpoint; return min sample cap. + + Per-peer transport or contract failures skip that peer. The pool fails + only when no peer is usable, or when usable peers disagree with each + other on model/checkpoint. + """ + result = self._health_gate.probe(self._peers) + self._routing.set_healthy(tuple(result.healthy_ids)) + return result.metadata + + def post_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """POST payload via round-robin, trying other peers on retryable failures.""" + return self.post_json_to(self.next_instance(), payload) + + def post_json_to( + self, instance: JointFMInstanceSettings, payload: Mapping[str, Any] + ) -> Mapping[str, Any]: + """POST to ``instance`` first; on retryable failure try remaining peers.""" + preferred = self._peers_by_id[instance.deployment_id] + candidates = self._failover_candidates(preferred) + last_error: BaseException | None = None + for candidate in candidates: + try: + result = candidate.post_json(payload) + except Exception as error: + if not _is_pool_retryable(error): + raise + last_error = error + self._routing.cool_down(candidate.deployment_id) + _log_unavailable(candidate.deployment_id, error) + continue + self._routing.reactivate(candidate.deployment_id) + return result + assert last_error is not None + raise last_error + + def _failover_candidates(self, preferred: PoolPeer) -> tuple[PoolPeer, ...]: + """Prefer health-eligible peers; then try health-excluded peers last.""" + eligible = self._routing.eligible(self._peers) + if preferred.deployment_id in {peer.deployment_id for peer in eligible}: + preferred_first = (preferred,) + tuple( + peer + for peer in eligible + if peer.deployment_id != preferred.deployment_id + ) + else: + preferred_first = eligible + seen = {peer.deployment_id for peer in preferred_first} + last_resort = tuple( + peer for peer in self._peers if peer.deployment_id not in seen + ) + return preferred_first + last_resort + + +def _reject_metadata_mismatch( + deployment_id: str, + metadata: HealthMetadata, + reference: HealthMetadata, +) -> None: + if metadata.model_version != reference.model_version: + raise UnsupportedModelVersionError( + "JointFM deployment pool model_version mismatch: " + f"{deployment_id!r} advertises {metadata.model_version!r}, " + f"expected {reference.model_version!r}" + ) + if metadata.checkpoint_version != reference.checkpoint_version: + raise UnsupportedServiceContractError( + "JointFM deployment pool checkpoint_version mismatch: " + f"{deployment_id!r} advertises {metadata.checkpoint_version!r}, " + f"expected {reference.checkpoint_version!r}" + ) + + +def _log_unavailable(deployment_id: str, error: BaseException) -> None: + logger.warning( + "JointFM instance unavailable: deployment_id=%s error=%s", deployment_id, error + ) + + +def _is_pool_retryable(error: BaseException) -> bool: + if isinstance(error, JointFMRequestError): + return True + return ( + isinstance(error, JointFMHTTPStatusError) + and error.status_code in _POOL_RETRYABLE_HTTP_STATUS_CODES + ) diff --git a/src/jointfm_client/settings.py b/src/jointfm_client/settings.py index ce3b0c4..78b5dd3 100644 --- a/src/jointfm_client/settings.py +++ b/src/jointfm_client/settings.py @@ -32,6 +32,7 @@ DEFAULT_CONFIG_PATH, EnvironmentVariableConfig, JOINTFM_DEPLOYMENT_ID_ENV, + JOINTFM_DEPLOYMENT_IDS_ENV, JOINTFM_DEPLOYMENT_TARGET_ENV, JOINTFM_DEPLOYMENT_URL_ENV, JOINTFM_LOCAL_BASE_URL_ENV, @@ -52,6 +53,7 @@ DeploymentSelector: TypeAlias = Literal[ "deployment_id", + "deployment_ids", "deployment_url", "predict_url", "pulumi_target", @@ -59,6 +61,14 @@ ] +@dataclass(frozen=True, slots=True) +class JointFMInstanceSettings: + """One hosted JointFM deployment target in a load-balanced pool.""" + + deployment_id: str + predict_url: str + + @dataclass(frozen=True, slots=True) class JointFMSettings: """Validated settings for one hosted or local JointFM service target.""" @@ -69,6 +79,7 @@ class JointFMSettings: predict_url: str deployment_selector: DeploymentSelector schema_version: str + instances: tuple[JointFMInstanceSettings, ...] = () model_version: str | None = None deployment_id: str | None = None deployment_url: str | None = None @@ -96,6 +107,16 @@ def load_settings( _required_env(env_values, environment.schema_version) ) model_version = _optional_model_version(env_values, environment.model_version) + deployment_ids_value = env_values.get(environment.deployment_ids) + if deployment_ids_value is not None and deployment_ids_value != "": + _reject_conflicting_selectors_with_deployment_ids(env_values, environment) + return _load_hosted_deployment_pool_settings( + env_values, + environment, + schema_version=schema_version, + model_version=model_version, + ) + selector_name = _resolve_single_deployment_selector(env_values, environment) if selector_name == environment.local_base_url: @@ -487,13 +508,96 @@ def _resolve_single_deployment_selector( if selector_name in env and env[selector_name] != "" ] if len(selector_names) != 1: - formatted_selectors = ", ".join(deployment_selector_envs) + formatted_selectors = ", ".join( + (*deployment_selector_envs, environment.deployment_ids) + ) raise JointFMConfigurationError( f"Exactly one deployment selector is required: {formatted_selectors}" ) return selector_names[0] +def _load_hosted_deployment_pool_settings( + env: Mapping[str, str], + environment: EnvironmentVariableConfig, + *, + schema_version: str, + model_version: str | None, +) -> JointFMSettings: + datarobot_endpoint = normalize_datarobot_endpoint( + _required_env(env, environment.datarobot_endpoint) + ) + datarobot_api_token = validate_datarobot_api_token( + _required_env(env, environment.datarobot_api_token) + ) + deployment_ids = _parse_deployment_ids( + _required_env(env, environment.deployment_ids) + ) + instances = tuple( + JointFMInstanceSettings( + deployment_id=deployment_id, + predict_url=build_hosted_predict_url(datarobot_endpoint, deployment_id), + ) + for deployment_id in deployment_ids + ) + primary = instances[0] + return JointFMSettings( + datarobot_endpoint=datarobot_endpoint, + datarobot_api_token=datarobot_api_token, + health_url=primary.predict_url, + predict_url=primary.predict_url, + deployment_selector="deployment_ids", + schema_version=schema_version, + instances=instances, + model_version=model_version, + deployment_id=primary.deployment_id, + deployment_url=build_hosted_deployment_url( + datarobot_endpoint, + primary.deployment_id, + ), + ) + + +def _parse_deployment_ids(value: str) -> tuple[str, ...]: + if value.strip() != value or "," not in value: + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must be a comma-separated list of " + "at least two unique deployment IDs" + ) + deployment_ids: list[str] = [] + seen: set[str] = set() + for part in value.split(","): + if part == "" or any(character.isspace() for character in part) or "/" in part: + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must contain non-empty deployment IDs " + "without whitespace or URL paths" + ) + if part in seen: + continue + seen.add(part) + deployment_ids.append(part) + if len(deployment_ids) < 2: + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must contain at least two unique deployment IDs" + ) + return tuple(deployment_ids) + + +def _reject_conflicting_selectors_with_deployment_ids( + env: Mapping[str, str], + environment: EnvironmentVariableConfig, +) -> None: + conflicting = [ + name + for name in environment.deployment_selector_names() + if name in env and env[name] != "" + ] + if conflicting: + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} cannot be combined with {', '.join(conflicting)}" + ) + + def _load_pulumi_target_outputs( outputs_path: str, deployment_target: str, diff --git a/tests/test_pool.py b/tests/test_pool.py new file mode 100644 index 0000000..a623ee4 --- /dev/null +++ b/tests/test_pool.py @@ -0,0 +1,285 @@ +# Copyright 2026 DataRobot, Inc. and its affiliates. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for JointFM instance pool load balancing.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +import threading +from typing import Any + +import pytest + +from jointfm_client import ( + JointFMHTTPStatusError, + JointFMInstancePool, + JointFMInstanceSettings, + UnsupportedModelVersionError, + UnsupportedServiceContractError, +) +from jointfm_client.transport import JSONTransport + + +def _instance(deployment_id: str) -> JointFMInstanceSettings: + """Instance.""" + return JointFMInstanceSettings( + deployment_id=deployment_id, + predict_url=( + "https://app.datarobot.com/api/v2/deployments/" + f"{deployment_id}/predictionsUnstructured" + ), + ) + + +def _health( + *, + model_version: str = "jointfm-inference:0.2.0+ckpt.sdk-test", + checkpoint_version: str = "sdk-test", + max_sample_count: int = 4096, +) -> dict[str, object]: + """Health.""" + return { + "status": "ok", + "schema_version": "v1", + "image_version": "0.2.0", + "model_version": model_version, + "checkpoint_version": checkpoint_version, + "checkpoint_path": "/models/jointfm.pt", + "device": "cpu", + "head": "studentt", + "decoding_strategy": "parallel_dense", + "supported_query_modes": ["forecast"], + "supported_return_modes": ["mean", "samples", "quantiles", "log_prob"], + "supported_time_index_modes": [ + "ordinal", + "continuous_float", + "absolute_datetime", + ], + "time_index_encoding": "legacy_discrete_grid", + "max_sample_count": max_sample_count, + } + + +class _Transport: + """Transport that serves health/predict per deployment id.""" + + def __init__( + self, + *, + health_by_id: Mapping[str, Mapping[str, Any]] | None = None, + fail_ids: frozenset[str] = frozenset(), + fail_status: int = 470, + ) -> None: + self.health_by_id = dict(health_by_id or {}) + self.fail_ids = fail_ids + self.fail_status = fail_status + + def get_json(self, url: str) -> Mapping[str, Any]: + """Get json.""" + raise AssertionError(f"unexpected GET {url}") + + def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Post json.""" + deployment_id = url.rstrip("/").split("/")[-2] + if deployment_id in self.fail_ids: + raise JointFMHTTPStatusError( + f"{deployment_id} unavailable", + status_code=self.fail_status, + response_body_excerpt="unavailable", + ) + if payload.get("request_type") == "health": + return self.health_by_id.get(deployment_id, _health()) + return {"ok": True, "deployment_id": deployment_id} + + +def _pool( + *, + instances: Sequence[JointFMInstanceSettings] | None = None, + transport: JSONTransport | None = None, + transports: Sequence[JSONTransport] | None = None, + peer_cooldown_seconds: float = 30.0, + expected_model_version: str | None = None, +) -> JointFMInstancePool: + peers = tuple(instances or (_instance("a"), _instance("b"))) + if transports is None: + shared = transport or _Transport() + transports = tuple(shared for _ in peers) + return JointFMInstancePool( + instances=peers, + transports=transports, + peer_cooldown_seconds=peer_cooldown_seconds, + expected_model_version=expected_model_version, + ) + + +def test_pool_retries_next_instance_on_470() -> None: + """Pool retries next instance on 470.""" + pool = _pool(transport=_Transport(fail_ids=frozenset({"a"}))) + assert pool.next_instance().deployment_id == "a" + assert pool.next_instance().deployment_id == "b" + assert pool.post_json({"schema_version": "v1"}) == { + "ok": True, + "deployment_id": "b", + } + + +def test_pool_raises_when_all_instances_unavailable() -> None: + """Pool raises when all instances unavailable.""" + pool = _pool(transport=_Transport(fail_ids=frozenset({"a", "b"}))) + with pytest.raises(JointFMHTTPStatusError, match="unavailable"): + pool.post_json({"schema_version": "v1"}) + + +def test_pool_health_rejects_mismatch_and_aligns_sample_cap() -> None: + """Pool health rejects mismatch and aligns sample cap.""" + mismatched = _pool( + transport=_Transport( + health_by_id={ + "a": _health(), + "b": _health(model_version="jointfm-inference:9.9.9+ckpt.other"), + } + ) + ) + with pytest.raises(UnsupportedModelVersionError, match="model_version"): + mismatched.probe_all_health() + assert {mismatched.instance_at(i).deployment_id for i in range(2)} == {"a", "b"} + + with pytest.raises(UnsupportedServiceContractError, match="checkpoint_version"): + _pool( + transport=_Transport( + health_by_id={ + "a": _health(checkpoint_version="ckpt-a"), + "b": _health(checkpoint_version="ckpt-b"), + } + ) + ).probe_all_health() + + metadata = _pool( + transport=_Transport( + health_by_id={ + "a": _health(max_sample_count=100), + "b": _health(max_sample_count=40), + } + ) + ).probe_all_health() + assert metadata.max_sample_count == 40 + + +def test_pool_posts_concurrent_across_peers() -> None: + """Distinct peer transports allow two POSTs to be in flight at once.""" + gate = threading.Barrier(2, timeout=2.0) + + class _BlockingTransport: + def get_json(self, url: str) -> Mapping[str, Any]: + raise AssertionError(f"unexpected GET {url}") + + def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: + del payload + gate.wait() + return {"ok": True, "url": url} + + instances = (_instance("a"), _instance("b")) + pool = _pool( + instances=instances, + transports=(_BlockingTransport(), _BlockingTransport()), + ) + + with ThreadPoolExecutor(max_workers=2) as executor: + futures = [ + executor.submit( + pool.post_json_to, + pool.instance_at(index), + {"schema_version": "v1"}, + ) + for index in range(2) + ] + results = [future.result(timeout=2.0) for future in futures] + + assert {result["url"] for result in results} == { + instances[0].predict_url, + instances[1].predict_url, + } + + +def test_pool_health_routes_only_reachable_peers() -> None: + """After health, sticky/RR skip peers that failed the probe.""" + pool = _pool( + transport=_Transport( + health_by_id={"b": _health()}, + fail_ids=frozenset({"a"}), + ) + ) + pool.probe_all_health() + assert pool.instance_at(0).deployment_id == "b" + assert pool.instance_at(1).deployment_id == "b" + assert pool.next_instance().deployment_id == "b" + assert pool.post_json({"schema_version": "v1"})["deployment_id"] == "b" + + +def test_pool_failover_retries_health_excluded_peer() -> None: + """When the last health-active peer fails, failover retries a recovered peer.""" + transport = _Transport( + health_by_id={"b": _health()}, + fail_ids=frozenset({"a"}), + ) + pool = _pool(transport=transport) + pool.probe_all_health() + assert pool.instance_at(0).deployment_id == "b" + + transport.fail_ids = frozenset({"b"}) + assert pool.post_json({"schema_version": "v1"})["deployment_id"] == "a" + assert pool.instance_at(0).deployment_id == "a" + + +def test_pool_health_skips_incompatible_peer_when_another_matches_pin() -> None: + """A pinned-incompatible backup is skipped; the matching primary stays usable.""" + pinned = "jointfm-inference:0.2.0+ckpt.sdk-test" + pool = _pool( + transport=_Transport( + health_by_id={ + "a": _health(model_version=pinned), + "b": _health(model_version="jointfm-inference:9.9.9+ckpt.other"), + } + ), + expected_model_version=pinned, + ) + metadata = pool.probe_all_health() + assert metadata.model_version == pinned + assert pool.instance_at(0).deployment_id == "a" + assert pool.instance_at(1).deployment_id == "a" + assert pool.post_json({"schema_version": "v1"})["deployment_id"] == "a" + + +def test_pool_cooldown_restores_peer_after_transient_failure( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Retryable POST failures cool a peer down; it returns after the cooldown.""" + clock = {"now": 100.0} + monkeypatch.setattr("jointfm_client.pool.time.monotonic", lambda: clock["now"]) + transport = _Transport(fail_ids=frozenset({"a"})) + pool = _pool(transport=transport, peer_cooldown_seconds=10.0) + + assert pool.post_json({"schema_version": "v1"})["deployment_id"] == "b" + assert pool.instance_at(0).deployment_id == "b" + assert pool.instance_at(1).deployment_id == "b" + + transport.fail_ids = frozenset() + clock["now"] = 109.0 + assert pool.instance_at(0).deployment_id == "b" + + clock["now"] = 110.0 + assert {pool.instance_at(i).deployment_id for i in range(2)} == {"a", "b"} + assert pool.post_json({"schema_version": "v1"})["deployment_id"] in {"a", "b"} diff --git a/tests/test_settings.py b/tests/test_settings.py index fa2623e..36825e3 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -22,6 +22,7 @@ DATAROBOT_API_TOKEN_ENV, DATAROBOT_ENDPOINT_ENV, JOINTFM_DEPLOYMENT_ID_ENV, + JOINTFM_DEPLOYMENT_IDS_ENV, JOINTFM_DEPLOYMENT_TARGET_ENV, JOINTFM_DEPLOYMENT_URL_ENV, JOINTFM_LOCAL_BASE_URL_ENV, @@ -502,3 +503,45 @@ def test_jointfm_client_from_env_attaches_non_secret_settings() -> None: assert client.settings is not None assert client.settings.predict_url.endswith("/predictionsUnstructured") assert "secret-token" not in repr(client) + + +def test_load_settings_with_deployment_ids_builds_instance_pool() -> None: + """Load settings with deployment ids builds instance pool.""" + env = _hosted_env(**{JOINTFM_DEPLOYMENT_IDS_ENV: "primary-id,backup-id"}) + del env[JOINTFM_DEPLOYMENT_ID_ENV] + + settings = load_settings(env=env, dotenv_path=None) + + assert settings.deployment_selector == "deployment_ids" + assert [instance.deployment_id for instance in settings.instances] == [ + "primary-id", + "backup-id", + ] + assert settings.instances[0].predict_url.endswith( + "/deployments/primary-id/predictionsUnstructured" + ) + assert settings.instances[1].predict_url.endswith( + "/deployments/backup-id/predictionsUnstructured" + ) + + +def test_load_settings_rejects_deployment_ids_combined_with_deployment_id() -> None: + """Load settings rejects deployment ids combined with deployment id.""" + with pytest.raises(JointFMConfigurationError, match="cannot be combined"): + load_settings( + env=_hosted_env(**{JOINTFM_DEPLOYMENT_IDS_ENV: "primary-id,backup-id"}), + dotenv_path=None, + ) + + +def test_load_settings_rejects_deployment_ids_with_fewer_than_two_unique_ids() -> None: + """Load settings rejects deployment ids with fewer than two unique ids.""" + env = _hosted_env(**{JOINTFM_DEPLOYMENT_IDS_ENV: "only-id"}) + del env[JOINTFM_DEPLOYMENT_ID_ENV] + with pytest.raises(JointFMConfigurationError, match="at least two unique"): + load_settings(env=env, dotenv_path=None) + + env = _hosted_env(**{JOINTFM_DEPLOYMENT_IDS_ENV: "same-id,same-id"}) + del env[JOINTFM_DEPLOYMENT_ID_ENV] + with pytest.raises(JointFMConfigurationError, match="at least two unique"): + load_settings(env=env, dotenv_path=None) diff --git a/tests/test_transport.py b/tests/test_transport.py index 4eb2d7d..0c96983 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -39,6 +39,7 @@ JointFMRequestError, JointFMHTTPStatusError, JointFMHTTPTransport, + JointFMInstanceSettings, JointFMResponseDecodeError, JointFMServiceError, JointFMRetryConfig, @@ -47,6 +48,7 @@ MeanForecastResult, SampleForecastResult, UnsupportedModelVersionError, + UnsupportedServiceContractError, ) @@ -143,7 +145,9 @@ def request(self, *args: Any, **kwargs: Any) -> requests.Response: def _health_payload( - *, model_version: str = "jointfm-inference:0.2.0+ckpt.sdk-test" + *, + model_version: str = "jointfm-inference:0.2.0+ckpt.sdk-test", + checkpoint_version: str = "sdk-test", ) -> dict[str, object]: """Health payload.""" return { @@ -151,7 +155,7 @@ def _health_payload( "schema_version": "v1", "image_version": "0.2.0", "model_version": model_version, - "checkpoint_version": "sdk-test", + "checkpoint_version": checkpoint_version, "checkpoint_path": "/models/jointfm.pt", "device": "cpu", "head": "studentt", @@ -1181,3 +1185,148 @@ def _server_url(server: ThreadingHTTPServer) -> str: host = server.server_address[0] port = server.server_address[1] return f"http://{host}:{port}/predict" + + +def _pool_settings(primary: str, backup: str) -> JointFMSettings: + """Pool settings.""" + return JointFMSettings( + datarobot_endpoint="https://app.datarobot.com/api/v2", + datarobot_api_token="secret-token", + health_url=primary, + predict_url=primary, + deployment_selector="deployment_ids", + schema_version="v1", + instances=( + JointFMInstanceSettings(deployment_id="primary-id", predict_url=primary), + JointFMInstanceSettings(deployment_id="backup-id", predict_url=backup), + ), + model_version="jointfm-inference:0.2.0+ckpt.sdk-test", + deployment_id="primary-id", + ) + + +def test_client_pool_health_gate_and_round_robin() -> None: + """Client pool health gate and round robin.""" + primary = ( + "https://app.datarobot.com/api/v2/deployments/" + "primary-id/predictionsUnstructured" + ) + backup = ( + "https://app.datarobot.com/api/v2/deployments/backup-id/predictionsUnstructured" + ) + settings = _pool_settings(primary, backup) + + class PoolTransport: + """Pool Transport (test helper).""" + + def __init__(self) -> None: + """Init.""" + self.urls: list[str] = [] + + def get_json(self, url: str) -> Mapping[str, Any]: + """Get json.""" + raise AssertionError(f"unexpected GET {url}") + + def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Post json.""" + self.urls.append(url) + if payload.get("request_type") == "health": + return _health_payload() + return _forecast_response_payload() + + transport = PoolTransport() + client = JointFMClient(settings=settings, transport=transport) + payload = { + "schema_version": "v1", + "model_version": "jointfm-inference:0.2.0+ckpt.sdk-test", + } + client.predict(payload) + client.predict(payload) + assert transport.urls[:2] == [primary, backup] + assert transport.urls[2:] == [primary, backup] + assert client._health_metadata is not None + + class MismatchTransport: + """Mismatch Transport (test helper).""" + + def get_json(self, url: str) -> Mapping[str, Any]: + """Get json.""" + raise AssertionError(f"unexpected GET {url}") + + def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Post json.""" + if payload.get("request_type") != "health": + raise AssertionError("predict must not run before pool health gate") + if url == primary: + return _health_payload(checkpoint_version="ckpt-a") + return _health_payload(checkpoint_version="ckpt-b") + + with pytest.raises(UnsupportedServiceContractError, match="checkpoint_version"): + JointFMClient(settings=settings, transport=MismatchTransport()).predict(payload) + + +def test_client_pool_forecast_samples_batches_across_peers() -> None: + """Pool sample batching pins batches to peers and merges the full sample count.""" + primary = ( + "https://app.datarobot.com/api/v2/deployments/" + "primary-id/predictionsUnstructured" + ) + backup = ( + "https://app.datarobot.com/api/v2/deployments/backup-id/predictionsUnstructured" + ) + settings = _pool_settings(primary, backup) + predict_urls: list[str] = [] + + class PoolSampleTransport: + """Pool Sample Transport (test helper).""" + + def get_json(self, url: str) -> Mapping[str, Any]: + """Get json.""" + raise AssertionError(f"unexpected GET {url}") + + def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """Post json.""" + if payload.get("request_type") == "health": + health = _health_payload() + health["max_sample_count"] = 2 + return health + predict_urls.append(url) + sample_count = cast(int, payload["n_samples"]) + seed = cast(int, payload["seed"]) + base = (seed - 7) * sample_count + samples = [ + [[float(sample_index)]] + for sample_index in range(base, base + sample_count) + ] + response_payload = _forecast_response_payload(return_mode="samples") + outputs = cast(dict[str, object], response_payload["outputs"]) + outputs["samples"] = samples + diagnostics = cast(dict[str, object], response_payload["diagnostics"]) + diagnostics["seed"] = payload.get("seed") + return response_payload + + client = JointFMClient(settings=settings, transport=PoolSampleTransport()) + schema = DataFrameSchema( + columns=(ColumnSpec(name="target", modality="numeric", role="target"),), + time_index_mode="ordinal", + ) + result = client.forecast_samples( + [{"target": 10.0}, {"target": 11.0}], + schema=schema, + query_times=[2], + requested_columns=["target"], + model_version="jointfm-inference:0.2.0+ckpt.sdk-test", + n_samples=4, + seed=7, + ) + + assert isinstance(result, SampleForecastResult) + assert len(result.samples) == 4 + assert result.samples == ( + ((0.0,),), + ((1.0,),), + ((2.0,),), + ((3.0,),), + ) + assert set(predict_urls) == {primary, backup} + assert len(predict_urls) == 2