mirror of
https://github.com/NicolasBohn/NexQuant.git
synced 2026-07-27 23:47:46 +00:00
Linlang
bbcee389b4
* implement runtime_env func for quant * add runtime_info code * add runtime env information to the prompt * format with black * optimize get_runtime_env code * delete unnecessary files * some refinement * fix fin_quant bugs --------- Co-authored-by: Xu Yang <peteryang@vip.qq.com>
44 lines
1.2 KiB
Python
44 lines
1.2 KiB
Python
from typing import Optional
|
|
|
|
from pydantic_settings import SettingsConfigDict
|
|
|
|
from rdagent.components.coder.CoSTEER.config import CoSTEERSettings
|
|
from rdagent.utils.env import (
|
|
Env,
|
|
QlibCondaConf,
|
|
QlibCondaEnv,
|
|
QTDockerEnv,
|
|
)
|
|
|
|
|
|
class ModelCoSTEERSettings(CoSTEERSettings):
|
|
model_config = SettingsConfigDict(env_prefix="MODEL_CoSTEER_")
|
|
|
|
env_type: str = "conda" # or "docker"
|
|
"""Environment to run model code in coder and runner: 'conda' for local conda env, 'docker' for Docker container"""
|
|
|
|
|
|
def get_model_env(
|
|
conf_type: Optional[str] = None,
|
|
extra_volumes: dict = {},
|
|
running_timeout_period: int = 600,
|
|
enable_cache: Optional[bool] = None,
|
|
) -> Env:
|
|
conf = ModelCoSTEERSettings()
|
|
if conf.env_type == "docker":
|
|
env = QTDockerEnv()
|
|
elif conf.env_type == "conda":
|
|
env = QlibCondaEnv(conf=QlibCondaConf())
|
|
else:
|
|
raise ValueError(f"Unknown env type: {conf.env_type}")
|
|
|
|
env.conf.extra_volumes = extra_volumes.copy()
|
|
env.conf.running_timeout_period = running_timeout_period
|
|
if enable_cache is not None:
|
|
env.conf.enable_cache = enable_cache
|
|
env.prepare()
|
|
return env
|
|
|
|
|
|
MODEL_COSTEER_SETTINGS = ModelCoSTEERSettings()
|