diff --git a/README.md b/README.md index de1c9e7..7e4395c 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ The SDK targets the DataRobot-hosted unstructured prediction route and the same - Current SDK package version: `0.4.0` - Current JointFM service schema: `schema_version="v1"` -The public API shape is a synchronous low-level `JointFMClient` with `health()` and `predict(payload)` methods plus high-level `forecast(...)`, `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` helpers. The SDK is not a proxy service; callers use it as a local Python library that talks to the hosted or local JointFM endpoint. +The public API shape is a synchronous low-level `JointFMClient` with `health()`, `health_instances()`, and `predict(payload)` methods plus high-level `forecast(...)`, `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` helpers. The SDK is not a proxy service; callers use it as a local Python library that talks to the hosted or local JointFM endpoint. SDK package versions are standard Python distribution versions: `[project].version` in `pyproject.toml` and `jointfm_client.__version__` describe the released client library. JointFM `schema_version`, `image_version`, `model_version`, and `checkpoint_version` are service compatibility identifiers carried in configuration, health metadata, requests, and responses. They are not SDK package versions, and changing a deployment pin does not by itself require changing the SDK package version. @@ -120,16 +120,20 @@ Direct local URL helpers are used by the local service selector: `build_local_he Hosted settings also derive `health_url` from the resolved deployment URL as `deployments/{deployment_id}/healthz`. `JointFMClient.health(cache=True)` stores typed `HealthMetadata` only when the caller asks for caching, and `JointFMClient.refresh_health()` fetches a fresh copy. +Each endpoint's health payload describes only that endpoint. With `JOINTFM_DEPLOYMENT_IDS`, the client probes every configured peer and aggregates locally: `health()` returns consensus metadata whose `max_sample_count` is the **minimum** reachable cap (the sample-batch size), while `health_instances()` returns one `InstanceHealth` per configured ID plus the **sum** of reachable caps as overall parallel capacity, a compact `topology` / `topology_label` (for example `2x5000` or `1x7000, 1x3000`), and errors for unavailable peers. + ## CLI Workflows The package installs a `jointfm-client` command. It reads `.env` by default, accepts `--dotenv ` for another file, and accepts `--no-dotenv` when the process environment should be the only source. -Validate credentials, resolve the deployment, call `/healthz`, and print non-secret service metadata: +Validate credentials, resolve the deployment, probe health, and print non-secret service metadata plus per-instance availability and sample topology: ```bash uv run jointfm-client health ``` +The health command includes consensus `service` metadata, an `instances` list (available/unavailable, per-instance sample cap, errors), `topology` (for example `1x7000, 1x3000`), overall `max_sample_count` (sum of reachable caps), and non-secret `deployment` settings when configured. + Submit one low-level JSON request file and write the JSON response file: ```bash diff --git a/docs/api-reference.md b/docs/api-reference.md index 380c9dc..688ff76 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -6,9 +6,9 @@ This reference covers the supported public Python surface exported by `jointfm_c | Name | Purpose | | --- | --- | -| `JointFMClient` | Synchronous client for hosted or local JointFM endpoints. Use `from_env()` for `.env` and `config.yaml` backed hosted settings, `health()` for typed service metadata, `predict(payload)` for low-level JSON prediction, `forecast(...)` for validated tabular forecasts, and the `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` convenience methods for typed forecast results. `health()` probes `GET /healthz` for local deployments and POSTs `{"request_type": "health"}` to `predict_url` for hosted DataRobot deployments because the DataRobot deployment gateway only proxies the unstructured prediction route. | +| `JointFMClient` | Synchronous client for hosted or local JointFM endpoints. Use `from_env()` for `.env` and `config.yaml` backed hosted settings, `health()` for consensus typed service metadata, `health_instances()` for per-deployment probe results and pooled sample topology, `predict(payload)` for low-level JSON prediction, `forecast(...)` for validated tabular forecasts, and the `forecast_mean(...)`, `forecast_samples(...)`, and `forecast_quantiles(...)` convenience methods for typed forecast results. `health()` probes `GET /healthz` for local deployments and POSTs `{"request_type": "health"}` to `predict_url` for hosted DataRobot deployments because the DataRobot deployment gateway only proxies the unstructured prediction route. | -`JointFMClient.from_env()` loads `config.yaml`, optional `.env` values, and process environment variables. `JointFMClient.health(cache=True)` caches health metadata only when requested. `JointFMClient.predict(payload)` requires `payload["model_version"]`; high-level forecast helpers resolve the configured model version when the caller does not pass one explicitly. When `forecast_samples(...)` requests an explicit `n_samples`, the client learns the deployment's `max_sample_count` from health metadata before the first prediction, splits oversized requests into capped prediction batches, and returns one merged `SampleForecastResult`. Clients configured without a reachable health route fall back to discovering the cap from the structured service error. +`JointFMClient.from_env()` loads `config.yaml`, optional `.env` values, and process environment variables. `JointFMClient.health(cache=True)` caches health metadata only when requested. `JointFMClient.health_instances()` returns the same probe as a `HealthInstances` object: one `InstanceHealth` per configured deployment (including failures), `max_sample_count` as the sum of reachable caps (overall parallel capacity), and `topology` / `topology_label` grouping those caps (unavailable peers are listed but excluded from the sum and topology). Each endpoint's health payload describes only that endpoint; the client aggregates by calling each configured peer. `health()` still exposes the minimum reachable `max_sample_count`, which is the sample-batch cap used by forecast helpers. `JointFMClient.predict(payload)` requires `payload["model_version"]`; high-level forecast helpers resolve the configured model version when the caller does not pass one explicitly. When `forecast_samples(...)` requests an explicit `n_samples`, the client learns the deployment's `max_sample_count` from health metadata before the first prediction, splits oversized requests into capped prediction batches, and returns one merged `SampleForecastResult`. Clients configured without a reachable health route fall back to discovering the cap from the structured service error. ## Contract Classes @@ -18,7 +18,9 @@ This reference covers the supported public Python surface exported by `jointfm_c | `DataFrameSchema` | Describes tabular history layout. Fields are `columns`, `time_index_mode`, `time_column`, `time_scale_seconds`, `use_local_normalized_time`, `calendar_id`, and `timezone`. | | `ForecastRequestMetadata` | Holds `schema_version`, `model_version`, `query_mode`, and `return_mode` for one forecast request. | | `ForecastRequest` | Validated request object that combines metadata, schema, history rows, query times, requested columns, sample or quantile controls, and `seed`, then emits a JSON-compatible payload with `to_payload()`. | -| `HealthMetadata` | Typed service-health payload with service status, schema and model versions, checkpoint metadata, device, head, `decoding_strategy`, advertised modes, time-index encoding, `max_sample_count`, and an optional `data_generation` block carrying advertised capacity limits. The container exposes it on `GET /healthz` for direct local access and as the response to `POST {"request_type": "health"}` on the unstructured prediction route for DataRobot-hosted deployments. | +| `HealthMetadata` | Typed service-health payload with service status, schema and model versions, checkpoint metadata, device, head, `decoding_strategy`, advertised modes, time-index encoding, `max_sample_count`, and an optional `data_generation` block carrying advertised capacity limits. The container exposes it on `GET /healthz` for direct local access and as the response to `POST {"request_type": "health"}` on the unstructured prediction route for DataRobot-hosted deployments. Each endpoint reports only its own capabilities. | +| `InstanceHealth` | One configured deployment's probe outcome: `deployment_id`, optional `metadata` (`HealthMetadata` when reachable), and optional `error` when the peer was skipped. | +| `HealthInstances` | Client aggregation of `health_instances()`: `instances` (one `InstanceHealth` per configured ID), `max_sample_count` (sum of reachable caps = overall parallel capacity), `topology` as `(count, cap)` pairs sorted by descending cap, and `topology_label` such as `2x5000` or `1x7000, 1x3000`. Unavailable peers stay in `instances` but are omitted from the sum and topology. | | `DataGenerationCapabilities` | Optional service-health block describing the deployed checkpoint's data-generation capacity. Fields are `sampler_type`, `min_features`, `max_features`, `min_targets`, `max_targets`, `t_input`, `t_output`, `n_input`, and `n_output`. | | `ForecastPlan` | Validated forecast plan returned by `plan_forecast_columns`. Fields are `columns` (ordered `ColumnSpec` tuple), `feature_columns`, `target_columns` (both reflect post-downgrade roles), and `requested_columns` (the caller's original target list). | | `StructuredError` | One structured JointFM service error with `code`, `message`, and optional `field`. | @@ -119,7 +121,7 @@ All SDK-specific exceptions inherit from `JointFMError`. | `JOINTFM_SCHEMA_VERSION` | Hosted calls | Request schema pin. The SDK supports only `v1`. | | `JOINTFM_MODEL_VERSION` | Hosted calls | Exact JointFM deployment model version expected from the service-health payload and prediction responses. | | `JOINTFM_DEPLOYMENT_ID` | One selector | Deployment ID used to build hosted health and prediction URLs. | -| `JOINTFM_DEPLOYMENT_IDS` | One selector | Comma-separated hosted deployment IDs for round-robin load balancing (at least two unique IDs). Mutually exclusive with other selectors. Peers must share `model_version` and `checkpoint_version`; the SDK uses the minimum `max_sample_count`. | +| `JOINTFM_DEPLOYMENT_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`. `health()` uses the minimum reachable `max_sample_count` as the sample-batch cap; `health_instances()` sums reachable caps for overall parallel capacity and reports topology. | | `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/__init__.py b/src/jointfm_client/__init__.py index d17d8fe..8029709 100644 --- a/src/jointfm_client/__init__.py +++ b/src/jointfm_client/__init__.py @@ -99,7 +99,7 @@ UnsupportedSchemaVersionError, UnsupportedServiceContractError, ) -from jointfm_client.pool import JointFMInstancePool +from jointfm_client.pool import HealthInstances, InstanceHealth, JointFMInstancePool from jointfm_client.notebooks import ( WORKSPACE_ROOT_MARKERS, bootstrap_notebook, @@ -169,7 +169,9 @@ "ForecastResponse", "HEALTH_REQUEST_TYPE", "HealthMetadata", + "HealthInstances", "HostedDeploymentConfig", + "InstanceHealth", "IMPORT_NAMESPACE", "MeanForecastResult", "JOINTFM_DEPLOYMENT_ID_ENV", diff --git a/src/jointfm_client/cli.py b/src/jointfm_client/cli.py index bab71f9..afdc99c 100644 --- a/src/jointfm_client/cli.py +++ b/src/jointfm_client/cli.py @@ -147,8 +147,24 @@ def _add_dotenv_arguments(parser: argparse.ArgumentParser) -> None: def _health_command(args: argparse.Namespace, stdout: TextIO) -> int: client = JointFMClient.from_env(dotenv_path=_dotenv_path(args)) - health = client.health() - payload: dict[str, Any] = {"service": asdict(health)} + health = client.health(cache=True) + instances = client.health_instances(cache=True) + payload: dict[str, Any] = { + "service": asdict(health), + "instances": [ + { + "deployment_id": entry.deployment_id, + "available": entry.metadata is not None, + "max_sample_count": ( + None if entry.metadata is None else entry.metadata.max_sample_count + ), + "error": entry.error, + } + for entry in instances.instances + ], + "topology": instances.topology_label, + "max_sample_count": instances.max_sample_count, + } if client.settings is not None: payload["deployment"] = { "selector": client.settings.deployment_selector, diff --git a/src/jointfm_client/client.py b/src/jointfm_client/client.py index c7c80fa..ef04f8a 100644 --- a/src/jointfm_client/client.py +++ b/src/jointfm_client/client.py @@ -18,6 +18,7 @@ from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass from pathlib import Path import re from typing import Any, Self, cast @@ -58,7 +59,7 @@ QuantileForecastResult, SampleForecastResult, ) -from jointfm_client.pool import JointFMInstancePool +from jointfm_client.pool import HealthInstances, InstanceHealth, JointFMInstancePool from jointfm_client.settings import ( JointFMSettings, load_settings, @@ -77,6 +78,14 @@ ) +@dataclass(frozen=True, slots=True) +class _ProbedHealth: + """Consensus health metadata and per-deployment probe outcomes.""" + + metadata: HealthMetadata + instances: tuple[InstanceHealth, ...] + + class JointFMClient: """Synchronous entrypoint for low-level predictions and forecast helpers.""" @@ -108,7 +117,7 @@ def __init__( self._retry_config = retry_config self._response_body_excerpt_characters = response_body_excerpt_characters self._datarobot_request_id_headers = datarobot_request_id_headers - self._health_metadata: HealthMetadata | None = None + self._probed_health: _ProbedHealth | None = None self._sample_batch_cap: int | None = None self._pool: JointFMInstancePool | None = None @@ -158,34 +167,73 @@ def health(self, *, cache: bool = False, refresh: bool = False) -> HealthMetadat 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. + the minimum ``max_sample_count`` across those peers. Use + ``health_instances()`` for the per-deployment results of this probe. + """ + return self._probe_health(cache=cache, refresh=refresh).metadata + + def health_instances( + self, *, cache: bool = False, refresh: bool = False + ) -> HealthInstances: + """Return per-deployment health from the same probe as ``health()``. + + A single-endpoint client yields one entry. A deployment-ID pool yields + one entry per configured ID, including peers skipped as unreachable or + incompatible. ``max_sample_count`` is the sum of each reachable instance's + ``max_sample_count`` (overall parallel capacity). ``topology`` / + ``topology_label`` group those caps (for example ``2x5000`` or + ``1x7000, 1x3000``); unavailable peers are listed but excluded from the + sum and topology. ``health()`` still returns the minimum reachable cap + used as the sample-batch size. """ - if cache and not refresh and self._health_metadata is not None: - return self._health_metadata + return HealthInstances.from_instances( + self._probe_health(cache=cache, refresh=refresh).instances + ) - 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 + def _probe_health(self, *, cache: bool, refresh: bool) -> _ProbedHealth: + """Probe endpoints and return consensus plus per-deployment results.""" + if cache and not refresh and self._probed_health is not None: + return self._probed_health + probed = ( + self._probe_pool_health() + if self._uses_pool() + else self._probe_single_health() + ) + + self._sample_batch_cap = probed.metadata.max_sample_count + if cache: + self._probed_health = probed + return probed + + @property + def _health_metadata(self) -> HealthMetadata | None: + """Return cached consensus health metadata, if a probe has been stored.""" + probed = self._probed_health + return None if probed is None else probed.metadata + + def _probe_pool_health(self) -> _ProbedHealth: + """Probe a ``JOINTFM_DEPLOYMENT_IDS`` pool via the shared pool client.""" + result = self._require_pool().probe_health() + return _ProbedHealth(result.metadata, result.instances) + + def _probe_single_health(self) -> _ProbedHealth: + """Probe the single configured endpoint (local or hosted).""" if self._uses_predict_route_for_health(): payload = self._fetch_hosted_health_payload() else: - health_url = self._require_health_url() - payload = self._transport_for_request().get_json(health_url) - expected_model_version = ( - None if self.settings is None else self.settings.model_version - ) + payload = self._transport_for_request().get_json(self._require_health_url()) + validate_service_metadata( - payload, expected_model_version=expected_model_version + payload, + expected_model_version=getattr(self.settings, "model_version", None), ) metadata = HealthMetadata.from_payload(payload) - self._sample_batch_cap = metadata.max_sample_count - if cache: - self._health_metadata = metadata - return metadata + deployment_id = getattr(self.settings, "deployment_id", None) + return _ProbedHealth( + metadata, + (InstanceHealth(deployment_id=deployment_id, metadata=metadata),), + ) def _uses_predict_route_for_health(self) -> bool: """Return whether hosted health probes must POST to the predict route.""" diff --git a/src/jointfm_client/pool.py b/src/jointfm_client/pool.py index a17443b..802398a 100644 --- a/src/jointfm_client/pool.py +++ b/src/jointfm_client/pool.py @@ -16,12 +16,13 @@ from __future__ import annotations +from collections import Counter from collections.abc import Mapping, Sequence from dataclasses import dataclass, replace import logging import threading import time -from typing import Any +from typing import Any, Self from jointfm_client.configuration import DEFAULT_RETRY_STATUS_CODES from jointfm_client.contract import ( @@ -68,12 +69,54 @@ def post_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: return self.transport.post_json(self.settings.predict_url, payload) +@dataclass(frozen=True, slots=True) +class InstanceHealth: + """Health probe outcome for one configured JointFM deployment.""" + + deployment_id: str | None + metadata: HealthMetadata | None = None + error: str | None = None + + +@dataclass(frozen=True, slots=True) +class HealthInstances: + """Per-deployment health results and the summed sample budget. + + ``max_sample_count`` is the sum of reachable peers' caps (overall parallel + capacity). ``topology`` groups those caps as ``(count, cap)`` pairs sorted by + descending cap; unavailable instances are omitted from both. + """ + + instances: tuple[InstanceHealth, ...] + max_sample_count: int + topology: tuple[tuple[int, int], ...] + topology_label: str + + @classmethod + def from_instances(cls, instances: tuple[InstanceHealth, ...]) -> Self: + """Build the response from per-deployment probe outcomes.""" + caps = tuple( + instance.metadata.max_sample_count + for instance in instances + if instance.metadata is not None + ) + counts = Counter(caps) + topology = tuple((counts[cap], cap) for cap in sorted(counts, reverse=True)) + return cls( + instances=instances, + max_sample_count=sum(caps), + topology=topology, + topology_label=", ".join(f"{count}x{cap}" for count, cap in topology), + ) + + @dataclass(frozen=True, slots=True) class HealthProbeResult: - """Merged health metadata and the peer IDs that passed the gate.""" + """Merged health metadata, healthy peer IDs, and per-deployment outcomes.""" metadata: HealthMetadata healthy_ids: frozenset[str] + instances: tuple[InstanceHealth, ...] class PeerRoutingState: @@ -150,6 +193,7 @@ def __init__(self, *, expected_model_version: str | None = None) -> None: def probe(self, peers: Sequence[PoolPeer]) -> HealthProbeResult: """Probe peers; skip per-peer failures; fail only when none are usable.""" + instances: list[InstanceHealth] = [] healthy: list[tuple[PoolPeer, HealthMetadata]] = [] last_error: BaseException | None = None for peer in peers: @@ -158,12 +202,25 @@ def probe(self, peers: Sequence[PoolPeer]) -> HealthProbeResult: validate_service_metadata( payload, expected_model_version=self._expected_model_version ) - healthy.append((peer, HealthMetadata.from_payload(payload))) + metadata = HealthMetadata.from_payload(payload) + healthy.append((peer, metadata)) + instances.append( + InstanceHealth( + deployment_id=peer.deployment_id, + metadata=metadata, + ) + ) 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) + instances.append( + InstanceHealth( + deployment_id=peer.deployment_id, + error=str(error), + ) + ) if not healthy: assert last_error is not None @@ -183,6 +240,7 @@ def probe(self, peers: Sequence[PoolPeer]) -> HealthProbeResult: return HealthProbeResult( metadata=metadata, healthy_ids=frozenset(peer.deployment_id for peer, _ in healthy), + instances=tuple(instances), ) @@ -241,6 +299,12 @@ def instance_at(self, index: int) -> JointFMInstanceSettings: """Return the eligible instance pinned for ``index`` (sticky batch mapping).""" return self._routing.at(self._peers, index).settings + def probe_health(self) -> HealthProbeResult: + """Probe peers and return consensus metadata plus per-peer results.""" + result = self._health_gate.probe(self._peers) + self._routing.set_healthy(tuple(result.healthy_ids)) + return result + def probe_all_health(self) -> HealthMetadata: """Probe peers; require matching model/checkpoint; return min sample cap. @@ -248,9 +312,7 @@ def probe_all_health(self) -> HealthMetadata: only when no peer is usable, or when usable peers disagree with each other on model/checkpoint. """ - result = self._health_gate.probe(self._peers) - self._routing.set_healthy(tuple(result.healthy_ids)) - return result.metadata + return self.probe_health().metadata def post_json(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: """POST payload via round-robin, trying other peers on retryable failures.""" diff --git a/tests/test_cli.py b/tests/test_cli.py index 0a8523d..20438c5 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -23,7 +23,13 @@ import pandas as pd -from jointfm_client import HealthMetadata, JointFMResponseDecodeError, JointFMSettings +from jointfm_client import ( + HealthInstances, + HealthMetadata, + InstanceHealth, + JointFMResponseDecodeError, + JointFMSettings, +) from jointfm_client import cli @@ -47,8 +53,9 @@ class FakeHealthClient: deployment_id="deployment-id", ) - def health(self) -> HealthMetadata: + def health(self, *, cache: bool = False, refresh: bool = False) -> HealthMetadata: """Health.""" + del cache, refresh return HealthMetadata( status="ok", schema_version="v1", @@ -70,6 +77,20 @@ def health(self) -> HealthMetadata: max_sample_count=4096, ) + def health_instances( + self, *, cache: bool = False, refresh: bool = False + ) -> HealthInstances: + """Health instances.""" + del cache, refresh + return HealthInstances.from_instances( + ( + InstanceHealth( + deployment_id="deployment-id", + metadata=self.health(), + ), + ) + ) + class FakePredictClient: """Fake Predict Client (test helper).""" @@ -128,6 +149,16 @@ def test_health_command_prints_non_secret_metadata(monkeypatch, capsys) -> None: assert payload["service"]["status"] == "ok" assert payload["service"]["decoding_strategy"] == "parallel_dense" assert payload["deployment"]["deployment_id"] == "deployment-id" + assert payload["instances"] == [ + { + "deployment_id": "deployment-id", + "available": True, + "max_sample_count": 4096, + "error": None, + } + ] + assert payload["topology"] == "1x4096" + assert payload["max_sample_count"] == 4096 assert "secret-token" not in output diff --git a/tests/test_pool.py b/tests/test_pool.py index a623ee4..3c2f906 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -24,6 +24,9 @@ import pytest from jointfm_client import ( + HealthInstances, + HealthMetadata, + InstanceHealth, JointFMHTTPStatusError, JointFMInstancePool, JointFMInstanceSettings, @@ -178,6 +181,75 @@ def test_pool_health_rejects_mismatch_and_aligns_sample_cap() -> None: assert metadata.max_sample_count == 40 +def test_health_instances_sums_caps_and_builds_topology() -> None: + """Reachable caps sum and group into a compact topology label.""" + result = _pool( + transport=_Transport( + health_by_id={ + "a": _health(max_sample_count=7000), + "b": _health(max_sample_count=3000), + } + ) + ).probe_health() + instances = HealthInstances.from_instances(result.instances) + + assert instances.max_sample_count == 10000 + assert instances.topology == ((1, 7000), (1, 3000)) + assert instances.topology_label == "1x7000, 1x3000" + + matching = HealthInstances.from_instances( + _pool( + transport=_Transport( + health_by_id={ + "a": _health(max_sample_count=5000), + "b": _health(max_sample_count=5000), + } + ) + ) + .probe_health() + .instances + ) + assert matching.max_sample_count == 10000 + assert matching.topology == ((2, 5000),) + assert matching.topology_label == "2x5000" + + +def test_health_instances_excludes_unavailable_peer_from_sum_and_topology() -> None: + """Failed peers stay listed with an error but do not contribute capacity.""" + result = _pool( + transport=_Transport( + health_by_id={"b": _health(max_sample_count=5000)}, + fail_ids=frozenset({"a"}), + ) + ).probe_health() + instances = HealthInstances.from_instances(result.instances) + + assert [entry.deployment_id for entry in instances.instances] == ["a", "b"] + assert instances.instances[0].metadata is None + assert instances.instances[0].error is not None + assert instances.instances[1].metadata is not None + assert instances.max_sample_count == 5000 + assert instances.topology == ((1, 5000),) + assert instances.topology_label == "1x5000" + + +def _typed_health(*, max_sample_count: int = 4096) -> HealthMetadata: + """Typed health.""" + return HealthMetadata.from_payload(_health(max_sample_count=max_sample_count)) + + +def test_health_instances_from_single_entry() -> None: + """A single reachable entry yields one instance and a 1x topology.""" + entry = InstanceHealth( + deployment_id="deployment-id", + metadata=_typed_health(max_sample_count=4096), + ) + instances = HealthInstances.from_instances((entry,)) + assert len(instances.instances) == 1 + assert instances.max_sample_count == 4096 + assert instances.topology_label == "1x4096" + + 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) diff --git a/tests/test_transport.py b/tests/test_transport.py index 0c96983..013bc74 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -618,6 +618,33 @@ def test_client_health_returns_typed_metadata_and_caches_only_when_requested() - assert transport.post_count == 4 +def test_client_health_instances_returns_one_entry_for_single_endpoint() -> None: + """Single-endpoint clients expose one InstanceHealth via health_instances().""" + settings = JointFMSettings( + datarobot_endpoint="https://app.datarobot.com/api/v2", + datarobot_api_token="secret-token", + health_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", + predict_url="https://app.datarobot.com/api/v2/deployments/deployment-id/predictionsUnstructured", + deployment_selector="deployment_id", + schema_version="v1", + model_version="jointfm-inference:0.2.0+ckpt.sdk-test", + deployment_id="deployment-id", + ) + transport = RecordingTransport() + client = JointFMClient(settings=settings, transport=transport) + + instances = client.health_instances() + + assert len(instances.instances) == 1 + assert instances.instances[0].deployment_id == "deployment-id" + assert instances.instances[0].metadata is not None + assert ( + instances.max_sample_count == instances.instances[0].metadata.max_sample_count + ) + assert instances.topology_label == f"1x{instances.max_sample_count}" + assert transport.post_count == 1 + + def test_client_from_env_forwards_timeout_and_retry_config( monkeypatch: pytest.MonkeyPatch, ) -> None: