feat: server side prefetch (#41)

* feat: serversice prefetch (wip)

* fix: pass test

* fix: prefetcher thread cleanup

* feat: use 202

* fix: log

Author: hanchchch

Makefile

@@ -5,9 +5,7 @@ all: openapi-python-client ui
 openapi-python-client:
 	openapi-python-client generate --url http://localhost:8000/openapi.json --config openapi.config.yaml --overwrite --meta poetry
 	sed -i -e 's/response_200 = File(payload=BytesIO(response.json()))/response_200 = File(payload=BytesIO(response.content))/g' openapi-lavender-data-rest/openapi_lavender_data_rest/api/iterations/get_next_iterations_iteration_id_next_get.py
-	sed -i -e 's/response_200 = File(payload=BytesIO(response.json()))/response_200 = File(payload=BytesIO(response.content))/g' openapi-lavender-data-rest/openapi_lavender_data_rest/api/iterations/get_submitted_result_iterations_iteration_id_next_cache_key_get.py
 	rm openapi-lavender-data-rest/openapi_lavender_data_rest/api/iterations/get_next_iterations_iteration_id_next_get.py-e 2> /dev/null
-	rm openapi-lavender-data-rest/openapi_lavender_data_rest/api/iterations/get_submitted_result_iterations_iteration_id_next_cache_key_get.py-e 2> /dev/null
 
 ui:
 	cd ./ui && pnpm build && cd ../

lavender_data/client/api.py

@@ -25,8 +25,6 @@
 from openapi_lavender_data_rest.api.iterations import (
     create_iteration_iterations_post,
     get_next_iterations_iteration_id_next_get,
-    submit_next_iterations_iteration_id_next_post,
-    get_submitted_result_iterations_iteration_id_next_cache_key_get,
     get_iteration_iterations_iteration_id_get,
     get_iterations_iterations_get,
     complete_index_iterations_iteration_id_complete_index_post,
@@ -304,10 +302,15 @@ def create_iteration(
         categorizer: Optional[IterationCategorizer] = None,
         collater: Optional[IterationCollater] = None,
         preprocessors: Optional[list[IterationPreprocessor]] = None,
+        max_retry_count: int = 0,
         rank: int = 0,
         world_size: Optional[int] = None,
         wait_participant_threshold: Optional[float] = None,
-        cluster_sync: bool = False,
+        no_cache: Optional[bool] = None,
+        num_workers: Optional[int] = None,
+        prefetch_factor: Optional[int] = None,
+        in_order: Optional[bool] = None,
+        cluster_sync: Optional[bool] = None,
     ):
         with self._get_client() as client:
             response = create_iteration_iterations_post.sync_detailed(
@@ -324,9 +327,14 @@ def create_iteration(
                     collater=collater,
                     preprocessors=preprocessors,
                     replication_pg=replication_pg,
+                    max_retry_count=max_retry_count,
                     rank=rank,
                     world_size=world_size,
                     wait_participant_threshold=wait_participant_threshold,
+                    no_cache=no_cache,
+                    num_workers=num_workers,
+                    prefetch_factor=prefetch_factor,
+                    in_order=in_order,
                     cluster_sync=cluster_sync,
                 ),
             )
@@ -366,61 +374,22 @@ def get_next_item(
         self,
         iteration_id: str,
         rank: int = 0,
-        no_cache: bool = False,
-        max_retry_count: int = 0,
         client: Optional[Client] = None,
     ):
         with self._get_client() if client is None else nullcontext() as _client:
             response = get_next_iterations_iteration_id_next_get.sync_detailed(
                 client=client or _client,
                 iteration_id=iteration_id,
                 rank=rank,
-                no_cache=no_cache,
-                max_retry_count=max_retry_count,
             )
-
         try:
             current = int(response.headers.get("X-Lavender-Data-Sample-Current"))
         except TypeError:
             current = None
-        return self._check_response(response).payload.read(), current
-
-    def submit_next_item(
-        self,
-        iteration_id: str,
-        rank: int = 0,
-        no_cache: bool = False,
-        max_retry_count: int = 0,
-        client: Optional[Client] = None,
-    ):
-        with self._get_client() if client is None else nullcontext() as _client:
-            response = submit_next_iterations_iteration_id_next_post.sync_detailed(
-                client=client or _client,
-                iteration_id=iteration_id,
-                rank=rank,
-                no_cache=no_cache,
-                max_retry_count=max_retry_count,
-            )
-        return self._check_response(response)
 
-    def get_submitted_result(
-        self,
-        iteration_id: str,
-        cache_key: str,
-        client: Optional[Client] = None,
-    ):
-        with self._get_client() if client is None else nullcontext() as _client:
-            response = get_submitted_result_iterations_iteration_id_next_cache_key_get.sync_detailed(
-                client=client or _client,
-                iteration_id=iteration_id,
-                cache_key=cache_key,
-            )
         if response.status_code == 202:
             raise LavenderDataApiError(response.content.decode("utf-8"))
-        try:
-            current = int(response.headers.get("X-Lavender-Data-Sample-Current"))
-        except TypeError:
-            current = None
+
         return self._check_response(response).payload.read(), current
 
     def complete_index(self, iteration_id: str, index: int):
@@ -602,10 +571,15 @@ def create_iteration(
     categorizer: Optional[IterationCategorizer] = None,
     collater: Optional[IterationCollater] = None,
     preprocessors: Optional[list[IterationPreprocessor]] = None,
+    max_retry_count: int = 0,
     rank: int = 0,
     world_size: Optional[int] = None,
     wait_participant_threshold: Optional[float] = None,
-    cluster_sync: bool = False,
+    no_cache: Optional[bool] = None,
+    num_workers: Optional[int] = None,
+    prefetch_factor: Optional[int] = None,
+    in_order: Optional[bool] = None,
+    cluster_sync: Optional[bool] = None,
 ):
     return _client_instance.create_iteration(
         dataset_id=dataset_id,
@@ -619,9 +593,14 @@ def create_iteration(
         categorizer=categorizer,
         collater=collater,
         preprocessors=preprocessors,
+        max_retry_count=max_retry_count,
         rank=rank,
         world_size=world_size,
         wait_participant_threshold=wait_participant_threshold,
+        no_cache=no_cache,
+        num_workers=num_workers,
+        prefetch_factor=prefetch_factor,
+        in_order=in_order,
         cluster_sync=cluster_sync,
     )
 
@@ -644,36 +623,10 @@ def get_iteration(iteration_id: str):
 def get_next_item(
     iteration_id: str,
     rank: int = 0,
-    no_cache: bool = False,
-    max_retry_count: int = 0,
 ):
     return _client_instance.get_next_item(
         iteration_id=iteration_id,
         rank=rank,
-        no_cache=no_cache,
-        max_retry_count=max_retry_count,
-    )
-
-
-@ensure_client()
-def submit_next_item(
-    iteration_id: str,
-    rank: int = 0,
-    no_cache: bool = False,
-    max_retry_count: int = 0,
-):
-    return _client_instance.submit_next_item(
-        iteration_id=iteration_id,
-        rank=rank,
-        no_cache=no_cache,
-        max_retry_count=max_retry_count,
-    )
-
-
-@ensure_client()
-def get_submitted_result(iteration_id: str, cache_key: str):
-    return _client_instance.get_submitted_result(
-        iteration_id=iteration_id, cache_key=cache_key
     )
 
 

lavender_data/client/cli/__init__.py

@@ -15,8 +15,6 @@
     get_iterations,
     get_iteration,
     get_next_item,
-    submit_next_item,
-    get_submitted_result,
     complete_index,
     pushback,
     get_progress,
@@ -92,22 +90,6 @@ def __init__(self, parent_parser: Optional[argparse.ArgumentParser] = None):
         self.iterations_next.add_argument("--no-cache", action="store_true")
         self.iterations_next.add_argument("--max-retry-count", type=int, default=0)
 
-        self.iterations_submit_next_item = self.iterations_command_parser.add_parser(
-            "async-next"
-        )
-        self.iterations_submit_next_item.add_argument("id", type=str)
-        self.iterations_submit_next_item.add_argument("--rank", type=int, default=0)
-        self.iterations_submit_next_item.add_argument("--no-cache", action="store_true")
-        self.iterations_submit_next_item.add_argument(
-            "--max-retry-count", type=int, default=0
-        )
-
-        self.iterations_async_result = self.iterations_command_parser.add_parser(
-            "async-result"
-        )
-        self.iterations_async_result.add_argument("id", type=str)
-        self.iterations_async_result.add_argument("key", type=str)
-
         self.iterations_complete_index = self.iterations_command_parser.add_parser(
             "complete-index"
         )
@@ -194,21 +176,6 @@ def main(self, args: Optional[argparse.Namespace] = None):
                     args.api_key,
                     args.id,
                     args.rank,
-                    args.no_cache,
-                    args.max_retry_count,
-                )
-            elif args.command == "async-next":
-                result = submit_next_item(
-                    args.api_url,
-                    args.api_key,
-                    args.id,
-                    args.rank,
-                    args.no_cache,
-                    args.max_retry_count,
-                )
-            elif args.command == "async-result":
-                result = get_submitted_result(
-                    args.api_url, args.api_key, args.id, args.key
                 )
             elif args.command == "complete-index":
                 result = complete_index(args.api_url, args.api_key, args.id, args.index)

lavender_data/client/cli/api_call.py

@@ -75,39 +75,11 @@ def get_next_item(
     api_key: str,
     iteration_id: str,
     rank: int,
-    no_cache: bool,
-    max_retry_count: int,
 ):
     return deserialize_sample(
         _api(api_url=api_url, api_key=api_key).get_next_item(
             iteration_id=iteration_id,
             rank=rank,
-            no_cache=no_cache,
-            max_retry_count=max_retry_count,
-        )[0]
-    )
-
-
-def submit_next_item(
-    api_url: str,
-    api_key: str,
-    iteration_id: str,
-    rank: int,
-    no_cache: bool,
-    max_retry_count: int,
-):
-    return _api(api_url=api_url, api_key=api_key).submit_next_item(
-        iteration_id=iteration_id,
-        rank=rank,
-        no_cache=no_cache,
-        max_retry_count=max_retry_count,
-    )
-
-
-def get_submitted_result(api_url: str, api_key: str, iteration_id: str, cache_key: str):
-    return deserialize_sample(
-        _api(api_url=api_url, api_key=api_key).get_submitted_result(
-            iteration_id=iteration_id, cache_key=cache_key
         )[0]
     )