feat: docker provisioning
Treat Docker as a first-class provisioning provider for the local e2e cluster and add the weight-shard fetch/validate/bind lifecycle behind it. - docker_cluster_provisioning: model Docker as a provider adapter behind a `DockerCli` boundary (`run_container`/`inspect_ssh_endpoint`/`remove_force`) exposing only per-node ownership primitives - local_e2e_cluster: wire Docker provisioning into the e2e driver (default `swactor-mvp-local-e2e-cluster` image) across the bootstrap/teardown flow - weight_shards: add `ModelArtifactRef` (parses `hf://repo@rev/path`), `ShardAssignment`, `ShardManifest`, and `ValidatedShard` with digest-based validation - shard_fetch: add `ShardLocator` (digest/split/stage -> uri + cache key), `ShardCache`/`ShardFetcher` traits, a `ShardFetchCoordinator`, and typed `FetchError`s - shard_weight_lifecycle: add the `ShardWeightLifecycle` state machine (Idle->Assigned->Located->Fetching->Fetched->Validating->Binding->Ready/Faulted) with a `WorkerShardBinder` trait - provisioner/telemetry: route provision logs onto the datastream via per-node/stream channels (`submit_bytes`) and add shard_fetch/shard_weight_lifecycle/weight_shards guarantee tests Signed-off-by: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
This commit is contained in:
parent
61aa5b972d
commit
8363be74cd
16 changed files with 2017 additions and 445 deletions
|
|
@ -11,7 +11,7 @@ use crate::provisioning::{
|
|||
NodeProvisionSpec, PluginNodeHandle, PluginObservation, PluginObservationSink, PluginSink,
|
||||
ProvisionEvent, ProvisionEventKind, ProvisionLogLine, ProvisionLogStream, ProvisionPlugin,
|
||||
};
|
||||
use crate::telemetry::{MvpProvisionEventRecord, MvpProvisionLogRecord};
|
||||
use crate::telemetry::{MvpProvisionEventRecord, MvpProvisionLogRecord, mvp_provision_log_channel};
|
||||
|
||||
use super::codec::JsonCodec;
|
||||
|
||||
|
|
@ -346,7 +346,10 @@ impl<P: ProvisionPlugin> ProvisionerActor<P> {
|
|||
|
||||
fn emit_log(&self, line: ProvisionLogLine) {
|
||||
if let Some(producer) = &self.telemetry {
|
||||
producer.submit_record(&MvpProvisionLogRecord::new(line));
|
||||
let channel = mvp_provision_log_channel(line.node_id, line.stream);
|
||||
let record = MvpProvisionLogRecord::new(line);
|
||||
let payload = serde_json::to_vec(&record).expect("serialize provisioning log record");
|
||||
producer.submit_bytes(channel, payload);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use std::collections::{HashMap, VecDeque};
|
||||
use std::collections::{BTreeMap, HashMap, VecDeque};
|
||||
use std::io::{BufRead, BufReader, Write};
|
||||
use std::path::PathBuf;
|
||||
use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio};
|
||||
|
|
@ -24,19 +24,20 @@ use mvp_system::actors::orchestrator::{
|
|||
LifecycleEventWire, OrchestratorActor, OrchestratorMsg, OrchestratorReport, RunCommandWire,
|
||||
StageRefWire,
|
||||
};
|
||||
use mvp_system::actors::provisioner::{ProvisionerActor, ProvisionerMsg, ProvisionerReport};
|
||||
use mvp_system::actors::register_mvp_actor_codecs;
|
||||
use mvp_system::arena_manager as arena;
|
||||
use mvp_system::dashboard_view::MvpClusterDashboardView;
|
||||
use mvp_system::distribution_stack::DistributionRuntimeStack;
|
||||
use mvp_system::docker_cluster_provisioning as docker_provision;
|
||||
use mvp_system::driver_pumps as driver_model;
|
||||
use mvp_system::edge_establisher as edge;
|
||||
use mvp_system::engine_builder as engine;
|
||||
use mvp_system::gpu_worker_ctl as worker_ctl;
|
||||
use mvp_system::gpu_worker_ingress_parser as ingress;
|
||||
use mvp_system::node_provisioning as node_provision;
|
||||
use mvp_system::observability_surface as obs;
|
||||
use mvp_system::orchestrator_run_fsm as fsm;
|
||||
use mvp_system::provisioning::{LocalDockerPlugin, NodeProvisionSpec, ProvisionLogStream};
|
||||
use mvp_system::provisioning::{NodeProvisionSpec, ProvisionLogStream};
|
||||
use mvp_system::run_plan as plan;
|
||||
use mvp_system::stage_controller as stage;
|
||||
use mvp_system::tx_rx_edge_actor as edge_actor;
|
||||
|
|
@ -58,6 +59,7 @@ const OBJECT_ALIGNMENT: u64 = 4;
|
|||
const ARENA_BYTES: usize = 16 * 1024;
|
||||
const RING_BYTES: usize = 4096;
|
||||
const EDGE_READY_TIMEOUT: Duration = Duration::from_secs(20);
|
||||
const DEFAULT_RUNTIME_SNAPSHOT_INTERVAL: Duration = Duration::from_millis(500);
|
||||
|
||||
struct MvpDashboard {
|
||||
url: String,
|
||||
|
|
@ -66,6 +68,8 @@ struct MvpDashboard {
|
|||
subscription: DatastreamSubscription,
|
||||
handle: dashboard::DashboardHandle,
|
||||
runtime_position: u64,
|
||||
runtime_snapshot_interval: Duration,
|
||||
last_runtime_snapshot: Option<Instant>,
|
||||
}
|
||||
|
||||
impl MvpDashboard {
|
||||
|
|
@ -78,6 +82,7 @@ impl MvpDashboard {
|
|||
.map_err(|e| format!("invalid MVP_DASHBOARD_PORT: {e}"))?;
|
||||
config.port = port;
|
||||
}
|
||||
let runtime_snapshot_interval = runtime_snapshot_interval_from_env()?;
|
||||
let url = format!("http://127.0.0.1:{}/view/datastream/live", config.port);
|
||||
let handle = dashboard::start_dashboard(config.clone());
|
||||
handle.register_view(Arc::new(MvpClusterDashboardView::new()));
|
||||
|
|
@ -96,15 +101,14 @@ impl MvpDashboard {
|
|||
subscription,
|
||||
handle,
|
||||
runtime_position: 0,
|
||||
runtime_snapshot_interval,
|
||||
last_runtime_snapshot: None,
|
||||
})
|
||||
}
|
||||
|
||||
fn url(&self) -> &str {
|
||||
&self.url
|
||||
}
|
||||
fn producer(&self) -> DatastreamProducer {
|
||||
self.producer.clone()
|
||||
}
|
||||
|
||||
fn drain(&mut self) {
|
||||
self.endpoint.tick();
|
||||
|
|
@ -119,6 +123,23 @@ impl MvpDashboard {
|
|||
self.drain();
|
||||
}
|
||||
|
||||
fn record_provision_log(&mut self, line: mvp_system::provisioning::ProvisionLogLine) {
|
||||
let channel = mvp_system::telemetry::mvp_provision_log_channel(line.node_id, line.stream);
|
||||
let record = mvp_system::telemetry::MvpProvisionLogRecord::new(line);
|
||||
let payload = serde_json::to_vec(&record).expect("serialize provisioning log record");
|
||||
self.producer.submit_bytes(channel, payload);
|
||||
self.drain();
|
||||
}
|
||||
|
||||
fn publish_runtime_snapshot_throttled(&mut self, stack: &DistributionRuntimeStack) {
|
||||
let due = self.last_runtime_snapshot.map_or(true, |last| {
|
||||
last.elapsed() >= self.runtime_snapshot_interval
|
||||
});
|
||||
if self.runtime_snapshot_interval.is_zero() || due {
|
||||
self.publish_runtime_snapshot(stack);
|
||||
}
|
||||
}
|
||||
|
||||
fn publish_runtime_snapshot(&mut self, stack: &DistributionRuntimeStack) {
|
||||
let stats = stack.runtime.stats();
|
||||
let actors = stats
|
||||
|
|
@ -149,6 +170,7 @@ impl MvpDashboard {
|
|||
if !stats.actor_details.is_empty() {
|
||||
self.ingest_runtime_json(RUNTIME_ACTORS, json!({ "actors": &stats.actor_details }));
|
||||
}
|
||||
self.last_runtime_snapshot = Some(Instant::now());
|
||||
}
|
||||
|
||||
fn ingest_runtime_json(&mut self, channel: &str, value: Value) {
|
||||
|
|
@ -163,6 +185,17 @@ impl MvpDashboard {
|
|||
}
|
||||
}
|
||||
|
||||
fn runtime_snapshot_interval_from_env() -> Result<Duration, String> {
|
||||
let Some(value) = std::env::var_os("MVP_RUNTIME_SNAPSHOT_MS") else {
|
||||
return Ok(DEFAULT_RUNTIME_SNAPSHOT_INTERVAL);
|
||||
};
|
||||
let millis = value
|
||||
.to_string_lossy()
|
||||
.parse::<u64>()
|
||||
.map_err(|e| format!("invalid MVP_RUNTIME_SNAPSHOT_MS: {e}"))?;
|
||||
Ok(Duration::from_millis(millis))
|
||||
}
|
||||
|
||||
fn main() -> ExitCode {
|
||||
let args = std::env::args().collect::<Vec<_>>();
|
||||
let result = if args.iter().any(|arg| arg == "--role=node") {
|
||||
|
|
@ -191,14 +224,25 @@ struct NodeStdoutLine {
|
|||
kind: String,
|
||||
stage_index: Option<u32>,
|
||||
event: Option<String>,
|
||||
endpoint: Option<EndpointAddr>,
|
||||
node_actor: Option<ActorAddress>,
|
||||
logical_node_id: Option<u64>,
|
||||
}
|
||||
#[derive(Clone, Debug)]
|
||||
|
||||
type LocalDockerNodeProvisioner = docker_provision::DockerNodeProvisioner<
|
||||
LocalE2eDockerCli,
|
||||
LocalE2eBootstrapFactory,
|
||||
LocalE2eBootstrapDatastream,
|
||||
>;
|
||||
|
||||
struct ProvisionedDockerNode {
|
||||
node_id: u64,
|
||||
stage_index: u32,
|
||||
endpoint: EndpointAddr,
|
||||
node_actor: ActorAddress,
|
||||
provider_process_id: Option<u32>,
|
||||
provisioner: LocalDockerNodeProvisioner,
|
||||
events: Receiver<LocalDockerNodeEvent>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
|
|
@ -208,6 +252,13 @@ struct ProvisionStats {
|
|||
stderr_line_count: usize,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
enum LocalDockerNodeEvent {
|
||||
Stdout(String),
|
||||
Stderr(String),
|
||||
Exited(Option<i32>),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
enum DriverIngressEvent {
|
||||
StreamArrived {
|
||||
|
|
@ -275,6 +326,23 @@ fn record_dashboard_event(dashboard: &mut Option<&mut MvpDashboard>, event: obs:
|
|||
dashboard.record_event(event);
|
||||
}
|
||||
}
|
||||
fn record_dashboard_provision_log(
|
||||
dashboard: &mut Option<&mut MvpDashboard>,
|
||||
run_id: u64,
|
||||
node_id: u64,
|
||||
stream: ProvisionLogStream,
|
||||
line: &str,
|
||||
) {
|
||||
if let Some(dashboard) = dashboard.as_deref_mut() {
|
||||
dashboard.record_provision_log(mvp_system::provisioning::ProvisionLogLine {
|
||||
run_id,
|
||||
node_id,
|
||||
stream,
|
||||
line: line.to_owned(),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
fn drain_dashboard(dashboard: &mut Option<&mut MvpDashboard>) {
|
||||
if let Some(dashboard) = dashboard.as_deref_mut() {
|
||||
dashboard.drain();
|
||||
|
|
@ -290,6 +358,15 @@ fn publish_runtime_snapshot(
|
|||
}
|
||||
}
|
||||
|
||||
fn publish_runtime_snapshot_throttled(
|
||||
dashboard: &mut Option<&mut MvpDashboard>,
|
||||
stack: &DistributionRuntimeStack,
|
||||
) {
|
||||
if let Some(dashboard) = dashboard.as_deref_mut() {
|
||||
dashboard.publish_runtime_snapshot_throttled(stack);
|
||||
}
|
||||
}
|
||||
|
||||
fn run_event(run_id: u64, kind: obs::EventKind) -> obs::Event {
|
||||
obs::Event::RunScoped {
|
||||
kind,
|
||||
|
|
@ -398,18 +475,6 @@ fn run_supervisor_once(
|
|||
))
|
||||
.map_err(|e| format!("spawn orchestrator actor: {e}"))?;
|
||||
stack.register_local_actor(driver.register_actor(orchestrator_addr, 1));
|
||||
let provisioner_report = stack
|
||||
.runtime
|
||||
.new_inbox::<ProvisionerReport>()
|
||||
.map_err(|e| format!("provisioner report inbox: {e}"))?;
|
||||
let provisioner_addr = stack
|
||||
.runtime
|
||||
.spawn(ProvisionerActor::new(
|
||||
LocalDockerPlugin::new(format!("mvp-local-e2e-cluster-{}", std::process::id())),
|
||||
stack.runtime.create_sender(),
|
||||
dashboard.as_ref().map(|dashboard| dashboard.producer()),
|
||||
))
|
||||
.map_err(|e| format!("spawn provisioner actor: {e}"))?;
|
||||
publish_runtime_snapshot(&mut dashboard, &stack);
|
||||
let mut provision_stats = ProvisionStats::default();
|
||||
|
||||
|
|
@ -433,9 +498,7 @@ fn run_supervisor_once(
|
|||
let orchestrator_actor_json = serde_json::to_string(&orchestrator_addr)
|
||||
.map_err(|e| format!("serialize orchestrator actor: {e}"))?;
|
||||
|
||||
let node1 = provision_local_docker_node(
|
||||
provisioner_addr,
|
||||
&provisioner_report,
|
||||
let mut node1 = provision_local_docker_node(
|
||||
&mut driver,
|
||||
&stack,
|
||||
&mut dashboard,
|
||||
|
|
@ -451,9 +514,7 @@ fn run_supervisor_once(
|
|||
)?;
|
||||
let node1_endpoint_json = serde_json::to_string(&node1.endpoint)
|
||||
.map_err(|e| format!("serialize node1 endpoint: {e}"))?;
|
||||
let node0 = provision_local_docker_node(
|
||||
provisioner_addr,
|
||||
&provisioner_report,
|
||||
let mut node0 = match provision_local_docker_node(
|
||||
&mut driver,
|
||||
&stack,
|
||||
&mut dashboard,
|
||||
|
|
@ -466,7 +527,13 @@ fn run_supervisor_once(
|
|||
&self_endpoint_json,
|
||||
&orchestrator_actor_json,
|
||||
),
|
||||
)?;
|
||||
) {
|
||||
Ok(node) => node,
|
||||
Err(error) => {
|
||||
let _ = stop_provisioned_nodes(&mut [&mut node1], &mut driver, &stack, &mut dashboard);
|
||||
return Err(error);
|
||||
}
|
||||
};
|
||||
|
||||
for stage in [&stage1, &stage0] {
|
||||
record_dashboard_event(
|
||||
|
|
@ -619,14 +686,14 @@ fn run_supervisor_once(
|
|||
}
|
||||
}
|
||||
}
|
||||
stage_ready_count += drain_provisioner_reports(
|
||||
&provisioner_report,
|
||||
stage_ready_count += drain_provisioned_node_events(
|
||||
&mut [&mut node0, &mut node1],
|
||||
run_id,
|
||||
&mut dashboard,
|
||||
&mut provision_stats,
|
||||
)?;
|
||||
drain_dashboard(&mut dashboard);
|
||||
publish_runtime_snapshot(&mut dashboard, &stack);
|
||||
publish_runtime_snapshot_throttled(&mut dashboard, &stack);
|
||||
|
||||
while let Some(report) = orchestrator_report.try_recv() {
|
||||
match report {
|
||||
|
|
@ -739,13 +806,10 @@ fn run_supervisor_once(
|
|||
&& sent_stop_to_node1
|
||||
{
|
||||
stop_provisioned_nodes(
|
||||
run_id,
|
||||
provisioner_addr,
|
||||
&provisioner_report,
|
||||
&mut [&mut node0, &mut node1],
|
||||
&mut driver,
|
||||
&stack,
|
||||
&mut dashboard,
|
||||
&mut provision_stats,
|
||||
)?;
|
||||
let builder_stage_assignments = topology
|
||||
.events
|
||||
|
|
@ -816,13 +880,10 @@ fn run_supervisor_once(
|
|||
|
||||
record_dashboard_event(&mut dashboard, run_fault_event(run_id));
|
||||
let _ = stop_provisioned_nodes(
|
||||
run_id,
|
||||
provisioner_addr,
|
||||
&provisioner_report,
|
||||
&mut [&mut node0, &mut node1],
|
||||
&mut driver,
|
||||
&stack,
|
||||
&mut dashboard,
|
||||
&mut provision_stats,
|
||||
);
|
||||
Err(format!(
|
||||
"timed out: injected={injected} token_received={token_received} completed={completed} torn_down={torn_down} stop0={sent_stop_to_node0} stop1={sent_stop_to_node1}"
|
||||
|
|
@ -2140,9 +2201,296 @@ fn local_docker_spec(
|
|||
}
|
||||
}
|
||||
|
||||
struct LocalE2eContainer {
|
||||
container_name: String,
|
||||
stdin: ChildStdin,
|
||||
}
|
||||
|
||||
struct LocalE2eDockerCli {
|
||||
spec: NodeProvisionSpec,
|
||||
events: Sender<LocalDockerNodeEvent>,
|
||||
nodes: HashMap<String, LocalE2eContainer>,
|
||||
provider_process_id: Option<u32>,
|
||||
}
|
||||
|
||||
impl LocalE2eDockerCli {
|
||||
fn new(spec: NodeProvisionSpec, events: Sender<LocalDockerNodeEvent>) -> Self {
|
||||
Self {
|
||||
spec,
|
||||
events,
|
||||
nodes: HashMap::new(),
|
||||
provider_process_id: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn provider_process_id(&self) -> Option<u32> {
|
||||
self.provider_process_id
|
||||
}
|
||||
}
|
||||
|
||||
impl docker_provision::DockerCli for LocalE2eDockerCli {
|
||||
fn run_container(
|
||||
&mut self,
|
||||
request: docker_provision::DockerRunRequest,
|
||||
) -> Result<docker_provision::DockerRunResult, docker_provision::DockerCliError> {
|
||||
let mut env = request.env.clone();
|
||||
for (key, value) in &self.spec.env {
|
||||
env.insert(key.clone(), value.clone());
|
||||
}
|
||||
|
||||
let mut command = Command::new("docker");
|
||||
command
|
||||
.arg("run")
|
||||
.arg("--rm")
|
||||
.arg("--add-host")
|
||||
.arg("host.docker.internal:host-gateway")
|
||||
.arg("--name")
|
||||
.arg(&request.container_name)
|
||||
.arg("-i");
|
||||
for (key, value) in &request.labels {
|
||||
command.arg("--label").arg(format!("{key}={value}"));
|
||||
}
|
||||
for (key, value) in &env {
|
||||
command.arg("-e").arg(format!("{key}={value}"));
|
||||
}
|
||||
command.arg(&self.spec.image);
|
||||
for arg in &self.spec.args {
|
||||
command.arg(arg);
|
||||
}
|
||||
|
||||
let mut child = command
|
||||
.stdin(Stdio::piped())
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped())
|
||||
.spawn()
|
||||
.map_err(|e| {
|
||||
docker_provision::DockerCliError::new(format!(
|
||||
"spawn Docker node {}: {e}",
|
||||
self.spec.node_id
|
||||
))
|
||||
})?;
|
||||
|
||||
let provider_process_id = child.id();
|
||||
self.provider_process_id = Some(provider_process_id);
|
||||
let stdin = child.stdin.take().ok_or_else(|| {
|
||||
docker_provision::DockerCliError::new(format!(
|
||||
"Docker node {} stdin missing",
|
||||
self.spec.node_id
|
||||
))
|
||||
})?;
|
||||
let stdout = child.stdout.take().ok_or_else(|| {
|
||||
docker_provision::DockerCliError::new(format!(
|
||||
"Docker node {} stdout missing",
|
||||
self.spec.node_id
|
||||
))
|
||||
})?;
|
||||
let stderr = child.stderr.take().ok_or_else(|| {
|
||||
docker_provision::DockerCliError::new(format!(
|
||||
"Docker node {} stderr missing",
|
||||
self.spec.node_id
|
||||
))
|
||||
})?;
|
||||
|
||||
self.nodes.insert(
|
||||
request.container_name.clone(),
|
||||
LocalE2eContainer {
|
||||
container_name: request.container_name.clone(),
|
||||
stdin,
|
||||
},
|
||||
);
|
||||
|
||||
spawn_local_docker_reader(stdout, self.events.clone(), true);
|
||||
spawn_local_docker_reader(stderr, self.events.clone(), false);
|
||||
let events = self.events.clone();
|
||||
thread::spawn(move || match child.wait() {
|
||||
Ok(status) => {
|
||||
let _ = events.send(LocalDockerNodeEvent::Exited(status.code()));
|
||||
}
|
||||
Err(error) => {
|
||||
let _ = events.send(LocalDockerNodeEvent::Stderr(format!(
|
||||
"wait Docker node: {error}"
|
||||
)));
|
||||
let _ = events.send(LocalDockerNodeEvent::Exited(None));
|
||||
}
|
||||
});
|
||||
|
||||
Ok(docker_provision::DockerRunResult {
|
||||
container_id: request.container_name,
|
||||
})
|
||||
}
|
||||
|
||||
fn inspect_ssh_endpoint(
|
||||
&mut self,
|
||||
container_id: &str,
|
||||
) -> Result<Option<node_provision::SshEndpoint>, docker_provision::DockerCliError> {
|
||||
Ok(Some(node_provision::SshEndpoint {
|
||||
host: "127.0.0.1".to_owned(),
|
||||
port: 0,
|
||||
user: "local-e2e".to_owned(),
|
||||
auth_ref: format!("local-docker:{container_id}"),
|
||||
}))
|
||||
}
|
||||
|
||||
fn remove_force(&mut self, container_id: &str) -> Result<(), docker_provision::DockerCliError> {
|
||||
let Some(mut node) = self.nodes.remove(container_id) else {
|
||||
return Ok(());
|
||||
};
|
||||
let _ = writeln!(node.stdin, "shutdown");
|
||||
let _ = node.stdin.flush();
|
||||
let status = Command::new("docker")
|
||||
.arg("stop")
|
||||
.arg("-t")
|
||||
.arg("2")
|
||||
.arg(&node.container_name)
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.map_err(|e| {
|
||||
docker_provision::DockerCliError::new(format!(
|
||||
"docker stop {}: {e}",
|
||||
node.container_name
|
||||
))
|
||||
})?;
|
||||
if status.success() {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(docker_provision::DockerCliError::new(format!(
|
||||
"docker stop {} exited with {status}",
|
||||
node.container_name
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct LocalE2eBootstrapFactory;
|
||||
|
||||
impl docker_provision::SshBootstrapClientFactory for LocalE2eBootstrapFactory {
|
||||
type Client = LocalE2eBootstrapClient;
|
||||
|
||||
fn client_for(&mut self, _spec: &node_provision::BootstrapSessionSpec) -> Self::Client {
|
||||
LocalE2eBootstrapClient
|
||||
}
|
||||
}
|
||||
|
||||
struct LocalE2eBootstrapClient;
|
||||
|
||||
impl docker_provision::BootstrapSshClient for LocalE2eBootstrapClient {
|
||||
fn connect(
|
||||
&mut self,
|
||||
_endpoint: &node_provision::SshEndpoint,
|
||||
) -> Result<(), docker_provision::BootstrapSshError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn probe_stdout(&mut self) -> Result<(), docker_provision::BootstrapSshError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn read_bootstrap_logs(
|
||||
&mut self,
|
||||
_stdout_sources: &[String],
|
||||
_stderr_sources: &[String],
|
||||
) -> Result<
|
||||
Vec<(node_provision::BootstrapLogStream, String)>,
|
||||
docker_provision::BootstrapSshError,
|
||||
> {
|
||||
Ok(Vec::new())
|
||||
}
|
||||
|
||||
fn run_verify_commands(
|
||||
&mut self,
|
||||
_commands: &[String],
|
||||
) -> Result<(), docker_provision::BootstrapSshError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn start_swactor(
|
||||
&mut self,
|
||||
_command: &str,
|
||||
_join: &node_provision::SwarmJoinSpec,
|
||||
) -> Result<(), docker_provision::BootstrapSshError> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn close(&mut self) {}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct LocalE2eBootstrapDatastream;
|
||||
|
||||
impl node_provision::BootstrapDatastreamSink for LocalE2eBootstrapDatastream {
|
||||
fn record(&mut self, _record: node_provision::BootstrapLogRecord) {}
|
||||
|
||||
fn flush(&mut self) {}
|
||||
}
|
||||
|
||||
fn spawn_local_docker_reader(
|
||||
stream: impl std::io::Read + Send + 'static,
|
||||
events: Sender<LocalDockerNodeEvent>,
|
||||
stdout: bool,
|
||||
) {
|
||||
thread::spawn(move || {
|
||||
let reader = BufReader::new(stream);
|
||||
for next in reader.lines() {
|
||||
let Ok(line) = next else {
|
||||
break;
|
||||
};
|
||||
let event = if stdout {
|
||||
LocalDockerNodeEvent::Stdout(line)
|
||||
} else {
|
||||
LocalDockerNodeEvent::Stderr(line)
|
||||
};
|
||||
if events.send(event).is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
fn logical_node_spec_from_local(spec: &NodeProvisionSpec) -> node_provision::LogicalNodeSpec {
|
||||
let logical_node_id = node_provision::LogicalNodeId(spec.node_id.to_string());
|
||||
node_provision::LogicalNodeSpec {
|
||||
run_id: node_provision::RunId(spec.run_id),
|
||||
logical_node_id: logical_node_id.clone(),
|
||||
group_id: node_provision::NodeGroupId("local-e2e".to_owned()),
|
||||
role: node_provision::RoleId("stage-worker".to_owned()),
|
||||
provider: node_provision::ProviderKind::Docker,
|
||||
shape: node_provision::DesiredNodeShape {
|
||||
image: spec.image.clone(),
|
||||
disk_gb: 0,
|
||||
gpu_name: None,
|
||||
min_gpu_ram_mb: None,
|
||||
min_down_mbps: None,
|
||||
min_up_mbps: None,
|
||||
min_reliability: None,
|
||||
require_verified: false,
|
||||
provider_labels: BTreeMap::from([
|
||||
("mvp.local_e2e".to_owned(), "true".to_owned()),
|
||||
("mvp.node_id".to_owned(), spec.node_id.to_string()),
|
||||
]),
|
||||
},
|
||||
boot: node_provision::BootSpec {
|
||||
ssh_user: "local-e2e".to_owned(),
|
||||
verify_commands: Vec::new(),
|
||||
start_swactor_command: spec.args.join(" "),
|
||||
stdout_sources: Vec::new(),
|
||||
stderr_sources: Vec::new(),
|
||||
timeout_policy: node_provision::BootstrapTimeoutPolicy {
|
||||
ssh_connect_secs: 1,
|
||||
boot_check_secs: 1,
|
||||
swactor_join_secs: 30,
|
||||
},
|
||||
},
|
||||
swarm_join: node_provision::SwarmJoinSpec {
|
||||
orch_swactor_addr: "local-e2e".to_owned(),
|
||||
join_token_ref: "local-e2e".to_owned(),
|
||||
expected_logical_node_id: logical_node_id,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn provision_local_docker_node(
|
||||
provisioner_addr: ActorAddress,
|
||||
provisioner_report: &swactor::runtime::Inbox<ProvisionerReport>,
|
||||
driver: &mut IrohDriver,
|
||||
stack: &DistributionRuntimeStack,
|
||||
dashboard: &mut Option<&mut MvpDashboard>,
|
||||
|
|
@ -2151,50 +2499,71 @@ fn provision_local_docker_node(
|
|||
) -> Result<ProvisionedDockerNode, String> {
|
||||
let expected_run_id = spec.run_id;
|
||||
let expected_node_id = spec.node_id;
|
||||
stack
|
||||
.runtime
|
||||
.send_to(
|
||||
provisioner_addr,
|
||||
ProvisionerMsg::StartNodes {
|
||||
nodes: vec![spec],
|
||||
reply_to: *provisioner_report.addr(),
|
||||
},
|
||||
)
|
||||
.map_err(|e| format!("start provisioning node {expected_node_id}: {e}"))?;
|
||||
let expected_stage_index = spec.stage_index.unwrap_or_default();
|
||||
let logical_spec = logical_node_spec_from_local(&spec);
|
||||
let (events_tx, events_rx) = mpsc::channel();
|
||||
let cli = LocalE2eDockerCli::new(spec, events_tx);
|
||||
let provider = docker_provision::DockerProvider::new(cli);
|
||||
let mut provisioner = docker_provision::DockerNodeProvisioner::new(
|
||||
provider,
|
||||
LocalE2eBootstrapFactory,
|
||||
LocalE2eBootstrapDatastream,
|
||||
);
|
||||
provisioner
|
||||
.start(logical_spec)
|
||||
.map_err(|e| format!("start Docker node {expected_node_id}: {e:?}"))?;
|
||||
let provider_process_id = provisioner.provider().cli().provider_process_id();
|
||||
|
||||
let started = Instant::now();
|
||||
while started.elapsed() < Duration::from_secs(30) {
|
||||
pump_network(driver, stack);
|
||||
drain_dashboard(dashboard);
|
||||
publish_runtime_snapshot(dashboard, stack);
|
||||
while let Some(report) = provisioner_report.try_recv() {
|
||||
match report {
|
||||
ProvisionerReport::NodeLive {
|
||||
run_id,
|
||||
node_id,
|
||||
stage_index,
|
||||
endpoint,
|
||||
node_actor,
|
||||
provider_process_id,
|
||||
} if run_id == expected_run_id && node_id == expected_node_id => {
|
||||
stats.node_live_count += 1;
|
||||
return Ok(ProvisionedDockerNode {
|
||||
node_id,
|
||||
stage_index: stage_index.unwrap_or_default(),
|
||||
endpoint,
|
||||
node_actor,
|
||||
provider_process_id,
|
||||
});
|
||||
publish_runtime_snapshot_throttled(dashboard, stack);
|
||||
while let Ok(event) = events_rx.try_recv() {
|
||||
match event {
|
||||
LocalDockerNodeEvent::Stdout(line) => {
|
||||
handle_local_node_stdout(
|
||||
expected_run_id,
|
||||
expected_node_id,
|
||||
&line,
|
||||
dashboard,
|
||||
stats,
|
||||
);
|
||||
if let Some((endpoint, node_actor, stage_index)) =
|
||||
ready_node_from_stdout(expected_node_id, expected_stage_index, &line)?
|
||||
{
|
||||
provisioner
|
||||
.observe_swactor_join(node_provision::SwactorId(format!(
|
||||
"local-e2e-node-{expected_node_id}"
|
||||
)))
|
||||
.map_err(|e| {
|
||||
format!("complete Docker node {expected_node_id} handoff: {e:?}")
|
||||
})?;
|
||||
stats.node_live_count += 1;
|
||||
return Ok(ProvisionedDockerNode {
|
||||
node_id: expected_node_id,
|
||||
stage_index,
|
||||
endpoint,
|
||||
node_actor,
|
||||
provider_process_id,
|
||||
provisioner,
|
||||
events: events_rx,
|
||||
});
|
||||
}
|
||||
}
|
||||
ProvisionerReport::NodeFailed {
|
||||
run_id,
|
||||
node_id,
|
||||
reason,
|
||||
} if run_id == expected_run_id && node_id == expected_node_id => {
|
||||
return Err(format!("node {node_id} provisioning failed: {reason}"));
|
||||
LocalDockerNodeEvent::Stderr(line) => {
|
||||
handle_local_node_stderr(
|
||||
expected_run_id,
|
||||
expected_node_id,
|
||||
&line,
|
||||
dashboard,
|
||||
stats,
|
||||
);
|
||||
}
|
||||
other => {
|
||||
let _ = handle_provisioner_report(other, expected_run_id, dashboard, stats)?;
|
||||
LocalDockerNodeEvent::Exited(status) => {
|
||||
return Err(format!(
|
||||
"node {expected_node_id} process exited before ready: {status:?}"
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -2203,49 +2572,83 @@ fn provision_local_docker_node(
|
|||
Err(format!("timed out provisioning node {expected_node_id}"))
|
||||
}
|
||||
|
||||
fn drain_provisioner_reports(
|
||||
provisioner_report: &swactor::runtime::Inbox<ProvisionerReport>,
|
||||
fn ready_node_from_stdout(
|
||||
expected_node_id: u64,
|
||||
default_stage_index: u32,
|
||||
line: &str,
|
||||
) -> Result<Option<(EndpointAddr, ActorAddress, u32)>, String> {
|
||||
let Ok(line) = serde_json::from_str::<NodeStdoutLine>(line) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if line.kind != "ready" {
|
||||
return Ok(None);
|
||||
}
|
||||
if line.logical_node_id != Some(expected_node_id) {
|
||||
return Ok(None);
|
||||
}
|
||||
let endpoint = line
|
||||
.endpoint
|
||||
.ok_or_else(|| format!("ready line for node {expected_node_id} missing endpoint"))?;
|
||||
let node_actor = line
|
||||
.node_actor
|
||||
.ok_or_else(|| format!("ready line for node {expected_node_id} missing node actor"))?;
|
||||
Ok(Some((
|
||||
endpoint,
|
||||
node_actor,
|
||||
line.stage_index.unwrap_or(default_stage_index),
|
||||
)))
|
||||
}
|
||||
|
||||
fn drain_provisioned_node_events(
|
||||
nodes: &mut [&mut ProvisionedDockerNode],
|
||||
run_id: u64,
|
||||
dashboard: &mut Option<&mut MvpDashboard>,
|
||||
stats: &mut ProvisionStats,
|
||||
) -> Result<usize, String> {
|
||||
let mut stage_ready_count = 0usize;
|
||||
while let Some(report) = provisioner_report.try_recv() {
|
||||
stage_ready_count += handle_provisioner_report(report, run_id, dashboard, stats)?;
|
||||
for node in nodes.iter_mut() {
|
||||
while let Ok(event) = node.events.try_recv() {
|
||||
match event {
|
||||
LocalDockerNodeEvent::Stdout(line) => {
|
||||
stage_ready_count +=
|
||||
handle_local_node_stdout(run_id, node.node_id, &line, dashboard, stats);
|
||||
}
|
||||
LocalDockerNodeEvent::Stderr(line) => {
|
||||
handle_local_node_stderr(run_id, node.node_id, &line, dashboard, stats);
|
||||
}
|
||||
LocalDockerNodeEvent::Exited(status) => {
|
||||
return Err(format!(
|
||||
"node {} process exited before stop: {status:?}",
|
||||
node.node_id
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(stage_ready_count)
|
||||
}
|
||||
|
||||
fn handle_provisioner_report(
|
||||
report: ProvisionerReport,
|
||||
fn handle_local_node_stdout(
|
||||
run_id: u64,
|
||||
node_id: u64,
|
||||
line: &str,
|
||||
dashboard: &mut Option<&mut MvpDashboard>,
|
||||
stats: &mut ProvisionStats,
|
||||
) -> Result<usize, String> {
|
||||
match report {
|
||||
ProvisionerReport::LogLine {
|
||||
run_id: line_run_id,
|
||||
stream,
|
||||
line,
|
||||
..
|
||||
} if line_run_id == run_id => {
|
||||
match stream {
|
||||
ProvisionLogStream::Stdout => stats.stdout_line_count += 1,
|
||||
ProvisionLogStream::Stderr => stats.stderr_line_count += 1,
|
||||
ProvisionLogStream::Provider => {}
|
||||
}
|
||||
Ok(record_stage_ready_from_stdout(
|
||||
run_id, stream, &line, dashboard,
|
||||
))
|
||||
}
|
||||
ProvisionerReport::NodeFailed {
|
||||
run_id: failed_run_id,
|
||||
node_id,
|
||||
reason,
|
||||
} if failed_run_id == run_id => Err(format!("node {node_id} failed: {reason}")),
|
||||
ProvisionerReport::NodeLive { .. } | ProvisionerReport::NodesStopped { .. } => Ok(0),
|
||||
ProvisionerReport::LogLine { .. } | ProvisionerReport::NodeFailed { .. } => Ok(0),
|
||||
}
|
||||
) -> usize {
|
||||
stats.stdout_line_count += 1;
|
||||
record_dashboard_provision_log(dashboard, run_id, node_id, ProvisionLogStream::Stdout, line);
|
||||
record_stage_ready_from_stdout(run_id, ProvisionLogStream::Stdout, line, dashboard)
|
||||
}
|
||||
|
||||
fn handle_local_node_stderr(
|
||||
run_id: u64,
|
||||
node_id: u64,
|
||||
line: &str,
|
||||
dashboard: &mut Option<&mut MvpDashboard>,
|
||||
stats: &mut ProvisionStats,
|
||||
) {
|
||||
stats.stderr_line_count += 1;
|
||||
record_dashboard_provision_log(dashboard, run_id, node_id, ProvisionLogStream::Stderr, line);
|
||||
}
|
||||
|
||||
fn record_stage_ready_from_stdout(
|
||||
|
|
@ -2282,42 +2685,20 @@ fn record_stage_ready_from_stdout(
|
|||
}
|
||||
|
||||
fn stop_provisioned_nodes(
|
||||
run_id: u64,
|
||||
provisioner_addr: ActorAddress,
|
||||
provisioner_report: &swactor::runtime::Inbox<ProvisionerReport>,
|
||||
nodes: &mut [&mut ProvisionedDockerNode],
|
||||
driver: &mut IrohDriver,
|
||||
stack: &DistributionRuntimeStack,
|
||||
dashboard: &mut Option<&mut MvpDashboard>,
|
||||
stats: &mut ProvisionStats,
|
||||
) -> Result<(), String> {
|
||||
stack
|
||||
.runtime
|
||||
.send_to(
|
||||
provisioner_addr,
|
||||
ProvisionerMsg::StopNodes {
|
||||
run_id,
|
||||
reply_to: *provisioner_report.addr(),
|
||||
},
|
||||
)
|
||||
.map_err(|e| format!("stop provisioned nodes: {e}"))?;
|
||||
let started = Instant::now();
|
||||
while started.elapsed() < Duration::from_secs(10) {
|
||||
pump_network(driver, stack);
|
||||
drain_dashboard(dashboard);
|
||||
publish_runtime_snapshot(dashboard, stack);
|
||||
while let Some(report) = provisioner_report.try_recv() {
|
||||
match report {
|
||||
ProvisionerReport::NodesStopped {
|
||||
run_id: stopped_run_id,
|
||||
} if stopped_run_id == run_id => return Ok(()),
|
||||
other => {
|
||||
let _ = handle_provisioner_report(other, run_id, dashboard, stats)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
thread::sleep(Duration::from_millis(10));
|
||||
for node in nodes.iter_mut().rev() {
|
||||
node.provisioner
|
||||
.stop()
|
||||
.map_err(|e| format!("stop node {}: {e:?}", node.node_id))?;
|
||||
}
|
||||
Err("timed out stopping provisioned nodes".to_owned())
|
||||
pump_network(driver, stack);
|
||||
drain_dashboard(dashboard);
|
||||
publish_runtime_snapshot(dashboard, stack);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn spawn_send_pump(
|
||||
|
|
|
|||
|
|
@ -17,11 +17,8 @@ use crate::telemetry::{
|
|||
MvpProvisionEventRecord, MvpProvisionLogRecord,
|
||||
};
|
||||
|
||||
const CHANNELS: &[&str] = &[
|
||||
MVP_LIFECYCLE,
|
||||
MVP_PROVISIONING_EVENTS,
|
||||
MVP_PROVISIONING_LOGS,
|
||||
];
|
||||
const CHANNELS: &[&str] = &[];
|
||||
const PROVISIONING_LOG_PREFIX: &str = "mvp.provisioning.logs.node.";
|
||||
const EVENT_LOG_CAP: usize = 256;
|
||||
const LOG_TAIL_CAP: usize = 128;
|
||||
|
||||
|
|
@ -141,7 +138,10 @@ impl DashboardView for MvpClusterDashboardView {
|
|||
state.apply_provision_event(frame, record);
|
||||
}
|
||||
}
|
||||
MVP_PROVISIONING_LOGS => {
|
||||
channel
|
||||
if channel == MVP_PROVISIONING_LOGS
|
||||
|| channel.starts_with(PROVISIONING_LOG_PREFIX) =>
|
||||
{
|
||||
if let Ok(record) = MvpProvisionLogRecord::decode(&frame.payload) {
|
||||
state.apply_provision_log(record);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,11 +1,12 @@
|
|||
//! Docker-backed provider and bootstrap wiring for node provisioning tests.
|
||||
//! Docker-backed provider and single-node bootstrap wiring.
|
||||
//!
|
||||
//! Docker is treated as a concrete provider adapter here: it creates and destroys
|
||||
//! real Docker container leases through a `DockerCli` boundary. Unit tests use a
|
||||
//! deterministic CLI implementation, but the provider behavior remains the same
|
||||
//! provider contract as a remote adapter.
|
||||
//! real Docker container leases through a `DockerCli` boundary. Higher-level
|
||||
//! callers may start many nodes, but this module intentionally exposes only
|
||||
//! per-node ownership primitives so deployment code does not grow a cluster
|
||||
//! provisioning layer.
|
||||
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::SystemTime;
|
||||
|
||||
use crate::node_provisioning::{
|
||||
|
|
@ -13,8 +14,7 @@ use crate::node_provisioning::{
|
|||
BootstrapObservation, BootstrapSessionEvent, BootstrapSessionSpec, BootstrapStage,
|
||||
CreateLeaseRequest, CreateLeaseResult, DesiredNodeShape, DestroyHandle, LeaseFacts,
|
||||
LogicalNodeId, NodeManager, NodeManagerCommand, NodeManagerMsg, NodeRecord, ProviderError,
|
||||
ProviderKind, ProviderLeaseId, ProviderPlugin, RunId, RunNodeGroupSpec, SshEndpoint, SwactorId,
|
||||
SwarmJoinSpec, expand_node_group,
|
||||
ProviderKind, ProviderLeaseId, ProviderPlugin, RunId, SshEndpoint, SwactorId, SwarmJoinSpec,
|
||||
};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
|
|
@ -391,65 +391,23 @@ impl<C: BootstrapSshClient> SshBootstrapSession<C> {
|
|||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct SwactorJoinEvent {
|
||||
pub logical_node_id: LogicalNodeId,
|
||||
pub swactor_id: SwactorId,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct SwactorJoinRouter {
|
||||
expected_nodes: BTreeSet<LogicalNodeId>,
|
||||
}
|
||||
|
||||
impl SwactorJoinRouter {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn register(&mut self, logical_node_id: LogicalNodeId) {
|
||||
self.expected_nodes.insert(logical_node_id);
|
||||
}
|
||||
|
||||
pub fn route(
|
||||
&self,
|
||||
event: SwactorJoinEvent,
|
||||
manager: &mut NodeManager,
|
||||
) -> Result<Vec<NodeManagerCommand>, DockerClusterError> {
|
||||
if !self.expected_nodes.contains(&event.logical_node_id) {
|
||||
return Err(DockerClusterError::UnknownNode(event.logical_node_id));
|
||||
}
|
||||
manager
|
||||
.handle(NodeManagerMsg::SwactorJoined {
|
||||
logical_node_id: event.logical_node_id,
|
||||
swactor_id: event.swactor_id,
|
||||
})
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))
|
||||
}
|
||||
}
|
||||
|
||||
pub trait SshBootstrapClientFactory {
|
||||
type Client: BootstrapSshClient;
|
||||
|
||||
fn client_for(&mut self, spec: &BootstrapSessionSpec) -> Self::Client;
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ManagedDockerNode<C> {
|
||||
pub manager: NodeManager,
|
||||
pub bootstrap: Option<SshBootstrapSession<C>>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum DockerClusterError {
|
||||
pub enum DockerNodeProvisionError {
|
||||
Provider(String),
|
||||
Node(String),
|
||||
UnknownNode(LogicalNodeId),
|
||||
EndpointUnavailable(ProviderLeaseId),
|
||||
MissingBootstrap(LogicalNodeId),
|
||||
Bootstrap(String),
|
||||
UnexpectedCommand(String),
|
||||
}
|
||||
|
||||
pub struct DockerClusterHarness<D, F, S>
|
||||
pub struct DockerNodeProvisioner<D, F, S>
|
||||
where
|
||||
D: DockerCli,
|
||||
F: SshBootstrapClientFactory,
|
||||
|
|
@ -458,12 +416,12 @@ where
|
|||
provider: DockerProvider<D>,
|
||||
client_factory: F,
|
||||
datastream: S,
|
||||
router: SwactorJoinRouter,
|
||||
nodes: BTreeMap<LogicalNodeId, ManagedDockerNode<F::Client>>,
|
||||
manager: NodeManager,
|
||||
bootstrap: Option<SshBootstrapSession<F::Client>>,
|
||||
teardown_complete: bool,
|
||||
}
|
||||
|
||||
impl<D, F, S> DockerClusterHarness<D, F, S>
|
||||
impl<D, F, S> DockerNodeProvisioner<D, F, S>
|
||||
where
|
||||
D: DockerCli,
|
||||
F: SshBootstrapClientFactory,
|
||||
|
|
@ -474,9 +432,9 @@ where
|
|||
provider,
|
||||
client_factory,
|
||||
datastream,
|
||||
router: SwactorJoinRouter::new(),
|
||||
nodes: BTreeMap::new(),
|
||||
teardown_complete: false,
|
||||
manager: NodeManager::new(),
|
||||
bootstrap: None,
|
||||
teardown_complete: true,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -484,131 +442,96 @@ where
|
|||
&self.provider
|
||||
}
|
||||
|
||||
pub fn provider_mut(&mut self) -> &mut DockerProvider<D> {
|
||||
&mut self.provider
|
||||
}
|
||||
|
||||
pub fn datastream(&self) -> &S {
|
||||
&self.datastream
|
||||
}
|
||||
|
||||
pub fn nodes(&self) -> &BTreeMap<LogicalNodeId, ManagedDockerNode<F::Client>> {
|
||||
&self.nodes
|
||||
pub fn manager(&self) -> &NodeManager {
|
||||
&self.manager
|
||||
}
|
||||
|
||||
pub fn records(&self) -> Vec<NodeRecord> {
|
||||
self.nodes
|
||||
.values()
|
||||
.filter_map(|node| node.manager.record().cloned())
|
||||
.collect()
|
||||
pub fn bootstrap(&self) -> Option<&SshBootstrapSession<F::Client>> {
|
||||
self.bootstrap.as_ref()
|
||||
}
|
||||
|
||||
pub fn all_ready(&self) -> bool {
|
||||
!self.nodes.is_empty() && self.nodes.values().all(|node| node.manager.is_ready())
|
||||
pub fn record(&self) -> Option<&NodeRecord> {
|
||||
self.manager.record()
|
||||
}
|
||||
|
||||
pub fn start_group(&mut self, group: &RunNodeGroupSpec) -> Result<(), DockerClusterError> {
|
||||
for spec in expand_node_group(group) {
|
||||
let logical_node_id = spec.logical_node_id.clone();
|
||||
self.router.register(logical_node_id.clone());
|
||||
let mut manager = NodeManager::new();
|
||||
let commands = manager
|
||||
.handle(NodeManagerMsg::Start(spec))
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))?;
|
||||
match self.process_start_commands(&mut manager, commands) {
|
||||
Ok(bootstrap) => {
|
||||
self.nodes
|
||||
.insert(logical_node_id, ManagedDockerNode { manager, bootstrap });
|
||||
}
|
||||
Err(error) => {
|
||||
self.nodes.insert(
|
||||
logical_node_id,
|
||||
ManagedDockerNode {
|
||||
manager,
|
||||
bootstrap: None,
|
||||
},
|
||||
);
|
||||
let _ = self.teardown();
|
||||
return Err(error);
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn is_ready(&self) -> bool {
|
||||
self.manager.is_ready()
|
||||
}
|
||||
|
||||
pub fn start(
|
||||
&mut self,
|
||||
spec: crate::node_provisioning::LogicalNodeSpec,
|
||||
) -> Result<(), DockerNodeProvisionError> {
|
||||
let commands = self
|
||||
.manager
|
||||
.handle(NodeManagerMsg::Start(spec))
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
self.teardown_complete = false;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn route_join(&mut self, event: SwactorJoinEvent) -> Result<(), DockerClusterError> {
|
||||
let logical_node_id = event.logical_node_id.clone();
|
||||
let node = self
|
||||
.nodes
|
||||
.get_mut(&logical_node_id)
|
||||
.ok_or_else(|| DockerClusterError::UnknownNode(logical_node_id.clone()))?;
|
||||
let commands = self.router.route(event, &mut node.manager)?;
|
||||
for command in commands {
|
||||
match command {
|
||||
NodeManagerCommand::BootstrapConvergenceObserved { swactor_id, .. } => {
|
||||
let bootstrap = node.bootstrap.as_mut().ok_or_else(|| {
|
||||
DockerClusterError::MissingBootstrap(logical_node_id.clone())
|
||||
})?;
|
||||
let events = bootstrap.convergence_observed(swactor_id, &mut self.datastream);
|
||||
Self::feed_bootstrap_events(&mut node.manager, events)?;
|
||||
}
|
||||
other => {
|
||||
return Err(DockerClusterError::Node(format!(
|
||||
"unexpected command {other:?}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
if let Err(error) = self.process_commands(commands) {
|
||||
let _ = self.stop();
|
||||
return Err(error);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn teardown(&mut self) -> Result<(), DockerClusterError> {
|
||||
for node in self.nodes.values_mut() {
|
||||
let commands = node
|
||||
.manager
|
||||
.handle(NodeManagerMsg::Destroy)
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))?;
|
||||
for command in commands {
|
||||
match command {
|
||||
NodeManagerCommand::CancelBootstrap { .. } => {
|
||||
if let Some(bootstrap) = node.bootstrap.as_mut() {
|
||||
let _ = bootstrap.cancel();
|
||||
}
|
||||
}
|
||||
NodeManagerCommand::DestroyLease(handle) => {
|
||||
self.provider
|
||||
.destroy_lease(&handle)
|
||||
.map_err(|error| DockerClusterError::Provider(error.reason))?;
|
||||
node.manager
|
||||
.handle(NodeManagerMsg::LeaseDestroyed)
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))?;
|
||||
}
|
||||
other => {
|
||||
return Err(DockerClusterError::Node(format!(
|
||||
"unexpected command {other:?}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
pub fn observe_swactor_join(
|
||||
&mut self,
|
||||
swactor_id: SwactorId,
|
||||
) -> Result<(), DockerNodeProvisionError> {
|
||||
let logical_node_id = self
|
||||
.manager
|
||||
.record()
|
||||
.ok_or_else(|| DockerNodeProvisionError::Node("node manager not started".into()))?
|
||||
.logical_node_id
|
||||
.clone();
|
||||
let commands = self
|
||||
.manager
|
||||
.handle(NodeManagerMsg::SwactorJoined {
|
||||
logical_node_id,
|
||||
swactor_id,
|
||||
})
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
self.process_commands(commands)
|
||||
}
|
||||
|
||||
pub fn stop(&mut self) -> Result<(), DockerNodeProvisionError> {
|
||||
if self.manager.record().is_none() {
|
||||
self.teardown_complete = true;
|
||||
return Ok(());
|
||||
}
|
||||
let commands = self
|
||||
.manager
|
||||
.handle(NodeManagerMsg::Destroy)
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
self.process_stop_commands(commands)?;
|
||||
self.teardown_complete = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn process_start_commands(
|
||||
fn process_commands(
|
||||
&mut self,
|
||||
manager: &mut NodeManager,
|
||||
commands: Vec<NodeManagerCommand>,
|
||||
) -> Result<Option<SshBootstrapSession<F::Client>>, DockerClusterError> {
|
||||
) -> Result<(), DockerNodeProvisionError> {
|
||||
let mut pending = commands;
|
||||
let mut bootstrap = None;
|
||||
while let Some(command) = pending.pop() {
|
||||
match command {
|
||||
NodeManagerCommand::CreateLease(request) => {
|
||||
let result = self
|
||||
.provider
|
||||
.create_lease(request)
|
||||
.map_err(|error| DockerClusterError::Provider(error.reason))?;
|
||||
let more = manager
|
||||
.map_err(|error| DockerNodeProvisionError::Provider(error.reason))?;
|
||||
let more = self
|
||||
.manager
|
||||
.handle(NodeManagerMsg::LeaseCreated(result))
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))?;
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
pending.extend(more);
|
||||
}
|
||||
NodeManagerCommand::LookupEndpoint(lease) => {
|
||||
|
|
@ -616,58 +539,115 @@ where
|
|||
let endpoint = match self
|
||||
.provider
|
||||
.lookup_endpoint(&lease)
|
||||
.map_err(|error| DockerClusterError::Provider(error.reason))?
|
||||
.map_err(|error| DockerNodeProvisionError::Provider(error.reason))?
|
||||
{
|
||||
Some(endpoint) => endpoint,
|
||||
None => {
|
||||
let _ = manager.handle(NodeManagerMsg::EndpointFailed(
|
||||
let _ = self.manager.handle(NodeManagerMsg::EndpointFailed(
|
||||
"docker ssh endpoint unavailable".into(),
|
||||
));
|
||||
return Err(DockerClusterError::EndpointUnavailable(lease_id));
|
||||
return Err(DockerNodeProvisionError::EndpointUnavailable(lease_id));
|
||||
}
|
||||
};
|
||||
let more = manager
|
||||
let more = self
|
||||
.manager
|
||||
.handle(NodeManagerMsg::EndpointKnown(endpoint))
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))?;
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
pending.extend(more);
|
||||
}
|
||||
NodeManagerCommand::StartBootstrap(spec) => {
|
||||
let logical_node_id = spec.logical_node_id.clone();
|
||||
let client = self.client_factory.client_for(&spec);
|
||||
let mut session = SshBootstrapSession::new(spec, client);
|
||||
let events = session.start(&mut self.datastream);
|
||||
Self::feed_bootstrap_events(manager, events)?;
|
||||
bootstrap = Some(session);
|
||||
let bootstrap_error = events.iter().find_map(|event| match event {
|
||||
BootstrapSessionEvent::Failed(reason) => Some(reason.clone()),
|
||||
BootstrapSessionEvent::Observed(_) | BootstrapSessionEvent::Closed => None,
|
||||
});
|
||||
Self::feed_bootstrap_events(&mut self.manager, events)?;
|
||||
self.bootstrap = Some(session);
|
||||
if let Some(reason) = bootstrap_error {
|
||||
return Err(DockerNodeProvisionError::Bootstrap(format!(
|
||||
"{}: {reason}",
|
||||
logical_node_id.0
|
||||
)));
|
||||
}
|
||||
}
|
||||
other => {
|
||||
return Err(DockerClusterError::Node(format!(
|
||||
"unexpected command {other:?}"
|
||||
NodeManagerCommand::BootstrapConvergenceObserved { swactor_id, .. } => {
|
||||
let logical_node_id = self
|
||||
.manager
|
||||
.record()
|
||||
.ok_or_else(|| {
|
||||
DockerNodeProvisionError::Node("node manager not started".into())
|
||||
})?
|
||||
.logical_node_id
|
||||
.clone();
|
||||
let bootstrap = self.bootstrap.as_mut().ok_or_else(|| {
|
||||
DockerNodeProvisionError::MissingBootstrap(logical_node_id.clone())
|
||||
})?;
|
||||
let events = bootstrap.convergence_observed(swactor_id, &mut self.datastream);
|
||||
Self::feed_bootstrap_events(&mut self.manager, events)?;
|
||||
}
|
||||
NodeManagerCommand::CancelBootstrap { .. }
|
||||
| NodeManagerCommand::DestroyLease(_) => {
|
||||
return Err(DockerNodeProvisionError::UnexpectedCommand(format!(
|
||||
"lifecycle command {command:?} outside stop"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(bootstrap)
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn process_stop_commands(
|
||||
&mut self,
|
||||
commands: Vec<NodeManagerCommand>,
|
||||
) -> Result<(), DockerNodeProvisionError> {
|
||||
for command in commands {
|
||||
match command {
|
||||
NodeManagerCommand::CancelBootstrap { .. } => {
|
||||
if let Some(bootstrap) = self.bootstrap.as_mut() {
|
||||
let _ = bootstrap.cancel();
|
||||
}
|
||||
}
|
||||
NodeManagerCommand::DestroyLease(handle) => {
|
||||
self.provider
|
||||
.destroy_lease(&handle)
|
||||
.map_err(|error| DockerNodeProvisionError::Provider(error.reason))?;
|
||||
self.manager
|
||||
.handle(NodeManagerMsg::LeaseDestroyed)
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
}
|
||||
other => {
|
||||
return Err(DockerNodeProvisionError::UnexpectedCommand(format!(
|
||||
"unexpected stop command {other:?}"
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn feed_bootstrap_events(
|
||||
manager: &mut NodeManager,
|
||||
events: Vec<BootstrapSessionEvent>,
|
||||
) -> Result<(), DockerClusterError> {
|
||||
) -> Result<(), DockerNodeProvisionError> {
|
||||
for event in events {
|
||||
match event {
|
||||
BootstrapSessionEvent::Observed(observation) => {
|
||||
manager
|
||||
.handle(NodeManagerMsg::BootstrapObserved(observation))
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))?;
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
}
|
||||
BootstrapSessionEvent::Failed(reason) => {
|
||||
manager
|
||||
.handle(NodeManagerMsg::BootstrapFailed(reason))
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))?;
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
}
|
||||
BootstrapSessionEvent::Closed => {
|
||||
manager
|
||||
.handle(NodeManagerMsg::BootstrapClosed)
|
||||
.map_err(|error| DockerClusterError::Node(error.reason))?;
|
||||
.map_err(|error| DockerNodeProvisionError::Node(error.reason))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -675,7 +655,7 @@ where
|
|||
}
|
||||
}
|
||||
|
||||
impl<D, F, S> Drop for DockerClusterHarness<D, F, S>
|
||||
impl<D, F, S> Drop for DockerNodeProvisioner<D, F, S>
|
||||
where
|
||||
D: DockerCli,
|
||||
F: SshBootstrapClientFactory,
|
||||
|
|
@ -683,7 +663,7 @@ where
|
|||
{
|
||||
fn drop(&mut self) {
|
||||
if !self.teardown_complete {
|
||||
let _ = self.teardown();
|
||||
let _ = self.stop();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -210,7 +210,6 @@ pub struct DeviceCopyLog {
|
|||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct ObjectRecordBuilder {
|
||||
spec: ObjectSpec,
|
||||
object_id: ObjectId,
|
||||
sequence: u64,
|
||||
extent: u64,
|
||||
|
|
@ -222,9 +221,8 @@ pub struct ObjectRecordBuilder {
|
|||
}
|
||||
|
||||
impl ObjectRecordBuilder {
|
||||
pub fn new(spec: ObjectSpec) -> Self {
|
||||
pub fn new(_spec: ObjectSpec) -> Self {
|
||||
Self {
|
||||
spec,
|
||||
object_id: ObjectId(9000),
|
||||
sequence: 0,
|
||||
extent: 0,
|
||||
|
|
@ -557,6 +555,7 @@ pub fn read_object_record(
|
|||
}))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn object_failure_metadata(
|
||||
bytes: &[u8],
|
||||
reason: ObjectFailureReason,
|
||||
|
|
|
|||
|
|
@ -24,11 +24,14 @@ pub mod orchestrator_token_endpoint;
|
|||
pub mod provisioning;
|
||||
pub mod resource_inventory;
|
||||
pub mod run_plan;
|
||||
pub mod shard_fetch;
|
||||
pub mod shard_weight_lifecycle;
|
||||
pub mod shared_ring_helper_abi;
|
||||
pub mod stage_controller;
|
||||
pub mod telemetry;
|
||||
pub mod tx_rx_edge_actor;
|
||||
pub mod weight_lifecycle;
|
||||
pub mod weight_shards;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
|
@ -36,3 +39,13 @@ mod tests;
|
|||
#[cfg(test)]
|
||||
#[path = "tests/driver_pumps_guarantees.rs"]
|
||||
mod driver_pumps_guarantees;
|
||||
|
||||
#[cfg(test)]
|
||||
#[path = "tests/shard_fetch_guarantees.rs"]
|
||||
mod shard_fetch_guarantees;
|
||||
#[cfg(test)]
|
||||
#[path = "tests/shard_weight_lifecycle_guarantees.rs"]
|
||||
mod shard_weight_lifecycle_guarantees;
|
||||
#[cfg(test)]
|
||||
#[path = "tests/weight_shards_guarantees.rs"]
|
||||
mod weight_shards_guarantees;
|
||||
|
|
|
|||
|
|
@ -86,6 +86,7 @@ pub enum ReadinessCommand {
|
|||
AssignObjectSpec { node_id: NodeId },
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
#[derive(Clone, Default)]
|
||||
struct NodeFacts {
|
||||
known: bool,
|
||||
|
|
|
|||
108
crates/mvp-system/src/shard_fetch.rs
Normal file
108
crates/mvp-system/src/shard_fetch.rs
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
use crate::weight_shards::{ShardAssignment, ShardManifest};
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ShardLocation {
|
||||
pub uri: String,
|
||||
pub cache_key: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct FetchShard {
|
||||
pub location: ShardLocation,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct FetchedShard {
|
||||
pub local_path: String,
|
||||
pub manifest: ShardManifest,
|
||||
}
|
||||
|
||||
impl FetchedShard {
|
||||
pub fn new(local_path: impl Into<String>, manifest: ShardManifest) -> Self {
|
||||
Self {
|
||||
local_path: local_path.into(),
|
||||
manifest,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ShardLocator;
|
||||
|
||||
impl ShardLocator {
|
||||
pub fn locate(assignment: &ShardAssignment) -> ShardLocation {
|
||||
let uri = assignment
|
||||
.model_ref
|
||||
.shard_uri(&assignment.split_id, assignment.stage_index);
|
||||
let cache_key = format!(
|
||||
"{}:{}:{:05}",
|
||||
assignment.expected_model_digest().as_str(),
|
||||
assignment.split_id.as_str(),
|
||||
assignment.stage_index,
|
||||
);
|
||||
ShardLocation { uri, cache_key }
|
||||
}
|
||||
}
|
||||
|
||||
pub trait ShardCache {
|
||||
fn get(&self, cache_key: &str) -> Option<FetchedShard>;
|
||||
fn insert(&mut self, cache_key: String, shard: FetchedShard);
|
||||
}
|
||||
|
||||
pub trait ShardFetcher {
|
||||
fn fetch(&mut self, request: &FetchShard) -> Result<FetchedShard, FetchError>;
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum FetchError {
|
||||
Unauthorized,
|
||||
NotFound,
|
||||
Unavailable,
|
||||
IntegrityMismatch,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ShardFetchStatus {
|
||||
CacheHit,
|
||||
Downloaded,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ShardFetchOutcome {
|
||||
pub shard: FetchedShard,
|
||||
pub location: ShardLocation,
|
||||
pub status: ShardFetchStatus,
|
||||
}
|
||||
|
||||
pub struct ShardFetchCoordinator;
|
||||
|
||||
impl ShardFetchCoordinator {
|
||||
pub fn get_or_fetch<C, F>(
|
||||
assignment: &ShardAssignment,
|
||||
cache: &mut C,
|
||||
fetcher: &mut F,
|
||||
) -> Result<ShardFetchOutcome, FetchError>
|
||||
where
|
||||
C: ShardCache,
|
||||
F: ShardFetcher,
|
||||
{
|
||||
let location = ShardLocator::locate(assignment);
|
||||
if let Some(shard) = cache.get(&location.cache_key) {
|
||||
return Ok(ShardFetchOutcome {
|
||||
shard,
|
||||
location,
|
||||
status: ShardFetchStatus::CacheHit,
|
||||
});
|
||||
}
|
||||
|
||||
let request = FetchShard {
|
||||
location: location.clone(),
|
||||
};
|
||||
let shard = fetcher.fetch(&request)?;
|
||||
cache.insert(location.cache_key.clone(), shard.clone());
|
||||
Ok(ShardFetchOutcome {
|
||||
shard,
|
||||
location,
|
||||
status: ShardFetchStatus::Downloaded,
|
||||
})
|
||||
}
|
||||
}
|
||||
165
crates/mvp-system/src/shard_weight_lifecycle.rs
Normal file
165
crates/mvp-system/src/shard_weight_lifecycle.rs
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
use crate::shard_fetch::{
|
||||
FetchError, ShardCache, ShardFetchCoordinator, ShardFetchStatus, ShardFetcher, ShardLocation,
|
||||
};
|
||||
use crate::weight_shards::{ShardAssignment, ShardValidationError, ValidatedShard};
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ShardLifecycleState {
|
||||
Idle,
|
||||
Assigned,
|
||||
Located,
|
||||
Fetching,
|
||||
Fetched,
|
||||
Validating,
|
||||
Binding,
|
||||
Ready,
|
||||
Faulted,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ShardLifecycleEvent {
|
||||
Assigned { assignment: ShardAssignment },
|
||||
Located { location: ShardLocation },
|
||||
Fetching { location: ShardLocation },
|
||||
CacheHit { cache_key: String },
|
||||
Fetched { uri: String, local_path: String },
|
||||
Validated { local_path: String },
|
||||
Binding { local_path: String },
|
||||
Ready { local_path: String },
|
||||
Faulted { reason: ShardLifecycleFault },
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum ShardLifecycleFault {
|
||||
Fetch(FetchError),
|
||||
Validation(ShardValidationError),
|
||||
Bind(BindError),
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum BindError {
|
||||
WorkerRejected,
|
||||
DeviceAllocationFailed,
|
||||
}
|
||||
|
||||
pub trait WorkerShardBinder {
|
||||
fn bind(&mut self, shard: &ValidatedShard) -> Result<(), BindError>;
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ShardWeightLifecycle {
|
||||
state: ShardLifecycleState,
|
||||
events: Vec<ShardLifecycleEvent>,
|
||||
}
|
||||
|
||||
impl ShardWeightLifecycle {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
state: ShardLifecycleState::Idle,
|
||||
events: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn state(&self) -> ShardLifecycleState {
|
||||
self.state
|
||||
}
|
||||
|
||||
pub fn events(&self) -> &[ShardLifecycleEvent] {
|
||||
&self.events
|
||||
}
|
||||
|
||||
pub fn load<C, F, B>(
|
||||
&mut self,
|
||||
assignment: ShardAssignment,
|
||||
cache: &mut C,
|
||||
fetcher: &mut F,
|
||||
binder: &mut B,
|
||||
) where
|
||||
C: ShardCache,
|
||||
F: ShardFetcher,
|
||||
B: WorkerShardBinder,
|
||||
{
|
||||
if matches!(
|
||||
self.state,
|
||||
ShardLifecycleState::Ready | ShardLifecycleState::Faulted
|
||||
) {
|
||||
return;
|
||||
}
|
||||
|
||||
self.state = ShardLifecycleState::Assigned;
|
||||
self.events.push(ShardLifecycleEvent::Assigned {
|
||||
assignment: assignment.clone(),
|
||||
});
|
||||
|
||||
let location = crate::shard_fetch::ShardLocator::locate(&assignment);
|
||||
self.state = ShardLifecycleState::Located;
|
||||
self.events.push(ShardLifecycleEvent::Located {
|
||||
location: location.clone(),
|
||||
});
|
||||
|
||||
self.state = ShardLifecycleState::Fetching;
|
||||
self.events.push(ShardLifecycleEvent::Fetching {
|
||||
location: location.clone(),
|
||||
});
|
||||
|
||||
let outcome = match ShardFetchCoordinator::get_or_fetch(&assignment, cache, fetcher) {
|
||||
Ok(outcome) => outcome,
|
||||
Err(error) => {
|
||||
self.fault(ShardLifecycleFault::Fetch(error));
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
match outcome.status {
|
||||
ShardFetchStatus::CacheHit => self.events.push(ShardLifecycleEvent::CacheHit {
|
||||
cache_key: outcome.location.cache_key,
|
||||
}),
|
||||
ShardFetchStatus::Downloaded => self.events.push(ShardLifecycleEvent::Fetched {
|
||||
uri: outcome.location.uri,
|
||||
local_path: outcome.shard.local_path.clone(),
|
||||
}),
|
||||
}
|
||||
|
||||
self.state = ShardLifecycleState::Fetched;
|
||||
self.state = ShardLifecycleState::Validating;
|
||||
let validated = match ValidatedShard::new(
|
||||
assignment,
|
||||
outcome.shard.manifest,
|
||||
outcome.shard.local_path.clone(),
|
||||
) {
|
||||
Ok(validated) => validated,
|
||||
Err(error) => {
|
||||
self.fault(ShardLifecycleFault::Validation(error));
|
||||
return;
|
||||
}
|
||||
};
|
||||
self.events.push(ShardLifecycleEvent::Validated {
|
||||
local_path: validated.local_path.clone(),
|
||||
});
|
||||
|
||||
self.state = ShardLifecycleState::Binding;
|
||||
self.events.push(ShardLifecycleEvent::Binding {
|
||||
local_path: validated.local_path.clone(),
|
||||
});
|
||||
if let Err(error) = binder.bind(&validated) {
|
||||
self.fault(ShardLifecycleFault::Bind(error));
|
||||
return;
|
||||
}
|
||||
|
||||
self.state = ShardLifecycleState::Ready;
|
||||
self.events.push(ShardLifecycleEvent::Ready {
|
||||
local_path: validated.local_path,
|
||||
});
|
||||
}
|
||||
|
||||
fn fault(&mut self, reason: ShardLifecycleFault) {
|
||||
self.state = ShardLifecycleState::Faulted;
|
||||
self.events.push(ShardLifecycleEvent::Faulted { reason });
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for ShardWeightLifecycle {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
|
@ -1,10 +1,11 @@
|
|||
//! MVP-system-owned datastream channel records.
|
||||
|
||||
use datastream::frame::ChannelId;
|
||||
use datastream::{ChannelRegistry, Record};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::observability_surface as obs;
|
||||
use crate::provisioning;
|
||||
use crate::provisioning::{self, ProvisionLogStream};
|
||||
|
||||
/// Structured MVP lifecycle facts: run, node, stage, edge, ring, object, step, and worker events.
|
||||
pub const MVP_LIFECYCLE: &str = "mvp.lifecycle";
|
||||
|
|
@ -67,6 +68,15 @@ impl MvpProvisionLogRecord {
|
|||
}
|
||||
}
|
||||
|
||||
pub fn mvp_provision_log_channel(node_id: u64, stream: ProvisionLogStream) -> ChannelId {
|
||||
let stream = match stream {
|
||||
ProvisionLogStream::Stdout => "stdout",
|
||||
ProvisionLogStream::Stderr => "stderr",
|
||||
ProvisionLogStream::Provider => "provider",
|
||||
};
|
||||
ChannelId::new(format!("mvp.provisioning.logs.node.{node_id}.{stream}"))
|
||||
}
|
||||
|
||||
impl Record for MvpProvisionLogRecord {
|
||||
const CHANNEL: &'static str = MVP_PROVISIONING_LOGS;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -205,30 +205,6 @@ impl docker::SshBootstrapClientFactory for FakeSshFactory {
|
|||
}
|
||||
}
|
||||
|
||||
fn started_manager() -> (provision::NodeManager, provision::LogicalNodeId) {
|
||||
let spec = one_spec();
|
||||
let logical_node_id = spec.logical_node_id.clone();
|
||||
let mut manager = provision::NodeManager::new();
|
||||
let start_commands = manager
|
||||
.handle(provision::NodeManagerMsg::Start(spec.clone()))
|
||||
.expect("start succeeds");
|
||||
assert!(matches!(
|
||||
start_commands.as_slice(),
|
||||
[provision::NodeManagerCommand::CreateLease(_)]
|
||||
));
|
||||
let lease = docker::DockerProvider::new(FakeDockerCli::new())
|
||||
.create_lease(provision::CreateLeaseRequest { spec })
|
||||
.expect("docker provider creates lease");
|
||||
let bootstrap_commands = manager
|
||||
.handle(provision::NodeManagerMsg::LeaseCreated(lease))
|
||||
.expect("lease accepted");
|
||||
assert!(matches!(
|
||||
bootstrap_commands.as_slice(),
|
||||
[provision::NodeManagerCommand::StartBootstrap(_)]
|
||||
));
|
||||
(manager, logical_node_id)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn docker_provider_maps_node_spec_to_container_lease_and_destroy_handle() {
|
||||
let spec = one_spec();
|
||||
|
|
@ -330,137 +306,110 @@ fn ssh_bootstrap_session_runs_pre_handoff_steps_and_closes_on_convergence() {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn join_router_routes_known_logical_nodes_and_rejects_unknown_nodes() {
|
||||
let (mut manager, logical_node_id) = started_manager();
|
||||
let mut router = docker::SwactorJoinRouter::new();
|
||||
router.register(logical_node_id.clone());
|
||||
|
||||
let unknown = router.route(
|
||||
docker::SwactorJoinEvent {
|
||||
logical_node_id: provision::LogicalNodeId("workers-99".into()),
|
||||
swactor_id: provision::SwactorId("swactor-wrong".into()),
|
||||
},
|
||||
&mut manager,
|
||||
fn docker_node_provisioner_rejects_join_for_wrong_logical_node() {
|
||||
let spec = one_spec();
|
||||
let mut node = docker::DockerNodeProvisioner::new(
|
||||
docker::DockerProvider::new(FakeDockerCli::new()),
|
||||
FakeSshFactory::default(),
|
||||
provision::InMemoryBootstrapDatastream::default(),
|
||||
);
|
||||
assert!(matches!(
|
||||
unknown,
|
||||
Err(docker::DockerClusterError::UnknownNode(id)) if id == provision::LogicalNodeId("workers-99".into())
|
||||
));
|
||||
|
||||
node.start(spec).expect("node starts");
|
||||
let mut manager = node.manager().clone();
|
||||
let wrong_join = manager.handle(provision::NodeManagerMsg::SwactorJoined {
|
||||
logical_node_id: provision::LogicalNodeId("workers-99".into()),
|
||||
swactor_id: provision::SwactorId("swactor-wrong".into()),
|
||||
});
|
||||
|
||||
assert!(wrong_join.is_err());
|
||||
assert_eq!(
|
||||
manager.record().expect("record exists").stage,
|
||||
node.record().expect("record exists").stage,
|
||||
provision::NodeStage::BootstrapRunning
|
||||
);
|
||||
|
||||
let commands = router
|
||||
.route(
|
||||
docker::SwactorJoinEvent {
|
||||
logical_node_id,
|
||||
swactor_id: provision::SwactorId("swactor-workers-0".into()),
|
||||
},
|
||||
&mut manager,
|
||||
)
|
||||
.expect("known join routes");
|
||||
assert!(matches!(
|
||||
commands.as_slice(),
|
||||
[provision::NodeManagerCommand::BootstrapConvergenceObserved { .. }]
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn docker_cluster_harness_wires_provider_bootstrap_join_and_teardown() {
|
||||
fn docker_node_provisioner_wires_provider_bootstrap_join_and_teardown() {
|
||||
let provider = docker::DockerProvider::new(FakeDockerCli::new());
|
||||
let datastream = provision::InMemoryBootstrapDatastream::default();
|
||||
let mut harness =
|
||||
docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream);
|
||||
let mut node =
|
||||
docker::DockerNodeProvisioner::new(provider, FakeSshFactory::default(), datastream);
|
||||
|
||||
harness
|
||||
.start_group(&docker_group_spec(2))
|
||||
.expect("cluster starts");
|
||||
assert!(!harness.all_ready());
|
||||
assert_eq!(harness.nodes().len(), 2);
|
||||
assert_eq!(harness.provider().cli().run_requests.len(), 2);
|
||||
assert_eq!(harness.datastream().records().len(), 4);
|
||||
node.start(one_spec()).expect("node starts");
|
||||
|
||||
let ids: Vec<_> = harness.nodes().keys().cloned().collect();
|
||||
for id in ids {
|
||||
harness
|
||||
.route_join(docker::SwactorJoinEvent {
|
||||
swactor_id: provision::SwactorId(format!("swactor-{}", id.0)),
|
||||
logical_node_id: id,
|
||||
})
|
||||
.expect("join routes");
|
||||
}
|
||||
|
||||
assert!(harness.all_ready());
|
||||
assert_eq!(harness.datastream().flush_count(), 2);
|
||||
for record in harness.records() {
|
||||
assert_eq!(record.stage, provision::NodeStage::Dormant);
|
||||
assert!(record.ready);
|
||||
assert_eq!(
|
||||
record.lease.as_ref().unwrap().provider,
|
||||
provision::ProviderKind::Docker
|
||||
);
|
||||
}
|
||||
|
||||
harness.teardown().expect("teardown succeeds");
|
||||
assert!(!node.is_ready());
|
||||
assert_eq!(node.provider().cli().run_requests.len(), 1);
|
||||
assert_eq!(node.datastream().records().len(), 2);
|
||||
assert_eq!(
|
||||
harness.provider().cli().removed,
|
||||
vec!["container-1", "container-2"]
|
||||
);
|
||||
for record in harness.records() {
|
||||
assert_eq!(record.stage, provision::NodeStage::Destroyed);
|
||||
assert!(!record.ready);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn docker_cluster_harness_teardown_cleans_known_leases_before_join() {
|
||||
let provider = docker::DockerProvider::new(FakeDockerCli::new());
|
||||
let datastream = provision::InMemoryBootstrapDatastream::default();
|
||||
let mut harness =
|
||||
docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream);
|
||||
|
||||
harness
|
||||
.start_group(&docker_group_spec(1))
|
||||
.expect("cluster starts");
|
||||
let id = provision::LogicalNodeId("workers-0".into());
|
||||
assert_eq!(
|
||||
harness.nodes()[&id].manager.record().unwrap().stage,
|
||||
node.record().expect("record exists").stage,
|
||||
provision::NodeStage::BootstrapRunning
|
||||
);
|
||||
|
||||
harness.teardown().expect("teardown succeeds before join");
|
||||
node.observe_swactor_join(provision::SwactorId("swactor-workers-0".into()))
|
||||
.expect("join completes handoff");
|
||||
|
||||
let node = &harness.nodes()[&id];
|
||||
assert!(node.is_ready());
|
||||
assert_eq!(node.datastream().flush_count(), 1);
|
||||
let record = node.record().expect("record exists");
|
||||
assert_eq!(record.stage, provision::NodeStage::Dormant);
|
||||
assert!(record.ready);
|
||||
assert_eq!(
|
||||
node.manager.record().unwrap().stage,
|
||||
record.lease.as_ref().unwrap().provider,
|
||||
provision::ProviderKind::Docker
|
||||
);
|
||||
assert!(node.bootstrap().expect("bootstrap exists").is_closed());
|
||||
assert!(node.bootstrap().unwrap().client().closed);
|
||||
|
||||
node.stop().expect("teardown succeeds");
|
||||
assert_eq!(node.provider().cli().removed, vec!["container-1"]);
|
||||
let record = node.record().expect("record exists");
|
||||
assert_eq!(record.stage, provision::NodeStage::Destroyed);
|
||||
assert!(!record.ready);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn docker_node_provisioner_teardown_cleans_known_lease_before_join() {
|
||||
let provider = docker::DockerProvider::new(FakeDockerCli::new());
|
||||
let datastream = provision::InMemoryBootstrapDatastream::default();
|
||||
let mut node =
|
||||
docker::DockerNodeProvisioner::new(provider, FakeSshFactory::default(), datastream);
|
||||
|
||||
node.start(one_spec()).expect("node starts");
|
||||
assert_eq!(
|
||||
node.record().unwrap().stage,
|
||||
provision::NodeStage::BootstrapRunning
|
||||
);
|
||||
|
||||
node.stop().expect("teardown succeeds before join");
|
||||
|
||||
assert_eq!(
|
||||
node.record().unwrap().stage,
|
||||
provision::NodeStage::Destroyed
|
||||
);
|
||||
assert!(node.bootstrap.as_ref().unwrap().is_closed());
|
||||
assert!(node.bootstrap.as_ref().unwrap().client().closed);
|
||||
assert_eq!(harness.provider().cli().removed, vec!["container-1"]);
|
||||
assert!(node.bootstrap().unwrap().is_closed());
|
||||
assert!(node.bootstrap().unwrap().client().closed);
|
||||
assert_eq!(node.provider().cli().removed, vec!["container-1"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn docker_cluster_harness_cleans_known_lease_when_endpoint_never_appears() {
|
||||
fn docker_node_provisioner_cleans_known_lease_when_endpoint_never_appears() {
|
||||
let mut cli = FakeDockerCli::new();
|
||||
cli.endpoints.insert("container-1".into(), None);
|
||||
let provider = docker::DockerProvider::new(cli);
|
||||
let datastream = provision::InMemoryBootstrapDatastream::default();
|
||||
let mut harness =
|
||||
docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream);
|
||||
let mut node =
|
||||
docker::DockerNodeProvisioner::new(provider, FakeSshFactory::default(), datastream);
|
||||
|
||||
let result = harness.start_group(&docker_group_spec(1));
|
||||
let result = node.start(one_spec());
|
||||
|
||||
assert!(matches!(
|
||||
result,
|
||||
Err(docker::DockerClusterError::EndpointUnavailable(id))
|
||||
Err(docker::DockerNodeProvisionError::EndpointUnavailable(id))
|
||||
if id == provision::ProviderLeaseId("docker:container-1".into())
|
||||
));
|
||||
assert_eq!(harness.provider().cli().removed, vec!["container-1"]);
|
||||
let id = provision::LogicalNodeId("workers-0".into());
|
||||
assert_eq!(node.provider().cli().removed, vec!["container-1"]);
|
||||
assert_eq!(
|
||||
harness.nodes()[&id].manager.record().unwrap().stage,
|
||||
node.record().unwrap().stage,
|
||||
provision::NodeStage::Destroyed
|
||||
);
|
||||
}
|
||||
|
|
|
|||
168
crates/mvp-system/src/tests/shard_fetch_guarantees.rs
Normal file
168
crates/mvp-system/src/tests/shard_fetch_guarantees.rs
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use mvp_system::shard_fetch as fetch;
|
||||
use mvp_system::weight_shards as shards;
|
||||
|
||||
fn assignment(stage_index: u32) -> shards::ShardAssignment {
|
||||
let model_ref =
|
||||
shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap();
|
||||
let split_scheme = shards::SplitScheme::GgufLayerContiguousV1;
|
||||
let split_id = shards::SplitId::derive(&model_ref, split_scheme);
|
||||
shards::ShardAssignment::new(
|
||||
model_ref,
|
||||
split_id,
|
||||
split_scheme,
|
||||
stage_index,
|
||||
8,
|
||||
shards::LayerRange::new(stage_index * 4, stage_index * 4 + 4).unwrap(),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn fetched_for(assignment: &shards::ShardAssignment) -> fetch::FetchedShard {
|
||||
fetch::FetchedShard::new(
|
||||
format!("/cache/stage-{:05}.gguf", assignment.stage_index),
|
||||
shards::ShardManifest::for_assignment(
|
||||
assignment,
|
||||
shards::ContentHash::literal(format!("sha256:stage-{}", assignment.stage_index))
|
||||
.unwrap(),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct MemoryCache {
|
||||
entries: BTreeMap<String, fetch::FetchedShard>,
|
||||
inserts: usize,
|
||||
}
|
||||
|
||||
impl fetch::ShardCache for MemoryCache {
|
||||
fn get(&self, cache_key: &str) -> Option<fetch::FetchedShard> {
|
||||
self.entries.get(cache_key).cloned()
|
||||
}
|
||||
|
||||
fn insert(&mut self, cache_key: String, shard: fetch::FetchedShard) {
|
||||
self.inserts += 1;
|
||||
self.entries.insert(cache_key, shard);
|
||||
}
|
||||
}
|
||||
|
||||
struct RecordingFetcher {
|
||||
calls: Vec<fetch::FetchShard>,
|
||||
result: Result<fetch::FetchedShard, fetch::FetchError>,
|
||||
}
|
||||
|
||||
impl fetch::ShardFetcher for RecordingFetcher {
|
||||
fn fetch(
|
||||
&mut self,
|
||||
request: &fetch::FetchShard,
|
||||
) -> Result<fetch::FetchedShard, fetch::FetchError> {
|
||||
self.calls.push(request.clone());
|
||||
self.result.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shard_locator_derives_expected_uri() {
|
||||
let assignment = assignment(3);
|
||||
let location = fetch::ShardLocator::locate(&assignment);
|
||||
|
||||
assert_eq!(
|
||||
location.uri,
|
||||
format!(
|
||||
"hf://org/repo@abcdef123456/shards/{}/stage-00003.gguf",
|
||||
assignment.split_id.as_str()
|
||||
)
|
||||
);
|
||||
assert!(location.cache_key.contains(assignment.split_id.as_str()));
|
||||
assert!(location.cache_key.ends_with(":00003"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_depends_on_model_split_and_stage() {
|
||||
let stage_three = assignment(3);
|
||||
let stage_four = assignment(4);
|
||||
|
||||
let different_model_ref =
|
||||
shards::ModelArtifactRef::parse("hf://other/repo@abcdef123456/model.gguf").unwrap();
|
||||
let split_scheme = shards::SplitScheme::GgufLayerContiguousV1;
|
||||
let different_model = shards::ShardAssignment::new(
|
||||
different_model_ref.clone(),
|
||||
shards::SplitId::derive(&different_model_ref, split_scheme),
|
||||
split_scheme,
|
||||
3,
|
||||
8,
|
||||
shards::LayerRange::new(12, 16).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let different_split = shards::ShardAssignment::new(
|
||||
stage_three.model_ref.clone(),
|
||||
shards::SplitId::literal("split-other").unwrap(),
|
||||
stage_three.split_scheme,
|
||||
3,
|
||||
8,
|
||||
stage_three.layer_range,
|
||||
)
|
||||
.unwrap();
|
||||
|
||||
let keys = [
|
||||
fetch::ShardLocator::locate(&stage_three).cache_key,
|
||||
fetch::ShardLocator::locate(&stage_four).cache_key,
|
||||
fetch::ShardLocator::locate(&different_model).cache_key,
|
||||
fetch::ShardLocator::locate(&different_split).cache_key,
|
||||
];
|
||||
|
||||
assert_eq!(
|
||||
keys.iter().collect::<std::collections::BTreeSet<_>>().len(),
|
||||
keys.len()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn coordinator_downloads_and_caches_missing_shard() {
|
||||
let assignment = assignment(2);
|
||||
let expected = fetched_for(&assignment);
|
||||
let mut cache = MemoryCache::default();
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Ok(expected.clone()),
|
||||
};
|
||||
|
||||
let outcome =
|
||||
fetch::ShardFetchCoordinator::get_or_fetch(&assignment, &mut cache, &mut fetcher).unwrap();
|
||||
|
||||
assert_eq!(outcome.status, fetch::ShardFetchStatus::Downloaded);
|
||||
assert_eq!(outcome.shard, expected);
|
||||
assert_eq!(fetcher.calls.len(), 1);
|
||||
assert_eq!(fetcher.calls[0].location, outcome.location);
|
||||
assert_eq!(cache.inserts, 1);
|
||||
assert_eq!(
|
||||
cache.entries.get(&outcome.location.cache_key),
|
||||
Some(&expected)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn coordinator_uses_cache_hit_without_fetching() {
|
||||
let assignment = assignment(5);
|
||||
let location = fetch::ShardLocator::locate(&assignment);
|
||||
let cached = fetched_for(&assignment);
|
||||
let mut cache = MemoryCache::default();
|
||||
cache
|
||||
.entries
|
||||
.insert(location.cache_key.clone(), cached.clone());
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Err(fetch::FetchError::Unavailable),
|
||||
};
|
||||
|
||||
let outcome =
|
||||
fetch::ShardFetchCoordinator::get_or_fetch(&assignment, &mut cache, &mut fetcher).unwrap();
|
||||
|
||||
assert_eq!(outcome.status, fetch::ShardFetchStatus::CacheHit);
|
||||
assert_eq!(outcome.location, location);
|
||||
assert_eq!(outcome.shard, cached);
|
||||
assert!(fetcher.calls.is_empty());
|
||||
assert_eq!(cache.inserts, 0);
|
||||
}
|
||||
308
crates/mvp-system/src/tests/shard_weight_lifecycle_guarantees.rs
Normal file
308
crates/mvp-system/src/tests/shard_weight_lifecycle_guarantees.rs
Normal file
|
|
@ -0,0 +1,308 @@
|
|||
use std::collections::BTreeMap;
|
||||
|
||||
use mvp_system::shard_fetch as fetch;
|
||||
use mvp_system::shard_weight_lifecycle as lifecycle;
|
||||
use mvp_system::weight_shards as shards;
|
||||
|
||||
fn assignment() -> shards::ShardAssignment {
|
||||
let model_ref =
|
||||
shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap();
|
||||
let split_scheme = shards::SplitScheme::GgufLayerContiguousV1;
|
||||
let split_id = shards::SplitId::derive(&model_ref, split_scheme);
|
||||
shards::ShardAssignment::new(
|
||||
model_ref,
|
||||
split_id,
|
||||
split_scheme,
|
||||
3,
|
||||
8,
|
||||
shards::LayerRange::new(12, 16).unwrap(),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn fetched_for(assignment: &shards::ShardAssignment) -> fetch::FetchedShard {
|
||||
fetch::FetchedShard::new(
|
||||
"/cache/stage-00003.gguf",
|
||||
shards::ShardManifest::for_assignment(
|
||||
assignment,
|
||||
shards::ContentHash::literal("sha256:stage-3").unwrap(),
|
||||
),
|
||||
)
|
||||
}
|
||||
|
||||
fn event_kinds(events: &[lifecycle::ShardLifecycleEvent]) -> Vec<&'static str> {
|
||||
events
|
||||
.iter()
|
||||
.map(|event| match event {
|
||||
lifecycle::ShardLifecycleEvent::Assigned { .. } => "assigned",
|
||||
lifecycle::ShardLifecycleEvent::Located { .. } => "located",
|
||||
lifecycle::ShardLifecycleEvent::Fetching { .. } => "fetching",
|
||||
lifecycle::ShardLifecycleEvent::CacheHit { .. } => "cache-hit",
|
||||
lifecycle::ShardLifecycleEvent::Fetched { .. } => "fetched",
|
||||
lifecycle::ShardLifecycleEvent::Validated { .. } => "validated",
|
||||
lifecycle::ShardLifecycleEvent::Binding { .. } => "binding",
|
||||
lifecycle::ShardLifecycleEvent::Ready { .. } => "ready",
|
||||
lifecycle::ShardLifecycleEvent::Faulted { .. } => "faulted",
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct MemoryCache {
|
||||
entries: BTreeMap<String, fetch::FetchedShard>,
|
||||
inserts: usize,
|
||||
}
|
||||
|
||||
impl fetch::ShardCache for MemoryCache {
|
||||
fn get(&self, cache_key: &str) -> Option<fetch::FetchedShard> {
|
||||
self.entries.get(cache_key).cloned()
|
||||
}
|
||||
|
||||
fn insert(&mut self, cache_key: String, shard: fetch::FetchedShard) {
|
||||
self.inserts += 1;
|
||||
self.entries.insert(cache_key, shard);
|
||||
}
|
||||
}
|
||||
|
||||
struct RecordingFetcher {
|
||||
calls: Vec<fetch::FetchShard>,
|
||||
result: Result<fetch::FetchedShard, fetch::FetchError>,
|
||||
}
|
||||
|
||||
impl fetch::ShardFetcher for RecordingFetcher {
|
||||
fn fetch(
|
||||
&mut self,
|
||||
request: &fetch::FetchShard,
|
||||
) -> Result<fetch::FetchedShard, fetch::FetchError> {
|
||||
self.calls.push(request.clone());
|
||||
self.result.clone()
|
||||
}
|
||||
}
|
||||
|
||||
struct RecordingBinder {
|
||||
calls: Vec<shards::ValidatedShard>,
|
||||
result: Result<(), lifecycle::BindError>,
|
||||
}
|
||||
|
||||
impl lifecycle::WorkerShardBinder for RecordingBinder {
|
||||
fn bind(&mut self, shard: &shards::ValidatedShard) -> Result<(), lifecycle::BindError> {
|
||||
self.calls.push(shard.clone());
|
||||
self.result.clone()
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_happy_path_reaches_ready() {
|
||||
let assignment = assignment();
|
||||
let fetched = fetched_for(&assignment);
|
||||
let mut cache = MemoryCache::default();
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Ok(fetched.clone()),
|
||||
};
|
||||
let mut binder = RecordingBinder {
|
||||
calls: Vec::new(),
|
||||
result: Ok(()),
|
||||
};
|
||||
let mut lifecycle = lifecycle::ShardWeightLifecycle::new();
|
||||
|
||||
lifecycle.load(assignment.clone(), &mut cache, &mut fetcher, &mut binder);
|
||||
|
||||
assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Ready);
|
||||
assert_eq!(
|
||||
event_kinds(lifecycle.events()),
|
||||
vec![
|
||||
"assigned",
|
||||
"located",
|
||||
"fetching",
|
||||
"fetched",
|
||||
"validated",
|
||||
"binding",
|
||||
"ready"
|
||||
]
|
||||
);
|
||||
assert_eq!(fetcher.calls.len(), 1);
|
||||
assert_eq!(cache.inserts, 1);
|
||||
assert_eq!(binder.calls.len(), 1);
|
||||
assert_eq!(binder.calls[0].assignment, assignment);
|
||||
assert_eq!(binder.calls[0].local_path, fetched.local_path);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_cache_hit_skips_fetch_but_still_validates_and_binds() {
|
||||
let assignment = assignment();
|
||||
let fetched = fetched_for(&assignment);
|
||||
let location = fetch::ShardLocator::locate(&assignment);
|
||||
let mut cache = MemoryCache::default();
|
||||
cache.entries.insert(location.cache_key.clone(), fetched);
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Err(fetch::FetchError::Unavailable),
|
||||
};
|
||||
let mut binder = RecordingBinder {
|
||||
calls: Vec::new(),
|
||||
result: Ok(()),
|
||||
};
|
||||
let mut lifecycle = lifecycle::ShardWeightLifecycle::new();
|
||||
|
||||
lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder);
|
||||
|
||||
assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Ready);
|
||||
assert_eq!(
|
||||
event_kinds(lifecycle.events()),
|
||||
vec![
|
||||
"assigned",
|
||||
"located",
|
||||
"fetching",
|
||||
"cache-hit",
|
||||
"validated",
|
||||
"binding",
|
||||
"ready"
|
||||
]
|
||||
);
|
||||
assert!(fetcher.calls.is_empty());
|
||||
assert_eq!(cache.inserts, 0);
|
||||
assert_eq!(binder.calls.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_fetch_failure_faults_without_validation_or_bind() {
|
||||
let assignment = assignment();
|
||||
let mut cache = MemoryCache::default();
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Err(fetch::FetchError::NotFound),
|
||||
};
|
||||
let mut binder = RecordingBinder {
|
||||
calls: Vec::new(),
|
||||
result: Ok(()),
|
||||
};
|
||||
let mut lifecycle = lifecycle::ShardWeightLifecycle::new();
|
||||
|
||||
lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder);
|
||||
|
||||
assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Faulted);
|
||||
assert_eq!(fetcher.calls.len(), 1);
|
||||
assert!(binder.calls.is_empty());
|
||||
assert!(matches!(
|
||||
lifecycle.events().last(),
|
||||
Some(lifecycle::ShardLifecycleEvent::Faulted {
|
||||
reason: lifecycle::ShardLifecycleFault::Fetch(fetch::FetchError::NotFound)
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_validation_failure_faults_without_bind() {
|
||||
let assignment = assignment();
|
||||
let mut bad_manifest = shards::ShardManifest::for_assignment(
|
||||
&assignment,
|
||||
shards::ContentHash::literal("sha256:stage-3").unwrap(),
|
||||
);
|
||||
bad_manifest.stage_index = 4;
|
||||
let mut cache = MemoryCache::default();
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Ok(fetch::FetchedShard::new(
|
||||
"/cache/bad-stage.gguf",
|
||||
bad_manifest,
|
||||
)),
|
||||
};
|
||||
let mut binder = RecordingBinder {
|
||||
calls: Vec::new(),
|
||||
result: Ok(()),
|
||||
};
|
||||
let mut lifecycle = lifecycle::ShardWeightLifecycle::new();
|
||||
|
||||
lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder);
|
||||
|
||||
assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Faulted);
|
||||
assert!(binder.calls.is_empty());
|
||||
assert!(matches!(
|
||||
lifecycle.events().last(),
|
||||
Some(lifecycle::ShardLifecycleEvent::Faulted {
|
||||
reason: lifecycle::ShardLifecycleFault::Validation(
|
||||
shards::ShardValidationError::StageIndexMismatch
|
||||
)
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn lifecycle_bind_failure_faults_after_validation() {
|
||||
let assignment = assignment();
|
||||
let fetched = fetched_for(&assignment);
|
||||
let mut cache = MemoryCache::default();
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Ok(fetched),
|
||||
};
|
||||
let mut binder = RecordingBinder {
|
||||
calls: Vec::new(),
|
||||
result: Err(lifecycle::BindError::WorkerRejected),
|
||||
};
|
||||
let mut lifecycle = lifecycle::ShardWeightLifecycle::new();
|
||||
|
||||
lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder);
|
||||
|
||||
assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Faulted);
|
||||
assert_eq!(binder.calls.len(), 1);
|
||||
assert!(event_kinds(lifecycle.events()).contains(&"validated"));
|
||||
assert!(matches!(
|
||||
lifecycle.events().last(),
|
||||
Some(lifecycle::ShardLifecycleEvent::Faulted {
|
||||
reason: lifecycle::ShardLifecycleFault::Bind(lifecycle::BindError::WorkerRejected)
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ready_terminal_state_is_idempotent() {
|
||||
let assignment = assignment();
|
||||
let fetched = fetched_for(&assignment);
|
||||
let mut cache = MemoryCache::default();
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Ok(fetched.clone()),
|
||||
};
|
||||
let mut binder = RecordingBinder {
|
||||
calls: Vec::new(),
|
||||
result: Ok(()),
|
||||
};
|
||||
let mut lifecycle = lifecycle::ShardWeightLifecycle::new();
|
||||
|
||||
lifecycle.load(assignment.clone(), &mut cache, &mut fetcher, &mut binder);
|
||||
let event_count = lifecycle.events().len();
|
||||
lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder);
|
||||
|
||||
assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Ready);
|
||||
assert_eq!(lifecycle.events().len(), event_count);
|
||||
assert_eq!(fetcher.calls.len(), 1);
|
||||
assert_eq!(binder.calls.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fault_terminal_state_is_idempotent() {
|
||||
let assignment = assignment();
|
||||
let fetched = fetched_for(&assignment);
|
||||
let mut cache = MemoryCache::default();
|
||||
let mut fetcher = RecordingFetcher {
|
||||
calls: Vec::new(),
|
||||
result: Err(fetch::FetchError::Unavailable),
|
||||
};
|
||||
let mut binder = RecordingBinder {
|
||||
calls: Vec::new(),
|
||||
result: Ok(()),
|
||||
};
|
||||
let mut lifecycle = lifecycle::ShardWeightLifecycle::new();
|
||||
|
||||
lifecycle.load(assignment.clone(), &mut cache, &mut fetcher, &mut binder);
|
||||
let event_count = lifecycle.events().len();
|
||||
fetcher.result = Ok(fetched);
|
||||
lifecycle.load(assignment, &mut cache, &mut fetcher, &mut binder);
|
||||
|
||||
assert_eq!(lifecycle.state(), lifecycle::ShardLifecycleState::Faulted);
|
||||
assert_eq!(lifecycle.events().len(), event_count);
|
||||
assert_eq!(fetcher.calls.len(), 1);
|
||||
assert!(binder.calls.is_empty());
|
||||
}
|
||||
|
|
@ -3,8 +3,9 @@ use mvp_system::observability_surface as obs;
|
|||
use mvp_system::provisioning::{
|
||||
ProvisionEvent, ProvisionEventKind, ProvisionLogLine, ProvisionLogStream,
|
||||
};
|
||||
use mvp_system::telemetry::{self, MvpLifecycleRecord};
|
||||
use mvp_system::telemetry::{MvpProvisionEventRecord, MvpProvisionLogRecord};
|
||||
use mvp_system::telemetry::{
|
||||
self, MvpLifecycleRecord, MvpProvisionEventRecord, MvpProvisionLogRecord,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn mvp_lifecycle_record_round_trips_on_owned_datastream_channel() {
|
||||
|
|
@ -63,6 +64,10 @@ fn provisioning_records_round_trip_on_owned_datastream_channels() {
|
|||
event
|
||||
);
|
||||
assert_eq!(MvpProvisionLogRecord::decode(&log.encode()).unwrap(), log);
|
||||
assert_eq!(
|
||||
telemetry::mvp_provision_log_channel(11, ProvisionLogStream::Stdout).as_str(),
|
||||
"mvp.provisioning.logs.node.11.stdout"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
|
|
|||
134
crates/mvp-system/src/tests/weight_shards_guarantees.rs
Normal file
134
crates/mvp-system/src/tests/weight_shards_guarantees.rs
Normal file
|
|
@ -0,0 +1,134 @@
|
|||
use mvp_system::weight_shards as shards;
|
||||
|
||||
fn model_ref() -> shards::ModelArtifactRef {
|
||||
shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap()
|
||||
}
|
||||
|
||||
fn assignment() -> shards::ShardAssignment {
|
||||
let model_ref = model_ref();
|
||||
let split_scheme = shards::SplitScheme::GgufLayerContiguousV1;
|
||||
let split_id = shards::SplitId::derive(&model_ref, split_scheme);
|
||||
shards::ShardAssignment::new(
|
||||
model_ref,
|
||||
split_id,
|
||||
split_scheme,
|
||||
3,
|
||||
8,
|
||||
shards::LayerRange::new(12, 16).unwrap(),
|
||||
)
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn content_hash() -> shards::ContentHash {
|
||||
shards::ContentHash::literal("sha256:test-content").unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn model_ref_canonicalization_is_stable() {
|
||||
let parsed = shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap();
|
||||
let from_parts =
|
||||
shards::ModelArtifactRef::hugging_face("/org/repo/", "abcdef123456", "/model.gguf")
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(parsed, from_parts);
|
||||
assert_eq!(parsed.as_str(), "hf://org/repo@abcdef123456/model.gguf");
|
||||
assert_eq!(parsed.repo(), "org/repo");
|
||||
assert_eq!(parsed.revision(), "abcdef123456");
|
||||
assert_eq!(parsed.path(), "model.gguf");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_id_is_deterministic_and_model_sensitive() {
|
||||
let first = model_ref();
|
||||
let same = shards::ModelArtifactRef::parse("hf://org/repo@abcdef123456/model.gguf").unwrap();
|
||||
let different =
|
||||
shards::ModelArtifactRef::parse("hf://org/repo@fedcba654321/model.gguf").unwrap();
|
||||
let scheme = shards::SplitScheme::GgufLayerContiguousV1;
|
||||
|
||||
assert_eq!(
|
||||
shards::SplitId::derive(&first, scheme),
|
||||
shards::SplitId::derive(&same, scheme)
|
||||
);
|
||||
assert_ne!(
|
||||
shards::SplitId::derive(&first, scheme),
|
||||
shards::SplitId::derive(&different, scheme)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn assignment_rejects_invalid_stage_shape_and_ranges() {
|
||||
let model_ref = model_ref();
|
||||
let scheme = shards::SplitScheme::GgufLayerContiguousV1;
|
||||
let split_id = shards::SplitId::derive(&model_ref, scheme);
|
||||
let range = shards::LayerRange::new(1, 2).unwrap();
|
||||
|
||||
assert_eq!(
|
||||
shards::ShardAssignment::new(model_ref.clone(), split_id.clone(), scheme, 0, 0, range),
|
||||
Err(shards::ShardAssignmentError::EmptyStageCount)
|
||||
);
|
||||
assert_eq!(
|
||||
shards::ShardAssignment::new(model_ref, split_id, scheme, 2, 2, range),
|
||||
Err(shards::ShardAssignmentError::StageIndexOutOfRange)
|
||||
);
|
||||
assert_eq!(
|
||||
shards::LayerRange::new(4, 4),
|
||||
Err(shards::LayerRangeError::EmptyOrInverted)
|
||||
);
|
||||
assert_eq!(
|
||||
shards::LayerRange::new(5, 4),
|
||||
Err(shards::LayerRangeError::EmptyOrInverted)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validator_accepts_matching_manifest() {
|
||||
let assignment = assignment();
|
||||
let manifest = shards::ShardManifest::for_assignment(&assignment, content_hash());
|
||||
|
||||
assert_eq!(
|
||||
shards::ShardValidator::validate(&assignment, &manifest),
|
||||
Ok(())
|
||||
);
|
||||
assert!(shards::ValidatedShard::new(assignment, manifest, "/cache/stage-00003.gguf").is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn validator_rejects_mismatched_manifest() {
|
||||
let assignment = assignment();
|
||||
let matching = shards::ShardManifest::for_assignment(&assignment, content_hash());
|
||||
|
||||
let mut wrong_model = matching.clone();
|
||||
wrong_model.model_digest = shards::ModelDigest::literal("wrong-model").unwrap();
|
||||
assert_eq!(
|
||||
shards::ShardValidator::validate(&assignment, &wrong_model),
|
||||
Err(shards::ShardValidationError::ModelDigestMismatch)
|
||||
);
|
||||
|
||||
let mut wrong_split = matching.clone();
|
||||
wrong_split.split_id = shards::SplitId::literal("split-wrong").unwrap();
|
||||
assert_eq!(
|
||||
shards::ShardValidator::validate(&assignment, &wrong_split),
|
||||
Err(shards::ShardValidationError::SplitIdMismatch)
|
||||
);
|
||||
|
||||
let mut wrong_stage = matching.clone();
|
||||
wrong_stage.stage_index = 4;
|
||||
assert_eq!(
|
||||
shards::ShardValidator::validate(&assignment, &wrong_stage),
|
||||
Err(shards::ShardValidationError::StageIndexMismatch)
|
||||
);
|
||||
|
||||
let mut wrong_count = matching.clone();
|
||||
wrong_count.stage_count = 9;
|
||||
assert_eq!(
|
||||
shards::ShardValidator::validate(&assignment, &wrong_count),
|
||||
Err(shards::ShardValidationError::StageCountMismatch)
|
||||
);
|
||||
|
||||
let mut wrong_range = matching;
|
||||
wrong_range.layer_range = shards::LayerRange::new(16, 20).unwrap();
|
||||
assert_eq!(
|
||||
shards::ShardValidator::validate(&assignment, &wrong_range),
|
||||
Err(shards::ShardValidationError::LayerRangeMismatch)
|
||||
);
|
||||
}
|
||||
348
crates/mvp-system/src/weight_shards.rs
Normal file
348
crates/mvp-system/src/weight_shards.rs
Normal file
|
|
@ -0,0 +1,348 @@
|
|||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ModelArtifactRef {
|
||||
canonical: String,
|
||||
repo: String,
|
||||
revision: String,
|
||||
path: String,
|
||||
}
|
||||
|
||||
impl ModelArtifactRef {
|
||||
pub fn parse(value: impl Into<String>) -> Result<Self, ModelArtifactRefError> {
|
||||
let value = value.into();
|
||||
let rest = value
|
||||
.strip_prefix("hf://")
|
||||
.ok_or(ModelArtifactRefError::InvalidScheme)?;
|
||||
let (repo, revision_and_path) = rest
|
||||
.split_once('@')
|
||||
.ok_or(ModelArtifactRefError::MissingRevision)?;
|
||||
let (revision, path) = revision_and_path
|
||||
.split_once('/')
|
||||
.ok_or(ModelArtifactRefError::MissingPath)?;
|
||||
Self::hugging_face(repo, revision, path)
|
||||
}
|
||||
|
||||
pub fn hugging_face(
|
||||
repo: impl Into<String>,
|
||||
revision: impl Into<String>,
|
||||
path: impl Into<String>,
|
||||
) -> Result<Self, ModelArtifactRefError> {
|
||||
let repo = repo.into().trim_matches('/').to_owned();
|
||||
let revision = revision.into();
|
||||
let path = path.into().trim_start_matches('/').to_owned();
|
||||
if repo.is_empty() {
|
||||
return Err(ModelArtifactRefError::MissingRepo);
|
||||
}
|
||||
if revision.is_empty() {
|
||||
return Err(ModelArtifactRefError::MissingRevision);
|
||||
}
|
||||
if revision.contains('/') {
|
||||
return Err(ModelArtifactRefError::RevisionMustBePathSegment);
|
||||
}
|
||||
if path.is_empty() {
|
||||
return Err(ModelArtifactRefError::MissingPath);
|
||||
}
|
||||
let canonical = format!("hf://{repo}@{revision}/{path}");
|
||||
Ok(Self {
|
||||
canonical,
|
||||
repo,
|
||||
revision,
|
||||
path,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.canonical
|
||||
}
|
||||
|
||||
pub fn repo(&self) -> &str {
|
||||
&self.repo
|
||||
}
|
||||
|
||||
pub fn revision(&self) -> &str {
|
||||
&self.revision
|
||||
}
|
||||
|
||||
pub fn path(&self) -> &str {
|
||||
&self.path
|
||||
}
|
||||
|
||||
pub fn model_digest(&self) -> ModelDigest {
|
||||
ModelDigest(stable_digest_hex(&["model", self.as_str()]))
|
||||
}
|
||||
|
||||
pub fn shard_uri(&self, split_id: &SplitId, stage_index: u32) -> String {
|
||||
format!(
|
||||
"hf://{}@{}/shards/{}/stage-{stage_index:05}.gguf",
|
||||
self.repo,
|
||||
self.revision,
|
||||
split_id.as_str(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ModelArtifactRefError {
|
||||
InvalidScheme,
|
||||
MissingRepo,
|
||||
MissingRevision,
|
||||
RevisionMustBePathSegment,
|
||||
MissingPath,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum SplitScheme {
|
||||
GgufLayerContiguousV1,
|
||||
}
|
||||
|
||||
impl SplitScheme {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
Self::GgufLayerContiguousV1 => "gguf-layer-contiguous-v1",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct SplitId(String);
|
||||
|
||||
impl SplitId {
|
||||
pub fn derive(model_ref: &ModelArtifactRef, scheme: SplitScheme) -> Self {
|
||||
Self(format!(
|
||||
"split-{}",
|
||||
stable_digest_hex(&["split", model_ref.as_str(), scheme.as_str()])
|
||||
))
|
||||
}
|
||||
|
||||
pub fn literal(value: impl Into<String>) -> Result<Self, SplitIdError> {
|
||||
let value = value.into();
|
||||
if value.is_empty() {
|
||||
return Err(SplitIdError::Empty);
|
||||
}
|
||||
if !value
|
||||
.bytes()
|
||||
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_')
|
||||
{
|
||||
return Err(SplitIdError::InvalidCharacter);
|
||||
}
|
||||
Ok(Self(value))
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum SplitIdError {
|
||||
Empty,
|
||||
InvalidCharacter,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct ModelDigest(String);
|
||||
|
||||
impl ModelDigest {
|
||||
pub fn literal(value: impl Into<String>) -> Result<Self, ModelDigestError> {
|
||||
let value = value.into();
|
||||
if value.is_empty() {
|
||||
return Err(ModelDigestError::Empty);
|
||||
}
|
||||
Ok(Self(value))
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ModelDigestError {
|
||||
Empty,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct ContentHash(String);
|
||||
|
||||
impl ContentHash {
|
||||
pub fn literal(value: impl Into<String>) -> Result<Self, ContentHashError> {
|
||||
let value = value.into();
|
||||
if value.is_empty() {
|
||||
return Err(ContentHashError::Empty);
|
||||
}
|
||||
Ok(Self(value))
|
||||
}
|
||||
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ContentHashError {
|
||||
Empty,
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
|
||||
pub struct LayerRange {
|
||||
pub start: u32,
|
||||
pub end_exclusive: u32,
|
||||
}
|
||||
|
||||
impl LayerRange {
|
||||
pub fn new(start: u32, end_exclusive: u32) -> Result<Self, LayerRangeError> {
|
||||
if start >= end_exclusive {
|
||||
return Err(LayerRangeError::EmptyOrInverted);
|
||||
}
|
||||
Ok(Self {
|
||||
start,
|
||||
end_exclusive,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum LayerRangeError {
|
||||
EmptyOrInverted,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ShardAssignment {
|
||||
pub model_ref: ModelArtifactRef,
|
||||
pub split_id: SplitId,
|
||||
pub split_scheme: SplitScheme,
|
||||
pub stage_index: u32,
|
||||
pub stage_count: u32,
|
||||
pub layer_range: LayerRange,
|
||||
}
|
||||
|
||||
impl ShardAssignment {
|
||||
pub fn new(
|
||||
model_ref: ModelArtifactRef,
|
||||
split_id: SplitId,
|
||||
split_scheme: SplitScheme,
|
||||
stage_index: u32,
|
||||
stage_count: u32,
|
||||
layer_range: LayerRange,
|
||||
) -> Result<Self, ShardAssignmentError> {
|
||||
if stage_count == 0 {
|
||||
return Err(ShardAssignmentError::EmptyStageCount);
|
||||
}
|
||||
if stage_index >= stage_count {
|
||||
return Err(ShardAssignmentError::StageIndexOutOfRange);
|
||||
}
|
||||
Ok(Self {
|
||||
model_ref,
|
||||
split_id,
|
||||
split_scheme,
|
||||
stage_index,
|
||||
stage_count,
|
||||
layer_range,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn expected_model_digest(&self) -> ModelDigest {
|
||||
self.model_ref.model_digest()
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ShardAssignmentError {
|
||||
EmptyStageCount,
|
||||
StageIndexOutOfRange,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ShardManifest {
|
||||
pub model_digest: ModelDigest,
|
||||
pub split_id: SplitId,
|
||||
pub stage_index: u32,
|
||||
pub stage_count: u32,
|
||||
pub layer_range: LayerRange,
|
||||
pub content_hash: ContentHash,
|
||||
}
|
||||
|
||||
impl ShardManifest {
|
||||
pub fn for_assignment(assignment: &ShardAssignment, content_hash: ContentHash) -> Self {
|
||||
Self {
|
||||
model_digest: assignment.expected_model_digest(),
|
||||
split_id: assignment.split_id.clone(),
|
||||
stage_index: assignment.stage_index,
|
||||
stage_count: assignment.stage_count,
|
||||
layer_range: assignment.layer_range,
|
||||
content_hash,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct ValidatedShard {
|
||||
pub assignment: ShardAssignment,
|
||||
pub manifest: ShardManifest,
|
||||
pub local_path: String,
|
||||
}
|
||||
|
||||
impl ValidatedShard {
|
||||
pub fn new(
|
||||
assignment: ShardAssignment,
|
||||
manifest: ShardManifest,
|
||||
local_path: impl Into<String>,
|
||||
) -> Result<Self, ShardValidationError> {
|
||||
ShardValidator::validate(&assignment, &manifest)?;
|
||||
Ok(Self {
|
||||
assignment,
|
||||
manifest,
|
||||
local_path: local_path.into(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ShardValidator;
|
||||
|
||||
impl ShardValidator {
|
||||
pub fn validate(
|
||||
assignment: &ShardAssignment,
|
||||
manifest: &ShardManifest,
|
||||
) -> Result<(), ShardValidationError> {
|
||||
if manifest.model_digest != assignment.expected_model_digest() {
|
||||
return Err(ShardValidationError::ModelDigestMismatch);
|
||||
}
|
||||
if manifest.split_id != assignment.split_id {
|
||||
return Err(ShardValidationError::SplitIdMismatch);
|
||||
}
|
||||
if manifest.stage_index != assignment.stage_index {
|
||||
return Err(ShardValidationError::StageIndexMismatch);
|
||||
}
|
||||
if manifest.stage_count != assignment.stage_count {
|
||||
return Err(ShardValidationError::StageCountMismatch);
|
||||
}
|
||||
if manifest.layer_range != assignment.layer_range {
|
||||
return Err(ShardValidationError::LayerRangeMismatch);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum ShardValidationError {
|
||||
ModelDigestMismatch,
|
||||
SplitIdMismatch,
|
||||
StageIndexMismatch,
|
||||
StageCountMismatch,
|
||||
LayerRangeMismatch,
|
||||
}
|
||||
|
||||
fn stable_digest_hex(parts: &[&str]) -> String {
|
||||
const FNV_OFFSET: u64 = 0xcbf29ce484222325;
|
||||
const FNV_PRIME: u64 = 0x100000001b3;
|
||||
|
||||
let mut hash = FNV_OFFSET;
|
||||
for part in parts {
|
||||
for byte in part.as_bytes() {
|
||||
hash ^= u64::from(*byte);
|
||||
hash = hash.wrapping_mul(FNV_PRIME);
|
||||
}
|
||||
hash ^= 0xff;
|
||||
hash = hash.wrapping_mul(FNV_PRIME);
|
||||
}
|
||||
format!("{hash:016x}")
|
||||
}
|
||||
Loading…
Reference in a new issue