feat: Add ResourceGraph, runner types, etc.

73ac8969bcf8328b32d9ae3e64c47f1e0f18c4f5 · 2025-07-01 13:14:26 +0100 · Arbion Halili
Files touched

A shared/types/graphs/resource_graph.py
M shared/types/models/sources.py
M shared/types/networking/edges.py
M shared/types/states/master.py
M shared/types/states/shared.py
M shared/types/worker/instances.py
A shared/types/worker/runners.py
M shared/types/worker/shards.py
Diff

commit 73ac8969bcf8328b32d9ae3e64c47f1e0f18c4f5
Author: Arbion Halili <99731180+ToxicPine@users.noreply.github.com>
Date:   Tue Jul 1 13:14:26 2025 +0100

    feat: Add ResourceGraph, runner types, etc.
---
 shared/types/graphs/resource_graph.py | 16 +++++++++
 shared/types/models/sources.py        |  2 +-
 shared/types/networking/edges.py      |  1 -
 shared/types/states/master.py         | 21 ++++++++++++
 shared/types/states/shared.py         |  6 ++--
 shared/types/worker/instances.py      | 24 ++++++++------
 shared/types/worker/runners.py        | 61 +++++++++++++++++++++++++++++++++++
 shared/types/worker/shards.py         | 10 ------
 8 files changed, 118 insertions(+), 23 deletions(-)

diff --git a/shared/types/graphs/resource_graph.py b/shared/types/graphs/resource_graph.py
new file mode 100644
index 00000000..6beca6a4
--- /dev/null
+++ b/shared/types/graphs/resource_graph.py
@@ -0,0 +1,16 @@
+from collections.abc import Mapping
+
+from pydantic import BaseModel
+
+from shared.types.common import NodeId
+from shared.types.networking.topology import Topology
+from shared.types.profiling.common import NodeProfile
+
+
+class ResourceGraph(BaseModel): ...
+
+
+def get_graph_of_compute_resources(
+    network_topology: Topology,
+    node_profiles: Mapping[NodeId, NodeProfile],
+) -> ResourceGraph: ...
diff --git a/shared/types/models/sources.py b/shared/types/models/sources.py
index 927a6ee2..419ed264 100644
--- a/shared/types/models/sources.py
+++ b/shared/types/models/sources.py
@@ -1,5 +1,5 @@
-from typing import Annotated, Any, Generic, Literal, TypeVar, Union, final
 from enum import Enum
+from typing import Annotated, Any, Generic, Literal, TypeVar, Union, final
 
 from pydantic import AnyHttpUrl, BaseModel, Field, TypeAdapter
 
diff --git a/shared/types/networking/edges.py b/shared/types/networking/edges.py
index 3866fc2e..bcf9f2ac 100644
--- a/shared/types/networking/edges.py
+++ b/shared/types/networking/edges.py
@@ -49,7 +49,6 @@ class EdgeMetadata(BaseModel, Generic[TE, TF]): ...
 
 
 @final
-@dataclass
 class MLXEdgeContext(EdgeMetadata[AddressingProtocol.IPvAny, ApplicationProtocol.MLX]):
     source_ip: IPvAnyAddress
     sink_ip: IPvAnyAddress
diff --git a/shared/types/states/master.py b/shared/types/states/master.py
index aef63941..752303df 100644
--- a/shared/types/states/master.py
+++ b/shared/types/states/master.py
@@ -1,11 +1,15 @@
+from collections.abc import Mapping, Sequence
 from queue import Queue
 
 from pydantic import BaseModel
 
 from shared.types.common import NodeId
+from shared.types.events.common import Event, EventTypes
+from shared.types.graphs.resource_graph import ResourceGraph
 from shared.types.networking.topology import NetworkState
 from shared.types.profiling.common import NodeProfile
 from shared.types.states.shared import SharedState
+from shared.types.worker.instances import InstanceData, InstanceId
 
 
 class ExternalCommand(BaseModel): ...
@@ -16,3 +20,20 @@ class MasterState(SharedState):
     node_profiles: dict[NodeId, NodeProfile]
     job_inbox: Queue[ExternalCommand]
     job_outbox: Queue[ExternalCommand]
+
+
+def get_inference_plan(
+    inbox: Queue[ExternalCommand],
+    outbox: Queue[ExternalCommand],
+    resource_graph: ResourceGraph,
+    current_instances: Mapping[InstanceId, InstanceData],
+) -> Mapping[InstanceId, InstanceData]: ...
+
+
+TransitionEventTypes = EventTypes
+
+
+def get_transition_events(
+    current_instances: Mapping[InstanceId, InstanceData],
+    target_instances: Mapping[InstanceId, InstanceData],
+) -> Sequence[Event[TransitionEventTypes]]: ...
diff --git a/shared/types/states/shared.py b/shared/types/states/shared.py
index f5c55c09..5f4fc3b5 100644
--- a/shared/types/states/shared.py
+++ b/shared/types/states/shared.py
@@ -1,10 +1,12 @@
+from collections.abc import Mapping
+
 from pydantic import BaseModel
 
 from shared.types.common import NodeId
 from shared.types.worker.common import InstanceId
-from shared.types.worker.shards import ShardPlacement
+from shared.types.worker.instances import InstanceData
 
 
 class SharedState(BaseModel):
     node_id: NodeId
-    compute_instances: dict[InstanceId, ShardPlacement]
+    compute_instances: Mapping[InstanceId, InstanceData]
diff --git a/shared/types/worker/instances.py b/shared/types/worker/instances.py
index 447552d7..d4ed748b 100644
--- a/shared/types/worker/instances.py
+++ b/shared/types/worker/instances.py
@@ -1,18 +1,24 @@
-from typing import Generic, TypeVar
+from collections.abc import Mapping
 
 from pydantic import BaseModel
 
 from shared.types.worker.common import InstanceId
-from shared.types.worker.downloads import BaseDownloadProgress, DownloadStatus
-from shared.types.worker.shards import ShardPlacement
+from shared.types.worker.runners import (
+    RunnerId,
+    RunnerPlacement,
+    RunnerState,
+    RunnerStateType,
+)
 
-DownloadStatusT = TypeVar("DownloadStatusT", bound=DownloadStatus)
 
-
-class Instance(ShardPlacement):
+class InstanceBase(BaseModel):
     instance_id: InstanceId
 
 
-class InstanceDownloadProgress(BaseModel, Generic[DownloadStatusT]):
-    instance_id: InstanceId
-    download_progress: BaseDownloadProgress[DownloadStatusT]
+class InstanceData(BaseModel):
+    runner_placements: RunnerPlacement
+    runner_states: Mapping[RunnerId, RunnerState[RunnerStateType]]
+
+
+class Instance(InstanceBase):
+    instance_data: InstanceData
diff --git a/shared/types/worker/runners.py b/shared/types/worker/runners.py
new file mode 100644
index 00000000..144a141f
--- /dev/null
+++ b/shared/types/worker/runners.py
@@ -0,0 +1,61 @@
+from collections.abc import Mapping, Sequence
+from enum import Enum
+from typing import Generic, Literal, TypeVar
+
+from pydantic import BaseModel
+
+from shared.types.common import NodeId
+from shared.types.models.common import ModelId
+from shared.types.worker.common import RunnerId
+from shared.types.worker.downloads import BaseDownloadProgress, DownloadStatus
+from shared.types.worker.shards import Shard, ShardType
+
+
+class RunnerStateType(str, Enum):
+    Rejected = "Rejected"
+    Starting = "Starting"
+    Downloading = "Downloading"
+    Running = "Running"
+    Failed = "Failed"
+
+
+RunnerStateTypeT = TypeVar("RunnerStateTypeT", bound=RunnerStateType)
+
+
+class RunnerState(BaseModel, Generic[RunnerStateTypeT]):
+    runner_state: RunnerStateTypeT
+
+
+class RejectedRunnerState(RunnerState[RunnerStateType.Rejected]):
+    runner_state: Literal[RunnerStateType.Rejected]
+
+
+class StartingRunnerState(RunnerState[RunnerStateType.Starting]):
+    runner_state: Literal[RunnerStateType.Starting]
+
+
+class DownloadingRunnerState(RunnerState[RunnerStateType.Downloading]):
+    runner_state: Literal[RunnerStateType.Downloading]
+    download_progress: BaseDownloadProgress[DownloadStatus]
+
+
+class RunningRunnerState(RunnerState[RunnerStateType.Running]):
+    runner_state: Literal[RunnerStateType.Running]
+
+
+class FailedRunnerState(RunnerState[RunnerStateType.Failed]):
+    runner_state: Literal[RunnerStateType.Failed]
+    error_message: str | None = None
+
+
+class RunnerData(BaseModel):
+    runner_id: RunnerId
+    runner_state: RunnerState[RunnerStateType] = RunnerState(
+        runner_state=RunnerStateType.Starting
+    )
+
+
+class RunnerPlacement(BaseModel):
+    model_id: ModelId
+    runner_to_shard: Mapping[RunnerId, Shard[ShardType]]
+    node_to_runner: Mapping[NodeId, Sequence[RunnerId]]
diff --git a/shared/types/worker/shards.py b/shared/types/worker/shards.py
index aa6df9ad..3e9055ae 100644
--- a/shared/types/worker/shards.py
+++ b/shared/types/worker/shards.py
@@ -3,10 +3,6 @@ from typing import Generic, TypeVar
 
 from pydantic import BaseModel
 
-from shared.types.common import NodeId
-from shared.types.models.common import ModelId
-from shared.types.worker.common import RunnerId
-
 
 class ShardType(str, Enum):
     PipelineParallel = "PipelineParallel"
@@ -21,9 +17,3 @@ class ShardData(BaseModel, Generic[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]]
← df824e2e fix: Ensure MasterState inherits from SharedState · back to Exo · feat: Update Interfaces 6de1f288 →