truefoundry 0.4.4rc12__py3-none-any.whl → 0.5.0rc1__py3-none-any.whl

truefoundry/ml/log_types/artifacts/utils.py

@@ -2,31 +2,31 @@ import json
 import logging
 import os
 import posixpath
-import shutil
-import tempfile
 from pathlib import Path
 from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.log_types.artifacts.constants import DESCRIPTION_MAX_LENGTH
 
-logger = logging.getLogger("truefoundry.ml")
+logger = logging.getLogger(__name__)
 
 
-def _copy_tree(src_path, dest_path, symlinks=False, ignore_dangling_symlinks=False):
+def _copy_tree(
+    root_dir: str, src_path: str, dest_path: str, dest_to_src: Dict[str, str]
+):
     os.makedirs(dest_path, exist_ok=True)
     for item in os.listdir(src_path):
         src = os.path.join(src_path, item)
         dest = os.path.join(dest_path, item)
         if os.path.isdir(src):
             _copy_tree(
-                src,
-                dest,
-                symlinks=symlinks,
-                ignore_dangling_symlinks=ignore_dangling_symlinks,
+                root_dir=root_dir,
+                src_path=src,
+                dest_path=dest,
+                dest_to_src=dest_to_src,
             )
         else:
-            shutil.copy2(src, dest, follow_symlinks=True)
+            dest_to_src[dest] = src
 
 
 def is_destination_path_dirlike(dest_path) -> bool:
@@ -48,9 +48,9 @@ def _copy_additional_files(
     model_dir: Optional[str],  # relative to files_dir e.g "model/"
     additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]],
     ignore_model_dir_dest_conflict: bool = False,
-):
+    existing_dest_to_src_map: Optional[Dict[str, str]] = None,
+) -> Dict[str, str]:
     """
-
     File copying examples:
         # non ambiguous
         # a.txt -> /tmp/                            result /tmp/a.txt
@@ -69,6 +69,7 @@ def _copy_additional_files(
         # .gitignore -> /tmp/.gitinclude             result /tmp/.gitinclude
         # a.txt -> /tmp/a                            result /tmp/a
     """
+    dest_to_src = existing_dest_to_src_map or {}
     for src_path, dest_path in additional_files:
         src_path = str(src_path)
         if not os.path.exists(src_path):
@@ -99,26 +100,47 @@ def _copy_additional_files(
             _src = src_path
             if is_destination_path_dirlike(dest_abs_path):
                 os.makedirs(dest_abs_path, exist_ok=True)
-                _dst = os.path.relpath(
-                    os.path.join(dest_abs_path, os.path.basename(_src)), files_abs_dir
-                )
+                dest_abs_path = os.path.join(dest_abs_path, os.path.basename(_src))
             else:
                 os.makedirs(os.path.dirname(dest_abs_path), exist_ok=True)
-                _dst = os.path.relpath(dest_abs_path, files_abs_dir)
+            _dst = os.path.relpath(dest_abs_path, files_abs_dir)
             logger.info(f"Adding file {_src} as /{_dst}")
-            shutil.copy2(src_path, dest_abs_path, follow_symlinks=True)
+            dest_to_src[dest_abs_path] = src_path
         elif os.path.isdir(src_path):
             os.makedirs(dest_abs_path, exist_ok=True)
             _src = src_path.rstrip("/")
             _dst = os.path.relpath(dest_abs_path, files_abs_dir).rstrip("/")
             logger.info(f"Adding contents of {_src}/ to /{_dst}/")
             _copy_tree(
+                root_dir=root_dir,
                 src_path=src_path,
                 dest_path=dest_abs_path,
-                symlinks=True,
-                ignore_dangling_symlinks=False,
+                dest_to_src=dest_to_src,
             )
 
+    return dest_to_src
+
+
+def _make_dest_to_src_map_from_dir(root_dir: str) -> Dict[str, str]:
+    dest_to_src_map = {}
+    for root, _, files in os.walk(root_dir):
+        for file in files:
+            src = os.path.join(root, file)
+            dest = src
+            dest_to_src_map[dest] = src
+    return dest_to_src_map
+
+
+def _get_src_dest_pairs(
+    root_dir: str,
+    dest_to_src_map: Dict[str, str],
+) -> Sequence[Tuple[str, str]]:
+    src_dest_pairs = [
+        (src_path, os.path.relpath(dest_abs_path, root_dir))
+        for dest_abs_path, src_path in dest_to_src_map.items()
+    ]
+    return src_dest_pairs
+
 
 def _validate_description(description: Optional[str]):
     if description is not None:
@@ -141,8 +163,8 @@ def _validate_artifact_metadata(metadata: Dict[str, Any]):
         raise MlFoundryException("`metadata` must be json serializable dict") from ve
 
 
-def calculate_local_directory_size(
-    directory: tempfile.TemporaryDirectory,  # type: ignore[type-arg]
+def calculate_total_size(
+    paths: Sequence[str],
 ):
     """
     Tells about the size of the artifact
@@ -153,9 +175,4 @@ def calculate_local_directory_size(
     Returns:
         total size of the artifact
     """
-    total_size = 0
-    for path, _dirs, files in os.walk(directory.name):
-        for f in files:
-            file_path = os.path.join(path, f)
-            total_size += os.stat(file_path).st_size
-    return total_size
+    return sum(os.stat(os.path.realpath(file_path)).st_size for file_path in paths)

truefoundry/ml/log_types/image/image.py

@@ -18,6 +18,7 @@ from truefoundry.ml.log_types.artifacts.constants import (
     FILES_DIR,
     INTERNAL_METADATA_PATH,
 )
+from truefoundry.ml.log_types.artifacts.utils import _make_dest_to_src_map_from_dir
 from truefoundry.ml.log_types.image.constants import (
     DEFAULT_IMAGE_FORMAT,
     IMAGE_KEY_REGEX,
@@ -344,6 +345,7 @@ class Image:
             name=key,
             artifact_type=ArtifactType.IMAGE,
             artifact_dir=temp_dir,
+            dest_to_src_map=_make_dest_to_src_map_from_dir(root_dir=temp_dir.name),
             internal_metadata=internal_metadata,
             step=step,
         )

truefoundry/ml/log_types/plot.py

@@ -18,6 +18,7 @@ from truefoundry.ml.log_types.artifacts.constants import (
     FILES_DIR,
     INTERNAL_METADATA_PATH,
 )
+from truefoundry.ml.log_types.artifacts.utils import _make_dest_to_src_map_from_dir
 from truefoundry.ml.log_types.pydantic_base import PydanticBase
 from truefoundry.ml.log_types.utils import validate_key_name
 from truefoundry.pydantic_v1 import BaseModel
@@ -186,6 +187,7 @@ class Plot:
             name=key,
             artifact_type=ArtifactType.PLOT,
             artifact_dir=temp_dir,
+            dest_to_src_map=_make_dest_to_src_map_from_dir(root_dir=temp_dir.name),
             internal_metadata=internal_metadata,
             step=step,
         )

truefoundry/ml/mlfoundry_api.py

@@ -2,7 +2,17 @@ import os
 import time
 import uuid
 from pathlib import Path
-from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import coolname
 import pandas as pd
@@ -42,7 +52,11 @@ from truefoundry.ml.log_types.artifacts.artifact import (
 )
 from truefoundry.ml.log_types.artifacts.dataset import DataDirectory
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
-from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
+from truefoundry.ml.log_types.artifacts.model import (
+    BlobStorageModelDirectory,
+    ModelVersion,
+    _log_model_version,
+)
 from truefoundry.ml.logger import logger
 from truefoundry.ml.mlfoundry_run import MlFoundryRun
 from truefoundry.ml.session import (
@@ -57,6 +71,9 @@ from truefoundry.ml.validation_utils import (
     _validate_run_name,
 )
 
+if TYPE_CHECKING:
+    from truefoundry.ml import ModelFrameworkType
+
 _SEARCH_MAX_RESULTS_DEFAULT = 1000
 
 _INTERNAL_ENV_VARS = [
@@ -1209,12 +1226,12 @@ class MlFoundry:
         *,
         ml_repo: str,
         name: str,
-        model_file_or_folder: str,
-        framework: Optional[Union[ModelFramework, str]],
+        model_file_or_folder: Union[str, BlobStorageModelDirectory],
         additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         progress: Optional[bool] = None,
+        framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
     ) -> ModelVersion:
         """
         Serialize and log a versioned model under the current ml_repo. Each logged model generates a new version
@@ -1226,12 +1243,24 @@ class MlFoundry:
             name (str): Name of the model. If a model with this name already exists under the current ML Repo,
                 the logged model will be added as a new version under that `name`. If no models exist with the given
                 `name`, the given model will be logged as version 1.
-            model_file_or_folder (str): Path to either a single file or a folder containing model files. This folder
-                is usually created using serialization methods of libraries or frameworks e.g. `joblib.dump`,
-                `model.save_pretrained(...)`, `torch.save(...)`, `model.save(...)`
-            framework (Union[enums.ModelFramework, str]): Model Framework. Ex:- pytorch, sklearn, tensorflow etc.
-                The full list of supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
-                Can also be `None` when `model` is `None`.
+
+            model_file_or_folder (Union[str, BlobStorageModelDirectory]):
+                str:
+                    Path to either a single file or a folder containing model files.
+                    This folder is typically created using serialization methods from libraries or frameworks,
+                    e.g., `joblib.dump`, `model.save_pretrained(...)`, `torch.save(...)`, or `model.save(...)`.
+                BlobStorageModelDirectory:
+                    uri (str): URI to the model file or folder in a storage integration associated with the specified ML Repo.
+                        The model files or folder must reside within the same storage integration as the specified ML Repo.
+                        Accepted URI formats include `s3://integration-bucket-name/prefix/path/to/model` or `gs://integration-bucket-name/prefix/path/to/model`.
+                        If the URI points to a model in a different storage integration, an error will be raised indicating "Invalid source URI."
+
+            framework (Optional[Union[ModelFramework, ModelFrameworkType]]): Framework used for model serialization.
+                Supported frameworks values (ModelFrameworkType) can be imported from `from truefoundry.ml import *`.
+                Supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
+                Can also be `None` if the framework is not known or not supported.
+                **Deprecated**: Prefer `ModelFrameworkType` over `enums.ModelFramework`.
+
             additional_files (Sequence[Tuple[Union[str, Path], Optional[str]]], optional): A list of pairs
                 of (source path, destination path) to add additional files and folders
                 to the model version contents. The first member of the pair should be a file or directory path
@@ -1245,10 +1274,12 @@ class MlFoundry:
                 You can also add additional files to model/ subdirectory by specifying the destination path as model/
 
                 ```python
+                from truefoundry.ml import TensorFlowFramework
+
                 run.log_model(
                     name="xyz",
                     model_file_or_folder="clf.joblib",
-                    framework="sklearn",
+                    framework=TensorFlowFramework(),
                     additional_files=[("foo.txt", "foo/bar/foo.txt"), ("tokenizer/", "foo/tokenizer/")]
                 )
                 ```
@@ -1281,8 +1312,7 @@ class MlFoundry:
             ### Sklearn
 
             ```python
-            from truefoundry.ml import get_client
-            from truefoundry.ml.enums import ModelFramework
+            from truefoundry.ml import get_client, SklearnFramework
 
             import joblib
             import numpy as np
@@ -1307,7 +1337,7 @@ class MlFoundry:
                 ml_repo="my-classification-project",
                 name="my-sklearn-model",
                 model_file_or_folder="sklearn-pipeline.joblib",
-                framework=ModelFramework.SKLEARN,
+                framework=SklearnFramework(),
                 metadata={"accuracy": 0.99, "f1": 0.80},
                 step=1,  # step number, useful when using iterative algorithms like SGD
             )
@@ -1317,8 +1347,7 @@ class MlFoundry:
             ### Huggingface Transformers
 
             ```python
-            from truefoundry.ml import get_client
-            from truefoundry.ml.enums import ModelFramework
+            from truefoundry.ml import get_client, TransformersFramework, LibraryName
 
             import torch
             from transformers import AutoTokenizer, AutoConfig, pipeline, AutoModelForCausalLM
@@ -1342,7 +1371,7 @@ class MlFoundry:
                 ml_repo="my-llm-project",
                 name="my-transformers-model",
                 model_file_or_folder="my-transformers-model/",
-                framework=ModelFramework.TRANSFORMERS
+                framework=TransformersFramework(library_name=LibraryName.TRANSFORMERS, pipeline_tag='text-generation')
             )
             print(model_version.fqn)
             ```
@@ -1356,12 +1385,12 @@ class MlFoundry:
             ml_repo_id=ml_repo_id,
             name=name,
             model_file_or_folder=model_file_or_folder,
-            framework=framework,
             additional_files=additional_files,
             description=description,
             metadata=metadata,
             step=None,
             progress=progress,
+            framework=framework,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version

truefoundry/ml/mlfoundry_run.py

@@ -20,7 +20,7 @@ from urllib.parse import urljoin, urlsplit
 
 from truefoundry import version
 from truefoundry.common.utils import relogin_error_message
-from truefoundry.ml import constants, enums
+from truefoundry.ml import constants
 from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     ArtifactType,
     DeleteRunRequest,
@@ -40,13 +40,17 @@ from truefoundry.ml.autogen.client import (  # type: ignore[attr-defined]
     UpdateRunRequestDto,
 )
 from truefoundry.ml.entities import Metric
-from truefoundry.ml.enums import RunStatus
+from truefoundry.ml.enums import ModelFramework, RunStatus
 from truefoundry.ml.exceptions import MlFoundryException
 from truefoundry.ml.internal_namespace import NAMESPACE
 from truefoundry.ml.log_types import Image, Plot
 from truefoundry.ml.log_types.artifacts.artifact import ArtifactPath, ArtifactVersion
 from truefoundry.ml.log_types.artifacts.general_artifact import _log_artifact_version
-from truefoundry.ml.log_types.artifacts.model import ModelVersion, _log_model_version
+from truefoundry.ml.log_types.artifacts.model import (
+    BlobStorageModelDirectory,
+    ModelVersion,
+    _log_model_version,
+)
 from truefoundry.ml.logger import logger
 from truefoundry.ml.run_utils import ParamsType, flatten_dict, process_params
 from truefoundry.ml.session import ACTIVE_RUNS, _get_api_client, get_active_session
@@ -61,6 +65,8 @@ if TYPE_CHECKING:
     import matplotlib
     import plotly
 
+    from truefoundry.ml import ModelFrameworkType
+
 
 def _ensure_not_deleted(method):
     @functools.wraps(method)
@@ -920,13 +926,13 @@ class MlFoundryRun:
         self,
         *,
         name: str,
-        model_file_or_folder: str,
-        framework: Optional[Union[enums.ModelFramework, str]],
+        model_file_or_folder: Union[str, BlobStorageModelDirectory],
         additional_files: Sequence[Tuple[Union[str, Path], Optional[str]]] = (),
         description: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None,
         step: int = 0,
         progress: Optional[bool] = None,
+        framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
     ) -> ModelVersion:
         # TODO (chiragjn): Document mapping of framework to list of valid model save kwargs
         # TODO (chiragjn): Add more examples
@@ -939,12 +945,21 @@ class MlFoundryRun:
             name (str): Name of the model. If a model with this name already exists under the current ML Repo,
                 the logged model will be added as a new version under that `name`. If no models exist with the given
                 `name`, the given model will be logged as version 1.
-            model_file_or_folder (str): Path to either a single file or a folder containing model files. This folder
-                is usually created using serialization methods of libraries or frameworks e.g. `joblib.dump`,
-                `model.save_pretrained(...)`, `torch.save(...)`, `model.save(...)`
-            framework (Union[enums.ModelFramework, str]): Model Framework. Ex:- pytorch, sklearn, tensorflow etc.
-                The full list of supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
-                Can also be `None` when `model` is `None`.
+            model_file_or_folder (Union[str, BlobStorageModelDirectory]):
+                str:
+                    Path to either a single file or a folder containing model files.
+                    This folder is typically created using serialization methods from libraries or frameworks,
+                    e.g., `joblib.dump`, `model.save_pretrained(...)`, `torch.save(...)`, or `model.save(...)`.
+                BlobStorageModelDirectory:
+                    uri (str): URI to the model file or folder in a storage integration associated with the specified ML Repo.
+                        The model files or folder must reside within the same storage integration as the specified ML Repo.
+                        Accepted URI formats include `s3://integration-bucket-name/prefix/path/to/model` or `gs://integration-bucket-name/prefix/path/to/model`.
+                        If the URI points to a model in a different storage integration, an error will be raised indicating "Invalid source URI."
+            framework (Optional[Union[ModelFramework, ModelFrameworkType]]): Framework used for model serialization.
+                Supported frameworks values (ModelFrameworkType) can be imported from `from truefoundry.ml import *`.
+                Supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
+                Can also be `None` if the framework is not known or not supported.
+                **Deprecated**: Prefer `ModelFrameworkType` over `enums.ModelFramework`.
             additional_files (Sequence[Tuple[Union[str, Path], Optional[str]]], optional): A list of pairs
                 of (source path, destination path) to add additional files and folders
                 to the model version contents. The first member of the pair should be a file or directory path
@@ -1064,12 +1079,12 @@ class MlFoundryRun:
             run=self,
             name=name,
             model_file_or_folder=model_file_or_folder,
-            framework=framework,
             additional_files=additional_files,
             description=description,
             metadata=metadata,
             step=step,
             progress=progress,
+            framework=framework,
         )
         logger.info(f"Logged model successfully with fqn {model_version.fqn!r}")
         return model_version

truefoundry/ml/model_framework.py

@@ -0,0 +1,169 @@
+import warnings
+from typing import Literal, Optional, Union, get_args
+
+from truefoundry.ml import ModelFramework
+from truefoundry.ml.autogen.entities import artifacts as autogen_artifacts
+from truefoundry.pydantic_v1 import BaseModel, Field
+
+
+class FastAIFramework(autogen_artifacts.FastAIFramework):
+    """FastAI model Framework"""
+
+    type: Literal["fastai"] = "fastai"
+
+
+class GluonFramework(autogen_artifacts.GluonFramework):
+    """Gluon model Framework"""
+
+    type: Literal["gluon"] = "gluon"
+
+
+class H2OFramework(autogen_artifacts.H2OFramework):
+    """H2O model Framework"""
+
+    type: Literal["h2o"] = "h2o"
+
+
+class KerasFramework(autogen_artifacts.KerasFramework):
+    """Keras model Framework"""
+
+    type: Literal["keras"] = "keras"
+
+
+class LightGBMFramework(autogen_artifacts.LightGBMFramework):
+    """LightGBM model Framework"""
+
+    type: Literal["lightgbm"] = "lightgbm"
+
+
+class ONNXFramework(autogen_artifacts.ONNXFramework):
+    """ONNX model Framework"""
+
+    type: Literal["onnx"] = "onnx"
+
+
+class PaddleFramework(autogen_artifacts.PaddleFramework):
+    """Paddle model Framework"""
+
+    type: Literal["paddle"] = "paddle"
+
+
+class PyTorchFramework(autogen_artifacts.PyTorchFramework):
+    """PyTorch model Framework"""
+
+    type: Literal["pytorch"] = "pytorch"
+
+
+class SklearnFramework(autogen_artifacts.SklearnFramework):
+    """Sklearn model Framework"""
+
+    type: Literal["sklearn"] = "sklearn"
+
+
+class SpaCyFramework(autogen_artifacts.SpaCyFramework):
+    """SpaCy model Framework"""
+
+    type: Literal["spacy"] = "spacy"
+
+
+class StatsModelsFramework(autogen_artifacts.StatsModelsFramework):
+    """StatsModels model Framework"""
+
+    type: Literal["statsmodels"] = "statsmodels"
+
+
+class TensorFlowFramework(autogen_artifacts.TensorFlowFramework):
+    """TensorFlow model Framework"""
+
+    type: Literal["tensorflow"] = "tensorflow"
+
+
+class TransformersFramework(autogen_artifacts.TransformersFramework):
+    """Transformers model Framework"""
+
+    type: Literal["transformers"] = "transformers"
+
+
+class XGBoostFramework(autogen_artifacts.XGBoostFramework):
+    """XGBoost model Framework"""
+
+    type: Literal["xgboost"] = "xgboost"
+
+
+# Union of all the model frameworks
+
+
+ModelFrameworkType = Union[
+    FastAIFramework,
+    GluonFramework,
+    H2OFramework,
+    KerasFramework,
+    LightGBMFramework,
+    ONNXFramework,
+    PaddleFramework,
+    PyTorchFramework,
+    SklearnFramework,
+    SpaCyFramework,
+    StatsModelsFramework,
+    TensorFlowFramework,
+    TransformersFramework,
+    XGBoostFramework,
+]
+
+
+class _ModelFramework(BaseModel):
+    __root__: ModelFrameworkType = Field(discriminator="type")
+
+    @classmethod
+    def to_model_framework_type(
+        cls,
+        framework: Optional[Union[str, ModelFramework, "ModelFrameworkType"]] = None,
+    ) -> Optional["ModelFrameworkType"]:
+        """
+        Converts a ModelFramework or string representation to a ModelFrameworkType object.
+
+        Args:
+            framework (Optional[Union[str, ModelFramework, ModelFrameworkType]]): ModelFrameworkType or equivalent input.
+                Supported frameworks can be found in `truefoundry.ml.enums.ModelFramework`.
+                May be `None` if the framework is unknown or unsupported.
+                **Deprecated**: Prefer passing a `ModelFrameworkType` instance.
+
+        Returns:
+            ModelFrameworkType corresponding to the input, or None if the input is None.
+        """
+        if framework is None:
+            return None
+
+        # Issue a deprecation warning for str and ModelFramework types
+        if isinstance(framework, (str, ModelFramework)):
+            warnings.warn(
+                "Passing a string or ModelFramework Enum is deprecated. Please use a ModelFrameworkType object.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+
+        # Convert string to ModelFramework
+        if isinstance(framework, str):
+            framework = ModelFramework(framework)
+
+        # Convert ModelFramework to ModelFrameworkType
+        if isinstance(framework, ModelFramework):
+            if framework == ModelFramework.UNKNOWN:
+                return None
+            return cls.parse_obj({"type": framework.value}).__root__
+
+        # Directly return if already a ModelFrameworkType
+        if isinstance(framework, get_args(ModelFrameworkType)):
+            return framework
+
+        raise ValueError(
+            "framework must be a string, ModelFramework enum, or ModelFrameworkType object"
+        )
+
+    @classmethod
+    def from_dict(cls, obj: dict) -> ModelFrameworkType:
+        """Create an instance of ModelFramework from a dict"""
+        if obj is None:
+            return None
+
+        return cls.parse_obj(obj).__root__

truefoundry/workflow/__init__.py

@@ -6,9 +6,9 @@ except ImportError:
 
 from flytekit import conditional
 from flytekit.types.directory import FlyteDirectory
+from flytekit.types.file import FlyteFile
 
 from truefoundry.common.constants import ENV_VARS
-from truefoundry.common.tfy_signed_url_fs import SignedURLFileSystem
 from truefoundry.deploy.v2.lib.patched_models import (
     ContainerTaskConfig,
     PythonTaskConfig,
@@ -18,6 +18,7 @@ from truefoundry.deploy.v2.lib.patched_models import (
 from truefoundry.workflow.container_task import ContainerTask
 from truefoundry.workflow.map_task import map_task
 from truefoundry.workflow.python_task import PythonFunctionTask
+from truefoundry.workflow.remote_filesystem.tfy_signed_url_fs import SignedURLFileSystem
 from truefoundry.workflow.task import task
 from truefoundry.workflow.workflow import ExecutionConfig, workflow
 
@@ -34,6 +35,7 @@ __all__ = [
     "ContainerTaskConfig",
     "PythonTaskConfig",
     "ExecutionConfig",
+    "FlyteFile",
 ]
 
 

truefoundry/workflow/remote_filesystem/__init__.py

@@ -0,0 +1,8 @@
+import logging
+
+from truefoundry.common.constants import ENV_VARS
+from truefoundry.workflow.remote_filesystem.logger import init_logger
+
+init_logger(
+    level=logging.getLevelNamesMapping[ENV_VARS.TFY_SIGNED_URL_CLIENT_LOG_LEVEL]
+)

truefoundry/workflow/remote_filesystem/logger.py

@@ -0,0 +1,36 @@
+import logging
+import sys
+from functools import wraps
+from timeit import default_timer
+
+logger = logging.getLogger("truefoundry.workflow.remote_filesystem")
+
+
+def init_logger(level=logging.WARNING):
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setLevel(level)
+    formatter = logging.Formatter(
+        "[%(name)s] %(asctime)s %(levelname)s %(message)s",
+        datefmt="%Y-%m-%dT%H:%M:%S%z",
+    )
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False
+
+
+def log_time(prefix: str = ""):
+    """Decorator to log the time taken by I/O operations."""
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            start_time = default_timer()
+            result = func(*args, **kwargs)
+            elapsed_time = default_timer() - start_time
+            logger.info(f"{prefix}{func.__name__} took {elapsed_time:.2f} seconds")
+            return result
+
+        return wrapper
+
+    return decorator

truefoundry/{common → workflow/remote_filesystem}/tfy_signed_url_client.py

@@ -11,9 +11,8 @@ from truefoundry.common.constants import (
     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
+from truefoundry.workflow.remote_filesystem.logger import log_time, logger
 
 LOG_PREFIX = "[tfy][fs]"
 DEFAULT_TTL = ENV_VARS.TFY_INTERNAL_SIGNED_URL_SERVER_DEFAULT_TTL