flyte 0.0.1b3__py3-none-any.whl → 0.2.0a0__py3-none-any.whl

flyte/remote/_secret.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import AsyncGenerator, Iterator, Literal, Union
+from typing import AsyncIterator, Literal, Union
 
 import rich.repr
 
-from flyte._api_commons import syncer
-from flyte._initialize import get_client, get_common_config, requires_client
+from flyte._initialize import ensure_client, get_client, get_common_config
 from flyte._protos.secret import definition_pb2, payload_pb2
+from flyte.syncify import syncify
 
 SecretTypes = Literal["regular", "image_pull"]
 
@@ -16,10 +16,10 @@ SecretTypes = Literal["regular", "image_pull"]
 class Secret:
     pb2: definition_pb2.Secret
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def create(cls, name: str, value: Union[str, bytes], type: SecretTypes = "regular"):
+        ensure_client()
         cfg = get_common_config()
         secret_type = (
             definition_pb2.SecretType.SECRET_TYPE_GENERIC
@@ -49,10 +49,10 @@ class Secret:
             ),
         )
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def get(cls, name: str) -> Secret:
+        ensure_client()
         cfg = get_common_config()
         resp = await get_client().secrets_service.GetSecret(
             request=payload_pb2.GetSecretRequest(
@@ -66,10 +66,10 @@ class Secret:
         )
         return Secret(pb2=resp.secret)
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
-    async def listall(cls, limit: int = 100) -> Union[Iterator[Secret], AsyncGenerator[Secret, None]]:
+    async def listall(cls, limit: int = 100) -> AsyncIterator[Secret]:
+        ensure_client()
         cfg = get_common_config()
         token = None
         while True:
@@ -88,10 +88,10 @@ class Secret:
             if not token:
                 break
 
+    @syncify
     @classmethod
-    @requires_client
-    @syncer.wrap
     async def delete(cls, name):
+        ensure_client()
         cfg = get_common_config()
         await get_client().secrets_service.DeleteSecret(  # type: ignore
             request=payload_pb2.DeleteSecretRequest(

flyte/remote/_task.py

@@ -3,17 +3,34 @@ from __future__ import annotations
 import functools
 from dataclasses import dataclass
 from threading import Lock
-from typing import Any, Callable, Coroutine, Dict, Literal, Optional, Union
+from typing import Any, AsyncIterator, Callable, Coroutine, Dict, Iterator, Literal, Optional, Tuple, Union
 
 import rich.repr
+from google.protobuf import timestamp
 
 import flyte
 import flyte.errors
-from flyte._api_commons import syncer
 from flyte._context import internal_ctx
-from flyte._datastructures import NativeInterface
-from flyte._initialize import get_client, get_common_config
+from flyte._initialize import ensure_client, get_client, get_common_config
+from flyte._logging import logger
+from flyte._protos.common import list_pb2
 from flyte._protos.workflow import task_definition_pb2, task_service_pb2
+from flyte.models import NativeInterface
+from flyte.syncify import syncify
+
+
+def _repr_task_metadata(metadata: task_definition_pb2.TaskMetadata) -> rich.repr.Result:
+    """
+    Rich representation of the task metadata.
+    """
+    if metadata.deployed_by:
+        if metadata.deployed_by.user:
+            yield "deployed_by", f"User: {metadata.deployed_by.user.spec.email}"
+        else:
+            yield "deployed_by", f"App: {metadata.deployed_by.application.spec.name}"
+    yield "short_name", metadata.short_name
+    yield "deployed_at", timestamp.to_datetime(metadata.deployed_at)
+    yield "environment_name", metadata.environment_name
 
 
 class LazyEntity:
@@ -22,8 +39,8 @@ class LazyEntity:
     The entity is derived from RemoteEntity so that it behaves exactly like the mimicked entity.
     """
 
-    def __init__(self, name: str, getter: Callable[..., Coroutine[Any, Any, Task]], *args, **kwargs):
-        self._task: Optional[Task] = None
+    def __init__(self, name: str, getter: Callable[..., Coroutine[Any, Any, TaskDetails]], *args, **kwargs):
+        self._task: Optional[TaskDetails] = None
         self._getter = getter
         self._name = name
         self._mutex = Lock()
@@ -32,8 +49,8 @@ class LazyEntity:
     def name(self) -> str:
         return self._name
 
-    @syncer.wrap
-    async def fetch(self) -> Task:
+    @syncify
+    async def fetch(self) -> TaskDetails:
         """
         Forwards all other attributes to task, causing the task to be fetched!
         """
@@ -48,7 +65,7 @@ class LazyEntity:
         """
         Forwards the call to the underlying task. The entity will be fetched if not already present
         """
-        tk = await self.fetch.aio(self)
+        tk = await self.fetch.aio()
         return await tk(*args, **kwargs)
 
     def __repr__(self) -> str:
@@ -62,7 +79,7 @@ AutoVersioning = Literal["latest", "current"]
 
 
 @dataclass
-class Task:
+class TaskDetails:
     pb2: task_definition_pb2.TaskDetails
 
     @classmethod
@@ -87,10 +104,19 @@ class Task:
         if version is None and auto_version not in ["latest", "current"]:
             raise ValueError("auto_version must be either 'latest' or 'current'.")
 
-        async def deferred_get(_version: str | None, _auto_version: AutoVersioning | None) -> Task:
+        async def deferred_get(_version: str | None, _auto_version: AutoVersioning | None) -> TaskDetails:
             if _version is None:
                 if _auto_version == "latest":
-                    raise NotImplementedError("auto_version=latest is not yet implemented.")
+                    tasks = []
+                    async for x in Task.listall.aio(
+                        by_task_name=name,
+                        sort_by=("created_at", "desc"),
+                        limit=1,
+                    ):
+                        tasks.append(x)
+                    if not tasks:
+                        raise flyte.errors.ReferenceTaskError(f"Task {name} not found.")
+                    _version = tasks[0].version
                 elif _auto_version == "current":
                     ctx = flyte.ctx()
                     if ctx is None:
@@ -136,6 +162,20 @@ class Task:
         """
         return self.pb2.spec.task_template.type
 
+    @property
+    def default_input_args(self) -> Tuple[str, ...]:
+        """
+        The default input arguments of the task.
+        """
+        return tuple(x.name for x in self.pb2.spec.default_inputs)
+
+    @property
+    def required_args(self) -> Tuple[str, ...]:
+        """
+        The required input arguments of the task.
+        """
+        return tuple(x for x, _ in self.interface.inputs.items() if x not in self.default_input_args)
+
     @functools.cached_property
     def interface(self) -> NativeInterface:
         """
@@ -143,7 +183,7 @@ class Task:
         """
         import flyte.types as types
 
-        return types.guess_interface(self.pb2.spec.task_template.interface)
+        return types.guess_interface(self.pb2.spec.task_template.interface, default_inputs=self.pb2.spec.default_inputs)
 
     @property
     def cache(self) -> flyte.Cache:
@@ -180,6 +220,19 @@ class Task:
         """
         Forwards the call to the underlying task. The entity will be fetched if not already present
         """
+        # TODO support kwargs, for this we need ordered inputs to be stored in the task spec.
+        if len(args) > 0:
+            raise flyte.errors.ReferenceTaskError(
+                f"Reference task {self.name} does not support positional arguments"
+                f"currently. Please use keyword arguments."
+            )
+        if len(self.required_args) > 0:
+            if len(args) + len(kwargs) < len(self.required_args):
+                raise ValueError(
+                    f"Task {self.name} requires at least {self.required_args} arguments, "
+                    f"but only received args:{args}  kwargs{kwargs}."
+                )
+
         ctx = internal_ctx()
         if ctx.is_task_context():
             # If we are in a task context, that implies we are executing a Run.
@@ -187,7 +240,7 @@ class Task:
             # We will also check if we are not initialized, It is not expected to be not initialized
             from flyte._internal.controllers import get_controller
 
-            controller = await get_controller()
+            controller = get_controller()
             if controller:
                 return await controller.submit_task_ref(self.pb2, *args, **kwargs)
         raise flyte.errors
@@ -205,13 +258,18 @@ class Task:
         env: Optional[Dict[str, str]] = None,
         secrets: Optional[flyte.SecretRequest] = None,
         **kwargs: Any,
-    ) -> Task:
+    ) -> TaskDetails:
         raise NotImplementedError
 
     def __rich_repr__(self) -> rich.repr.Result:
         """
         Rich representation of the task.
         """
+        yield "friendly_name", self.pb2.spec.short_name
+        yield "environment", self.pb2.spec.environment
+        yield "default_inputs_keys", self.default_input_args
+        yield "required_args", self.required_args
+        yield "raw_default_inputs", [str(x) for x in self.pb2.spec.default_inputs]
         yield "project", self.pb2.task_id.project
         yield "domain", self.pb2.task_id.domain
         yield "name", self.name
@@ -223,5 +281,111 @@ class Task:
         yield "resources", self.resources
 
 
+@dataclass
+class Task:
+    pb2: task_definition_pb2.Task
+
+    def __init__(self, pb2: task_definition_pb2.Task):
+        self.pb2 = pb2
+
+    @property
+    def name(self) -> str:
+        """
+        The name of the task.
+        """
+        return self.pb2.task_id.name
+
+    @property
+    def version(self) -> str:
+        """
+        The version of the task.
+        """
+        return self.pb2.task_id.version
+
+    @classmethod
+    def get(cls, name: str, version: str | None = None, auto_version: AutoVersioning | None = None) -> LazyEntity:
+        """
+        Get a task by its ID or name. If both are provided, the ID will take precedence.
+
+        Either version or auto_version are required parameters.
+
+        :param name: The name of the task.
+        :param version: The version of the task.
+        :param auto_version: If set to "latest", the latest-by-time ordered from now, version of the task will be used.
+         If set to "current", the version will be derived from the callee tasks context. This is useful if you are
+         deploying all environments with the same version. If auto_version is current, you can only access the task from
+         within a task context.
+        """
+        return TaskDetails.get(name, version=version, auto_version=auto_version)
+
+    @syncify
+    @classmethod
+    async def listall(
+        cls,
+        by_task_name: str | None = None,
+        sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
+        limit: int = 100,
+    ) -> Union[AsyncIterator[Task], Iterator[Task]]:
+        """
+        Get all runs for the current project and domain.
+
+        :param by_task_name: If provided, only tasks with this name will be returned.
+        :param sort_by: The sorting criteria for the project list, in the format (field, order).
+        :param limit: The maximum number of tasks to return.
+        :return: An iterator of runs.
+        """
+        ensure_client()
+        token = None
+        sort_by = sort_by or ("created_at", "asc")
+        sort_pb2 = list_pb2.Sort(
+            key=sort_by[0], direction=list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING
+        )
+        cfg = get_common_config()
+        filters = []
+        if by_task_name:
+            filters.append(
+                list_pb2.Filter(
+                    function=list_pb2.Filter.Function.EQUAL,
+                    field="name",
+                    values=[by_task_name],
+                )
+            )
+        original_limit = limit
+        if limit > cfg.batch_size:
+            limit = cfg.batch_size
+        retrieved = 0
+        while True:
+            resp = await get_client().task_service.ListTasks(
+                task_service_pb2.ListTasksRequest(
+                    org=cfg.org,
+                    request=list_pb2.ListRequest(
+                        sort_by=sort_pb2,
+                        filters=filters,
+                        limit=limit,
+                        token=token,
+                    ),
+                )
+            )
+            token = resp.token
+            for t in resp.tasks:
+                retrieved += 1
+                yield cls(t)
+            if not token or retrieved >= original_limit:
+                logger.debug(f"Retrieved {retrieved} tasks, stopping iteration.")
+                break
+
+    def __rich_repr__(self) -> rich.repr.Result:
+        """
+        Rich representation of the task.
+        """
+        yield "project", self.pb2.task_id.project
+        yield "domain", self.pb2.task_id.domain
+        yield "name", self.pb2.task_id.name
+        yield "version", self.pb2.task_id.version
+        yield "short_name", self.pb2.metadata.short_name
+        for t in _repr_task_metadata(self.pb2.metadata):
+            yield t
+
+
 if __name__ == "__main__":
     tk = Task.get(name="example_task")

flyte/report/_report.py

@@ -4,10 +4,10 @@ import string
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Dict, List, Union
 
-from flyte._api_commons import syncer
 from flyte._internal.runtime import io
 from flyte._logging import logger
 from flyte._tools import ipython_check
+from flyte.syncify import syncify
 
 if TYPE_CHECKING:
     from IPython.core.display import HTML
@@ -112,7 +112,7 @@ def get_tab(name: str, /, create_if_missing: bool = True) -> Tab:
     return report.get_tab(name, create_if_missing=create_if_missing)
 
 
-@syncer.wrap
+@syncify
 async def log(content: str, do_flush: bool = False):
     """
     Log content to the main tab. The content should be a valid HTML string, but not a complete HTML document,
@@ -126,7 +126,7 @@ async def log(content: str, do_flush: bool = False):
         await flush.aio()
 
 
-@syncer.wrap
+@syncify
 async def flush():
     """
     Flush the report.
@@ -149,7 +149,7 @@ async def flush():
     logger.debug(f"Report flushed to {final_path}")
 
 
-@syncer.wrap
+@syncify
 async def replace(content: str, do_flush: bool = False):
     """
     Get the report. Replaces the content of the main tab.

flyte/storage/__init__.py

@@ -1,4 +1,8 @@
 __all__ = [
+    "ABFS",
+    "GCS",
+    "S3",
+    "Storage",
     "get",
     "get_random_local_directory",
     "get_random_local_path",
@@ -11,6 +15,7 @@ __all__ = [
     "put_stream",
 ]
 
+from ._config import ABFS, GCS, S3, Storage
 from ._storage import (
     get,
     get_random_local_directory,

flyte/storage/_config.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+import datetime
+import os
+import typing
+from dataclasses import dataclass
+from typing import ClassVar
+
+from flyte.config import set_if_exists
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class Storage(object):
+    """
+    Data storage configuration that applies across any provider.
+    """
+
+    retries: int = 3
+    backoff: datetime.timedelta = datetime.timedelta(seconds=5)
+    enable_debug: bool = False
+    attach_execution_metadata: bool = True
+
+    _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
+        "enable_debug": "UNION_STORAGE_DEBUG",
+        "retries": "UNION_STORAGE_RETRIES",
+        "backoff": "UNION_STORAGE_BACKOFF_SECONDS",
+    }
+
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        """
+        Returns the configuration as kwargs for constructing an fsspec filesystem.
+        """
+        return {}
+
+    @classmethod
+    def _auto_as_kwargs(cls) -> typing.Dict[str, typing.Any]:
+        retries = os.getenv(cls._KEY_ENV_VAR_MAPPING["retries"])
+        backoff = os.getenv(cls._KEY_ENV_VAR_MAPPING["backoff"])
+        enable_debug = os.getenv(cls._KEY_ENV_VAR_MAPPING["enable_debug"])
+
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "enable_debug", enable_debug)
+        kwargs = set_if_exists(kwargs, "retries", retries)
+        kwargs = set_if_exists(kwargs, "backoff", backoff)
+        return kwargs
+
+    @classmethod
+    def auto(cls) -> Storage:
+        """
+        Construct the config object automatically from environment variables.
+        """
+        return cls(**cls._auto_as_kwargs())
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class S3(Storage):
+    """
+    S3 specific configuration
+    """
+
+    endpoint: typing.Optional[str] = None
+    access_key_id: typing.Optional[str] = None
+    secret_access_key: typing.Optional[str] = None
+
+    _KEY_ENV_VAR_MAPPING: ClassVar[typing.Dict[str, str]] = {
+        "endpoint": "FLYTE_AWS_ENDPOINT",
+        "access_key_id": "FLYTE_AWS_ACCESS_KEY_ID",
+        "secret_access_key": "FLYTE_AWS_SECRET_ACCESS_KEY",
+    } | Storage._KEY_ENV_VAR_MAPPING
+
+    # Refer to https://github.com/developmentseed/obstore/blob/33654fc37f19a657689eb93327b621e9f9e01494/obstore/python/obstore/store/_aws.pyi#L11
+    # for key and secret
+    _CONFIG_KEY_FSSPEC_S3_KEY_ID: ClassVar = "access_key_id"
+    _CONFIG_KEY_FSSPEC_S3_SECRET: ClassVar = "secret_access_key"
+    _CONFIG_KEY_ENDPOINT: ClassVar = "endpoint_url"
+    _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
+
+    @classmethod
+    def auto(cls) -> S3:
+        """
+        :return: Config
+        """
+        endpoint = os.getenv(cls._KEY_ENV_VAR_MAPPING["endpoint"], None)
+        access_key_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["access_key_id"], None)
+        secret_access_key = os.getenv(cls._KEY_ENV_VAR_MAPPING["secret_access_key"], None)
+
+        kwargs = super()._auto_as_kwargs()
+        kwargs = set_if_exists(kwargs, "endpoint", endpoint)
+        kwargs = set_if_exists(kwargs, "access_key_id", access_key_id)
+        kwargs = set_if_exists(kwargs, "secret_access_key", secret_access_key)
+
+        return S3(**kwargs)
+
+    @classmethod
+    def for_sandbox(cls) -> S3:
+        """
+        :return:
+        """
+        kwargs = super()._auto_as_kwargs()
+        final_kwargs = kwargs | {
+            "endpoint": "http://localhost:4566",
+            "access_key_id": "minio",
+            "secret_access_key": "miniostorage",
+        }
+        return S3(**final_kwargs)
+
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        # Construct the config object
+        kwargs.pop("anonymous", None)  # Remove anonymous if it exists, as we handle it separately
+        config: typing.Dict[str, typing.Any] = {}
+        if self._CONFIG_KEY_FSSPEC_S3_KEY_ID in kwargs or self.access_key_id:
+            config[self._CONFIG_KEY_FSSPEC_S3_KEY_ID] = kwargs.pop(
+                self._CONFIG_KEY_FSSPEC_S3_KEY_ID, self.access_key_id
+            )
+        if self._CONFIG_KEY_FSSPEC_S3_SECRET in kwargs or self.secret_access_key:
+            config[self._CONFIG_KEY_FSSPEC_S3_SECRET] = kwargs.pop(
+                self._CONFIG_KEY_FSSPEC_S3_SECRET, self.secret_access_key
+            )
+        if self._CONFIG_KEY_ENDPOINT in kwargs or self.endpoint:
+            config["endpoint_url"] = kwargs.pop(self._CONFIG_KEY_ENDPOINT, self.endpoint)
+
+        retries = kwargs.pop("retries", self.retries)
+        backoff = kwargs.pop("backoff", self.backoff)
+
+        if anonymous:
+            config[self._KEY_SKIP_SIGNATURE] = True
+
+        retry_config = {
+            "max_retries": retries,
+            "backoff": {
+                "base": 2,
+                "init_backoff": backoff,
+                "max_backoff": datetime.timedelta(seconds=16),
+            },
+            "retry_timeout": datetime.timedelta(minutes=3),
+        }
+
+        client_options = {"timeout": "99999s", "allow_http": True}
+
+        if config:
+            kwargs["config"] = config
+        kwargs["client_options"] = client_options or None
+        kwargs["retry_config"] = retry_config or None
+
+        return kwargs
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class GCS(Storage):
+    """
+    Any GCS specific configuration.
+    """
+
+    gsutil_parallelism: bool = False
+
+    _KEY_ENV_VAR_MAPPING: ClassVar[dict[str, str]] = {
+        "gsutil_parallelism": "GCP_GSUTIL_PARALLELISM",
+    }
+
+    @classmethod
+    def auto(cls) -> GCS:
+        gsutil_parallelism = os.getenv(cls._KEY_ENV_VAR_MAPPING["gsutil_parallelism"], None)
+
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "gsutil_parallelism", gsutil_parallelism)
+        return GCS(**kwargs)
+
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        kwargs.pop("anonymous", None)
+        return kwargs
+
+
+@dataclass(init=True, repr=True, eq=True, frozen=True)
+class ABFS(Storage):
+    """
+    Any Azure Blob Storage specific configuration.
+    """
+
+    account_name: typing.Optional[str] = None
+    account_key: typing.Optional[str] = None
+    tenant_id: typing.Optional[str] = None
+    client_id: typing.Optional[str] = None
+    client_secret: typing.Optional[str] = None
+
+    _KEY_ENV_VAR_MAPPING: ClassVar[dict[str, str]] = {
+        "account_name": "AZURE_STORAGE_ACCOUNT_NAME",
+        "account_key": "AZURE_STORAGE_ACCOUNT_KEY",
+        "tenant_id": "AZURE_TENANT_ID",
+        "client_id": "AZURE_CLIENT_ID",
+        "client_secret": "AZURE_CLIENT_SECRET",
+    }
+    _KEY_SKIP_SIGNATURE: ClassVar = "skip_signature"
+
+    @classmethod
+    def auto(cls) -> ABFS:
+        account_name = os.getenv(cls._KEY_ENV_VAR_MAPPING["account_name"], None)
+        account_key = os.getenv(cls._KEY_ENV_VAR_MAPPING["account_key"], None)
+        tenant_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["tenant_id"], None)
+        client_id = os.getenv(cls._KEY_ENV_VAR_MAPPING["client_id"], None)
+        client_secret = os.getenv(cls._KEY_ENV_VAR_MAPPING["client_secret"], None)
+
+        kwargs: typing.Dict[str, typing.Any] = {}
+        kwargs = set_if_exists(kwargs, "account_name", account_name)
+        kwargs = set_if_exists(kwargs, "account_key", account_key)
+        kwargs = set_if_exists(kwargs, "tenant_id", tenant_id)
+        kwargs = set_if_exists(kwargs, "client_id", client_id)
+        kwargs = set_if_exists(kwargs, "client_secret", client_secret)
+        return ABFS(**kwargs)
+
+    def get_fsspec_kwargs(self, anonymous: bool = False, **kwargs) -> typing.Dict[str, typing.Any]:
+        kwargs.pop("anonymous", None)
+        config: typing.Dict[str, typing.Any] = {}
+        if "account_name" in kwargs or self.account_name:
+            config["account_name"] = kwargs.get("account_name", self.account_name)
+        if "account_key" in kwargs or self.account_key:
+            config["account_key"] = kwargs.get("account_key", self.account_key)
+        if "client_id" in kwargs or self.client_id:
+            config["client_id"] = kwargs.get("client_id", self.client_id)
+        if "client_secret" in kwargs or self.client_secret:
+            config["client_secret"] = kwargs.get("client_secret", self.client_secret)
+        if "tenant_id" in kwargs or self.tenant_id:
+            config["tenant_id"] = kwargs.get("tenant_id", self.tenant_id)
+
+        if anonymous:
+            config[self._KEY_SKIP_SIGNATURE] = True
+
+        client_options = {"timeout": "99999s", "allow_http": "true"}
+
+        if config:
+            kwargs["config"] = config
+        kwargs["client_options"] = client_options
+
+        return kwargs

flyte/storage/_storage.py

@@ -74,7 +74,27 @@ def get_underlying_filesystem(
 
     storage_config = get_storage()
     if storage_config:
-        kwargs.update(storage_config.get_fsspec_kwargs(anonymous, **kwargs))
+        kwargs = storage_config.get_fsspec_kwargs(anonymous, **kwargs)
+    elif protocol:
+        match protocol:
+            case "s3":
+                # If the protocol is s3, we can use the s3 filesystem
+                from flyte.storage import S3
+
+                kwargs = S3.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+            case "gs":
+                # If the protocol is gs, we can use the gs filesystem
+                from flyte.storage import GCS
+
+                kwargs = GCS.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+            case "abfs" | "abfss":
+                # If the protocol is abfs or abfss, we can use the abfs filesystem
+                from flyte.storage import ABFS
+
+                kwargs = ABFS.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+            case _:
+                pass
+
     return fsspec.filesystem(protocol, **kwargs)
 
 
@@ -127,7 +147,7 @@ async def _get_from_filesystem(
     return to_path
 
 
-async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = False, **kwargs):
+async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = False, **kwargs) -> str:
     if not to_path:
         from flyte._context import internal_ctx
 
@@ -142,7 +162,7 @@ async def put(from_path: str, to_path: Optional[str] = None, recursive: bool = F
     else:
         dst = file_system.put(from_path, to_path, recursive=recursive, **kwargs)
     if isinstance(dst, (str, pathlib.Path)):
-        return dst
+        return str(dst)
     else:
         return to_path