feat: sglang backend support

1 files changed, 207 insertions, 0 deletions

diff --git a/app/backends/sglomni/status.py b/app/backends/sglomni/status.py
new file mode 100644
index 0000000..099817e
--- /dev/null
+++ b/app/backends/sglomni/status.py
@@ -0,0 +1,207 @@
+"""detect() for the hub's backend menu: how far sglang-omni is set up."""
+
+from pathlib import Path
+from typing import List, Optional
+
+from backends import BackendStatus, ServerSpec, envs, format_launch_hint, \
+    probe, servers
+from backends.sglomni import gpu
+from backends.sglomni.catalog import ModelEntry, entry_by_repo, \
+    fallback_config_path, config_path
+from backends.sglomni.constants import DEFAULT_PORT, SERVER_NAME, \
+    SERVER_START_TIMEOUT, SGLOMNI_PIP_PKG
+from backends.sglomni.models import installed_entries
+from backends.sglomni.pythonenv import SGLOMNI_ENV, env_version
+from converter import config
+
+
+def _is_installed() -> bool:
+    if envs.env_script("sgl-omni", SGLOMNI_ENV).is_file():
+        return True
+    return envs.module_available("sglang_omni", SGLOMNI_ENV)
+
+
+def build_spec(entry: ModelEntry) -> ServerSpec:
+    """The managed ServerSpec hosting ENTRY on the configured port.
+
+    Public because the run preparation (hub) and the CLI's managed-server
+    bootstrap need to boot exactly the model their run selected, which can
+    differ from detect()'s default. The config yaml is the vendored copy
+    from the catalog; models without one (Higgs, ZONOS2) run from
+    --model-path alone — except that a GPU the FP8 kernels cannot run
+    launches ZONOS2's vendored bf16 config instead (launch_config_path).
+    """
+    url = config.SGLOMNI_API_URL
+    argv: List[str] = [
+        str(envs.env_script("sgl-omni", SGLOMNI_ENV)),
+        "serve", "--model-path", entry.repo,
+    ]
+    yaml_path = launch_config_path(entry)
+    if yaml_path is not None:
+        argv += ["--config", str(yaml_path)]
+    argv += ["--port", str(_port())]
+    return ServerSpec(SERVER_NAME, url, argv,
+                      identity=probe.IDENTITY_SGLOMNI,
+                      start_timeout=SERVER_START_TIMEOUT)
+
+
+def needs_fp8_fallback(entry: ModelEntry) -> bool:
+    """True when ENTRY's default FP8-quantized pipeline cannot run here.
+
+    Only models with a declared fp8_min_compute_capability are candidates,
+    and only when an NVIDIA GPU actually answers: a GPU this tool cannot
+    read keeps upstream defaults rather than second-guessing the host."""
+    if not (entry.fp8_moe and entry.fp8_min_compute_capability):
+        return False
+    capability = gpu.compute_capability()
+    if capability is None:
+        return False
+    return capability < entry.fp8_min_compute_capability
+
+
+def launch_config_path(entry: ModelEntry) -> Optional[Path]:
+    """The vendored config yaml ENTRY's server should launch with (None =
+    --model-path alone, the upstream default).
+
+    A model whose default pipeline quantizes its MoE experts to FP8
+    launches its vendored bf16 config instead on GPUs the FP8 Triton
+    kernels cannot compile on (needs_fp8_fallback) — the server then runs
+    the model in bf16 at about twice the MoE VRAM. A missing fallback
+    file degrades to the model's normal config (the server reports the
+    FP8 failure itself; the boot-failure hint names it)."""
+    if needs_fp8_fallback(entry) and entry.bf16_config:
+        path = fallback_config_path(entry)
+        if path is not None and path.is_file():
+            return path
+    return config_path(entry)
+
+
+def gpu_fallback_note(entry: ModelEntry) -> Optional[str]:
+    """A human-readable note when ENTRY launches its bf16 fallback here.
+
+    Printed by the install and boot flows so the (small) VRAM cost and
+    the reason are on the record before the server starts. None when the
+    model runs its default (FP8) pipeline, or when no GPU answered."""
+    if not needs_fp8_fallback(entry):
+        return None
+    minimum = (".".join(str(part)
+                        for part in entry.fp8_min_compute_capability))
+    where = gpu.describe() or "unknown GPU"
+    path = fallback_config_path(entry)
+    if path is None or not path.is_file():
+        return (f"{entry.label}'s default pipeline quantizes its MoE "
+                f"experts to FP8, which needs compute capability "
+                f"{minimum}+; this GPU ({where}) cannot run it, and the "
+                "vendored bf16 fallback config is missing — the server "
+                "will fail to start this model.")
+    return (f"{entry.label}'s default pipeline quantizes its MoE experts "
+            f"to FP8, which needs compute capability {minimum}+; this GPU "
+            f"({where}) cannot run it — launching the vendored bf16 "
+            "config instead (about twice the MoE VRAM).")
+
+
+def _port() -> int:
+    return _port_of(config.SGLOMNI_API_URL)
+
+
+def _port_of(url: str) -> int:
+    import urllib.parse
+    try:
+        return urllib.parse.urlsplit(url).port or DEFAULT_PORT
+    except ValueError:
+        return DEFAULT_PORT
+
+
+def _managed_running_entry() -> Optional[ModelEntry]:
+    """The catalog model a locally-managed, running server hosts."""
+    if servers.pid_for(SERVER_NAME) is None:
+        return None
+    if not servers.alive(SERVER_NAME):
+        return None
+    return entry_by_repo(
+        probe.sglomni_served_model(config.SGLOMNI_API_URL))
+
+
+def detect() -> BackendStatus:
+    """Detect whether sglang-omni is installed, plus the launch command.
+
+    One managed spec exists per detection, hosting the first *installed*
+    catalog model (catalog order) on the single configured port — runs
+    needing another model boot it via their own spec (the Generate form's
+    Model menu / CLI --model), restarting a managed server that hosts
+    something else. Which model currently answers is read via the probe
+    (local pid alive → our URL; otherwise the remote URL) so the status
+    names the *running* model even when it differs from the default one.
+    """
+    installed = _is_installed()
+    present = installed_entries()
+    configured = installed and bool(present)
+    details: List[str] = []
+    details.append("pip: installed" if installed else
+                   f"not installed — run setup to pip install "
+                   f"{SGLOMNI_PIP_PKG}")
+    if installed:
+        version = env_version()
+        if version is not None:
+            details.append(f"python: {version[0]}.{version[1]}")
+        else:
+            details.append("python: unknown version")
+    if present:
+        # A model launching its bf16 fallback on this GPU says so, so the
+        # status line explains why its server boots with a config file.
+        def _model_tag(model_entry: ModelEntry) -> str:
+            if needs_fp8_fallback(model_entry):
+                return f"{model_entry.key} (bf16 fallback)"
+            return model_entry.key
+        details.append(f"models: {', '.join(_model_tag(e) for e in present)}")
+    else:
+        details.append("no models downloaded — run setup (or Configure) to "
+                       "install one")
+    details.append(f"port: {_port()}")
+    specs: List[ServerSpec] = [build_spec(present[0])] if configured else []
+    managed = servers.manages(specs)
+    local_models: List[str] = []
+    if managed and specs and servers.alive(specs[0].name):
+        entry = _managed_running_entry()
+        if entry is not None:
+            local_models.append(entry.label)
+    remote_models, remote_urls = _detect_remote(managed)
+    running_models = list(dict.fromkeys(local_models + remote_models))
+    return BackendStatus(
+        SERVER_NAME, "SGLang-Omni",
+        installed=installed, configured=configured,
+        running=managed or bool(remote_urls),
+        details=details,
+        launch_hint=format_launch_hint(specs),
+        servers=specs,
+        managed=managed,
+        remote=bool(remote_urls),
+        remote_urls=remote_urls,
+        remote_models=remote_models,
+        running_models=running_models,
+        partial="installed (no models)" if installed and not configured
+        else "")
+
+
+def _detect_remote(managed: bool = False):
+    """Detect an externally-run sglang-omni server at the remote URL.
+
+    Returns ``([model_label, ...], {spec_name: url})``. The remote URL must
+    answer as sglang-omni (probe identity); when it equals the local URL
+    and this tool started that server, it is ignored (already reported as
+    "[local]").
+    """
+    remote_models: List[str] = []
+    remote_urls: dict = {}
+    url = (config.SGLOMNI_REMOTE_URL or "").strip()
+    if not url:
+        return remote_models, remote_urls
+    if managed and probe.same_endpoint(url, config.SGLOMNI_API_URL):
+        return remote_models, remote_urls
+    if probe.identify_server(url) != probe.IDENTITY_SGLOMNI:
+        return remote_models, remote_urls
+    remote_urls[SERVER_NAME] = url
+    entry = entry_by_repo(probe.sglomni_served_model(url))
+    remote_models.append(entry.label if entry
+                         else probe.sglomni_served_model(url) or "unknown")
+    return remote_models, remote_urls