wehub-resource-sync
91 lines
3.1 KiB
Python
91 lines
3.1 KiB
Python
"""Environment variables used by the MLC LLM."""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import List # noqa: UP035
|
|
|
|
MLC_CHAT_CONFIG_VERSION = "0.1.0"
|
|
|
|
|
|
def _check():
|
|
if MLC_JIT_POLICY not in ["ON", "OFF", "REDO", "READONLY"]:
|
|
raise ValueError(
|
|
'Invalid MLC_JIT_POLICY. It has to be one of "ON", "OFF", "REDO", "READONLY"'
|
|
f"but got {MLC_JIT_POLICY}."
|
|
)
|
|
|
|
if MLC_DOWNLOAD_CACHE_POLICY not in ["ON", "OFF", "REDO", "READONLY"]:
|
|
raise ValueError(
|
|
"Invalid MLC_AUTO_DOWNLOAD_POLICY. "
|
|
'It has to be one of "ON", "OFF", "REDO", "READONLY"'
|
|
f"but got {MLC_DOWNLOAD_CACHE_POLICY}."
|
|
)
|
|
|
|
|
|
def _get_cache_dir() -> Path:
|
|
if "MLC_LLM_HOME" in os.environ:
|
|
result = Path(os.environ["MLC_LLM_HOME"])
|
|
elif sys.platform == "win32":
|
|
result = Path(os.environ["LOCALAPPDATA"])
|
|
result = result / "mlc_llm"
|
|
elif os.getenv("XDG_CACHE_HOME", None) is not None:
|
|
result = Path(os.getenv("XDG_CACHE_HOME"))
|
|
result = result / "mlc_llm"
|
|
else:
|
|
result = Path(os.path.expanduser("~/.cache"))
|
|
result = result / "mlc_llm"
|
|
result.mkdir(parents=True, exist_ok=True)
|
|
if not result.is_dir():
|
|
raise ValueError(
|
|
f"The default cache directory is not a directory: {result}. "
|
|
"Use environment variable MLC_LLM_HOME to specify a valid cache directory."
|
|
)
|
|
(result / "model_weights").mkdir(parents=True, exist_ok=True)
|
|
(result / "model_lib").mkdir(parents=True, exist_ok=True)
|
|
return result
|
|
|
|
|
|
def _get_dso_suffix() -> str:
|
|
if "MLC_DSO_SUFFIX" in os.environ:
|
|
return os.environ["MLC_DSO_SUFFIX"]
|
|
if sys.platform == "win32":
|
|
return "dll"
|
|
if sys.platform == "darwin":
|
|
return "dylib"
|
|
return "so"
|
|
|
|
|
|
def _get_test_model_path() -> List[Path]: # noqa: UP006
|
|
paths = []
|
|
if "MLC_LLM_TEST_MODEL_PATH" in os.environ:
|
|
paths += [Path(p) for p in os.environ["MLC_LLM_TEST_MODEL_PATH"].split(os.pathsep)]
|
|
# by default, we reuse the cache dir via mlc_llm chat
|
|
# note that we do not auto download for testcase
|
|
# to avoid networking dependencies
|
|
base_list = ["hf"]
|
|
paths += [_get_cache_dir() / "model_weights" / base / "mlc-ai" for base in base_list] + [
|
|
Path(os.path.abspath(os.path.curdir)),
|
|
Path(os.path.abspath(os.path.curdir)) / "dist",
|
|
]
|
|
return paths
|
|
|
|
|
|
def _get_read_only_weight_caches() -> List[Path]: # noqa: UP006
|
|
if "MLC_LLM_READONLY_WEIGHT_CACHE" in os.environ:
|
|
return [Path(p) for p in os.environ["MLC_LLM_READONLY_WEIGHT_CACHE"].split(os.pathsep)]
|
|
return []
|
|
|
|
|
|
MLC_TEMP_DIR = os.getenv("MLC_TEMP_DIR", None)
|
|
MLC_MULTI_ARCH = os.environ.get("MLC_MULTI_ARCH", None)
|
|
MLC_JIT_POLICY = os.environ.get("MLC_JIT_POLICY", "ON")
|
|
MLC_DSO_SUFFIX = _get_dso_suffix()
|
|
MLC_TEST_MODEL_PATH: List[Path] = _get_test_model_path() # noqa: UP006
|
|
|
|
MLC_DOWNLOAD_CACHE_POLICY = os.environ.get("MLC_DOWNLOAD_CACHE_POLICY", "ON")
|
|
MLC_LLM_HOME: Path = _get_cache_dir()
|
|
MLC_LLM_READONLY_WEIGHT_CACHE = _get_read_only_weight_caches()
|
|
|
|
_check()
|