Reload offloaded gpu memory occupation for certain parts, e.g. weight, cache. Args: control_request: Control request object containing parameters for reloading memory tags: list of tags to reload, supported values: ["weight", "kv_cache"]
(self, control_request: ControlRequest)
| 1608 | return asyncio.run(self._wait_for_control_responses(control_request.request_id, 60, executors=executors)) |
| 1609 | |
| 1610 | def _control_wakeup(self, control_request: ControlRequest): |
| 1611 | """ |
| 1612 | Reload offloaded gpu memory occupation for certain parts, e.g. weight, cache. |
| 1613 | |
| 1614 | Args: |
| 1615 | control_request: Control request object containing parameters for reloading memory |
| 1616 | tags: list of tags to reload, supported values: ["weight", "kv_cache"] |
| 1617 | """ |
| 1618 | # Args check |
| 1619 | tags = self._parse_tags(control_request) |
| 1620 | control_request.args["tags"] = tags |
| 1621 | |
| 1622 | # Determine which executors are needed for the wakeup command |
| 1623 | executors = set() |
| 1624 | if "weight" in tags: |
| 1625 | executors.add("worker") |
| 1626 | if "kv_cache" in tags: |
| 1627 | executors.add("worker") |
| 1628 | if self.cfg.cache_config.num_cpu_blocks > 0 or self.cfg.cache_config.kvcache_storage_backend: |
| 1629 | executors.add("cache_transfer") |
| 1630 | |
| 1631 | # Dispatch wakeup request to executors |
| 1632 | self.llm_logger.info(f"Dispatch wakeup request to executors: {list(executors)}") |
| 1633 | self._dispatch_control_request(control_request, executors) |
| 1634 | result = asyncio.run(self._wait_for_control_responses(control_request.request_id, 300, executors=executors)) |
| 1635 | |
| 1636 | # Resume the engine after wakeup |
| 1637 | self._control_resume(None) |
| 1638 | |
| 1639 | return result |
| 1640 | |
| 1641 | def _dispatch_control_request(self, control_request: ControlRequest, executors: List[str]): |
| 1642 | """ |
nothing calls this directly
no test coverage detected