(engine_dir: Path)
| 773 | model_format = self.llm_args.model_format |
| 774 | |
| 775 | def build_task(engine_dir: Path): |
| 776 | if model_format is not _ModelFormatKind.TLLM_ENGINE: |
| 777 | model_loader_kwargs = { |
| 778 | 'llm_args': self.llm_args, |
| 779 | 'workspace': str(self.workspace), |
| 780 | 'llm_build_stats': self.llm_build_stats, |
| 781 | } |
| 782 | |
| 783 | if self.llm_args.parallel_config.is_multi_gpu: |
| 784 | assert self.mpi_session |
| 785 | |
| 786 | #mpi_session cannot be pickled so remove from self.llm_args |
| 787 | if self.llm_args.mpi_session: |
| 788 | del self.llm_args.mpi_session |
| 789 | |
| 790 | # The engine_dir:Path will be stored to MPINodeState.state |
| 791 | build_infos = self.mpi_session.submit_sync( |
| 792 | CachedModelLoader._node_build_task, |
| 793 | engine_dir=engine_dir, |
| 794 | **model_loader_kwargs) |
| 795 | self.llm_build_stats.build_steps_info = build_infos[0] |
| 796 | |
| 797 | else: # single-gpu |
| 798 | with ModelLoader(**model_loader_kwargs) as model_loader: |
| 799 | model_loader(engine_dir=engine_dir) |
| 800 | |
| 801 | release_gc() |
| 802 | |
| 803 | has_storage = True |
| 804 | if self.build_cache_enabled: |
nothing calls this directly
no test coverage detected