truefoundry 0.2.10__py3-none-any.whl → 0.3.0__py3-none-any.whl

truefoundry/deploy/io/rich_output_callback.py

@@ -0,0 +1,27 @@
+from rich.console import Console
+from rich.highlighter import ReprHighlighter
+from rich.panel import Panel
+from rich.text import Text
+
+from truefoundry.deploy.io.output_callback import OutputCallBack
+
+
+def _text(line):
+    return Text.from_ansi(str(line)) if "\x1b" in line else str(line)
+
+
+class RichOutputCallBack(OutputCallBack):
+    console = Console(soft_wrap=True)
+    highlighter = ReprHighlighter()
+
+    def print_header(self, line):
+        self.console.rule(_text(line), style="cyan")
+
+    def print_line(self, line):
+        self.console.print(_text(line))
+
+    def print_lines_in_panel(self, lines, header=None):
+        self.console.print(Panel(self.highlighter("\n".join(lines)), title=header))
+
+    def print_code_in_panel(self, lines, header=None, lang="python"):
+        self.print_lines_in_panel(lines, header)

truefoundry/deploy/json_util.py

@@ -0,0 +1,7 @@
+import datetime
+
+
+def json_default_encoder(o):
+    if isinstance(o, datetime.datetime):
+        return o.isoformat()
+    raise TypeError(f"Cannot json encode {type(o)}: {o}")

truefoundry/deploy/lib/auth/auth_service_client.py

@@ -0,0 +1,181 @@
+import time
+from abc import ABC, abstractmethod
+
+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.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
+
+    @classmethod
+    def from_base_url(cls, base_url: str) -> "AuthServiceClient":
+        from truefoundry.deploy.lib.clients.servicefoundry_client import (
+            ServiceFoundryServiceClient,
+        )
+
+        client = ServiceFoundryServiceClient(init_session=False, base_url=base_url)
+        if client.python_sdk_config.use_sfy_server_auth_apis:
+            return ServiceFoundryServerAuthServiceClient(base_url)
+        return AuthServerServiceClient(base_url)
+
+    @abstractmethod
+    def refresh_token(self, token: Token, host: str = None) -> Token: ...
+
+    @abstractmethod
+    def get_device_code(self) -> DeviceCode: ...
+
+    @abstractmethod
+    def get_token_from_device_code(
+        self, device_code: str, timeout: float = 60, poll_interval_seconds: int = 1
+    ) -> Token: ...
+
+
+class ServiceFoundryServerAuthServiceClient(AuthServiceClient):
+    def __init__(self, base_url):
+        super().__init__(base_url)
+
+    def refresh_token(self, token: Token, host: 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
+            raise Exception(
+                f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
+            )
+        url = f"{self._api_server_url}/{VERSION_PREFIX}/oauth2/token"
+        data = {
+            "tenantName": token.tenant_name,
+            "refreshToken": token.refresh_token,
+            "grantType": "refresh_token",
+            "returnJWT": True,
+        }
+        res = requests.post(url, json=data)
+        try:
+            res = request_handling(res)
+            return Token.parse_obj(res)
+        except BadRequestException as ex:
+            raise Exception(
+                f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
+            ) from ex
+
+    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)
+
+    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._api_server_url}/{VERSION_PREFIX}/oauth2/token"
+        data = {
+            "tenantName": self._tenant_name,
+            "deviceCode": device_code,
+            "grantType": "device_code",
+            "returnJWT": True,
+        }
+        response = requests.post(url=url, json=data)
+        start_time = time.monotonic()
+
+        for response in poll_for_function(
+            requests.post, poll_after_secs=poll_interval_seconds, url=url, json=data
+        ):
+            if response.status_code == 201:
+                response = response.json()
+                return Token.parse_obj(response)
+            elif response.status_code == 202:
+                logger.debug("User has not authorized yet. Checking again.")
+            else:
+                raise Exception(
+                    "Failed to get token using device code. "
+                    f"status_code {response.status_code},\n {response.text}"
+                )
+            time_elapsed = time.monotonic() - start_time
+            if time_elapsed > timeout:
+                logger.warning("Polled server for %s secs.", int(time_elapsed))
+                break
+
+        raise Exception(f"Did not get authorized within {timeout} seconds.")
+
+
+class AuthServerServiceClient(AuthServiceClient):
+    def __init__(self, base_url):
+        super().__init__(base_url)
+
+    def refresh_token(self, token: Token, host: 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
+            raise Exception(
+                f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
+            )
+        url = f"{self._auth_server_url}/api/{VERSION_PREFIX}/oauth/token/refresh"
+        data = {
+            "tenantName": token.tenant_name,
+            "refreshToken": token.refresh_token,
+        }
+        res = requests.post(url, json=data)
+        try:
+            res = request_handling(res)
+            return Token.parse_obj(res)
+        except BadRequestException as ex:
+            raise Exception(
+                f"Unable to resume login session. Please log in again using `tfy login {host_arg_str} --relogin`"
+            ) from ex
+
+    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)
+        # TODO: temporary cleanup of incorrect attributes
+        res = {"userCode": res.get("userCode"), "deviceCode": res.get("deviceCode")}
+        return DeviceCode.parse_obj(res)
+
+    def get_token_from_device_code(
+        self, device_code: str, timeout: float = 60, poll_interval_seconds: int = 1
+    ) -> Token:
+        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
+        ):
+            if response.status_code == 201:
+                response = response.json()
+                return Token.parse_obj(response)
+            elif response.status_code == 202:
+                logger.debug("User has not authorized yet. Checking again.")
+            else:
+                raise Exception(
+                    "Failed to get token using device code. "
+                    f"status_code {response.status_code},\n {response.text}"
+                )
+            time_elapsed = time.monotonic() - start_time
+            if time_elapsed > timeout:
+                logger.warning("Polled server for %s secs.", int(time_elapsed))
+                break
+
+        raise Exception(f"Did not get authorized within {timeout} seconds.")

truefoundry/deploy/lib/auth/credential_file_manager.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import os
+import threading
+from functools import lru_cache, wraps
+from pathlib import Path
+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.logger import logger
+
+
+def _ensure_lock_taken(method):
+    @wraps(method)
+    def lock_guard(self, *method_args, **method_kwargs):
+        if not self.lock_taken():
+            raise Exception(
+                "Trying to write to credential file without using with block"
+            )
+        return method(self, *method_args, **method_kwargs)
+
+    return lock_guard
+
+
+CRED_FILE_THREAD_LOCK = threading.RLock()
+
+
+@lru_cache(maxsize=None)
+def get_file_lock(lock_file_path: str) -> FileLock:
+    return FileLock(lock_file_path)
+
+
+class CredentialsFileManager:
+    def __init__(
+        self,
+        credentials_file_path: Path = CREDENTIAL_FILEPATH,
+        lock_timeout: float = 60.0,
+    ) -> None:
+        credentials_file_path = credentials_file_path.absolute()
+        logger.debug("credential file path %r", credentials_file_path)
+
+        credentials_lock_file_path = f"{credentials_file_path}.lock"
+        logger.debug("credential lock file path %r", credentials_lock_file_path)
+
+        self._credentials_file_path = credentials_file_path
+        cred_file_dir = credentials_file_path.parent
+        cred_file_dir.mkdir(exist_ok=True, parents=True)
+
+        self._file_lock = get_file_lock(credentials_lock_file_path)
+        self._lock_timeout = lock_timeout
+        self._lock_owner: Optional[int] = None
+
+    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:
+            raise Exception(
+                "Could not aquire 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"
+                "Is any other process trying to login?"
+            ) from ex
+        logger.debug("Acquired file and thread lock to access credential file")
+        self._lock_owner = threading.get_ident()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self._file_lock.release()
+        CRED_FILE_THREAD_LOCK.release()
+        logger.debug("Released file and thread lock to access credential file")
+        self._lock_owner = None
+
+    def lock_taken(self) -> bool:
+        return self._lock_owner == threading.get_ident()
+
+    @_ensure_lock_taken
+    def read(self) -> CredentialsFileContent:
+        try:
+            return CredentialsFileContent.parse_file(self._credentials_file_path)
+        except Exception as ex:
+            raise Exception(
+                "Error while reading the credentials file "
+                f"{self._credentials_file_path}. Please login again "
+                "using `tfy login --relogin` or `tfy.login(relogin=True)` function"
+            ) from ex
+
+    @_ensure_lock_taken
+    def write(self, credentials_file_content: CredentialsFileContent) -> None:
+        if not isinstance(credentials_file_content, CredentialsFileContent):
+            raise Exception(
+                "Only object of type `CredentialsFileContent` is allowed. "
+                f"Got {type(credentials_file_content)}"
+            )
+        logger.debug("Updating the credential file content")
+        with open(self._credentials_file_path, "w", encoding="utf8") as file:
+            file.write(credentials_file_content.json())
+
+    @_ensure_lock_taken
+    def delete(self) -> bool:
+        if not os.path.exists(self._credentials_file_path):
+            return False
+        os.remove(self._credentials_file_path)
+        return True
+
+    @_ensure_lock_taken
+    def exists(self) -> bool:
+        return self._credentials_file_path.exists()

truefoundry/deploy/lib/auth/credential_provider.py

@@ -0,0 +1,131 @@
+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.logger import logger
+
+TOKEN_REFRESH_LOCK = threading.RLock()
+
+
+class CredentialProvider(ABC):
+    @property
+    @abstractmethod
+    def token(self) -> Token: ...
+
+    @property
+    @abstractmethod
+    def base_url(self) -> str: ...
+
+    @staticmethod
+    @abstractmethod
+    def can_provide() -> bool: ...
+
+
+class EnvCredentialProvider(CredentialProvider):
+    def __init__(self) -> None:
+        from truefoundry.deploy.lib.clients.servicefoundry_client import (
+            ServiceFoundryServiceClient,
+        )
+
+        logger.debug("Using env var credential provider")
+        api_key = os.getenv(API_KEY_ENV_NAME)
+        if not api_key:
+            raise Exception(
+                f"Value of {API_KEY_ENV_NAME} 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)
+
+        servicefoundry_client = ServiceFoundryServiceClient(
+            init_session=False, base_url=base_url
+        )
+        self._token: Token = servicefoundry_client.get_token_from_api_key(
+            api_key=api_key
+        )
+
+    @staticmethod
+    def can_provide() -> bool:
+        return API_KEY_ENV_NAME in os.environ
+
+    @property
+    def token(self) -> Token:
+        with TOKEN_REFRESH_LOCK:
+            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
+                )
+            return self._token
+
+    @property
+    def base_url(self) -> str:
+        return self._host
+
+
+class FileCredentialProvider(CredentialProvider):
+    def __init__(self) -> None:
+        logger.debug("Using file credential provider")
+        self._cred_file = CredentialsFileManager()
+
+        with self._cred_file:
+            self._last_cred_file_content = self._cred_file.read()
+            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)
+
+    @staticmethod
+    def can_provide() -> bool:
+        with CredentialsFileManager() as cred_file:
+            return cred_file.exists()
+
+    @property
+    def token(self) -> Token:
+        with TOKEN_REFRESH_LOCK:
+            if not self._token.is_going_to_be_expired():
+                return self._token
+
+            logger.info("Refreshing access token")
+            with self._cred_file:
+                new_cred_file_content = self._cred_file.read()
+                new_token = new_cred_file_content.to_token()
+                new_host = new_cred_file_content.host
+
+                if new_cred_file_content == self._last_cred_file_content:
+                    self._token = self._auth_service.refresh_token(
+                        self._token, self.base_url
+                    )
+                    self._last_cred_file_content = CredentialsFileContent(
+                        host=self._host,
+                        access_token=self._token.access_token,
+                        refresh_token=self._token.refresh_token,
+                    )
+                    self._cred_file.write(self._last_cred_file_content)
+                    return self._token
+
+                if (
+                    new_host == self._host
+                    and new_token.to_user_info() == self._token.to_user_info()
+                ):
+                    self._last_cred_file_content = new_cred_file_content
+                    self._token = new_token
+                    # recursive
+                    return self.token
+
+                raise Exception(
+                    "Credentials on disk changed while mlfoundry was running."
+                )
+
+    @property
+    def base_url(self) -> str:
+        return self._host

truefoundry/deploy/lib/auth/servicefoundry_session.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from truefoundry.deploy.lib.auth.credential_provider import (
+    CredentialProvider,
+    EnvCredentialProvider,
+    FileCredentialProvider,
+)
+from truefoundry.deploy.lib.model.entity import UserInfo
+from truefoundry.logger import logger
+
+ACTIVE_SESSION: Optional[ServiceFoundrySession] = None
+
+
+class ServiceFoundrySession:
+    def __init__(self) -> None:
+        self._cred_provider = self._get_cred_provider()
+        self._user_info: UserInfo = self._cred_provider.token.to_user_info()
+
+        global ACTIVE_SESSION
+        if (ACTIVE_SESSION is None) or (
+            ACTIVE_SESSION
+            and ACTIVE_SESSION.base_url != self.base_url
+            and ACTIVE_SESSION.user_info != self.user_info
+        ):
+            logger.info(
+                "Logged in to %r as %r (%s)",
+                self.base_url,
+                self.user_info.user_id,
+                self.user_info.email or self.user_info.user_type.value,
+            )
+        ACTIVE_SESSION = self
+
+    @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(
+                "Please login again using `tfy login --relogin`"
+                "or `tfy.login(relogin=True)` function"
+            )
+        return final_cred_provider
+
+    @property
+    def access_token(self) -> str:
+        return self._cred_provider.token.access_token
+
+    @property
+    def base_url(self) -> str:
+        return self._cred_provider.base_url
+
+    @property
+    def user_info(self) -> UserInfo:
+        return self._user_info