diff --git a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py index 1971ee587..c6f634a08 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -289,7 +289,7 @@ class OpenAIVectorStoreMixin(ABC): self.openai_vector_stores = await self._load_openai_vector_stores() self.openai_file_batches = await self._load_openai_vector_store_file_batches() self._file_batch_tasks = {} - # TODO: Enable resume for multi-worker deployments, only works for single worker for now + # TODO: Resume only works for single worker deployment. Jobs with multiple workers will need to be handled differently. await self._resume_incomplete_batches() self._last_file_batch_cleanup_time = 0 diff --git a/tests/integration/vector_io/test_openai_vector_stores.py b/tests/integration/vector_io/test_openai_vector_stores.py index 572642733..8321c0dec 100644 --- a/tests/integration/vector_io/test_openai_vector_stores.py +++ b/tests/integration/vector_io/test_openai_vector_stores.py @@ -1049,10 +1049,12 @@ def test_openai_vector_store_file_batch_cancel(compat_client_with_empty_stores, # Create a vector store vector_store = compat_client.vector_stores.create(name="batch_cancel_test_store") - # Create multiple files + # Create a larger batch to ensure we have time to cancel before completion file_ids = [] - for i in range(3): - with BytesIO(f"This is batch cancel test file {i}".encode()) as file_buffer: + for i in range(50): # Large batch to slow down processing + with BytesIO( + f"This is batch cancel test file {i} with more content to process slowly".encode() * 100 + ) as file_buffer: file_buffer.name = f"batch_cancel_test_{i}.txt" file = compat_client.files.create(file=file_buffer, purpose="assistants") file_ids.append(file.id) @@ -1062,7 +1064,8 @@ def test_openai_vector_store_file_batch_cancel(compat_client_with_empty_stores, vector_store_id=vector_store.id, file_ids=file_ids, ) - # Cancel the batch immediately after creation (before processing can complete) + + # Cancel the batch immediately after creation (large batch gives us time) cancelled_batch = compat_client.vector_stores.file_batches.cancel( vector_store_id=vector_store.id, batch_id=batch.id,