Pause scheduling after the current forward step drains.
(self)
| 153 | self.has_runable_event.set() |
| 154 | |
| 155 | async def drain_for_sleep(self): |
| 156 | """Pause scheduling after the current forward step drains.""" |
| 157 | if self._sleep_requested: |
| 158 | logger.info('EngineLoop sleep drain already requested; waiting for drain point.') |
| 159 | await self._main_sleep_drain_event.wait() |
| 160 | if self.config.role != EngineRole.Hybrid: |
| 161 | await self._migration_sleep_drain_event.wait() |
| 162 | return |
| 163 | logger.info('EngineLoop sleep drain requested.') |
| 164 | self._sleep_requested = True |
| 165 | self._main_sleep_drain_event.clear() |
| 166 | if self.config.role != EngineRole.Hybrid: |
| 167 | self._migration_sleep_drain_event.clear() |
| 168 | self._sleep_resume_event.clear() |
| 169 | # Wake main_loop if it is idle waiting for runnable work, so it can |
| 170 | # observe _sleep_requested and acknowledge the drain. |
| 171 | self.has_runable_event.event.set() |
| 172 | # Wake migration_loop if it is idle on migration_event; it has its own |
| 173 | # drain acknowledgement because migration can also touch KV resources. |
| 174 | if self.config.role != EngineRole.Hybrid: |
| 175 | self.migration_event.set() |
| 176 | await self._main_sleep_drain_event.wait() |
| 177 | if self.config.role != EngineRole.Hybrid: |
| 178 | await self._migration_sleep_drain_event.wait() |
| 179 | logger.info('EngineLoop reached sleep drain point.') |
| 180 | |
| 181 | def resume_from_sleep(self): |
| 182 | """Resume scheduling after wakeup restores engine resources.""" |