| 510 | keep_last=True) |
| 511 | |
| 512 | def dynamic_load(self): |
| 513 | if hasattr(self, 'inference_pids') and self.inference_pids is not None: |
| 514 | return |
| 515 | gpu_infer = os.environ.get( |
| 516 | 'LOCAL_WORLD_SIZE') or torch.cuda.device_count() |
| 517 | pmi_rank = int(os.environ['RANK']) |
| 518 | pmi_world_size = int(os.environ['WORLD_SIZE']) |
| 519 | in_q_list = [ |
| 520 | torch.multiprocessing.Manager().Queue() for _ in range(gpu_infer) |
| 521 | ] |
| 522 | out_q = torch.multiprocessing.Manager().Queue() |
| 523 | initialized_events = [ |
| 524 | torch.multiprocessing.Manager().Event() for _ in range(gpu_infer) |
| 525 | ] |
| 526 | context = mp.spawn( |
| 527 | self.mp_worker, |
| 528 | nprocs=gpu_infer, |
| 529 | args=(gpu_infer, pmi_rank, pmi_world_size, in_q_list, out_q, |
| 530 | initialized_events, self), |
| 531 | join=False) |
| 532 | all_initialized = False |
| 533 | while not all_initialized: |
| 534 | all_initialized = all( |
| 535 | event.is_set() for event in initialized_events) |
| 536 | if not all_initialized: |
| 537 | time.sleep(0.1) |
| 538 | print('Inference model is initialized', flush=True) |
| 539 | self.in_q_list = in_q_list |
| 540 | self.out_q = out_q |
| 541 | self.inference_pids = context.pids() |
| 542 | self.initialized_events = initialized_events |
| 543 | |
| 544 | def transfer_data_to_cuda(self, data, device): |
| 545 | if data is None: |