Class DBBenchTask

src/server/tasks/dbbench/task.py:37–274

class DBBenchTask(Task):

    def __init__(self,
                 data_file: str,
                 db_file: Optional[str] = None,
                 db_password: str = 'password',
                 max_round: int = 20,
                 env_driver: str = 'docker',
                 env_options: Optional[dict] = None,
                 **configs):
        super().__init__(**configs)
        self.full_async = True
        self.logger = logging.getLogger(__name__)

        self.max_round = max_round
        self.data_file = data_file
        self.db_root_dir = db_file

        self.dataset = []
        # Load dataset
        with open(self.data_file) as f:
            raw_data = f.read()
            if self.data_file.endswith("json"):
                data = json.loads(raw_data)
            else:  # Assuming jsonl
                data = [json.loads(line) for line in raw_data.strip().split('\n')]

        for entry in data:
            ans_key = "answer_md5" if entry["type"][0] in ("INSERT", "DELETE", "UPDATE") else "label"
            ans = entry.pop(ans_key, None)  # Use pop with default
            inp = entry
            self.dataset.append((inp, ans))

        self.env_delegation = DBBenchEnvironmentDelegation(db_password)
        self.env_controller = create_controller(env_driver, self.env_delegation, **env_options)
        self.env_controller_background_task = None

        self.logger.info(f"DBBench initialized with {len(self.dataset)} samples. Root dir: {self.db_root_dir}")

    def get_indices(self) -> List[SampleIndex]:
        return list(range(len(self.dataset)))

    async def start_sample(self, index: int, session: Session) -> TaskSampleExecutionResult:
        self.env_controller.loop = asyncio.get_running_loop()
        if not self.env_controller_background_task:
            self.env_controller_background_task = asyncio.create_task(self.env_controller.background_task())
            weakref.finalize(self, self.env_controller_background_task.cancel)

        database: Optional[Database] = None
        try:
            entry = self.dataset[index][0]
            ground_truth = self.dataset[index][1]

            use_sqlite = entry.get("user_sqlite", False)
            if use_sqlite:
                db_dir = entry['create']['database']
                init_file = entry['create']['init']
                sqlite_path = os.path.join(self.db_root_dir, db_dir, init_file)