diff --git a/master/api.py b/master/api.py new file mode 100644 index 00000000..b5455c1d --- /dev/null +++ b/master/api.py @@ -0,0 +1,39 @@ +from typing import NewType +from shared.types.model import ModelId, DownloadProgress, ModelMeta, ModelSource +from shared.types.instance import InstanceId, Instance +from shared.types.model import Topology + + +RequestId = NewType("RequestId", str) + + +class ControlPlaneAPI: + def get_cluster_state() -> ClusterState: + ... + + def get_topology() -> Topology: + ... + + def get_instances() -> list[Instance]: + ... + + def get_instance(instance_id: InstanceId) -> Instance: + ... + + def create_instance(model_id: ModelId) -> InstanceId: + ... + + def delete_instance(instance_id: InstanceId) -> None: + ... + + def get_model_meta(model_id: ModelId) -> ModelMeta: + ... + + def download_model(model_source: ModelSource) -> InstanceId: + ... + + def get_model_download_progress(instance_id: InstanceId) -> DownloadProgress: + ... + + def create_chat_completion_request() -> RequestId: + ... diff --git a/shared/types/instance.py b/shared/types/instance.py new file mode 100644 index 00000000..67f48285 --- /dev/null +++ b/shared/types/instance.py @@ -0,0 +1,90 @@ +from typing import Annotated, Literal, Generic, TypeVar, Union +from enum import Enum +from pydantic import BaseModel, UUID4, PositiveInt, Field + +from shared.types.model import ModelId + +InstanceId = Annotated[str, UUID4] +NodeId = Annotated[str, UUID4] +RunnerId = Annotated[str, UUID4] + + +class ShardType(str, Enum): + PipelineParallel = "PipelineParallel" + + +ShardTypeT = TypeVar("ShardTypeT", bound=ShardType) + + +class ShardData(BaseModel, Generic[ShardTypeT]): + shard_type: ShardTypeT + + +class Shard(BaseModel, Generic[ShardTypeT]): + shard_data: ShardData[ShardTypeT] + runner_id: RunnerId + + +class ShardPlacement(BaseModel): + model_id: ModelId + shard_assignments: dict[NodeId, Shard[ShardType]] + + +class DownloadProgressData(BaseModel): + total_bytes: Annotated[int, PositiveInt] + downloaded_bytes: Annotated[int, PositiveInt] + + +class DownloadStatus(str, Enum): + Pending = "Pending" + Downloading = "Downloading" + Completed = "Completed" + Failed = "Failed" + + +DownloadStatusT = TypeVar("DownloadStatusT", bound=DownloadStatus) + + +class BaseDownloadProgress(BaseModel, Generic[DownloadStatusT]): + node_id: NodeId + download_status: DownloadStatusT + + +class DownloadPending(BaseDownloadProgress[DownloadStatus.Pending]): + download_status: Literal[DownloadStatus.Pending] = Field(DownloadStatus.Pending) + + +class DownloadCompleted(BaseDownloadProgress[DownloadStatus.Completed]): + download_status: Literal[DownloadStatus.Completed] = Field(DownloadStatus.Completed) + + +class DownloadFailed(BaseDownloadProgress[DownloadStatus.Failed]): + download_status: Literal[DownloadStatus.Failed] = Field(DownloadStatus.Failed) + error_message: str + + +class DownloadOngoing(BaseDownloadProgress[DownloadStatus.Downloading]): + download_status: Literal[DownloadStatus.Downloading] = Field( + DownloadStatus.Downloading + ) + download_progress: DownloadProgressData + + +DownloadProgress = Annotated[ + Union[ + DownloadPending, + DownloadCompleted, + DownloadFailed, + DownloadOngoing, + ], + Field(discriminator="download_status"), +] + + +class Instance(ShardPlacement): + instance_id: InstanceId + + +class InstanceDownloadProgress(BaseModel, Generic[DownloadStatusT]): + instance_id: InstanceId + download_progress: BaseDownloadProgress[DownloadStatusT] diff --git a/shared/types/model.py b/shared/types/model.py index d0e11ed6..907366c7 100644 --- a/shared/types/model.py +++ b/shared/types/model.py @@ -1,7 +1,7 @@ -from typing import Annotated, Any, Generic, Literal, TypeVar, final +from typing import Annotated, Any, Generic, Literal, Sequence, TypeVar, Union, final from uuid import UUID -from pydantic import AnyHttpUrl, BaseModel, Field, TypeAdapter +from pydantic import AnyHttpUrl, BaseModel, Field, PositiveInt, TypeAdapter from pydantic.types import UuidVersion SourceType = Literal["HuggingFace", "GitHub"] @@ -13,7 +13,6 @@ ModelId = type("ModelId", (UUID,), {}) ModelIdParser: TypeAdapter[ModelId] = TypeAdapter(_ModelId) RepoPath = Annotated[str, Field(pattern=r"^[^/]+/[^/]+$")] -RepoURL = Annotated[str, AnyHttpUrl] class BaseModelSource(BaseModel, Generic[T]): @@ -44,8 +43,27 @@ class GitHubModelSource(BaseModelSource[Literal["GitHub"]]): source_data: GitHubModelSourceData -RepoType = BaseModelSource[SourceType] +class ModelMetadata(BaseModel): + storage_size_kilobytes: Annotated[int, PositiveInt] + n_layers: Annotated[int, PositiveInt] -RepoValidatorThing = Annotated[RepoType, Field(discriminator="source_type")] -RepoValidator: TypeAdapter[RepoValidatorThing] = TypeAdapter(RepoValidatorThing) +_ModelSource = Annotated[ + Union[ + HuggingFaceModelSource, + GitHubModelSource, + ], + Field(discriminator="source_type"), +] +ModelSource = BaseModelSource[SourceType] + + +@final +class Model(BaseModel): + model_id: ModelId + model_sources: Sequence[ModelSource] + model_metadata: ModelMetadata + + +ModelIdAdapter: TypeAdapter[ModelId] = TypeAdapter(_ModelId) +ModelSourceAdapter: TypeAdapter[ModelSource] = TypeAdapter(_ModelSource)