NexQuant/rdagent/scenarios/qlib/experiment/workspace.py

from pathlib import Path
from typing import Any

import pandas as pd

from rdagent.core.experiment import FBWorkspace
from rdagent.log import rdagent_logger as logger
from rdagent.utils.env import QTDockerEnv


class QlibFBWorkspace(FBWorkspace):
    def __init__(self, template_folder_path: Path, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.inject_code_from_folder(template_folder_path)

    def execute(self, qlib_config_name: str = "conf.yaml", run_env: dict = {}, *args, **kwargs) -> str:
        qtde = QTDockerEnv()
        qtde.prepare()

        # Run the Docker command
        execute_log = qtde.run(
            local_path=str(self.workspace_path),
            entry="rm -r mlruns",
            env=run_env,
        )
        # Run the Qlib backtest
        execute_log = qtde.run(
            local_path=str(self.workspace_path),
            entry=f"qrun {qlib_config_name}",
            env=run_env,
        )

        execute_log = qtde.run(
            local_path=str(self.workspace_path),
            entry="python read_exp_res.py",
            env=run_env,
        )

        ret_df = pd.read_pickle(self.workspace_path / "ret.pkl")
        logger.log_object(ret_df, tag="returns")  # TODO: tag

        csv_path = self.workspace_path / "qlib_res.csv"

        if not csv_path.exists():
            logger.error(f"File {csv_path} does not exist.")
            return None

        return pd.read_csv(csv_path, index_col=0).iloc[:, 0]