phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-04 20:14:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

improve cancel test

Browse source
This commit is contained in:
Swapna Lekkala 2025-10-03 15:12:41 -07:00
parent 510ace263b
commit 878f679e4a
2 changed files with 8 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
llama_stack/providers/utils/memory/openai_vector_store_mixin.py
View file

@ -289,7 +289,7 @@ class OpenAIVectorStoreMixin(ABC):
self.openai_vector_stores = await self._load_openai_vector_stores() self.openai_vector_stores = await self._load_openai_vector_stores()
self.openai_file_batches = await self._load_openai_vector_store_file_batches() self.openai_file_batches = await self._load_openai_vector_store_file_batches()
self._file_batch_tasks = {} 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() await self._resume_incomplete_batches()
self._last_file_batch_cleanup_time = 0 self._last_file_batch_cleanup_time = 0

11
tests/integration/vector_io/test_openai_vector_stores.py
View file

@ -1049,10 +1049,12 @@ def test_openai_vector_store_file_batch_cancel(compat_client_with_empty_stores,
# Create a vector store # Create a vector store
vector_store = compat_client.vector_stores.create(name="batch_cancel_test_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 = [] file_ids = []
for i in range(3): for i in range(50): # Large batch to slow down processing
with BytesIO(f"This is batch cancel test file {i}".encode()) as file_buffer: 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_buffer.name = f"batch_cancel_test_{i}.txt"
file = compat_client.files.create(file=file_buffer, purpose="assistants") file = compat_client.files.create(file=file_buffer, purpose="assistants")
file_ids.append(file.id) 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, vector_store_id=vector_store.id,
file_ids=file_ids, 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( cancelled_batch = compat_client.vector_stores.file_batches.cancel(
vector_store_id=vector_store.id, vector_store_id=vector_store.id,
batch_id=batch.id, batch_id=batch.id,