Skip to content

Commit

Permalink
Set cache_regions=True in the case of s3fs (#1592)
Browse files Browse the repository at this point in the history
* Set cache_regions=True in the case of s3fs

* Linting

Signed-off-by: eduardo apolinario <eapolinario@users.noreply.github.com>

* Fix tests

Signed-off-by: eduardo apolinario <eapolinario@users.noreply.github.com>

* More linting

Signed-off-by: eduardo apolinario <eapolinario@users.noreply.github.com>

---------

Signed-off-by: eduardo apolinario <eapolinario@users.noreply.github.com>
Co-authored-by: eduardo apolinario <eapolinario@users.noreply.github.com>
Signed-off-by: Fabio Graetz <fabiograetz@googlemail.com>
  • Loading branch information
2 people authored and fg91 committed Apr 23, 2023
1 parent 538a7f3 commit 50c2e38
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 6 deletions.
6 changes: 4 additions & 2 deletions flytekit/core/data_persistence.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down
8 changes: 4 additions & 4 deletions tests/flytekit/unit/core/test_data.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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")
Expand All @@ -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")
Expand All @@ -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):
Expand Down

0 comments on commit 50c2e38

Please sign in to comment.