diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index dd843fd85..37c16209c 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -268,6 +268,8 @@ class Envs: # Hi-Cache SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None) + SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None) + SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None) # Mooncake KV Transfer SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index 38df15262..daa79f93f 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -7,6 +7,7 @@ from typing import Any, List, Optional import torch +from sglang.srt.environ import envs from sglang.srt.mem_cache.memory_pool_host import HostKVCache logger = logging.getLogger(__name__) @@ -186,7 +187,7 @@ class HiCacheFile(HiCacheStorage): def __init__( self, storage_config: HiCacheStorageConfig, file_path: str = "/tmp/hicache" ): - self.file_path = os.getenv("SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR", file_path) + self.file_path = envs.SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR.get() or file_path tp_rank, tp_size, model_name, is_mla_model = ( storage_config.tp_rank, diff --git a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py index 854ead41f..abb26c4ed 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py +++ b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py @@ -1,11 +1,11 @@ import logging -import os import time import uuid from typing import Any, List, Optional, Union import torch +from sglang.srt.environ import envs from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig from .nixl_utils import ( @@ -44,7 +44,7 @@ class HiCacheNixl(HiCacheStorage): plugin = nixlconfig.get_specified_plugin() # Might be better to be unified across HiCache backends and moved to HiCacheController - file_path = os.getenv("SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR", file_path) + file_path = envs.SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR.get() or file_path self.file_manager = ( NixlFileManager(file_path) if plugin not in NixlBackendSelection.OBJ_PLUGINS