truefoundry 0.4.0rc1__py3-none-any.whl → 0.4.0rc3__py3-none-any.whl

truefoundry/ml/login.py

@@ -1,241 +0,0 @@
-import os
-import threading
-from functools import lru_cache, wraps
-from pathlib import Path
-from typing import Optional
-
-import click
-from filelock import FileLock, Timeout
-
-from truefoundry.ml.env_vars import API_KEY_GLOBAL, TRACKING_HOST_GLOBAL
-from truefoundry.ml.exceptions import MlFoundryException
-from truefoundry.ml.logger import logger
-from truefoundry.ml.run_utils import resolve_tracking_uri
-from truefoundry.ml.services.auth_service import get_auth_service
-from truefoundry.ml.services.entities import Token
-from truefoundry.pydantic_v1 import BaseModel, Field, NonEmptyStr
-
-CREDENTIALS_DIR = Path.home() / ".truefoundry"
-CREDENTIALS_FILE = CREDENTIALS_DIR / "credentials.json"
-
-
-OLD_CREDENTIALS_DIR = Path.home() / ".mlfoundry"
-OLD_CREDENTIALS_FILE = OLD_CREDENTIALS_DIR / "credentials.netrc"
-
-if OLD_CREDENTIALS_FILE.exists():
-    logger.warning(
-        "%s file is deprecated. You can delete this file now.", OLD_CREDENTIALS_FILE
-    )
-
-
-class CredentialsFileContent(BaseModel):
-    access_token: NonEmptyStr = Field(repr=False)
-    refresh_token: Optional[NonEmptyStr] = Field(repr=False)
-    host: NonEmptyStr
-
-    class Config:
-        allow_mutation = False
-
-    def to_token(self) -> Token:
-        return Token(
-            access_token=self.access_token,  # type: ignore[call-arg]
-            refresh_token=self.refresh_token,
-        )
-
-
-def _ensure_lock_taken(method):
-    @wraps(method)
-    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"
-            )
-        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 = CREDENTIALS_FILE,
-        lock_timeout: float = 60.0,
-    ):
-        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 MlFoundryException(
-                "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 MlFoundryException(
-                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):
-        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 MlFoundryException(
-                "Error while reading the credentials file "
-                f"{self._credentials_file_path}. Please login again "
-                "using `tfy login` command or `truefoundry.login()` function"
-            ) from ex
-
-    @_ensure_lock_taken
-    def write(self, credentials_file_content: CredentialsFileContent):
-        if not isinstance(credentials_file_content, CredentialsFileContent):
-            raise MlFoundryException(
-                "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 exists(self) -> bool:
-        return self._credentials_file_path.exists()
-
-
-def login(
-    tracking_uri: Optional[str] = None,
-    relogin: bool = False,
-    api_key: Optional[str] = None,
-) -> bool:
-    """Save API key in local file system for a given `tracking_uri`.
-
-    Args:
-        tracking_uri (Optional[str], optional): tracking_uri for the given API key
-        relogin (bool, optional): Overwrites the existing API key for the `tracking_uri` if
-            set to `True`. If set to `False` and an API key is already present for
-            the given `tracking_uri`, then the existing API key is kept untouched.
-            Default is `False`.
-        api_key (Optional[str], optional): The API key for the given `tracking_uri`.
-            If `api_key` is not passed, this function prompts for the API key.
-
-    Returns:
-        bool: Returns `True` if any credential was persisted.
-    """
-    from truefoundry.ml.session import EnvCredentialProvider
-
-    if API_KEY_GLOBAL in os.environ and TRACKING_HOST_GLOBAL in os.environ:
-        logger.warning(
-            "Skipping login because environment variables %s and "
-            "%s are set and will be used when running truefoundry. "
-            "If you want to relogin then unset these environment keys.",
-            TRACKING_HOST_GLOBAL,
-            API_KEY_GLOBAL,
-        )
-        return False
-
-    if EnvCredentialProvider.can_provide():
-        logger.warning(
-            "TFY_API_KEY env var is already set. "
-            "When running mlfoundry, it will use the api key to authorize.\n"
-            "Login will just save the credentials on disk."
-        )
-
-    tracking_uri = resolve_tracking_uri(tracking_uri).strip("/")
-    auth_service = get_auth_service(tracking_uri=tracking_uri)
-
-    cred_file = CredentialsFileManager()
-
-    with cred_file:
-        if not relogin and cred_file.exists():
-            cred_file_content = cred_file.read()
-            if tracking_uri != cred_file_content.host:
-                if click.confirm(
-                    f"Already logged in to {cred_file_content.host!r}\n"
-                    f"Do you want to relogin to {tracking_uri!r}?"
-                ):
-                    return login(
-                        tracking_uri=tracking_uri, relogin=True, api_key=api_key
-                    )
-            user_info = cred_file_content.to_token().to_user_info()
-            user_name_display_info = user_info.email or user_info.user_type.value
-            print(
-                f"Already logged in to {cred_file_content.host!r} as "
-                f"{user_info.user_id!r} ({user_name_display_info!r})\n"
-                "Please use `tfy login --relogin` or `truefoundry.login(relogin=True)` "
-                "to force relogin"
-            )
-            return False
-
-        if api_key:
-            logger.debug("Logging in with api key")
-            token = Token(access_token=api_key, refresh_token=None)  # type: ignore[call-arg]
-            cred_file_content = CredentialsFileContent(
-                access_token=token.access_token,
-                refresh_token=token.refresh_token,
-                host=tracking_uri,
-            )
-            cred_file.write(cred_file_content)
-        else:
-            device_code = auth_service.get_device_code()
-            url_to_go = device_code.get_user_clickable_url(tracking_uri=tracking_uri)
-            print(f"Opening:- {url_to_go}")
-            print(
-                "Please click on the above link if it is not "
-                "automatically opened in a browser window."
-            )
-            click.launch(url_to_go)
-            token = auth_service.get_token_from_device_code(
-                device_code=device_code.device_code, timeout=120
-            )
-            cred_file_content = CredentialsFileContent(
-                access_token=token.access_token,
-                refresh_token=token.refresh_token,
-                host=tracking_uri,
-            )
-            cred_file.write(cred_file_content)
-
-    user_info = token.to_user_info()
-    user_name_display_info = user_info.email or user_info.user_type.value
-    print(
-        f"Logged in to {cred_file_content.host!r} as {user_info.user_id!r} ({user_name_display_info})"
-    )
-    return True

truefoundry/ml/services/auth_service.py

@@ -1,109 +0,0 @@
-import time
-from urllib.parse import urlparse
-
-from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
-    ApiClient,
-    AuthApi,
-    Configuration,
-)
-from truefoundry.ml.exceptions import MlFoundryException
-from truefoundry.ml.logger import logger
-from truefoundry.ml.run_utils import append_path_to_rest_tracking_uri
-from truefoundry.ml.services.entities import (
-    AuthServerInfo,
-    DeviceCode,
-    HostCreds,
-    Token,
-)
-from truefoundry.ml.services.utils import http_request, http_request_safe
-
-# TODO: This will eventually go away, this is duplicate of AuthServiceClient
-
-
-class AuthService:
-    def __init__(self, url: str, tenant_name: str):
-        self._host_creds = HostCreds(host=url.rstrip("/"), token=None)
-        self._tenant_name = tenant_name
-
-    def refresh_token(self, token: Token) -> Token:
-        if not token.refresh_token:
-            # TODO: Add a way to propagate error messages without traceback to the output interface side
-            raise MlFoundryException(
-                "Unable to resume login session. Please log in again using `tfy login [--host HOST] --relogin`"
-            )
-        try:
-            response = http_request_safe(
-                method="post",
-                host_creds=self._host_creds,
-                endpoint="api/v1/oauth/token/refresh",
-                json={
-                    "tenantName": token.tenant_name,
-                    "refreshToken": token.refresh_token,
-                },
-                timeout=3,
-            )
-        except MlFoundryException as e:
-            if e.status_code and (400 <= e.status_code < 500):
-                raise MlFoundryException(
-                    "Unable to resume login session. "
-                    "Please log in again using `tfy login [--host HOST] --relogin`"
-                ) from None
-            raise
-        return Token.parse_obj(response)
-
-    def get_device_code(self) -> DeviceCode:
-        response = http_request_safe(
-            method="post",
-            host_creds=self._host_creds,
-            endpoint="api/v1/oauth/device",
-            json={"tenantName": self._tenant_name},
-            timeout=3,
-        )
-        return DeviceCode.parse_obj(response)
-
-    def get_token_from_device_code(
-        self, device_code: str, timeout: float = 60
-    ) -> Token:
-        start_time = time.monotonic()
-        while (time.monotonic() - start_time) <= timeout:
-            response = http_request(
-                method="post",
-                host_creds=self._host_creds,
-                endpoint="api/v1/oauth/device/token",
-                json={"tenantName": self._tenant_name, "deviceCode": device_code},
-                timeout=3,
-            )
-            if response.status_code == 202:
-                logger.debug("User has not authorized yet. Checking again.")
-                time.sleep(1.0)
-                continue
-            if response.status_code == 201:
-                response = response.json()
-                return Token.parse_obj(response)
-            raise MlFoundryException(
-                "Failed to get token using device code.\n"
-                f"Status Code: {response.status_code},\nResponse: {response.text}"
-            )
-        raise MlFoundryException(f"Did not get authorized within {timeout} seconds.")
-
-
-def get_auth_service(tracking_uri: str) -> AuthService:
-    tracking_uri = append_path_to_rest_tracking_uri(tracking_uri)
-    parsed_tracking_uri = urlparse(tracking_uri)
-    host = parsed_tracking_uri.netloc
-    # Anonymous api
-    api_client = ApiClient(
-        configuration=Configuration(
-            host=tracking_uri.rstrip("/"),
-            access_token=None,
-        )
-    )
-    auth_api = AuthApi(api_client=api_client)
-    auth_server_info = auth_api.get_tenant_id_get(
-        host_name=host,
-        _request_timeout=3,
-    )
-    tenant_info = AuthServerInfo.parse_obj(auth_server_info.dict())
-    return AuthService(
-        url=tenant_info.auth_server_url, tenant_name=tenant_info.tenant_name
-    )