Set cache_regions=True in the case of s3fs (#1592)

* Set cache_regions=True in the case of s3fs * Linting Signed-off-by: eduardo apolinario <[email protected]> * Fix tests Signed-off-by: eduardo apolinario <[email protected]> * More linting Signed-off-by: eduardo apolinario <[email protected]> --------- Signed-off-by: eduardo apolinario <[email protected]> Co-authored-by: eduardo apolinario <[email protected]>
flyteorg · May 16, 2023 · 8e26a59 · 8e26a59
1 parent f8cc630
commit 8e26a59
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py
@@ -25,7 +25,7 @@
 import pathlib
 import tempfile
 import typing
-from typing import Union, cast
+from typing import Any, Dict, Union, cast
 from uuid import UUID
 
 import fsspec
@@ -46,7 +46,9 @@
 
 
 def s3_setup_args(s3_cfg: configuration.S3Config, anonymous: bool = False):
-    kwargs = {}
+    kwargs: Dict[str, Any] = {
+        "cache_regions": True,
+    }
     if s3_cfg.access_key_id:
         kwargs[_FSSPEC_S3_KEY_ID] = s3_cfg.access_key_id
 

diff --git a/tests/flytekit/unit/core/test_data.py b/tests/flytekit/unit/core/test_data.py
@@ -176,7 +176,7 @@ def test_s3_setup_args_env_empty(mock_os, mock_get_config_file):
     mock_os.get.return_value = None
     s3c = S3Config.auto()
     kwargs = s3_setup_args(s3c)
-    assert kwargs == {}
+    assert kwargs == {"cache_regions": True}
 
 
 @mock.patch("flytekit.configuration.get_config_file")
@@ -191,7 +191,7 @@ def test_s3_setup_args_env_both(mock_os, mock_get_config_file):
     }
     mock_os.get.side_effect = lambda x, y: ee.get(x)
     kwargs = s3_setup_args(S3Config.auto())
-    assert kwargs == {"key": "flyte", "secret": "flyte-secret"}
+    assert kwargs == {"key": "flyte", "secret": "flyte-secret", "cache_regions": True}
 
 
 @mock.patch("flytekit.configuration.get_config_file")
@@ -204,7 +204,7 @@ def test_s3_setup_args_env_flyte(mock_os, mock_get_config_file):
     }
     mock_os.get.side_effect = lambda x, y: ee.get(x)
     kwargs = s3_setup_args(S3Config.auto())
-    assert kwargs == {"key": "flyte", "secret": "flyte-secret"}
+    assert kwargs == {"key": "flyte", "secret": "flyte-secret", "cache_regions": True}
 
 
 @mock.patch("flytekit.configuration.get_config_file")
@@ -218,7 +218,7 @@ def test_s3_setup_args_env_aws(mock_os, mock_get_config_file):
     mock_os.get.side_effect = lambda x, y: ee.get(x)
     kwargs = s3_setup_args(S3Config.auto())
     # not explicitly in kwargs, since fsspec/boto3 will use these env vars by default
-    assert kwargs == {}
+    assert kwargs == {"cache_regions": True}
 
 
 def test_crawl_local_nt(source_folder):