Repository: vllm-project/aibrix
File: python/aibrix_kvcache/aibrix_kvcache/envs.py (Python)

Code before:
    AIBRIX_KV_CACHE_OL_L2_CACHE_NUM_ASYNC_WORKERS: int = 8

    # Mock Connector
    AIBRIX_KV_CACHE_OL_MOCK_USE_RDMA: bool = False
    AIBRIX_KV_CACHE_OL_MOCK_USE_MPUT_MGET: bool = False

    # RocksDB Env Vars
    AIBRIX_KV_CACHE_OL_ROCKSDB_ROOT: str = os.path.expanduser(
        os.path.join(os.path.expanduser("~"), ".kv_cache_ol", "rocksdb")
    )
    AIBRIX_KV_CACHE_OL_ROCKSDB_TTL_S: int = 600
    AIBRIX_KV_CACHE_OL_ROCKSDB_WRITE_BUFFER_SIZE: int = 64 * 1024 * 1024
    AIBRIX_KV_CACHE_OL_ROCKSDB_TARGET_FILE_SIZE_BASE: int = 64 * 1024 * 1024
    AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_WRITE_BUFFER_NUMBER: int = 3
    AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_TOTAL_WAL_SIZE: int = 128 * 1024 * 1024
    AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_BACKGROUND_JOBS: int = 8

    # InfiniStore Env Vars
    AIBRIX_KV_CACHE_OL_INFINISTORE_HOST_ADDR: str = "127.0.0.1"
    AIBRIX_KV_CACHE_OL_INFINISTORE_SERVICE_PORT: int = 12345
    AIBRIX_KV_CACHE_OL_INFINISTORE_CONNECTION_TYPE: str = "RDMA"
    AIBRIX_KV_CACHE_OL_INFINISTORE_IB_PORT: int = 1
    AIBRIX_KV_CACHE_OL_INFINISTORE_LINK_TYPE: str = "Ethernet"
    AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST: List[str] = ["mlx5_0"]
    AIBRIX_KV_CACHE_OL_INFINISTORE_USE_GDR: bool = True

    # HPKV Env Vars
    AIBRIX_KV_CACHE_OL_HPKV_REMOTE_ADDR: str = "127.0.0.1"
    AIBRIX_KV_CACHE_OL_HPKV_REMOTE_PORT: int = 12346
    AIBRIX_KV_CACHE_OL_HPKV_LOCAL_ADDR: str = "127.0.0.1"
    AIBRIX_KV_CACHE_OL_HPKV_LOCAL_PORT: int = 12345
    AIBRIX_KV_CACHE_OL_HPKV_USE_GDR: bool = True

# The begin-* and end* here are used by the documentation generator
# to extract the used env vars.

# begin-env-vars-definition

kv_cache_ol_environment_variables: Dict[str, Callable[[], Any]] = {
    "AIBRIX_KV_CACHE_OL_DOUBLE_GET_THRESHOLD": lambda: tuple(
        map(
            lambda x, t: t(x),
            map(
                str.strip,
                os.getenv(
                    "AIBRIX_KV_CACHE_OL_DOUBLE_GET_THRESHOLD", "4,0.1"
                ).split(","),
            ),
            (int, float),
        )
    ),

Review comment (suggestion):
I kind of feel this fashion is a little bit conflict with `NCCL_IB_HCA` setting, however, their `:1` is port, but not gid. Let's stick with this way at this moment and consider whether need to do internal search or env auto replacement later.

Code after:
    AIBRIX_KV_CACHE_OL_L2_CACHE_NUM_ASYNC_WORKERS: int = 8

    # Mock Connector
    AIBRIX_KV_CACHE_OL_MOCK_USE_RDMA: bool = False
    AIBRIX_KV_CACHE_OL_MOCK_USE_MPUT_MGET: bool = False

    # RocksDB Env Vars
    AIBRIX_KV_CACHE_OL_ROCKSDB_ROOT: str = os.path.expanduser(
        os.path.join(os.path.expanduser("~"), ".kv_cache_ol", "rocksdb")
    )
    AIBRIX_KV_CACHE_OL_ROCKSDB_TTL_S: int = 600
    AIBRIX_KV_CACHE_OL_ROCKSDB_WRITE_BUFFER_SIZE: int = 64 * 1024 * 1024
    AIBRIX_KV_CACHE_OL_ROCKSDB_TARGET_FILE_SIZE_BASE: int = 64 * 1024 * 1024
    AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_WRITE_BUFFER_NUMBER: int = 3
    AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_TOTAL_WAL_SIZE: int = 128 * 1024 * 1024
    AIBRIX_KV_CACHE_OL_ROCKSDB_MAX_BACKGROUND_JOBS: int = 8

    # InfiniStore Env Vars
    AIBRIX_KV_CACHE_OL_INFINISTORE_HOST_ADDR: str = "127.0.0.1"
    AIBRIX_KV_CACHE_OL_INFINISTORE_SERVICE_PORT: int = 12345
    AIBRIX_KV_CACHE_OL_INFINISTORE_CONNECTION_TYPE: str = "RDMA"
    AIBRIX_KV_CACHE_OL_INFINISTORE_IB_PORT: int = 1
    AIBRIX_KV_CACHE_OL_INFINISTORE_LINK_TYPE: str = "Ethernet"
    # Since 0.2.42, InfiniStore supports RDMA GID index in client config,
    # users can specify the GID index of each device in this format:
    # "mlx5_0:gid0,mlx5_1:gid1,mlx5_2:gid2"
    AIBRIX_KV_CACHE_OL_INFINISTORE_VISIBLE_DEV_LIST: List[str] = ["mlx5_0"]
    AIBRIX_KV_CACHE_OL_INFINISTORE_USE_GDR: bool = True

    # HPKV Env Vars
    AIBRIX_KV_CACHE_OL_HPKV_REMOTE_ADDR: str = "127.0.0.1"
    AIBRIX_KV_CACHE_OL_HPKV_REMOTE_PORT: int = 12346
    AIBRIX_KV_CACHE_OL_HPKV_LOCAL_ADDR: str = "127.0.0.1"
    AIBRIX_KV_CACHE_OL_HPKV_LOCAL_PORT: int = 12345
    AIBRIX_KV_CACHE_OL_HPKV_USE_GDR: bool = True

# The begin-* and end* here are used by the documentation generator
# to extract the used env vars.

# begin-env-vars-definition

kv_cache_ol_environment_variables: Dict[str, Callable[[], Any]] = {
    "AIBRIX_KV_CACHE_OL_DOUBLE_GET_THRESHOLD": lambda: tuple(
        map(
            lambda x, t: t(x),
            map(
                str.strip,
                os.getenv(
                    "AIBRIX_KV_CACHE_OL_DOUBLE_GET_THRESHOLD", "4,0.1"
                ).split(","),
            ),