From 900d924a7fe2cfb7ecb3a0fe294777d98f684347 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Sat, 27 Jun 2026 11:45:57 +0200 Subject: [PATCH] Fix datastore cache download retries Signed-off-by: Minh Vu --- nemo/utils/data_utils.py | 10 ++++++---- tests/utils/test_utils.py | 24 ++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/nemo/utils/data_utils.py b/nemo/utils/data_utils.py index 56c27b556347..67d387ffbb86 100644 --- a/nemo/utils/data_utils.py +++ b/nemo/utils/data_utils.py @@ -234,12 +234,14 @@ def open_datastore_object_with_binary(path: str, num_retries: int = 5): return None -def open_best(path: str, mode: str = "rb"): +def open_best(path: str, mode: str = "rb", num_retries: int = 5): """Open a file using the best available method (Lhotse, datastore binary, or standard open). Args: path: path to the file or datastore object mode: file opening mode (default: "rb") + num_retries: number of retries if the get command fails with ais binary, + as AIS Python SDK has its own retry mechanism Returns: File-like object @@ -247,7 +249,7 @@ def open_best(path: str, mode: str = "rb"): if LHOTSE_AVAILABLE: return lhotse_open_best(path, mode=mode) if is_datastore_path(path): - return open_datastore_object_with_binary(path) + return open_datastore_object_with_binary(path, num_retries=num_retries) return open(path, mode=mode, encoding='utf-8' if 'b' not in mode else None) @@ -276,8 +278,8 @@ def get_datastore_object(path: str, force: bool = False, num_retries: int = 5) - if not os.path.isdir(local_dir): os.makedirs(local_dir, exist_ok=True) - with open(local_path, 'wb') as f: - f.write(open_best(path).read(), num_retries=num_retries) + with open(local_path, 'wb') as f, open_best(path, num_retries=num_retries) as stream: + f.write(stream.read()) return local_path diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index a012fd18e763..5ad214858c29 100644 --- a/tests/utils/test_utils.py +++ b/tests/utils/test_utils.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import io import os from unittest import mock @@ -22,6 +23,7 @@ ais_binary, ais_endpoint_to_dir, bucket_and_object_from_uri, + get_datastore_object, is_datastore_path, resolve_cache_dir, ) @@ -90,3 +92,25 @@ def test_ais_binary(self): with mock.patch('shutil.which', lambda x: None), mock.patch('os.path.isfile', lambda x: None): ais_binary.cache_clear() assert ais_binary() is None + + @pytest.mark.unit + def test_get_datastore_object_passes_num_retries_to_datastore_open_on_cache_miss(self, tmp_path): + """Test datastore cache misses preserve caller retries when downloading through the AIS fallback.""" + local_path = tmp_path / 'cache' / 'object.bin' + opened_with = {} + + def fake_open_datastore_object_with_binary(path, num_retries=5): + opened_with['path'] = path + opened_with['num_retries'] = num_retries + return io.BytesIO(b'payload') + + with ( + mock.patch('nemo.utils.data_utils.LHOTSE_AVAILABLE', False), + mock.patch('nemo.utils.data_utils.open_datastore_object_with_binary', side_effect=fake_open_datastore_object_with_binary), + mock.patch('nemo.utils.data_utils.datastore_path_to_local_path', return_value=str(local_path)), + ): + resolved_path = get_datastore_object('ais://bucket/object', num_retries=7) + + assert resolved_path == str(local_path) + assert local_path.read_bytes() == b'payload' + assert opened_with == {'path': 'ais://bucket/object', 'num_retries': 7}