From e681b40b8f5d02ee7c2b87bcc0c108e08a27445e Mon Sep 17 00:00:00 2001 From: SunXiaoye <31361630+JingofXin@users.noreply.github.com> Date: Wed, 21 Aug 2024 18:52:18 +0800 Subject: [PATCH] Clear EV before test_AutoModel and restore them after test (#149) --- lazyllm/configs.py | 33 +++++++++++++++++-- tests/advanced_tests/full_test/test_deploy.py | 30 +++++++++++++++++ tests/basic_tests/test_config.py | 12 +++++++ 3 files changed, 73 insertions(+), 2 deletions(-) diff --git a/lazyllm/configs.py b/lazyllm/configs.py index 5d9badba..0f73e309 100644 --- a/lazyllm/configs.py +++ b/lazyllm/configs.py @@ -1,6 +1,7 @@ import os from enum import Enum import json +from typing import List, Union class Mode(Enum): @@ -11,6 +12,8 @@ class Mode(Enum): class Config(object): def __init__(self, prefix='LAZYLLM', home='~/.lazyllm/'): + self._config_params = dict() + self._env_map_name = dict() self.prefix = prefix self.impl, self.cfgs = dict(), dict() self.add('home', str, home, 'HOME') @@ -38,6 +41,21 @@ def get_all_configs(self): return self.impl def add(self, name, type, default=None, env=None): + update_params = (type, default, env) + if name not in self._config_params or self._config_params[name] != update_params: + if name in self._config_params: + print(f"Warning: The default configuration parameter {name}({self._config_params[name]}) " + f"has been added, but a new {name}({update_params}) has been added repeatedly.") + self._config_params.update({name: update_params}) + if isinstance(env, str): + self._env_map_name[('lazyllm_' + env).upper()] = name + elif isinstance(env, dict): + for k in env.keys(): + self._env_map_name[('lazyllm_' + k).upper()] = name + self._update_impl(name, type, default, env) + return self + + def _update_impl(self, name, type, default=None, env=None): self.impl[name] = self.cfgs.pop(name) if name in self.cfgs else default if isinstance(env, dict): for k, v in env.items(): @@ -48,7 +66,6 @@ def add(self, name, type, default=None, env=None): self.impl[name] = self.getenv(env, type, self.impl[name]) if not isinstance(self.impl[name], type): raise TypeError( f'Invalid config type for {name}, type is {type}') - return self def __getitem__(self, name): return self.impl[name] @@ -56,10 +73,22 @@ def __getitem__(self, name): def __str__(self): return str(self.impl) + def refresh(self, targets: Union[str, List[str]] = None) -> None: + names = targets + if isinstance(targets, str): + names = targets.lower() + if names.startswith('lazyllm_'): + names = names[8:] + names = [names] + elif targets is None: + curr_envs = [key for key in os.environ.keys() if key.startswith('LAZYLLM_')] + names = list(set([self._env_map_name[key] for key in curr_envs])) + assert isinstance(names, list) + for name in names: + self._update_impl(name, *self._config_params[name]) config = Config().add('mode', Mode, Mode.Normal, dict(DISPLAY=Mode.Display, DEBUG=Mode.Debug) ).add('repr_ml', bool, False, 'REPR_USE_ML' ).add('rag_store', str, 'none', 'RAG_STORE' - ).add('redis_url', str, 'none', 'REDIS_URL' ).add('gpu_type', str, 'A100', 'GPU_TYPE' ) diff --git a/tests/advanced_tests/full_test/test_deploy.py b/tests/advanced_tests/full_test/test_deploy.py index 4453e13b..c75f4d79 100644 --- a/tests/advanced_tests/full_test/test_deploy.py +++ b/tests/advanced_tests/full_test/test_deploy.py @@ -1,14 +1,43 @@ +import os import json import time import pytest import httpx import random +from functools import wraps from gradio_client import Client import lazyllm from lazyllm import deploy from lazyllm.launcher import cleanup +def reset_env(func): + + env_vars_to_reset = [ + "LAZYLLM_OPENAI_API_KEY", + "LAZYLLM_KIMI_API_KEY", + "LAZYLLM_GLM_API_KEY", + "LAZYLLM_QWEN_API_KEY", + "LAZYLLM_SENSENOVA_API_KEY", + "LAZYLLM_SENSENOVA_SECRET_KEY", + ] + + @wraps(func) + def wrapper(*args, **kwargs): + original_values = {var: os.environ.get(var, None) for var in env_vars_to_reset} + for var in env_vars_to_reset: + os.environ.pop(var, None) + lazyllm.config.refresh(var) + result = func(*args, **kwargs) + for var, value in original_values.items(): + if value is None: + os.environ.pop(var, None) + else: + os.environ[var] = value + lazyllm.config.refresh(var) + return result + return wrapper + class TestDeploy(object): def setup_method(self): @@ -82,6 +111,7 @@ def test_bark(self): res = m('你好啊，很高兴认识你。') assert "sounds" in json.loads(res) + @reset_env def test_AutoModel(self): # No model_name and key chat = lazyllm.AutoModel() diff --git a/tests/basic_tests/test_config.py b/tests/basic_tests/test_config.py index 1c39e406..316d89bf 100644 --- a/tests/basic_tests/test_config.py +++ b/tests/basic_tests/test_config.py @@ -1,6 +1,7 @@ import lazyllm from lazyllm.configs import Mode import os +import copy import pytest import contextlib import inspect @@ -40,6 +41,17 @@ def run_subprocess(self, pytester): class TestConfig(object): + def test_refresh(self): + origin = copy.deepcopy(lazyllm.config.impl) + os.environ['LAZYLLM_GPU_TYPE'] = 'H100' + lazyllm.config.refresh('LAZYLLM_GPU_TYPE') + assert lazyllm.config.impl['gpu_type'] == 'H100' + os.environ['LAZYLLM_GPU_TYPE'] = origin['gpu_type'] + lazyllm.config.refresh('gpu_type') + assert lazyllm.config.impl['gpu_type'] == origin['gpu_type'] + lazyllm.config.refresh() + assert lazyllm.config.impl == origin + def test_config_mode(self): print(os.environ.get('LAZYLLM_DISPLAY')) assert lazyllm.config['mode'] == Mode.Normal