| 128 | return count |
| 129 | |
| 130 | class CodeGen: |
| 131 | def __init__(self, args): |
| 132 | self.args = args |
| 133 | self.contents = {} # dictionary: maps task_id to contents (mix of tests and codes) |
| 134 | self.code_scores = {} # dictionary: maps task_id to scores (mix of tests and codes) |
| 135 | self.tests = [] |
| 136 | self.codes = [] |
| 137 | self.pair_scores = {} # dictionary: maps task_id to (code_id, test_id, score) |
| 138 | self.total_codes_requested = 0 |
| 139 | self.total_tests_requested = 0 |
| 140 | |
| 141 | logging.info("Setting up VM...") |
| 142 | |
| 143 | self.docker_execute = DockerExecute(sources_dirname=args.sources_dirname) |
| 144 | |
| 145 | logging.info("Starting LLM workers...") |
| 146 | |
| 147 | self.manager = JobManager(args) |
| 148 | |
| 149 | def add_code_or_test_job(self): |
| 150 | test_count = len(self.tests) |
| 151 | code_count = len(self.codes) |
| 152 | |
| 153 | add_code = True |
| 154 | |
| 155 | if test_count == 0 and code_count == 0: |
| 156 | add_code = self.total_tests_requested > self.total_codes_requested |
| 157 | else: |
| 158 | add_code = test_count > code_count |
| 159 | |
| 160 | if add_code: |
| 161 | logging.info(f"Adding a job to write more code (tests asked/completed={self.total_codes_requested}/{test_count}, codes asked/completed={self.total_tests_requested}/{code_count})") |
| 162 | self.manager.add_code_job() |
| 163 | self.total_codes_requested += 1 |
| 164 | else: |
| 165 | logging.info(f"Adding a job to write more tests (tests asked/completed={self.total_codes_requested}/{test_count}, codes asked/completed={self.total_tests_requested}/{code_count})") |
| 166 | self.manager.add_test_job() |
| 167 | self.total_tests_requested += 1 |
| 168 | |
| 169 | def test_pair(self, code_id, test_id): |
| 170 | exit_code, logs = copy_and_run_pytest( |
| 171 | args.sources_dirname, |
| 172 | args.function_name, |
| 173 | code_id, |
| 174 | test_id, |
| 175 | self.docker_execute) |
| 176 | |
| 177 | if exit_code != 0: |
| 178 | logging.info(f"Test failed: code {code_id} <-> test {test_id}: exit_code={exit_code} logs={logs}") |
| 179 | |
| 180 | if len(logs) == "": |
| 181 | logging.info("Test failed really badly somehow. Deleting {code_id} and {test_id} to avoid repeating this error.") |
| 182 | self.codes.remove(code_id) |
| 183 | self.tests.remove(test_id) |
| 184 | return False |
| 185 | |
| 186 | logging.info(f"Test passed: code {code_id} <-> test {test_id} - Asking judge if we are done") |
| 187 | |