File size: 1,650 Bytes
472bb49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Per-request qua_sdk DeployConfig selection (WHERE/HOW the SDK stages run).

GPU path = the SDK's SPACE_ZEROGPU preset (float16, AOTI band, dynamic
batching) with ASR torch.compile forced OFF: ZeroGPU forks per lease so the
compile cache never amortizes, and dynamo crashes on the wav2vec2 mask path
(torch 2.8 + transformers 5.0 ConstantVariable assertion). CPU path = the
local-subprocess / worker / dev fallback: ``config.CPU_DTYPE`` (bfloat16
dodges the SDPA QK^T cache cliff), no compile, tighter 300s batch cap. The
choice is made inside the stage functions because the same code body runs in
three contexts: a ZeroGPU lease, a forced-CPU subprocess/worker (per-thread
flag or no CUDA), and plain local dev.
"""

from __future__ import annotations

from qua_sdk.deploy.presets import SPACE_ZEROGPU, BatchingConfig, DeployConfig

CPU_MAX_BATCH_SECONDS = 300


def gpu_deploy() -> DeployConfig:
    cfg = SPACE_ZEROGPU.model_copy(deep=True)
    cfg.torch_compile = False
    return cfg


def cpu_deploy() -> DeployConfig:
    from config import CPU_DTYPE
    return DeployConfig(
        name="space_cpu", device="cpu", dtype=CPU_DTYPE, torch_compile=False,
        batching=BatchingConfig(max_batch_seconds=CPU_MAX_BATCH_SECONDS),
    )


def select_deploy() -> DeployConfig:
    """CPU when the request forced it or no CUDA is visible; GPU otherwise."""
    from src.core.zero_gpu import is_user_forced_cpu
    if is_user_forced_cpu():
        return cpu_deploy()
    try:
        import torch
        if not torch.cuda.is_available():
            return cpu_deploy()
    except Exception:
        return cpu_deploy()
    return gpu_deploy()