mainsequence 9.0.6.dev40__py3-none-any.whl → 9.0.6.dev41__py3-none-any.whl

mainsequence/client/_transport.py

@@ -0,0 +1,154 @@
+"""One bounded retry policy for SDK HTTP sessions."""
+
+import math
+import time
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+import requests
+from requests.adapters import HTTPAdapter
+from urllib3.exceptions import InvalidHeader
+from urllib3.util import Retry, Timeout
+
+DEFAULT_ALLOWED_METHODS = frozenset({"HEAD", "GET", "OPTIONS"})
+DEFAULT_STATUS_FORCELIST = (429, 500, 502, 503, 504)
+DEFAULT_TIMEOUT = (5.0, 120.0)
+
+
+class RequestBudget:
+    def __init__(self, timeout):
+        if timeout is None:
+            timeout = DEFAULT_TIMEOUT
+        if isinstance(timeout, Timeout):
+            connect, read = timeout.connect_timeout, timeout.read_timeout
+            total = timeout.total
+        elif isinstance(timeout, tuple):
+            connect, read = timeout
+            total = None
+        else:
+            connect = read = total = timeout
+        self.connect = self._seconds(connect)
+        self.read = self._seconds(read)
+        self.deadline = time.monotonic() + self._seconds(
+            self.connect + self.read if total is None else total
+        )
+
+    @staticmethod
+    def _seconds(value):
+        seconds = float(value)
+        if not math.isfinite(seconds) or seconds <= 0:
+            raise ValueError("Request timeouts must be positive, finite seconds.")
+        return seconds
+
+    @property
+    def remaining(self):
+        return self.deadline - time.monotonic()
+
+    def timeout(self):
+        remaining = self.remaining
+        if remaining <= 0:
+            raise requests.Timeout("The SDK request's total timeout budget is exhausted.")
+        return Timeout(
+            total=remaining,
+            connect=min(self.connect, remaining),
+            read=min(self.read, remaining),
+        )
+
+
+_ACTIVE_BUDGET: ContextVar[RequestBudget | None] = ContextVar("sdk_request_budget", default=None)
+
+
+@contextmanager
+def request_budget(timeout):
+    budget = RequestBudget(timeout)
+    parent = _ACTIVE_BUDGET.get()
+    if parent is not None:
+        budget.deadline = min(budget.deadline, parent.deadline)
+    token = _ACTIVE_BUDGET.set(budget)
+    try:
+        yield budget
+    finally:
+        _ACTIVE_BUDGET.reset(token)
+
+
+def remaining_timeout(timeout):
+    budget = _ACTIVE_BUDGET.get()
+    if budget is None:
+        return timeout
+    connect, read = timeout
+    remaining = budget.timeout().total
+    return Timeout(total=remaining, connect=min(connect, remaining), read=min(read, remaining))
+
+
+class DeadlineHTTPAdapter(HTTPAdapter):
+    def __init__(self, *, retries, backoff_factor):
+        if not isinstance(retries, int) or retries < 0:
+            raise ValueError("retries must be a non-negative integer.")
+        if not math.isfinite(backoff_factor) or backoff_factor < 0:
+            raise ValueError("backoff_factor must be non-negative and finite.")
+        # urllib3 must never add a second retry loop or resend a write on connect failure.
+        super().__init__(
+            max_retries=Retry(
+                total=0,
+                connect=0,
+                read=False,
+                status=0,
+                other=0,
+                allowed_methods=DEFAULT_ALLOWED_METHODS,
+            )
+        )
+        self.retries = retries
+        self.backoff_factor = backoff_factor
+
+    def send(self, request, stream=False, timeout=None, verify=True, cert=None, proxies=None):
+        budget = RequestBudget(timeout)
+        parent = _ACTIVE_BUDGET.get()
+        if parent is not None:
+            budget.deadline = min(budget.deadline, parent.deadline)
+        retryable = request.method in DEFAULT_ALLOWED_METHODS
+        for attempt in range(self.retries + 1):
+            response = None
+            error = None
+            try:
+                response = super().send(
+                    request,
+                    stream=stream,
+                    timeout=budget.timeout(),
+                    verify=verify,
+                    cert=cert,
+                    proxies=proxies,
+                )
+                if not stream:
+                    # Read before deciding to retry, so body failures share the same budget.
+                    _ = response.content
+                if not retryable or response.status_code not in DEFAULT_STATUS_FORCELIST:
+                    return response
+            except requests.exceptions.SSLError:
+                raise
+            except (
+                requests.ConnectionError,
+                requests.Timeout,
+                requests.exceptions.ChunkedEncodingError,
+            ) as exc:
+                error = exc
+
+            delay = min(self.backoff_factor * 2**attempt, Retry.DEFAULT_BACKOFF_MAX)
+            if error is None:
+                try:
+                    retry_after = self.max_retries.get_retry_after(response)
+                except InvalidHeader:
+                    retry_after = None
+                if retry_after is not None:
+                    delay = max(delay, retry_after)
+
+            if not retryable or attempt == self.retries or delay >= budget.remaining:
+                if error is not None:
+                    if response is not None:
+                        response.close()
+                    raise error
+                return response
+            if response is not None:
+                response.close()
+            time.sleep(delay)
+
+        raise AssertionError("Unreachable retry state")

mainsequence/client/client.py

@@ -9,7 +9,7 @@ import requests
 from mainsequence.defaults import resolve_backend_endpoint
 
 from .exceptions import raise_for_response
-from .utils import AuthLoaders, build_session
+from .utils import AuthLoaders, build_session, make_request
 from .utils import loaders as _global_loaders
 from .utils import session as _global_session
 
@@ -35,6 +35,7 @@ class MainSequenceClient:
     - Uses ONE shared loaders + session (from utils) by default
     - Can be reconfigured once at startup
     """
+
     _default: MainSequenceClient | None = None
 
     def __init__(
@@ -72,15 +73,18 @@ class MainSequenceClient:
         timeout: tuple[float, float] | float | None = None,
     ) -> Any:
         url = self.config.api_root.rstrip("/") + "/" + path.lstrip("/")
-        to = timeout or self.config.timeout
-
-        r = self.session.request(method.upper(), url, params=params, json=json, files=files, timeout=to)
-
-        # Refresh auth once if needed
-        if r.status_code in (401, 403):
-            self.loaders.refresh_headers()
-            self.session.headers.update(self.loaders.auth_headers)
-            r = self.session.request(method.upper(), url, params=params, json=json, files=files, timeout=to)
+        to = self.config.timeout if timeout is None else timeout
+        payload = {"params": params, "json": json}
+        if files is not None:
+            payload["files"] = files
+        r = make_request(
+            self.session,
+            method,
+            url,
+            self.loaders,
+            payload=payload,
+            time_out=to,
+        )
 
         raise_for_response(r, payload={"params": params, "json": json})
 
@@ -95,7 +99,9 @@ class MainSequenceClient:
         """
         if cls._default is None:
             # Reuse the existing global singletons from utils to avoid two sessions in prod.
-            cls._default = cls(MainSequenceClientConfig(), loaders=_global_loaders, session=_global_session)
+            cls._default = cls(
+                MainSequenceClientConfig(), loaders=_global_loaders, session=_global_session
+            )
         return cls._default
 
     @classmethod

mainsequence/client/inference.py

@@ -69,8 +69,8 @@ class InferenceClient:
         if self.organization_environment_uid is not None:
             query["organization_environment_uid"] = self.organization_environment_uid
         url = f"{utils.API_ENDPOINT.rstrip('/')}/{path}"
-        # The shared make_request loop retries network failures. Inference uses
-        # the same auth loader with a dedicated zero-retry transport instead.
+        # Inference opts out of safe-read retries as well: an uncertain outcome
+        # must preserve the caller's explicit idempotency key.
         with utils.build_session(loaders=utils.loaders, retries=0) as session:
             headers = dict(utils.loaders.refresh_headers(force=False, session=session))
             if key is not None:

mainsequence/client/utils.py

@@ -17,26 +17,21 @@ from uuid import UUID, getnode
 
 import psutil
 import requests
-from requests.adapters import HTTPAdapter
 from requests.structures import CaseInsensitiveDict
-from urllib3.util.retry import Retry
 
 from mainsequence.defaults import resolve_backend_endpoint
 from mainsequence.logconf import logger
 
+from ._transport import DEFAULT_ALLOWED_METHODS as DEFAULT_ALLOWED_METHODS
+from ._transport import DEFAULT_STATUS_FORCELIST as DEFAULT_STATUS_FORCELIST
+from ._transport import DEFAULT_TIMEOUT as DEFAULT_TIMEOUT
+from ._transport import DeadlineHTTPAdapter, remaining_timeout, request_budget
+
 # ---- Backend defaults (single source of truth) ----
 MAINSEQUENCE_ENDPOINT = resolve_backend_endpoint()
 API_ENDPOINT = f"{MAINSEQUENCE_ENDPOINT}/api/v1"
 AUTH_ENDPOINT = MAINSEQUENCE_ENDPOINT.rstrip("/")
 
-DEFAULT_STATUS_FORCELIST = (429, 500, 502, 503, 504)
-DEFAULT_ALLOWED_METHODS = frozenset(["HEAD", "GET", "OPTIONS", "POST", "PUT", "PATCH", "DELETE"])
-
-# requests supports either a float or (connect, read). Preserve previous ~120s read behavior,
-# but add a sane connect timeout.
-DEFAULT_TIMEOUT: tuple[float, float] = (5.0, 120.0)
-
-
 DATE_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"
 
 
@@ -383,7 +378,7 @@ class RuntimeCredentialAuthProvider(BaseAuthProvider):
                     "credential_secret": credential_secret,
                 },
                 headers={"Content-Type": "application/json"},
-                timeout=self.timeout,
+                timeout=remaining_timeout(self.timeout),
                 allow_redirects=False,
             )
             if response.status_code < 200 or response.status_code >= 300:
@@ -515,7 +510,7 @@ class JWTAuthProvider(BaseAuthProvider):
                 self.refresh_url,
                 json={"refresh": self.refresh_token},
                 headers={"Content-Type": "application/json"},
-                timeout=self.timeout,
+                timeout=remaining_timeout(self.timeout),
             )
 
             if r.status_code in _REFUSED_RENEWAL_STATUSES:
@@ -564,7 +559,7 @@ class JWTAuthProvider(BaseAuthProvider):
             self.obtain_url,
             json={"username": username, "password": password},
             headers={"Content-Type": "application/json"},
-            timeout=self.timeout,
+            timeout=remaining_timeout(self.timeout),
         )
         r.raise_for_status()
         data = r.json()
@@ -671,63 +666,46 @@ def make_request(
 ):
     from requests.models import Response
 
-    TIMEOFF = 0.25
-    TRIES = int(15 // TIMEOFF)
     timeout = DEFAULT_TIMEOUT if time_out is None else time_out
     payload = {} if payload is None else payload
-
-    def get_req(session):
-        if r_type == "GET":
-            return session.get
-        elif r_type == "POST":
-            return session.post
-        elif r_type == "PUT":
-            return session.put
-        elif r_type == "PATCH":
-            return session.patch
-        elif r_type == "DELETE":
-            return session.delete
-        else:
-            raise NotImplementedError(f"Unsupported method: {r_type}")
+    r_type = r_type.upper()
+    if r_type not in DEFAULT_ALLOWED_METHODS | {"POST", "PUT", "PATCH", "DELETE"}:
+        raise NotImplementedError(f"Unsupported method: {r_type}")
 
     request_kwargs = {}
     if r_type in ("POST", "PATCH") and "files" in payload:
-        request_kwargs["data"] = payload.get("json", {})
-        request_kwargs["files"] = payload["files"]
+        request_kwargs = dict(payload)
+        request_kwargs["data"] = request_kwargs.pop("json", {})
         s.headers.pop("Content-Type", None)
     else:
-        request_kwargs = payload
-
-    req = get_req(session=s)
-    keep_request = True
-    counter = 0
-    auth_retried = False
+        request_kwargs = dict(payload)
+    request_kwargs.setdefault("allow_redirects", False)
+    req = getattr(s, r_type.lower())
 
     if accept_gzip:
         s.headers.setdefault("Accept-Encoding", "gzip")
 
-    while keep_request:
+    with request_budget(timeout) as budget:
         try:
-            if loaders is not None:
-                s.headers.update(loaders.refresh_headers(force=False, session=s))
-
-            start_time = time.perf_counter()
-            logger.debug(f"Requesting {r_type} from {url}")
-            r = req(url, timeout=timeout, **request_kwargs)
-            duration = time.perf_counter() - start_time
-            logger.debug(f"{url} took {duration:.4f} seconds.")
-
-            if r.status_code == 401 and loaders is not None and not auth_retried:
-                logger.warning(f"Error {r.status_code}; forcing auth refresh once")
-                # A refused renewal raises AuthError: it is reported below with the
-                # reason and the repair, not as a traceback.
-                s.headers.update(loaders.refresh_headers(force=True, session=s))
-                req = get_req(session=s)
-                auth_retried = True
-                continue
-
-            keep_request = False
-            break
+            for auth_attempt in range(2):
+                if loaders is not None:
+                    budget.timeout()
+                    s.headers.update(loaders.refresh_headers(force=bool(auth_attempt), session=s))
+
+                start_time = time.perf_counter()
+                logger.debug(f"Requesting {r_type} from {url}")
+                r = req(url, timeout=budget.timeout(), **request_kwargs)
+                duration = time.perf_counter() - start_time
+                logger.debug(f"{url} took {duration:.4f} seconds.")
+                if (
+                    r.status_code != 401
+                    or loaders is None
+                    or auth_attempt
+                    or r_type not in DEFAULT_ALLOWED_METHODS
+                ):
+                    return r
+                r.close()
+                logger.warning("Error 401; forcing auth refresh once within the request budget")
 
         except AuthError as e:
             logger.warning(f"Auth error for {url}: {e}")
@@ -736,32 +714,15 @@ def make_request(
             r.error_type = "auth_error"
             r.status_code = 401
             r._content = str(e).encode("utf-8")
-            keep_request = False
-            break
-
-        except requests.exceptions.ConnectionError:
-            logger.exception(f"Error connecting {url}")
-        except TypeError as e:
-            logger.exception(f"Type error for {url} exception {e}")
-            raise e
-        except Exception as e:
-            logger.exception(f"Error connecting {url} exception {e}")
-
-        counter += 1
-        if counter >= TRIES:
-            keep_request = False
+            return r
+        except requests.RequestException as e:
+            logger.warning(f"Request {r_type} to {url} failed: {e}")
             r = Response()
             r.code = "expired"
             r.error_type = "expired"
             r.status_code = 500
-            break
-
-        logger.debug(
-            f"Trying request again after {TIMEOFF}s - Counter: {counter}/{TRIES} - URL: {url}"
-        )
-        time.sleep(TIMEOFF)
-
-    return r
+            r._content = str(e).encode("utf-8")
+            return r
 
 
 def build_session(
@@ -779,23 +740,7 @@ def build_session(
     if accept_gzip:
         s.headers.setdefault("Accept-Encoding", "gzip")
 
-    retry_kwargs = dict(
-        total=retries,
-        connect=retries,
-        read=retries,
-        status=retries,
-        backoff_factor=backoff_factor,
-        status_forcelist=DEFAULT_STATUS_FORCELIST,
-        respect_retry_after_header=True,
-        raise_on_status=False,
-    )
-
-    try:
-        retry_cfg = Retry(allowed_methods=DEFAULT_ALLOWED_METHODS, **retry_kwargs)
-    except TypeError:
-        retry_cfg = Retry(method_whitelist=DEFAULT_ALLOWED_METHODS, **retry_kwargs)
-
-    adapter = HTTPAdapter(max_retries=retry_cfg)
+    adapter = DeadlineHTTPAdapter(retries=retries, backoff_factor=backoff_factor)
     s.mount("https://", adapter)
     s.mount("http://", adapter)
     return s
@@ -929,24 +874,7 @@ def _install_retry_adapters_in_place(
     Configure retry adapters on an EXISTING session object (do not rebind 'session').
     This is critical so 'from utils import session' users still get the updated behavior.
     """
-    retry_kwargs = dict(
-        total=retries,
-        connect=retries,
-        read=retries,
-        status=retries,
-        backoff_factor=backoff_factor,
-        status_forcelist=DEFAULT_STATUS_FORCELIST,
-        respect_retry_after_header=True,
-        raise_on_status=False,
-    )
-
-    # urllib3 compatibility across versions
-    try:
-        retry_cfg = Retry(allowed_methods=DEFAULT_ALLOWED_METHODS, **retry_kwargs)
-    except TypeError:
-        retry_cfg = Retry(method_whitelist=DEFAULT_ALLOWED_METHODS, **retry_kwargs)
-
-    adapter = HTTPAdapter(max_retries=retry_cfg)
+    adapter = DeadlineHTTPAdapter(retries=retries, backoff_factor=backoff_factor)
 
     # Close old adapters' pools (best-effort), then mount new ones
     for prefix in ("https://", "http://"):

{mainsequence-9.0.6.dev40.dist-info → mainsequence-9.0.6.dev41.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mainsequence
-Version: 9.0.6.dev40
+Version: 9.0.6.dev41
 Summary: Main Sequence SDK
 Author-email: Main Sequence GmbH <dev@main-sequence.io>
 License-Expression: MIT

{mainsequence-9.0.6.dev40.dist-info → mainsequence-9.0.6.dev41.dist-info}/RECORD

@@ -34,27 +34,28 @@ mainsequence/cli/sdk_utils.py,sha256=SFuSzh8WDNf21mLiDa51K8Fqe5XiLWh2KhHj8zUjCXo
 mainsequence/cli/ssh_utils.py,sha256=m7HgFHZQ545sxY9igfQLpaNjWHm4SiN9VwCkAuGy2yw,11124
 mainsequence/cli/ui.py,sha256=memd6pmjALfGj8zndGCx_z80ztQeqTcOvcofU0un34U,2277
 mainsequence/client/__init__.py,sha256=lsr9JWxeVqOG0KRoHErW5M40ACPSR7SAEE508dTN5XM,672
+mainsequence/client/_transport.py,sha256=l2xkbJ21SbrJhzgOf58Yikxg-uzOQPyo-H5hENdCBo8,5374
 mainsequence/client/agent_runtime_models.py,sha256=lay0DN3lN5p6-3ajVsaHvPgcmrURFwrNfTb6upoImm4,66642
 mainsequence/client/base.py,sha256=T3k3xptM09TreiVogrxAPWoIy8zM6Dtn3Cbwsnb33ag,42678
-mainsequence/client/client.py,sha256=sUsJhjwKq4oZsEX2Ozm-7owcyQq0kgFNKUrtEYwf8cU,3907
+mainsequence/client/client.py,sha256=Jaw9RYObdVDS5J2O0pln99-XpuW9e4vW-SsbJRNldfw,3866
 mainsequence/client/compute_validation.py,sha256=ByeX0CkIE7qq06ggZFfN92U3Qp2e9cvo2RTIGWa5_Vw,6662
 mainsequence/client/exceptions.py,sha256=xx2L4l3OMNsTybbYkiTKK398ahnutEascM0It35tGr0,3815
 mainsequence/client/github_issues.py,sha256=pSTUw3UPrmE6wqnpoJctartTaNSq6iu6PQRrd8hip6I,14568
-mainsequence/client/inference.py,sha256=iMg3l5DGW6PQeVkz38Z3lTpxryMO14CTK6gIRHLYUBE,7368
+mainsequence/client/inference.py,sha256=g24H9LWumMRABIevho8FS7BBOuyYvf80pRwCmJZpLeY,7353
 mainsequence/client/models_foundry.py,sha256=iN1jJzDokfniEQF0Ac5eepAB7L5pVjJnKavaHLP7BHA,36310
 mainsequence/client/models_helpers.py,sha256=y7cgcsH8yO3dfIWPhRdcYam7ItWHw3Ps63-WYDe4YHE,46262
 mainsequence/client/models_user.py,sha256=BU_hfn4-4KqYz3nsaPlXuYZmUG6_-7YM892mkM92Zyo,37613
 mainsequence/client/observability.py,sha256=_WrDKCaNk_TXmKElKCXC1inork6Exxxe8FRY4m1GI9s,16008
-mainsequence/client/utils.py,sha256=fQur4p-ojLWTMcNWlYWx7uZKbGNEc4c2kJq1Wl_w070,31925
+mainsequence/client/utils.py,sha256=c4JjMIqKfwSz8DdkwLGmGtQ9lSmVn_KEw9KWGDwz_T0,30133
 mainsequence/client/value_sets.py,sha256=7TA-mpuAd3hOduNgO1WTQQTuNOGaFgue-VqCsDWKaC8,6181
 mainsequence/instrumentation/__init__.py,sha256=J_9XmfjprElvz0u2t9rTyJT0kFdyKpxdbkdQiaw28_U,316
 mainsequence/instrumentation/utils.py,sha256=b-ihjSioAepeLfMpudM2SqCtvnxMqIXileHx3klBCQU,3377
 mainsequence/server/__init__.py,sha256=BHUAhtdWPDRdxX07NYO_6wrBv5mf0GsXflecO52PmUc,73
 mainsequence/server/caller_assertions.py,sha256=hZW1pkkvQ9gaFL075bxUbbctTRrgMMQRTSc2Da0vKCM,9705
 mainsequence/server/fastapi.py,sha256=tD1aVuJOgoBTu7H8bs9IldPNYRT4fhkr1leN3KZP1pI,10521
-mainsequence-9.0.6.dev40.dist-info/licenses/LICENSE,sha256=b6Ecsk8wjsEby3lbSzdljR6hHiFGHBizt9iWV2p1c_I,1074
-mainsequence-9.0.6.dev40.dist-info/METADATA,sha256=qtgjSWqmCAC819Qw7KGhUW9lv_HTqZ0O3Mv-5hDfcqo,7402
-mainsequence-9.0.6.dev40.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
-mainsequence-9.0.6.dev40.dist-info/entry_points.txt,sha256=XGAOq3-mF2lWWpeRJvoaWyMarRcK167mH1g124aGdUs,88
-mainsequence-9.0.6.dev40.dist-info/top_level.txt,sha256=BFeblM_Q9q0cUWBvTxnfWcWQEHtf5XZi9nOS-Gdfwus,28
-mainsequence-9.0.6.dev40.dist-info/RECORD,,
+mainsequence-9.0.6.dev41.dist-info/licenses/LICENSE,sha256=b6Ecsk8wjsEby3lbSzdljR6hHiFGHBizt9iWV2p1c_I,1074
+mainsequence-9.0.6.dev41.dist-info/METADATA,sha256=0Dok8XfsCUQJJDt_13aGFAKOU8jLBJzFUbqHm-yh7Uo,7402
+mainsequence-9.0.6.dev41.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+mainsequence-9.0.6.dev41.dist-info/entry_points.txt,sha256=XGAOq3-mF2lWWpeRJvoaWyMarRcK167mH1g124aGdUs,88
+mainsequence-9.0.6.dev41.dist-info/top_level.txt,sha256=BFeblM_Q9q0cUWBvTxnfWcWQEHtf5XZi9nOS-Gdfwus,28
+mainsequence-9.0.6.dev41.dist-info/RECORD,,