Stash initial attempt at implementing process prompt before being sad

src/discovery/mod.rs

@@ -69,6 +69,7 @@ impl Default for NodeInfo {
     }
 }
 
+#[derive(Debug)]
 pub struct UdpDiscovery {
     discovery_handle: JoinHandle<()>,
     presence_handle: JoinHandle<()>,

src/main.rs

@@ -11,24 +11,53 @@ use tonic::{transport::Server, Request, Response, Status};
 
 use crate::discovery::{NodeInfo, UdpDiscovery};
 use crate::node_service::{
-    CollectTopologyRequest, Empty, ExampleRequest, HealthCheckRequest, HealthCheckResponse, Loss,
-    PromptRequest, SendOpaqueStatusRequest, SendResultRequest, Tensor, Topology as TopologyProto,
+    CollectTopologyRequest, Empty, ExampleRequest, HealthCheckRequest, HealthCheckResponse,
+    InferenceState, Loss, PromptRequest, SendOpaqueStatusRequest, SendResultRequest, Tensor,
+    Topology as TopologyProto,
 };
 use node_service::node_service_server::{NodeService, NodeServiceServer};
 use node_service::TensorRequest;
 use std::collections::HashSet;
 use topology::Topology;
+use uuid::Uuid;
 
 pub mod node_service {
     tonic::include_proto!("node_service"); // The string specified here must match the proto package name
 }
 
+#[derive(Debug)]
 struct Node {
     node_info: NodeInfo,
     current_topology: Topology,
     udp_discovery: UdpDiscovery,
 }
 
+impl Node {
+    #[tracing::instrument]
+    pub async fn process_prompt(
+        &self,
+        base_shard: Shard,
+        prompt: String,
+        request_id: String,
+        inference_state: Option<InferenceState>,
+    ) {
+        let shard = self
+            .current_topology
+            .get_shard_for_node(base_shard, &self.node_info.node_id);
+
+        todo!();
+        // if shard.is_first_layer() {
+        //     let result = self
+        //         .inference_engine
+        //         .infer_prompt(request_id, shard, prompt, inference_state)
+        //         .await;
+        //     self.process_inference_result(shard, result, request_id, inference_state)
+        // } else {
+        //     self.forward_prompt(shard, prompt, request_id, inference_state)
+        // }
+    }
+}
+
 impl Default for Node {
     fn default() -> Self {
         let node_info = NodeInfo::default();
@@ -52,9 +81,35 @@ enum OpaqueStatus {
 #[derive(Debug, Deserialize, Serialize, Clone)]
 struct Shard {
     pub model_id: String,
-    pub start_layer: i32,
-    pub end_layer: i32,
-    pub n_layers: i32,
+    pub start_layer: u32,
+    pub end_layer: u32,
+    #[serde(rename = "n_layers")]
+    pub total_layers: u32,
+}
+
+impl Shard {
+    pub fn is_first_layer(&self) -> bool {
+        self.start_layer == 0
+    }
+
+    pub fn is_last_layer(&self) -> bool {
+        self.end_layer == self.total_layers - 1
+    }
+
+    pub fn len(&self) -> u32 {
+        self.end_layer - self.start_layer + 1
+    }
+}
+
+impl From<node_service::Shard> for Shard {
+    fn from(proto: node_service::Shard) -> Self {
+        Self {
+            model_id: proto.model_id,
+            start_layer: proto.start_layer as u32,
+            end_layer: proto.end_layer as u32,
+            total_layers: proto.n_layers as u32,
+        }
+    }
 }
 
 #[derive(Debug, Deserialize, Serialize, Clone)]
@@ -141,7 +196,22 @@ impl NodeService for Node {
         &self,
         request: Request<PromptRequest>,
     ) -> Result<Response<Tensor>, Status> {
-        todo!()
+        let request = request.into_inner();
+        let request_id = request
+            .request_id
+            .unwrap_or_else(|| Uuid::new_v4().to_string());
+
+        let result = self.process_prompt(
+            request
+                .shard
+                .expect("No shard given. ExoPy does not allow this")
+                .into(),
+            request.prompt,
+            request_id,
+            request.inference_state,
+        );
+
+        todo!();
     }
 
     async fn send_tensor(

src/partitioning.rs

@@ -1,3 +1,4 @@
+use crate::Shard;
 use crate::topology::Topology;
 
 pub enum PartitionStrategy {
@@ -5,9 +6,9 @@ pub enum PartitionStrategy {
 }
 
 pub struct Partition {
-    node_id: String,
-    start: f32,
-    end: f32,
+    pub node_id: String,
+    pub start: f32,
+    pub end: f32,
 }
 
 impl PartitionStrategy {
@@ -34,30 +35,23 @@ impl PartitionStrategy {
     }
 }
 
-pub struct ModelShard<'a> {
-    model_id: &'a str,
-    start_layer: u8,
-    end_layer: u8,
-    total_layers: u8,
-}
-
 pub fn shard_model_by_partition(
-    partitions: Vec<Partition>,
-    total_layers: u8,
+    partition_set: &[Partition],
+    total_layers: u32,
     model_id: &str,
-) -> Vec<ModelShard> {
-    let mut shards: Vec<ModelShard<'_>> = Vec::with_capacity(partitions.len());
+) -> Vec<Shard> {
+    let mut shards: Vec<Shard> = Vec::with_capacity(partition_set.len());
 
-    for partition in partitions {
-        let start_layer = (partition.start * total_layers as f32).round() as u8;
-        let mut end_layer = (partition.end * total_layers as f32).round() as u8 - 1;
+    for partition in partition_set {
+        let start_layer = (partition.start * total_layers as f32).round() as u32;
+        let mut end_layer = (partition.end * total_layers as f32).round() as u32 - 1;
 
         if end_layer < start_layer {
             end_layer = total_layers - 1;
         }
 
-        shards.push(ModelShard {
-            model_id,
+        shards.push(Shard {
+            model_id: model_id.to_string(),
             start_layer,
             end_layer,
             total_layers,

src/topology.rs

@@ -1,4 +1,5 @@
-use crate::{device_capability_data, node_service};
+use crate::partitioning::{shard_model_by_partition, PartitionStrategy};
+use crate::{device_capability_data, node_service, Shard};
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::process::Command;
@@ -10,6 +11,26 @@ pub struct Topology {
     pub active_node_id: Option<String>,
 }
 
+impl Topology {
+    pub fn get_shard_for_node(&self, base_shard: Shard, node_id: &str) -> Shard {
+        let partition_set = PartitionStrategy::RingMemoryWeighted.partition(&self);
+
+        // TODO: This feels like it could be a better data structure
+        let partition_index = partition_set
+            .iter()
+            .position(|s| s.node_id == node_id)
+            .expect("Did not find node in partition set");
+
+        let shards = shard_model_by_partition(
+            &partition_set,
+            base_shard.total_layers.try_into().unwrap(),
+            base_shard.model_id.as_str(),
+        );
+
+        shards[partition_index].clone()
+    }
+}
+
 impl Topology {
     pub fn update_node(&mut self, node_id: String, device_capabilities: DeviceCapabilities) {
         self.nodes.insert(node_id, device_capabilities);
@@ -48,8 +69,8 @@ impl Topology {
     }
 }
 
-impl From<crate::node_service::Topology> for Topology {
-    fn from(proto: crate::node_service::Topology) -> Self {
+impl From<node_service::Topology> for Topology {
+    fn from(proto: node_service::Topology) -> Self {
         let nodes = proto
             .nodes
             .into_iter()
@@ -83,8 +104,8 @@ impl From<crate::node_service::Topology> for Topology {
     }
 }
 
-impl Into<crate::node_service::Topology> for Topology {
-    fn into(self) -> crate::node_service::Topology {
+impl Into<node_service::Topology> for Topology {
+    fn into(self) -> node_service::Topology {
         let nodes = self
             .nodes
             .iter()
@@ -210,8 +231,8 @@ impl DeviceCapabilities {
     }
 }
 
-impl From<crate::node_service::DeviceCapabilities> for DeviceCapabilities {
-    fn from(value: crate::node_service::DeviceCapabilities) -> Self {
+impl From<node_service::DeviceCapabilities> for DeviceCapabilities {
+    fn from(value: node_service::DeviceCapabilities) -> Self {
         DeviceCapabilities {
             model: value.model,
             chip: value.chip,
@@ -228,8 +249,8 @@ pub struct DeviceFlops {
     pub int8: f64,
 }
 
-impl From<crate::node_service::DeviceFlops> for DeviceFlops {
-    fn from(value: crate::node_service::DeviceFlops) -> Self {
+impl From<node_service::DeviceFlops> for DeviceFlops {
+    fn from(value: node_service::DeviceFlops) -> Self {
         DeviceFlops {
             fp32: value.fp32,
             fp16: value.fp16,