[Feat] Support OpenAI Files API and refactor storage library to support local, s3, tos, and redis (for metadata) (#1583)

* Basic metadata server with new batchjob work with legacy local batch driver.
* Connect job_manager with batch API
* Finish e2e batch server test.
* lint-fix
* Review fix.
* Refactored storage by add general storage adapters optimized for batch job.
Original storage in batch folder now behavior as a group of helper that bridge raw data and batch request.
Local driver (request_proxy) has more logic to work with storage adapters.
* Fix local storage errors to make e2e tests pass.
* Provide a special implementation for small file multipart upload.
* Support redis as background and add pagination to storage list.
* Conditionally enabling redis tests.
* Adjust unit tests' folder.
* Lint fix
* bug fix
* Support metadata file API, including OpenAI files compatible get file and using Head method.
* Refactor multi-upload logic to unify small and large file multi-upload
* Format
* Fix test_path_traversal_protection in test_local_storage.py
* Fix generate_filename bug.
* Lint fix

Signed-off-by: Jingyuan Zhang <[email protected]>
Co-authored-by: Jingyuan Zhang <[email protected]>

Author: zhangjyr

python/aibrix/aibrix/batch/driver.py

@@ -13,30 +13,38 @@
 # limitations under the License.
 
 import asyncio
-from typing import Optional
+from typing import Any, Dict, List, Optional
 
 import aibrix.batch.storage as _storage
 from aibrix.batch.constant import DEFAULT_JOB_POOL_SIZE
 from aibrix.batch.job_entity import JobEntityManager
 from aibrix.batch.job_manager import JobManager
 from aibrix.batch.request_proxy import RequestProxy
 from aibrix.batch.scheduler import JobScheduler
+from aibrix.batch.storage.batch_metastore import initialize_batch_metastore
 from aibrix.metadata.logger import init_logger
+from aibrix.storage import StorageType
 
 logger = init_logger(__name__)
 
 
 class BatchDriver:
-    def __init__(self, job_entity_manager: Optional[JobEntityManager] = None):
+    def __init__(
+        self,
+        job_entity_manager: Optional[JobEntityManager] = None,
+        storage_type: StorageType = StorageType.AUTO,
+        metastore_type: StorageType = StorageType.AUTO,
+    ):
         """
         This is main entrance to bind all components to serve job requests.
         """
-        _storage.initialize_storage()
+        _storage.initialize_storage(storage_type)
+        initialize_batch_metastore(metastore_type)
         self._storage = _storage
         self._job_manager: JobManager = JobManager(job_entity_manager)
         self._scheduler: Optional[JobScheduler] = None
         self._scheduling_task: Optional[asyncio.Task] = None
-        self._proxy: RequestProxy = RequestProxy(self._storage, self._job_manager)
+        self._proxy: RequestProxy = RequestProxy(self._job_manager)
         # Only create jobs_running_loop if JobEntityManager does not have its own sched
         if not job_entity_manager or not job_entity_manager.is_scheduler_enabled():
             self._scheduler = JobScheduler(self._job_manager, DEFAULT_JOB_POOL_SIZE)
@@ -47,14 +55,11 @@ def __init__(self, job_entity_manager: Optional[JobEntityManager] = None):
     def job_manager(self) -> JobManager:
         return self._job_manager
 
-    def upload_batch_data(self, input_file_name):
-        file_id = self._storage.submit_job_input(input_file_name)
-        return file_id
+    async def upload_job_data(self, input_file_name) -> str:
+        return await self._storage.upload_input_data(input_file_name)
 
-    def retrieve_job_result(self, file_id):
-        num_requests = _storage.get_job_num_request(file_id)
-        req_results = _storage.get_job_results(file_id, 0, num_requests)
-        return req_results
+    async def retrieve_job_result(self, file_id) -> List[Dict[str, Any]]:
+        return await self._storage.download_output_data(file_id)
 
     async def jobs_running_loop(self):
         """
@@ -66,7 +71,17 @@ async def jobs_running_loop(self):
         while True:
             one_job = await self._scheduler.round_robin_get_job()
             if one_job:
-                await self._proxy.execute_queries(one_job)
+                try:
+                    await self._proxy.execute_queries(one_job)
+                except Exception as e:
+                    job = self._job_manager.mark_job_failed(one_job)
+                    logger.error(
+                        "Failed to execute job",
+                        job_id=one_job,
+                        status=job.status.state.value,
+                        error=e,
+                    )
+                    raise
             await asyncio.sleep(0)
 
     async def close(self):
@@ -84,15 +99,15 @@ async def close(self):
         if self._scheduler:
             await self._scheduler.close()
 
-    def clear_job(self, job_id):
+    async def clear_job(self, job_id):
         job = self._job_manager.get_job(job_id)
         if job is None:
             return
 
         self._job_manager.job_deleted_handler(job)
         if self._job_manager.get_job(job_id) is None:
-            self._storage.delete_job(job.spec.input_file_id)
+            await self._storage.remove_job_data(job.spec.input_file_id)
             if job.status.output_file_id is not None:
-                self._storage.delete_job(job.status.output_file_id)
+                await self._storage.remove_job_data(job.status.output_file_id)
             if job.status.error_file_id is not None:
-                self._storage.delete_job(job.status.error_file_id)
+                await self._storage.remove_job_data(job.status.error_file_id)

python/aibrix/aibrix/batch/job_entity/batch_job.py

@@ -58,6 +58,14 @@ class BatchJobState(str, Enum):
     CANCELLING = "cancelling"
     CANCELED = "canceled"
 
+    def is_finished(self):
+        return self in [
+            BatchJobState.COMPLETED,
+            BatchJobState.FAILED,
+            BatchJobState.CANCELED,
+            BatchJobState.EXPIRED,
+        ]
+
 
 class BatchJobErrorCode(str, Enum):
     """Error codes for batch job."""
@@ -228,6 +236,9 @@ class RequestCountStats(NoExtraBaseModel):
     """Holds the statistics on the processing of the batch."""
 
     total: int = Field(default=0, description="Total number of requests in the batch")
+    launched: int = Field(
+        default=0, description="Number of requests that have been launched"
+    )
     completed: int = Field(
         default=0,
         description="Number of requests that have been successfully completed",
@@ -299,6 +310,17 @@ class BatchJobStatus(NoExtraBaseModel):
         description="List of errors that occurred during the batch job processing",
     )
 
+    temp_output_file_id: Optional[str] = Field(
+        default=None,
+        alias="tempOutputFileID",
+        description="The ID of the file containing the results of successfully completed requests",
+    )
+    temp_error_file_id: Optional[str] = Field(
+        default=None,
+        alias="tempErrorFileID",
+        description="The ID of the file containing details for any failed requests",
+    )
+
     output_file_id: Optional[str] = Field(
         default=None,
         alias="outputFileID",
@@ -310,8 +332,8 @@ class BatchJobStatus(NoExtraBaseModel):
         description="The ID of the file containing details for any failed requests",
     )
 
-    request_counts: Optional[RequestCountStats] = Field(
-        default=None,
+    request_counts: RequestCountStats = Field(
+        default_factory=RequestCountStats,
         alias="requestCounts",
         description="Statistics on the processing of the batch",
     )