From 1c090405119c95eeb764543b96d86b75dbd6c147 Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Tue, 25 Aug 2026 12:29:15 +0200 Subject: [PATCH 01/10] failover infra --- .env.sample | 4 + config.sample.yaml | 2 + src/jointfm_client/__init__.py | 6 ++ src/jointfm_client/client.py | 56 ++++++++--- src/jointfm_client/configuration.py | 5 + src/jointfm_client/pool.py | 149 ++++++++++++++++++++++++++++ src/jointfm_client/settings.py | 113 +++++++++++++++++++++ 7 files changed, 322 insertions(+), 13 deletions(-) create mode 100644 src/jointfm_client/pool.py diff --git a/.env.sample b/.env.sample index d72da9a..83b41d4 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). +# 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/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/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..6182f0b 100644 --- a/src/jointfm_client/client.py +++ b/src/jointfm_client/client.py @@ -57,6 +57,7 @@ QuantileForecastResult, SampleForecastResult, ) +from jointfm_client.pool import JointFMInstancePool from jointfm_client.settings import ( JointFMSettings, load_settings, @@ -108,6 +109,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 +154,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 multiple hosted deployments are configured via + ``JOINTFM_DEPLOYMENT_IDS``, every instance is probed and must advertise + the same ``model_version``. """ 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 +208,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 +256,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 +316,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,7 +495,6 @@ 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: @@ -498,10 +508,7 @@ def _forecast_sample_batches( 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, - ) + response_payload = self._post_predict_json(batch_payload) batch_result = _forecast_response_from_payload( response_payload, batch_payload, @@ -522,6 +529,29 @@ def _forecast_sample_batches( f"JointFM forecast response violated the V1 contract: {error}" ) from error + 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 + self._pool = JointFMInstancePool( + instances=self.settings.instances, + transport=self._transport_for_request(), + expected_model_version=self.settings.model_version, + ) + return self._pool + def _require_settings(self, method_name: str) -> JointFMSettings: if self.settings is None: raise JointFMConfigurationError( 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..dea319a --- /dev/null +++ b/src/jointfm_client/pool.py @@ -0,0 +1,149 @@ +# 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 +import logging +import threading +from typing import Any + +from jointfm_client.contract import ( + HEALTH_REQUEST_TYPE, + HealthMetadata, + validate_service_metadata, +) +from jointfm_client.exceptions import ( + JointFMHTTPStatusError, + JointFMRequestError, + UnsupportedModelVersionError, +) +from jointfm_client.settings import JointFMInstanceSettings +from jointfm_client.transport import JSONTransport + +logger = logging.getLogger(__name__) + +_POOL_RETRYABLE_HTTP_STATUS_CODES = frozenset({502, 503, 504}) + + +class JointFMInstancePool: + """Distributes JointFM requests across multiple hosted deployment instances.""" + + def __init__( + self, + *, + instances: Sequence[JointFMInstanceSettings], + transport: JSONTransport, + expected_model_version: str | None = None, + ) -> None: + if len(instances) < 2: + raise ValueError("JointFMInstancePool requires at least two instances") + self._instances = tuple(instances) + self._transport = transport + self._expected_model_version = expected_model_version + self._lock = threading.Lock() + self._index = 0 + + def next_instance(self) -> JointFMInstanceSettings: + """Return the next instance using round-robin selection.""" + with self._lock: + instance = self._instances[self._index] + self._index = (self._index + 1) % len(self._instances) + return instance + + def probe_all_health(self) -> HealthMetadata: + """Probe instances; log unavailable ones and require matching model_version.""" + metadata_by_instance: list[tuple[JointFMInstanceSettings, HealthMetadata]] = [] + last_error: BaseException | None = None + for instance in self._instances: + try: + payload = self._fetch_hosted_health_payload(instance) + validate_service_metadata( + payload, + expected_model_version=self._expected_model_version, + ) + metadata_by_instance.append( + (instance, HealthMetadata.from_payload(payload)) + ) + except Exception as error: + if not _should_retry_on_next_instance(error): + raise + last_error = error + logger.warning( + "JointFM instance unavailable: deployment_id=%s error=%s", + instance.deployment_id, + error, + ) + + if not metadata_by_instance: + assert last_error is not None + raise last_error + + reference = metadata_by_instance[0][1] + for instance, metadata in metadata_by_instance[1:]: + if metadata.model_version != reference.model_version: + raise UnsupportedModelVersionError( + "JointFM deployment pool model_version mismatch: " + f"{instance.deployment_id!r} advertises {metadata.model_version!r}, " + f"expected {reference.model_version!r}" + ) + return reference + + def post_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + """POST one payload, trying untried instances on retryable failures.""" + tried: set[str] = set() + last_error: BaseException | None = None + while len(tried) < len(self._instances): + instance = self._next_untried_instance(tried) + tried.add(instance.deployment_id) + try: + return self._transport.post_json(instance.predict_url, payload) + except Exception as error: + if not _should_retry_on_next_instance(error): + raise + last_error = error + logger.warning( + "JointFM instance unavailable: deployment_id=%s error=%s", + instance.deployment_id, + error, + ) + assert last_error is not None + raise last_error + + def _next_untried_instance(self, tried: set[str]) -> JointFMInstanceSettings: + for _ in range(len(self._instances)): + instance = self.next_instance() + if instance.deployment_id not in tried: + return instance + raise RuntimeError("JointFMInstancePool has no untried instances") + + def _fetch_hosted_health_payload( + self, + instance: JointFMInstanceSettings, + ) -> Mapping[str, Any]: + return self._transport.post_json( + instance.predict_url, + {"request_type": HEALTH_REQUEST_TYPE}, + ) + + +def _should_retry_on_next_instance(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..ede7f47 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,15 @@ ] +@dataclass(frozen=True, slots=True) +class JointFMInstanceSettings: + """One hosted JointFM deployment target in a load-balanced pool.""" + + deployment_id: str + predict_url: str + health_url: str + + @dataclass(frozen=True, slots=True) class JointFMSettings: """Validated settings for one hosted or local JointFM service target.""" @@ -69,6 +80,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 +108,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: @@ -494,6 +516,97 @@ def _resolve_single_deployment_selector( 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_list: list[JointFMInstanceSettings] = [] + for deployment_id in deployment_ids: + predict_url = build_hosted_predict_url(datarobot_endpoint, deployment_id) + instances_list.append( + JointFMInstanceSettings( + deployment_id=deployment_id, + predict_url=predict_url, + health_url=predict_url, + ) + ) + instances = tuple(instances_list) + primary = instances[0] + return JointFMSettings( + datarobot_endpoint=datarobot_endpoint, + datarobot_api_token=datarobot_api_token, + health_url=primary.health_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: + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must not contain leading or trailing whitespace" + ) + deployment_ids: list[str] = [] + seen: set[str] = set() + for part in value.split(","): + if part == "": + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must not contain empty deployment IDs" + ) + if any(character.isspace() for character in part): + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must not contain whitespace around deployment IDs" + ) + if "/" in part: + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must contain deployment IDs, not URLs" + ) + if part in seen: + continue + seen.add(part) + deployment_ids.append(part) + if not deployment_ids: + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must contain at least one deployment ID" + ) + return tuple(deployment_ids) + + +def _reject_conflicting_selectors_with_deployment_ids( + env: Mapping[str, str], environment: EnvironmentVariableConfig +) -> None: + conflicting_selectors = [ + selector_name + for selector_name in environment.deployment_selector_names() + if selector_name in env and env[selector_name] != "" + ] + if conflicting_selectors: + formatted_selectors = ", ".join(conflicting_selectors) + raise JointFMConfigurationError( + f"{JOINTFM_DEPLOYMENT_IDS_ENV} cannot be combined with {formatted_selectors}" + ) + + def _load_pulumi_target_outputs( outputs_path: str, deployment_target: str, From c39e2e163494962d1227654fac63c85ef4488582 Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Tue, 25 Aug 2026 12:48:19 +0200 Subject: [PATCH 02/10] bug fixes --- README.md | 16 ++++++++++++++-- docs/api-reference.md | 5 +++-- src/jointfm_client/client.py | 22 +++++++++++++++++++--- src/jointfm_client/pool.py | 33 ++++++++++++++++++++++++++++++--- 4 files changed, 66 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index cc0248a..3d71cc5 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. Use either a single-target selector (`JOINTFM_DEPLOYMENT_ID`, `JOINTFM_DEPLOYMENT_URL`, `JOINTFM_PREDICT_URL`, `JOINTFM_DEPLOYMENT_TARGET`, or `JOINTFM_LOCAL_BASE_URL`) or the load-balanced hosted selector `JOINTFM_DEPLOYMENT_IDS` (comma-separated deployment IDs of the same checkpoint). `JOINTFM_DEPLOYMENT_IDS` cannot be combined with another selector. `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,11 +86,21 @@ 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 for round-robin load balancing. Peers must advertise the same `model_version` and `checkpoint_version`; unavailable peers are logged and skipped. Mutually exclusive with the 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` - `JOINTFM_LOCAL_BASE_URL`: builds direct local `GET /healthz` and `POST /predict` URLs without DataRobot authentication +Load-balanced hosted example: + +```dotenv +DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2 +DATAROBOT_API_TOKEN= +JOINTFM_SCHEMA_VERSION=v1 +JOINTFM_DEPLOYMENT_IDS=chevron-id,research-id +``` + Pulumi deployment discovery is explicit and file-backed. Export stack outputs to a JSON object keyed by target name, then set `JOINTFM_DEPLOYMENT_TARGET` to the key and `JOINTFM_PULUMI_OUTPUTS_PATH` to that JSON file: ```json @@ -238,13 +248,15 @@ uv run python -c "import jointfm_client; print(jointfm_client.__version__)" ### Configure A Deployment -Create `.env` from `.env.sample` or set the same values in your shell. A hosted forecast needs the DataRobot API v2 endpoint, token, schema pin, and exactly one deployment selector: +Create `.env` from `.env.sample` or set the same values in your shell. A hosted forecast needs the DataRobot API v2 endpoint, token, schema pin, and exactly one deployment selector (`JOINTFM_DEPLOYMENT_ID` or `JOINTFM_DEPLOYMENT_IDS`, among the other selectors listed above): ```dotenv DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2 DATAROBOT_API_TOKEN= JOINTFM_SCHEMA_VERSION=v1 JOINTFM_DEPLOYMENT_ID= +# Or load-balance same-checkpoint peers: +# 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/docs/api-reference.md b/docs/api-reference.md index 6f10e74..69d13b1 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. 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/client.py b/src/jointfm_client/client.py index 6182f0b..4b08fcd 100644 --- a/src/jointfm_client/client.py +++ b/src/jointfm_client/client.py @@ -156,8 +156,10 @@ def health(self, *, cache: bool = False, refresh: bool = False) -> HealthMetadat validation and returns the same typed health payload. When multiple hosted deployments are configured via - ``JOINTFM_DEPLOYMENT_IDS``, every instance is probed and must advertise - the same ``model_version``. + ``JOINTFM_DEPLOYMENT_IDS``, every reachable instance is probed and must + advertise the same ``model_version`` and ``checkpoint_version``. The + cached 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 @@ -545,9 +547,23 @@ def _require_pool(self) -> JointFMInstancePool: ) if self._pool is None: assert self.settings is not None + # Per-instance transport retries would delay moving to the next peer + # under outage; the pool itself retries across instances. + if self._transport is None: + pool_transport: JSONTransport = 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, + ) + else: + pool_transport = self._transport self._pool = JointFMInstancePool( instances=self.settings.instances, - transport=self._transport_for_request(), + transport=pool_transport, expected_model_version=self.settings.model_version, ) return self._pool diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py index dea319a..3917a56 100644 --- a/src/jointfm_client/pool.py +++ b/src/jointfm_client/pool.py @@ -16,6 +16,7 @@ from __future__ import annotations +from dataclasses import replace from collections.abc import Mapping, Sequence import logging import threading @@ -30,6 +31,7 @@ JointFMHTTPStatusError, JointFMRequestError, UnsupportedModelVersionError, + UnsupportedServiceContractError, ) from jointfm_client.settings import JointFMInstanceSettings from jointfm_client.transport import JSONTransport @@ -40,7 +42,12 @@ class JointFMInstancePool: - """Distributes JointFM requests across multiple hosted deployment instances.""" + """Distributes JointFM requests across multiple hosted deployment instances. + + Callers should pass a fail-fast transport (``max_attempts=1``). Per-instance + transport retries would delay moving to the next peer under outage; this pool + retries across instances after logging unavailable ones. + """ def __init__( self, @@ -65,7 +72,13 @@ def next_instance(self) -> JointFMInstanceSettings: return instance def probe_all_health(self) -> HealthMetadata: - """Probe instances; log unavailable ones and require matching model_version.""" + """Probe instances; require matching checkpoint identity across healthy peers. + + Unavailable instances are logged and skipped. Healthy peers must share + ``model_version`` and ``checkpoint_version``. The returned metadata uses + the minimum ``max_sample_count`` across those peers so sample batching + stays within every instance's advertised cap. + """ metadata_by_instance: list[tuple[JointFMInstanceSettings, HealthMetadata]] = [] last_error: BaseException | None = None for instance in self._instances: @@ -93,6 +106,7 @@ def probe_all_health(self) -> HealthMetadata: raise last_error reference = metadata_by_instance[0][1] + aligned_max_sample_count = reference.max_sample_count for instance, metadata in metadata_by_instance[1:]: if metadata.model_version != reference.model_version: raise UnsupportedModelVersionError( @@ -100,7 +114,20 @@ def probe_all_health(self) -> HealthMetadata: f"{instance.deployment_id!r} advertises {metadata.model_version!r}, " f"expected {reference.model_version!r}" ) - return reference + if metadata.checkpoint_version != reference.checkpoint_version: + raise UnsupportedServiceContractError( + "JointFM deployment pool checkpoint_version mismatch: " + f"{instance.deployment_id!r} advertises " + f"{metadata.checkpoint_version!r}, " + f"expected {reference.checkpoint_version!r}" + ) + aligned_max_sample_count = min( + aligned_max_sample_count, + metadata.max_sample_count, + ) + if aligned_max_sample_count == reference.max_sample_count: + return reference + return replace(reference, max_sample_count=aligned_max_sample_count) def post_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: """POST one payload, trying untried instances on retryable failures.""" From 929ed718692253b147f779c527b64ad61241bbce Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Tue, 25 Aug 2026 12:57:36 +0200 Subject: [PATCH 03/10] reject single value for JOINTFM_DEPLOYMENT_IDS --- .env.sample | 2 +- README.md | 2 +- docs/api-reference.md | 2 +- src/jointfm_client/pool.py | 5 ++++- src/jointfm_client/settings.py | 8 +++++--- 5 files changed, 12 insertions(+), 7 deletions(-) diff --git a/.env.sample b/.env.sample index 83b41d4..395cb55 100644 --- a/.env.sample +++ b/.env.sample @@ -13,7 +13,7 @@ JOINTFM_SCHEMA_VERSION=v1 # JOINTFM_DEPLOYMENT_ID= # 1b. Comma-separated deployment IDs for load-balanced hosted calls (same checkpoint). -# Mutually exclusive with JOINTFM_DEPLOYMENT_ID and other selectors. +# 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. diff --git a/README.md b/README.md index 3d71cc5..182a6b1 100644 --- a/README.md +++ b/README.md @@ -86,7 +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 for round-robin load balancing. Peers must advertise the same `model_version` and `checkpoint_version`; unavailable peers are logged and skipped. Mutually exclusive with the other selectors. +- `JOINTFM_DEPLOYMENT_IDS`: comma-separated hosted deployment IDs for round-robin load balancing (at least two unique IDs). Peers must advertise the same `model_version` and `checkpoint_version`; unavailable peers are logged and skipped. Mutually exclusive with the 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` diff --git a/docs/api-reference.md b/docs/api-reference.md index 69d13b1..380c9dc 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -119,7 +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. Mutually exclusive with other selectors. Peers must share `model_version` and `checkpoint_version`; the SDK uses the minimum `max_sample_count`. | +| `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. | diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py index 3917a56..1802f69 100644 --- a/src/jointfm_client/pool.py +++ b/src/jointfm_client/pool.py @@ -22,6 +22,7 @@ import threading from typing import Any +from jointfm_client.configuration import DEFAULT_RETRY_STATUS_CODES from jointfm_client.contract import ( HEALTH_REQUEST_TYPE, HealthMetadata, @@ -38,7 +39,9 @@ logger = logging.getLogger(__name__) -_POOL_RETRYABLE_HTTP_STATUS_CODES = frozenset({502, 503, 504}) +# Match the transport's retryable statuses so pool failover covers DR 470 +# (stopped / warming deployment) and other transient codes, not only gateways. +_POOL_RETRYABLE_HTTP_STATUS_CODES = frozenset(DEFAULT_RETRY_STATUS_CODES) class JointFMInstancePool: diff --git a/src/jointfm_client/settings.py b/src/jointfm_client/settings.py index ede7f47..89f5591 100644 --- a/src/jointfm_client/settings.py +++ b/src/jointfm_client/settings.py @@ -509,7 +509,9 @@ 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}" ) @@ -585,9 +587,9 @@ def _parse_deployment_ids(value: str) -> tuple[str, ...]: continue seen.add(part) deployment_ids.append(part) - if not deployment_ids: + if len(deployment_ids) < 2: raise JointFMConfigurationError( - f"{JOINTFM_DEPLOYMENT_IDS_ENV} must contain at least one deployment ID" + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must contain at least two unique deployment IDs" ) return tuple(deployment_ids) From 8c913cf531df127576f3d8d90d8ebdfbf117d5d8 Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Tue, 25 Aug 2026 13:00:21 +0200 Subject: [PATCH 04/10] add tests --- tests/test_pool.py | 183 ++++++++++++++++++++++++++++++++++++++++ tests/test_settings.py | 43 ++++++++++ tests/test_transport.py | 64 ++++++++++++++ 3 files changed, 290 insertions(+) create mode 100644 tests/test_pool.py diff --git a/tests/test_pool.py b/tests/test_pool.py new file mode 100644 index 0000000..7d1d3d7 --- /dev/null +++ b/tests/test_pool.py @@ -0,0 +1,183 @@ +# 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 +from typing import Any + +import pytest + +from jointfm_client import ( + JointFMHTTPStatusError, + JointFMInstancePool, + JointFMInstanceSettings, + JointFMRequestError, + UnsupportedModelVersionError, + UnsupportedServiceContractError, +) + + +def _instance(deployment_id: str) -> JointFMInstanceSettings: + """Instance.""" + url = ( + "https://app.datarobot.com/api/v2/deployments/" + f"{deployment_id}/predictionsUnstructured" + ) + return JointFMInstanceSettings( + deployment_id=deployment_id, + predict_url=url, + health_url=url, + ) + + +def _health_payload( + *, + 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 payload.""" + 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 _PoolTransport: + """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 + 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) + deployment_id = url.rstrip("/").split("/")[-2] + if deployment_id in self.fail_ids: + if self.fail_status < 0: + raise JointFMRequestError(f"{deployment_id} unreachable") + 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_payload()) + return {"ok": True, "deployment_id": deployment_id} + + +def test_pool_round_robin_and_retries_next_instance_on_470() -> None: + """Pool round robin and retries next instance on 470.""" + pool = JointFMInstancePool( + instances=(_instance("a"), _instance("b")), + transport=_PoolTransport(fail_ids=frozenset({"a"}), fail_status=470), + ) + + assert pool.next_instance().deployment_id == "a" + assert pool.next_instance().deployment_id == "b" + assert pool.next_instance().deployment_id == "a" + + result = pool.post_json({"schema_version": "v1"}) + + assert result == {"ok": True, "deployment_id": "b"} + + +def test_pool_raises_when_all_instances_unavailable() -> None: + """Pool raises when all instances unavailable.""" + transport = _PoolTransport(fail_ids=frozenset({"a", "b"}), fail_status=470) + pool = JointFMInstancePool( + instances=(_instance("a"), _instance("b")), + transport=transport, + ) + + with pytest.raises(JointFMHTTPStatusError, match="unavailable"): + pool.post_json({"schema_version": "v1"}) + + +def test_pool_health_rejects_model_or_checkpoint_mismatch() -> None: + """Pool health rejects model or checkpoint mismatch.""" + pool = JointFMInstancePool( + instances=(_instance("a"), _instance("b")), + transport=_PoolTransport( + health_by_id={ + "a": _health_payload(), + "b": _health_payload( + model_version="jointfm-inference:9.9.9+ckpt.other" + ), + } + ), + ) + with pytest.raises(UnsupportedModelVersionError, match="model_version"): + pool.probe_all_health() + + pool = JointFMInstancePool( + instances=(_instance("a"), _instance("b")), + transport=_PoolTransport( + health_by_id={ + "a": _health_payload(checkpoint_version="ckpt-a"), + "b": _health_payload(checkpoint_version="ckpt-b"), + } + ), + ) + with pytest.raises(UnsupportedServiceContractError, match="checkpoint_version"): + pool.probe_all_health() + + +def test_pool_health_uses_minimum_max_sample_count() -> None: + """Pool health uses minimum max sample count.""" + pool = JointFMInstancePool( + instances=(_instance("a"), _instance("b")), + transport=_PoolTransport( + health_by_id={ + "a": _health_payload(max_sample_count=100), + "b": _health_payload(max_sample_count=40), + } + ), + ) + + metadata = pool.probe_all_health() + + assert metadata.max_sample_count == 40 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..eb29c48 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -39,6 +39,7 @@ JointFMRequestError, JointFMHTTPStatusError, JointFMHTTPTransport, + JointFMInstanceSettings, JointFMResponseDecodeError, JointFMServiceError, JointFMRetryConfig, @@ -1181,3 +1182,66 @@ def _server_url(server: ThreadingHTTPServer) -> str: host = server.server_address[0] port = server.server_address[1] return f"http://{host}:{port}/predict" + + +def test_client_predict_round_robins_across_pool_instances() -> None: + """Client predict round robins across pool instances.""" + primary = ( + "https://app.datarobot.com/api/v2/deployments/" + "primary-id/predictionsUnstructured" + ) + backup = ( + "https://app.datarobot.com/api/v2/deployments/backup-id/predictionsUnstructured" + ) + settings = 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, + health_url=primary, + ), + JointFMInstanceSettings( + deployment_id="backup-id", + predict_url=backup, + health_url=backup, + ), + ), + model_version="jointfm-inference:0.2.0+ckpt.sdk-test", + deployment_id="primary-id", + ) + + 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 == [primary, backup] From 2f1bd7563663e940cd5a40c93f58733198f93bb2 Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Tue, 25 Aug 2026 13:13:34 +0200 Subject: [PATCH 05/10] health prob before traffic --- src/jointfm_client/client.py | 4 +++ tests/test_transport.py | 69 ++++++++++++++++++++++++++++++++++-- 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/src/jointfm_client/client.py b/src/jointfm_client/client.py index 4b08fcd..739adb6 100644 --- a/src/jointfm_client/client.py +++ b/src/jointfm_client/client.py @@ -622,6 +622,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/tests/test_transport.py b/tests/test_transport.py index eb29c48..e2c8f07 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -48,6 +48,7 @@ MeanForecastResult, SampleForecastResult, UnsupportedModelVersionError, + UnsupportedServiceContractError, ) @@ -144,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 { @@ -152,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", @@ -1244,4 +1247,64 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: client.predict(payload) client.predict(payload) - assert transport.urls == [primary, backup] + assert transport.urls[:2] == [primary, backup] + assert transport.urls[2:] == [primary, backup] + assert client._health_metadata is not None + + +def test_client_pool_predict_with_pin_probes_health_before_traffic() -> None: + """Client pool predict with pin probes health before traffic.""" + primary = ( + "https://app.datarobot.com/api/v2/deployments/" + "primary-id/predictionsUnstructured" + ) + backup = ( + "https://app.datarobot.com/api/v2/deployments/backup-id/predictionsUnstructured" + ) + settings = 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, + health_url=primary, + ), + JointFMInstanceSettings( + deployment_id="backup-id", + predict_url=backup, + health_url=backup, + ), + ), + model_version="jointfm-inference:0.2.0+ckpt.sdk-test", + deployment_id="primary-id", + ) + + 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 "primary-id" in url: + return _health_payload(checkpoint_version="ckpt-a") + return _health_payload(checkpoint_version="ckpt-b") + + client = JointFMClient(settings=settings, transport=MismatchTransport()) + + with pytest.raises(UnsupportedServiceContractError, match="checkpoint_version"): + client.predict( + { + "schema_version": "v1", + "model_version": "jointfm-inference:0.2.0+ckpt.sdk-test", + } + ) From cbfc2a8d5add672824344f7af3e87cb1f7203ea4 Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Tue, 25 Aug 2026 13:21:47 +0200 Subject: [PATCH 06/10] cleanup --- README.md | 18 ++---- src/jointfm_client/client.py | 20 +++---- src/jointfm_client/pool.py | 85 ++++++++++---------------- src/jointfm_client/settings.py | 53 +++++++---------- tests/test_pool.py | 105 +++++++++++++-------------------- tests/test_transport.py | 83 ++++++-------------------- 6 files changed, 126 insertions(+), 238 deletions(-) diff --git a/README.md b/README.md index 182a6b1..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. Use either a single-target selector (`JOINTFM_DEPLOYMENT_ID`, `JOINTFM_DEPLOYMENT_URL`, `JOINTFM_PREDICT_URL`, `JOINTFM_DEPLOYMENT_TARGET`, or `JOINTFM_LOCAL_BASE_URL`) or the load-balanced hosted selector `JOINTFM_DEPLOYMENT_IDS` (comma-separated deployment IDs of the same checkpoint). `JOINTFM_DEPLOYMENT_IDS` cannot be combined with another selector. `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,21 +86,12 @@ 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 for round-robin load balancing (at least two unique IDs). Peers must advertise the same `model_version` and `checkpoint_version`; unavailable peers are logged and skipped. Mutually exclusive with the other selectors. +- `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` - `JOINTFM_LOCAL_BASE_URL`: builds direct local `GET /healthz` and `POST /predict` URLs without DataRobot authentication -Load-balanced hosted example: - -```dotenv -DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2 -DATAROBOT_API_TOKEN= -JOINTFM_SCHEMA_VERSION=v1 -JOINTFM_DEPLOYMENT_IDS=chevron-id,research-id -``` - Pulumi deployment discovery is explicit and file-backed. Export stack outputs to a JSON object keyed by target name, then set `JOINTFM_DEPLOYMENT_TARGET` to the key and `JOINTFM_PULUMI_OUTPUTS_PATH` to that JSON file: ```json @@ -248,15 +239,14 @@ uv run python -c "import jointfm_client; print(jointfm_client.__version__)" ### Configure A Deployment -Create `.env` from `.env.sample` or set the same values in your shell. A hosted forecast needs the DataRobot API v2 endpoint, token, schema pin, and exactly one deployment selector (`JOINTFM_DEPLOYMENT_ID` or `JOINTFM_DEPLOYMENT_IDS`, among the other selectors listed above): +Create `.env` from `.env.sample` or set the same values in your shell. A hosted forecast needs the DataRobot API v2 endpoint, token, schema pin, and exactly one deployment selector: ```dotenv DATAROBOT_ENDPOINT=https://app.datarobot.com/api/v2 DATAROBOT_API_TOKEN= JOINTFM_SCHEMA_VERSION=v1 JOINTFM_DEPLOYMENT_ID= -# Or load-balance same-checkpoint peers: -# JOINTFM_DEPLOYMENT_IDS=chevron-id,research-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/src/jointfm_client/client.py b/src/jointfm_client/client.py index 739adb6..b279856 100644 --- a/src/jointfm_client/client.py +++ b/src/jointfm_client/client.py @@ -155,11 +155,9 @@ def health(self, *, cache: bool = False, refresh: bool = False) -> HealthMetadat container short-circuits that body before any schema or model version validation and returns the same typed health payload. - When multiple hosted deployments are configured via - ``JOINTFM_DEPLOYMENT_IDS``, every reachable instance is probed and must - advertise the same ``model_version`` and ``checkpoint_version``. The - cached sample-batch cap is the minimum ``max_sample_count`` across those - peers. + 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 @@ -547,10 +545,10 @@ def _require_pool(self) -> JointFMInstancePool: ) if self._pool is None: assert self.settings is not None - # Per-instance transport retries would delay moving to the next peer - # under outage; the pool itself retries across instances. - if self._transport is None: - pool_transport: JSONTransport = JointFMHTTPTransport.from_settings( + transport = self._transport + if transport is None: + # Fail fast per peer; the pool retries across instances. + transport = JointFMHTTPTransport.from_settings( self.settings, timeout=self._timeout, retry_config=JointFMRetryConfig(max_attempts=1), @@ -559,11 +557,9 @@ def _require_pool(self) -> JointFMInstancePool: ), datarobot_request_id_headers=self._datarobot_request_id_headers, ) - else: - pool_transport = self._transport self._pool = JointFMInstancePool( instances=self.settings.instances, - transport=pool_transport, + transport=transport, expected_model_version=self.settings.model_version, ) return self._pool diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py index 1802f69..0d2883e 100644 --- a/src/jointfm_client/pool.py +++ b/src/jointfm_client/pool.py @@ -16,8 +16,8 @@ from __future__ import annotations -from dataclasses import replace from collections.abc import Mapping, Sequence +from dataclasses import replace import logging import threading from typing import Any @@ -39,17 +39,13 @@ logger = logging.getLogger(__name__) -# Match the transport's retryable statuses so pool failover covers DR 470 -# (stopped / warming deployment) and other transient codes, not only gateways. _POOL_RETRYABLE_HTTP_STATUS_CODES = frozenset(DEFAULT_RETRY_STATUS_CODES) class JointFMInstancePool: - """Distributes JointFM requests across multiple hosted deployment instances. + """Round-robin JointFM requests across hosted deployments. - Callers should pass a fail-fast transport (``max_attempts=1``). Per-instance - transport retries would delay moving to the next peer under outage; this pool - retries across instances after logging unavailable ones. + Use a fail-fast transport (``max_attempts=1``); this pool retries peers. """ def __init__( @@ -75,42 +71,33 @@ def next_instance(self) -> JointFMInstanceSettings: return instance def probe_all_health(self) -> HealthMetadata: - """Probe instances; require matching checkpoint identity across healthy peers. - - Unavailable instances are logged and skipped. Healthy peers must share - ``model_version`` and ``checkpoint_version``. The returned metadata uses - the minimum ``max_sample_count`` across those peers so sample batching - stays within every instance's advertised cap. - """ - metadata_by_instance: list[tuple[JointFMInstanceSettings, HealthMetadata]] = [] + """Probe peers; require matching model/checkpoint; return min sample cap.""" + healthy: list[tuple[JointFMInstanceSettings, HealthMetadata]] = [] last_error: BaseException | None = None for instance in self._instances: try: - payload = self._fetch_hosted_health_payload(instance) + payload = self._transport.post_json( + instance.predict_url, + {"request_type": HEALTH_REQUEST_TYPE}, + ) validate_service_metadata( payload, expected_model_version=self._expected_model_version, ) - metadata_by_instance.append( - (instance, HealthMetadata.from_payload(payload)) - ) + healthy.append((instance, HealthMetadata.from_payload(payload))) except Exception as error: - if not _should_retry_on_next_instance(error): + if not _is_pool_retryable(error): raise last_error = error - logger.warning( - "JointFM instance unavailable: deployment_id=%s error=%s", - instance.deployment_id, - error, - ) + _log_unavailable(instance.deployment_id, error) - if not metadata_by_instance: + if not healthy: assert last_error is not None raise last_error - reference = metadata_by_instance[0][1] - aligned_max_sample_count = reference.max_sample_count - for instance, metadata in metadata_by_instance[1:]: + reference = healthy[0][1] + max_samples = reference.max_sample_count + for instance, metadata in healthy[1:]: if metadata.model_version != reference.model_version: raise UnsupportedModelVersionError( "JointFM deployment pool model_version mismatch: " @@ -124,53 +111,45 @@ def probe_all_health(self) -> HealthMetadata: f"{metadata.checkpoint_version!r}, " f"expected {reference.checkpoint_version!r}" ) - aligned_max_sample_count = min( - aligned_max_sample_count, - metadata.max_sample_count, - ) - if aligned_max_sample_count == reference.max_sample_count: + max_samples = min(max_samples, metadata.max_sample_count) + if max_samples == reference.max_sample_count: return reference - return replace(reference, max_sample_count=aligned_max_sample_count) + return replace(reference, max_sample_count=max_samples) def post_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: - """POST one payload, trying untried instances on retryable failures.""" + """POST payload, trying untried peers on retryable failures.""" tried: set[str] = set() last_error: BaseException | None = None while len(tried) < len(self._instances): - instance = self._next_untried_instance(tried) + instance = self._next_untried(tried) tried.add(instance.deployment_id) try: return self._transport.post_json(instance.predict_url, payload) except Exception as error: - if not _should_retry_on_next_instance(error): + if not _is_pool_retryable(error): raise last_error = error - logger.warning( - "JointFM instance unavailable: deployment_id=%s error=%s", - instance.deployment_id, - error, - ) + _log_unavailable(instance.deployment_id, error) assert last_error is not None raise last_error - def _next_untried_instance(self, tried: set[str]) -> JointFMInstanceSettings: + def _next_untried(self, tried: set[str]) -> JointFMInstanceSettings: for _ in range(len(self._instances)): instance = self.next_instance() if instance.deployment_id not in tried: return instance raise RuntimeError("JointFMInstancePool has no untried instances") - def _fetch_hosted_health_payload( - self, - instance: JointFMInstanceSettings, - ) -> Mapping[str, Any]: - return self._transport.post_json( - instance.predict_url, - {"request_type": HEALTH_REQUEST_TYPE}, - ) + +def _log_unavailable(deployment_id: str, error: BaseException) -> None: + logger.warning( + "JointFM instance unavailable: deployment_id=%s error=%s", + deployment_id, + error, + ) -def _should_retry_on_next_instance(error: BaseException) -> bool: +def _is_pool_retryable(error: BaseException) -> bool: if isinstance(error, JointFMRequestError): return True return ( diff --git a/src/jointfm_client/settings.py b/src/jointfm_client/settings.py index 89f5591..78b5dd3 100644 --- a/src/jointfm_client/settings.py +++ b/src/jointfm_client/settings.py @@ -67,7 +67,6 @@ class JointFMInstanceSettings: deployment_id: str predict_url: str - health_url: str @dataclass(frozen=True, slots=True) @@ -534,22 +533,18 @@ def _load_hosted_deployment_pool_settings( deployment_ids = _parse_deployment_ids( _required_env(env, environment.deployment_ids) ) - instances_list: list[JointFMInstanceSettings] = [] - for deployment_id in deployment_ids: - predict_url = build_hosted_predict_url(datarobot_endpoint, deployment_id) - instances_list.append( - JointFMInstanceSettings( - deployment_id=deployment_id, - predict_url=predict_url, - health_url=predict_url, - ) + instances = tuple( + JointFMInstanceSettings( + deployment_id=deployment_id, + predict_url=build_hosted_predict_url(datarobot_endpoint, deployment_id), ) - instances = tuple(instances_list) + for deployment_id in deployment_ids + ) primary = instances[0] return JointFMSettings( datarobot_endpoint=datarobot_endpoint, datarobot_api_token=datarobot_api_token, - health_url=primary.health_url, + health_url=primary.predict_url, predict_url=primary.predict_url, deployment_selector="deployment_ids", schema_version=schema_version, @@ -564,24 +559,18 @@ def _load_hosted_deployment_pool_settings( def _parse_deployment_ids(value: str) -> tuple[str, ...]: - if value.strip() != value: + if value.strip() != value or "," not in value: raise JointFMConfigurationError( - f"{JOINTFM_DEPLOYMENT_IDS_ENV} must not contain leading or trailing whitespace" + 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 == "": - raise JointFMConfigurationError( - f"{JOINTFM_DEPLOYMENT_IDS_ENV} must not contain empty deployment IDs" - ) - if any(character.isspace() for character in part): + if part == "" or any(character.isspace() for character in part) or "/" in part: raise JointFMConfigurationError( - f"{JOINTFM_DEPLOYMENT_IDS_ENV} must not contain whitespace around deployment IDs" - ) - if "/" in part: - raise JointFMConfigurationError( - f"{JOINTFM_DEPLOYMENT_IDS_ENV} must contain deployment IDs, not URLs" + f"{JOINTFM_DEPLOYMENT_IDS_ENV} must contain non-empty deployment IDs " + "without whitespace or URL paths" ) if part in seen: continue @@ -595,17 +584,17 @@ def _parse_deployment_ids(value: str) -> tuple[str, ...]: def _reject_conflicting_selectors_with_deployment_ids( - env: Mapping[str, str], environment: EnvironmentVariableConfig + env: Mapping[str, str], + environment: EnvironmentVariableConfig, ) -> None: - conflicting_selectors = [ - selector_name - for selector_name in environment.deployment_selector_names() - if selector_name in env and env[selector_name] != "" + conflicting = [ + name + for name in environment.deployment_selector_names() + if name in env and env[name] != "" ] - if conflicting_selectors: - formatted_selectors = ", ".join(conflicting_selectors) + if conflicting: raise JointFMConfigurationError( - f"{JOINTFM_DEPLOYMENT_IDS_ENV} cannot be combined with {formatted_selectors}" + f"{JOINTFM_DEPLOYMENT_IDS_ENV} cannot be combined with {', '.join(conflicting)}" ) diff --git a/tests/test_pool.py b/tests/test_pool.py index 7d1d3d7..9279c0e 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -25,7 +25,6 @@ JointFMHTTPStatusError, JointFMInstancePool, JointFMInstanceSettings, - JointFMRequestError, UnsupportedModelVersionError, UnsupportedServiceContractError, ) @@ -33,24 +32,22 @@ def _instance(deployment_id: str) -> JointFMInstanceSettings: """Instance.""" - url = ( - "https://app.datarobot.com/api/v2/deployments/" - f"{deployment_id}/predictionsUnstructured" - ) return JointFMInstanceSettings( deployment_id=deployment_id, - predict_url=url, - health_url=url, + predict_url=( + "https://app.datarobot.com/api/v2/deployments/" + f"{deployment_id}/predictionsUnstructured" + ), ) -def _health_payload( +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 payload.""" + """Health.""" return { "status": "ok", "schema_version": "v1", @@ -73,7 +70,7 @@ def _health_payload( } -class _PoolTransport: +class _Transport: """Transport that serves health/predict per deployment id.""" def __init__( @@ -86,7 +83,6 @@ def __init__( self.health_by_id = dict(health_by_id or {}) self.fail_ids = fail_ids self.fail_status = fail_status - self.urls: list[str] = [] def get_json(self, url: str) -> Mapping[str, Any]: """Get json.""" @@ -94,90 +90,73 @@ def get_json(self, url: str) -> Mapping[str, Any]: def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: """Post json.""" - self.urls.append(url) deployment_id = url.rstrip("/").split("/")[-2] if deployment_id in self.fail_ids: - if self.fail_status < 0: - raise JointFMRequestError(f"{deployment_id} unreachable") 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_payload()) + return self.health_by_id.get(deployment_id, _health()) return {"ok": True, "deployment_id": deployment_id} -def test_pool_round_robin_and_retries_next_instance_on_470() -> None: - """Pool round robin and retries next instance on 470.""" +def test_pool_retries_next_instance_on_470() -> None: + """Pool retries next instance on 470.""" pool = JointFMInstancePool( instances=(_instance("a"), _instance("b")), - transport=_PoolTransport(fail_ids=frozenset({"a"}), fail_status=470), + transport=_Transport(fail_ids=frozenset({"a"})), ) - assert pool.next_instance().deployment_id == "a" assert pool.next_instance().deployment_id == "b" - assert pool.next_instance().deployment_id == "a" - - result = pool.post_json({"schema_version": "v1"}) - - assert result == {"ok": True, "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.""" - transport = _PoolTransport(fail_ids=frozenset({"a", "b"}), fail_status=470) pool = JointFMInstancePool( instances=(_instance("a"), _instance("b")), - transport=transport, + transport=_Transport(fail_ids=frozenset({"a", "b"})), ) - with pytest.raises(JointFMHTTPStatusError, match="unavailable"): pool.post_json({"schema_version": "v1"}) -def test_pool_health_rejects_model_or_checkpoint_mismatch() -> None: - """Pool health rejects model or checkpoint mismatch.""" - pool = JointFMInstancePool( - instances=(_instance("a"), _instance("b")), - transport=_PoolTransport( - health_by_id={ - "a": _health_payload(), - "b": _health_payload( - model_version="jointfm-inference:9.9.9+ckpt.other" - ), - } - ), - ) +def test_pool_health_rejects_mismatch_and_aligns_sample_cap() -> None: + """Pool health rejects mismatch and aligns sample cap.""" with pytest.raises(UnsupportedModelVersionError, match="model_version"): - pool.probe_all_health() + JointFMInstancePool( + instances=(_instance("a"), _instance("b")), + transport=_Transport( + health_by_id={ + "a": _health(), + "b": _health(model_version="jointfm-inference:9.9.9+ckpt.other"), + } + ), + ).probe_all_health() - pool = JointFMInstancePool( - instances=(_instance("a"), _instance("b")), - transport=_PoolTransport( - health_by_id={ - "a": _health_payload(checkpoint_version="ckpt-a"), - "b": _health_payload(checkpoint_version="ckpt-b"), - } - ), - ) with pytest.raises(UnsupportedServiceContractError, match="checkpoint_version"): - pool.probe_all_health() - - -def test_pool_health_uses_minimum_max_sample_count() -> None: - """Pool health uses minimum max sample count.""" - pool = JointFMInstancePool( + JointFMInstancePool( + instances=(_instance("a"), _instance("b")), + transport=_Transport( + health_by_id={ + "a": _health(checkpoint_version="ckpt-a"), + "b": _health(checkpoint_version="ckpt-b"), + } + ), + ).probe_all_health() + + metadata = JointFMInstancePool( instances=(_instance("a"), _instance("b")), - transport=_PoolTransport( + transport=_Transport( health_by_id={ - "a": _health_payload(max_sample_count=100), - "b": _health_payload(max_sample_count=40), + "a": _health(max_sample_count=100), + "b": _health(max_sample_count=40), } ), - ) - - metadata = pool.probe_all_health() - + ).probe_all_health() assert metadata.max_sample_count == 40 diff --git a/tests/test_transport.py b/tests/test_transport.py index e2c8f07..036b7a4 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -1187,16 +1187,9 @@ def _server_url(server: ThreadingHTTPServer) -> str: return f"http://{host}:{port}/predict" -def test_client_predict_round_robins_across_pool_instances() -> None: - """Client predict round robins across pool instances.""" - primary = ( - "https://app.datarobot.com/api/v2/deployments/" - "primary-id/predictionsUnstructured" - ) - backup = ( - "https://app.datarobot.com/api/v2/deployments/backup-id/predictionsUnstructured" - ) - settings = JointFMSettings( +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, @@ -1204,21 +1197,25 @@ def test_client_predict_round_robins_across_pool_instances() -> None: deployment_selector="deployment_ids", schema_version="v1", instances=( - JointFMInstanceSettings( - deployment_id="primary-id", - predict_url=primary, - health_url=primary, - ), - JointFMInstanceSettings( - deployment_id="backup-id", - predict_url=backup, - health_url=backup, - ), + 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).""" @@ -1243,47 +1240,12 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: "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 - -def test_client_pool_predict_with_pin_probes_health_before_traffic() -> None: - """Client pool predict with pin probes health before traffic.""" - primary = ( - "https://app.datarobot.com/api/v2/deployments/" - "primary-id/predictionsUnstructured" - ) - backup = ( - "https://app.datarobot.com/api/v2/deployments/backup-id/predictionsUnstructured" - ) - settings = 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, - health_url=primary, - ), - JointFMInstanceSettings( - deployment_id="backup-id", - predict_url=backup, - health_url=backup, - ), - ), - model_version="jointfm-inference:0.2.0+ckpt.sdk-test", - deployment_id="primary-id", - ) - class MismatchTransport: """Mismatch Transport (test helper).""" @@ -1295,16 +1257,9 @@ 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 "primary-id" in url: + if url == primary: return _health_payload(checkpoint_version="ckpt-a") return _health_payload(checkpoint_version="ckpt-b") - client = JointFMClient(settings=settings, transport=MismatchTransport()) - with pytest.raises(UnsupportedServiceContractError, match="checkpoint_version"): - client.predict( - { - "schema_version": "v1", - "model_version": "jointfm-inference:0.2.0+ckpt.sdk-test", - } - ) + JointFMClient(settings=settings, transport=MismatchTransport()).predict(payload) From c3ba633122b3b6cad2a56f0688eaa11d4915e9ab Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Tue, 25 Aug 2026 13:30:55 +0200 Subject: [PATCH 07/10] pool.py formatting --- src/jointfm_client/pool.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py index 0d2883e..527661c 100644 --- a/src/jointfm_client/pool.py +++ b/src/jointfm_client/pool.py @@ -77,12 +77,10 @@ def probe_all_health(self) -> HealthMetadata: for instance in self._instances: try: payload = self._transport.post_json( - instance.predict_url, - {"request_type": HEALTH_REQUEST_TYPE}, + instance.predict_url, {"request_type": HEALTH_REQUEST_TYPE} ) validate_service_metadata( - payload, - expected_model_version=self._expected_model_version, + payload, expected_model_version=self._expected_model_version ) healthy.append((instance, HealthMetadata.from_payload(payload))) except Exception as error: @@ -143,9 +141,7 @@ def _next_untried(self, tried: set[str]) -> JointFMInstanceSettings: def _log_unavailable(deployment_id: str, error: BaseException) -> None: logger.warning( - "JointFM instance unavailable: deployment_id=%s error=%s", - deployment_id, - error, + "JointFM instance unavailable: deployment_id=%s error=%s", deployment_id, error ) From 4415f536523182ec26d23eed2c1665ce90a88515 Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Wed, 26 Aug 2026 09:12:01 +0200 Subject: [PATCH 08/10] parallel runs --- src/jointfm_client/client.py | 93 +++++++++++++++------ src/jointfm_client/pool.py | 141 +++++++++++++++++++++++++------ tests/test_pool.py | 156 +++++++++++++++++++++++++++++------ tests/test_transport.py | 67 +++++++++++++++ 4 files changed, 385 insertions(+), 72 deletions(-) diff --git a/src/jointfm_client/client.py b/src/jointfm_client/client.py index b279856..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 @@ -500,28 +501,26 @@ def _forecast_sample_batches( ) -> 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._post_predict_json(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: @@ -529,6 +528,45 @@ 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) @@ -545,25 +583,32 @@ def _require_pool(self) -> JointFMInstancePool: ) if self._pool is None: assert self.settings is not None - transport = self._transport - if transport is None: - # Fail fast per peer; the pool retries across instances. - transport = 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, + # 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, - transport=transport, + 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( diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py index 527661c..c4f2a76 100644 --- a/src/jointfm_client/pool.py +++ b/src/jointfm_client/pool.py @@ -20,6 +20,7 @@ from dataclasses import replace import logging import threading +import time from typing import Any from jointfm_client.configuration import DEFAULT_RETRY_STATUS_CODES @@ -40,52 +41,90 @@ logger = logging.getLogger(__name__) _POOL_RETRYABLE_HTTP_STATUS_CODES = frozenset(DEFAULT_RETRY_STATUS_CODES) +_DEFAULT_PEER_COOLDOWN_SECONDS = 30.0 class JointFMInstancePool: """Round-robin JointFM requests across hosted deployments. - Use a fail-fast transport (``max_attempts=1``); this pool retries peers. + 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], - transport: JSONTransport, + 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") self._instances = tuple(instances) - self._transport = transport + self._transports = { + instance.deployment_id: transport + for instance, transport in zip(instances, transports, strict=True) + } + # Key locks by transport identity so a shared injected Session serializes, + # while distinct per-peer Sessions stay concurrent. + self._post_locks = { + id(transport): threading.Lock() for transport in self._transports.values() + } self._expected_model_version = expected_model_version + self._peer_cooldown_seconds = peer_cooldown_seconds self._lock = threading.Lock() self._index = 0 + self._active_ids = {instance.deployment_id for instance in self._instances} + self._cooldown_until: dict[str, float] = {} + + @property + def instance_count(self) -> int: + """Number of peers currently eligible for routing.""" + return len(self._eligible_instances()) def next_instance(self) -> JointFMInstanceSettings: - """Return the next instance using round-robin selection.""" + """Return the next eligible instance using round-robin selection.""" with self._lock: - instance = self._instances[self._index] - self._index = (self._index + 1) % len(self._instances) + active = self._eligible_instances_unlocked() + instance = active[self._index % len(active)] + self._index = (self._index + 1) % len(active) return instance + def instance_at(self, index: int) -> JointFMInstanceSettings: + """Return the eligible instance pinned for ``index`` (sticky batch mapping).""" + active = self._eligible_instances() + return active[index % len(active)] + def probe_all_health(self) -> HealthMetadata: - """Probe peers; require matching model/checkpoint; return min sample cap.""" + """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. + """ healthy: list[tuple[JointFMInstanceSettings, HealthMetadata]] = [] last_error: BaseException | None = None for instance in self._instances: try: - payload = self._transport.post_json( - instance.predict_url, {"request_type": HEALTH_REQUEST_TYPE} + payload = self._post_json( + instance, {"request_type": HEALTH_REQUEST_TYPE} ) validate_service_metadata( payload, expected_model_version=self._expected_model_version ) healthy.append((instance, HealthMetadata.from_payload(payload))) except Exception as error: - if not _is_pool_retryable(error): - raise + # Skip unreachable or incompatible peers; a bad backup must not + # take down a healthy primary. last_error = error _log_unavailable(instance.deployment_id, error) @@ -110,33 +149,87 @@ def probe_all_health(self) -> HealthMetadata: f"expected {reference.checkpoint_version!r}" ) max_samples = min(max_samples, metadata.max_sample_count) + + with self._lock: + self._active_ids = {instance.deployment_id for instance, _ in healthy} + self._index = 0 + for instance, _ in healthy: + self._cooldown_until.pop(instance.deployment_id, None) + if max_samples == reference.max_sample_count: return reference return replace(reference, max_sample_count=max_samples) def post_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: - """POST payload, trying untried peers on retryable failures.""" - tried: set[str] = set() + """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.""" + candidates = self._failover_candidates(instance) last_error: BaseException | None = None - while len(tried) < len(self._instances): - instance = self._next_untried(tried) - tried.add(instance.deployment_id) + for candidate in candidates: try: - return self._transport.post_json(instance.predict_url, payload) + result = self._post_json(candidate, payload) except Exception as error: if not _is_pool_retryable(error): raise last_error = error - _log_unavailable(instance.deployment_id, error) + self._cool_down(candidate.deployment_id) + _log_unavailable(candidate.deployment_id, error) + continue + self._clear_cooldown(candidate.deployment_id) + return result assert last_error is not None raise last_error - def _next_untried(self, tried: set[str]) -> JointFMInstanceSettings: - for _ in range(len(self._instances)): - instance = self.next_instance() - if instance.deployment_id not in tried: - return instance - raise RuntimeError("JointFMInstancePool has no untried instances") + def _failover_candidates( + self, preferred: JointFMInstanceSettings + ) -> tuple[JointFMInstanceSettings, ...]: + active = self._eligible_instances() + if preferred.deployment_id in {peer.deployment_id for peer in active}: + return (preferred,) + tuple( + peer for peer in active if peer.deployment_id != preferred.deployment_id + ) + return active + + def _eligible_instances(self) -> tuple[JointFMInstanceSettings, ...]: + with self._lock: + return self._eligible_instances_unlocked() + + def _eligible_instances_unlocked(self) -> tuple[JointFMInstanceSettings, ...]: + now = time.monotonic() + active = tuple( + instance + for instance in self._instances + if instance.deployment_id in self._active_ids + ) + not_cooling = tuple( + instance + for instance in active + if self._cooldown_until.get(instance.deployment_id, 0.0) <= now + ) + # All cooling: still try health-active peers rather than stall. + return not_cooling or active or self._instances + + def _cool_down(self, deployment_id: str) -> None: + with self._lock: + self._cooldown_until[deployment_id] = ( + time.monotonic() + self._peer_cooldown_seconds + ) + + def _clear_cooldown(self, deployment_id: str) -> None: + with self._lock: + self._cooldown_until.pop(deployment_id, None) + + def _post_json( + self, instance: JointFMInstanceSettings, payload: Mapping[str, Any] + ) -> Mapping[str, Any]: + transport = self._transports[instance.deployment_id] + with self._post_locks[id(transport)]: + return transport.post_json(instance.predict_url, payload) def _log_unavailable(deployment_id: str, error: BaseException) -> None: diff --git a/tests/test_pool.py b/tests/test_pool.py index 9279c0e..1a3ea96 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -16,7 +16,9 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +import threading from typing import Any import pytest @@ -28,6 +30,7 @@ UnsupportedModelVersionError, UnsupportedServiceContractError, ) +from jointfm_client.transport import JSONTransport def _instance(deployment_id: str) -> JointFMInstanceSettings: @@ -102,12 +105,29 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: 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 = JointFMInstancePool( - instances=(_instance("a"), _instance("b")), - transport=_Transport(fail_ids=frozenset({"a"})), - ) + 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"}) == { @@ -118,45 +138,133 @@ def test_pool_retries_next_instance_on_470() -> None: def test_pool_raises_when_all_instances_unavailable() -> None: """Pool raises when all instances unavailable.""" - pool = JointFMInstancePool( - instances=(_instance("a"), _instance("b")), - transport=_Transport(fail_ids=frozenset({"a", "b"})), - ) + 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"): - JointFMInstancePool( - instances=(_instance("a"), _instance("b")), - transport=_Transport( - health_by_id={ - "a": _health(), - "b": _health(model_version="jointfm-inference:9.9.9+ckpt.other"), - } - ), - ).probe_all_health() + 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"): - JointFMInstancePool( - instances=(_instance("a"), _instance("b")), + _pool( transport=_Transport( health_by_id={ "a": _health(checkpoint_version="ckpt-a"), "b": _health(checkpoint_version="ckpt-b"), } - ), + ) ).probe_all_health() - metadata = JointFMInstancePool( - instances=(_instance("a"), _instance("b")), + 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_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_transport.py b/tests/test_transport.py index 036b7a4..0c96983 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -1263,3 +1263,70 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: 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 From cbad79464f6b9e32827ec14f794a2ee5e68714c2 Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Wed, 26 Aug 2026 11:57:23 +0200 Subject: [PATCH 09/10] fix: ensure failover reaches health-excluded peers --- src/jointfm_client/pool.py | 24 +++++++++++++++++------- tests/test_pool.py | 15 +++++++++++++++ 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py index c4f2a76..97fcd74 100644 --- a/src/jointfm_client/pool.py +++ b/src/jointfm_client/pool.py @@ -180,7 +180,7 @@ def post_json_to( self._cool_down(candidate.deployment_id) _log_unavailable(candidate.deployment_id, error) continue - self._clear_cooldown(candidate.deployment_id) + self._reactivate(candidate.deployment_id) return result assert last_error is not None raise last_error @@ -188,12 +188,21 @@ def post_json_to( def _failover_candidates( self, preferred: JointFMInstanceSettings ) -> tuple[JointFMInstanceSettings, ...]: - active = self._eligible_instances() - if preferred.deployment_id in {peer.deployment_id for peer in active}: - return (preferred,) + tuple( - peer for peer in active if peer.deployment_id != preferred.deployment_id + """Prefer health-eligible peers; then try health-excluded peers last.""" + eligible = self._eligible_instances() + 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 ) - return active + else: + preferred_first = eligible + seen = {peer.deployment_id for peer in preferred_first} + last_resort = tuple( + peer for peer in self._instances if peer.deployment_id not in seen + ) + return preferred_first + last_resort def _eligible_instances(self) -> tuple[JointFMInstanceSettings, ...]: with self._lock: @@ -220,8 +229,9 @@ def _cool_down(self, deployment_id: str) -> None: time.monotonic() + self._peer_cooldown_seconds ) - def _clear_cooldown(self, deployment_id: str) -> None: + def _reactivate(self, deployment_id: str) -> None: with self._lock: + self._active_ids.add(deployment_id) self._cooldown_until.pop(deployment_id, None) def _post_json( diff --git a/tests/test_pool.py b/tests/test_pool.py index 1a3ea96..a623ee4 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -229,6 +229,21 @@ def test_pool_health_routes_only_reachable_peers() -> None: 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" From 34ac5adbf4842a722c7d284e035f77bedeea9ecc Mon Sep 17 00:00:00 2001 From: Yahor Kavaliou Date: Wed, 26 Aug 2026 15:53:30 +0200 Subject: [PATCH 10/10] fix: refactor pool classes --- src/jointfm_client/pool.py | 299 +++++++++++++++++++++++-------------- 1 file changed, 186 insertions(+), 113 deletions(-) diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py index 97fcd74..a17443b 100644 --- a/src/jointfm_client/pool.py +++ b/src/jointfm_client/pool.py @@ -17,7 +17,7 @@ from __future__ import annotations from collections.abc import Mapping, Sequence -from dataclasses import replace +from dataclasses import dataclass, replace import logging import threading import time @@ -44,6 +44,148 @@ _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. @@ -69,40 +211,35 @@ def __init__( ) if peer_cooldown_seconds < 0: raise ValueError("peer_cooldown_seconds must be >= 0") - self._instances = tuple(instances) - self._transports = { - instance.deployment_id: transport - for instance, transport in zip(instances, transports, strict=True) - } # Key locks by transport identity so a shared injected Session serializes, # while distinct per-peer Sessions stay concurrent. - self._post_locks = { - id(transport): threading.Lock() for transport in self._transports.values() - } - self._expected_model_version = expected_model_version - self._peer_cooldown_seconds = peer_cooldown_seconds - self._lock = threading.Lock() - self._index = 0 - self._active_ids = {instance.deployment_id for instance in self._instances} - self._cooldown_until: dict[str, float] = {} + 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._eligible_instances()) + return len(self._routing.eligible(self._peers)) def next_instance(self) -> JointFMInstanceSettings: """Return the next eligible instance using round-robin selection.""" - with self._lock: - active = self._eligible_instances_unlocked() - instance = active[self._index % len(active)] - self._index = (self._index + 1) % len(active) - return instance + return self._routing.next(self._peers).settings def instance_at(self, index: int) -> JointFMInstanceSettings: """Return the eligible instance pinned for ``index`` (sticky batch mapping).""" - active = self._eligible_instances() - return active[index % len(active)] + return self._routing.at(self._peers, index).settings def probe_all_health(self) -> HealthMetadata: """Probe peers; require matching model/checkpoint; return min sample cap. @@ -111,54 +248,9 @@ def probe_all_health(self) -> HealthMetadata: only when no peer is usable, or when usable peers disagree with each other on model/checkpoint. """ - healthy: list[tuple[JointFMInstanceSettings, HealthMetadata]] = [] - last_error: BaseException | None = None - for instance in self._instances: - try: - payload = self._post_json( - instance, {"request_type": HEALTH_REQUEST_TYPE} - ) - validate_service_metadata( - payload, expected_model_version=self._expected_model_version - ) - healthy.append((instance, 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(instance.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 instance, metadata in healthy[1:]: - if metadata.model_version != reference.model_version: - raise UnsupportedModelVersionError( - "JointFM deployment pool model_version mismatch: " - f"{instance.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"{instance.deployment_id!r} advertises " - f"{metadata.checkpoint_version!r}, " - f"expected {reference.checkpoint_version!r}" - ) - max_samples = min(max_samples, metadata.max_sample_count) - - with self._lock: - self._active_ids = {instance.deployment_id for instance, _ in healthy} - self._index = 0 - for instance, _ in healthy: - self._cooldown_until.pop(instance.deployment_id, None) - - if max_samples == reference.max_sample_count: - return reference - return replace(reference, max_sample_count=max_samples) + 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.""" @@ -168,28 +260,27 @@ def post_json_to( self, instance: JointFMInstanceSettings, payload: Mapping[str, Any] ) -> Mapping[str, Any]: """POST to ``instance`` first; on retryable failure try remaining peers.""" - candidates = self._failover_candidates(instance) + preferred = self._peers_by_id[instance.deployment_id] + candidates = self._failover_candidates(preferred) last_error: BaseException | None = None for candidate in candidates: try: - result = self._post_json(candidate, payload) + result = candidate.post_json(payload) except Exception as error: if not _is_pool_retryable(error): raise last_error = error - self._cool_down(candidate.deployment_id) + self._routing.cool_down(candidate.deployment_id) _log_unavailable(candidate.deployment_id, error) continue - self._reactivate(candidate.deployment_id) + self._routing.reactivate(candidate.deployment_id) return result assert last_error is not None raise last_error - def _failover_candidates( - self, preferred: JointFMInstanceSettings - ) -> tuple[JointFMInstanceSettings, ...]: + def _failover_candidates(self, preferred: PoolPeer) -> tuple[PoolPeer, ...]: """Prefer health-eligible peers; then try health-excluded peers last.""" - eligible = self._eligible_instances() + eligible = self._routing.eligible(self._peers) if preferred.deployment_id in {peer.deployment_id for peer in eligible}: preferred_first = (preferred,) + tuple( peer @@ -200,46 +291,28 @@ def _failover_candidates( preferred_first = eligible seen = {peer.deployment_id for peer in preferred_first} last_resort = tuple( - peer for peer in self._instances if peer.deployment_id not in seen + peer for peer in self._peers if peer.deployment_id not in seen ) return preferred_first + last_resort - def _eligible_instances(self) -> tuple[JointFMInstanceSettings, ...]: - with self._lock: - return self._eligible_instances_unlocked() - def _eligible_instances_unlocked(self) -> tuple[JointFMInstanceSettings, ...]: - now = time.monotonic() - active = tuple( - instance - for instance in self._instances - if instance.deployment_id in self._active_ids +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}" ) - not_cooling = tuple( - instance - for instance in active - if self._cooldown_until.get(instance.deployment_id, 0.0) <= now + 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}" ) - # All cooling: still try health-active peers rather than stall. - return not_cooling or active or self._instances - - def _cool_down(self, deployment_id: str) -> None: - with self._lock: - self._cooldown_until[deployment_id] = ( - time.monotonic() + self._peer_cooldown_seconds - ) - - def _reactivate(self, deployment_id: str) -> None: - with self._lock: - self._active_ids.add(deployment_id) - self._cooldown_until.pop(deployment_id, None) - - def _post_json( - self, instance: JointFMInstanceSettings, payload: Mapping[str, Any] - ) -> Mapping[str, Any]: - transport = self._transports[instance.deployment_id] - with self._post_locks[id(transport)]: - return transport.post_json(instance.predict_url, payload) def _log_unavailable(deployment_id: str, error: BaseException) -> None: