Clear EV before test_AutoModel and restore them after test (LazyAGI#149)

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,18 +66,29 @@ 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]
 
     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'
                 )

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()

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