Merge Seth's Control Plane API Work into Alex's Events Branch

Co-authored-by: Seth Howes <[email protected]>

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:
+        ...

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]

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)