lightseekorg--tokenspeed/python/tokenspeed/runtime/utils/env.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import warnings
from contextlib import contextmanager
from typing import Any

from tokenspeed.runtime.utils.pdl import pdl_enabled
from tokenspeed.runtime.utils.server_args import ServerArgs

global_server_args_dict: dict = {
    "attention_backend": ServerArgs.attention_backend,
    "sampling_backend": ServerArgs.sampling_backend,
    "attention_use_fp4_indexer_cache": ServerArgs.attention_use_fp4_indexer_cache,
    "deepseek_v4_mega_moe_max_num_tokens": ServerArgs.deepseek_v4_mega_moe_max_num_tokens,
    "deepseek_v4_indexer_prefill_max_logits_mb": ServerArgs.deepseek_v4_indexer_prefill_max_logits_mb,
    "deepseek_v4_prefill_chunk_size": ServerArgs.deepseek_v4_prefill_chunk_size,
    "triton_attention_reduce_in_fp32": ServerArgs.triton_attention_reduce_in_fp32,
    "kv_cache_dtype": ServerArgs.kv_cache_dtype,
    "enable_nan_detection": ServerArgs.enable_nan_detection,
    "enable_p2p_check": ServerArgs.enable_p2p_check,
    "mapping": ServerArgs.mapping,
    "force_deterministic_rsag": ServerArgs.force_deterministic_rsag,
    "low_latency_max_num_tokens_per_gpu": ServerArgs.low_latency_max_num_tokens_per_gpu,
    "device": ServerArgs.device,
    "draft_model_path_use_base": ServerArgs.draft_model_path_use_base,
    "disable_pdl": ServerArgs.disable_pdl,
    "enable_prefix_caching": ServerArgs.enable_prefix_caching,
    "mla_disable_ragged": ServerArgs.mla_disable_ragged,
    "chunked_prefill_size": ServerArgs.chunked_prefill_size,
    "mla_chunk_multiplier": ServerArgs.mla_chunk_multiplier,
    "ep_num_redundant_experts": ServerArgs.ep_num_redundant_experts,
    "ep_dispatch_algorithm": ServerArgs.ep_dispatch_algorithm,
    "enable_eplb": ServerArgs.enable_eplb,
    "mm_attention_backend": ServerArgs.mm_attention_backend,
    "comm_fusion_max_num_tokens": ServerArgs.comm_fusion_max_num_tokens,
    "enable_allreduce_fusion": ServerArgs.enable_allreduce_fusion,
    "max_prefill_tokens": ServerArgs.max_prefill_tokens,
    "max_model_len": ServerArgs.max_model_len,
    "max_num_seqs": ServerArgs.max_num_seqs,
    "moe_backend": ServerArgs.moe_backend,
    "enforce_eager": ServerArgs.enforce_eager,
    "max_cudagraph_capture_size": ServerArgs.max_cudagraph_capture_size,
    "cudagraph_capture_sizes": ServerArgs.cudagraph_capture_sizes,
    "disable_prefill_graph": ServerArgs.disable_prefill_graph,
    "prefill_graph_max_tokens": ServerArgs.prefill_graph_max_tokens,
    "mamba_track_interval": ServerArgs.mamba_track_interval,
    "all2all_backend": ServerArgs.all2all_backend,
}


def global_server_args_dict_update(server_args: ServerArgs):
    global_server_args_dict.update(
        {
            "attention_backend": server_args.attention_backend,
            "sampling_backend": server_args.sampling_backend,
            "attention_use_fp4_indexer_cache": server_args.attention_use_fp4_indexer_cache,
            "deepseek_v4_mega_moe_max_num_tokens": server_args.deepseek_v4_mega_moe_max_num_tokens,
            "deepseek_v4_indexer_prefill_max_logits_mb": server_args.deepseek_v4_indexer_prefill_max_logits_mb,
            "deepseek_v4_prefill_chunk_size": server_args.deepseek_v4_prefill_chunk_size,
            "triton_attention_reduce_in_fp32": server_args.triton_attention_reduce_in_fp32,
            "kv_cache_dtype": server_args.kv_cache_dtype,
            "enable_nan_detection": server_args.enable_nan_detection,
            "enable_p2p_check": server_args.enable_p2p_check,
            "mapping": server_args.mapping,
            "force_deterministic_rsag": server_args.force_deterministic_rsag,
            "low_latency_max_num_tokens_per_gpu": server_args.low_latency_max_num_tokens_per_gpu,
            "device": server_args.device,
            "draft_model_path_use_base": server_args.draft_model_path_use_base,
            "speculative_algorithm": server_args.speculative_algorithm,
            "speculative_num_draft_tokens": server_args.speculative_num_draft_tokens,
            "disable_pdl": server_args.disable_pdl,
            "enable_prefix_caching": server_args.enable_prefix_caching,
            "mla_disable_ragged": server_args.mla_disable_ragged,
            "chunked_prefill_size": server_args.chunked_prefill_size,
            "mla_chunk_multiplier": server_args.mla_chunk_multiplier,
            "ep_num_redundant_experts": server_args.ep_num_redundant_experts,
            "ep_dispatch_algorithm": server_args.ep_dispatch_algorithm,
            "enable_eplb": server_args.enable_eplb,
            "mm_attention_backend": server_args.mm_attention_backend,
            "comm_fusion_max_num_tokens": server_args.comm_fusion_max_num_tokens,
            "enable_allreduce_fusion": server_args.enable_allreduce_fusion,
            "max_prefill_tokens": server_args.max_prefill_tokens,
            "max_model_len": server_args.max_model_len,
            "max_num_seqs": server_args.max_num_seqs,
            "moe_backend": server_args.moe_backend,
            "enforce_eager": server_args.enforce_eager,
            "max_cudagraph_capture_size": server_args.max_cudagraph_capture_size,
            "cudagraph_capture_sizes": server_args.cudagraph_capture_sizes,
            "disable_prefill_graph": server_args.disable_prefill_graph,
            "prefill_graph_max_tokens": server_args.prefill_graph_max_tokens,
            "all2all_backend": server_args.all2all_backend,
        }
    )
    pdl_enabled.cache_clear()


class EnvField:
    def __init__(self, default: Any):
        self.default = default
        #  we use None to indicate whether the value is set or not
        # If the value is manually set to None, we need mark it as _set_to_none.
        # Always use clear() to reset the value, which leads to the default fallback.
        self._set_to_none = False

    def __set_name__(self, owner, name):
        self.name = name

    def parse(self, value: str) -> Any:
        raise NotImplementedError()

    def get(self) -> Any:
        value = os.getenv(self.name)
        if self._set_to_none:
            if value is not None:
                raise RuntimeError(f"{self.name} is set while marked as None.")
            return None

        if value is None:
            return self.default

        try:
            return self.parse(value)
        except ValueError as e:
            warnings.warn(
                f'Invalid value for {self.name}: {e}, using default "{self.default}"'
            )
            return self.default

    def is_set(self):
        #  If None is manually set, it is considered as set.
        return self.name in os.environ or self._set_to_none

    def get_set_value_or(self, or_value: Any):
        #  Ugly usage, but only way to get custom default value.
        return self.get() if self.is_set() else or_value

    def set(self, value: Any):
        if value is None:
            self._set_to_none = True
            os.environ.pop(self.name, None)
        else:
            self._set_to_none = False
            os.environ[self.name] = str(value)

    @contextmanager
    def override(self, value: Any):
        backup_present = self.name in os.environ
        backup_value = os.environ.get(self.name)
        backup_set_to_none = self._set_to_none
        self.set(value)
        yield
        if backup_present:
            os.environ[self.name] = backup_value
        else:
            os.environ.pop(self.name, None)
        self._set_to_none = backup_set_to_none

    def clear(self):
        os.environ.pop(self.name, None)
        self._set_to_none = False

    @property
    def value(self):
        return self.get()

    def __bool__(self):
        raise RuntimeError(
            "Please use `envs.YOUR_FLAG.get()` instead of `envs.YOUR_FLAG`"
        )

    def __len__(self):
        raise RuntimeError(
            "Please use `envs.YOUR_FLAG.get()` instead of `envs.YOUR_FLAG`"
        )


class EnvStr(EnvField):
    def parse(self, value: str) -> str:
        return value


class EnvBool(EnvField):
    def parse(self, value: str) -> bool:
        value = value.lower()
        if value in ["true", "1", "yes", "y"]:
            return True
        if value in ["false", "0", "no", "n"]:
            return False
        raise ValueError(f'"{value}" is not a valid boolean value')


class EnvInt(EnvField):
    def parse(self, value: str) -> int:
        try:
            return int(value)
        except ValueError:
            raise ValueError(f'"{value}" is not a valid integer value')


class EnvFloat(EnvField):
    def parse(self, value: str) -> float:
        try:
            return float(value)
        except ValueError:
            raise ValueError(f'"{value}" is not a valid float value')


class Envs:
    # fmt: off

    # Model download
    TOKENSPEED_USE_MODELSCOPE = EnvBool(False)

    # Test and debug
    TOKENSPEED_CUDA_COREDUMP = EnvBool(False)
    TOKENSPEED_CUDA_COREDUMP_DIR = EnvStr("/tmp/tokenspeed_cuda_coredumps")
    TOKENSPEED_PROFILE_WITH_STACK = EnvBool(True)
    TOKENSPEED_TEST_REQUEST_TIME_STATS = EnvBool(False)
    TOKENSPEED_PROFILER_DIR = EnvStr("/tmp")
    TOKENSPEED_CI_SMALL_KV_SIZE = EnvInt(-1)
    TOKENSPEED_NVTX = EnvBool(False)
    TOKENSPEED_DP_SAMPLING_BACKEND = EnvStr(None)

    # Scheduler
    TOKENSPEED_BLOCK_NONZERO_RANK_CHILDREN = EnvBool(True)

    # Mooncake
    TOKENSPEED_KVSTORE_MOONCAKE_CONFIG_PATH = EnvStr(None)
    MOONCAKE_MASTER = EnvStr(None)
    MOONCAKE_LOCAL_HOSTNAME = EnvStr("localhost")
    MOONCAKE_TE_META_DATA_SERVER = EnvStr("P2PHANDSHAKE")
    MOONCAKE_GLOBAL_SEGMENT_SIZE = EnvStr("4gb")
    MOONCAKE_PROTOCOL = EnvStr("tcp")
    MOONCAKE_DEVICE = EnvStr("")
    MOONCAKE_MASTER_METRICS_PORT = EnvInt(9003)
    MOONCAKE_CHECK_SERVER = EnvBool(False)
    TOKENSPEED_DISAGGREGATION_FAILED_SESSION_TTL = EnvInt(30)
    TOKENSPEED_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0)
    TOKENSPEED_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2)
    TOKENSPEED_DISAGGREGATION_QUEUE_SIZE = EnvInt(4)
    TOKENSPEED_DISAGGREGATION_THREAD_POOL_SIZE = EnvInt(-1)
    TOKENSPEED_DISAGGREGATION_BOOTSTRAP_TIMEOUT = EnvInt(120)
    TOKENSPEED_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300)
    TOKENSPEED_EPD_ENCODE_RING_SLOTS = EnvInt(64)
    TOKENSPEED_EPD_ENCODE_RING_SLOT_MB = EnvInt(None)
    TOKENSPEED_EPD_ENCODE_EMBED_CACHE_MB = EnvInt(4096)
    TOKENSPEED_EPD_ENCODE_EMBED_CACHE_DRAM_MB = EnvInt(0)
    TOKENSPEED_EPD_RECV_POOL_SLOTS = EnvInt(16)
    TOKENSPEED_EPD_RECV_POOL_SLOT_MB = EnvInt(256)
    TOKENSPEED_EPD_EMBEDDING_SHARD = EnvBool(True)
    TOKENSPEED_PD_LAYERWISE_DEBUG = EnvBool(False)
    TOKENSPEED_PD_PREFILL_METADATA_TIMEOUT = EnvFloat(5.0)

    # Quantization
    TOKENSPEED_NVFP4_GEMM_SWIGLU_NVFP4_QUANT = EnvBool(True)
    TOKENSPEED_MINIMAX_AR_USE_TRITON = EnvBool(False)

    # EPLB
    TOKENSPEED_EXPERT_DISTRIBUTION_RECORDER_DIR = EnvStr("/tmp")

    # Runtime behavior
    TOKENSPEED_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN = EnvBool(False)
    TOKENSPEED_DETOKENIZER_MAX_STATES = EnvInt(1 << 16)
    TOKENSPEED_FORCE_FAKE_FULL_NVLINK = EnvBool(False)
    TOKENSPEED_HEALTH_CHECK_TIMEOUT = EnvInt(20)
    TOKENSPEED_HOST_IP = EnvStr("")
    TOKENSPEED_LOGGING_CONFIG_PATH = EnvStr(None)
    TOKENSPEED_MAMBA_SSM_DTYPE = EnvStr("float32")
    TOKENSPEED_MODEL_REDIRECT_PATH = EnvStr(None)
    TOKENSPEED_MOE_PADDING = EnvBool(False)
    TOKENSPEED_MOE_CONFIG_DIR = EnvStr(None)
    TOKENSPEED_ENABLE_TORCH_INFERENCE_MODE = EnvBool(True)
    TOKENSPEED_NUMA_AWARE_WORKER_AFFINITY = EnvBool(True)
    TOKENSPEED_REQUEST_CONVERSION_WORKERS = EnvInt(8)

    # Multimodal / VLM
    TOKENSPEED_LOG_MM_TIMING = EnvBool(False)
    TOKENSPEED_MM_ENABLE_ENCODER_CUDA_GRAPH = EnvBool(False)
    TOKENSPEED_MM_VIDEO_ENCODER_CUDA_GRAPH_MAX_SEQUENCES_PER_BATCH = EnvInt(None)
    TOKENSPEED_MM_SKIP_COMPUTE_HASH = EnvBool(False)

    # fmt: on


envs = Envs()