feat: set max memory avaiable while download

vllm-project · Dec 9, 2024 · f61c845 · f61c845
1 parent 5c2c220
commit f61c845
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/python/aibrix/aibrix/downloader/s3.py b/python/aibrix/aibrix/downloader/s3.py
@@ -155,6 +155,8 @@ def download(
         config_kwargs = {
             "max_concurrency": envs.DOWNLOADER_NUM_THREADS,
             "use_threads": enable_range,
+            "max_io_queue": envs.DOWNLOADER_S3_MAX_IO_QUEUE,
+            "io_chunksize": envs.DOWNLOADER_S3_IO_CHUNKSIZE,
         }
         if envs.DOWNLOADER_PART_THRESHOLD is not None:
             config_kwargs["multipart_threshold"] = envs.DOWNLOADER_PART_THRESHOLD

diff --git a/python/aibrix/aibrix/envs.py b/python/aibrix/aibrix/envs.py
@@ -48,6 +48,12 @@ def _parse_int_or_none(value: Optional[str]) -> Optional[int]:
 DOWNLOADER_ALLOW_FILE_SUFFIX = _parse_list_str(
     os.getenv("DOWNLOADER_ALLOW_FILE_SUFFIX")
 )
+DOWNLOADER_S3_MAX_IO_QUEUE = _parse_int_or_none(
+    os.getenv("DOWNLOADER_S3_MAX_IO_QUEUE", "100")
+)
+DOWNLOADER_S3_IO_CHUNKSIZE = _parse_int_or_none(
+    os.getenv("DOWNLOADER_S3_IO_CHUNKSIZE", "16777216")
+)
 
 DOWNLOADER_FORCE_DOWNLOAD = _is_true(os.getenv("DOWNLOADER_FORCE_DOWNLOAD", "0"))
 DOWNLOADER_CHECK_FILE_EXIST = _is_true(os.getenv("DOWNLOADER_CHECK_FILE_EXIST", "1"))