truefoundry 0.5.3rc4__py3-none-any.whl → 0.5.4__py3-none-any.whl

truefoundry/__init__.py

@@ -1,3 +1,12 @@
+from truefoundry.common.warnings import (
+    suppress_truefoundry_deprecation_warnings,
+    surface_truefoundry_deprecation_warnings,
+)
 from truefoundry.deploy.core import login, logout
 
-__all__ = ["login", "logout"]
+surface_truefoundry_deprecation_warnings()
+__all__ = [
+    "login",
+    "logout",
+    "suppress_truefoundry_deprecation_warnings",
+]

truefoundry/autodeploy/cli.py

@@ -32,8 +32,8 @@ from truefoundry.autodeploy.tools.commit import CommitConfirmation
 from truefoundry.autodeploy.tools.docker_run import DockerRun, DockerRunLog
 from truefoundry.autodeploy.utils.client import get_git_binary
 from truefoundry.cli.const import COMMAND_CLS
+from truefoundry.common.session import Session
 from truefoundry.deploy import Build, DockerFileBuild, Job, LocalSource, Port, Service
-from truefoundry.deploy.lib.auth.servicefoundry_session import ServiceFoundrySession
 from truefoundry.deploy.lib.clients.servicefoundry_client import (
     ServiceFoundryServiceClient,
 )
@@ -45,7 +45,7 @@ def _get_openai_client() -> OpenAI:
             api_key=AUTODEPLOY_OPENAI_API_KEY, base_url=AUTODEPLOY_OPENAI_BASE_URL
         )
     try:
-        session = ServiceFoundrySession()
+        session = Session.new()
         resp = requests.get(
             f"{AUTODEPLOY_TFY_BASE_URL}/api/svc/v1/llm-gateway/access-details",
             headers={

truefoundry/cli/__main__.py

@@ -4,7 +4,6 @@ import sys
 import rich_click as click
 
 from truefoundry import logger
-from truefoundry.cli.commands.pat import get_generate_pat_command
 from truefoundry.cli.config import CliConfig
 from truefoundry.cli.const import GROUP_CLS
 from truefoundry.cli.util import setup_rich_click
@@ -22,7 +21,6 @@ from truefoundry.deploy.cli.commands import (
     get_terminate_command,
     get_trigger_command,
 )
-from truefoundry.gateway.cli.cli import get_gateway_cli
 from truefoundry.ml.cli.cli import get_ml_cli
 from truefoundry.version import __version__
 
@@ -90,8 +88,6 @@ def create_truefoundry_cli() -> click.MultiCommand:
     cli.add_command(get_trigger_command())
     cli.add_command(get_terminate_command())
     cli.add_command(get_ml_cli())
-    cli.add_command(get_gateway_cli())
-    cli.add_command(get_generate_pat_command())
 
     if not (sys.platform.startswith("win32") or sys.platform.startswith("cygwin")):
         cli.add_command(get_patch_command())

truefoundry/cli/util.py

@@ -5,13 +5,13 @@ from typing import Dict
 
 import rich_click as click
 from packaging.version import parse as parse_version
-from requests.exceptions import ConnectionError
+from requests.exceptions import ConnectionError, Timeout
 from rich.padding import Padding
 from rich.panel import Panel
 from rich.table import Table
 
 from truefoundry.cli.console import console
-from truefoundry.common.exceptions import BadRequestException
+from truefoundry.common.exceptions import HttpRequestException
 from truefoundry.common.utils import is_debug_env_set
 
 
@@ -25,12 +25,21 @@ def setup_rich_click():
 def handle_exception(exception):
     if is_debug_env_set():
         console.print_exception(show_locals=True)
-    if isinstance(exception, BadRequestException):
+    if isinstance(exception, HttpRequestException):
         print_dict_as_table_panel(
             {"Status Code": str(exception.status_code), "Error": exception.message},
             title="Command Failed",
             border_style="red",
         )
+    elif isinstance(exception, Timeout):
+        loc = ""
+        if exception.request:
+            loc = f" at {exception.request.url}"
+        print_dict_as_table_panel(
+            {"Error": f"Request to TrueFoundry{loc} timed out."},
+            title="Command Failed",
+            border_style="red",
+        )
     elif isinstance(exception, ConnectionError):
         loc = ""
         if exception.request:

truefoundry/common/auth_service_client.py

@@ -7,7 +7,10 @@ import requests
 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.request_utils import (
+    request_handling,
+    requests_retry_session,
+)
 from truefoundry.common.utils import poll_for_function, relogin_error_message
 from truefoundry.logger import logger
 
@@ -16,14 +19,13 @@ class AuthServiceClient(ABC):
     def __init__(self, tenant_name: str):
         self._tenant_name = tenant_name
 
-    # TODO (chiragjn): Rename base_url to tfy_host
     @classmethod
-    def from_base_url(cls, base_url: str) -> "AuthServiceClient":
+    def from_tfy_host(cls, tfy_host: str) -> "AuthServiceClient":
         from truefoundry.common.servicefoundry_client import (
             ServiceFoundryServiceClient,
         )
 
-        client = ServiceFoundryServiceClient(base_url=base_url)
+        client = ServiceFoundryServiceClient(tfy_host=tfy_host)
         if client.python_sdk_config.use_sfy_server_auth_apis:
             return ServiceFoundryServerAuthServiceClient(
                 tenant_name=client.tenant_info.tenant_name, url=client._api_server_url
@@ -159,6 +161,7 @@ class AuthServerServiceClient(AuthServiceClient):
         session = requests_retry_session(retries=2)
         response = session.post(url, json=data)
         response_data = request_handling(response)
+        assert isinstance(response_data, dict)
         # TODO: temporary cleanup of incorrect attributes
         return DeviceCode.parse_obj(
             {

truefoundry/common/constants.py

@@ -38,6 +38,9 @@ class TrueFoundrySdkEnv(BaseSettings):
     )
     TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT: int = 5  # default: 5 seconds
     TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL: int = 3600  # default: 1 hour
+    TFY_INTERNAL_MULTIPART_UPLOAD_FINALIZE_SIGNED_URL_TIMEOUT: int = (
+        24 * 60 * 60
+    )  # default: 24 hour
     TFY_INTERNAL_SIGNED_URL_REQUEST_TIMEOUT: int = 3600  # default: 1 hour
     TFY_INTERNAL_SIGNED_URL_CLIENT_LOG_LEVEL: str = "WARNING"
 
@@ -63,4 +66,3 @@ API_SERVER_RELATIVE_PATH = "api/svc"
 MLFOUNDRY_SERVER_RELATIVE_PATH = "api/ml"
 VERSION_PREFIX = "v1"
 SERVICEFOUNDRY_CLIENT_MAX_RETRIES = 2
-GATEWAY_SERVER_RELATIVE_PATH = "api/llm"

truefoundry/common/credential_provider.py

@@ -17,10 +17,9 @@ class CredentialProvider(ABC):
     @abstractmethod
     def token(self) -> Token: ...
 
-    # TODO (chiragjn): Rename base_url to tfy_host
     @property
     @abstractmethod
-    def base_url(self) -> str: ...
+    def tfy_host(self) -> str: ...
 
     @staticmethod
     @abstractmethod
@@ -38,7 +37,7 @@ class EnvCredentialProvider(CredentialProvider):
                 f"Value of {TFY_API_KEY_ENV_KEY} env var should be non-empty string"
             )
         self._host = resolve_tfy_host()
-        self._auth_service = AuthServiceClient.from_base_url(base_url=self._host)
+        self._auth_service = AuthServiceClient.from_tfy_host(tfy_host=self._host)
         self._token: Token = Token(access_token=api_key, refresh_token=None)  # type: ignore[call-arg]
 
     @staticmethod
@@ -51,12 +50,12 @@ class EnvCredentialProvider(CredentialProvider):
             if self._token.is_going_to_be_expired():
                 logger.info("Refreshing access token")
                 self._token = self._auth_service.refresh_token(
-                    self._token, self.base_url
+                    self._token, self.tfy_host
                 )
             return self._token
 
     @property
-    def base_url(self) -> str:
+    def tfy_host(self) -> str:
         return self._host
 
 
@@ -70,7 +69,7 @@ class FileCredentialProvider(CredentialProvider):
             self._token = self._last_cred_file_content.to_token()
             self._host = self._last_cred_file_content.host
 
-        self._auth_service = AuthServiceClient.from_base_url(base_url=self._host)
+        self._auth_service = AuthServiceClient.from_tfy_host(tfy_host=self._host)
 
     @staticmethod
     def can_provide() -> bool:
@@ -91,7 +90,7 @@ class FileCredentialProvider(CredentialProvider):
 
                 if new_cred_file_content == self._last_cred_file_content:
                     self._token = self._auth_service.refresh_token(
-                        self._token, self.base_url
+                        self._token, self.tfy_host
                     )
                     self._last_cred_file_content = CredentialsFileContent(
                         host=self._host,
@@ -115,5 +114,5 @@ class FileCredentialProvider(CredentialProvider):
                 )
 
     @property
-    def base_url(self) -> str:
+    def tfy_host(self) -> str:
         return self._host

truefoundry/common/exceptions.py

@@ -1,12 +1,16 @@
 from typing import Optional
 
 
-class BadRequestException(Exception):
-    def __init__(self, status_code: int, message: Optional[str] = None):
-        super().__init__()
-        self.status_code = status_code
-        self.message = message
+# TODO (chiragjn): We need to establish uniform exception handling across codebase
+class HttpRequestException(Exception):
+    def __init__(self, message: str, status_code: Optional[int] = None):
+        self.message = str(message)
         self.status_code = status_code
+        super().__init__(message)
+
+    def __str__(self) -> str:
+        return self.message or ""
+
 
-    def __str__(self):
-        return self.message
+class BadRequestException(HttpRequestException):
+    pass

truefoundry/common/request_utils.py

@@ -1,30 +1,40 @@
-from typing import Optional
+import json
+from contextlib import contextmanager
+from typing import Any, Dict, Optional, Type, Union
 
 import requests
 from requests import Response
 from requests.adapters import HTTPAdapter
 from urllib3 import Retry
 
-from truefoundry.common.exceptions import BadRequestException
+from truefoundry.common.exceptions import BadRequestException, HttpRequestException
 
 
-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
+# TODO (chiragjn): Rename this to json_response_handling
+def request_handling(response: Response) -> Optional[Any]:
+    status_code = response.status_code
     if 200 <= status_code <= 299:
         if response.content == b"":
+            # TODO (chiragjn): Do we really need this empty check?
             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)
+            return response.json()
+        except json.JSONDecodeError as je:
+            raise ValueError(
+                f"Failed to parse response as json. Response: {response.text}",
+            ) from je
+
+    try:
+        message = str(response.json()["message"])
+    except Exception:
+        message = response.text
+
+    if 400 <= status_code <= 499:
+        raise BadRequestException(message=message, status_code=response.status_code)
     if 500 <= status_code <= 599:
-        raise Exception(response.content)
+        raise HttpRequestException(message=message, status_code=response.status_code)
+
+    raise HttpRequestException("Unknown error occurred", status_code=status_code)
 
 
 def urllib3_retry(
@@ -82,3 +92,75 @@ def requests_retry_session(
     session.mount("http://", adapter)
     session.mount("https://", adapter)
     return session
+
+
+def http_request(
+    *,
+    method: str,
+    url: str,
+    token: Optional[str] = None,
+    timeout=None,
+    headers: Optional[Dict[str, str]] = None,
+    session: Optional[requests.Session] = None,
+    exception_class: Type[HttpRequestException] = HttpRequestException,
+    **kwargs,
+) -> requests.Response:
+    session = session or requests.Session()
+    headers = headers or {}
+    if token is not None:
+        headers["Authorization"] = f"Bearer {token}"
+    try:
+        response = session.request(
+            method=method, url=url, headers=headers, timeout=timeout, **kwargs
+        )
+        return response
+    except requests.exceptions.ConnectionError as ce:
+        raise exception_class("Failed to connect to TrueFoundry") from ce
+    except requests.exceptions.Timeout as te:
+        raise exception_class(f"Request to {url} timed out") from te
+    except Exception as e:
+        raise exception_class(f"Request to {url} failed with error {str(e)}") from e
+
+
+@contextmanager
+def cloud_storage_http_request(
+    *,
+    method: str,
+    url: str,
+    session: Optional[requests.Session] = None,
+    timeout: Optional[Union[int, float]] = None,
+    exception_class: Type[HttpRequestException] = HttpRequestException,
+    **kwargs,
+):
+    """
+    Performs an HTTP PUT/GET request using Python's `requests` module with automatic retry.
+    """
+    # Note: This does not support auth and is only meant to be used for pre-signed URLs.
+    session = session or requests_retry_session(retries=5, backoff_factor=0.5)
+    headers = kwargs.get("headers", {}) or {}
+    if "blob.core.windows.net" in url:
+        headers.update({"x-ms-blob-type": "BlockBlob"})
+    if method.lower() not in ("put", "get"):
+        raise ValueError(f"Illegal http method: {method}")
+    yield http_request(
+        method=method,
+        url=url,
+        session=session,
+        timeout=timeout,
+        headers=headers,
+        exception_class=exception_class,
+        **kwargs,
+    )
+
+
+# TODO: Try and eliminate this
+def augmented_raise_for_status(
+    response, exception_class: Type[HttpRequestException] = HttpRequestException
+):
+    try:
+        response.raise_for_status()
+    except requests.exceptions.HTTPError as he:
+        raise exception_class(
+            message=f"Request failed with status code {he.response.status_code}. Response: {he.response.text}",
+            status_code=he.response.status_code,
+        ) from he

truefoundry/common/servicefoundry_client.py

@@ -13,7 +13,10 @@ from truefoundry.common.entities import (
     PythonSDKConfig,
     TenantInfo,
 )
-from truefoundry.common.request_utils import request_handling, requests_retry_session
+from truefoundry.common.request_utils import (
+    request_handling,
+    requests_retry_session,
+)
 from truefoundry.common.utils import (
     get_tfy_servers_config,
     timed_lru_cache,
@@ -25,19 +28,8 @@ from truefoundry.version import __version__
 def check_min_cli_version(fn):
     @functools.wraps(fn)
     def inner(*args, **kwargs):
-        if __version__ != "0.0.0":
-            client: "ServiceFoundryServiceClient" = args[0]
-            # "0.0.0" indicates dev version
-            # noinspection PyProtectedMember
-            min_cli_version_required = client._min_cli_version_required
-            if version.parse(__version__) < version.parse(min_cli_version_required):
-                raise Exception(
-                    "You are using an outdated version of `truefoundry`.\n"
-                    f"Run `pip install truefoundry>={min_cli_version_required}` to install the supported version.",
-                )
-        else:
-            logger.debug("Ignoring minimum cli version check")
-
+        client: "ServiceFoundryServiceClient" = args[0]
+        client.check_min_cli_version()
         return fn(*args, **kwargs)
 
     return inner
@@ -51,34 +43,33 @@ def session_with_retries(
 
 @timed_lru_cache(seconds=30 * 60)
 def _cached_get_tenant_info(api_server_url: str, tenant_host: str) -> TenantInfo:
-    res = session_with_retries().get(
+    response = session_with_retries().get(
         url=f"{api_server_url}/{VERSION_PREFIX}/tenant-id",
         params={"hostName": tenant_host},
     )
-    res = request_handling(res)
-    return TenantInfo.parse_obj(res)
+    response_data = request_handling(response)
+    return TenantInfo.parse_obj(response_data)
 
 
 @timed_lru_cache(seconds=30 * 60)
 def _cached_get_python_sdk_config(api_server_url: str) -> PythonSDKConfig:
-    res = session_with_retries().get(
+    response = session_with_retries().get(
         url=f"{api_server_url}/{VERSION_PREFIX}/min-cli-version"
     )
-    res = request_handling(res)
-    return PythonSDKConfig.parse_obj(res)
+    response_data = request_handling(response)
+    return PythonSDKConfig.parse_obj(response_data)
 
 
 class ServiceFoundryServiceClient:
-    # TODO (chiragjn): Rename base_url to tfy_host
-    def __init__(self, base_url: str):
-        self._base_url = base_url.strip("/")
-        tfy_servers_config = get_tfy_servers_config(self._base_url)
+    def __init__(self, tfy_host: str):
+        self._tfy_host = tfy_host.strip("/")
+        tfy_servers_config = get_tfy_servers_config(self._tfy_host)
         self._tenant_host = tfy_servers_config.tenant_host
         self._api_server_url = tfy_servers_config.servicefoundry_server_url
 
     @property
-    def base_url(self) -> str:
-        return self._base_url
+    def tfy_host(self) -> str:
+        return self._tfy_host
 
     @property
     def tenant_info(self) -> TenantInfo:
@@ -95,6 +86,17 @@ class ServiceFoundryServiceClient:
 
     @functools.cached_property
     def _min_cli_version_required(self) -> str:
-        return _cached_get_python_sdk_config(
-            self._api_server_url
-        ).truefoundry_cli_min_version
+        return self.python_sdk_config.truefoundry_cli_min_version
+
+    def check_min_cli_version(self) -> None:
+        if __version__ != "0.0.0":
+            # "0.0.0" indicates dev version
+            # noinspection PyProtectedMember
+            min_cli_version_required = self._min_cli_version_required
+            if version.parse(__version__) < version.parse(min_cli_version_required):
+                raise Exception(
+                    "You are using an outdated version of `truefoundry`.\n"
+                    f"Run `pip install truefoundry>={min_cli_version_required}` to install the supported version.",
+                )
+        else:
+            logger.debug("Ignoring minimum cli version check")

truefoundry/common/session.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import threading
+from typing import Optional
+
+from truefoundry.common.credential_provider import (
+    CredentialProvider,
+    EnvCredentialProvider,
+    FileCredentialProvider,
+)
+from truefoundry.common.entities import UserInfo
+from truefoundry.common.utils import relogin_error_message
+from truefoundry.logger import logger
+
+SESSION_LOCK = threading.RLock()
+ACTIVE_SESSION: Optional["Session"] = None
+
+
+class Session:
+    def __init__(self) -> None:
+        self._closed = False
+        self._cred_provider: Optional[CredentialProvider] = self._get_cred_provider()
+        self._user_info: Optional[UserInfo] = self._cred_provider.token.to_user_info()
+
+    @staticmethod
+    def _get_cred_provider() -> CredentialProvider:
+        final_cred_provider = None
+        for cred_provider in [EnvCredentialProvider, FileCredentialProvider]:
+            if cred_provider.can_provide():
+                final_cred_provider = cred_provider()
+                break
+        if final_cred_provider is None:
+            raise Exception(
+                relogin_error_message(
+                    "No active session found. Perhaps you are not logged in?",
+                )
+            )
+        return final_cred_provider
+
+    @classmethod
+    def new(cls) -> "Session":
+        global ACTIVE_SESSION
+        with SESSION_LOCK:
+            new_session = cls()
+            if ACTIVE_SESSION and ACTIVE_SESSION == new_session:
+                return ACTIVE_SESSION
+
+            if ACTIVE_SESSION:
+                ACTIVE_SESSION.close()
+
+            ACTIVE_SESSION = new_session
+            logger.info(
+                "Logged in to %r as %r (%s)",
+                new_session.tfy_host,
+                new_session.user_info.user_id,
+                new_session.user_info.email or new_session.user_info.user_type.value,
+            )
+
+            return ACTIVE_SESSION
+
+    def close(self):
+        self._closed = True
+        self._user_info = None
+        self._cred_provider = None
+
+    def _assert_not_closed(self):
+        if self._closed:
+            raise Exception("This session has been deactivated.")
+
+    @property
+    def access_token(self) -> str:
+        assert self._cred_provider is not None
+        return self._cred_provider.token.access_token
+
+    @property
+    def tfy_host(self) -> str:
+        assert self._cred_provider is not None
+        return self._cred_provider.tfy_host
+
+    @property
+    def user_info(self) -> UserInfo:
+        self._assert_not_closed()
+        assert self._user_info is not None
+        return self._user_info
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Session):
+            return False
+        return (
+            type(self._cred_provider) == type(other._cred_provider)  # noqa: E721
+            and self.user_info == other.user_info
+            and self.tfy_host == other.tfy_host
+        )