From 5d13e3fd495595783f5dc3bd6f9f0870186fd14e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Sat, 2 May 2026 11:07:44 +0000 Subject: [PATCH] fix(batch): finalize jobs with existing temp uploads Co-authored-by: Varun Gupta --- python/aibrix/aibrix/batch/job_driver.py | 5 +-- python/aibrix/tests/batch/test_driver.py | 51 +++++++++++++++++++++++- 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/python/aibrix/aibrix/batch/job_driver.py b/python/aibrix/aibrix/batch/job_driver.py index 81fda9a95..395ccdcbd 100644 --- a/python/aibrix/aibrix/batch/job_driver.py +++ b/python/aibrix/aibrix/batch/job_driver.py @@ -193,7 +193,7 @@ async def execute_job(self, job_id): logger.warning("Job not found", job_id=job_id) return - # Check if temp file IDs exist to determine if we should skip steps 1 and 3 + # Check if temp file IDs exist to determine if we should skip preparation. has_temp_files = ( job.status.temp_output_file_id and job.status.temp_error_file_id ) @@ -228,8 +228,7 @@ async def execute_job(self, job_id): # Step 3: Aggregate outputs if job.status.state == BatchJobState.FINALIZING: - if not has_temp_files: - await storage.finalize_job_output_data(job) + await storage.finalize_job_output_data(job) logger.debug("Completed job", job_id=job_id) job = await self._sync_job_status(job_id) diff --git a/python/aibrix/tests/batch/test_driver.py b/python/aibrix/tests/batch/test_driver.py index 9d7f80dcd..b65525c7e 100644 --- a/python/aibrix/tests/batch/test_driver.py +++ b/python/aibrix/tests/batch/test_driver.py @@ -17,13 +17,20 @@ import os import tempfile from pathlib import Path +from unittest.mock import AsyncMock, patch import pytest import aibrix.batch.constant as constant from aibrix.batch.driver import BatchDriver -from aibrix.batch.job_driver import EchoInferenceEngineClient -from aibrix.batch.job_entity import BatchJobErrorCode, BatchJobState, BatchJobStatus +from aibrix.batch.job_driver import EchoInferenceEngineClient, JobDriver +from aibrix.batch.job_entity import ( + BatchJob, + BatchJobErrorCode, + BatchJobSpec, + BatchJobState, + BatchJobStatus, +) from aibrix.storage import StorageType constant.EXPIRE_INTERVAL = 0.1 @@ -280,6 +287,46 @@ async def test_batch_driver_resuming(): Path(temp_path).unlink(missing_ok=True) +@pytest.mark.asyncio +async def test_execute_job_finalizes_job_with_existing_temp_files(): + job = BatchJob.new_local( + BatchJobSpec.from_strings( + input_file_id="input-file-id", + endpoint="/v1/chat/completions", + ), + request_count=1, + ) + job.status.state = BatchJobState.IN_PROGRESS + job.status.output_file_id = "output-file-id" + job.status.error_file_id = "error-file-id" + job.status.temp_output_file_id = "temp-output-file-id" + job.status.temp_error_file_id = "temp-error-file-id" + + progress_manager = AsyncMock() + progress_manager.get_job.return_value = job + progress_manager.mark_job_done.return_value = job + + async def finish_job(_job_id): + job.status.state = BatchJobState.FINALIZING + return job + + job_driver = JobDriver(progress_manager, EchoInferenceEngineClient()) + with patch.object( + job_driver, "execute_worker", new=AsyncMock(side_effect=finish_job) + ), patch( + "aibrix.batch.job_driver.storage.prepare_job_ouput_files", + new_callable=AsyncMock, + ) as prepare_job_output_files, patch( + "aibrix.batch.job_driver.storage.finalize_job_output_data", + new_callable=AsyncMock, + ) as finalize_job_output_data: + await job_driver.execute_job(job.status.job_id) + + prepare_job_output_files.assert_not_awaited() + finalize_job_output_data.assert_awaited_once_with(job) + progress_manager.mark_job_done.assert_awaited_once_with(job.status.job_id) + + @pytest.mark.asyncio async def test_batch_driver_validation_failed() -> None: """