highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-10-09 20:55:06 +02:00
Code Issues Packages Projects Releases Wiki Activity

Speed up docker launch (#2099)

Browse Source 
* use move instead of copy

* added logging

* fix overwrites

* tested throughly

* fixes

* clearer commenting
This commit is contained in:
hagen-danswer
2024-08-12 17:45:05 -07:00
committed by GitHub
parent 83e945ba57
commit bbb8c5ff0b
1 changed files with 26 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
backend/model_server/main.py
View File

@@ -1,4 +1,3 @@
import asyncio
import os import os
import shutil import shutil
from collections.abc import AsyncGenerator from collections.abc import AsyncGenerator
@@ -24,27 +23,32 @@ from shared_configs.configs import MODEL_SERVER_PORT
os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1" os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
HF_CACHE_PATH = Path("/root/.cache/huggingface/")
TEMP_HF_CACHE_PATH = Path("/root/.cache/temp_huggingface/")
transformer_logging.set_verbosity_error() transformer_logging.set_verbosity_error()
logger = setup_logger() logger = setup_logger()
async def manage_huggingface_cache() -> None: def _move_files_recursively(source: Path, dest: Path, overwrite: bool = False) -> None:
temp_hf_cache = Path("/root/.cache/temp_huggingface") """
hf_cache = Path("/root/.cache/huggingface") This moves the files from the temp huggingface cache to the huggingface cache
if temp_hf_cache.is_dir() and any(temp_hf_cache.iterdir()):
hf_cache.mkdir(parents=True, exist_ok=True) We have to move each file individually because the directories might
for item in temp_hf_cache.iterdir(): have the same name but not the same contents and we dont want to remove
if item.is_dir(): the files in the existing huggingface cache that don't exist in the temp
await asyncio.to_thread( huggingface cache.
shutil.copytree, item, hf_cache / item.name, dirs_exist_ok=True """
) for item in source.iterdir():
else: target_path = dest / item.relative_to(source)
await asyncio.to_thread(shutil.copy2, item, hf_cache) if item.is_dir():
await asyncio.to_thread(shutil.rmtree, temp_hf_cache) _move_files_recursively(item, target_path, overwrite)
logger.info("Copied contents of temp_huggingface and deleted the directory.") else:
else: target_path.parent.mkdir(parents=True, exist_ok=True)
logger.info("Source directory is empty or does not exist. Skipping copy.") if target_path.exists() and not overwrite:
continue
shutil.move(str(item), str(target_path))
@asynccontextmanager @asynccontextmanager
@@ -54,7 +58,11 @@ async def lifespan(app: FastAPI) -> AsyncGenerator:
else: else:
logger.info("GPU is not available") logger.info("GPU is not available")
await manage_huggingface_cache() if TEMP_HF_CACHE_PATH.is_dir():
logger.info("Moving contents of temp_huggingface to huggingface cache.")
_move_files_recursively(TEMP_HF_CACHE_PATH, HF_CACHE_PATH)
shutil.rmtree(TEMP_HF_CACHE_PATH, ignore_errors=True)
logger.info("Moved contents of temp_huggingface to huggingface cache.")
torch.set_num_threads(max(MIN_THREADS_ML_MODELS, torch.get_num_threads())) torch.set_num_threads(max(MIN_THREADS_ML_MODELS, torch.get_num_threads()))
logger.info(f"Torch Threads: {torch.get_num_threads()}") logger.info(f"Torch Threads: {torch.get_num_threads()}")