From 8e26a599007f5887432aceffcab42652df4f4175 Mon Sep 17 00:00:00 2001 From: Eduardo Apolinario <653394+eapolinario@users.noreply.github.com> Date: Mon, 17 Apr 2023 17:21:04 -0700 Subject: [PATCH] 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 * Fix tests Signed-off-by: eduardo apolinario * More linting Signed-off-by: eduardo apolinario --------- Signed-off-by: eduardo apolinario Co-authored-by: eduardo apolinario --- flytekit/core/data_persistence.py | 6 ++++-- tests/flytekit/unit/core/test_data.py | 8 ++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/flytekit/core/data_persistence.py b/flytekit/core/data_persistence.py index 068771afa6..2b849e8993 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,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 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):