NexQuant/rdagent/components/coder/factor_coder/factor.py

import logging
from __future__ import annotations

import subprocess  # nosec B404
import uuid
from pathlib import Path
from typing import Tuple, Union

import pandas as pd
from filelock import FileLock

from rdagent.app.kaggle.conf import KAGGLE_IMPLEMENT_SETTING
from rdagent.components.coder.CoSTEER.task import CoSTEERTask
from rdagent.components.coder.factor_coder.config import FACTOR_COSTEER_SETTINGS
from rdagent.core.exception import CodeFormatError, CustomRuntimeError, NoOutputError
from rdagent.core.experiment import Experiment, FBWorkspace
from rdagent.core.utils import cache_with_pickle  # nosec
from rdagent.oai.llm_utils import md5_hash


class FactorTask(CoSTEERTask):
    # TODO:  generalized the attributes into the Task
    # - factor_* -> *
    def __init__(
        self,
        factor_name,
        factor_description,
        factor_formulation,
        *args,
        variables: dict = {},
        resource: str = None,
        factor_implementation: bool = False,
        **kwargs,
    ) -> None:
        self.factor_name = (
            factor_name  # TODO: remove it in the later version. Keep it only for pickle version compatibility  # nosec
        )
        self.factor_formulation = factor_formulation
        self.variables = variables
        self.factor_resources = resource
        self.factor_implementation = factor_implementation
        super().__init__(name=factor_name, description=factor_description, *args, **kwargs)

    @property
    def factor_description(self):
        """for compatibility"""
        return self.description

    def get_task_information(self):
        return f"""factor_name: {self.factor_name}
factor_description: {self.factor_description}
factor_formulation: {self.factor_formulation}
variables: {str(self.variables)}"""

    def get_task_brief_information(self):
        return f"""factor_name: {self.factor_name}
factor_description: {self.factor_description}
factor_formulation: {self.factor_formulation}
variables: {str(self.variables)}"""

    def get_task_information_and_implementation_result(self):
        return {
            "factor_name": self.factor_name,
            "factor_description": self.factor_description,
            "factor_formulation": self.factor_formulation,
            "variables": str(self.variables),
            "factor_implementation": str(self.factor_implementation),
        }

    @staticmethod
    def from_dict(dict):
        return FactorTask(**dict)

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}[{self.factor_name}]>"


class FactorFBWorkspace(FBWorkspace):
    """
    This class is used to implement a factor by writing the code to a file.
    Input data and output factor value are also written to files.
    """

    # TODO: (Xiao) think raising errors may get better information for processing
    FB_EXEC_SUCCESS = "Execution succeeded without error."
    FB_CODE_NOT_SET = "code is not set."
    FB_EXECUTION_SUCCEEDED = "Execution succeeded without error."
    FB_OUTPUT_FILE_NOT_FOUND = "\nExpected output file not found."
    FB_OUTPUT_FILE_FOUND = "\nExpected output file found."

    def __init__(
        self,
        *args,
        raise_exception: bool = False,
        **kwargs,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.raise_exception = raise_exception

    def hash_func(self, data_type: str = "Debug") -> str:
        return (
            md5_hash(data_type + self.file_dict["factor.py"])
            if ("factor.py" in self.file_dict and not self.raise_exception)
            else None
        )

    @cache_with_pickle(hash_func)  # nosec
    def execute(self, data_type: str = "Debug") -> Tuple[str, pd.DataFrame]:  # nosec
        """
        execute the implementation and get the factor value by the following steps:  # nosec
        1. make the directory in workspace path
        2. write the code to the file in the workspace path
        3. link all the source data to the workspace path folder
        if call_factor_py is True:
            4. execute the code  # nosec
        else:
            4. generate a script from template to import the factor.py dump get the factor value to result.h5
        5. read the factor value from the output file in the workspace path folder
        returns the execution feedback as a string and the factor value as a pandas dataframe  # nosec


        Regarding the cache mechanism:
        1. We will store the function's return value to ensure it behaves as expected.
        - The cached information will include a tuple with the following: (execution_feedback, executed_factor_value_dataframe, Optional[Exception])  # nosec

        """
        self.before_execute()  # nosec
        if self.file_dict is None or "factor.py" not in self.file_dict:
            if self.raise_exception:
                raise CodeFormatError(self.FB_CODE_NOT_SET)
            else:
                return self.FB_CODE_NOT_SET, None
        with FileLock(self.workspace_path / "execution.lock"):  # nosec
            if self.target_task.version == 1:
                source_data_path = (
                    Path(
                        FACTOR_COSTEER_SETTINGS.data_folder_debug,
                    )
                    if data_type == "Debug"  # FIXME: (yx) don't think we should use a debug tag for this.
                    else Path(
                        FACTOR_COSTEER_SETTINGS.data_folder,
                    )
                )
            elif self.target_task.version == 2:
                # TODO you can change the name of the data folder for a better understanding
                source_data_path = Path(KAGGLE_IMPLEMENT_SETTING.local_data_path) / KAGGLE_IMPLEMENT_SETTING.competition

            source_data_path.mkdir(exist_ok=True, parents=True)
            code_path = self.workspace_path / f"factor.py"

            self.link_all_files_in_folder_to_workspace(source_data_path, self.workspace_path)

            execution_feedback = self.FB_EXECUTION_SUCCEEDED  # nosec
            execution_success = False  # nosec
            execution_error = None  # nosec

            if self.target_task.version == 1:
                execution_code_path = code_path  # nosec
            elif self.target_task.version == 2:
                execution_code_path = self.workspace_path / f"{uuid.uuid4()}.py"  # nosec
                execution_code_path.write_text((Path(__file__).parent / "factor_execution_template.txt").read_text())  # nosec

            try:
                subprocess.check_output(  # nosec
                    [FACTOR_COSTEER_SETTINGS.python_bin, str(execution_code_path)],  # nosec
                    shell=False,
                    cwd=self.workspace_path,
                    stderr=subprocess.STDOUT,  # nosec
                    timeout=FACTOR_COSTEER_SETTINGS.file_based_execution_timeout,  # nosec
                )
                execution_success = True  # nosec
            except subprocess.CalledProcessError as e:  # nosec
                import site

                execution_feedback = (  # nosec
                    e.output.decode()
                    .replace(str(execution_code_path.parent.absolute()), r"/path/to")  # nosec
                    .replace(str(site.getsitepackages()[0]), r"/path/to/site-packages")
                )
                if len(execution_feedback) > 2000:  # nosec
                    execution_feedback = (  # nosec
                        execution_feedback[:1000] + "....hidden long error message...." + execution_feedback[-1000:]  # nosec
                    )
                if self.raise_exception:
                    raise CustomRuntimeError(execution_feedback)  # nosec
                else:
                    execution_error = CustomRuntimeError(execution_feedback)  # nosec
            except subprocess.TimeoutExpired:  # nosec
                execution_feedback += f"Execution timeout error and the timeout is set to {FACTOR_COSTEER_SETTINGS.file_based_execution_timeout} seconds."  # nosec
                if self.raise_exception:
                    raise CustomRuntimeError(execution_feedback)  # nosec
                else:
                    execution_error = CustomRuntimeError(execution_feedback)  # nosec

            workspace_output_file_path = self.workspace_path / "result.h5"
            if workspace_output_file_path.exists() and execution_success:  # nosec
                try:
                    executed_factor_value_dataframe = pd.read_hdf(workspace_output_file_path)  # nosec
                    execution_feedback += self.FB_OUTPUT_FILE_FOUND  # nosec
                except Exception as e:
                    execution_feedback += f"Error found when reading hdf file: {e}"[:1000]  # nosec
                    executed_factor_value_dataframe = None  # nosec
            else:
                execution_feedback += self.FB_OUTPUT_FILE_NOT_FOUND  # nosec
                executed_factor_value_dataframe = None  # nosec
                if self.raise_exception:
                    raise NoOutputError(execution_feedback)  # nosec
                else:
                    execution_error = NoOutputError(execution_feedback)  # nosec

        return execution_feedback, executed_factor_value_dataframe  # nosec

    def __str__(self) -> str:
        # NOTE:
        # If the code cache works, the workspace will be None.
        return f"File Factor[{self.target_task.factor_name}]: {self.workspace_path}"

    def __repr__(self) -> str:
        return self.__str__()

    @staticmethod
    def from_folder(task: FactorTask, path: Union[str, Path], **kwargs):
        path = Path(path)
        code_dict = {}
        for file_path in path.iterdir():
            if file_path.suffix == ".py":
                code_dict[file_path.name] = file_path.read_text()
        return FactorFBWorkspace(target_task=task, code_dict=code_dict, **kwargs)


FactorExperiment = Experiment
FeatureExperiment = Experiment