chore: implement runtime_env func (#1104)

* 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>

rdagent/components/coder/factor_coder/config.py

@@ -1,6 +1,14 @@
+import os
+from typing import Optional
+
 from pydantic_settings import SettingsConfigDict
 
 from rdagent.components.coder.CoSTEER.config import CoSTEERSettings
+from rdagent.utils.env import (
+    CondaConf,
+    Env,
+    LocalEnv,
+)
 
 
 class FactorCoSTEERSettings(CoSTEERSettings):
@@ -25,4 +33,21 @@ class FactorCoSTEERSettings(CoSTEERSettings):
     """Path to the Python binary"""
 
 
+def get_factor_env(
+    conf_type: Optional[str] = None,
+    extra_volumes: dict = {},
+    running_timeout_period: int = 600,
+    enable_cache: Optional[bool] = None,
+) -> Env:
+    conf = FactorCoSTEERSettings()
+    if hasattr(conf, "python_bin"):
+        env = LocalEnv(conf=(CondaConf(conda_env_name=os.environ.get("CONDA_DEFAULT_ENV"))))
+    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
+
+
 FACTOR_COSTEER_SETTINGS = FactorCoSTEERSettings()

rdagent/components/coder/model_coder/conf.py

@@ -1,6 +1,14 @@
+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):
@@ -10,4 +18,26 @@ class ModelCoSTEERSettings(CoSTEERSettings):
     """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()

rdagent/scenarios/data_science/scen/__init__.py

@@ -16,6 +16,7 @@ from rdagent.scenarios.kaggle.kaggle_crawler import (
     download_data,
     get_metric_direction,
 )
+from rdagent.scenarios.shared.get_runtime_info import get_runtime_environment_by_env
 from rdagent.utils.agent.tpl import T
 
 
@@ -169,13 +170,8 @@ class DataScienceScen(Scenario):
     def get_runtime_environment(self) -> str:
         # TODO:  add it into base class.  Environment should(i.e. `DSDockerConf`) should be part of the scenario class.
         """Return runtime environment information."""
-        env = get_ds_env()
-        implementation = FBWorkspace()
-        fname = "runtime_info.py"
-        implementation.inject_files(
-            **{fname: (Path(__file__).absolute().resolve().parent / "runtime_info.py").read_text()}
-        )
-        stdout = implementation.execute(env=env, entry=f"python {fname}")
+        ds_env = get_ds_env()
+        stdout = get_runtime_environment_by_env(env=ds_env)
         return stdout
 
     def _get_data_folder_description(self) -> str:

rdagent/scenarios/general_model/scenario.py

@@ -50,3 +50,6 @@ The output of your code should be in the format:
 The simulator user can use to test your model:
 {self.simulator}
 """
+
+    def get_runtime_environment(self):
+        return None

rdagent/scenarios/qlib/experiment/factor_experiment.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 from pathlib import Path
 
+from rdagent.components.coder.factor_coder.config import get_factor_env
 from rdagent.components.coder.factor_coder.factor import (
     FactorExperiment,
     FactorFBWorkspace,
@@ -10,6 +11,7 @@ from rdagent.core.experiment import Task
 from rdagent.core.scenario import Scenario
 from rdagent.scenarios.qlib.experiment.utils import get_data_folder_intro
 from rdagent.scenarios.qlib.experiment.workspace import QlibFBWorkspace
+from rdagent.scenarios.shared.get_runtime_info import get_runtime_environment_by_env
 from rdagent.utils.agent.tpl import T
 
 
@@ -23,7 +25,11 @@ class QlibFactorExperiment(FactorExperiment[FactorTask, QlibFBWorkspace, FactorF
 class QlibFactorScenario(Scenario):
     def __init__(self) -> None:
         super().__init__()
-        self._background = deepcopy(T(".prompts:qlib_factor_background").r())
+        self._background = deepcopy(
+            T(".prompts:qlib_factor_background").r(
+                runtime_environment=self.get_runtime_environment(),
+            )
+        )
         self._source_data = deepcopy(get_data_folder_intro())
         self._output_format = deepcopy(T(".prompts:qlib_factor_output_format").r())
         self._interface = deepcopy(T(".prompts:qlib_factor_interface").r())
@@ -77,3 +83,8 @@ The output of your code should be in the format:
 The simulator user can use to test your factor:
 {self.simulator}
 """
+
+    def get_runtime_environment(self):
+        factor_env = get_factor_env()
+        stdout = get_runtime_environment_by_env(env=factor_env)
+        return stdout

rdagent/scenarios/qlib/experiment/model_experiment.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 from pathlib import Path
 
+from rdagent.components.coder.model_coder.conf import get_model_env
 from rdagent.components.coder.model_coder.model import (
     ModelExperiment,
     ModelFBWorkspace,
@@ -9,6 +10,7 @@ from rdagent.components.coder.model_coder.model import (
 from rdagent.core.experiment import Task
 from rdagent.core.scenario import Scenario
 from rdagent.scenarios.qlib.experiment.workspace import QlibFBWorkspace
+from rdagent.scenarios.shared.get_runtime_info import get_runtime_environment_by_env
 from rdagent.utils.agent.tpl import T
 
 
@@ -22,7 +24,11 @@ class QlibModelExperiment(ModelExperiment[ModelTask, QlibFBWorkspace, ModelFBWor
 class QlibModelScenario(Scenario):
     def __init__(self) -> None:
         super().__init__()
-        self._background = deepcopy(T(".prompts:qlib_model_background").r())
+        self._background = deepcopy(
+            T(".prompts:qlib_model_background").r(
+                runtime_environment=self.get_runtime_environment(),
+            )
+        )
         self._output_format = deepcopy(T(".prompts:qlib_model_output_format").r())
         self._interface = deepcopy(T(".prompts:qlib_model_interface").r())
         self._simulator = deepcopy(T(".prompts:qlib_model_simulator").r())
@@ -69,3 +75,8 @@ The output of your code should be in the format:
 The simulator user can use to test your model:
 {self.simulator}
 """
+
+    def get_runtime_environment(self):
+        model_env = get_model_env()
+        stdout = get_runtime_environment_by_env(env=model_env)
+        return stdout

rdagent/scenarios/qlib/experiment/prompts.yaml

@@ -3,6 +3,11 @@ qlib_quant_background: |-
   
   You are one of the most authoritative quantitative researchers at a top Wall Street hedge fund. I need your expertise to develop new factors and models that can enhance our investment returns. Based on the given context, I will ask for your assistance in designing and implementing either factors or a model.
 
+  {% if runtime_environment is not none %}
+  ====== Runtime Environment ======
+  You have following environment to run the code:
+  {{ runtime_environment }}
+  {% endif %}
 
 qlib_factor_background: |-
   The factor is a characteristic or variable used in quant investment that can help explain the returns and risks of a portfolio or a single asset. Factors are used by investors to identify and exploit sources of excess returns, and they are central to many quantitative investment strategies.
@@ -16,6 +21,12 @@ qlib_factor_background: |-
   The factor might not provide all the parts of the information above since some might not be applicable.
   Please specifically give all the hyperparameter in the factors like the window size, look back period, and so on. One factor should statically defines one output with a static source data. For example, last 10 days momentum and last 20 days momentum should be two different factors.
 
+  {% if runtime_environment is not none %}
+  ====== Runtime Environment ======
+  You have following environment to run the code:
+  {{ runtime_environment }}
+  {% endif %}
+
 qlib_factor_interface: |-
   Your python code should follow the interface to better interact with the user's system.
   Your python code should contain the following part: the import part, the function part, and the main part. You should write a main function name: "calculate_{function_name}" and call this function in "if __name__ == __main__" part. Don't write any try-except block in your python code. The user will catch the exception message and provide the feedback to you.
@@ -165,6 +176,12 @@ qlib_model_background: |-
   6. ModelType: The type of the model, "Tabular" for tabular model and "TimeSeries" for time series model.
   The model should provide clear and detailed documentation of its architecture and hyperparameters. One model should statically define one output with a fixed architecture and hyperparameters.
 
+  {% if runtime_environment is not none %}
+  ====== Runtime Environment ======
+  You have following environment to run the code:
+  {{ runtime_environment }}
+  {% endif %}
+
 qlib_model_interface: |-
   Your python code should follow the interface to better interact with the user's system.
   You code should contain several parts:

rdagent/scenarios/qlib/experiment/quant_experiment.py

@@ -2,6 +2,7 @@ from copy import deepcopy
 from pathlib import Path
 
 # Factor
+from rdagent.components.coder.factor_coder.config import get_factor_env
 from rdagent.components.coder.factor_coder.factor import (
     FactorExperiment,
     FactorFBWorkspace,
@@ -9,6 +10,7 @@ from rdagent.components.coder.factor_coder.factor import (
 )
 
 # Model
+from rdagent.components.coder.model_coder.conf import get_model_env
 from rdagent.components.coder.model_coder.model import (
     ModelExperiment,
     ModelFBWorkspace,
@@ -18,6 +20,7 @@ from rdagent.core.experiment import Task
 from rdagent.core.scenario import Scenario
 from rdagent.scenarios.qlib.experiment.utils import get_data_folder_intro
 from rdagent.scenarios.qlib.experiment.workspace import QlibFBWorkspace
+from rdagent.scenarios.shared.get_runtime_info import get_runtime_environment_by_env
 from rdagent.utils.agent.tpl import T
 
 
@@ -43,14 +46,18 @@ class QlibQuantScenario(Scenario):
 
     def background(self, tag=None) -> str:
         assert tag in [None, "factor", "model"]
-        quant_background = "The background of the scenario is as follows:\n" + T(".prompts:qlib_quant_background").r()
-        factor_background = (
-            "This time, I need your help with the research and development of the factor. The background of the factor scenario is as follows:\n"
-            + T(".prompts:qlib_factor_background").r()
+        quant_background = "The background of the scenario is as follows:\n" + T(".prompts:qlib_quant_background").r(
+            runtime_environment=self.get_runtime_environment(),
         )
-        model_background = (
-            "This time, I need your help with the research and development of the model. The background of the model scenario is as follows:\n"
-            + T(".prompts:qlib_model_background").r()
+        factor_background = "This time, I need your help with the research and development of the factor. The background of the factor scenario is as follows:\n" + T(
+            ".prompts:qlib_factor_background"
+        ).r(
+            runtime_environment=self.get_runtime_environment(tag="factor"),
+        )
+        model_background = "This time, I need your help with the research and development of the model. The background of the model scenario is as follows:\n" + T(
+            ".prompts:qlib_model_background"
+        ).r(
+            runtime_environment=self.get_runtime_environment(tag="model"),
         )
 
         # TODO: There are some issues here
@@ -165,3 +172,31 @@ class QlibQuantScenario(Scenario):
             return common_description("model") + interface("model") + output("model") + simulator("model")
         elif action == "factor" or action == "model":
             return common_description(action) + interface(action) + output(action) + simulator(action)
+
+    def get_runtime_environment(self, tag: str = None) -> str:
+        assert tag in [None, "factor", "model"]
+
+        if tag is None or tag == "factor":
+            # Use factor env to get the runtime environment
+            factor_env = get_factor_env()
+            factor_stdout = get_runtime_environment_by_env(env=factor_env)
+            if tag == "factor":
+                stdout = factor_stdout
+
+        if tag is None or tag == "model":
+            # Use model env to get the runtime environment
+            model_env = get_model_env()
+            model_stdout = get_runtime_environment_by_env(env=model_env)
+            if tag == "model":
+                stdout = model_stdout
+
+        if tag is None:
+            # Combine the outputs from both environments
+            stdout = (
+                "=== [Environment to generate the factors] ===\n"
+                + factor_stdout.strip()
+                + "\n\n=== [Environment to train the models] ===\n"
+                + model_stdout.strip()
+            )
+
+        return stdout

rdagent/scenarios/shared/get_runtime_info.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+
+from rdagent.core.experiment import FBWorkspace
+from rdagent.utils.env import Env
+
+
+def get_runtime_environment_by_env(env: Env) -> str:
+    implementation = FBWorkspace()
+    fname = "runtime_info.py"
+    implementation.inject_files(**{fname: (Path(__file__).absolute().resolve().parent / "runtime_info.py").read_text()})
+    stdout = implementation.execute(env=env, entry=f"python {fname}")
+    return stdout