707 lines
20 KiB
Rust
707 lines
20 KiB
Rust
pub const MO01_HEADER_BYTES: u64 = 40;
|
|
const TOKEN_ID_WIDTH_BYTES: u32 = 4;
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct RunId(pub u64);
|
|
|
|
impl From<u64> for RunId {
|
|
fn from(value: u64) -> Self {
|
|
Self(value)
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct NodeId(pub u64);
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct EdgeId(pub u64);
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct EdgeAllocator {
|
|
next: u64,
|
|
}
|
|
|
|
impl EdgeAllocator {
|
|
pub fn new() -> Self {
|
|
Self { next: 1 }
|
|
}
|
|
|
|
pub fn alloc(&mut self) -> EdgeId {
|
|
let edge_id = EdgeId(self.next);
|
|
self.next += 1;
|
|
edge_id
|
|
}
|
|
}
|
|
|
|
impl Default for EdgeAllocator {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum DTypeFamily {
|
|
BFloat,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct ModelFacts {
|
|
pub model_id: String,
|
|
pub gguf_source: GgufSource,
|
|
pub num_layers: u32,
|
|
pub hidden_dim: u64,
|
|
pub dtype_family: DTypeFamily,
|
|
pub dtype_width_bytes: u64,
|
|
pub max_seq_len: u64,
|
|
pub eos_token_id: u32,
|
|
pub tokenizer: TokenizerSource,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
|
pub enum GgufSource {
|
|
LocalPath(String),
|
|
HuggingFaceGguf {
|
|
repo: String,
|
|
file: String,
|
|
revision: Option<String>,
|
|
},
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
|
pub enum TokenizerSource {
|
|
EmbeddedGguf,
|
|
LocalPath(String),
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub enum PromptSource {
|
|
Inline(String),
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub struct SamplingPolicy {
|
|
pub temperature_millis: u32,
|
|
pub top_k: u32,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum TokenOutputPolicy {
|
|
EmitAll,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct RoleId(pub u64);
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub struct PortId(pub String);
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct GgufModelPlan {
|
|
pub model_id: String,
|
|
pub gguf_source: GgufSource,
|
|
pub num_layers: u32,
|
|
pub hidden_dim: u32,
|
|
pub dtype_family: DTypeFamily,
|
|
pub dtype_width_bytes: u32,
|
|
pub max_seq_len: u32,
|
|
pub eos_token_id: u32,
|
|
pub tokenizer: TokenizerSource,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct RuntimePlan {
|
|
pub prompt: PromptSource,
|
|
pub sampling: SamplingPolicy,
|
|
pub token_output_policy: TokenOutputPolicy,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct RuntimeConfig {
|
|
pub max_tokens: u32,
|
|
pub prompt: PromptSource,
|
|
pub sampling: SamplingPolicy,
|
|
pub token_output_policy: TokenOutputPolicy,
|
|
}
|
|
|
|
impl RuntimeConfig {
|
|
pub fn test_default() -> Self {
|
|
Self {
|
|
max_tokens: 4,
|
|
prompt: PromptSource::Inline("test prompt".to_owned()),
|
|
sampling: SamplingPolicy {
|
|
temperature_millis: 0,
|
|
top_k: 1,
|
|
},
|
|
token_output_policy: TokenOutputPolicy::EmitAll,
|
|
}
|
|
}
|
|
|
|
fn plan(&self) -> RuntimePlan {
|
|
RuntimePlan {
|
|
prompt: self.prompt.clone(),
|
|
sampling: self.sampling,
|
|
token_output_policy: self.token_output_policy,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub struct StagePlacement {
|
|
pub stage_index: u32,
|
|
pub node_id: NodeId,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub enum PlacementInput {
|
|
FixedLinear(Vec<StagePlacement>),
|
|
FixedLinearWithEdgeOverride {
|
|
stages: Vec<StagePlacement>,
|
|
forced_activation_edges: Vec<(u32, u32)>,
|
|
},
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum RingDirection {
|
|
Ingress,
|
|
Egress,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum HostPinning {
|
|
Pageable,
|
|
PinnedRequired,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum WakeCoalescing {
|
|
PendingBit,
|
|
ReadySet,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub struct RingSpec {
|
|
pub data_capacity: u64,
|
|
pub alignment: u32,
|
|
pub direction: RingDirection,
|
|
pub host_pinning: HostPinning,
|
|
pub wake_coalescing: WakeCoalescing,
|
|
}
|
|
|
|
impl RingSpec {
|
|
pub fn test_default_activation() -> Self {
|
|
Self {
|
|
data_capacity: 1 << 20,
|
|
alignment: 64,
|
|
direction: RingDirection::Egress,
|
|
host_pinning: HostPinning::Pageable,
|
|
wake_coalescing: WakeCoalescing::PendingBit,
|
|
}
|
|
}
|
|
|
|
pub fn test_default_token() -> Self {
|
|
Self {
|
|
data_capacity: 4096,
|
|
alignment: 8,
|
|
direction: RingDirection::Egress,
|
|
host_pinning: HostPinning::Pageable,
|
|
wake_coalescing: WakeCoalescing::PendingBit,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct PlannerInput {
|
|
pub run_id: RunId,
|
|
pub orchestrator_node_id: NodeId,
|
|
pub model: ModelFacts,
|
|
pub runtime: RuntimeConfig,
|
|
pub candidate_pool: Vec<NodeId>,
|
|
pub stage_count: u32,
|
|
pub placement: PlacementInput,
|
|
pub activation_ring: RingSpec,
|
|
pub token_ring: RingSpec,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub enum EdgeKind {
|
|
TokenIn,
|
|
Activation,
|
|
TokenOut,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum ObjectKind {
|
|
Token,
|
|
Activation,
|
|
Weight,
|
|
ModelShard,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum ShapeRule {
|
|
TokenIds,
|
|
ActivationRows { max_seq_len: u32, hidden_dim: u32 },
|
|
WeightTensor,
|
|
ModelShardBytes,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum LayoutRule {
|
|
Contiguous,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum SequencePolicy {
|
|
Ordered,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub struct ObjectSpec {
|
|
pub kind: ObjectKind,
|
|
pub max_extent: u64,
|
|
pub dtype_family: DTypeFamily,
|
|
pub dtype_width_bytes: u32,
|
|
pub shape: ShapeRule,
|
|
pub layout: LayoutRule,
|
|
pub alignment: u32,
|
|
pub sequence_policy: SequencePolicy,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
|
pub enum EdgeEndpoint {
|
|
Orchestrator { node_id: NodeId },
|
|
Stage { node_id: NodeId, stage_index: u32 },
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct EdgePlan {
|
|
pub run_id: RunId,
|
|
pub edge_id: EdgeId,
|
|
pub kind: EdgeKind,
|
|
pub producer: EdgeEndpoint,
|
|
pub consumer: EdgeEndpoint,
|
|
pub object_spec: ObjectSpec,
|
|
pub ring_spec: RingSpec,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct StagePlan {
|
|
pub run_id: RunId,
|
|
pub stage_index: u32,
|
|
pub stage_count: u32,
|
|
pub node_id: NodeId,
|
|
pub gguf_source: GgufSource,
|
|
pub layer_start: u32,
|
|
pub layer_end_exclusive: u32,
|
|
pub inbound_edge: EdgeId,
|
|
pub outbound_edge: EdgeId,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct RunPlan {
|
|
pub run_id: RunId,
|
|
pub model: GgufModelPlan,
|
|
pub runtime: RuntimePlan,
|
|
pub stages: Vec<StagePlan>,
|
|
pub edges: Vec<EdgePlan>,
|
|
pub max_tokens: u32,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct InboundEdgeProvision {
|
|
pub edge_id: EdgeId,
|
|
pub kind: EdgeKind,
|
|
pub object_spec: ObjectSpec,
|
|
pub ring_spec: RingSpec,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct OutboundEdgeProvision {
|
|
pub edge_id: EdgeId,
|
|
pub kind: EdgeKind,
|
|
pub consumer_node_id: NodeId,
|
|
pub object_spec: ObjectSpec,
|
|
pub ring_spec: RingSpec,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct StageModelFacts {
|
|
pub model_id: String,
|
|
pub hidden_dim: u32,
|
|
pub dtype_family: DTypeFamily,
|
|
pub dtype_width_bytes: u32,
|
|
pub max_seq_len: u32,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct StageRuntimeFacts {
|
|
pub role_id: RoleId,
|
|
pub input_port: PortId,
|
|
pub output_port: PortId,
|
|
pub sampling: Option<SamplingPolicy>,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct ProvisionStage {
|
|
pub run_id: RunId,
|
|
pub node_id: NodeId,
|
|
pub stage_index: u32,
|
|
pub stage_count: u32,
|
|
pub gguf_source: GgufSource,
|
|
pub tokenizer: TokenizerSource,
|
|
pub layer_start: u32,
|
|
pub layer_end_exclusive: u32,
|
|
pub inbound: InboundEdgeProvision,
|
|
pub outbound: OutboundEdgeProvision,
|
|
pub model: StageModelFacts,
|
|
pub runtime: StageRuntimeFacts,
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum PlanRejectionKind {
|
|
UnknownNode,
|
|
DuplicateStageAssignment,
|
|
MissingStage,
|
|
InvalidStageCount,
|
|
EdgeEndpointMismatch,
|
|
ModelStageLayoutMismatch,
|
|
InvalidObjectSpec,
|
|
UnsupportedShapeOrLayout,
|
|
InvalidRingSpec,
|
|
}
|
|
|
|
#[derive(Clone, Debug, PartialEq, Eq)]
|
|
pub struct PlanRejection {
|
|
kind: PlanRejectionKind,
|
|
}
|
|
|
|
impl PlanRejection {
|
|
pub fn kind(&self) -> PlanRejectionKind {
|
|
self.kind
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
|
pub enum ProjectionRejection {
|
|
UnknownStage,
|
|
MissingEdge,
|
|
}
|
|
|
|
pub fn plan_run(input: PlannerInput) -> Result<RunPlan, PlanRejection> {
|
|
validate_global_input(&input)?;
|
|
let placements = validated_placements(&input)?;
|
|
let model = model_plan(&input.model)?;
|
|
let runtime = input.runtime.plan();
|
|
let max_tokens = input.runtime.max_tokens;
|
|
let gguf_source = model.gguf_source.clone();
|
|
let hidden_dim = model.hidden_dim;
|
|
let dtype_width_bytes = model.dtype_width_bytes;
|
|
let max_seq_len = model.max_seq_len;
|
|
|
|
let activation_extent = input
|
|
.model
|
|
.max_seq_len
|
|
.checked_mul(input.model.hidden_dim)
|
|
.and_then(|value| value.checked_mul(input.model.dtype_width_bytes))
|
|
.ok_or_else(|| reject(PlanRejectionKind::InvalidObjectSpec))?;
|
|
if activation_extent == 0 {
|
|
return Err(reject(PlanRejectionKind::InvalidObjectSpec));
|
|
}
|
|
|
|
let token_extent = input
|
|
.model
|
|
.max_seq_len
|
|
.checked_mul(u64::from(TOKEN_ID_WIDTH_BYTES))
|
|
.ok_or_else(|| reject(PlanRejectionKind::InvalidObjectSpec))?;
|
|
if token_extent == 0 {
|
|
return Err(reject(PlanRejectionKind::InvalidObjectSpec));
|
|
}
|
|
let token_spec = ObjectSpec {
|
|
kind: ObjectKind::Token,
|
|
max_extent: token_extent,
|
|
dtype_family: input.model.dtype_family,
|
|
dtype_width_bytes: TOKEN_ID_WIDTH_BYTES,
|
|
shape: ShapeRule::TokenIds,
|
|
layout: LayoutRule::Contiguous,
|
|
alignment: TOKEN_ID_WIDTH_BYTES,
|
|
sequence_policy: SequencePolicy::Ordered,
|
|
};
|
|
let activation_spec = ObjectSpec {
|
|
kind: ObjectKind::Activation,
|
|
max_extent: activation_extent,
|
|
dtype_family: input.model.dtype_family,
|
|
dtype_width_bytes,
|
|
shape: ShapeRule::ActivationRows {
|
|
max_seq_len,
|
|
hidden_dim,
|
|
},
|
|
layout: LayoutRule::Contiguous,
|
|
alignment: dtype_width_bytes,
|
|
sequence_policy: SequencePolicy::Ordered,
|
|
};
|
|
let token_data_capacity = MO01_HEADER_BYTES
|
|
.checked_add(token_extent)
|
|
.ok_or_else(|| reject(PlanRejectionKind::InvalidObjectSpec))?;
|
|
let mut token_ring = input.token_ring;
|
|
token_ring.data_capacity = token_ring.data_capacity.max(token_data_capacity);
|
|
|
|
let mut edge_allocator = EdgeAllocator::new();
|
|
let token_in_edge = edge_allocator.alloc();
|
|
let mut activation_edges = Vec::with_capacity(input.stage_count.saturating_sub(1) as usize);
|
|
for _ in 0..input.stage_count.saturating_sub(1) {
|
|
activation_edges.push(edge_allocator.alloc());
|
|
}
|
|
let token_out_edge = edge_allocator.alloc();
|
|
|
|
let mut edges = Vec::with_capacity(input.stage_count as usize + 1);
|
|
edges.push(EdgePlan {
|
|
run_id: input.run_id,
|
|
edge_id: token_in_edge,
|
|
kind: EdgeKind::TokenIn,
|
|
producer: EdgeEndpoint::Orchestrator {
|
|
node_id: input.orchestrator_node_id,
|
|
},
|
|
consumer: EdgeEndpoint::Stage {
|
|
node_id: placements[0].node_id,
|
|
stage_index: 0,
|
|
},
|
|
object_spec: token_spec,
|
|
ring_spec: token_ring,
|
|
});
|
|
|
|
for stage_index in 0..input.stage_count.saturating_sub(1) {
|
|
edges.push(EdgePlan {
|
|
run_id: input.run_id,
|
|
edge_id: activation_edges[stage_index as usize],
|
|
kind: EdgeKind::Activation,
|
|
producer: EdgeEndpoint::Stage {
|
|
node_id: placements[stage_index as usize].node_id,
|
|
stage_index,
|
|
},
|
|
consumer: EdgeEndpoint::Stage {
|
|
node_id: placements[stage_index as usize + 1].node_id,
|
|
stage_index: stage_index + 1,
|
|
},
|
|
object_spec: activation_spec,
|
|
ring_spec: input.activation_ring,
|
|
});
|
|
}
|
|
|
|
edges.push(EdgePlan {
|
|
run_id: input.run_id,
|
|
edge_id: token_out_edge,
|
|
kind: EdgeKind::TokenOut,
|
|
producer: EdgeEndpoint::Stage {
|
|
node_id: placements[input.stage_count as usize - 1].node_id,
|
|
stage_index: input.stage_count - 1,
|
|
},
|
|
consumer: EdgeEndpoint::Orchestrator {
|
|
node_id: input.orchestrator_node_id,
|
|
},
|
|
object_spec: token_spec,
|
|
ring_spec: token_ring,
|
|
});
|
|
|
|
let mut stages = Vec::with_capacity(input.stage_count as usize);
|
|
for placement in &placements {
|
|
let stage_index = placement.stage_index;
|
|
let (start, end) = layer_range(input.model.num_layers, input.stage_count, stage_index);
|
|
let inbound_edge = if stage_index == 0 {
|
|
token_in_edge
|
|
} else {
|
|
activation_edges[stage_index as usize - 1]
|
|
};
|
|
let outbound_edge = if stage_index + 1 == input.stage_count {
|
|
token_out_edge
|
|
} else {
|
|
activation_edges[stage_index as usize]
|
|
};
|
|
stages.push(StagePlan {
|
|
run_id: input.run_id,
|
|
stage_index,
|
|
stage_count: input.stage_count,
|
|
node_id: placement.node_id,
|
|
gguf_source: gguf_source.clone(),
|
|
layer_start: start,
|
|
layer_end_exclusive: end,
|
|
inbound_edge,
|
|
outbound_edge,
|
|
});
|
|
}
|
|
|
|
Ok(RunPlan {
|
|
run_id: input.run_id,
|
|
model,
|
|
runtime,
|
|
stages,
|
|
edges,
|
|
max_tokens,
|
|
})
|
|
}
|
|
|
|
pub fn derive_stage_provision(
|
|
plan: &RunPlan,
|
|
stage_index: u32,
|
|
) -> Result<ProvisionStage, ProjectionRejection> {
|
|
let stage = plan
|
|
.stages
|
|
.iter()
|
|
.find(|stage| stage.stage_index == stage_index)
|
|
.ok_or(ProjectionRejection::UnknownStage)?;
|
|
let inbound = plan
|
|
.edges
|
|
.iter()
|
|
.find(|edge| edge.edge_id == stage.inbound_edge)
|
|
.ok_or(ProjectionRejection::MissingEdge)?;
|
|
let outbound = plan
|
|
.edges
|
|
.iter()
|
|
.find(|edge| edge.edge_id == stage.outbound_edge)
|
|
.ok_or(ProjectionRejection::MissingEdge)?;
|
|
Ok(ProvisionStage {
|
|
run_id: plan.run_id,
|
|
node_id: stage.node_id,
|
|
stage_index,
|
|
stage_count: stage.stage_count,
|
|
gguf_source: stage.gguf_source.clone(),
|
|
tokenizer: plan.model.tokenizer.clone(),
|
|
layer_start: stage.layer_start,
|
|
layer_end_exclusive: stage.layer_end_exclusive,
|
|
inbound: InboundEdgeProvision {
|
|
edge_id: inbound.edge_id,
|
|
kind: inbound.kind,
|
|
object_spec: inbound.object_spec,
|
|
ring_spec: ring_spec_for_direction(inbound.ring_spec, RingDirection::Ingress),
|
|
},
|
|
outbound: OutboundEdgeProvision {
|
|
edge_id: outbound.edge_id,
|
|
kind: outbound.kind,
|
|
consumer_node_id: endpoint_node_id(&outbound.consumer),
|
|
object_spec: outbound.object_spec,
|
|
ring_spec: ring_spec_for_direction(outbound.ring_spec, RingDirection::Egress),
|
|
},
|
|
model: StageModelFacts {
|
|
model_id: plan.model.model_id.clone(),
|
|
hidden_dim: plan.model.hidden_dim,
|
|
dtype_family: plan.model.dtype_family,
|
|
dtype_width_bytes: plan.model.dtype_width_bytes,
|
|
max_seq_len: plan.model.max_seq_len,
|
|
},
|
|
runtime: StageRuntimeFacts {
|
|
role_id: RoleId(u64::from(stage.stage_index)),
|
|
input_port: PortId("input".to_owned()),
|
|
output_port: PortId("output".to_owned()),
|
|
sampling: if stage.stage_index + 1 == stage.stage_count {
|
|
Some(plan.runtime.sampling)
|
|
} else {
|
|
None
|
|
},
|
|
},
|
|
})
|
|
}
|
|
|
|
fn validate_global_input(input: &PlannerInput) -> Result<(), PlanRejection> {
|
|
if input.stage_count == 0 {
|
|
return Err(reject(PlanRejectionKind::InvalidStageCount));
|
|
}
|
|
if input.model.num_layers < input.stage_count {
|
|
return Err(reject(PlanRejectionKind::ModelStageLayoutMismatch));
|
|
}
|
|
if input.model.max_seq_len == 0 || input.model.dtype_width_bytes == 0 {
|
|
return Err(reject(PlanRejectionKind::InvalidObjectSpec));
|
|
}
|
|
if input.model.hidden_dim == 0 {
|
|
return Err(reject(PlanRejectionKind::UnsupportedShapeOrLayout));
|
|
}
|
|
if !valid_ring(input.activation_ring) || !valid_ring(input.token_ring) {
|
|
return Err(reject(PlanRejectionKind::InvalidRingSpec));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn validated_placements(input: &PlannerInput) -> Result<Vec<StagePlacement>, PlanRejection> {
|
|
let (stages, forced_edges) = match &input.placement {
|
|
PlacementInput::FixedLinear(stages) => (stages.as_slice(), &[][..]),
|
|
PlacementInput::FixedLinearWithEdgeOverride {
|
|
stages,
|
|
forced_activation_edges,
|
|
} => (stages.as_slice(), forced_activation_edges.as_slice()),
|
|
};
|
|
|
|
if forced_edges
|
|
.iter()
|
|
.any(|(producer, consumer)| *consumer != producer.saturating_add(1))
|
|
{
|
|
return Err(reject(PlanRejectionKind::EdgeEndpointMismatch));
|
|
}
|
|
|
|
let candidate_nodes = input
|
|
.candidate_pool
|
|
.iter()
|
|
.copied()
|
|
.collect::<std::collections::BTreeSet<_>>();
|
|
let mut by_stage = std::collections::BTreeMap::new();
|
|
for placement in stages {
|
|
if !candidate_nodes.contains(&placement.node_id) {
|
|
return Err(reject(PlanRejectionKind::UnknownNode));
|
|
}
|
|
if by_stage.insert(placement.stage_index, *placement).is_some() {
|
|
return Err(reject(PlanRejectionKind::DuplicateStageAssignment));
|
|
}
|
|
}
|
|
|
|
let mut dense = Vec::with_capacity(input.stage_count as usize);
|
|
for stage_index in 0..input.stage_count {
|
|
let placement = by_stage
|
|
.remove(&stage_index)
|
|
.ok_or_else(|| reject(PlanRejectionKind::MissingStage))?;
|
|
dense.push(placement);
|
|
}
|
|
Ok(dense)
|
|
}
|
|
|
|
fn model_plan(model: &ModelFacts) -> Result<GgufModelPlan, PlanRejection> {
|
|
let hidden_dim = u32::try_from(model.hidden_dim)
|
|
.map_err(|_| reject(PlanRejectionKind::UnsupportedShapeOrLayout))?;
|
|
let dtype_width_bytes = u32::try_from(model.dtype_width_bytes)
|
|
.map_err(|_| reject(PlanRejectionKind::InvalidObjectSpec))?;
|
|
let max_seq_len = u32::try_from(model.max_seq_len)
|
|
.map_err(|_| reject(PlanRejectionKind::InvalidObjectSpec))?;
|
|
|
|
Ok(GgufModelPlan {
|
|
model_id: model.model_id.clone(),
|
|
gguf_source: model.gguf_source.clone(),
|
|
num_layers: model.num_layers,
|
|
hidden_dim,
|
|
dtype_family: model.dtype_family,
|
|
dtype_width_bytes,
|
|
max_seq_len,
|
|
eos_token_id: model.eos_token_id,
|
|
tokenizer: model.tokenizer.clone(),
|
|
})
|
|
}
|
|
|
|
fn ring_spec_for_direction(mut spec: RingSpec, direction: RingDirection) -> RingSpec {
|
|
spec.direction = direction;
|
|
spec
|
|
}
|
|
|
|
fn layer_range(num_layers: u32, stage_count: u32, stage_index: u32) -> (u32, u32) {
|
|
let start = (u64::from(num_layers) * u64::from(stage_index) / u64::from(stage_count)) as u32;
|
|
let end = (u64::from(num_layers) * u64::from(stage_index + 1) / u64::from(stage_count)) as u32;
|
|
(start, end)
|
|
}
|
|
|
|
fn endpoint_node_id(endpoint: &EdgeEndpoint) -> NodeId {
|
|
match endpoint {
|
|
EdgeEndpoint::Orchestrator { node_id } | EdgeEndpoint::Stage { node_id, .. } => *node_id,
|
|
}
|
|
}
|
|
|
|
fn valid_ring(spec: RingSpec) -> bool {
|
|
spec.data_capacity > 0 && spec.alignment > 0 && spec.alignment.is_power_of_two()
|
|
}
|
|
|
|
fn reject(kind: PlanRejectionKind) -> PlanRejection {
|
|
PlanRejection { kind }
|
|
}
|