From 84b335aa42b62df8ac8d445d69b1a55b9e65a8b3 Mon Sep 17 00:00:00 2001 From: Soham Manoli Date: Wed, 21 Aug 2024 11:13:19 -0700 Subject: [PATCH] python/pytorch: Fix WorkerSessionManager returning None when using samplers with no workers When using a dynamic batch sampler without a dataloader (e.g to just get the batch indices), we trigger a bug where session manager returns None. Signed-off-by: Soham Manoli --- python/aistore/pytorch/worker_session_manager.py | 4 +++- python/tests/integration/pytorch/test_samplers.py | 9 +++++++++ 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/python/aistore/pytorch/worker_session_manager.py b/python/aistore/pytorch/worker_session_manager.py index 837e5062f68..7dd28e661f8 100644 --- a/python/aistore/pytorch/worker_session_manager.py +++ b/python/aistore/pytorch/worker_session_manager.py @@ -33,11 +33,13 @@ def __init__(self, session_manager: SessionManager): @property def session(self): """ - Returns: Active request session acquired for a specific Pytorch dataloader worker + Returns an active request session acquired for a specific Pytorch dataloader worker. """ # sessions are not thread safe, so we must return different sessions for each worker worker_info = get_worker_info() if worker_info is None: + if self._session is None: + self._session = self._create_session() return self._session # if we only have one session but multiple workers, create more if worker_info.id not in self._worker_sessions: diff --git a/python/tests/integration/pytorch/test_samplers.py b/python/tests/integration/pytorch/test_samplers.py index 0a7ea3a9362..cbdc18eb5db 100644 --- a/python/tests/integration/pytorch/test_samplers.py +++ b/python/tests/integration/pytorch/test_samplers.py @@ -76,3 +76,12 @@ def test_dynamic_sampler(self): # Test that all objects are included in our batch self.assertEqual(num_objects, NUM_OBJECTS) + + def test_sampler_no_workers(self): + sampler = DynamicBatchSampler( + data_source=self.dataset, + max_batch_size=MAX_BATCH_SIZE, + ) + + for _ in sampler: + continue