跳到内容

vllm_gaudi.extension.environment

_VLLM_VALUES 模块属性 

_VLLM_VALUES = {}

VllmValue 

VllmValue(name, env_var_type, depend=None)
源代码在 vllm_gaudi/extension/environment.py 
 99
100
101
102
103
def VllmValue(name, env_var_type, depend=None):
    if depend is not None:
        return Value(name, env_var_type=env_var_type, dependencies=depend)
    global _VLLM_VALUES
    return Value(name, lambda _: _VLLM_VALUES[name], env_var_type=env_var_type)

_get_build 

_get_build(_)
源代码在 vllm_gaudi/extension/environment.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def _get_build(_):
    import re
    import subprocess
    output = subprocess.run("pip show habana-torch-plugin",
                            shell=True,
                            text=True,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.PIPE)
    version_re = re.compile(r'Version:\s*(?P<version>.*)')
    match = version_re.search(output.stdout)
    if output.returncode == 0 and match:
        return match.group('version')
    # In cpu-test environment we don't have access to habana-torch-plugin
    from vllm_gaudi.extension.utils import is_fake_hpu
    result = '0.0.0.0' if is_fake_hpu() else None
    logger().warning(f"Unable to detect habana-torch-plugin version! Returning: {result}")
    return result

_get_hw 

_get_hw(_)
源代码在 vllm_gaudi/extension/environment.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
def _get_hw(_):
    import habana_frameworks.torch.utils.experimental as htexp
    device_type = htexp._get_device_type()
    match device_type:
        case htexp.synDeviceType.synDeviceGaudi2:
            return "gaudi2"
        case htexp.synDeviceType.synDeviceGaudi3:
            return "gaudi3"
    from vllm_gaudi.extension.utils import is_fake_hpu
    if is_fake_hpu():
        return "cpu"
    logger().warning(f'Unknown device type: {device_type}')
    return None

_get_prefix 

_get_prefix(_)
源代码在 vllm_gaudi/extension/environment.py 
31
32
33
34
35
36
37
38
39
def _get_prefix(_):
    if os.environ.get('VLLM_UNIFIED_ATTN'):
        return True
    conti_pa = os.environ.get('VLLM_CONTIGUOUS_PA')
    if conti_pa is None:
        return True
    elif boolean(conti_pa) is True:
        return False
    return True

_get_pt_bridge_mode 

_get_pt_bridge_mode(_)
源代码在 vllm_gaudi/extension/environment.py 
94
95
96
def _get_pt_bridge_mode(_):
    import habana_frameworks.torch as htorch
    return 'lazy' if htorch.utils.internal.is_lazy() else 'eager'

_get_vllm_engine_version 

_get_vllm_engine_version(_)
源代码在 vllm_gaudi/extension/environment.py 
85
86
87
88
89
90
91
def _get_vllm_engine_version(_):
    try:
        import vllm.envs as envs
        return 'v1'
    except ImportError:
        logger().info("vllm module not installed, returning 'unknown' for engine version")
        return 'unknown'

_get_vllm_hash 

_get_vllm_hash(_)
源代码在 vllm_gaudi/extension/environment.py 
42
43
44
45
46
47
48
49
50
51
52
53
def _get_vllm_hash(_):
    import subprocess

    try:
        commit_hash = subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD']).decode('utf-8').strip()
        try:
            branch = subprocess.check_output(['git', 'rev-parse', '--abbrev-ref', 'HEAD']).decode('utf-8').strip()
        except subprocess.CalledProcessError:
            branch = "Error getting branch name"
        return branch + "+" + commit_hash
    except subprocess.CalledProcessError as e:
        return "Error getting commit hash"

get_environment 

get_environment()
源代码在 vllm_gaudi/extension/environment.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def get_environment():
    values = [
        Value('hw', _get_hw, env_var_type=str, check=choice('cpu', 'gaudi', 'gaudi2', 'gaudi3')),
        Value('build',
              _get_build,
              env_var_type=str,
              check=regex(r'^\d+\.\d+\.\d+\.\d+$',
                          hint='You can override detected build by specifying VLLM_BUILD env variable')),
        Value('engine_version', _get_vllm_engine_version, env_var_type=str),
        Value('bridge_mode', _get_pt_bridge_mode, env_var_type=str, check=choice('eager', 'lazy')),
        VllmValue('model_type', str),
        VllmValue('prefix_caching', boolean, depend=_get_prefix),
        Value('vllm_gaudi_commit', _get_vllm_hash, env_var_type=str)
    ]
    return split_values_and_flags(values)

set_vllm_config 

set_vllm_config(cfg)
源代码在 vllm_gaudi/extension/environment.py 
75
76
77
78
79
80
81
82
def set_vllm_config(cfg):
    global _VLLM_VALUES

    if hasattr(cfg.model_config, 'hf_config'):
        _VLLM_VALUES['model_type'] = cfg.model_config.hf_config.model_type
    else:
        _VLLM_VALUES['model_type'] = cfg.model_config.model_type
    _VLLM_VALUES['prefix_caching'] = cfg.cache_config.enable_prefix_caching