truefoundry 0.4.4rc6__tar.gz → 0.4.4rc8__tar.gz

{truefoundry-0.4.4rc6 → truefoundry-0.4.4rc8}/PKG-INFO

@@ -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 → truefoundry-0.4.4rc8}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truefoundry"
-version = "0.4.4rc6"
+version = "0.4.4rc8"
 description = "Truefoundry CLI"
 authors = ["Abhishek Choudhary <abhishek@truefoundry.com>"]
 readme = "README.md"
@@ -54,8 +54,6 @@ pandas = [
     { version = ">=1.0.0,<3.0.0", python = "<3.10" },
     { version = ">=1.4.0,<3.0.0", python = ">=3.10" },
 ]
-fsspec = "^2024.9.0"
-
 
 [tool.poetry.extras]
 workflow = ["flytekit"]

{truefoundry-0.4.4rc6 → truefoundry-0.4.4rc8}/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-0.4.4rc6 → truefoundry-0.4.4rc8}/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-0.4.4rc6 → truefoundry-0.4.4rc8}/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