truefoundry 0.4.3__tar.gz → 0.4.4__tar.gz

{truefoundry-0.4.3 → truefoundry-0.4.4}/PKG-INFO

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

{truefoundry-0.4.3 → truefoundry-0.4.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "truefoundry"
-version = "0.4.3"
+version = "0.4.4"
 description = "Truefoundry CLI"
 authors = ["Abhishek Choudhary <abhishek@truefoundry.com>"]
 readme = "README.md"
@@ -158,5 +158,5 @@ warn_required_dynamic_aliases = true
 
 [tool.pytest.ini_options]
 env = [
-    "TFY_CLI_PYTHONBUILD_PYTHON_IMAGE_REPO=docker.io/truefoundrycloud/test/python",
+    "TFY_PYTHONBUILD_PYTHON_IMAGE_REPO=docker.io/truefoundrycloud/test/python",
 ]

truefoundry-0.4.4/truefoundry/common/constants.py

@@ -0,0 +1,49 @@
+import os
+from pathlib import Path
+from typing import Optional
+
+from truefoundry.pydantic_v1 import BaseSettings, Field, SecretStr
+
+TFY_CONFIG_DIR = Path.home() / ".truefoundry"
+CREDENTIAL_FILEPATH = TFY_CONFIG_DIR / "credentials.json"
+
+# These keys are kept separately because we use them in error messages and some checks
+TFY_HOST_ENV_KEY = "TFY_HOST"
+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"
+
+
+class TrueFoundrySdkEnv(BaseSettings):
+    # Note: Every field in this class should have a default value
+    # 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_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)
+    TFY_ARTIFACTS_DISABLE_MULTIPART_UPLOAD: bool = False
+    TFY_ARTIFACTS_DOWNLOAD_FSYNC_CHUNKS: bool = False
+    TFY_CLI_LOCAL_DEV_MODE: bool = False
+    TFY_PYTHONBUILD_PYTHON_IMAGE_REPO: str = "public.ecr.aws/docker/library/python"
+
+
+ENV_VARS = TrueFoundrySdkEnv()
+API_SERVER_RELATIVE_PATH = "api/svc"
+MLFOUNDRY_SERVER_RELATIVE_PATH = "api/ml"
+VERSION_PREFIX = "v1"
+SERVICEFOUNDRY_CLIENT_MAX_RETRIES = 2

{truefoundry-0.4.3 → truefoundry-0.4.4}/truefoundry/common/credential_provider.py

@@ -3,7 +3,7 @@ import threading
 from abc import ABC, abstractmethod
 
 from truefoundry.common.auth_service_client import AuthServiceClient
-from truefoundry.common.constants import TFY_API_KEY_ENV_KEY
+from truefoundry.common.constants import ENV_VARS, TFY_API_KEY_ENV_KEY
 from truefoundry.common.credential_file_manager import CredentialsFileManager
 from truefoundry.common.entities import CredentialsFileContent, Token
 from truefoundry.common.utils import resolve_tfy_host
@@ -30,7 +30,9 @@ class CredentialProvider(ABC):
 class EnvCredentialProvider(CredentialProvider):
     def __init__(self) -> None:
         logger.debug("Using env var credential provider")
-        api_key = os.getenv(TFY_API_KEY_ENV_KEY)
+        api_key = (
+            ENV_VARS.TFY_API_KEY.get_secret_value() if ENV_VARS.TFY_API_KEY else None
+        )
         if not api_key:
             raise Exception(
                 f"Value of {TFY_API_KEY_ENV_KEY} env var should be non-empty string"

{truefoundry-0.4.3 → truefoundry-0.4.4}/truefoundry/common/request_utils.py

@@ -51,7 +51,7 @@ def urllib3_retry(
 def requests_retry_session(
     retries: int = 2,
     backoff_factor: float = 0.3,
-    status_forcelist=(408, 429, 500, 502, 503, 504, 524),
+    status_forcelist=(408, 429, 417, 500, 502, 503, 504, 524),
     method_whitelist=frozenset({"GET", "POST"}),
     raise_on_status: bool = False,
     session: Optional[requests.Session] = None,

{truefoundry-0.4.3 → truefoundry-0.4.4}/truefoundry/common/servicefoundry_client.py

@@ -43,8 +43,10 @@ def check_min_cli_version(fn):
     return inner
 
 
-def session_with_retries() -> requests.Session:
-    return requests_retry_session(retries=SERVICEFOUNDRY_CLIENT_MAX_RETRIES)
+def session_with_retries(
+    retries: int = SERVICEFOUNDRY_CLIENT_MAX_RETRIES,
+) -> requests.Session:
+    return requests_retry_session(retries=retries)
 
 
 @timed_lru_cache(seconds=30 * 60)

truefoundry-0.4.4/truefoundry/common/tfy_signed_url_client.py

@@ -0,0 +1,260 @@
+# file: client.py
+from enum import Enum
+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,
+    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
+
+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):
+    uri: str
+    signed_url: str
+    headers: Optional[Dict[str, Any]] = None
+
+
+class SignedURLIsDirectoryAPIResponseDto(BaseModel):
+    is_directory: bool = Field(..., alias="isDirectory")
+
+
+class SignedURLExistsAPIResponseDto(BaseModel):
+    exists: bool
+
+
+class FileInfo(BaseModel):
+    path: str
+    is_directory: bool = Field(..., alias="isDirectory")
+    bytes: Optional[int] = None
+    signed_url: Optional[str] = None
+
+
+class PagedList(BaseModel):
+    items: List[FileInfo]
+    token: Optional[str] = None
+
+
+class SignedURLServerEndpoint(str, Enum):
+    """Enumeration for Signed URL Server endpoints."""
+
+    READ = "/v1/signed-uri/read"
+    WRITE = "/v1/signed-uri/write"
+    EXISTS = "/v1/exists"
+    IS_DIRECTORY = "/v1/is-dir"
+    LIST_FILES = "/v1/list-files"
+
+
+class SignedURLClient:
+    """Client to interact with the Signed URL Server for file operations."""
+
+    def __init__(
+        self,
+        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."""
+        # 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 = 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
+        )
+
+    @log_time
+    def _make_request(
+        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."""
+        url = urljoin(self.base_url, endpoint)
+        try:
+            response = self.session.request(
+                method, url, headers=headers, json=payload, timeout=MAX_TIMEOUT
+            )
+            response.raise_for_status()
+            return response.json()
+        except RequestException as e:
+            raise RuntimeError(f"Error during request to {url}: {e}") from e
+
+    def _make_server_api_call(
+        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=endpoint_with_params, method="GET", payload=None, headers=headers
+        )
+
+    @log_time
+    def _upload_data(self, signed_url: str, data: Any) -> None:
+        """
+        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(
+                url=signed_url,
+                headers={"Content-Type": "application/octet-stream"},
+                data=data,
+                timeout=REQUEST_TIMEOUT,
+            )
+            response.raise_for_status()
+        except Exception as e:
+            raise RuntimeError(f"Failed to upload data: {e}") from e
+
+    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(
+            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)
+        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."""
+        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:
+            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]:
+        """Common method to download a file using a signed URL."""
+        try:
+            response = self.session.get(
+                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
+            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."""
+        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."""
+        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)
+
+    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},
+            headers=self.signed_url_server_headers,
+        )
+        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},
+            headers=self.signed_url_server_headers,
+        )
+        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
+    ) -> Union[List[FileInfo], List[str]]:
+        """List files in the specified directory."""
+        token = ""
+        items: List[FileInfo] = []
+        # Fetch all files in the specified path, in pages of max_results
+        while True:
+            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)
+            token = response_obj.token
+            if not token:
+                break
+
+        # Return the items or paths based on the detail flag
+        return items if detail else [item.path for item in items]

truefoundry-0.4.4/truefoundry/common/tfy_signed_url_fs.py

@@ -0,0 +1,244 @@
+# file: tfy_signed_url_fs.py
+# pylint: disable=W0223
+import io
+import os
+from pathlib import Path
+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.utils import log_time
+
+
+class SignedURLFileSystem(AbstractFileSystem):
+    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):
+        """Check if a file exists at the given path."""
+        return self.client.exists(path)
+
+    @log_time
+    def get(
+        self,
+        rpath,
+        lpath,
+        recursive=False,
+        callback=DEFAULT_CALLBACK,
+        maxdepth=None,
+        **kwargs,
+    ):
+        """Get file(s) to local"""
+        # TODO: Add support for ThreadPoolExecutor here
+        # TODO: Add support for async download
+        # TODO: Do a proper error handling here
+        if self.isdir(rpath):
+            if not recursive:
+                raise ValueError(
+                    f"{rpath} is a directory, but recursive is not enabled."
+                )
+
+            # Handle recursive download
+            files = self.ls(rpath, detail=True)
+
+            for file_info in files:
+                file_path = file_info.path.rstrip("/").rsplit("/")[-1]
+
+                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
+
+                if is_directory:
+                    # If it's a directory, create the directory locally
+                    Path(target_local_path).mkdir(parents=True, exist_ok=True)
+                    relative_path = relative_path + "/"
+                    if recursive:
+                        # Recursively download the contents of the directory
+                        self.get(
+                            relative_path,
+                            target_local_path,
+                            recursive=True,
+                            maxdepth=maxdepth,
+                            **kwargs,
+                        )
+                else:
+                    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("/"):
+                # 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
+    def put(
+        self,
+        lpath,
+        rpath,
+        recursive=False,
+        callback=DEFAULT_CALLBACK,
+        maxdepth=None,
+        **kwargs,
+    ):
+        local_path = Path(lpath)
+        if local_path.is_dir():
+            if not recursive:
+                raise ValueError(
+                    f"{lpath} is a directory, but recursive is set to False."
+                )
+
+            # Optionally limit recursion depth
+            max_depth = maxdepth if maxdepth is not None else float("inf")
+
+            # Walk through the directory structure
+            for root, _, files in os.walk(lpath):
+                current_depth = Path(root).relative_to(local_path).parts
+                if len(current_depth) > max_depth:
+                    continue  # Skip files deeper than the max depth
+
+                rel_dir = Path(root).relative_to(local_path)
+                remote_dir = (
+                    rpath.rstrip("/")
+                    if rel_dir == Path(".")
+                    else rpath.rstrip("/") + "/" + str(rel_dir)
+                )
+
+                # Upload each file
+                for file in files:
+                    local_file_path = Path(root) / file
+                    remote_file_path = f"{remote_dir}/{file}"
+                    self.client.upload(
+                        file_path=str(local_file_path),
+                        storage_uri=remote_file_path,
+                    )
+            return None
+        else:
+            if rpath.endswith("/"):
+                rpath = os.path.join(rpath, local_path.name)
+            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)
+
+    def open(
+        self,
+        path,
+        mode="rb",
+        block_size=None,
+        cache_options=None,
+        compression=None,
+        **kwargs,
+    ):
+        """
+        Open a file for reading or writing.
+        """
+        if "r" in mode:
+            # Reading mode
+            file_content = self.client.download_to_bytes(path)
+            return io.BytesIO(file_content)
+        elif "w" in mode or "a" in mode:
+            # Writing mode (appending treated as writing)
+            buffer = io.BytesIO()
+            buffer.seek(0)
+
+            def on_close(buffer=buffer, path=path):
+                """
+                Callback when file is closed, automatically upload the content.
+                """
+                buffer.seek(0)
+                self.client.upload_from_bytes(buffer.read(), storage_uri=path)
+
+            # Wrapping the buffer to automatically upload on close
+            return io.BufferedWriter(buffer, on_close)
+
+    def write(self, path, data, **kwargs):
+        """
+        Write data to the file at the specified path (this could be tied to open's close).
+        """
+        if isinstance(data, io.BytesIO):
+            data.seek(0)
+            content = data.read()
+        elif isinstance(data, str):
+            content = data.encode()  # Encode to bytes
+        else:
+            raise ValueError("Unsupported data type for writing")
+
+        # 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)
+
+
+class SignedURLBufferedFile(AbstractBufferedFile):
+    """
+    Buffered file implementation for Signed URL-based file system.
+    # TODO: Need to test this implementation
+    """
+
+    def __init__(
+        self, fs: SignedURLFileSystem, path: str, mode: str, block_size: int, **kwargs
+    ):
+        """
+        Initialize the buffered file, determining the mode (read/write).
+        """
+        super().__init__(fs, path, mode, block_size, **kwargs)
+        self.buffer = io.BytesIO()
+        self.client = fs.client
+
+        if "r" in mode:
+            # Download the file content for reading
+            file_content = fs.client.download_to_bytes(path)
+            self.buffer.write(file_content)
+            self.buffer.seek(0)  # Reset buffer after writing content
+
+    def _upload_on_close(self):
+        """
+        Upload content back to the remote store when the file is closed.
+        """
+        self.buffer.seek(0)
+        self.client.upload_from_bytes(self.buffer.read(), storage_uri=self.path)
+
+    def close(self):
+        """
+        Close the file, ensuring the content is uploaded for write/append modes.
+        """
+        if self.writable():
+            self._upload_on_close()
+        self.buffer.close()
+        super().close()
+
+    def _fetch_range(self, start, end):
+        """
+        Fetch a specific byte range from the file. Useful for large files and range reads.
+        """
+        self.buffer.seek(start)
+        return self.buffer.read(end - start)
+
+    def _upload_chunk(self, final=False):
+        """
+        Upload a chunk of the file. For larger files, data may be uploaded in chunks.
+        """
+        if final:
+            self._upload_on_close()

{truefoundry-0.4.3 → truefoundry-0.4.4}/truefoundry/common/utils.py

@@ -1,4 +1,3 @@
-import os
 import time
 from functools import lru_cache, wraps
 from time import monotonic_ns
@@ -7,8 +6,8 @@ from urllib.parse import urljoin, urlparse
 
 from truefoundry.common.constants import (
     API_SERVER_RELATIVE_PATH,
+    ENV_VARS,
     MLFOUNDRY_SERVER_RELATIVE_PATH,
-    TFY_CLI_LOCAL_DEV_MODE_ENV_KEY,
     TFY_HOST_ENV_KEY,
 )
 from truefoundry.pydantic_v1 import BaseSettings
@@ -41,7 +40,7 @@ _tfy_servers_config = None
 def get_tfy_servers_config(base_url: str) -> _TFYServersConfig:
     global _tfy_servers_config
     if _tfy_servers_config is None:
-        if os.getenv(TFY_CLI_LOCAL_DEV_MODE_ENV_KEY):
+        if ENV_VARS.TFY_CLI_LOCAL_DEV_MODE:
             _tfy_servers_config = _TFYServersConfig()  # type: ignore[call-arg]
         else:
             _tfy_servers_config = _TFYServersConfig.from_base_url(base_url)
@@ -95,11 +94,25 @@ def validate_tfy_host(tfy_host: str) -> None:
 
 
 def resolve_tfy_host(tfy_host: Optional[str] = None) -> str:
-    if not tfy_host and not os.getenv(TFY_HOST_ENV_KEY):
+    if not tfy_host and not ENV_VARS.TFY_HOST:
         raise ValueError(
             f"Either `host` should be provided using `--host <value>`, or `{TFY_HOST_ENV_KEY}` env must be set"
         )
-    tfy_host = tfy_host or os.getenv(TFY_HOST_ENV_KEY)
+    tfy_host = tfy_host or ENV_VARS.TFY_HOST
     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