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

truefoundry/common/tfy_signed_url_client.py

@@ -92,7 +92,7 @@ class SignedURLClient:
             retries=self.max_retries, backoff_factor=self.retry_backoff_factor
         )
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def _make_request(
         self,
         endpoint: str,
@@ -111,6 +111,7 @@ class SignedURLClient:
         except RequestException as e:
             raise RuntimeError(f"Error during request to {url}: {e}") from e
 
+    @log_time(prefix=LOG_PREFIX)
     def _make_server_api_call(
         self,
         endpoint: SignedURLServerEndpoint,
@@ -124,7 +125,7 @@ class SignedURLClient:
             endpoint=endpoint_with_params, method="GET", payload=None, headers=headers
         )
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def _upload_data(self, signed_url: str, data: Any) -> None:
         """
         Upload data to the specified storage path using a signed URL.
@@ -144,6 +145,7 @@ class SignedURLClient:
         except Exception as e:
             raise RuntimeError(f"Failed to upload data: {e}") from e
 
+    @log_time(prefix=LOG_PREFIX)
     def upload_from_bytes(self, data: bytes, storage_uri: str) -> str:
         """Upload bytes to the specified storage path using a signed URL."""
         signed_object = self._make_server_api_call(
@@ -155,7 +157,7 @@ class SignedURLClient:
         self._upload_data(pre_signed_object_dto.signed_url, data)
         return storage_uri
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def upload(self, file_path: str, storage_uri: str) -> str:
         """Upload a file to the specified storage path using a signed URL."""
         logger.info(f"{LOG_PREFIX} Uploading {file_path} to {storage_uri}")
@@ -169,7 +171,7 @@ class SignedURLClient:
             self._upload_data(pre_signed_object_dto.signed_url, file)
         return storage_uri
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def _download_file(
         self, signed_url: str, local_path: Optional[str] = None
     ) -> Optional[bytes]:
@@ -191,7 +193,7 @@ class SignedURLClient:
         except RequestException as e:
             raise RuntimeError(f"Failed to download file from {signed_url}: {e}") from e
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     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."""
         logger.info(f"{LOG_PREFIX} Downloading {storage_uri} to {local_path}")
@@ -204,6 +206,7 @@ class SignedURLClient:
         self._download_file(presigned_object.signed_url, local_path)
         return local_path
 
+    @log_time(prefix=LOG_PREFIX)
     def download_to_bytes(self, storage_uri: str) -> bytes:
         """Download a file from the specified storage path and return it as bytes."""
         response = self._make_server_api_call(
@@ -214,6 +217,7 @@ class SignedURLClient:
         presigned_object = SignedURLAPIResponseDto.parse_obj(response)
         return self._download_file(presigned_object.signed_url)
 
+    @log_time(prefix=LOG_PREFIX)
     def exists(self, uri: str) -> bool:
         """Check if a file exists at the specified path."""
         response = self._make_server_api_call(
@@ -223,6 +227,7 @@ class SignedURLClient:
         )
         return SignedURLExistsAPIResponseDto.parse_obj(response).exists
 
+    @log_time(prefix=LOG_PREFIX)
     def is_directory(self, uri: str) -> bool:
         """Check if the specified URI is a directory."""
         response = self._make_server_api_call(
@@ -236,7 +241,7 @@ class SignedURLClient:
         logger.info(f"{LOG_PREFIX} Path {uri} is a directory: {is_directory}")
         return is_directory
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def list_files(
         self, path: str, detail: bool = False, max_results: int = 1000
     ) -> Union[List[FileInfo], List[str]]:

truefoundry/common/tfy_signed_url_fs.py

@@ -8,7 +8,7 @@ from typing import Optional
 from fsspec.spec import DEFAULT_CALLBACK, AbstractBufferedFile, AbstractFileSystem
 
 from truefoundry.common.constants import ENV_VARS
-from truefoundry.common.tfy_signed_url_client import SignedURLClient
+from truefoundry.common.tfy_signed_url_client import LOG_PREFIX, SignedURLClient
 from truefoundry.common.utils import log_time
 
 
@@ -21,12 +21,12 @@ class SignedURLFileSystem(AbstractFileSystem):
         token = token or ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN
         self.client = SignedURLClient(base_url, token)
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def exists(self, path, **kwargs):
         """Check if a file exists at the given path."""
         return self.client.exists(path)
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def get(
         self,
         rpath,
@@ -87,7 +87,7 @@ class SignedURLFileSystem(AbstractFileSystem):
             Path(os.path.dirname(target_local_path)).mkdir(parents=True, exist_ok=True)
             self.client.download(storage_uri=rpath, local_path=target_local_path)
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def put(
         self,
         lpath,
@@ -134,7 +134,7 @@ class SignedURLFileSystem(AbstractFileSystem):
                 rpath = os.path.join(rpath, local_path.name)
             return self.client.upload(file_path=lpath, storage_uri=rpath)
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def isdir(self, path):
         """Is this entry directory-like?"""
         return self.client.is_directory(path)
@@ -170,6 +170,7 @@ class SignedURLFileSystem(AbstractFileSystem):
             # Wrapping the buffer to automatically upload on close
             return io.BufferedWriter(buffer, on_close)
 
+    @log_time(prefix=LOG_PREFIX)
     def write(self, path, data, **kwargs):
         """
         Write data to the file at the specified path (this could be tied to open's close).
@@ -185,7 +186,7 @@ class SignedURLFileSystem(AbstractFileSystem):
         # Upload the content to the remote file system
         self.client.upload_from_bytes(content, storage_uri=path)
 
-    @log_time
+    @log_time(prefix=LOG_PREFIX)
     def ls(self, path, detail=True, **kwargs):
         """List objects at path."""
         return self.client.list_files(path, detail=detail)

truefoundry/common/utils.py

@@ -10,6 +10,7 @@ from truefoundry.common.constants import (
     MLFOUNDRY_SERVER_RELATIVE_PATH,
     TFY_HOST_ENV_KEY,
 )
+from truefoundry.logger import logger
 from truefoundry.pydantic_v1 import BaseSettings
 
 T = TypeVar("T")
@@ -104,15 +105,18 @@ def resolve_tfy_host(tfy_host: Optional[str] = None) -> str:
     return tfy_host
 
 
-def log_time(func):
+def log_time(prefix: str = ""):
     """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
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            result = func(*args, **kwargs)
+            elapsed_time = time.time() - start_time
+            logger.info(f"{prefix}{func.__name__} took {elapsed_time:.2f} seconds")
+            return result
+
+        return wrapper
 
-    return wrapper
+    return decorator