(config: Config)
| 20 | |
| 21 | |
| 22 | async def create_debug_models(config: Config) -> None: |
| 23 | from trinity.common.models.allocator import Allocator |
| 24 | |
| 25 | allocator = Allocator(config.explorer) |
| 26 | rollout_models, auxiliary_models = await allocator.create_all_models() |
| 27 | logger.info( |
| 28 | "----------------------------------------------------\n" |
| 29 | "Inference models started successfully for debugging.\n" |
| 30 | "Press Ctrl+C to exit.\n" |
| 31 | "----------------------------------------------------" |
| 32 | ) |
| 33 | try: |
| 34 | while True: |
| 35 | await asyncio.sleep(1) |
| 36 | except KeyboardInterrupt: |
| 37 | logger.info("Shutting down debug models...") |
| 38 | # wait for all models to shutdown |
| 39 | models = rollout_models + [model for sublist in auxiliary_models for model in sublist] |
| 40 | await asyncio.gather(*[model.shutdown() for model in models]) |
| 41 | ray.shutdown() |
| 42 | |
| 43 | |
| 44 | async def get_debug_models(config: Config) -> Tuple["ModelWrapper", List["ModelWrapper"]]: |
no test coverage detected