feat: add Crusoe managed inference backend (#20475)
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Sonnet 4.6
parent
49ac447c94
commit
b49d05fd0e
@@ -69,6 +69,7 @@ from sglang.utils import LazyImport
|
|||||||
from sglang.version import __version__
|
from sglang.version import __version__
|
||||||
|
|
||||||
Anthropic = LazyImport("sglang.lang.backend.anthropic", "Anthropic")
|
Anthropic = LazyImport("sglang.lang.backend.anthropic", "Anthropic")
|
||||||
|
Crusoe = LazyImport("sglang.lang.backend.crusoe", "Crusoe")
|
||||||
LiteLLM = LazyImport("sglang.lang.backend.litellm", "LiteLLM")
|
LiteLLM = LazyImport("sglang.lang.backend.litellm", "LiteLLM")
|
||||||
OpenAI = LazyImport("sglang.lang.backend.openai", "OpenAI")
|
OpenAI = LazyImport("sglang.lang.backend.openai", "OpenAI")
|
||||||
VertexAI = LazyImport("sglang.lang.backend.vertexai", "VertexAI")
|
VertexAI = LazyImport("sglang.lang.backend.vertexai", "VertexAI")
|
||||||
@@ -106,6 +107,7 @@ __all__ = [
|
|||||||
"unconditional_likelihood_normalized",
|
"unconditional_likelihood_normalized",
|
||||||
"ServerArgs",
|
"ServerArgs",
|
||||||
"Anthropic",
|
"Anthropic",
|
||||||
|
"Crusoe",
|
||||||
"LiteLLM",
|
"LiteLLM",
|
||||||
"OpenAI",
|
"OpenAI",
|
||||||
"VertexAI",
|
"VertexAI",
|
||||||
|
|||||||
@@ -0,0 +1,43 @@
|
|||||||
|
import os
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
|
from sglang.lang.backend.openai import OpenAI
|
||||||
|
from sglang.lang.chat_template import ChatTemplate
|
||||||
|
|
||||||
|
CRUSOE_BASE_URL = "https://managed-inference-api-proxy.crusoecloud.com/v1/"
|
||||||
|
|
||||||
|
|
||||||
|
class Crusoe(OpenAI):
|
||||||
|
"""SGLang backend for Crusoe managed inference.
|
||||||
|
|
||||||
|
Crusoe exposes an OpenAI-compatible API, so this is a thin wrapper
|
||||||
|
around the OpenAI backend that handles Crusoe-specific defaults.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
model_name: The model to use, e.g. "meta-llama/Llama-3.1-8B-Instruct".
|
||||||
|
api_key: Crusoe API key. Defaults to CRUSOE_API_KEY env var.
|
||||||
|
base_url: Override the Crusoe endpoint. Defaults to the Crusoe API.
|
||||||
|
chat_template: Optional custom chat template.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
model_name: str,
|
||||||
|
api_key: Optional[str] = None,
|
||||||
|
base_url: Optional[str] = None,
|
||||||
|
chat_template: Optional[ChatTemplate] = None,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
resolved_api_key = api_key or os.environ.get("CRUSOE_API_KEY")
|
||||||
|
if not resolved_api_key:
|
||||||
|
raise ValueError(
|
||||||
|
"Crusoe API key required. Pass api_key= or set CRUSOE_API_KEY."
|
||||||
|
)
|
||||||
|
|
||||||
|
super().__init__(
|
||||||
|
model_name=model_name,
|
||||||
|
chat_template=chat_template,
|
||||||
|
api_key=resolved_api_key,
|
||||||
|
base_url=base_url or CRUSOE_BASE_URL,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
@@ -0,0 +1,79 @@
|
|||||||
|
"""
|
||||||
|
Manual tests for the Crusoe managed inference backend.
|
||||||
|
|
||||||
|
Requires CRUSOE_API_KEY to be set in the environment.
|
||||||
|
|
||||||
|
Run all tests:
|
||||||
|
python3 -m unittest test/manual/test_crusoe_backend.py
|
||||||
|
|
||||||
|
Run a single test:
|
||||||
|
python3 -m unittest test_crusoe_backend.TestCrusoeBackend.test_mt_bench
|
||||||
|
"""
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
from sglang import Crusoe, set_default_backend
|
||||||
|
from sglang.test.test_programs import (
|
||||||
|
test_mt_bench,
|
||||||
|
test_parallel_decoding,
|
||||||
|
test_parallel_encoding,
|
||||||
|
test_stream,
|
||||||
|
)
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
# Default model available on Crusoe managed inference.
|
||||||
|
DEFAULT_CRUSOE_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B"
|
||||||
|
|
||||||
|
|
||||||
|
class TestCrusoeBackend(CustomTestCase):
|
||||||
|
backend = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def setUpClass(cls):
|
||||||
|
cls.backend = Crusoe(DEFAULT_CRUSOE_MODEL)
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
set_default_backend(self.backend)
|
||||||
|
|
||||||
|
def test_mt_bench(self):
|
||||||
|
test_mt_bench()
|
||||||
|
|
||||||
|
def test_stream(self):
|
||||||
|
test_stream()
|
||||||
|
|
||||||
|
def test_parallel_decoding(self):
|
||||||
|
test_parallel_decoding()
|
||||||
|
|
||||||
|
def test_parallel_encoding(self):
|
||||||
|
test_parallel_encoding()
|
||||||
|
|
||||||
|
|
||||||
|
class TestCrusoeBackendInit(CustomTestCase):
|
||||||
|
"""Unit tests for Crusoe backend initialisation — no network required."""
|
||||||
|
|
||||||
|
def test_raises_without_api_key(self):
|
||||||
|
import os
|
||||||
|
|
||||||
|
key = os.environ.pop("CRUSOE_API_KEY", None)
|
||||||
|
try:
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
Crusoe(DEFAULT_CRUSOE_MODEL, api_key=None)
|
||||||
|
finally:
|
||||||
|
if key is not None:
|
||||||
|
os.environ["CRUSOE_API_KEY"] = key
|
||||||
|
|
||||||
|
def test_accepts_explicit_api_key(self):
|
||||||
|
backend = Crusoe(DEFAULT_CRUSOE_MODEL, api_key="test-key")
|
||||||
|
self.assertIsNotNone(backend)
|
||||||
|
|
||||||
|
def test_custom_base_url(self):
|
||||||
|
backend = Crusoe(
|
||||||
|
DEFAULT_CRUSOE_MODEL,
|
||||||
|
api_key="test-key",
|
||||||
|
base_url="https://managed-inference-api-proxy.crusoecloud.com/v1/",
|
||||||
|
)
|
||||||
|
self.assertIsNotNone(backend)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user