[JFMGC-54] Report pooled health topology via health_instances

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 <path>` 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

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. |

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",

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,

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."""