Hongsheng Liuandhsliuustc0106
8360346a66
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com> Co-authored-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>
60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Generic startup handoff between model loaders and offload backends."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
|
|
from torch import nn
|
|
|
|
from vllm_omni.diffusion.model_loader.host_weight_plan import HostWeightPlan
|
|
|
|
_STARTUP_STATE_ATTR = "_omni_offload_startup_state"
|
|
|
|
|
|
@dataclass
|
|
class OffloadStartupState:
|
|
"""Loader-owned state consumed by the offloader startup boundary."""
|
|
|
|
host_weight_plan: HostWeightPlan | None = None
|
|
fresh_model_loader: Callable[[], nn.Module] | None = None
|
|
allow_fresh_retry: bool = False
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.allow_fresh_retry != (self.fresh_model_loader is not None):
|
|
raise ValueError("fresh-model retry policy requires exactly one retry callback")
|
|
|
|
def close_loader_ownership(self) -> None:
|
|
"""Release a plan that never reached a backend."""
|
|
if self.host_weight_plan is None:
|
|
return
|
|
carrier = self.host_weight_plan.lease_carrier
|
|
if carrier is not None:
|
|
carrier.close()
|
|
|
|
|
|
def attach_offload_startup_state(model: nn.Module, state: OffloadStartupState) -> None:
|
|
"""Attach one process-local startup handoff to a loaded pipeline."""
|
|
previous = getattr(model, _STARTUP_STATE_ATTR, None)
|
|
if previous is not None:
|
|
previous.close_loader_ownership()
|
|
setattr(model, _STARTUP_STATE_ATTR, state)
|
|
|
|
|
|
def take_offload_startup_state(model: nn.Module) -> OffloadStartupState | None:
|
|
"""Take and remove the loader handoff from a pipeline exactly once."""
|
|
state = getattr(model, _STARTUP_STATE_ATTR, None)
|
|
if state is None:
|
|
return None
|
|
delattr(model, _STARTUP_STATE_ATTR)
|
|
return state
|
|
|
|
|
|
__all__ = [
|
|
"OffloadStartupState",
|
|
"attach_offload_startup_state",
|
|
"take_offload_startup_state",
|
|
]
|