truefoundry 0.3.0rc10__py3-none-any.whl → 0.3.0rc12__py3-none-any.whl

truefoundry/deploy/lib/auth/auth_service_client.py

@@ -1,4 +1,5 @@
 import time
+from abc import ABC, abstractmethod
 
 import requests
 
@@ -9,17 +10,114 @@ from truefoundry.deploy.lib.model.entity import DeviceCode, Token
 from truefoundry.logger import logger
 
 
-class AuthServiceClient:
+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)
-        tenant_info = client.get_tenant_info()
 
-        self._auth_server_url = tenant_info.auth_server_url
-        self._tenant_name = tenant_info.tenant_name
+        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"
@@ -33,7 +131,7 @@ class AuthServiceClient:
             "tenantName": token.tenant_name,
             "refreshToken": token.refresh_token,
         }
-        res = requests.post(url, data=data)
+        res = requests.post(url, json=data)
         try:
             res = request_handling(res)
             return Token.parse_obj(res)
@@ -45,12 +143,14 @@ class 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, data=data)
+        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
+        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 = {
@@ -61,7 +161,7 @@ class AuthServiceClient:
         poll_interval_seconds = 1
 
         for response in poll_for_function(
-            requests.post, poll_after_secs=poll_interval_seconds, url=url, data=data
+            requests.post, poll_after_secs=poll_interval_seconds, url=url, json=data
         ):
             if response.status_code == 201:
                 response = response.json()

truefoundry/deploy/lib/auth/credential_provider.py

@@ -44,7 +44,7 @@ class EnvCredentialProvider(CredentialProvider):
         # TODO: Read host from cred file as well.
         base_url = resolve_base_url().strip("/")
         self._host = base_url
-        self._auth_service = AuthServiceClient(base_url=base_url)
+        self._auth_service = AuthServiceClient.from_base_url(base_url=base_url)
 
         servicefoundry_client = ServiceFoundryServiceClient(
             init_session=False, base_url=base_url
@@ -82,7 +82,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(base_url=self._host)
+        self._auth_service = AuthServiceClient.from_base_url(base_url=self._host)
 
     @staticmethod
     def can_provide() -> bool:

truefoundry/deploy/lib/clients/servicefoundry_client.py

@@ -26,12 +26,14 @@ from truefoundry.deploy.lib.model.entity import (
     Deployment,
     DockerRegistryCredentials,
     JobRun,
+    PythonSDKConfig,
     TenantInfo,
     Token,
     TriggerJobResult,
     Workspace,
     WorkspaceResources,
 )
+from truefoundry.deploy.lib.util import timed_lru_cache
 from truefoundry.deploy.lib.win32 import allow_interrupt
 from truefoundry.deploy.v2.lib.models import (
     AppDeploymentStatusResponse,
@@ -93,6 +95,24 @@ def check_min_cli_version(fn):
     return inner
 
 
+@timed_lru_cache(seconds=30 * 60)
+def _cached_get_tenant_info(api_server_url: str) -> TenantInfo:
+    res = requests.get(
+        url=f"{api_server_url}/v1/tenant-id",
+        params={"hostName": urlparse(api_server_url).netloc},
+    )
+    res = request_handling(res)
+    return TenantInfo.parse_obj(res)
+
+
+@timed_lru_cache(seconds=30 * 60)
+def _cached_get_python_sdk_config(api_server_url: str) -> PythonSDKConfig:
+    url = f"{api_server_url}/v1/min-cli-version"
+    res = requests.get(url)
+    res = request_handling(res)
+    return PythonSDKConfig.parse_obj(res)
+
+
 class ServiceFoundryServiceClient:
     def __init__(self, init_session: bool = True, base_url: Optional[str] = None):
         self._session: Optional[ServiceFoundrySession] = None
@@ -111,20 +131,19 @@ class ServiceFoundryServiceClient:
     def base_url(self) -> str:
         return self._base_url
 
-    def get_tenant_info(self) -> TenantInfo:
-        res = requests.get(
-            url=f"{self._api_server_url}/v1/tenant-id",
-            params={"hostName": urlparse(self._api_server_url).netloc},
-        )
-        res = request_handling(res)
-        return TenantInfo.parse_obj(res)
+    @property
+    def tenant_info(self) -> TenantInfo:
+        return _cached_get_tenant_info(self._api_server_url)
+
+    @property
+    def python_sdk_config(self) -> PythonSDKConfig:
+        return _cached_get_python_sdk_config(self._api_server_url)
 
     @functools.cached_property
     def _min_cli_version_required(self) -> str:
-        url = f"{self._api_server_url}/v1/min-cli-version"
-        res = requests.get(url)
-        res = request_handling(res)
-        return res["truefoundryCliMinVersion"]
+        return _cached_get_python_sdk_config(
+            self._api_server_url
+        ).truefoundry_cli_min_version
 
     def _get_header(self):
         if not self._session:

truefoundry/deploy/lib/model/entity.py

@@ -297,6 +297,11 @@ class CredentialsFileContent(BaseModel):
 class DeviceCode(BaseModel):
     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")
 
     class Config:
         allow_population_by_field_name = True
@@ -306,6 +311,18 @@ class DeviceCode(BaseModel):
         return f"{auth_host}/authorize/device?userCode={self.user_code}"
 
 
+class PythonSDKConfig(BaseModel):
+    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 Config:
+        allow_population_by_field_name = True
+        allow_mutation = False
+
+
 class JobRun(Base):
     name: str
     applicationName: str

truefoundry/deploy/lib/session.py

@@ -92,7 +92,7 @@ def login(
                 api_key=api_key, servicefoundry_client=servicefoundry_client
             )
         else:
-            auth_service = AuthServiceClient(base_url=host)
+            auth_service = AuthServiceClient.from_base_url(base_url=host)
             # interactive login
             token = _login_with_device_code(base_url=host, auth_service=auth_service)
 
@@ -136,11 +136,23 @@ def _login_with_device_code(
 ) -> Token:
     logger.debug("Logging in with device code")
     device_code = auth_service.get_device_code()
-    url_to_go = device_code.get_user_clickable_url(auth_host=base_url)
-    output_hook.print_line(f"Opening:- {url_to_go}")
-    output_hook.print(
-        "Please click on the above link if it is not "
-        "automatically opened in a browser window."
+    auto_open_url = None
+    message = "Please click on the above link if it is not automatically opened in a browser window."
+    if device_code.complete_verification_url:
+        auto_open_url = device_code.complete_verification_url
+    elif device_code.verification_url:
+        if device_code.message:
+            message = device_code.message
+        else:
+            message = f"Please open the following URL in a browser and enter the code {device_code.user_code} when prompted: {device_code.verification_url}"
+    else:
+        auto_open_url = device_code.get_user_clickable_url(auth_host=base_url)
+    output_hook.print_line(message)
+    if auto_open_url:
+        output_hook.print_line(f"Opening:- {auto_open_url}")
+        click.launch(auto_open_url)
+    return auth_service.get_token_from_device_code(
+        device_code=device_code.device_code,
+        timeout=device_code.expires_in_seconds,
+        poll_interval_seconds=device_code.interval_in_seconds,
     )
-    click.launch(url_to_go)
-    return auth_service.get_token_from_device_code(device_code=device_code.device_code)

truefoundry/deploy/lib/util.py

@@ -1,5 +1,7 @@
 import os
 import re
+from functools import lru_cache, wraps
+from time import monotonic_ns
 from typing import Union
 
 from truefoundry.deploy.lib.const import (
@@ -68,3 +70,21 @@ def find_list_paths(data, parent_key="", sep="."):
             new_key = f"{parent_key}[{i}]"
             list_paths.extend(find_list_paths(value, new_key, sep))
     return list_paths
+
+
+def timed_lru_cache(seconds: int = 300, maxsize: int = None):
+    def wrapper_cache(func):
+        func = lru_cache(maxsize=maxsize)(func)
+        func.delta = seconds * 10**9
+        func.expiration = monotonic_ns() + func.delta
+
+        @wraps(func)
+        def wrapped_func(*args, **kwargs):
+            if monotonic_ns() >= func.expiration:
+                func.cache_clear()
+                func.expiration = monotonic_ns() + func.delta
+            return func(*args, **kwargs)
+
+        return wrapped_func
+
+    return wrapper_cache

{truefoundry-0.3.0rc10.dist-info → truefoundry-0.3.0rc12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.3.0rc10
+Version: 0.3.0rc12
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com
@@ -24,7 +24,7 @@ Requires-Dist: flytekit (==1.12.2) ; extra == "workflow"
 Requires-Dist: gitignorefile (>=1.1.2,<1.2.0)
 Requires-Dist: importlib-metadata (>=6.0.1,<8.0.0)
 Requires-Dist: importlib-resources (>=5.2.0,<6.0.0)
-Requires-Dist: mlfoundry (==0.11.3) ; extra == "ml"
+Requires-Dist: mlfoundry (==0.11.4) ; extra == "ml"
 Requires-Dist: openai (>=1.16.2,<2.0.0)
 Requires-Dist: packaging (>=20.0,<25.0)
 Requires-Dist: pydantic (>=1.10.0,<3)

{truefoundry-0.3.0rc10.dist-info → truefoundry-0.3.0rc12.dist-info}/RECORD

@@ -78,12 +78,12 @@ truefoundry/deploy/io/output_callback.py,sha256=V2YwUFec4G4a67lM4r-x_64AqdOVNo_9
 truefoundry/deploy/io/rich_output_callback.py,sha256=TJLiRD-EnFVwgcepxR7WN0koKqW1X2DevETPhNPi_nU,829
 truefoundry/deploy/json_util.py,sha256=x_-7YYQ4_HUIJ8ofOcclAp9JWhgTWjR9Th6Q0FuRqGk,175
 truefoundry/deploy/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/deploy/lib/auth/auth_service_client.py,sha256=rrz2t6lkNgVBoWsIrdyaK9OHN5XZJtA4AAqaDx3PQyc,3299
+truefoundry/deploy/lib/auth/auth_service_client.py,sha256=C3sgoDZ2RUijJY_nStqt-BmSwUAdjVa8yAhazOGz710,7386
 truefoundry/deploy/lib/auth/credential_file_manager.py,sha256=DXeXWoVakfZI2Geu8Futwef_eVMHrZ1et7d9MIXOhhY,4233
-truefoundry/deploy/lib/auth/credential_provider.py,sha256=MwTN8170TXi7g9m2Fw4VRReZqk1DzBFG1bMVgd7jYZk,4375
+truefoundry/deploy/lib/auth/credential_provider.py,sha256=eXgfA2-q20XxslOmlocmi-rKQ5AtBxVMGuftfrLVlb4,4403
 truefoundry/deploy/lib/auth/servicefoundry_session.py,sha256=2OahwRg-8l3QUwpC2iLA8lHavkeTxKXUP_0AN27HPS8,1859
 truefoundry/deploy/lib/clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/deploy/lib/clients/servicefoundry_client.py,sha256=NASbj67lrQKZwhu7mVu6ngE9A8C76FqqUpfFYQKjVAU,26394
+truefoundry/deploy/lib/clients/servicefoundry_client.py,sha256=yqeaF8nFeqR01q9d7E2s7sDQVpgF1ES9CECXpKNMT5g,26945
 truefoundry/deploy/lib/clients/shell_client.py,sha256=tMrc0Ha1DmGtUCJrZD8eusOzfe8R_WIe6AAH7nxL0xA,461
 truefoundry/deploy/lib/clients/utils.py,sha256=rK7DrvA71kSTjy23-Bk4LTQjgViBWeHtcV_SlBLZw6M,1282
 truefoundry/deploy/lib/const.py,sha256=Yk_nXeZWzwKs-6hEXDjVDyjwEGh35T5TWaBxyJgP-Zw,1395
@@ -96,9 +96,9 @@ truefoundry/deploy/lib/exceptions.py,sha256=ZW44bSwmlcr3GcB10E8Jz-zlgu9wyi8tE7xA
 truefoundry/deploy/lib/logs_utils.py,sha256=SQxRv3jDDmgHdOUMhlMaAPGYskybnBUMpst7QU_i_sc,1469
 truefoundry/deploy/lib/messages.py,sha256=nhp0bCYf_XpUM68hTq5lBY-__vtEyV2uP7NgnJXJ_Vg,925
 truefoundry/deploy/lib/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-truefoundry/deploy/lib/model/entity.py,sha256=gFlSG6RVuaWQQ2v3YtONhnTbIRzKUIryMgcJN6HmZV4,11309
-truefoundry/deploy/lib/session.py,sha256=0S9foO78BXYSwCh6nPaWG-QuwgxaK4VZqzshjNkyXls,5229
-truefoundry/deploy/lib/util.py,sha256=3Pc0gQgomapTSjMQo743STShBn1B6PNdBpw3bTpdxOw,2200
+truefoundry/deploy/lib/model/entity.py,sha256=aNiQ7kK_V1kH7C_VC6HP7ibas_9Ohkoqzs-WdTfrrf4,12037
+truefoundry/deploy/lib/session.py,sha256=WumH3cfVr8DqcuY5ngdnPBI_oeMUxG_JqyONQYKlM6k,5844
+truefoundry/deploy/lib/util.py,sha256=sNMv9peDU7RTCDaAfyWPpPpCs_WcZGqOvc13Jqx9y6w,2809
 truefoundry/deploy/lib/win32.py,sha256=1RcvPTdlOAJ48rt8rCbE2Ufha2ztRqBAE9dueNXArrY,5009
 truefoundry/deploy/v2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 truefoundry/deploy/v2/lib/__init__.py,sha256=WEiVMZXOVljzEE3tpGJil14liIn_PCDoACJ6b3tZ6sI,188
@@ -130,7 +130,7 @@ truefoundry/workflow/map_task.py,sha256=2m3qGXQ90k9LdS45q8dqCCECc3qr8t2m_LMCVd1m
 truefoundry/workflow/python_task.py,sha256=SRXRLC4vdBqGjhkwuaY39LEWN6iPCpJAuW17URRdWTY,1128
 truefoundry/workflow/task.py,sha256=ToitYiKcNzFCtOVQwz1W8sRjbR97eVS7vQBdbgUQtKg,1779
 truefoundry/workflow/workflow.py,sha256=WaTqUjhwfAXDWu4E5ehuwAxrCbDJkoAf1oWmR2E9Qy0,4575
-truefoundry-0.3.0rc10.dist-info/METADATA,sha256=ycIiBcPg1h3rNv8TSK3OHc4PHWDH6DGVP9f0MC5acSo,2694
-truefoundry-0.3.0rc10.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-truefoundry-0.3.0rc10.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
-truefoundry-0.3.0rc10.dist-info/RECORD,,
+truefoundry-0.3.0rc12.dist-info/METADATA,sha256=UHVlgtge2QopJNqTeUfoy51XbmFNbTrCKSMC5HUieug,2694
+truefoundry-0.3.0rc12.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+truefoundry-0.3.0rc12.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
+truefoundry-0.3.0rc12.dist-info/RECORD,,