| 351 | |
| 352 | |
| 353 | class QueryWorkerToServerTask(SubmititWorker): |
| 354 | def graceful_exit( |
| 355 | self, |
| 356 | monitor: ResourceMonitor | None, |
| 357 | reason: Optional[str] = None, |
| 358 | ) -> None: |
| 359 | self.report_dead(reason=reason) |
| 360 | |
| 361 | if monitor is not None: |
| 362 | if monitor.has_started: |
| 363 | util_report: UtilizationReport = monitor.stop_and_report( |
| 364 | num_examples=self.num_examples_processed |
| 365 | ) |
| 366 | util_report.job_type = JobType.GENERATION.value |
| 367 | self.submit_util_report(util_report) |
| 368 | # Erase the monitor |
| 369 | monitor = None |
| 370 | |
| 371 | sys.exit(1) |
| 372 | |
| 373 | def __call__( |
| 374 | self, |
| 375 | model_config: ModelConfig, |
| 376 | buffer_size: int = 128, |
| 377 | gpu_id: Optional[int] = None, |
| 378 | ) -> None: |
| 379 | self.exit_reason = "Unkown" |
| 380 | |
| 381 | if gpu_id is not None: |
| 382 | # Set CUDA_VISIBLE_DEVICES to only show the specific GPU for this subprocess |
| 383 | os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) |
| 384 | # Set the GPU for this process |
| 385 | torch.cuda.set_device( |
| 386 | 0 |
| 387 | ) # Use GPU 0 since CUDA_VISIBLE_DEVICES makes it the only visible GPU |
| 388 | |
| 389 | try: |
| 390 | self.monitor = ResourceMonitor( |
| 391 | interval_sec=10, |
| 392 | track_children=False, |
| 393 | track_cgroup=False, |
| 394 | profile_gpu=False, |
| 395 | timings_only=not self.do_in_depth_monitoring, |
| 396 | ) |
| 397 | self.monitor.start() |
| 398 | |
| 399 | vllm_model = LLM( |
| 400 | model=model_config.model_name, |
| 401 | dtype=model_config.dtype, |
| 402 | max_model_len=model_config.max_tokens, |
| 403 | ) |
| 404 | |
| 405 | sampling_params = SamplingParams( |
| 406 | temperature=model_config.temperature, |
| 407 | max_tokens=model_config.max_tokens, |
| 408 | top_p=model_config.top_p, |
| 409 | top_k=model_config.top_k, |
| 410 | # skip_special_tokens=False, |
no outgoing calls
no test coverage detected