Spaces:
Running on Zero
Running on Zero
File size: 5,050 Bytes
eb808a5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """Serializable inference dispatch across local, GPU, and ZeroGPU runtimes."""
from __future__ import annotations
import logging
import time
from dataclasses import dataclass
from typing import Any, Callable, Literal
from core.errors import GatewayError
from core.loader import ModelLoader
from core.runtime import inference_context, memory_stats
logger = logging.getLogger(__name__)
def _gpu_task(**options: Any) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
"""Use the Space scheduler in deployment and a direct call in local tests."""
try:
import spaces
except ImportError:
return lambda function: function
return spaces.GPU(**options)
@dataclass(frozen=True, slots=True)
class InferenceCommand:
"""Pickle-safe description of one model method invocation."""
model_name: str
method_name: str
arguments: dict[str, Any]
request_id: str
duration_seconds: int
gpu_size: Literal["large", "xlarge"] = "large"
@dataclass(frozen=True, slots=True)
class CommandError:
"""Serializable error returned from a forked ZeroGPU worker."""
message: str
status_code: int
code: str
@dataclass(frozen=True, slots=True)
class CommandResult:
"""Serializable success or failure from a GPU allocation."""
value: Any = None
error: CommandError | None = None
def _duration(command: InferenceCommand) -> int:
return command.duration_seconds
def _zero_gpu_active() -> bool:
try:
from spaces.config import Config
except ImportError:
return False
return bool(Config.zero_gpu)
def _invoke(command: InferenceCommand) -> CommandResult:
"""Execute one command inside an allocated device process."""
loader = ModelLoader()
started = time.perf_counter()
try:
with loader.use_model(command.model_name) as model:
method = getattr(model, command.method_name, None)
if not callable(method):
raise RuntimeError(
f"Model {command.model_name} does not implement "
f"{command.method_name}"
)
with inference_context(loader.settings, loader.device):
value = method(**command.arguments)
except GatewayError as exc:
return CommandResult(
error=CommandError(exc.message, exc.status_code, exc.code)
)
except Exception as exc:
logger.exception(
"model inference failed",
extra={
"model": command.model_name,
"request_id": command.request_id,
**memory_stats(),
},
)
message = str(exc)
if "out of memory" in message.lower():
error = CommandError(
f"{command.model_name} ran out of memory",
507,
"out_of_memory",
)
elif isinstance(exc, OSError):
error = CommandError(
f"File operation failed during {command.model_name} inference",
500,
"file_error",
)
else:
error = CommandError(
f"{command.model_name} inference failed: {message}",
500,
"inference_failed",
)
return CommandResult(error=error)
finally:
if _zero_gpu_active():
loader.close()
logger.info(
"model inference completed",
extra={
"model": command.model_name,
"request_id": command.request_id,
"execution_time": round(time.perf_counter() - started, 3),
**memory_stats(),
},
)
return CommandResult(value=value)
@_gpu_task(duration=_duration, size="large")
def _invoke_large(command: InferenceCommand) -> CommandResult:
return _invoke(command)
@_gpu_task(duration=_duration, size="xlarge")
def _invoke_xlarge(command: InferenceCommand) -> CommandResult:
return _invoke(command)
def execute_inference(command: InferenceCommand) -> Any:
"""Acquire the requested ZeroGPU tier and restore domain errors in the parent."""
try:
result = (
_invoke_xlarge(command)
if command.gpu_size == "xlarge"
else _invoke_large(command)
)
except Exception as exc:
message = str(exc)
normalized = message.lower()
if "quota" in normalized or "no gpu" in normalized or "zerogpu" in normalized:
raise GatewayError(
"ZeroGPU is unavailable or its quota was exceeded",
status_code=503,
code="gpu_unavailable",
) from exc
raise GatewayError(
"ZeroGPU scheduling failed",
status_code=503,
code="gpu_scheduling_failed",
) from exc
if result.error is not None:
raise GatewayError(
result.error.message,
status_code=result.error.status_code,
code=result.error.code,
)
return result.value
|