diff --git a/changelog/395.added.md b/changelog/395.added.md new file mode 100644 index 00000000..25a7200a --- /dev/null +++ b/changelog/395.added.md @@ -0,0 +1 @@ +The Azure AI Foundry `TabPFNClassifier` and `TabPFNRegressor` gained `collect_timeout_s`, bounding how long `predict()` waits for a result the endpoint is still computing (default two hours). On expiry the fit continues server-side, so predicting again collects it. diff --git a/changelog/395.fixed.md b/changelog/395.fixed.md new file mode 100644 index 00000000..4c31a490 --- /dev/null +++ b/changelog/395.fixed.md @@ -0,0 +1 @@ +Thinking mode against an Azure AI Foundry endpoint now returns a result instead of failing. A fit that takes longer than the platform's request window is handed back to be collected, and the client collects it automatically rather than mistaking the hand-off for a prediction. Collecting re-sends only the fit's id and the test rows, so the training data is not uploaded again on each attempt. diff --git a/src/tabpfn_client/foundry/estimator.py b/src/tabpfn_client/foundry/estimator.py index ae7c4a55..b53d4bc7 100644 --- a/src/tabpfn_client/foundry/estimator.py +++ b/src/tabpfn_client/foundry/estimator.py @@ -11,10 +11,17 @@ This client sends requests as `application/json` only (Foundry also accepts `multipart/form-data`, but we don't use it here). + +A fit that outlives the platform's request window answers HTTP 202 naming +the fit rather than the prediction; `predict*` collects it by re-sending +that id until the result is ready, so a long Thinking fit completes in one +call from the caller's point of view. Only the id and the test rows go on +the wire for those follow-ups, not the training data again. """ from __future__ import annotations +import time from typing import Any, Dict, Literal, Optional, Union, cast import httpx @@ -97,6 +104,28 @@ def _raise_for_status(resp: httpx.Response) -> None: ) +#: Fallback pause between collect attempts when the endpoint sends no +#: `Retry-After`. Short enough to feel responsive, long enough that polling a +#: multi-minute fit does not hammer the endpoint. +_DEFAULT_RETRY_AFTER_S = 10.0 + + +def _retry_after_s(resp: httpx.Response) -> float: + """How long the endpoint asked us to wait, in seconds. + + Only the delta-seconds form is honoured; an HTTP-date, or anything + unparseable, falls back to the default rather than failing a collect that + is otherwise progressing. + """ + raw = resp.headers.get("Retry-After") + if raw is None: + return _DEFAULT_RETRY_AFTER_S + try: + return max(0.0, float(raw)) + except ValueError: + return _DEFAULT_RETRY_AFTER_S + + def _to_jsonable(X: Any) -> list: """Coerce numpy / pandas inputs to plain Python lists for JSON.""" if isinstance(X, pd.DataFrame): @@ -201,6 +230,7 @@ def __init__( use_kv_cache: bool = False, fit_mode: Optional[FitModeLiteral] = None, timeout_s: float = 300.0, + collect_timeout_s: float = 7200.0, ): self.endpoint_url = endpoint_url self.api_key = api_key @@ -221,6 +251,7 @@ def __init__( self.use_kv_cache = use_kv_cache self.fit_mode = fit_mode self.timeout_s = timeout_s + self.collect_timeout_s = collect_timeout_s self._validate_args() def _validate_args(self) -> None: @@ -265,6 +296,13 @@ def _validate_args(self) -> None: f"limit); got {self.thinking_timeout_s!r}." ) + if self.collect_timeout_s < 0: + raise ValueError( + f"collect_timeout_s must be >= 0 (0 does not wait at all: a " + f"result that is not ready raises immediately); got " + f"{self.collect_timeout_s!r}." + ) + @property def _thinking_active(self) -> bool: return self.thinking_mode or self.thinking_effort is not None @@ -385,16 +423,52 @@ def _invoke( cached_model_id=self._cached_model_id if self._cache_active else None, thinking_block=self._build_thinking_block(), ) - resp = self._http_client().post( - self.endpoint_url, - json=body, - headers=self._headers(), - ) - _raise_for_status(resp) - payload = resp.json() - if self._cache_active: - self._cached_model_id = payload.get("model_id") or self._cached_model_id - return payload + deadline = time.monotonic() + self.collect_timeout_s + while True: + resp = self._http_client().post( + self.endpoint_url, + json=body, + headers=self._headers(), + ) + _raise_for_status(resp) + payload = resp.json() + if self._cache_active: + self._cached_model_id = payload.get("model_id") or self._cached_model_id + + if resp.status_code != 202: + return payload + + # 202 means the work outlived the request window and is still + # running server-side. It names the fit, so the follow-up carries + # that id and the test rows alone — re-sending the training data + # would cost the whole payload again on every attempt. + model_id = payload.get("model_id") + if model_id is None: + raise FoundryEndpointError( + "Endpoint answered HTTP 202 without a model_id, so the " + "result cannot be collected.", + request=resp.request, + response=resp, + error_code=payload.get("error_code"), + trace_id=payload.get("trace_id"), + ) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError( + f"Endpoint still working after collect_timeout_s=" + f"{self.collect_timeout_s:g}s. The fit continues server-side; " + f"predict again with the same data to collect it, or raise " + f"collect_timeout_s." + ) + body = _build_request_body( + task=self._task, + tabpfn_config=self._build_tabpfn_config(), + predict_params=params, + X_test=X_test, + cached_model_id=model_id, + thinking_block=self._build_thinking_block(), + ) + time.sleep(min(_retry_after_s(resp), remaining)) class TabPFNClassifier(_FoundryBase, ClassifierMixin): diff --git a/tests/unit/test_foundry_collect.py b/tests/unit/test_foundry_collect.py new file mode 100644 index 00000000..b2dd827c --- /dev/null +++ b/tests/unit/test_foundry_collect.py @@ -0,0 +1,139 @@ +# Copyright (c) Prior Labs GmbH 2026. +# Licensed under the Apache License, Version 2.0 +"""Collecting a result the endpoint could not return in one request. + +A fit that outlives the platform's request window answers HTTP 202 naming the +fit instead of the prediction. `predict*` collects it by re-sending that id, +so the caller still makes a single call — and the follow-ups carry only the id +and the test rows, never the training data again, which at a large context is +the whole payload on every attempt. +""" + +import json + +import httpx +import numpy as np +import pytest +import respx + +from tabpfn_client.foundry import FoundryEndpointError, TabPFNClassifier +from tabpfn_client.foundry.estimator import _DEFAULT_RETRY_AFTER_S, _retry_after_s + + +URL = "https://example.inference.ml.azure.com/invocations" +KEY = "test-key" + + +def _fitted(**kwargs): + """A thinking classifier with training data attached, plus that data.""" + clf = TabPFNClassifier( + endpoint_url=URL, api_key=KEY, thinking_effort="medium", **kwargs + ) + rng = np.random.default_rng(0) + X = rng.normal(size=(8, 3)) + y = (X[:, 0] > 0).astype(int) + clf.fit(X, y) + return clf, X + + +def _in_progress(model_id="mid-1", retry_after="0"): + return httpx.Response( + 202, + headers={"Retry-After": retry_after}, + json={ + "message": "The thinking fit for this dataset is still running.", + "error_code": "FIT_IN_PROGRESS", + "model_id": model_id, + }, + ) + + +def _done(model_id="mid-1"): + return httpx.Response( + 200, json={"prediction": [0, 1], "metadata": {}, "model_id": model_id} + ) + + +def _body(call): + return json.loads(call.request.content) + + +class TestCollect: + @respx.mock + def test_collects_the_result_after_a_202(self): + route = respx.post(URL).mock(side_effect=[_in_progress(), _done()]) + clf, X = _fitted() + + assert list(clf.predict(X[:2])) == [0, 1] + assert route.call_count == 2 + + @respx.mock + def test_resend_drops_the_training_data(self): + """The reason the id is worth returning at all.""" + route = respx.post(URL).mock(side_effect=[_in_progress(), _done()]) + clf, X = _fitted() + clf.predict(X[:2]) + + first, second = _body(route.calls[0]), _body(route.calls[1]) + assert "x_train" in first and "y_train" in first + assert "x_train" not in second and "y_train" not in second + assert second["context"] == {"model_id": "mid-1"} + assert "x_test" in second + + @respx.mock + def test_several_202s_are_all_collected(self): + route = respx.post(URL).mock( + side_effect=[_in_progress(), _in_progress(), _in_progress(), _done()] + ) + clf, X = _fitted() + + assert list(clf.predict(X[:2])) == [0, 1] + assert route.call_count == 4 + + @respx.mock + def test_202_without_a_model_id_is_reported(self): + """Nothing to collect with, so say so rather than looping forever.""" + respx.post(URL).mock( + return_value=httpx.Response( + 202, json={"message": "still running", "error_code": "FIT_IN_PROGRESS"} + ) + ) + clf, X = _fitted() + + with pytest.raises(FoundryEndpointError, match="without a model_id"): + clf.predict(X[:2]) + + @respx.mock + def test_gives_up_once_the_collect_budget_is_spent(self): + respx.post(URL).mock(return_value=_in_progress()) + clf, X = _fitted(collect_timeout_s=0) + + with pytest.raises(TimeoutError, match="collect_timeout_s"): + clf.predict(X[:2]) + + def test_rejects_a_negative_collect_budget(self): + with pytest.raises(ValueError, match="collect_timeout_s"): + TabPFNClassifier(endpoint_url=URL, api_key=KEY, collect_timeout_s=-1) + + +class TestRetryAfter: + @staticmethod + def _resp(**kwargs): + return httpx.Response(202, request=httpx.Request("POST", URL), **kwargs) + + def test_defaults_when_the_endpoint_says_nothing(self): + assert _retry_after_s(self._resp()) == _DEFAULT_RETRY_AFTER_S + + def test_honours_delta_seconds(self): + assert _retry_after_s(self._resp(headers={"Retry-After": "3"})) == 3.0 + + def test_falls_back_on_an_http_date(self): + # Only the delta-seconds form is honoured; an unparseable value must not + # fail a collect that is otherwise progressing. + stamp = "Wed, 21 Oct 2026 07:28:00 GMT" + assert _retry_after_s(self._resp(headers={"Retry-After": stamp})) == ( + _DEFAULT_RETRY_AFTER_S + ) + + def test_clamps_a_negative_delta(self): + assert _retry_after_s(self._resp(headers={"Retry-After": "-5"})) == 0.0