NexQuant/rdagent/core/experiment.py

from __future__ import annotations

import shutil
import uuid
from abc import ABC, abstractmethod
from copy import deepcopy
from pathlib import Path
from typing import Any, Dict, Generic, Optional, Sequence, TypeVar

from rdagent.core.conf import RD_AGENT_SETTINGS

"""
This file contains the all the class about organizing the task in RD-Agent.
"""


class Task(ABC):
    @abstractmethod
    def get_task_information(self):
        """
        Get the task information string to build the unique key
        """
        pass


ASpecificTask = TypeVar("ASpecificTask", bound=Task)


class Workspace(ABC, Generic[ASpecificTask]):
    """
    A workspace is a place to store the task implementation. It evolves as the developer implements the task.
    To get a snapshot of the workspace, make sure call `copy` to get a copy of the workspace.
    """

    def __init__(self, target_task: ASpecificTask = None) -> None:
        self.target_task: ASpecificTask = target_task

    @abstractmethod
    def execute(self, *args, **kwargs) -> object:
        raise NotImplementedError("execute method is not implemented.")

    @abstractmethod
    def copy(self) -> Workspace:
        raise NotImplementedError("copy method is not implemented.")


ASpecificWS = TypeVar("ASpecificWS", bound=Workspace)


class WsLoader(ABC, Generic[ASpecificTask, ASpecificWS]):
    @abstractmethod
    def load(self, task: ASpecificTask) -> ASpecificWS:
        raise NotImplementedError("load method is not implemented.")


class FBWorkspace(Workspace):
    """
    File-based task workspace

    The implemented task will be a folder which contains related elements.
    - Data
    - Code Workspace
    - Output
        - After execution, it will generate the final output as file.

    A typical way to run the pipeline of FBWorkspace will be
    (We didn't add it as a method due to that we may pass arguments into `prepare` or `execute` based on our requirements.)

    .. code-block:: python

        def run_pipeline(self, **files: str):
            self.prepare()
            self.inject_code(**files)
            self.execute()

    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.code_dict = (
            {}
        )  # The code injected into the folder, store them in the variable to reproduce the former result
        self.workspace_path: Path = RD_AGENT_SETTINGS.workspace_path / uuid.uuid4().hex

    @property
    def code(self) -> str:
        code_string = ""
        for file_name, code in self.code_dict.items():
            code_string += f"File: {file_name}\n{code}\n"
        return code_string

    def prepare(self, *args, **kwargs):
        """
        Prepare the workspace except the injected code
        - Data
        - Documentation
            typical usage of `*args, **kwargs`:
                Different methods shares the same data. The data are passed by the arguments.
        """
        self.workspace_path.mkdir(parents=True, exist_ok=True)

    def inject_code(self, **files: str):
        """
        Inject the code into the folder.
        {
            <file name>: <code>
        }
        """
        self.prepare()
        for k, v in files.items():
            self.code_dict[k] = v
            with open(self.workspace_path / k, "w") as f:
                f.write(v)

    def get_files(self) -> list[Path]:
        """
        Get the environment description.

        To be general, we only return a list of filenames.
        How to summarize the environment is the responsibility of the Developer.
        """
        return list(self.workspace_path.iterdir())

    def inject_code_from_folder(self, folder_path: Path):
        """
        Load the workspace from the folder
        """
        for file_path in folder_path.iterdir():
            if file_path.suffix == ".py" or file_path.suffix == ".yaml":
                self.inject_code(**{file_path.name: file_path.read_text()})

    def copy(self) -> FBWorkspace:
        """
        copy the workspace from the original one
        """
        return deepcopy(self)

    def clear(self) -> None:
        """
        Clear the workspace
        """
        shutil.rmtree(self.workspace_path)
        self.code_dict = {}

    @abstractmethod
    def execute(self, *args, **kwargs) -> object:
        """
        Before each execution, make sure to prepare and inject code
        """
        self.prepare()
        self.inject_code(**self.code_dict)


ASpecificWSForExperiment = TypeVar("ASpecificWSForExperiment", bound=Workspace)
ASpecificWSForSubTasks = TypeVar("ASpecificWSForSubTasks", bound=Workspace)


class Experiment(ABC, Generic[ASpecificTask, ASpecificWSForExperiment, ASpecificWSForSubTasks]):
    """
    The experiment is a sequence of tasks and the implementations of the tasks after generated by the Developer.
    """

    def __init__(self, sub_tasks: Sequence[ASpecificTask]) -> None:
        self.sub_tasks = sub_tasks
        self.sub_workspace_list: Sequence[ASpecificWSForSubTasks] = [None for _ in self.sub_tasks]
        self.based_experiments: Sequence[ASpecificWSForExperiment] = []
        self.result: object = None  # The result of the experiment, can be different types in different scenarios.
        self.experiment_workspace: ASpecificWSForExperiment = None


ASpecificExp = TypeVar("ASpecificExp", bound=Experiment)

TaskOrExperiment = TypeVar("TaskOrExperiment", Task, Experiment)


class Loader(ABC, Generic[TaskOrExperiment]):
    @abstractmethod
    def load(self, *args, **kwargs) -> TaskOrExperiment:
        raise NotImplementedError("load method is not implemented.")