truefoundry 0.4.0rc2__tar.gz → 0.4.0rc3__tar.gz

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.4.0rc2
+Version: 0.4.0rc3
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truefoundry"
-version = "0.4.0rc2"
+version = "0.4.0rc3"
 description = "Truefoundry CLI"
 authors = ["Abhishek Choudhary <abhishek@truefoundry.com>"]
 readme = "README.md"
@@ -77,6 +77,17 @@ shap =  [
     { version = ">=0.40.0", python = "<3.12" },
     { version = ">=0.45.0", python = ">=3.12" },
 ]
+pandas-stubs = "*"
+types-Pygments = "*"
+types-cffi = "*"
+types-colorama = "*"
+types-decorator = "*"
+types-protobuf = "*"
+types-python-dateutil = "*"
+types-requests = "*"
+types-six = "*"
+types-tqdm = "*"
+
 
 [tool.poetry-dynamic-versioning]
 enable = false
@@ -137,6 +148,7 @@ warn_unreachable = true
 warn_unused_ignores = true
 
 packages = [
+    "truefoundry.common",
     "truefoundry.ml",
 ]
 exclude = [

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/truefoundry/__init__.py

@@ -1 +1,3 @@
 from truefoundry.deploy.core import login, logout
+
+__all__ = ["login", "logout"]

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/truefoundry/autodeploy/agents/developer.py

@@ -23,7 +23,7 @@ from truefoundry.autodeploy.tools import (
 class Developer(Agent):
     system_prompt = """
 You are a software engineer.
-You goal is to do whatever you have to do to succesfully run a project.
+You goal is to do whatever you have to do to successfully run a project.
 If the project already contains a dockerfile, you can use that.
 If the project does not contain a dockerfile, you have to create one.
 You need to fix any issue to ensure that the image builds succesfully.

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/truefoundry/autodeploy/agents/project_identifier.py

@@ -38,7 +38,7 @@ Prefer using file type counts over list files.
 use the knowledge of file type counts to create the right glob pattern while listing files.
 Prefer yarn over npm if the project is using yarn.
 If there is  a lock file, the project maybe using yarn.
-Look for *.lock if there is a lock type file persent.
+Look for *.lock if there is a lock type file present.
 Always response with a function call.
     """
     max_iter = 10
@@ -50,7 +50,7 @@ Always response with a function call.
             None,
             description="""
 A Service is designed to run always and should never terminate.
-A Job is desiged to finish after sometime.
+A Job is designed to finish after sometime.
 """,
         )
         primary_programming_language: str = Field(

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/truefoundry/autodeploy/agents/tester.py

@@ -29,7 +29,7 @@ If you find any port number in the logs. Re run the docker image by exposing tha
 
 If the image is not running a service, try to identify whether there is any issue.
 Your goal is not to fix the issue. Your goal is to create a very detailed justification and report whether
-the testing was succesful.
+the testing was successful.
 Always response with a function call.
 Return response once you are done testing.
     """

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/truefoundry/autodeploy/cli.py

@@ -292,7 +292,7 @@ def cli(project_root_path: str, deploy: bool, workspace_fqn: str = None):
         )
         if not re.match(r"^[a-z][a-z0-9\-]{1,30}[a-z0-9]$", name):
             console.print(
-                "[bold magenta]TrueFoundry:[/] The name should be between 2-30 alphaneumaric"
+                "[bold magenta]TrueFoundry:[/] The name should be between 2-30 alphanumeric"
                 " characters and '-'. The first character should not be a digit."
             )
         else:

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/truefoundry/autodeploy/tools/list_files.py

@@ -20,7 +20,7 @@ from truefoundry.autodeploy.tools.base import (
 class ListFiles(Tool):
     description = """
 List all files.
-If you want to find all json files recuresively under directory a/b
+If you want to find all json files recursively under directory a/b
 the subdir should be a/b and pattern will be *.json
 
 If you want to find all json files recuresively under current directory

{truefoundry-0.4.0rc2/truefoundry/deploy/lib/auth → truefoundry-0.4.0rc3/truefoundry/common}/auth_service_client.py

@@ -1,40 +1,39 @@
 import time
 from abc import ABC, abstractmethod
+from typing import Optional
 
 import requests
 
-from truefoundry.deploy.lib.clients.utils import poll_for_function, request_handling
-from truefoundry.deploy.lib.const import VERSION_PREFIX
-from truefoundry.deploy.lib.exceptions import BadRequestException
-from truefoundry.deploy.lib.model.entity import DeviceCode, Token
+from truefoundry.common.constants import VERSION_PREFIX
+from truefoundry.common.entities import DeviceCode, Token
+from truefoundry.common.exceptions import BadRequestException
+from truefoundry.common.request_utils import request_handling, requests_retry_session
+from truefoundry.common.utils import poll_for_function
 from truefoundry.logger import logger
 
 
 class AuthServiceClient(ABC):
-    def __init__(self, base_url):
-        from truefoundry.deploy.lib.clients.servicefoundry_client import (
-            ServiceFoundryServiceClient,
-        )
-
-        client = ServiceFoundryServiceClient(init_session=False, base_url=base_url)
-
-        self._api_server_url = client._api_server_url
-        self._auth_server_url = client.tenant_info.auth_server_url
-        self._tenant_name = client.tenant_info.tenant_name
+    def __init__(self, tenant_name: str):
+        self._tenant_name = tenant_name
 
     @classmethod
     def from_base_url(cls, base_url: str) -> "AuthServiceClient":
-        from truefoundry.deploy.lib.clients.servicefoundry_client import (
+        from truefoundry.common.servicefoundry_client import (
             ServiceFoundryServiceClient,
         )
 
-        client = ServiceFoundryServiceClient(init_session=False, base_url=base_url)
+        client = ServiceFoundryServiceClient(base_url=base_url)
         if client.python_sdk_config.use_sfy_server_auth_apis:
-            return ServiceFoundryServerAuthServiceClient(base_url)
-        return AuthServerServiceClient(base_url)
+            return ServiceFoundryServerAuthServiceClient(
+                tenant_name=client.tenant_info.tenant_name, url=client._api_server_url
+            )
+        return AuthServerServiceClient(
+            tenant_name=client.tenant_info.tenant_name,
+            url=client.tenant_info.auth_server_url,
+        )
 
     @abstractmethod
-    def refresh_token(self, token: Token, host: str = None) -> Token: ...
+    def refresh_token(self, token: Token, host: Optional[str] = None) -> Token: ...
 
     @abstractmethod
     def get_device_code(self) -> DeviceCode: ...
@@ -46,10 +45,11 @@ class AuthServiceClient(ABC):
 
 
 class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
-    def __init__(self, base_url):
-        super().__init__(base_url)
+    def __init__(self, tenant_name: str, url: str):
+        super().__init__(tenant_name=tenant_name)
+        self._api_server_url = url
 
-    def refresh_token(self, token: Token, host: str = None) -> Token:
+    def refresh_token(self, token: Token, host: Optional[str] = None) -> Token:
         host_arg_str = f"--host {host}" if host else "--host HOST"
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
@@ -63,10 +63,11 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
             "grantType": "refresh_token",
             "returnJWT": True,
         }
-        res = requests.post(url, json=data)
+        session = requests_retry_session(retries=2)
+        response = session.post(url, json=data)
         try:
-            res = request_handling(res)
-            return Token.parse_obj(res)
+            response_data = request_handling(response)
+            return Token.parse_obj(response_data)
         except BadRequestException as ex:
             raise Exception(
                 f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
@@ -75,9 +76,10 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
     def get_device_code(self) -> DeviceCode:
         url = f"{self._api_server_url}/{VERSION_PREFIX}/oauth2/device-authorize"
         data = {"tenantName": self._tenant_name}
-        res = requests.post(url, json=data)
-        res = request_handling(res)
-        return DeviceCode.parse_obj(res)
+        session = requests_retry_session(retries=2)
+        response = session.post(url, json=data)
+        response_data = request_handling(response)
+        return DeviceCode.parse_obj(response_data)
 
     def get_token_from_device_code(
         self, device_code: str, timeout: float = 60, poll_interval_seconds: int = 1
@@ -91,7 +93,6 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
             "grantType": "device_code",
             "returnJWT": True,
         }
-        response = requests.post(url=url, json=data)
         start_time = time.monotonic()
 
         for response in poll_for_function(
@@ -116,10 +117,12 @@ class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
 
 
 class AuthServerServiceClient(AuthServiceClient):
-    def __init__(self, base_url):
-        super().__init__(base_url)
+    def __init__(self, tenant_name: str, url: str):
+        super().__init__(tenant_name=tenant_name)
+
+        self._auth_server_url = url
 
-    def refresh_token(self, token: Token, host: str = None) -> Token:
+    def refresh_token(self, token: Token, host: Optional[str] = None) -> Token:
         host_arg_str = f"--host {host}" if host else "--host HOST"
         if not token.refresh_token:
             # TODO: Add a way to propagate error messages without traceback to the output interface side
@@ -131,10 +134,11 @@ class AuthServerServiceClient(AuthServiceClient):
             "tenantName": token.tenant_name,
             "refreshToken": token.refresh_token,
         }
-        res = requests.post(url, json=data)
+        session = requests_retry_session(retries=2)
+        response = session.post(url, json=data)
         try:
-            res = request_handling(res)
-            return Token.parse_obj(res)
+            response_data = request_handling(response)
+            return Token.parse_obj(response_data)
         except BadRequestException as ex:
             raise Exception(
                 f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
@@ -143,22 +147,28 @@ class AuthServerServiceClient(AuthServiceClient):
     def get_device_code(self) -> DeviceCode:
         url = f"{self._auth_server_url}/api/{VERSION_PREFIX}/oauth/device"
         data = {"tenantName": self._tenant_name}
-        res = requests.post(url, json=data)
-        res = request_handling(res)
+        session = requests_retry_session(retries=2)
+        response = session.post(url, json=data)
+        response_data = request_handling(response)
         # TODO: temporary cleanup of incorrect attributes
-        res = {"userCode": res.get("userCode"), "deviceCode": res.get("deviceCode")}
-        return DeviceCode.parse_obj(res)
+        return DeviceCode.parse_obj(
+            {
+                "userCode": response_data.get("userCode"),
+                "deviceCode": response_data.get("deviceCode"),
+            }
+        )
 
     def get_token_from_device_code(
         self, device_code: str, timeout: float = 60, poll_interval_seconds: int = 1
     ) -> Token:
+        timeout = timeout or 60
+        poll_interval_seconds = poll_interval_seconds or 1
         url = f"{self._auth_server_url}/api/{VERSION_PREFIX}/oauth/device/token"
         data = {
             "tenantName": self._tenant_name,
             "deviceCode": device_code,
         }
         start_time = time.monotonic()
-        poll_interval_seconds = 1
 
         for response in poll_for_function(
             requests.post, poll_after_secs=poll_interval_seconds, url=url, json=data

truefoundry-0.4.0rc3/truefoundry/common/constants.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+
+TFY_CONFIG_DIR = Path.home() / ".truefoundry"
+CREDENTIAL_FILEPATH = TFY_CONFIG_DIR / "credentials.json"
+
+TFY_HOST_ENV_KEY = "TFY_HOST"
+TFY_API_KEY_ENV_KEY = "TFY_API_KEY"
+SERVICEFOUNDRY_SERVER_URL_ENV_KEY = "SERVICEFOUNDRY_SERVER_URL"
+
+API_SERVER_RELATIVE_PATH = "api/svc"
+VERSION_PREFIX = "v1"
+SERVICEFOUNDRY_CLIENT_MAX_RETRIES = 2

{truefoundry-0.4.0rc2/truefoundry/deploy/lib/auth → truefoundry-0.4.0rc3/truefoundry/common}/credential_file_manager.py

@@ -8,14 +8,14 @@ from typing import Optional
 
 from filelock import FileLock, Timeout
 
-from truefoundry.deploy.lib.const import CREDENTIAL_FILEPATH
-from truefoundry.deploy.lib.model.entity import CredentialsFileContent
+from truefoundry.common.constants import CREDENTIAL_FILEPATH
+from truefoundry.common.entities import CredentialsFileContent
 from truefoundry.logger import logger
 
 
 def _ensure_lock_taken(method):
     @wraps(method)
-    def lock_guard(self, *method_args, **method_kwargs):
+    def lock_guard(self: "CredentialsFileManager", *method_args, **method_kwargs):
         if not self.lock_taken():
             raise Exception(
                 "Trying to write to credential file without using with block"
@@ -55,17 +55,17 @@ class CredentialsFileManager:
 
     def __enter__(self) -> CredentialsFileManager:
         # The lock objects are recursive locks, which means that once acquired, they will not block on successive lock requests:
-        lock_aquired = CRED_FILE_THREAD_LOCK.acquire(timeout=self._lock_timeout)
-        if not lock_aquired:
+        lock_acquired = CRED_FILE_THREAD_LOCK.acquire(timeout=self._lock_timeout)
+        if not lock_acquired:
             raise Exception(
-                "Could not aquire CRED_FILE_THREAD_LOCK"
+                "Could not acquire CRED_FILE_THREAD_LOCK"
                 f" in {self._lock_timeout} seconds"
             )
         try:
             self._file_lock.acquire(timeout=self._lock_timeout)
         except Timeout as ex:
             raise Exception(
-                f"Failed to aquire lock on credential file within {self._lock_timeout} seconds.\n"
+                f"Failed to acquire lock on credential file within {self._lock_timeout} seconds.\n"
                 "Is any other process trying to login?"
             ) from ex
         logger.debug("Acquired file and thread lock to access credential file")

{truefoundry-0.4.0rc2/truefoundry/deploy/lib/auth → truefoundry-0.4.0rc3/truefoundry/common}/credential_provider.py

@@ -2,14 +2,11 @@ import os
 import threading
 from abc import ABC, abstractmethod
 
-from truefoundry.deploy.lib.auth.auth_service_client import AuthServiceClient
-from truefoundry.deploy.lib.auth.credential_file_manager import (
-    CredentialsFileContent,
-    CredentialsFileManager,
-)
-from truefoundry.deploy.lib.clients.utils import resolve_base_url
-from truefoundry.deploy.lib.const import API_KEY_ENV_NAME
-from truefoundry.deploy.lib.model.entity import Token
+from truefoundry.common.auth_service_client import AuthServiceClient
+from truefoundry.common.constants import TFY_API_KEY_ENV_KEY
+from truefoundry.common.credential_file_manager import CredentialsFileManager
+from truefoundry.common.entities import CredentialsFileContent, Token
+from truefoundry.common.utils import resolve_base_url
 from truefoundry.logger import logger
 
 TOKEN_REFRESH_LOCK = threading.RLock()
@@ -32,20 +29,20 @@ class CredentialProvider(ABC):
 class EnvCredentialProvider(CredentialProvider):
     def __init__(self) -> None:
         logger.debug("Using env var credential provider")
-        api_key = os.getenv(API_KEY_ENV_NAME)
+        api_key = os.getenv(TFY_API_KEY_ENV_KEY)
         if not api_key:
             raise Exception(
-                f"Value of {API_KEY_ENV_NAME} env var should be non-empty string"
+                f"Value of {TFY_API_KEY_ENV_KEY} env var should be non-empty string"
             )
         # TODO: Read host from cred file as well.
         base_url = resolve_base_url().strip("/")
         self._host = base_url
         self._auth_service = AuthServiceClient.from_base_url(base_url=base_url)
-        self._token: Token = Token(access_token=api_key, refresh_token=None)
+        self._token: Token = Token(access_token=api_key, refresh_token=None)  # type: ignore[call-arg]
 
     @staticmethod
     def can_provide() -> bool:
-        return API_KEY_ENV_NAME in os.environ
+        return TFY_API_KEY_ENV_KEY in os.environ
 
     @property
     def token(self) -> Token:

{truefoundry-0.4.0rc2/truefoundry/ml/services → truefoundry-0.4.0rc3/truefoundry/common}/entities.py

@@ -1,7 +1,6 @@
-import enum
 import time
+from enum import Enum
 from typing import Optional
-from urllib.parse import urlparse
 
 import jwt
 from typing_extensions import TypedDict
@@ -9,40 +8,21 @@ from typing_extensions import TypedDict
 from truefoundry.pydantic_v1 import BaseModel, Field, NonEmptyStr, validator
 
 
-class HostCreds(BaseModel):
-    host: str
-    token: Optional[str] = None
-
-
-class AuthServerInfo(BaseModel):
-    tenant_name: NonEmptyStr
-    auth_server_url: str
-
-    class Config:
-        allow_mutation = False
-
-
-class ArtifactCredential(BaseModel):
-    run_id: str
-    path: str
-    signed_uri: str
-
-
-class UserType(enum.Enum):
+class UserType(Enum):
     user = "user"
     serviceaccount = "serviceaccount"
 
 
 class UserInfo(BaseModel):
+    class Config:
+        allow_population_by_field_name = True
+        allow_mutation = False
+
     user_id: NonEmptyStr
     user_type: UserType = UserType.user
     email: Optional[str] = None
     tenant_name: NonEmptyStr = Field(alias="tenantName")
 
-    class Config:
-        allow_population_by_field_name = True
-        allow_mutation = False
-
 
 class _DecodedToken(TypedDict):
     tenantName: str
@@ -62,7 +42,7 @@ class Token(BaseModel):
         allow_mutation = False
 
     @validator("decoded_value", always=True, pre=True)
-    def _decode_jwt(cls, v, values, **kwargs):
+    def _decode_jwt(cls, v, values, **kwargs) -> _DecodedToken:
         access_token = values["access_token"]
         return jwt.decode(
             access_token,
@@ -75,34 +55,70 @@ class Token(BaseModel):
 
     @property
     def tenant_name(self) -> str:
-        assert self.decoded_value is not None
         return self.decoded_value["tenantName"]
 
     def is_going_to_be_expired(self, buffer_in_seconds: int = 120) -> bool:
-        assert self.decoded_value is not None
         exp = int(self.decoded_value["exp"])
         return (exp - time.time()) < buffer_in_seconds
 
     def to_user_info(self) -> UserInfo:
-        assert self.decoded_value is not None
-        return UserInfo(
+        return UserInfo(  # type: ignore[call-arg]
             user_id=self.decoded_value["username"],
-            email=(
-                self.decoded_value["email"] if "email" in self.decoded_value else None
-            ),
+            email=self.decoded_value["email"]
+            if "email" in self.decoded_value
+            else None,
             user_type=UserType(self.decoded_value.get("userType", UserType.user.value)),
-            tenant_name=self.tenant_name,  # type: ignore[call-arg]
+            tenant_name=self.tenant_name,
         )
 
 
-class DeviceCode(BaseModel):
-    user_code: str = Field(alias="userCode")
-    device_code: str = Field(alias="deviceCode")
+class CredentialsFileContent(BaseModel):
+    class Config:
+        allow_mutation = False
+
+    access_token: NonEmptyStr = Field(repr=False)
+    refresh_token: Optional[NonEmptyStr] = Field(repr=False)
+    host: NonEmptyStr
+
+    def to_token(self) -> Token:
+        return Token(  # type: ignore[call-arg]
+            access_token=self.access_token, refresh_token=self.refresh_token
+        )
+
+
+class TenantInfo(BaseModel):
+    class Config:
+        allow_population_by_field_name = True
+        allow_mutation = False
+
+    tenant_name: NonEmptyStr = Field(alias="tenantName")
+    auth_server_url: str
+
 
+class PythonSDKConfig(BaseModel):
     class Config:
         allow_population_by_field_name = True
         allow_mutation = False
 
-    def get_user_clickable_url(self, tracking_uri: str) -> str:
-        parsed_tracking_uri = urlparse(tracking_uri)
-        return f"{parsed_tracking_uri.scheme}://{parsed_tracking_uri.netloc}/authorize/device?userCode={self.user_code}"
+    min_version: str = Field(alias="minVersion")
+    truefoundry_cli_min_version: str = Field(alias="truefoundryCliMinVersion")
+    use_sfy_server_auth_apis: Optional[bool] = Field(
+        alias="useSFYServerAuthAPIs", default=False
+    )
+
+
+class DeviceCode(BaseModel):
+    class Config:
+        allow_population_by_field_name = True
+        allow_mutation = False
+
+    user_code: str = Field(alias="userCode")
+    device_code: str = Field(alias="deviceCode")
+    verification_url: Optional[str] = Field(alias="verificationURI")
+    complete_verification_url: Optional[str] = Field(alias="verificationURIComplete")
+    expires_in_seconds: int = Field(alias="expiresInSeconds", default=60)
+    interval_in_seconds: int = Field(alias="intervalInSeconds", default=1)
+    message: Optional[str] = Field(alias="message")
+
+    def get_user_clickable_url(self, auth_host: str) -> str:
+        return f"{auth_host}/authorize/device?userCode={self.user_code}"

{truefoundry-0.4.0rc2/truefoundry/deploy/lib → truefoundry-0.4.0rc3/truefoundry/common}/exceptions.py

@@ -10,12 +10,3 @@ class BadRequestException(Exception):
 
     def __str__(self):
         return self.message
-
-
-class ConfigurationException(Exception):
-    def __init__(self, message: Optional[str] = None):
-        super().__init__()
-        self.message = message
-
-    def __str__(self):
-        return self.message

{truefoundry-0.4.0rc2 → truefoundry-0.4.0rc3}/truefoundry/common/request_utils.py

@@ -1,13 +1,38 @@
+from typing import Optional
+
 import requests
+from requests import Response
 from requests.adapters import HTTPAdapter
 from urllib3 import Retry
 
+from truefoundry.common.exceptions import BadRequestException
+
+
+def request_handling(response: Response):
+    try:
+        status_code = response.status_code
+    except Exception as e:
+        raise Exception("Unknown error occurred. Couldn't get status code.") from e
+    if 200 <= status_code <= 299:
+        if response.content == b"":
+            return None
+        return response.json()
+    if 400 <= status_code <= 499:
+        try:
+            message = str(response.json()["message"])
+        except Exception:
+            message = response.text
+        raise BadRequestException(status_code=response.status_code, message=message)
+    if 500 <= status_code <= 599:
+        raise Exception(response.content)
+
 
 def urllib3_retry(
-    retries=5,
-    backoff_factor=0.3,
-    status_forcelist=(408, 429, 500, 502, 503, 504),
+    retries: int = 2,
+    backoff_factor: float = 0.3,
+    status_forcelist=(408, 429, 500, 502, 503, 504, 524),
     method_whitelist=frozenset({"GET", "POST"}),
+    raise_on_status: bool = False,
 ):
     retry = Retry(
         total=retries,
@@ -18,17 +43,19 @@ def urllib3_retry(
         allowed_methods=method_whitelist,
         status_forcelist=status_forcelist,
         respect_retry_after_header=True,
+        raise_on_status=raise_on_status,
     )
     return retry
 
 
 def requests_retry_session(
-    retries=5,
-    backoff_factor=0.3,
-    status_forcelist=(408, 429, 500, 502, 503, 504),
+    retries: int = 2,
+    backoff_factor: float = 0.3,
+    status_forcelist=(408, 429, 500, 502, 503, 504, 524),
     method_whitelist=frozenset({"GET", "POST"}),
-    session=None,
-):
+    raise_on_status: bool = False,
+    session: Optional[requests.Session] = None,
+) -> requests.Session:
     """
     Returns a `requests` session with retry capabilities for certain HTTP status codes.
 
@@ -49,6 +76,7 @@ def requests_retry_session(
         backoff_factor=backoff_factor,
         status_forcelist=status_forcelist,
         method_whitelist=method_whitelist,
+        raise_on_status=raise_on_status,
     )
     adapter = HTTPAdapter(max_retries=retry)
     session.mount("http://", adapter)