feat: Split pool sample forecasts across all endpoints and drop the single-request utilization stage

README.md

@@ -157,7 +157,7 @@ Check that the deployment serves the largest request its health envelope adverti
 uv run jointfm-client max-utilization
 ```
 
-`max-utilization` sends a seeded synthetic sample forecast with `max_series` target columns, an `n_input`-row history, an `n_output`-step horizon, and `max_sample_count` samples. It then fills every request slot at once: each reachable endpoint (every pool instance with `JOINTFM_DEPLOYMENT_IDS`) receives as many simultaneous maximal requests as its `max_concurrent_requests` admits. It prints the wall-clock time of each stage and a final `VERDICT: PASS` or `VERDICT: FAIL` line. It exits with status 1 on FAIL: when a stage raises a service or transport error, when the pool fails a request over to another endpoint, or when an endpoint fails the health probe. Retries are disabled for this command, so a request that fails under load is reported instead of retried. Each run uses real GPU time on the deployment. `task max-utilization` runs the same command, and `notebooks/max_utilization.ipynb` shows the per-stage timings as a table.
+`max-utilization` fills every request slot at once with a seeded synthetic sample forecast of `max_series` target columns, an `n_input`-row history, an `n_output`-step horizon, and `max_sample_count` samples: each reachable endpoint (every pool instance with `JOINTFM_DEPLOYMENT_IDS`) receives as many simultaneous maximal requests as its `max_concurrent_requests` admits. It prints the burst's wall-clock time and a final `VERDICT: PASS` or `VERDICT: FAIL` line. It exits with status 1 on FAIL: when a request raises a service or transport error or returns the wrong sample count, or when an endpoint fails the health probe. Retries are disabled for this command, so a request that fails under load is reported instead of retried. Each run uses real GPU time on the deployment. `task max-utilization` runs the same command, and `notebooks/max_utilization.ipynb` shows the burst timing as a table.
 
 ## Notebook Workflows
 
@@ -183,7 +183,7 @@ The current forecast request contract is:
 - `time_column`: required for `"absolute_datetime"`, and used for ordered ordinal or continuous histories when supplied
 - `query_times`: non-empty future forecast times only
 - `requested_columns`: optional column names or integer column indices, with duplicates rejected
-- `n_samples`: positive sample count for sampled forecasts and quantile estimation. When `return_mode="samples"` exceeds the `max_sample_count` advertised by the deployment's health metadata, `forecast_samples(...)` splits the request into capped prediction batches up front and returns one merged `SampleForecastResult`.
+- `n_samples`: positive sample count for sampled forecasts and quantile estimation. When `return_mode="samples"` exceeds the `max_sample_count` advertised by the deployment's health metadata, `forecast_samples(...)` splits the request into capped prediction batches up front and returns one merged `SampleForecastResult`. With `JOINTFM_DEPLOYMENT_IDS`, every sample request is split evenly across all reachable endpoints, whatever its size, and each endpoint serves at most one of its batches at a time. Batch `i` uses seed `seed + i`, so a seeded result depends on how many endpoints were reachable.
 - `condition`: required with `query_mode="condition"` and forbidden otherwise. One `EqualityCondition` or `IntervalCondition`, or a list of them, each covering the future positions its `query_time_indices` names; sent on the wire as the `conditions` list, see below.
 
 ### Conditional Queries

notebooks/forecast_mean.ipynb

@@ -113,9 +113,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "joint-client-python (3.13.3)",
+   "display_name": "Python (joint-client-python)",
    "language": "python",
-   "name": "python3"
+   "name": "joint-client-python"
   },
   "language_info": {
    "codemirror_mode": {

notebooks/max_utilization.ipynb

@@ -47,12 +47,9 @@
    },
    "source": [
     "# Maximum Utilization\n",
-    "Send the configured JointFM deployment the largest request its `/healthz` capacity envelope admits, and report a PASS/FAIL verdict with wall-clock times. The request is sized from the `data_generation` block: `max_series` columns, an `n_input`-row history, an `n_output`-step horizon, and `max_sample_count` sample paths. Every column is a target, so the response is also the largest the deployment can be asked for. The history is a seeded synthetic random walk; only the request's size matters here, not its values.\n",
+    "Fill every request slot of the configured JointFM deployment at once with the largest request its `/healthz` capacity envelope admits, and report a PASS/FAIL verdict with the wall-clock time. The request is sized from the `data_generation` block: `max_series` columns, an `n_input`-row history, an `n_output`-step horizon, and `max_sample_count` sample paths. Every column is a target, so the response is also the largest the deployment can be asked for. The history is a seeded synthetic random walk; only the request's size matters here, not its values.\n",
     "\n",
-    "Two stages run:\n",
-    "\n",
-    "- `single_request` sends one maximal request to one endpoint.\n",
-    "- `concurrent_saturation` fills every request slot at once. Each reachable endpoint (every pool instance when `JOINTFM_DEPLOYMENT_IDS` configures a pool) advertises `max_concurrent_requests`, the number of requests it serves at the same time, and receives that many maximal requests simultaneously, each for its own `max_sample_count`. The stage runs only when the endpoints offer more than one request slot in total.\n",
+    "Each reachable endpoint (every pool instance when `JOINTFM_DEPLOYMENT_IDS` configures a pool) advertises `max_concurrent_requests`, the number of requests it serves at the same time, and receives that many maximal requests simultaneously, each for its own `max_sample_count`. This is the one place the client occupies more than one slot per endpoint; ordinary forecasts send at most one request per endpoint at a time.\n",
     "\n",
     "Each run spends real GPU time on the hosted deployment. The terminal equivalent is `uv run jointfm-client max-utilization`, which exits with status 1 on FAIL."
    ]
@@ -83,10 +80,10 @@
     "language": "markdown"
    },
    "source": [
-    "## Stage Timings\n",
-    "A stage fails when it raises a JointFM service or transport error, returns the wrong sample count, or triggers a pool failover. A failover means the pool rerouted a failed request to another endpoint: the forecast still arrives, but one endpoint could not serve the load. An endpoint that fails the health probe fails the whole check.\n",
+    "## Burst Timing\n",
+    "The burst fails when any request raises a JointFM service or transport error or returns the wrong sample count. An endpoint that fails the health probe fails the whole check.\n",
     "\n",
-    "`wallclock_s` spans request building, every round trip, and response parsing. It is therefore an upper bound on each single request's duration, and `read_timeout_share` relates that bound to the client's per-request read timeout. A share close to 1 on `single_request` means a slightly larger request, or a busier deployment, would time out. A deployment that queues its slots instead of serving them in parallel takes about `requests` times the `single_request` time for `concurrent_saturation`. The client also parses all concurrent responses in one Python process, so the time can grow even when the deployment serves the slots in parallel; only a time near `single_request` rules queueing out."
+    "`wallclock_s` spans request building, every round trip, and response parsing. It is therefore an upper bound on each single request's duration, and `read_timeout_share` relates that bound to the client's per-request read timeout. A share close to 1 means a slightly larger request, or a busier deployment, would time out. The client parses all concurrent responses in one Python process, so the time also grows with the number of slots when the deployment serves them in parallel."
    ]
   },
   {
@@ -101,6 +98,7 @@
    "source": [
     "import pandas as pd\n",
     "\n",
+    "stage = report.concurrent_saturation\n",
     "pd.DataFrame(\n",
     "    [\n",
     "        {\n",
@@ -112,10 +110,8 @@
     "            \"read_timeout_share\": round(\n",
     "                stage.wallclock_seconds / report.read_timeout_seconds, 3\n",
     "            ),\n",
-    "            \"failovers\": len(stage.failovers),\n",
     "            \"errors\": \"; \".join(stage.errors),\n",
     "        }\n",
-    "        for stage in report.stages\n",
     "    ]\n",
     ")"
    ]

src/jointfm_client/client.py

@@ -446,9 +446,9 @@ def forecast(
         hinted_columns = [column for column in payload["columns"] if "traits" in column]
         if hinted_columns:
             require_metadata_trait_support(self.health(cache=True), hinted_columns)
-        sample_cap = self._resolve_sample_batch_cap(payload)
-        if sample_cap is not None:
-            return self._forecast_sample_batches(payload, sample_cap)
+        batch_sizes = self._resolve_sample_batch_sizes(payload)
+        if batch_sizes is not None:
+            return self._forecast_sample_batches(payload, batch_sizes)
 
         try:
             response_payload = self._post_predict_json(payload)
@@ -457,7 +457,10 @@ def forecast(
             if sample_cap is None:
                 raise
             self._sample_batch_cap = sample_cap
-            return self._forecast_sample_batches(payload, sample_cap)
+            return self._forecast_sample_batches(
+                payload,
+                self._sample_batch_sizes(cast(int, payload["n_samples"]), sample_cap),
+            )
 
         return _forecast_response_from_payload(response_payload, payload)
 
@@ -810,40 +813,60 @@ def _forecast_payload_from_rows(
             query_rows=query_rows,
         )
 
-    def _resolve_sample_batch_cap(self, payload: Mapping[str, Any]) -> int | None:
-        """Return the batch size for an oversized sample request, if one applies.
+    def _resolve_sample_batch_sizes(
+        self, payload: Mapping[str, Any]
+    ) -> list[int] | None:
+        """Return per-request sample counts when a sample request is split.
 
         The deployment advertises its sampling budget as ``max_sample_count`` in
         `/healthz`, so an explicit sample request probes health once and batches
         locally instead of sending a request the service is guaranteed to reject.
         The cap is remembered for the client's lifetime. Clients configured
         without a reachable health route fall back to learning the cap from the
-        service's ``INPUT_SIZE_EXCEEDED`` response.
+        service's ``INPUT_SIZE_EXCEEDED`` response. Returns ``None`` when the
+        request goes out as one prediction.
         """
-        if _requested_sample_count(payload) is None:
+        requested_samples = _requested_sample_count(payload)
+        if requested_samples is None:
             return None
         if self._sample_batch_cap is None and self._health_route_configured():
             self.health()
-        return _known_sample_batch_cap(payload, self._sample_batch_cap)
+        if self._sample_batch_cap is None:
+            return None
+        batch_sizes = self._sample_batch_sizes(
+            requested_samples, self._sample_batch_cap
+        )
+        return batch_sizes if len(batch_sizes) > 1 else None
+
+    def _sample_batch_sizes(self, requested_samples: int, sample_cap: int) -> list[int]:
+        """Split ``requested_samples`` evenly into batches of at most ``sample_cap``.
+
+        A pool spreads every sample request over all routable peers, at least one
+        sample per batch, because downloading the response dominates a sample
+        request and concurrent downloads from separate peers overlap. A single
+        endpoint batches only what exceeds its cap.
+        """
+        peer_count = self._require_pool().instance_count if self._uses_pool() else 1
+        batch_count = max(
+            -(-requested_samples // sample_cap), min(requested_samples, peer_count)
+        )
+        base_size, remainder = divmod(requested_samples, batch_count)
+        return [
+            base_size + (1 if batch_index < remainder else 0)
+            for batch_index in range(batch_count)
+        ]
 
     def _forecast_sample_batches(
         self,
         payload: Mapping[str, Any],
-        sample_cap: int,
+        batch_sizes: Sequence[int],
     ) -> SampleForecastResult:
-        requested_samples = cast(int, payload["n_samples"])
-        remaining_samples = requested_samples
         batch_payloads: list[dict[str, Any]] = []
-        batch_index = 0
-
-        while remaining_samples > 0:
-            batch_samples = min(sample_cap, remaining_samples)
+        for batch_index, batch_samples in enumerate(batch_sizes):
             batch_payload = dict(payload)
             batch_payload["n_samples"] = batch_samples
             _set_batch_seed(batch_payload, batch_index)
             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)
@@ -864,6 +887,8 @@ def _forecast_sample_batches_parallel(
         self, batch_payloads: Sequence[Mapping[str, Any]]
     ) -> list[SampleForecastResult]:
         pool = self._require_pool()
+        # Batches beyond the peer count queue on the per-peer POST lock, so no
+        # peer ever serves two of this client's batches at once.
         max_workers = min(len(batch_payloads), pool.instance_count)
 
         def _run_batch(item: tuple[int, Mapping[str, Any]]) -> SampleForecastResult:
@@ -1178,20 +1203,6 @@ def _sample_batch_cap_from_error(
     return sample_cap
 
 
-def _known_sample_batch_cap(
-    payload: Mapping[str, Any],
-    sample_cap: int | None,
-) -> int | None:
-    if sample_cap is None:
-        return None
-    requested_samples = _requested_sample_count(payload)
-    if requested_samples is None:
-        return None
-    if requested_samples <= sample_cap:
-        return None
-    return sample_cap
-
-
 def _requested_sample_count(payload: Mapping[str, Any]) -> int | None:
     if payload.get("return_mode") != "samples":
         return None