The engine_dir is the path to save the built engine.
(self, engine_dir: Optional[Path] = None)
| 237 | self.counter += 1 |
| 238 | |
| 239 | def __call__(self, engine_dir: Optional[Path] = None) -> Path: |
| 240 | ''' |
| 241 | The engine_dir is the path to save the built engine. |
| 242 | ''' |
| 243 | if self.llm_args.model_format is _ModelFormatKind.TLLM_ENGINE: |
| 244 | return self.model_obj.model_dir |
| 245 | |
| 246 | if self.llm_args.parallel_config.is_multi_gpu: |
| 247 | torch.cuda.set_device(self.global_rank % self.mapping.gpus_per_node) |
| 248 | |
| 249 | pipeline = ModelLoader.BuildPipeline( |
| 250 | self.llm_args.enable_tqdm, |
| 251 | [label for label, _ in self._build_pipeline], |
| 252 | [handler for _, handler in self._build_pipeline], |
| 253 | llm_build_stats=self.llm_build_stats, |
| 254 | ) |
| 255 | pipeline() |
| 256 | |
| 257 | assert engine_dir |
| 258 | |
| 259 | runtime_context = _ModelRuntimeContext( |
| 260 | engine=self._engine, |
| 261 | mapping=self.mapping, |
| 262 | model_info=self._model_info, |
| 263 | ) |
| 264 | self.save(runtime_context, self.model_obj.model_dir, engine_dir) |
| 265 | return engine_dir |
| 266 | |
| 267 | def __enter__(self): |
| 268 | return self |
nothing calls this directly
no test coverage detected