From 16b6b9655dff047de450d774c9c14f0e6f68390a Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Fri, 14 Apr 2023 14:55:52 -0700 Subject: [PATCH 1/4] Set cache_regions=True in the case of s3fs --- flytekit/core/data_persistence.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index ea36689874..344ac92b14 100644 --- a/flytekit/core/data_persistence.py +++ b/flytekit/core/data_persistence.py @@ -46,7 +46,9 @@ def s3_setup_args(s3_cfg: configuration.S3Config, anonymous: bool = False): - kwargs = {} + kwargs = { + "cache_regions": True, + } if s3_cfg.access_key_id: kwargs[_FSSPEC_S3_KEY_ID] = s3_cfg.access_key_id From 84440559f1c024eea5f859090b428d5246dc36c9 Mon Sep 17 00:00:00 2001 From: eduardo apolinario Date: Fri, 14 Apr 2023 15:42:18 -0700 Subject: [PATCH 2/4] Linting Signed-off-by: eduardo apolinario --- flytekit/core/data_persistence.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index 344ac92b14..7e5201560a 100644 --- a/flytekit/core/data_persistence.py +++ b/flytekit/core/data_persistence.py @@ -47,7 +47,7 @@ def s3_setup_args(s3_cfg: configuration.S3Config, anonymous: bool = False): kwargs = { - "cache_regions": True, + "cache_regions": True, } if s3_cfg.access_key_id: kwargs[_FSSPEC_S3_KEY_ID] = s3_cfg.access_key_id From ee22e670e56422b8b00769eb7256eed656dbe6c6 Mon Sep 17 00:00:00 2001 From: eduardo apolinario Date: Fri, 14 Apr 2023 16:15:54 -0700 Subject: [PATCH 3/4] Fix tests Signed-off-by: eduardo apolinario --- tests/flytekit/unit/core/test_data.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/flytekit/unit/core/test_data.py b/tests/flytekit/unit/core/test_data.py index 1b33ad2923..2d61b58d8c 100644 --- 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): From 88f26de18824176dcef29afc0f0e0ddd3918dde8 Mon Sep 17 00:00:00 2001 From: eduardo apolinario Date: Mon, 17 Apr 2023 10:52:42 -0700 Subject: [PATCH 4/4] More linting Signed-off-by: eduardo apolinario --- flytekit/core/data_persistence.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index 7e5201560a..84effdf5a1 100644 --- 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,7 @@ 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: