truefoundry 0.4.4rc6__py3-none-any.whl → 0.4.4rc8__py3-none-any.whl

truefoundry/common/tfy_signed_url_client.py

@@ -3,12 +3,17 @@ from enum import Enum
 from typing import Any, Dict, List, Optional
 from urllib.parse import urlencode, urljoin
 
-import requests
 from requests.exceptions import RequestException
 
 from truefoundry.common.constants import ENV_VARS
+from truefoundry.common.request_utils import requests_retry_session
+from truefoundry.common.utils import log_time
 from truefoundry.pydantic_v1 import BaseModel, Field
 
+# TODO: Move print statements to logger
+# TODO: Move these constants to the constants module
+REQUEST_TIMEOUT = 3600
+
 DEFAULT_TTL = ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL
 MAX_TIMEOUT = ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT
 
@@ -69,12 +74,18 @@ class SignedURLClient:
         self.base_url = base_url or ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_HOST
         self.token = token or ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN
         self.ttl = ttl
+        self.max_retries: int = 3
+        self.retry_backoff_factor: float = 0.3
+        self.session = requests_retry_session(
+            retries=self.max_retries, backoff_factor=self.retry_backoff_factor
+        )
 
         self.headers = {
             "Authorization": f"Bearer {self.token}",
             "Content-Type": "application/json",
         }
 
+    @log_time
     def _make_request(
         self, endpoint: str, method: str = "GET", payload: Optional[Dict] = None
     ) -> Dict:
@@ -94,7 +105,7 @@ class SignedURLClient:
         """
         url = urljoin(self.base_url, endpoint)
         try:
-            response = requests.request(
+            response = self.session.request(
                 method, url, headers=self.headers, json=payload, timeout=MAX_TIMEOUT
             )
             response.raise_for_status()
@@ -110,14 +121,15 @@ class SignedURLClient:
         endpoint_with_params = f"{endpoint.value}?{query_string}"
         return self._make_request(endpoint_with_params)
 
+    @log_time
     def _upload_data(self, signed_url: str, data: Any) -> None:
         """Helper method to upload data to the signed URL."""
         try:
-            # TODO: Add TIMEOUT to this request
-            response = requests.put(
+            response = self.session.put(
                 signed_url,
                 headers={"Content-Type": "application/octet-stream"},
                 data=data,
+                timeout=REQUEST_TIMEOUT,
             )
             response.raise_for_status()
         except Exception as e:
@@ -133,6 +145,7 @@ class SignedURLClient:
         self._upload_data(pre_signed_object_dto.signed_url, data)
         return storage_uri
 
+    @log_time
     def upload(self, file_path: str, storage_uri: str) -> str:
         """Upload a file to the specified storage path using a signed URL."""
         print(f"Uploading {file_path} to {storage_uri}")
@@ -145,6 +158,7 @@ class SignedURLClient:
             self._upload_data(pre_signed_object_dto.signed_url, file)
         return storage_uri
 
+    @log_time
     def _download_file(
         self, signed_url: str, local_path: Optional[str] = None
     ) -> Optional[bytes]:
@@ -158,25 +172,21 @@ class SignedURLClient:
             Optional[bytes]: The content of the downloaded file if local_path is None; otherwise, None.
         """
         try:
-            # TODO: Add TIMEOUT to this request
-            response = requests.get(
-                signed_url,
-                stream=True,
-                headers=self.headers,
+            response = self.session.get(
+                signed_url, stream=True, headers=self.headers, timeout=REQUEST_TIMEOUT
             )
             response.raise_for_status()
-
             if local_path:
                 with open(local_path, "wb") as file:
                     for chunk in response.iter_content(chunk_size=8192):
                         file.write(chunk)
                 return None  # File saved successfully; no need to return content.
-
             return response.content  # Return the content if not saving to a local path.
 
-        except requests.RequestException as e:
+        except RequestException as e:
             raise RuntimeError(f"Failed to download file from {signed_url}: {e}") from e
 
+    @log_time
     def download(self, storage_uri: str, local_path: str) -> Optional[str]:
         """Download a file from the specified storage path to a local path using a signed URL.
 
@@ -194,7 +204,7 @@ class SignedURLClient:
         )
         presigned_object = SignedURLAPIResponseDto.parse_obj(response)
         self._download_file(presigned_object.signed_url, local_path)
-        return local_path  # Indicate successful download.
+        return local_path
 
     def download_to_bytes(self, storage_uri: str) -> bytes:
         """Download a file from the specified storage path and return it as bytes.
@@ -229,6 +239,7 @@ class SignedURLClient:
         print(f"Path {uri} is_directory: {response_obj.is_directory}")
         return response_obj.is_directory
 
+    @log_time
     def list_files(self, path: str, detail: bool = False, max_results: int = 1000):
         """List files in the specified directory."""
         token = ""

truefoundry/common/tfy_signed_url_fs.py

@@ -9,19 +9,24 @@ from fsspec.spec import DEFAULT_CALLBACK, AbstractBufferedFile, AbstractFileSyst
 
 from truefoundry.common.constants import ENV_VARS
 from truefoundry.common.tfy_signed_url_client import SignedURLClient
+from truefoundry.common.utils import log_time
 
 
 class SignedURLFileSystem(AbstractFileSystem):
-    def __init__(self, base_url: Optional[str] = None, token: Optional[str] = None):
+    def __init__(
+        self, base_url: Optional[str] = None, token: Optional[str] = None, **kwargs
+    ):
         super().__init__()
         base_url = base_url or ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_HOST
         token = token or ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN
         self.client = SignedURLClient(base_url, token)
 
+    @log_time
     def exists(self, path, **kwargs):
         """Is there a file at the given path"""
         return self.client.exists(path)
 
+    @log_time
     def get(
         self,
         rpath,
@@ -78,6 +83,7 @@ class SignedURLFileSystem(AbstractFileSystem):
                 target_local_path = target_local_path + rpath.rsplit("/")[-1]
             self.client.download(storage_uri=rpath, local_path=target_local_path)
 
+    @log_time
     def put(
         self,
         lpath,
@@ -123,6 +129,7 @@ class SignedURLFileSystem(AbstractFileSystem):
             # Handle single file upload
             return self.client.upload(file_path=lpath, storage_uri=rpath)
 
+    @log_time
     def isdir(self, path):
         """Is this entry directory-like?"""
         return self.client.is_directory(path)
@@ -173,6 +180,7 @@ class SignedURLFileSystem(AbstractFileSystem):
         # Upload the content to the remote file system
         self.client.upload_from_bytes(content, storage_uri=path)
 
+    @log_time
     def ls(self, path, detail=True, **kwargs):
         """List objects at path."""
         return self.client.list_files(path, detail=detail)

truefoundry/common/utils.py

@@ -102,3 +102,17 @@ def resolve_tfy_host(tfy_host: Optional[str] = None) -> str:
     tfy_host = tfy_host.strip("/")
     validate_tfy_host(tfy_host)
     return tfy_host
+
+
+def log_time(func):
+    """Decorator to log the time taken by I/O operations."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        elapsed_time = time.time() - start_time
+        print(f"{func.__name__} took {elapsed_time:.2f} seconds")
+        return result
+
+    return wrapper

{truefoundry-0.4.4rc6.dist-info → truefoundry-0.4.4rc8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.4.4rc6
+Version: 0.4.4rc8
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com
@@ -22,7 +22,6 @@ Requires-Dist: docker (>=6.1.2,<8.0.0)
 Requires-Dist: fastapi (>=0.56.0,<0.200.0)
 Requires-Dist: filelock (>=3.8.0,<4.0.0)
 Requires-Dist: flytekit (==1.12.2) ; extra == "workflow"
-Requires-Dist: fsspec (>=2024.9.0,<2025.0.0)
 Requires-Dist: gitignorefile (>=1.1.2,<2.0.0)
 Requires-Dist: importlib-metadata (>=4.11.3,<9.0.0)
 Requires-Dist: importlib-resources (>=5.2.0,<7.0.0)

{truefoundry-0.4.4rc6.dist-info → truefoundry-0.4.4rc8.dist-info}/RECORD

@@ -34,9 +34,9 @@ truefoundry/common/entities.py,sha256=8O-EGPk4PKqnyoFMKUTxISCU19rz0KBnfRDJU695Dh
 truefoundry/common/exceptions.py,sha256=ePpiQ_zmWe4e94gOgeMiyP_AZnKwjEBfyXsB5ScGYcI,329
 truefoundry/common/request_utils.py,sha256=5xw4YGUcMf71Ncal3OfFCa-PoWDIvG3hYGCDa4Da4OI,2854
 truefoundry/common/servicefoundry_client.py,sha256=2fxmgCM-ckFHpnm6n_mL-5Z8RWN_q-dYVvFC29bkYSg,3120
-truefoundry/common/tfy_signed_url_client.py,sha256=zqGIG7L6roPdSSdZ5IVTIVjYsOJsWvMa8SMkYttkg70,9401
-truefoundry/common/tfy_signed_url_fs.py,sha256=14RJWXeXy2ADqHTfVvBtb7qtQwULxOMF8a1Jrgmgb04,7989
-truefoundry/common/utils.py,sha256=MYFjNtHGqauqhj9tmbdErCJR49AfXDwg-5kYbBh8HpI,3258
+truefoundry/common/tfy_signed_url_client.py,sha256=wD8R88wxECvHsHo_BJSeOODLylmJjksZt8iF9pUAd4w,9822
+truefoundry/common/tfy_signed_url_fs.py,sha256=sMrpRnlmVin9q8kfv5iVJZGD6aZUgYKlCXxjDx8lKm4,8129
+truefoundry/common/utils.py,sha256=vacYbVTEvcEY4niH4C9Wb_aCeaGgh2un3_ApAt556VQ,3621
 truefoundry/deploy/__init__.py,sha256=ugawKF2G02EmEXX35oZ2tec12d9oWN28Sf6mtGGIERY,2281
 truefoundry/deploy/auto_gen/models.py,sha256=4MaxkG2_5Wg6avaZRlK0D4JiVEM5rk3NU0BCiTx8VyU,82477
 truefoundry/deploy/builder/__init__.py,sha256=1qjHMNBE1poRCZW0WrG46dFM1f1IlivD5352qzsioMU,4953
@@ -342,7 +342,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.4.4rc6.dist-info/METADATA,sha256=YuV_hYNyXMHOBpjsQSnBnVgRt0bcj3fstht3SkKhPsQ,3146
-truefoundry-0.4.4rc6.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-truefoundry-0.4.4rc6.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
-truefoundry-0.4.4rc6.dist-info/RECORD,,
+truefoundry-0.4.4rc8.dist-info/METADATA,sha256=umX9ghRm2Gvbw8WS3JM_yFm__OVuvFtvvZB18oLE0a4,3101
+truefoundry-0.4.4rc8.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+truefoundry-0.4.4rc8.dist-info/entry_points.txt,sha256=TXvUxQkI6zmqJuycPsyxEIMr3oqfDjgrWj0m_9X12x4,95
+truefoundry-0.4.4rc8.dist-info/RECORD,,