| 182 | |
| 183 | |
| 184 | class QueryBatchLocalCheckpointableTask(submitit.helpers.Checkpointable): |
| 185 | def __init__(self, num_checkpoints: int = 10): |
| 186 | self.reset_state(num_checkpoints) |
| 187 | self.monitor_runner: ResourceMonitor | None = None |
| 188 | |
| 189 | def reset_state(self, num_checkpoints: int): |
| 190 | self.prompts_to_process: List[List[str]] = [] |
| 191 | self.processed_outputs: List[List[QueryResult]] = [] |
| 192 | self.num_prompts: int = 0 |
| 193 | self.num_checkpoints = num_checkpoints |
| 194 | self.original_num_chunks: int | None = None |
| 195 | self.setup_done: bool = False |
| 196 | # This is a list of util reports as if we have interrupts we save multiple reports |
| 197 | self.util_reports: List[UtilizationReport] = [] |
| 198 | |
| 199 | self.num_examples_processed_in_current_job: int = 0 |
| 200 | |
| 201 | def __call__( |
| 202 | self, prompts: List[str], model_config: ModelConfig |
| 203 | ) -> Tuple[List[QueryResult], List[UtilizationReport]]: |
| 204 | export() |
| 205 | |
| 206 | self.monitor_runner = ResourceMonitor( |
| 207 | interval_sec=5.0, |
| 208 | track_children=False, |
| 209 | per_process_top_n=0, |
| 210 | track_cgroup=True, |
| 211 | profile_gpu=False, |
| 212 | timings_only=True, |
| 213 | ) |
| 214 | self.monitor_runner.start() |
| 215 | |
| 216 | if not self.setup_done: |
| 217 | # This is the first time running, so set prompts to process to prompts |
| 218 | self.reset_state(self.num_checkpoints) |
| 219 | |
| 220 | prompts_chunks: List[List[str]] = chunk_list(prompts, self.num_checkpoints) |
| 221 | self.prompts_to_process = prompts_chunks |
| 222 | self.num_prompts = len(prompts) |
| 223 | self.original_num_chunks = len(prompts_chunks) |
| 224 | |
| 225 | self.setup_done = True |
| 226 | |
| 227 | vllm_model = LLM( |
| 228 | model=model_config.model_name, |
| 229 | dtype=model_config.dtype, |
| 230 | ) |
| 231 | |
| 232 | sampling_params = SamplingParams( |
| 233 | temperature=model_config.temperature, |
| 234 | max_tokens=model_config.max_tokens, |
| 235 | top_p=model_config.top_p, |
| 236 | top_k=model_config.top_k, |
| 237 | # include_stop_str_in_output=True, |
| 238 | # skip_special_tokens=False, |
| 239 | logprobs=10, # For entropy calculation |
| 240 | ) |
| 241 | |