Remove open_async usage in put raw data

flytekit/core/data_persistence.py

@@ -199,7 +199,8 @@ def get_filesystem(
         elif protocol == "s3":
             s3kwargs = s3_setup_args(self._data_config.s3, anonymous=anonymous)
             s3kwargs.update(kwargs)
-            return fsspec.filesystem(protocol, **s3kwargs)  # type: ignore
+            fs = fsspec.filesystem(protocol, **s3kwargs)  # type: ignore
+            return fs
         elif protocol == "gs":
             if anonymous:
                 kwargs["token"] = _ANON
@@ -220,8 +221,8 @@ async def get_async_filesystem_for_path(
     ) -> Union[AsyncFileSystem, fsspec.AbstractFileSystem]:
         protocol = get_protocol(path)
         loop = asyncio.get_running_loop()
-
-        return self.get_filesystem(protocol, anonymous=anonymous, path=path, asynchronous=True, loop=loop, **kwargs)
+        fs = self.get_filesystem(protocol, anonymous=anonymous, path=path, asynchronous=True, loop=loop, **kwargs)
+        return fs
 
     def get_filesystem_for_path(self, path: str = "", anonymous: bool = False, **kwargs) -> fsspec.AbstractFileSystem:
         protocol = get_protocol(path)
@@ -425,45 +426,30 @@ async def async_put_raw_data(
 
         # raw bytes
         if isinstance(lpath, bytes):
-            fs = await self.get_async_filesystem_for_path(to_path)
-            if isinstance(fs, AsyncFileSystem):
-                async with fs.open_async(to_path, "wb", **kwargs) as s:
-                    s.write(lpath)
-            else:
-                with fs.open(to_path, "wb", **kwargs) as s:
-                    s.write(lpath)
-
+            fs = self.get_filesystem_for_path(to_path)
+            with fs.open(to_path, "wb", **kwargs) as s:
+                s.write(lpath)
             return to_path
 
         # If lpath is a buffered reader of some kind
         if isinstance(lpath, io.BufferedReader) or isinstance(lpath, io.BytesIO):
             if not lpath.readable():
                 raise FlyteAssertion("Buffered reader must be readable")
-            fs = await self.get_async_filesystem_for_path(to_path)
+            fs = self.get_filesystem_for_path(to_path)
             lpath.seek(0)
-            if isinstance(fs, AsyncFileSystem):
-                async with fs.open_async(to_path, "wb", **kwargs) as s:
-                    while data := lpath.read(read_chunk_size_bytes):
-                        s.write(data)
-            else:
-                with fs.open(to_path, "wb", **kwargs) as s:
-                    while data := lpath.read(read_chunk_size_bytes):
-                        s.write(data)
+            with fs.open(to_path, "wb", **kwargs) as s:
+                while data := lpath.read(read_chunk_size_bytes):
+                    s.write(data)
             return to_path
 
         if isinstance(lpath, io.StringIO):
             if not lpath.readable():
                 raise FlyteAssertion("Buffered reader must be readable")
-            fs = await self.get_async_filesystem_for_path(to_path)
+            fs = self.get_filesystem_for_path(to_path)
             lpath.seek(0)
-            if isinstance(fs, AsyncFileSystem):
-                async with fs.open_async(to_path, "wb", **kwargs) as s:
-                    while data_str := lpath.read(read_chunk_size_bytes):
-                        s.write(data_str.encode(encoding))
-            else:
-                with fs.open(to_path, "wb", **kwargs) as s:
-                    while data_str := lpath.read(read_chunk_size_bytes):
-                        s.write(data_str.encode(encoding))
+            with fs.open(to_path, "wb", **kwargs) as s:
+                while data_str := lpath.read(read_chunk_size_bytes):
+                    s.write(data_str.encode(encoding))
             return to_path
 
         raise FlyteAssertion(f"Unsupported lpath type {type(lpath)}")

tests/flytekit/unit/core/test_data_persistence.py

@@ -1,16 +1,17 @@
 import io
 import os
-import fsspec
 import pathlib
 import random
 import string
 import sys
 import tempfile
 
+import fsspec
 import mock
 import pytest
 from azure.identity import ClientSecretCredential, DefaultAzureCredential
 
+from flytekit.configuration import Config
 from flytekit.core.data_persistence import FileAccessProvider
 from flytekit.core.local_fsspec import FlyteLocalFileSystem
 
@@ -207,3 +208,18 @@ def __init__(self, *args, **kwargs):
 
     fp = FileAccessProvider("/tmp", "s3://my-bucket")
     fp.get_filesystem("testgetfs", test_arg="test_arg")
+
+
+@pytest.mark.sandbox_test
+def test_put_raw_data_bytes():
+    dc = Config.for_sandbox().data_config
+    raw_output = f"s3://my-s3-bucket/"
+    provider = FileAccessProvider(local_sandbox_dir="/tmp/unittest", raw_output_prefix=raw_output, data_config=dc)
+    prefix = provider.get_random_string()
+    provider.put_raw_data(lpath=b"hello", upload_prefix=prefix, file_name="hello_bytes")
+    provider.put_raw_data(lpath=io.BytesIO(b"hello"), upload_prefix=prefix, file_name="hello_bytes_io")
+    provider.put_raw_data(lpath=io.StringIO("hello"), upload_prefix=prefix, file_name="hello_string_io")
+
+    fs = provider.get_filesystem("s3")
+    listing = fs.ls(f"{raw_output}{prefix}/")
+    assert len(listing) == 3