truefoundry 0.4.4rc8__tar.gz → 0.4.4rc10__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: truefoundry
-Version: 0.4.4rc8
+Version: 0.4.4rc10
 Summary: Truefoundry CLI
 Author: Abhishek Choudhary
 Author-email: abhishek@truefoundry.com

{truefoundry-0.4.4rc8 → truefoundry-0.4.4rc10}/pyproject.toml

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

{truefoundry-0.4.4rc8 → truefoundry-0.4.4rc10}/truefoundry/common/constants.py

@@ -13,12 +13,6 @@ TFY_API_KEY_ENV_KEY = "TFY_API_KEY"
 
 TFY_INTERNAL_SIGNED_URL_SERVER_HOST_ENV_KEY = "TFY_INTERNAL_SIGNED_URL_SERVER_HOST"
 TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN_ENV_KEY = "TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN"
-TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL_ENV_KEY = (
-    "TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL"
-)
-TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT_ENV_KEY = (
-    "TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT"
-)
 
 
 class TrueFoundrySdkEnv(BaseSettings):
@@ -26,18 +20,19 @@ class TrueFoundrySdkEnv(BaseSettings):
     # Never expect the user to set these values
     TFY_HOST: Optional[str] = Field(default=None, env=TFY_HOST_ENV_KEY)
     TFY_API_KEY: Optional[SecretStr] = Field(default=None, env=TFY_API_KEY_ENV_KEY)
+
+    # Internal Signed URL Server
     TFY_INTERNAL_SIGNED_URL_SERVER_HOST: Optional[str] = Field(
         default=None, env=TFY_INTERNAL_SIGNED_URL_SERVER_HOST_ENV_KEY
     )
     TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN: Optional[str] = Field(
         default=None, env=TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN_ENV_KEY
     )
-    TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL: int = Field(
-        default=3600, env=TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL_ENV_KEY
-    )  # 1 hour
-    TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT: int = Field(
-        default=5, env=TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT_ENV_KEY
-    )  # 5 seconds
+    TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT: int = 5  # default: 5 seconds
+    TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL: int = 3600  # default: 1 hour
+    TFY_INTERNAL_SIGNED_URL_REQUEST_TIMEOUT: int = 3600  # default: 1 hour
+
+    # Artifacts
     TFY_ARTIFACTS_DOWNLOAD_CHUNK_SIZE_BYTES: int = 100 * 1000 * 1000
     TFY_ARTIFACTS_DOWNLOAD_MAX_WORKERS: int = max(min(32, (os.cpu_count() or 2) * 2), 4)
     TFY_ARTIFACTS_UPLOAD_MAX_WORKERS: int = max(min(32, (os.cpu_count() or 2) * 2), 4)

{truefoundry-0.4.4rc8 → truefoundry-0.4.4rc10}/truefoundry/common/tfy_signed_url_client.py

@@ -1,21 +1,24 @@
 # file: client.py
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 from urllib.parse import urlencode, urljoin
 
 from requests.exceptions import RequestException
 
-from truefoundry.common.constants import ENV_VARS
+from truefoundry.common.constants import (
+    ENV_VARS,
+    TFY_INTERNAL_SIGNED_URL_SERVER_HOST_ENV_KEY,
+    TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN_ENV_KEY,
+)
 from truefoundry.common.request_utils import requests_retry_session
 from truefoundry.common.utils import log_time
+from truefoundry.logger import logger
 from truefoundry.pydantic_v1 import BaseModel, Field
 
-# TODO: Move print statements to logger
-# TODO: Move these constants to the constants module
-REQUEST_TIMEOUT = 3600
-
+LOG_PREFIX = "[tfy][fs]"
 DEFAULT_TTL = ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL
 MAX_TIMEOUT = ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_MAX_TIMEOUT
+REQUEST_TIMEOUT = ENV_VARS.TFY_INTERNAL_SIGNED_URL_REQUEST_TIMEOUT
 
 
 class SignedURLAPIResponseDto(BaseModel):
@@ -34,9 +37,9 @@ class SignedURLExistsAPIResponseDto(BaseModel):
 
 class FileInfo(BaseModel):
     path: str
-    is_directoy: bool = Field(..., alias="isDirectory")
+    is_directory: bool = Field(..., alias="isDirectory")
     bytes: Optional[int] = None
-    signedUrl: Optional[str] = None
+    signed_url: Optional[str] = None
 
 
 class PagedList(BaseModel):
@@ -62,51 +65,46 @@ class SignedURLClient:
         base_url: Optional[str] = None,
         token: Optional[str] = None,
         ttl: int = DEFAULT_TTL,
+        max_retries: int = 3,
+        retry_backoff_factor: float = 0.3,
     ):
-        """
-        Initialize the SignedURLClient.
-
-        Args:
-            base_url (str): Base URL of the signed URL server (optional).
-            token (str): Token for authentication (optional).
-            ttl (int): Default time-to-live for signed URLs in seconds.
-        """
+        """Initialize the SignedURLClient."""
+        # Set the base URL and token from the environment variables
         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.signed_url_server_headers = {
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
+        }
+
+        if not self.base_url or not self.token:
+            raise ValueError(
+                f"Set `base_url` and `token` or define env vars "
+                f"{TFY_INTERNAL_SIGNED_URL_SERVER_HOST_ENV_KEY} and "
+                f"{TFY_INTERNAL_SIGNED_URL_SERVER_TOKEN_ENV_KEY}"
+            )
+
+        # Set the TTL, max retries, and backoff factor
         self.ttl = ttl
-        self.max_retries: int = 3
-        self.retry_backoff_factor: float = 0.3
+        self.max_retries: int = max_retries
+        self.retry_backoff_factor: float = retry_backoff_factor
         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
+        self,
+        endpoint: str,
+        method: str = "GET",
+        payload: Optional[Dict] = None,
+        headers: Optional[Dict] = None,
     ) -> Dict:
-        """
-        Internal method to handle requests to the signed URL server.
-
-        Args:
-            endpoint (str): The endpoint for the request.
-            method (str): HTTP method (GET, POST, etc.).
-            payload (Dict, optional): JSON payload for the request.
-
-        Returns:
-            Dict: JSON response from the server.
-
-        Raises:
-            RuntimeError: For network errors or invalid responses.
-        """
+        """Internal method to handle requests to the signed URL server."""
         url = urljoin(self.base_url, endpoint)
         try:
             response = self.session.request(
-                method, url, headers=self.headers, json=payload, timeout=MAX_TIMEOUT
+                method, url, headers=headers, json=payload, timeout=MAX_TIMEOUT
             )
             response.raise_for_status()
             return response.json()
@@ -114,19 +112,30 @@ class SignedURLClient:
             raise RuntimeError(f"Error during request to {url}: {e}") from e
 
     def _make_server_api_call(
-        self, endpoint: SignedURLServerEndpoint, params: Optional[Dict] = None
+        self,
+        endpoint: SignedURLServerEndpoint,
+        params: Optional[Dict] = None,
+        headers: Optional[Dict] = None,
     ) -> Dict:
         """Get a signed URL for the specified operation and URI."""
         query_string = urlencode(params or {})
         endpoint_with_params = f"{endpoint.value}?{query_string}"
-        return self._make_request(endpoint_with_params)
+        return self._make_request(
+            endpoint=endpoint_with_params, method="GET", payload=None, headers=headers
+        )
 
     @log_time
     def _upload_data(self, signed_url: str, data: Any) -> None:
-        """Helper method to upload data to the signed URL."""
+        """
+        Upload data to the specified storage path using a signed URL.
+
+        Args:
+            signed_url: str: The signed URL to upload the data to.
+            data: Bytes or IO: The data to upload.
+        """
         try:
             response = self.session.put(
-                signed_url,
+                url=signed_url,
                 headers={"Content-Type": "application/octet-stream"},
                 data=data,
                 timeout=REQUEST_TIMEOUT,
@@ -140,6 +149,7 @@ class SignedURLClient:
         signed_object = self._make_server_api_call(
             endpoint=SignedURLServerEndpoint.WRITE,
             params={"uri": storage_uri, "expiryInSeconds": self.ttl},
+            headers=self.signed_url_server_headers,
         )
         pre_signed_object_dto = SignedURLAPIResponseDto.parse_obj(signed_object)
         self._upload_data(pre_signed_object_dto.signed_url, data)
@@ -148,10 +158,11 @@ class SignedURLClient:
     @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}")
+        logger.info(f"{LOG_PREFIX} Uploading {file_path} to {storage_uri}")
         response = self._make_server_api_call(
             endpoint=SignedURLServerEndpoint.WRITE,
             params={"uri": storage_uri, "expiryInSeconds": self.ttl},
+            headers=self.signed_url_server_headers,
         )
         pre_signed_object_dto = SignedURLAPIResponseDto.parse_obj(response)
         with open(file_path, "rb") as file:
@@ -162,62 +173,43 @@ class SignedURLClient:
     def _download_file(
         self, signed_url: str, local_path: Optional[str] = None
     ) -> Optional[bytes]:
-        """Common method to download a file using a signed URL.
-
-        Args:
-            signed_url (str): The signed URL to download from.
-            local_path (Optional[str]): The local path to save the file. If None, the file will not be saved.
-
-        Returns:
-            Optional[bytes]: The content of the downloaded file if local_path is None; otherwise, None.
-        """
+        """Common method to download a file using a signed URL."""
         try:
             response = self.session.get(
-                signed_url, stream=True, headers=self.headers, timeout=REQUEST_TIMEOUT
+                signed_url,
+                stream=True,
+                headers={"Content-Type": "application/octet-stream"},
+                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.
-
+                return None
+            return response.content
         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.
-
-        Args:
-            storage_uri (str): The storage URI to download the file from.
-            local_path (str): The local path to save the downloaded file.
-
-        Returns:
-            Optional[str]: The local path if successful; None if the file is saved.
-        """
-        print(f"Dowloading {storage_uri} to {local_path}")
+        """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}")
         response = self._make_server_api_call(
             endpoint=SignedURLServerEndpoint.READ,
             params={"uri": storage_uri, "expiryInSeconds": self.ttl},
+            headers=self.signed_url_server_headers,
         )
         presigned_object = SignedURLAPIResponseDto.parse_obj(response)
         self._download_file(presigned_object.signed_url, local_path)
         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.
-
-        Args:
-            storage_uri (str): The storage URI to download the file from.
-
-        Returns:
-            bytes: The content of the downloaded file.
-        """
+        """Download a file from the specified storage path and return it as bytes."""
         response = self._make_server_api_call(
             endpoint=SignedURLServerEndpoint.READ,
             params={"uri": storage_uri, "expiryInSeconds": self.ttl},
+            headers=self.signed_url_server_headers,
         )
         presigned_object = SignedURLAPIResponseDto.parse_obj(response)
         return self._download_file(presigned_object.signed_url)
@@ -225,22 +217,29 @@ class SignedURLClient:
     def exists(self, uri: str) -> bool:
         """Check if a file exists at the specified path."""
         response = self._make_server_api_call(
-            endpoint=SignedURLServerEndpoint.EXISTS, params={"uri": uri}
+            endpoint=SignedURLServerEndpoint.EXISTS,
+            params={"uri": uri},
+            headers=self.signed_url_server_headers,
         )
-        response_obj = SignedURLExistsAPIResponseDto.parse_obj(response)
-        return response_obj.exists
+        return SignedURLExistsAPIResponseDto.parse_obj(response).exists
 
     def is_directory(self, uri: str) -> bool:
         """Check if the specified URI is a directory."""
         response = self._make_server_api_call(
-            endpoint=SignedURLServerEndpoint.IS_DIRECTORY, params={"path": uri}
+            endpoint=SignedURLServerEndpoint.IS_DIRECTORY,
+            params={"path": uri},
+            headers=self.signed_url_server_headers,
         )
-        response_obj = SignedURLIsDirectoryAPIResponseDto.parse_obj(response)
-        print(f"Path {uri} is_directory: {response_obj.is_directory}")
-        return response_obj.is_directory
+        is_directory = SignedURLIsDirectoryAPIResponseDto.parse_obj(
+            response
+        ).is_directory
+        logger.info(f"{LOG_PREFIX} Path {uri} is a directory: {is_directory}")
+        return is_directory
 
     @log_time
-    def list_files(self, path: str, detail: bool = False, max_results: int = 1000):
+    def list_files(
+        self, path: str, detail: bool = False, max_results: int = 1000
+    ) -> Union[List[FileInfo], List[str]]:
         """List files in the specified directory."""
         token = ""
         items: List[FileInfo] = []
@@ -249,6 +248,7 @@ class SignedURLClient:
             response = self._make_server_api_call(
                 endpoint=SignedURLServerEndpoint.LIST_FILES,
                 params={"path": path, "maxResults": max_results, "pageToken": token},
+                headers=self.signed_url_server_headers,
             )
             response_obj = PagedList.parse_obj(response)
             items.extend(response_obj.items)
@@ -257,6 +257,4 @@ class SignedURLClient:
                 break
 
         # Return the items or paths based on the detail flag
-        if detail:
-            return items
-        return [item.path for item in items]
+        return items if detail else [item.path for item in items]

{truefoundry-0.4.4rc8 → truefoundry-0.4.4rc10}/truefoundry/common/tfy_signed_url_fs.py

@@ -23,7 +23,7 @@ class SignedURLFileSystem(AbstractFileSystem):
 
     @log_time
     def exists(self, path, **kwargs):
-        """Is there a file at the given path"""
+        """Check if a file exists at the given path."""
         return self.client.exists(path)
 
     @log_time
@@ -38,8 +38,8 @@ class SignedURLFileSystem(AbstractFileSystem):
     ):
         """Get file(s) to local"""
         # TODO: Add support for ThreadPoolExecutor here
-        # TODO: Do a proper error handling here
         # TODO: Add support for async download
+        # TODO: Do a proper error handling here
         if self.isdir(rpath):
             if not recursive:
                 raise ValueError(
@@ -52,7 +52,7 @@ class SignedURLFileSystem(AbstractFileSystem):
             for file_info in files:
                 file_path = file_info.path.rstrip("/").rsplit("/")[-1]
 
-                is_directory = file_info.is_directoy
+                is_directory = file_info.is_directory
                 # Construct the relative path for local download
                 relative_path = rpath.rstrip("/") + "/" + file_path
                 target_local_path = lpath.rstrip("/") + "/" + file_path
@@ -71,16 +71,20 @@ class SignedURLFileSystem(AbstractFileSystem):
                             **kwargs,
                         )
                 else:
-                    # Download the file
                     self.client.download(
                         storage_uri=relative_path, local_path=target_local_path
                     )
 
         else:
+            # Ensure the directory exists first
             target_local_path = lpath
             if target_local_path.endswith("/"):
-                Path(target_local_path).mkdir(parents=True, exist_ok=True)
-                target_local_path = target_local_path + rpath.rsplit("/")[-1]
+                # If it ends with "/", it's a directory, so create the directory first
+                target_local_path = os.path.join(
+                    target_local_path, rpath.rsplit("/", 1)[-1]
+                )
+            # Create the directory for the target file path (common for both cases)
+            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
@@ -126,7 +130,8 @@ class SignedURLFileSystem(AbstractFileSystem):
                     )
             return None
         else:
-            # Handle single file upload
+            if rpath.endswith("/"):
+                rpath = os.path.join(rpath, local_path.name)
             return self.client.upload(file_path=lpath, storage_uri=rpath)
 
     @log_time
@@ -189,6 +194,7 @@ class SignedURLFileSystem(AbstractFileSystem):
 class SignedURLBufferedFile(AbstractBufferedFile):
     """
     Buffered file implementation for Signed URL-based file system.
+    # TODO: Need to test this implementation
     """
 
     def __init__(