Speed up docker launch (#2099)

* 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

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
the files in the existing huggingface cache that don't exist in the temp
huggingface cache.
"""
for item in source.iterdir():
target_path = dest / item.relative_to(source)
if item.is_dir(): if item.is_dir():
await asyncio.to_thread( _move_files_recursively(item, target_path, overwrite)
shutil.copytree, item, hf_cache / item.name, dirs_exist_ok=True
)
else: else:
await asyncio.to_thread(shutil.copy2, item, hf_cache) target_path.parent.mkdir(parents=True, exist_ok=True)
await asyncio.to_thread(shutil.rmtree, temp_hf_cache) if target_path.exists() and not overwrite:
logger.info("Copied contents of temp_huggingface and deleted the directory.") continue
else: shutil.move(str(item), str(target_path))
logger.info("Source directory is empty or does not exist. Skipping copy.")
@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()}")