move chunks to core
This commit is contained in:
Evan
@@ -7,7 +7,9 @@ authors = [
|
||||
{ name = "Evan", email = "[email protected]" }
|
||||
]
|
||||
requires-python = ">=3.13"
|
||||
dependencies = []
|
||||
dependencies = [
|
||||
"pydantic",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["uv_build>=0.9.24,<0.10.0"]
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
from enum import Enum
|
||||
from typing import Annotated, Any
|
||||
|
||||
import aiofiles
|
||||
import aiofiles.os as aios
|
||||
import tomlkit
|
||||
from anyio import Path, open_file
|
||||
from huggingface_hub import model_info
|
||||
from loguru import logger
|
||||
from pydantic import (
|
||||
AliasChoices,
|
||||
BaseModel,
|
||||
Field,
|
||||
PositiveInt,
|
||||
ValidationError,
|
||||
field_validator,
|
||||
model_validator,
|
||||
)
|
||||
from tomlkit.exceptions import TOMLKitError
|
||||
|
||||
from exo.shared.constants import (
|
||||
EXO_CUSTOM_MODEL_CARDS_DIR,
|
||||
EXO_ENABLE_IMAGE_MODELS,
|
||||
RESOURCES_DIR,
|
||||
)
|
||||
from exo.shared.types.common import ModelId
|
||||
from exo.shared.types.memory import Memory
|
||||
from exo.utils.pydantic_ext import CamelCaseModel
|
||||
|
||||
# kinda ugly...
|
||||
# TODO: load search path from config.toml
|
||||
_custom_cards_dir = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR))
|
||||
CARD_SEARCH_PATH = [
|
||||
Path(RESOURCES_DIR) / "inference_model_cards",
|
||||
Path(RESOURCES_DIR) / "image_model_cards",
|
||||
_custom_cards_dir,
|
||||
]
|
||||
|
||||
_card_cache: dict[ModelId, "ModelCard"] = {}
|
||||
|
||||
|
||||
async def _refresh_card_cache():
|
||||
for path in CARD_SEARCH_PATH:
|
||||
async for toml_file in path.rglob("*.toml"):
|
||||
try:
|
||||
card = await ModelCard.load_from_path(toml_file)
|
||||
if card.model_id not in _card_cache:
|
||||
_card_cache[card.model_id] = card
|
||||
except (ValidationError, TOMLKitError):
|
||||
pass
|
||||
|
||||
|
||||
def _is_image_card(card: "ModelCard") -> bool:
|
||||
return any(t in (ModelTask.TextToImage, ModelTask.ImageToImage) for t in card.tasks)
|
||||
|
||||
|
||||
async def get_model_cards() -> list["ModelCard"]:
|
||||
if len(_card_cache) == 0:
|
||||
await _refresh_card_cache()
|
||||
if EXO_ENABLE_IMAGE_MODELS:
|
||||
return list(_card_cache.values())
|
||||
return [c for c in _card_cache.values() if not _is_image_card(c)]
|
||||
|
||||
|
||||
class ModelTask(str, Enum):
|
||||
TextGeneration = "TextGeneration"
|
||||
TextToImage = "TextToImage"
|
||||
ImageToImage = "ImageToImage"
|
||||
|
||||
|
||||
class ComponentInfo(CamelCaseModel):
|
||||
component_name: str
|
||||
component_path: str
|
||||
storage_size: Memory
|
||||
n_layers: PositiveInt | None = None
|
||||
can_shard: bool
|
||||
safetensors_index_filename: str | None = None
|
||||
|
||||
|
||||
class ModelCard(CamelCaseModel):
|
||||
model_id: ModelId
|
||||
storage_size: Memory
|
||||
n_layers: PositiveInt
|
||||
hidden_size: PositiveInt
|
||||
supports_tensor: bool
|
||||
num_key_value_heads: PositiveInt | None = None
|
||||
tasks: list[ModelTask]
|
||||
components: list[ComponentInfo] | None = None
|
||||
family: str = ""
|
||||
quantization: str = ""
|
||||
base_model: str = ""
|
||||
capabilities: list[str] = []
|
||||
uses_cfg: bool = False
|
||||
trust_remote_code: bool = True
|
||||
|
||||
@field_validator("tasks", mode="before")
|
||||
@classmethod
|
||||
def _validate_tasks(cls, v: list[str | ModelTask]) -> list[ModelTask]:
|
||||
return [item if isinstance(item, ModelTask) else ModelTask(item) for item in v]
|
||||
|
||||
async def save(self, path: Path) -> None:
|
||||
async with await open_file(path, "w") as f:
|
||||
py = self.model_dump(exclude_none=True)
|
||||
data = tomlkit.dumps(py) # pyright: ignore[reportUnknownMemberType]
|
||||
await f.write(data)
|
||||
|
||||
async def save_to_custom_dir(self) -> None:
|
||||
await aios.makedirs(str(_custom_cards_dir), exist_ok=True)
|
||||
await self.save(_custom_cards_dir / (self.model_id.normalize() + ".toml"))
|
||||
|
||||
@staticmethod
|
||||
async def load_from_path(path: Path) -> "ModelCard":
|
||||
async with await open_file(path, "r") as f:
|
||||
py = tomlkit.loads(await f.read())
|
||||
return ModelCard.model_validate(py)
|
||||
|
||||
# Is it okay that model card.load defaults to network access if the card doesn't exist? do we want to be more explicit here?
|
||||
@staticmethod
|
||||
async def load(model_id: ModelId) -> "ModelCard":
|
||||
if model_id not in _card_cache:
|
||||
await _refresh_card_cache()
|
||||
if (mc := _card_cache.get(model_id)) is not None:
|
||||
return mc
|
||||
|
||||
return await ModelCard.fetch_from_hf(model_id)
|
||||
|
||||
@staticmethod
|
||||
async def fetch_from_hf(model_id: ModelId) -> "ModelCard":
|
||||
"""Fetches storage size and number of layers for a Hugging Face model, returns Pydantic ModelMeta."""
|
||||
# TODO: failure if files do not exist
|
||||
config_data = await fetch_config_data(model_id)
|
||||
num_layers = config_data.layer_count
|
||||
mem_size_bytes = await fetch_safetensors_size(model_id)
|
||||
|
||||
mc = ModelCard(
|
||||
model_id=ModelId(model_id),
|
||||
storage_size=mem_size_bytes,
|
||||
n_layers=num_layers,
|
||||
hidden_size=config_data.hidden_size or 0,
|
||||
supports_tensor=config_data.supports_tensor,
|
||||
num_key_value_heads=config_data.num_key_value_heads,
|
||||
tasks=[ModelTask.TextGeneration],
|
||||
trust_remote_code=False,
|
||||
)
|
||||
await mc.save_to_custom_dir()
|
||||
_card_cache[model_id] = mc
|
||||
return mc
|
||||
|
||||
|
||||
async def delete_custom_card(model_id: ModelId) -> bool:
|
||||
"""Delete a user-added custom model card. Returns True if deleted."""
|
||||
card_path = _custom_cards_dir / (ModelId(model_id).normalize() + ".toml")
|
||||
if await card_path.exists():
|
||||
await card_path.unlink()
|
||||
_card_cache.pop(model_id, None)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_custom_card(model_id: ModelId) -> bool:
|
||||
"""Check if a model card exists in the custom cards directory."""
|
||||
import os
|
||||
|
||||
card_path = Path(str(EXO_CUSTOM_MODEL_CARDS_DIR)) / (
|
||||
ModelId(model_id).normalize() + ".toml"
|
||||
)
|
||||
return os.path.isfile(str(card_path))
|
||||
|
||||
|
||||
class ConfigData(BaseModel):
|
||||
model_config = {"extra": "ignore"} # Allow unknown fields
|
||||
|
||||
architectures: list[str] | None = None
|
||||
hidden_size: Annotated[int, Field(ge=0)] | None = None
|
||||
num_key_value_heads: PositiveInt | None = None
|
||||
layer_count: int = Field(
|
||||
validation_alias=AliasChoices(
|
||||
"num_hidden_layers",
|
||||
"num_layers",
|
||||
"n_layer",
|
||||
"n_layers",
|
||||
"num_decoder_layers",
|
||||
"decoder_layers",
|
||||
)
|
||||
)
|
||||
|
||||
@property
|
||||
def supports_tensor(self) -> bool:
|
||||
return self.architectures in [
|
||||
["Glm4MoeLiteForCausalLM"],
|
||||
["GlmMoeDsaForCausalLM"],
|
||||
["DeepseekV32ForCausalLM"],
|
||||
["DeepseekV3ForCausalLM"],
|
||||
["Qwen3NextForCausalLM"],
|
||||
["Qwen3MoeForCausalLM"],
|
||||
["Qwen3_5MoeForConditionalGeneration"],
|
||||
["Qwen3_5ForConditionalGeneration"],
|
||||
["MiniMaxM2ForCausalLM"],
|
||||
["LlamaForCausalLM"],
|
||||
["GptOssForCausalLM"],
|
||||
["Step3p5ForCausalLM"],
|
||||
["NemotronHForCausalLM"],
|
||||
]
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def defer_to_text_config(cls, data: dict[str, Any]):
|
||||
text_config = data.get("text_config")
|
||||
if text_config is None:
|
||||
return data
|
||||
|
||||
for field in [
|
||||
"architectures",
|
||||
"hidden_size",
|
||||
"num_key_value_heads",
|
||||
"num_hidden_layers",
|
||||
"num_layers",
|
||||
"n_layer",
|
||||
"n_layers",
|
||||
"num_decoder_layers",
|
||||
"decoder_layers",
|
||||
]:
|
||||
if (val := text_config.get(field)) is not None: # pyright: ignore[reportAny]
|
||||
data[field] = val
|
||||
|
||||
return data
|
||||
|
||||
|
||||
async def fetch_config_data(model_id: ModelId) -> ConfigData:
|
||||
"""Downloads and parses config.json for a model."""
|
||||
from exo.download.download_utils import (
|
||||
download_file_with_retry,
|
||||
ensure_models_dir,
|
||||
)
|
||||
|
||||
target_dir = (await ensure_models_dir()) / model_id.normalize()
|
||||
await aios.makedirs(target_dir, exist_ok=True)
|
||||
config_path = await download_file_with_retry(
|
||||
model_id,
|
||||
"main",
|
||||
"config.json",
|
||||
target_dir,
|
||||
lambda curr_bytes, total_bytes, is_renamed: logger.debug(
|
||||
f"Downloading config.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
|
||||
),
|
||||
)
|
||||
async with aiofiles.open(config_path, "r") as f:
|
||||
return ConfigData.model_validate_json(await f.read())
|
||||
|
||||
|
||||
async def fetch_safetensors_size(model_id: ModelId) -> Memory:
|
||||
"""Gets model size from safetensors index or falls back to HF API."""
|
||||
from exo.download.download_utils import (
|
||||
download_file_with_retry,
|
||||
ensure_models_dir,
|
||||
)
|
||||
from exo.shared.types.worker.downloads import ModelSafetensorsIndex
|
||||
|
||||
target_dir = (await ensure_models_dir()) / model_id.normalize()
|
||||
await aios.makedirs(target_dir, exist_ok=True)
|
||||
try:
|
||||
index_path = await download_file_with_retry(
|
||||
model_id,
|
||||
"main",
|
||||
"model.safetensors.index.json",
|
||||
target_dir,
|
||||
lambda curr_bytes, total_bytes, is_renamed: logger.debug(
|
||||
f"Downloading model.safetensors.index.json for {model_id}: {curr_bytes}/{total_bytes} ({is_renamed=})"
|
||||
),
|
||||
)
|
||||
async with aiofiles.open(index_path, "r") as f:
|
||||
index_data = ModelSafetensorsIndex.model_validate_json(await f.read())
|
||||
|
||||
metadata = index_data.metadata
|
||||
if metadata is not None:
|
||||
return Memory.from_bytes(metadata.total_size)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
info = model_info(model_id)
|
||||
if info.safetensors is None:
|
||||
raise ValueError(f"No safetensors info found for {model_id}")
|
||||
return Memory.from_bytes(info.safetensors.total)
|
||||
@@ -0,0 +1,51 @@
|
||||
# pyright: reportAny=false, reportUnknownArgumentType=false, reportUnknownVariableType=false
|
||||
|
||||
from typing import Any, Self
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, model_serializer, model_validator
|
||||
from pydantic.alias_generators import to_camel
|
||||
from pydantic_core.core_schema import (
|
||||
SerializerFunctionWrapHandler,
|
||||
ValidatorFunctionWrapHandler,
|
||||
)
|
||||
|
||||
|
||||
class CamelCaseModel(BaseModel):
|
||||
"""
|
||||
A model whose fields are aliased to camel-case from snake-case.
|
||||
"""
|
||||
|
||||
model_config = ConfigDict(
|
||||
alias_generator=to_camel,
|
||||
validate_by_name=True,
|
||||
extra="forbid",
|
||||
strict=True,
|
||||
)
|
||||
|
||||
|
||||
class FrozenModel(BaseModel):
|
||||
model_config = ConfigDict(
|
||||
alias_generator=to_camel,
|
||||
validate_by_name=True,
|
||||
extra="forbid",
|
||||
strict=True,
|
||||
frozen=True,
|
||||
)
|
||||
|
||||
|
||||
class TaggedModel(CamelCaseModel):
|
||||
@model_serializer(mode="wrap")
|
||||
def _serialize(self, handler: SerializerFunctionWrapHandler):
|
||||
inner = handler(self)
|
||||
return {self.__class__.__name__: inner}
|
||||
|
||||
@model_validator(mode="wrap")
|
||||
@classmethod
|
||||
def _validate(cls, v: Any, handler: ValidatorFunctionWrapHandler) -> Self:
|
||||
if isinstance(v, dict) and len(v) == 1 and cls.__name__ in v:
|
||||
return handler(v[cls.__name__])
|
||||
|
||||
return handler(v)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}({super().__str__()})"
|
||||
@@ -0,0 +1,89 @@
|
||||
from collections.abc import Generator
|
||||
from typing import Any, Literal
|
||||
|
||||
from exo.api.types import (
|
||||
FinishReason,
|
||||
GenerationStats,
|
||||
ImageGenerationStats,
|
||||
ToolCallItem,
|
||||
TopLogprobItem,
|
||||
Usage,
|
||||
)
|
||||
from exo_core.model_cards import ModelId
|
||||
from exo_core.pydantic import TaggedModel
|
||||
|
||||
from .common import CommandId
|
||||
|
||||
|
||||
class BaseChunk(TaggedModel):
|
||||
model: ModelId
|
||||
|
||||
|
||||
class TokenChunk(BaseChunk):
|
||||
text: str
|
||||
token_id: int
|
||||
usage: Usage | None
|
||||
finish_reason: Literal["stop", "length", "content_filter"] | None = None
|
||||
stats: GenerationStats | None = None
|
||||
logprob: float | None = None
|
||||
top_logprobs: list[TopLogprobItem] | None = None
|
||||
is_thinking: bool = False
|
||||
|
||||
|
||||
class ErrorChunk(BaseChunk):
|
||||
error_message: str
|
||||
finish_reason: Literal["error"] = "error"
|
||||
|
||||
|
||||
class ToolCallChunk(BaseChunk):
|
||||
tool_calls: list[ToolCallItem]
|
||||
usage: Usage | None
|
||||
finish_reason: Literal["tool_calls"] = "tool_calls"
|
||||
stats: GenerationStats | None = None
|
||||
|
||||
|
||||
class ImageChunk(BaseChunk):
|
||||
data: str
|
||||
chunk_index: int
|
||||
total_chunks: int
|
||||
image_index: int
|
||||
is_partial: bool = False
|
||||
partial_index: int | None = None
|
||||
total_partials: int | None = None
|
||||
stats: ImageGenerationStats | None = None
|
||||
format: Literal["png", "jpeg", "webp"] | None = None
|
||||
finish_reason: FinishReason | None = None
|
||||
error_message: str | None = None
|
||||
|
||||
def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
|
||||
for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
|
||||
if name == "data" and hasattr(value, "__len__"): # pyright: ignore[reportAny]
|
||||
yield name, f"<{len(self.data)} chars>"
|
||||
elif name is not None:
|
||||
yield name, value
|
||||
|
||||
|
||||
class InputImageChunk(BaseChunk):
|
||||
command_id: CommandId
|
||||
data: str
|
||||
chunk_index: int
|
||||
total_chunks: int
|
||||
|
||||
def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
|
||||
for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
|
||||
if name == "data" and hasattr(value, "__len__"): # pyright: ignore[reportAny]
|
||||
yield name, f"<{len(self.data)} chars>"
|
||||
elif name is not None:
|
||||
yield name, value
|
||||
|
||||
|
||||
class PrefillProgressChunk(BaseChunk):
|
||||
"""Data class for prefill progress events during streaming."""
|
||||
|
||||
processed_tokens: int
|
||||
total_tokens: int
|
||||
|
||||
|
||||
GenerationChunk = (
|
||||
TokenChunk | ImageChunk | ToolCallChunk | ErrorChunk | PrefillProgressChunk
|
||||
)
|
||||
@@ -0,0 +1,61 @@
|
||||
from typing import Self
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import GetCoreSchemaHandler, field_validator
|
||||
from pydantic_core import core_schema
|
||||
|
||||
from exo_core.pydantic import CamelCaseModel
|
||||
|
||||
|
||||
class Id(str):
|
||||
def __new__(cls, value: str | None = None) -> Self:
|
||||
return super().__new__(cls, value or str(uuid4()))
|
||||
|
||||
@classmethod
|
||||
def __get_pydantic_core_schema__(
|
||||
cls, _source: type, handler: GetCoreSchemaHandler
|
||||
) -> core_schema.CoreSchema:
|
||||
# Just use a plain string schema
|
||||
return core_schema.no_info_after_validator_function(
|
||||
cls, core_schema.str_schema()
|
||||
)
|
||||
|
||||
|
||||
class NodeId(Id):
|
||||
pass
|
||||
|
||||
|
||||
class SystemId(Id):
|
||||
pass
|
||||
|
||||
|
||||
class ModelId(Id):
|
||||
def normalize(self) -> str:
|
||||
return self.replace("/", "--")
|
||||
|
||||
def short(self) -> str:
|
||||
return self.split("/")[-1]
|
||||
|
||||
|
||||
class SessionId(CamelCaseModel):
|
||||
master_node_id: NodeId
|
||||
election_clock: int
|
||||
|
||||
|
||||
class CommandId(Id):
|
||||
pass
|
||||
|
||||
|
||||
class Host(CamelCaseModel):
|
||||
ip: str
|
||||
port: int
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.ip}:{self.port}"
|
||||
|
||||
@field_validator("port")
|
||||
@classmethod
|
||||
def check_port(cls, v: int) -> int:
|
||||
if not (0 <= v <= 65535):
|
||||
raise ValueError("Port must be between 0 and 65535")
|
||||
return v
|
||||
@@ -0,0 +1,72 @@
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import model_validator
|
||||
|
||||
from exo_core.model_cards import ModelTask
|
||||
from exo.shared.types.common import Host, Id, NodeId
|
||||
from exo.shared.types.worker.runners import RunnerId, ShardAssignments, ShardMetadata
|
||||
from exo_core.pydantic_ext import CamelCaseModel, TaggedModel
|
||||
|
||||
|
||||
class InstanceId(Id):
|
||||
pass
|
||||
|
||||
|
||||
class InstanceMeta(str, Enum):
|
||||
MlxRing = "MlxRing"
|
||||
MlxJaccl = "MlxJaccl"
|
||||
Vllm = "Vllm"
|
||||
|
||||
|
||||
class BaseInstance(TaggedModel):
|
||||
instance_id: InstanceId
|
||||
shard_assignments: ShardAssignments
|
||||
|
||||
def shard(self, runner_id: RunnerId) -> ShardMetadata | None:
|
||||
return self.shard_assignments.runner_to_shard.get(runner_id, None)
|
||||
|
||||
|
||||
class MlxRingInstance(BaseInstance):
|
||||
hosts_by_node: dict[NodeId, list[Host]]
|
||||
ephemeral_port: int
|
||||
|
||||
|
||||
class MlxJacclInstance(BaseInstance):
|
||||
jaccl_devices: list[list[str | None]]
|
||||
jaccl_coordinators: dict[NodeId, str]
|
||||
|
||||
|
||||
class VllmInstance(BaseInstance):
|
||||
pass
|
||||
|
||||
|
||||
# TODO: Single node instance
|
||||
Instance = MlxRingInstance | MlxJacclInstance | VllmInstance
|
||||
|
||||
|
||||
class BoundInstance(CamelCaseModel):
|
||||
instance: Instance
|
||||
bound_runner_id: RunnerId
|
||||
bound_node_id: NodeId
|
||||
|
||||
@property
|
||||
def bound_shard(self) -> ShardMetadata:
|
||||
shard = self.instance.shard(self.bound_runner_id)
|
||||
assert shard is not None
|
||||
return shard
|
||||
|
||||
@property
|
||||
def is_image_model(self) -> bool:
|
||||
return (
|
||||
ModelTask.TextToImage in self.bound_shard.model_card.tasks
|
||||
or ModelTask.ImageToImage in self.bound_shard.model_card.tasks
|
||||
)
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_shard_exists(self) -> "BoundInstance":
|
||||
assert (
|
||||
self.bound_runner_id in self.instance.shard_assignments.runner_to_shard
|
||||
), (
|
||||
"Bound Instance must be constructed with a runner_id that is in the instances assigned shards"
|
||||
)
|
||||
return self
|
||||
@@ -0,0 +1,75 @@
|
||||
from collections.abc import Generator
|
||||
from typing import Any, Literal
|
||||
|
||||
from exo.api.types import (
|
||||
FinishReason,
|
||||
GenerationStats,
|
||||
ImageGenerationStats,
|
||||
ToolCallItem,
|
||||
TopLogprobItem,
|
||||
Usage,
|
||||
)
|
||||
from exo.utils.pydantic_ext import TaggedModel
|
||||
|
||||
|
||||
class BaseRunnerResponse(TaggedModel):
|
||||
pass
|
||||
|
||||
|
||||
class TokenizedResponse(BaseRunnerResponse):
|
||||
prompt_tokens: int
|
||||
|
||||
|
||||
class GenerationResponse(BaseRunnerResponse):
|
||||
text: str
|
||||
token: int
|
||||
logprob: float | None = None
|
||||
top_logprobs: list[TopLogprobItem] | None = None
|
||||
finish_reason: FinishReason | None = None
|
||||
stats: GenerationStats | None = None
|
||||
usage: Usage | None
|
||||
is_thinking: bool = False
|
||||
|
||||
|
||||
class ImageGenerationResponse(BaseRunnerResponse):
|
||||
image_data: bytes
|
||||
format: Literal["png", "jpeg", "webp"] = "png"
|
||||
stats: ImageGenerationStats | None = None
|
||||
image_index: int = 0
|
||||
|
||||
def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
|
||||
for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
|
||||
if name == "image_data":
|
||||
yield name, f"<{len(self.image_data)} bytes>"
|
||||
elif name is not None:
|
||||
yield name, value
|
||||
|
||||
|
||||
class PartialImageResponse(BaseRunnerResponse):
|
||||
image_data: bytes
|
||||
format: Literal["png", "jpeg", "webp"] = "png"
|
||||
partial_index: int
|
||||
total_partials: int
|
||||
image_index: int = 0
|
||||
|
||||
def __repr_args__(self) -> Generator[tuple[str, Any], None, None]:
|
||||
for name, value in super().__repr_args__(): # pyright: ignore[reportAny]
|
||||
if name == "image_data":
|
||||
yield name, f"<{len(self.image_data)} bytes>"
|
||||
elif name is not None:
|
||||
yield name, value
|
||||
|
||||
|
||||
class ToolCallResponse(BaseRunnerResponse):
|
||||
tool_calls: list[ToolCallItem]
|
||||
usage: Usage | None
|
||||
stats: GenerationStats | None = None
|
||||
|
||||
|
||||
class FinishedResponse(BaseRunnerResponse):
|
||||
pass
|
||||
|
||||
|
||||
class PrefillProgressResponse(BaseRunnerResponse):
|
||||
processed_tokens: int
|
||||
total_tokens: int
|
||||
@@ -0,0 +1,96 @@
|
||||
from collections.abc import Mapping
|
||||
|
||||
from pydantic import model_validator
|
||||
|
||||
from exo_core.model_cards import ModelId
|
||||
from exo_core.types.common import Id, NodeId
|
||||
from exo_core.types.shards import ShardMetadata
|
||||
from exo_core.pydantic import CamelCaseModel, TaggedModel
|
||||
|
||||
|
||||
class RunnerId(Id):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class BaseRunnerStatus(TaggedModel):
|
||||
def is_running(self):
|
||||
return isinstance(self, RunnerRunning)
|
||||
|
||||
|
||||
class RunnerIdle(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerConnecting(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerConnected(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerLoading(BaseRunnerStatus):
|
||||
layers_loaded: int = 0
|
||||
total_layers: int = 0
|
||||
|
||||
|
||||
class RunnerLoaded(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerWarmingUp(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerReady(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerRunning(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerShuttingDown(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerShutdown(BaseRunnerStatus):
|
||||
pass
|
||||
|
||||
|
||||
class RunnerFailed(BaseRunnerStatus):
|
||||
error_message: str | None = None
|
||||
|
||||
|
||||
RunnerStatus = (
|
||||
RunnerIdle
|
||||
| RunnerConnecting
|
||||
| RunnerConnected
|
||||
| RunnerLoading
|
||||
| RunnerLoaded
|
||||
| RunnerWarmingUp
|
||||
| RunnerReady
|
||||
| RunnerRunning
|
||||
| RunnerShuttingDown
|
||||
| RunnerShutdown
|
||||
| RunnerFailed
|
||||
)
|
||||
|
||||
|
||||
class ShardAssignments(CamelCaseModel):
|
||||
model_id: ModelId
|
||||
runner_to_shard: Mapping[RunnerId, ShardMetadata]
|
||||
node_to_runner: Mapping[NodeId, RunnerId]
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_runners_exist(self) -> "ShardAssignments":
|
||||
for runner_id in self.node_to_runner.values():
|
||||
if runner_id not in self.runner_to_shard:
|
||||
raise ValueError(
|
||||
f"Runner {runner_id} in node_to_runner does not exist in runner_to_shard"
|
||||
)
|
||||
return self
|
||||
@@ -0,0 +1,84 @@
|
||||
from enum import Enum
|
||||
from typing import TypeAlias, final
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from exo_core.model_cards import ModelCard
|
||||
from exo_core.pydantic import TaggedModel
|
||||
|
||||
|
||||
class Sharding(str, Enum):
|
||||
Tensor = "Tensor"
|
||||
Pipeline = "Pipeline"
|
||||
|
||||
|
||||
class BaseShardMetadata(TaggedModel):
|
||||
"""
|
||||
Defines a specific shard of the model that is ready to be run on a device.
|
||||
Replaces previous `Shard` object.
|
||||
"""
|
||||
|
||||
model_card: ModelCard
|
||||
device_rank: int
|
||||
world_size: int
|
||||
|
||||
# Error handling; equivalent to monkey-patch, but we can't monkey-patch runner.py
|
||||
# This is kinda annoying because it allocates memory in the ShardMetadata object. Can be rethought after Shanghai.
|
||||
immediate_exception: bool = False
|
||||
should_timeout: float | None = None
|
||||
|
||||
start_layer: int = Field(ge=0)
|
||||
end_layer: int = Field(ge=0)
|
||||
n_layers: int = Field(ge=0)
|
||||
|
||||
@property
|
||||
def is_first_layer(self) -> bool:
|
||||
return self.start_layer == 0
|
||||
|
||||
@property
|
||||
def is_last_layer(self) -> bool:
|
||||
return self.end_layer == self.n_layers
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(
|
||||
(
|
||||
self.model_card.model_id,
|
||||
self.start_layer,
|
||||
self.end_layer,
|
||||
self.n_layers,
|
||||
self.device_rank,
|
||||
self.world_size,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@final
|
||||
class PipelineShardMetadata(BaseShardMetadata):
|
||||
"""
|
||||
Pipeline parallelism shard meta.
|
||||
|
||||
Layers are represented as a half-open interval [start_layer, end_layer),
|
||||
where start_layer is inclusive and end_layer is exclusive.
|
||||
"""
|
||||
|
||||
|
||||
@final
|
||||
class CfgShardMetadata(BaseShardMetadata):
|
||||
"""Shard metadata for CFG-parallel image generation models."""
|
||||
|
||||
cfg_rank: int # 0 = positive branch, 1 = negative branch
|
||||
cfg_world_size: int = 2
|
||||
|
||||
# Pipeline-relative coordinates (computed at placement time)
|
||||
pipeline_rank: int # rank within the pipeline group (0, 1, 2, ...)
|
||||
pipeline_world_size: int # number of nodes per pipeline group
|
||||
|
||||
|
||||
@final
|
||||
class TensorShardMetadata(BaseShardMetadata):
|
||||
pass
|
||||
|
||||
|
||||
ShardMetadata: TypeAlias = (
|
||||
PipelineShardMetadata | CfgShardMetadata | TensorShardMetadata
|
||||
)
|
||||
@@ -0,0 +1,103 @@
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from exo.api.types import (
|
||||
ImageEditsTaskParams,
|
||||
ImageGenerationTaskParams,
|
||||
)
|
||||
from exo_core.types.common import CommandId, Id
|
||||
from exo_core.types.text_generation import TextGenerationTaskParams
|
||||
from exo_core.types.instances import BoundInstance, InstanceId
|
||||
from exo_core.types.runners import RunnerId
|
||||
from exo_core.types.shards import ShardMetadata
|
||||
from exo_core.pydantic import TaggedModel
|
||||
|
||||
|
||||
class TaskId(Id):
|
||||
pass
|
||||
|
||||
|
||||
CANCEL_ALL_TASKS = TaskId("CANCEL_ALL_TASKS")
|
||||
|
||||
|
||||
class TaskStatus(str, Enum):
|
||||
Pending = "Pending"
|
||||
Running = "Running"
|
||||
Complete = "Complete"
|
||||
TimedOut = "TimedOut"
|
||||
Failed = "Failed"
|
||||
Cancelled = "Cancelled"
|
||||
|
||||
|
||||
class BaseTask(TaggedModel):
|
||||
task_id: TaskId = Field(default_factory=TaskId)
|
||||
task_status: TaskStatus = Field(default=TaskStatus.Pending)
|
||||
instance_id: InstanceId
|
||||
|
||||
|
||||
class CreateRunner(BaseTask): # emitted by Worker
|
||||
bound_instance: BoundInstance
|
||||
|
||||
|
||||
class DownloadModel(BaseTask): # emitted by Worker
|
||||
shard_metadata: ShardMetadata
|
||||
|
||||
|
||||
class LoadModel(BaseTask): # emitted by Worker
|
||||
pass
|
||||
|
||||
|
||||
class ConnectToGroup(BaseTask): # emitted by Worker
|
||||
pass
|
||||
|
||||
|
||||
class StartWarmup(BaseTask): # emitted by Worker
|
||||
pass
|
||||
|
||||
|
||||
class TextGeneration(BaseTask): # emitted by Master
|
||||
command_id: CommandId
|
||||
task_params: TextGenerationTaskParams
|
||||
|
||||
error_type: str | None = Field(default=None)
|
||||
error_message: str | None = Field(default=None)
|
||||
|
||||
|
||||
class CancelTask(BaseTask):
|
||||
cancelled_task_id: TaskId
|
||||
runner_id: RunnerId
|
||||
|
||||
|
||||
class ImageGeneration(BaseTask): # emitted by Master
|
||||
command_id: CommandId
|
||||
task_params: ImageGenerationTaskParams
|
||||
|
||||
error_type: str | None = Field(default=None)
|
||||
error_message: str | None = Field(default=None)
|
||||
|
||||
|
||||
class ImageEdits(BaseTask): # emitted by Master
|
||||
command_id: CommandId
|
||||
task_params: ImageEditsTaskParams
|
||||
|
||||
error_type: str | None = Field(default=None)
|
||||
error_message: str | None = Field(default=None)
|
||||
|
||||
|
||||
class Shutdown(BaseTask): # emitted by Worker
|
||||
runner_id: RunnerId
|
||||
|
||||
|
||||
Task = (
|
||||
CreateRunner
|
||||
| DownloadModel
|
||||
| ConnectToGroup
|
||||
| LoadModel
|
||||
| StartWarmup
|
||||
| TextGeneration
|
||||
| CancelTask
|
||||
| ImageGeneration
|
||||
| ImageEdits
|
||||
| Shutdown
|
||||
)
|
||||
@@ -0,0 +1,72 @@
|
||||
"""Canonical internal type for text generation task parameters.
|
||||
|
||||
All external API formats (Chat Completions, Claude Messages, OpenAI Responses)
|
||||
are converted to TextGenerationTaskParams at the API boundary via adapters.
|
||||
"""
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from exo_core.types.common import ModelId
|
||||
|
||||
MessageRole = Literal["user", "assistant", "system", "developer"]
|
||||
ReasoningEffort = Literal["none", "minimal", "low", "medium", "high", "xhigh"]
|
||||
|
||||
|
||||
def resolve_reasoning_params(
|
||||
reasoning_effort: ReasoningEffort | None,
|
||||
enable_thinking: bool | None,
|
||||
) -> tuple[ReasoningEffort | None, bool | None]:
|
||||
"""
|
||||
enable_thinking=True -> reasoning_effort="medium"
|
||||
enable_thinking=False -> reasoning_effort="none"
|
||||
reasoning_effort="none" -> enable_thinking=False
|
||||
reasoning_effort=<anything else> -> enable_thinking=True
|
||||
"""
|
||||
resolved_effort: ReasoningEffort | None = reasoning_effort
|
||||
resolved_thinking: bool | None = enable_thinking
|
||||
|
||||
if reasoning_effort is None and enable_thinking is not None:
|
||||
resolved_effort = "medium" if enable_thinking else "none"
|
||||
|
||||
if enable_thinking is None and reasoning_effort is not None:
|
||||
resolved_thinking = reasoning_effort != "none"
|
||||
|
||||
return resolved_effort, resolved_thinking
|
||||
|
||||
|
||||
class InputMessage(BaseModel, frozen=True):
|
||||
"""Internal message for text generation pipelines."""
|
||||
|
||||
role: MessageRole
|
||||
content: str
|
||||
|
||||
|
||||
class TextGenerationTaskParams(BaseModel, frozen=True):
|
||||
"""Canonical internal task params for text generation.
|
||||
|
||||
Every API adapter converts its wire type into this before handing
|
||||
off to the master/worker pipeline.
|
||||
"""
|
||||
|
||||
model: ModelId
|
||||
input: list[InputMessage]
|
||||
instructions: str | None = None
|
||||
max_output_tokens: int | None = None
|
||||
temperature: float | None = None
|
||||
top_p: float | None = None
|
||||
stream: bool = False
|
||||
tools: list[dict[str, Any]] | None = None
|
||||
bench: bool = False
|
||||
top_k: int | None = None
|
||||
stop: str | list[str] | None = None
|
||||
seed: int | None = None
|
||||
chat_template_messages: list[dict[str, Any]] | None = None
|
||||
reasoning_effort: ReasoningEffort | None = None
|
||||
enable_thinking: bool | None = None
|
||||
logprobs: bool = False
|
||||
top_logprobs: int | None = None
|
||||
min_p: float | None = None
|
||||
repetition_penalty: float | None = None
|
||||
repetition_context_size: int | None = None
|
||||
@@ -0,0 +1,151 @@
|
||||
from math import ceil
|
||||
from typing import Self, overload
|
||||
|
||||
from exo_core.pydantic import FrozenModel
|
||||
|
||||
|
||||
class Memory(FrozenModel):
|
||||
in_bytes: int = 0
|
||||
|
||||
@classmethod
|
||||
def from_bytes(cls, val: int) -> Self:
|
||||
"""Construct a new Memory object from a number of bytes"""
|
||||
return cls(in_bytes=val)
|
||||
|
||||
@property
|
||||
def in_kb(self) -> int:
|
||||
"""The approximate kilobytes this memory represents, rounded up. Setting this property rounds to the nearest byte."""
|
||||
return ceil(self.in_bytes / 1024)
|
||||
|
||||
@in_kb.setter
|
||||
def in_kb(self, val: int):
|
||||
"""Set this memorys value in kilobytes."""
|
||||
self.in_bytes = val * 1024
|
||||
|
||||
@classmethod
|
||||
def from_kb(cls, val: int) -> Self:
|
||||
"""Construct a new Memory object from a number of kilobytes"""
|
||||
return cls(in_bytes=val * 1024)
|
||||
|
||||
@classmethod
|
||||
def from_float_kb(cls, val: float) -> Self:
|
||||
"""Construct a new Memory object from a number of kilobytes, rounding where appropriate"""
|
||||
return cls(in_bytes=round(val * 1024))
|
||||
|
||||
@property
|
||||
def in_mb(self) -> int:
|
||||
"""The approximate megabytes this memory represents, rounded to nearest MB. Setting this property rounds to the nearest byte."""
|
||||
return round(self.in_bytes / (1024**2))
|
||||
|
||||
@in_mb.setter
|
||||
def in_mb(self, val: int):
|
||||
"""Set the megabytes for this memory."""
|
||||
self.in_bytes = val * (1024**2)
|
||||
|
||||
@property
|
||||
def in_float_mb(self) -> float:
|
||||
"""The megabytes this memory represents as a float. Setting this property rounds to the nearest byte."""
|
||||
return self.in_bytes / (1024**2)
|
||||
|
||||
@in_float_mb.setter
|
||||
def in_float_mb(self, val: float):
|
||||
"""Set the megabytes for this memory, rounded to the nearest byte."""
|
||||
self.in_bytes = round(val * (1024**2))
|
||||
|
||||
@classmethod
|
||||
def from_mb(cls, val: float) -> Self:
|
||||
"""Construct a new Memory object from a number of megabytes"""
|
||||
return cls(in_bytes=round(val * (1024**2)))
|
||||
|
||||
@classmethod
|
||||
def from_gb(cls, val: float) -> Self:
|
||||
"""Construct a new Memory object from a number of megabytes"""
|
||||
return cls(in_bytes=round(val * (1024**3)))
|
||||
|
||||
@property
|
||||
def in_gb(self) -> float:
|
||||
"""The approximate gigabytes this memory represents."""
|
||||
return self.in_bytes / (1024**3)
|
||||
|
||||
def __add__(self, other: object) -> "Memory":
|
||||
if isinstance(other, Memory):
|
||||
return Memory.from_bytes(self.in_bytes + other.in_bytes)
|
||||
return NotImplemented
|
||||
|
||||
def __radd__(self, other: object) -> "Memory":
|
||||
if other == 0:
|
||||
return self
|
||||
return NotImplemented
|
||||
|
||||
def __sub__(self, other: object) -> "Memory":
|
||||
if isinstance(other, Memory):
|
||||
return Memory.from_bytes(self.in_bytes - other.in_bytes)
|
||||
return NotImplemented
|
||||
|
||||
def __mul__(self, other: int | float):
|
||||
return Memory.from_bytes(round(self.in_bytes * other))
|
||||
|
||||
def __rmul__(self, other: int | float):
|
||||
return self * other
|
||||
|
||||
@overload
|
||||
def __truediv__(self, other: "Memory") -> float: ...
|
||||
@overload
|
||||
def __truediv__(self, other: int) -> "Memory": ...
|
||||
@overload
|
||||
def __truediv__(self, other: float) -> "Memory": ...
|
||||
def __truediv__(self, other: object) -> "Memory | float":
|
||||
if isinstance(other, Memory):
|
||||
return self.in_bytes / other.in_bytes
|
||||
if isinstance(other, (int, float)):
|
||||
return Memory.from_bytes(round(self.in_bytes / other))
|
||||
return NotImplemented
|
||||
|
||||
def __floordiv__(self, other: object) -> "Memory":
|
||||
if isinstance(other, (int, float)):
|
||||
return Memory.from_bytes(int(self.in_bytes // other))
|
||||
return NotImplemented
|
||||
|
||||
def __lt__(self, other: object) -> bool:
|
||||
if isinstance(other, Memory):
|
||||
return self.in_bytes < other.in_bytes
|
||||
return NotImplemented
|
||||
|
||||
def __le__(self, other: object) -> bool:
|
||||
if isinstance(other, Memory):
|
||||
return self.in_bytes <= other.in_bytes
|
||||
return NotImplemented
|
||||
|
||||
def __gt__(self, other: object) -> bool:
|
||||
if isinstance(other, Memory):
|
||||
return self.in_bytes > other.in_bytes
|
||||
return NotImplemented
|
||||
|
||||
def __ge__(self, other: object) -> bool:
|
||||
if isinstance(other, Memory):
|
||||
return self.in_bytes >= other.in_bytes
|
||||
return NotImplemented
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if isinstance(other, Memory):
|
||||
return self.in_bytes == other.in_bytes
|
||||
return NotImplemented
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Memory.from_bytes({self.in_bytes})"
|
||||
|
||||
def __str__(self) -> str:
|
||||
if self.in_gb > 2:
|
||||
val = self.in_gb
|
||||
unit = "GiB"
|
||||
elif self.in_mb > 2:
|
||||
val = self.in_mb
|
||||
unit = "MiB"
|
||||
elif self.in_kb > 3:
|
||||
val = self.in_kb
|
||||
unit = "KiB"
|
||||
else:
|
||||
val = self.in_bytes
|
||||
unit = "B"
|
||||
|
||||
return f"{val:.2f} {unit}".rstrip("0").rstrip(".") + f" {unit}"
|
||||
@@ -1,5 +1,5 @@
|
||||
[project]
|
||||
name = "vllm-runner"
|
||||
name = "vllm-engine"
|
||||
version = "0.1.0"
|
||||
description = "Add your description here"
|
||||
readme = "README.md"
|
||||
|
||||
Reference in New Issue
Block a user