diff --git a/README.md b/README.md index a5dc07f..02f10ec 100644 --- a/README.md +++ b/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 diff --git a/docs/api-reference.md b/docs/api-reference.md index 5057cc9..734942d 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -8,7 +8,7 @@ This reference covers the supported public Python surface exported by `jointfm_c | --- | --- | | `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, `endpoint_client(deployment_id)` for an independent single-endpoint client bound to one configured deployment, `predict(payload)` for low-level JSON prediction, `forecast(...)` for validated tabular forecasts, and the `forecast_mean(...)`, `forecast_samples(...)`, `forecast_quantiles(...)`, and `forecast_log_prob(...)` convenience methods for typed forecast results. `forecast_log_prob(...)` is the one that asks about values the caller already holds: it takes `query_rows`, one observed row per entry of `query_times` carrying every declared column, and returns their log density under the model's joint. Each forecast method accepts `condition=`, one `EqualityCondition` or `IntervalCondition` or a list of them, which switches the request to the `condition` query mode and, before anything is sent, checks that the deployment's health metadata advertises the mode and every condition kind the conditions use. `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. `feature_importance(...)` runs permutation feature importance: one baseline `forecast_samples` call plus one per shuffled feature column, returning a list of `{"feature", "mean", "distance"}` dicts, each holding that feature's absolute forecast-mean shift and centered squared 2-Wasserstein distance indexed by target and horizon. | -`JointFMClient.from_env()` loads `config.yaml`, optional `.env` values, and process environment variables. The read-only `timeout` and `retry_config` properties return the `JointFMTimeoutConfig` and `JointFMRetryConfig` the client was built with. `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 (one request per instance; each instance's `max_concurrent_requests` says how many such requests it serves at once), 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, and the minimum reachable `max_concurrent_requests`. `JointFMClient.endpoint_client(deployment_id)` returns a new client bound to one configured deployment (a pool instance, or the client's own `settings.deployment_id`) with this client's timeout and retry policy but its own HTTP session and health cache, so several such clients can send requests to the same deployment at once; it raises `JointFMConfigurationError` for any other deployment ID. `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. The read-only `timeout` and `retry_config` properties return the `JointFMTimeoutConfig` and `JointFMRetryConfig` the client was built with. `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 (one request per instance; each instance's `max_concurrent_requests` says how many such requests it serves at once), 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, and the minimum reachable `max_concurrent_requests`. `JointFMClient.endpoint_client(deployment_id)` returns a new client bound to one configured deployment (a pool instance, or the client's own `settings.deployment_id`) with this client's timeout and retry policy but its own HTTP session and health cache, so several such clients can send requests to the same deployment at once; it raises `JointFMConfigurationError` for any other deployment ID. `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`. A `JOINTFM_DEPLOYMENT_IDS` pool splits every sample request evenly across all routable endpoints, at least one sample per batch, because downloading the response dominates a sample request and downloads from separate endpoints overlap; each endpoint serves at most one batch at a time. Batch `i` uses seed `seed + i`, so a seeded result depends on the number of routable endpoints. Clients configured without a reachable health route fall back to discovering the cap from the structured service error. ## Contract Classes @@ -27,8 +27,8 @@ This reference covers the supported public Python surface exported by `jointfm_c | `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, one request per instance), `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_series`, `max_series`, `t_input`, `t_output`, `n_input`, and `n_output`. `min_series` and `max_series` bound the total number of series in one request, features and targets counted together — there is no separate per-role cap. | -| `UtilizationStage` | One stage of `run_max_utilization_check`: `name`, the total requested `n_samples`, the `request_count` of maximal requests sent at once, `wallclock_seconds` (an upper bound on each request's duration), the pool `failovers` logged while it ran, and `errors`, one entry per request that ended in a JointFM error or a wrong sample count. `passed` is true without errors or failovers. | -| `MaxUtilizationReport` | Result of `run_max_utilization_check`: the `envelope` the requests were sized from, `decoding_strategy`, `topology_label`, `request_slots` (the sum of `max_concurrent_requests` over reachable endpoints), `unavailable_instances`, `read_timeout_seconds`, `health_seconds`, the `single_request` stage, and the `concurrent_saturation` stage (`None` when there is only one request slot). `passed` requires every endpoint to be reachable and every stage to pass; `format_lines()` renders the report ending in a `VERDICT:` line. | +| `UtilizationStage` | The burst of `run_max_utilization_check`: `name`, the total requested `n_samples`, the `request_count` of maximal requests sent at once, `wallclock_seconds` (an upper bound on each request's duration), and `errors`, one entry per request that ended in a JointFM error or a wrong sample count. `passed` is true without errors. | +| `MaxUtilizationReport` | Result of `run_max_utilization_check`: the `envelope` the requests were sized from, `decoding_strategy`, `topology_label`, `request_slots` (the sum of `max_concurrent_requests` over reachable endpoints), `unavailable_instances`, `read_timeout_seconds`, `health_seconds`, and the `concurrent_saturation` stage. `passed` requires every endpoint to be reachable and the stage to pass; `format_lines()` renders the report ending in a `VERDICT:` line. | | `ForecastPlan` | Validated forecast plan returned by `plan_forecast_columns`. Fields are `columns` (ordered `ColumnSpec` tuple), `feature_columns`, `target_columns` (both as the caller declared them), and `requested_columns` (the caller's target list). | | `StructuredError` | One structured JointFM service error with `code`, `message`, and optional `field`. | | `ForecastDiagnostics` | Response diagnostics containing `history_rows`, `horizon_count`, optional `seed`, and, on condition responses, `condition_draws` and an optional `interval_estimator`. | @@ -122,7 +122,7 @@ All SDK-specific exceptions inherit from `JointFMError`. | `validate_forecast_horizon(history_times, query_times, time_index_mode=...)` | Validate that query times are future, increasing, and encoded for the selected time-index mode. | | `resolve_notebook_project_root(start_dir=None)` | Resolve the nearest src-layout Python project root for notebooks started inside a repository tree. | | `bootstrap_notebook(add_src_root=False)` | Switch notebook working directory to the project root and optionally prepend the local `src` directory. | -| `run_max_utilization_check(client, *, seed=DEFAULT_UTILIZATION_SEED)` | Send the largest sample forecast the deployment's `data_generation` envelope admits (all `max_series` columns as targets, `n_input` history rows, `n_output` query times, `max_sample_count` samples). Then fill every request slot at once: each reachable endpoint receives `max_concurrent_requests` simultaneous requests for its own `max_sample_count`, each through its own `endpoint_client`. Returns a `MaxUtilizationReport`. A JointFM error during a stage is recorded in that stage, not raised. Raises `JointFMConfigurationError` unless the client was built with `JointFMRetryConfig(max_attempts=1)`, and `JointFMCapacityError` when health metadata has no `data_generation` block. | +| `run_max_utilization_check(client, *, seed=DEFAULT_UTILIZATION_SEED)` | Fill every request slot at once with the largest sample forecast the deployment's `data_generation` envelope admits (all `max_series` columns as targets, `n_input` history rows, `n_output` query times): each reachable endpoint receives `max_concurrent_requests` simultaneous requests for its own `max_sample_count`, each through its own `endpoint_client`. Returns a `MaxUtilizationReport`. A JointFM error during the burst is recorded in the stage, not raised. Raises `JointFMConfigurationError` unless the client was built with `JointFMRetryConfig(max_attempts=1)`, and `JointFMCapacityError` when health metadata has no `data_generation` block. | | `plan_forecast_columns(*, health, feature_columns, target_columns, history_length, query_times_length)` | Build a `ForecastPlan` from `HealthMetadata`. Caps `history_length` at `data_generation.n_input` and `query_times_length` at `n_output` (both upper bounds; smaller requests are allowed), checks the combined feature and target count against `data_generation.min_series` and `max_series`, and rejects duplicate column names. Raises `JointFMCapacityError` when any bound is exceeded or required metadata is missing. | ## Environment Variables @@ -134,7 +134,7 @@ All SDK-specific exceptions inherit from `JointFMError`. | `JOINTFM_SCHEMA_VERSION` | Hosted calls | Request schema pin. The SDK supports only `v5`. | | `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`. `health()` uses the minimum reachable `max_sample_count` as the sample-batch cap; `health_instances()` sums reachable caps, one request per instance, and reports topology. | +| `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, one request per instance, and reports topology. Sample forecasts are split evenly across all reachable instances. | | `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. | @@ -173,7 +173,7 @@ The string literals are exposed as `PREDICT_REQUEST_TYPE`, `HEALTH_REQUEST_TYPE` | `time_column` | For absolute datetime, optional otherwise | Name of the history time column. It must not duplicate a modeled column name. | | `requested_columns` | Optional | Output column names or integer indices, answered in the order given. Duplicates are rejected. Any declared column may be named, whatever its `role` and whatever a condition says about it. Defaults to every declared column in declared order, so a condition response and a forecast response over the same schema line up column for column. Under `return_mode="log_prob"` an explicit list must name every declared column in declared order, because the score is joint over them. | | `query_rows` | `log_prob` mode | Observed rows to score, one per entry of `query_times`, each carrying a value for every declared column. Required by `return_mode="log_prob"` and rejected for every other mode. A row that contradicts a condition — a pinned column given another value, or a bounded column outside its range — is refused by the service instead of scored. | -| `n_samples` | Samples and quantiles controls | Positive sample count when sampling controls are needed. Oversized sample forecasts are batched automatically against the cap advertised in health metadata. | +| `n_samples` | Samples and quantiles controls | Positive sample count when sampling controls are needed. Oversized sample forecasts are batched automatically against the cap advertised in health metadata; a deployment-ID pool splits every sample forecast evenly across its reachable instances. | | `quantiles` | Quantiles mode | Quantile levels in `(0, 1)`, required for `return_mode="quantiles"`. | | `seed` | Optional | Integer random seed for reproducible stochastic outputs. | | `conditions` | With `query_mode="condition"` | Non-empty list of `{"column", "kind": "equality", "value", "query_time_indices"}` or `{"column", "kind": "interval", "lower", "upper", "query_time_indices"}` entries, `null` marking an open bound. `query_time_indices` is a list of distinct indices into `query_times`, or `null` for every position. Two conditions on one column must cover disjoint positions, and every covered position must leave at least one column unconditioned. Forbidden with any other query mode. | diff --git a/notebooks/forecast_mean.ipynb b/notebooks/forecast_mean.ipynb index 6f60442..5167366 100644 --- a/notebooks/forecast_mean.ipynb +++ b/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": { diff --git a/notebooks/max_utilization.ipynb b/notebooks/max_utilization.ipynb index 0d2a5fb..4e3b590 100644 --- a/notebooks/max_utilization.ipynb +++ b/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", ")" ] diff --git a/src/jointfm_client/client.py b/src/jointfm_client/client.py index d1d846a..788183d 100644 --- a/src/jointfm_client/client.py +++ b/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 diff --git a/src/jointfm_client/utilization.py b/src/jointfm_client/utilization.py index 22af9ec..8b409c8 100644 --- a/src/jointfm_client/utilization.py +++ b/src/jointfm_client/utilization.py @@ -15,14 +15,14 @@ """Maximum-utilization check against a live JointFM deployment. The check reads the capacity envelope a deployment advertises in `/healthz` -and sends the largest sample forecast that envelope admits: ``max_series`` +and sizes the largest sample forecast that envelope admits: ``max_series`` target columns, an ``n_input``-row history, an ``n_output``-step horizon, and -the full sample budget. Every column is a target, because targets are what the -response carries, so this is also the largest response the deployment can be -asked for. Every request slot the reachable endpoints advertise -(``max_concurrent_requests`` per endpoint) is then filled at once with one such -request each. The outcome is a pass or fail verdict plus the wall-clock time -of every stage. +the endpoint's full sample budget. Every column is a target, because targets +are what the response carries, so this is also the largest response the +deployment can be asked for. Every request slot the reachable endpoints +advertise (``max_concurrent_requests`` per endpoint) is filled at once with +one such request each, so every resource of every endpoint is in use. The +outcome is a pass or fail verdict plus the wall-clock time of that burst. """ from __future__ import annotations @@ -30,7 +30,6 @@ from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass -import logging import random import time @@ -43,8 +42,6 @@ JointFMError, ) -POOL_LOGGER_NAME = "jointfm_client.pool" -SINGLE_REQUEST_STAGE = "single_request" CONCURRENT_SATURATION_STAGE = "concurrent_saturation" SYNTHETIC_COLUMN_PREFIX = "series_" SYNTHETIC_START_VALUE = 100.0 @@ -57,28 +54,23 @@ class UtilizationStage: """Outcome of one maximal forecast stage. ``n_samples`` is the total sample count the stage asked for and - ``request_count`` how many maximal requests it sent at once: one for - ``single_request``, one per request slot for ``concurrent_saturation``. - ``wallclock_seconds`` spans request building, every round trip, and - response parsing, so it bounds each single request's duration from above. - ``failovers`` holds the pool's "instance unavailable" warnings logged - during the stage: the pool reroutes a failed request to another endpoint, - so the forecast can still succeed although one endpoint could not serve - it, and the stage counts that as a failure. ``errors`` holds one entry per - request that ended in a JointFM error or a wrong sample count. + ``request_count`` how many maximal requests it sent at once, one per + request slot. ``wallclock_seconds`` spans request building, every round + trip, and response parsing, so it bounds each single request's duration + from above. ``errors`` holds one entry per request that ended in a + JointFM error or a wrong sample count. """ name: str n_samples: int request_count: int wallclock_seconds: float - failovers: tuple[str, ...] errors: tuple[str, ...] @property def passed(self) -> bool: - """Return whether every request completed without error or failover.""" - return not self.errors and not self.failovers + """Return whether every request completed without error.""" + return not self.errors @dataclass(frozen=True, slots=True) @@ -90,8 +82,6 @@ class MaxUtilizationReport: health probe; any entry fails the check, because the pool then runs below its configured capacity. ``request_slots`` is the sum of ``max_concurrent_requests`` over the reachable endpoints. - ``concurrent_saturation`` is ``None`` when that sum is one, where it would - repeat ``single_request``. """ envelope: DataGenerationCapabilities @@ -101,22 +91,12 @@ class MaxUtilizationReport: unavailable_instances: tuple[str, ...] read_timeout_seconds: float health_seconds: float - single_request: UtilizationStage - concurrent_saturation: UtilizationStage | None - - @property - def stages(self) -> tuple[UtilizationStage, ...]: - """Return the stages that ran, in execution order.""" - if self.concurrent_saturation is None: - return (self.single_request,) - return (self.single_request, self.concurrent_saturation) + concurrent_saturation: UtilizationStage @property def passed(self) -> bool: - """Return whether every endpoint was reachable and every stage passed.""" - return not self.unavailable_instances and all( - stage.passed for stage in self.stages - ) + """Return whether every endpoint was reachable and the burst passed.""" + return not self.unavailable_instances and self.concurrent_saturation.passed def format_lines(self) -> list[str]: """Return a human-readable report ending in the verdict line.""" @@ -134,26 +114,21 @@ def format_lines(self) -> list[str]: f"unavailable: {deployment_id}" for deployment_id in self.unavailable_instances ) - if self.concurrent_saturation is None: - lines.append( - f"{CONCURRENT_SATURATION_STAGE}: not applicable (one request slot)" - ) - for stage in self.stages: - status = "PASS" if stage.passed else "FAIL" - lines.append( - f"{stage.name}: {status} n_samples={stage.n_samples} " - f"requests={stage.request_count} " - f"wallclock={stage.wallclock_seconds:.2f}s" - ) - lines.extend(f" failover: {message}" for message in stage.failovers) - lines.extend(f" error: {error}" for error in stage.errors) + stage = self.concurrent_saturation + status = "PASS" if stage.passed else "FAIL" + lines.append( + f"{stage.name}: {status} n_samples={stage.n_samples} " + f"requests={stage.request_count} " + f"wallclock={stage.wallclock_seconds:.2f}s" + ) + lines.extend(f" error: {error}" for error in stage.errors) lines.append(f"VERDICT: {'PASS' if self.passed else 'FAIL'}") return lines @dataclass(frozen=True, slots=True) class _MaximalRequest: - """The envelope-sized forecast inputs every stage sends.""" + """The envelope-sized forecast inputs every request slot sends.""" history: Sequence[Mapping[str, float]] plan: ForecastPlan @@ -191,25 +166,24 @@ def run_max_utilization_check( *, seed: int = DEFAULT_UTILIZATION_SEED, ) -> MaxUtilizationReport: - """Send the largest advertised sample forecast and report the verdict. + """Fill every advertised request slot at once and report the verdict. - ``single_request`` sends one maximal request. ``concurrent_saturation`` - then fills every request slot of every reachable endpoint at once: each - endpoint receives ``max_concurrent_requests`` simultaneous requests for - its own ``max_sample_count``, each through its own single-endpoint client - (see :meth:`JointFMClient.endpoint_client`), so no client-side lock - serializes them. Slot ``i`` uses seed ``seed + i``. + ``concurrent_saturation`` fills every request slot of every reachable + endpoint at once: each endpoint receives ``max_concurrent_requests`` + simultaneous maximal requests for its own ``max_sample_count``, each + through its own single-endpoint client (see + :meth:`JointFMClient.endpoint_client`), so neither the pool's sample + split nor its per-endpoint lock reshapes or serializes them. Slot ``i`` + uses seed ``seed + i``. The client must be built with ``JointFMRetryConfig(max_attempts=1)``: a retried request would turn a failure at maximum load into a slower - success. Pool endpoints already send each request once and fail over - instead, and those failovers are recorded per stage. + success. - Raises :class:`JointFMConfigurationError` when the client retries or when - the pool logger cannot emit warnings (failovers would go unseen), and + Raises :class:`JointFMConfigurationError` when the client retries and :class:`JointFMCapacityError` when the deployment advertises no - ``data_generation`` block. Errors raised while a stage runs are recorded - in that stage instead, so the report still shows every stage. + ``data_generation`` block. Errors raised during the burst are recorded in + the stage instead, so the report still shows them. """ if client.retry_config.max_attempts != 1: raise JointFMConfigurationError( @@ -217,12 +191,6 @@ def run_max_utilization_check( "JointFMRetryConfig(max_attempts=1); got " f"max_attempts={client.retry_config.max_attempts}" ) - pool_logger = logging.getLogger(POOL_LOGGER_NAME) - if not pool_logger.isEnabledFor(logging.WARNING): - raise JointFMConfigurationError( - f"logger {POOL_LOGGER_NAME!r} must emit WARNING records so pool " - "failovers are visible to the utilization check" - ) health_started = time.perf_counter() health = client.health(cache=True, refresh=True) @@ -264,18 +232,9 @@ def run_max_utilization_check( for index, (deployment_id, metadata) in enumerate(endpoint_slots) ) - single_request = _run_single_request( - client, - pool_logger=pool_logger, - request=request, - n_samples=health.max_sample_count, - seed=seed, + concurrent_saturation = _run_concurrent_saturation( + client, request=request, slots=slots ) - concurrent_saturation = None - if len(slots) > 1: - concurrent_saturation = _run_concurrent_saturation( - client, pool_logger=pool_logger, request=request, slots=slots - ) return MaxUtilizationReport( envelope=envelope, @@ -289,57 +248,13 @@ def run_max_utilization_check( ), read_timeout_seconds=client.timeout.read_seconds, health_seconds=health_seconds, - single_request=single_request, concurrent_saturation=concurrent_saturation, ) -class _FailoverRecorder(logging.Handler): - """Collect pool warnings emitted while one stage runs.""" - - def __init__(self) -> None: - super().__init__(level=logging.WARNING) - self.messages: list[str] = [] - - def emit(self, record: logging.LogRecord) -> None: - """Store the formatted warning message.""" - self.messages.append(record.getMessage()) - - -def _run_single_request( - client: JointFMClient, - *, - pool_logger: logging.Logger, - request: _MaximalRequest, - n_samples: int, - seed: int, -) -> UtilizationStage: - recorder = _FailoverRecorder() - pool_logger.addHandler(recorder) - started = time.perf_counter() - try: - error = request.send(client, n_samples=n_samples, seed=seed) - # A service or transport failure under maximum load is the verdict itself, - # so it is recorded rather than propagated; other errors still raise. - except JointFMError as caught: - error = f"{type(caught).__name__}: {caught}" - finally: - wallclock_seconds = time.perf_counter() - started - pool_logger.removeHandler(recorder) - return UtilizationStage( - name=SINGLE_REQUEST_STAGE, - n_samples=n_samples, - request_count=1, - wallclock_seconds=wallclock_seconds, - failovers=tuple(recorder.messages), - errors=() if error is None else (error,), - ) - - def _run_concurrent_saturation( client: JointFMClient, *, - pool_logger: logging.Logger, request: _MaximalRequest, slots: Sequence[_RequestSlot], ) -> UtilizationStage: @@ -354,14 +269,14 @@ def _run_concurrent_saturation( try: for slot_client in slot_clients: slot_client.health(cache=True) - # Recorded for the same reason as in the single-request stage. + # A service or transport failure under maximum load is the verdict itself, + # so it is recorded rather than propagated; other errors still raise. except JointFMError as caught: return UtilizationStage( name=CONCURRENT_SATURATION_STAGE, n_samples=total_samples, request_count=len(slots), wallclock_seconds=0.0, - failovers=(), errors=( f"health probe before the burst: {type(caught).__name__}: {caught}", ), @@ -373,28 +288,22 @@ def _send_slot(index: int) -> str | None: error = request.send( slot_clients[index], n_samples=slot.n_samples, seed=slot.seed ) - # Recorded per slot for the same reason as in the single-request stage. + # Recorded per slot for the same reason as a failed health probe. except JointFMError as caught: error = f"{type(caught).__name__}: {caught}" return ( None if error is None else f"slot {index} ({slot.deployment_id}): {error}" ) - recorder = _FailoverRecorder() - pool_logger.addHandler(recorder) started = time.perf_counter() - try: - with ThreadPoolExecutor(max_workers=len(slots)) as executor: - outcomes = tuple(executor.map(_send_slot, range(len(slots)))) - finally: - wallclock_seconds = time.perf_counter() - started - pool_logger.removeHandler(recorder) + with ThreadPoolExecutor(max_workers=len(slots)) as executor: + outcomes = tuple(executor.map(_send_slot, range(len(slots)))) + wallclock_seconds = time.perf_counter() - started return UtilizationStage( name=CONCURRENT_SATURATION_STAGE, n_samples=total_samples, request_count=len(slots), wallclock_seconds=wallclock_seconds, - failovers=tuple(recorder.messages), errors=tuple(error for error in outcomes if error is not None), ) diff --git a/tests/test_pool.py b/tests/test_pool.py index 377a385..9e74f82 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -19,6 +19,7 @@ from collections.abc import Mapping, Sequence from concurrent.futures import ThreadPoolExecutor import threading +import time from typing import Any import pytest @@ -35,6 +36,8 @@ ) from jointfm_client.transport import JSONTransport +_OVERLAP_PROBE_SECONDS = 0.02 + def _instance(deployment_id: str) -> JointFMInstanceSettings: """Instance.""" @@ -288,6 +291,46 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: } +def test_pool_never_overlaps_two_posts_on_one_peer() -> None: + """More concurrent batches than peers queue per peer instead of overlapping.""" + in_flight: dict[str, int] = {} + peak: dict[str, int] = {} + counter_lock = threading.Lock() + + class _SlowTransport: + 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 + with counter_lock: + in_flight[url] = in_flight.get(url, 0) + 1 + peak[url] = max(peak.get(url, 0), in_flight[url]) + time.sleep(_OVERLAP_PROBE_SECONDS) + with counter_lock: + in_flight[url] -= 1 + return {"ok": True} + + instances = (_instance("a"), _instance("b")) + pool = _pool( + instances=instances, + transports=(_SlowTransport(), _SlowTransport()), + ) + + batch_count = 6 + with ThreadPoolExecutor(max_workers=batch_count) as executor: + futures = [ + executor.submit( + pool.post_json_to, pool.instance_at(index), {"schema_version": "v5"} + ) + for index in range(batch_count) + ] + for future in futures: + future.result(timeout=2.0) + + assert peak == {instance.predict_url: 1 for instance in instances} + + def test_pool_health_routes_only_reachable_peers() -> None: """After health, sticky/RR skip peers that failed the probe.""" pool = _pool( diff --git a/tests/test_transport.py b/tests/test_transport.py index 86de109..ecd0505 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -1420,3 +1420,64 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: ) assert set(predict_urls) == {primary, backup} assert len(predict_urls) == 2 + + +@pytest.mark.parametrize( + ("n_samples", "expected_sizes", "expected_seeds"), + [(3, [2, 1], [7, 8]), (1, [1], [7])], +) +def test_client_pool_splits_sample_requests_below_cap_across_all_peers( + n_samples: int, expected_sizes: list[int], expected_seeds: list[int] +) -> None: + """A pool splits a sample request evenly over every peer even below the cap.""" + 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) + predicts: list[tuple[str, int, int]] = [] + + 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": + return _health_payload() + sample_count = cast(int, payload["n_samples"]) + seed = cast(int, payload["seed"]) + predicts.append((url, sample_count, seed)) + response_payload = _forecast_response_payload(return_mode="samples") + outputs = cast(dict[str, object], response_payload["outputs"]) + outputs["samples"] = [[[float(seed)]] for _ in range(sample_count)] + diagnostics = cast(dict[str, object], response_payload["diagnostics"]) + diagnostics["seed"] = 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.3.0+ckpt.sdk-test", + n_samples=n_samples, + seed=7, + ) + + assert len(result.samples) == n_samples + assert sorted(size for _, size, _ in predicts) == sorted(expected_sizes) + assert sorted(seed for _, _, seed in predicts) == expected_seeds + assert len({url for url, _, _ in predicts}) == len(expected_sizes) + assert result.diagnostics.seed == 7 diff --git a/tests/test_utilization.py b/tests/test_utilization.py index 615778a..08816cc 100644 --- a/tests/test_utilization.py +++ b/tests/test_utilization.py @@ -121,8 +121,7 @@ class _EnvelopeTransport: """Serve health and echo-shaped sample forecasts; fail chosen deployments. ``sample_caps`` and ``concurrency`` set one deployment's advertised - ``max_sample_count`` and ``max_concurrent_requests``. With - ``simultaneous_after`` set, every forecast after the first that many must + ``max_sample_count`` and ``max_concurrent_requests``. Every forecast must reach a barrier together with ``barrier_parties - 1`` others, so a stage that sends them one at a time breaks the barrier instead of passing. """ @@ -135,7 +134,6 @@ def __init__( health_failures: frozenset[str] = frozenset(), sample_caps: Mapping[str, int] | None = None, concurrency: Mapping[str, int] | None = None, - simultaneous_after: int | None = None, barrier_parties: int = 1, ) -> None: self.data_generation = data_generation @@ -143,7 +141,6 @@ def __init__( self.health_failures = health_failures self.sample_caps = dict(sample_caps or {}) self.concurrency = dict(concurrency or {}) - self.simultaneous_after = simultaneous_after self.barrier = threading.Barrier( barrier_parties, timeout=_BARRIER_TIMEOUT_SECONDS ) @@ -167,12 +164,7 @@ def post_json(self, url: str, payload: Mapping[str, Any]) -> Mapping[str, Any]: ) with self.lock: self.forecasts.append((deployment_id, payload)) - forecast_index = len(self.forecasts) - 1 - if ( - self.simultaneous_after is not None - and forecast_index >= self.simultaneous_after - ): - self.barrier.wait() + self.barrier.wait() if deployment_id in self.forecast_failures: raise _http_error(deployment_id, self.forecast_failures[deployment_id]) query_times = cast(list[int], payload["query_times"]) @@ -231,9 +223,8 @@ def test_single_endpoint_sends_the_full_envelope() -> None: assert report.passed assert report.request_slots == 1 - assert report.concurrent_saturation is None - assert report.single_request.n_samples == _SAMPLE_CAP - assert report.single_request.request_count == 1 + assert report.concurrent_saturation.n_samples == _SAMPLE_CAP + assert report.concurrent_saturation.request_count == 1 [(deployment_id, payload)] = transport.forecasts assert deployment_id == "single-id" assert payload["n_samples"] == _SAMPLE_CAP @@ -246,19 +237,16 @@ def test_single_endpoint_sends_the_full_envelope() -> None: def test_saturation_fills_every_request_slot_of_one_endpoint_at_once() -> None: """An endpoint serving two requests at once receives two maximal requests together.""" - transport = _EnvelopeTransport( - concurrency={"single-id": 2}, simultaneous_after=1, barrier_parties=2 - ) + transport = _EnvelopeTransport(concurrency={"single-id": 2}, barrier_parties=2) report = run_max_utilization_check(_client(transport, "single-id")) assert report.passed assert report.request_slots == 2 saturation = report.concurrent_saturation - assert saturation is not None assert saturation.request_count == 2 assert saturation.n_samples == 2 * _SAMPLE_CAP - burst = transport.forecasts[1:] + burst = transport.forecasts assert [deployment_id for deployment_id, _ in burst] == ["single-id"] * 2 assert sorted(payload["seed"] for _, payload in burst) == [7, 8] assert all(payload["n_samples"] == _SAMPLE_CAP for _, payload in burst) @@ -269,21 +257,18 @@ def test_pool_saturation_sends_each_endpoint_its_slots_and_own_cap() -> None: transport = _EnvelopeTransport( sample_caps={"a-id": 2, "b-id": 3}, concurrency={"a-id": 2}, - simultaneous_after=1, barrier_parties=3, ) report = run_max_utilization_check(_client(transport, "a-id", "b-id")) assert report.passed - assert report.single_request.n_samples == 2 saturation = report.concurrent_saturation - assert saturation is not None assert saturation.request_count == 3 assert saturation.n_samples == 2 + 2 + 3 burst = sorted( (deployment_id, payload["n_samples"]) - for deployment_id, payload in transport.forecasts[1:] + for deployment_id, payload in transport.forecasts ) assert burst == [("a-id", 2), ("a-id", 2), ("b-id", 3)] @@ -297,32 +282,18 @@ def test_every_failed_slot_is_reported() -> None: report = run_max_utilization_check(_client(transport, "single-id")) saturation = report.concurrent_saturation - assert saturation is not None assert len(saturation.errors) == 2 assert all("JointFMHTTPStatusError" in error for error in saturation.errors) assert not report.passed -def test_pool_failover_fails_the_stage_although_the_forecast_succeeds() -> None: - """A request rerouted to a healthy peer still fails the verdict.""" - retryable = next(iter(DEFAULT_RETRY_STATUS_CODES)) - transport = _EnvelopeTransport(forecast_failures={"a-id": retryable}) - - report = run_max_utilization_check(_client(transport, "a-id", "b-id")) - - assert report.single_request.errors == () - assert report.single_request.failovers - assert not report.passed - assert report.format_lines()[-1] == "VERDICT: FAIL" - - def test_service_error_is_recorded_as_a_failed_stage() -> None: """A non-retryable service error ends the stage and fails the verdict.""" transport = _EnvelopeTransport(forecast_failures={"single-id": 500}) report = run_max_utilization_check(_client(transport, "single-id")) - [error] = report.single_request.errors + [error] = report.concurrent_saturation.errors assert "JointFMHTTPStatusError" in error assert not report.passed