(self)
| 770 | trust_remote_code=self.llm_args.trust_remote_code) |
| 771 | |
| 772 | def _build_model(self) -> 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: |
| 805 | try: |
| 806 | # TODO[chunweiy]: Cover the case when the model is from HF model hub. |
| 807 | if self.model_loader.model_obj.is_local_model: |
| 808 | # This is not perfect, but will make build-cache much more robust. |
| 809 | free_storage = self.engine_cache_stage.parent.free_storage_in_gb( |
| 810 | ) |
| 811 | model_size = get_directory_size_in_gb( |
| 812 | self.model_loader.model_obj.model_dir) |
| 813 | require_size = model_size * 1.3 |
| 814 | has_storage = free_storage >= require_size |
| 815 | |
| 816 | if not has_storage: |
| 817 | print_colored( |
| 818 | f"Build cache is disabled since the cache storage is too small.\n ", |
| 819 | 'yellow') |
| 820 | print_colored( |
| 821 | f"Free storage: {free_storage}GB, Required storage: {require_size}GB\n", |
| 822 | 'grey') |
| 823 | except ValueError: |
| 824 | has_storage = False |
| 825 | except Exception as e: |
| 826 | logger.error(e) |
| 827 | has_storage = False |
| 828 | |
| 829 | if enable_llm_debug(): |
no test coverage detected