254 lines
8.6 KiB
Rust
254 lines
8.6 KiB
Rust
use makepad_ai_hub::client::{ArtifactBytes, ContentProvider};
|
|
use makepad_ai_hub::error::AssetAiError;
|
|
use makepad_ai_hub::lease::{LeaseTable, Origin, Renew};
|
|
use makepad_ai_hub::protocol::{
|
|
ArtifactRefJson, GenerateRequestJson, HealthJson, JobStatusJson, ModelInfoJson,
|
|
JOB_STATE_CANCELLED, JOB_STATE_DONE, JOB_STATE_RUNNING,
|
|
};
|
|
use makepad_ai_hub::registry::Domain;
|
|
use makepad_flow::engine::executors::gen::GenExecutor;
|
|
use makepad_ai_hub::job::{GenPick, GenRouter};
|
|
use makepad_flow::engine::executors::{Executor, Poll};
|
|
use makepad_flow::graph::evaluate;
|
|
use makepad_micro_serde::SerJson;
|
|
use std::collections::BTreeMap;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
const BYTES: &[u8] = b"test image";
|
|
|
|
#[derive(Default)]
|
|
struct Hub {
|
|
leases: LeaseTable,
|
|
jobs: BTreeMap<String, JobStatusJson>,
|
|
requests: Vec<GenerateRequestJson>,
|
|
bye_calls: usize,
|
|
}
|
|
|
|
impl Hub {
|
|
fn complete_active(&mut self) {
|
|
for (id, status) in &mut self.jobs {
|
|
if status.state == JOB_STATE_RUNNING {
|
|
status.state = JOB_STATE_DONE.to_string();
|
|
status.artifacts = vec![ArtifactRefJson {
|
|
id: "artifact".into(),
|
|
url: "/artifact/artifact".into(),
|
|
content_type: "image/png".into(),
|
|
sha256: Some(makepad_ai_hub::sha256::sha256_hex(BYTES)),
|
|
byte_len: Some(BYTES.len() as u64),
|
|
}];
|
|
self.leases.release(id);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone, Default)]
|
|
struct Farm {
|
|
hubs: [Arc<Mutex<Hub>>; 2],
|
|
}
|
|
|
|
impl GenRouter for Farm {
|
|
|
|
fn route(
|
|
&self,
|
|
_domain: &str,
|
|
request: &GenerateRequestJson,
|
|
excluded: &[String],
|
|
) -> Result<GenPick, String> {
|
|
let (index, url) = ["http://hub-a", "http://hub-b"]
|
|
.into_iter()
|
|
.enumerate()
|
|
.find(|(_, url)| !excluded.iter().any(|excluded| excluded == url))
|
|
.ok_or_else(|| "no alternate provider".to_string())?;
|
|
Ok(GenPick {
|
|
provider: Box::new(Provider {
|
|
hub: self.hubs[index].clone(),
|
|
origin: Mutex::new(None),
|
|
reject_for_local_use: index == 0 && request.prompt.as_deref() == Some("retry"),
|
|
}),
|
|
base_url: url.into(),
|
|
model: request.model.clone(),
|
|
model_state: Some("ready".into()),
|
|
})
|
|
}
|
|
}
|
|
|
|
struct Provider {
|
|
hub: Arc<Mutex<Hub>>,
|
|
origin: Mutex<Option<Origin>>,
|
|
reject_for_local_use: bool,
|
|
}
|
|
|
|
impl ContentProvider for Provider {
|
|
fn health(&self) -> Result<HealthJson, AssetAiError> {
|
|
Err(AssetAiError::Unavailable("unused".into()))
|
|
}
|
|
|
|
fn list_models(&self) -> Result<Vec<ModelInfoJson>, AssetAiError> {
|
|
Err(AssetAiError::Unavailable("unused".into()))
|
|
}
|
|
|
|
fn request(
|
|
&self,
|
|
_domain: Domain,
|
|
request: &GenerateRequestJson,
|
|
) -> Result<String, AssetAiError> {
|
|
let origin = Origin {
|
|
node_key: request.origin_key.clone().expect("leased generation"),
|
|
epoch: request.origin_epoch.expect("leased generation epoch"),
|
|
};
|
|
assert!(origin.node_key.len() <= 128, "AIHub origin key limit");
|
|
assert_eq!(origin.epoch, 7);
|
|
*self.origin.lock().unwrap() = Some(origin.clone());
|
|
let mut hub = self.hub.lock().unwrap();
|
|
let job_id = format!("job-{}", hub.jobs.len() + 1);
|
|
assert_eq!(hub.leases.register(&job_id, origin, 0), Renew::Renewed);
|
|
let state = if self.reject_for_local_use {
|
|
hub.leases.release(&job_id);
|
|
JOB_STATE_CANCELLED
|
|
} else {
|
|
JOB_STATE_RUNNING
|
|
};
|
|
hub.jobs.insert(
|
|
job_id.clone(),
|
|
JobStatusJson {
|
|
job_id: job_id.clone(),
|
|
state: state.into(),
|
|
stage: Some("generating".into()),
|
|
progress: Some(0.5),
|
|
artifacts: Vec::new(),
|
|
error: self
|
|
.reject_for_local_use
|
|
.then(|| "local-use: gpu-counter-unavailable".into()),
|
|
model: Some(request.model.clone()),
|
|
queued_ms: None,
|
|
started_ms: None,
|
|
finished_ms: None,
|
|
log: None,
|
|
partial_text: None,
|
|
partial_artifacts: None,
|
|
live: None,
|
|
serving: None,
|
|
text: None,
|
|
},
|
|
);
|
|
hub.requests.push(request.clone());
|
|
Ok(job_id)
|
|
}
|
|
|
|
fn poll(&self, job_id: &str) -> Result<JobStatusJson, AssetAiError> {
|
|
Ok(self.hub.lock().unwrap().jobs[job_id].clone())
|
|
}
|
|
|
|
fn fetch_artifact(&self, _artifact_id: &str) -> Result<ArtifactBytes, AssetAiError> {
|
|
Ok(ArtifactBytes {
|
|
content_type: "image/png".into(),
|
|
bytes: BYTES.to_vec(),
|
|
})
|
|
}
|
|
|
|
fn cancel(&self, job_id: &str) -> Result<JobStatusJson, AssetAiError> {
|
|
let mut hub = self.hub.lock().unwrap();
|
|
hub.leases.release(job_id);
|
|
hub.jobs.get_mut(job_id).unwrap().state = JOB_STATE_CANCELLED.into();
|
|
Ok(hub.jobs[job_id].clone())
|
|
}
|
|
|
|
fn bye(&self) -> Result<(), AssetAiError> {
|
|
let origin = self
|
|
.origin
|
|
.lock()
|
|
.unwrap()
|
|
.clone()
|
|
.expect("submitted origin");
|
|
let mut hub = self.hub.lock().unwrap();
|
|
hub.bye_calls += 1;
|
|
// Exercise AIHub's actual origin-wide lease semantics, including the
|
|
// sibling cancellation that a provider-local fake would conceal.
|
|
for (job_id, _) in hub.leases.bye(&origin.node_key) {
|
|
hub.jobs.get_mut(&job_id).unwrap().state = JOB_STATE_CANCELLED.into();
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
fn node(prompt: &str) -> makepad_flow::Node {
|
|
evaluate(
|
|
&format!("use mod.flow.*\nlet image = Image{{model: \"flux1-schnell\" prompt: \"{prompt}\" width: 64 height: 64 steps: 4 seed: 77}}\nFlow{{image}}\n"),
|
|
"lease-isolation.splash",
|
|
).unwrap().nodes.into_iter().find(|node| node.id == "image").unwrap()
|
|
}
|
|
|
|
fn executor(farm: &Farm) -> GenExecutor {
|
|
// Every concept card is owned by the same Flow host and epoch. Use the
|
|
// maximum wire key length so a derived identity cannot append past it.
|
|
GenExecutor::new(Arc::new(farm.clone()), ("h".repeat(128), 7))
|
|
}
|
|
|
|
#[test]
|
|
fn cancelling_one_generation_leaves_a_sibling_on_the_same_hub_running() {
|
|
let farm = Farm::default();
|
|
let mut cancelled = executor(&farm);
|
|
let mut sibling = executor(&farm);
|
|
cancelled.start(&node("cancel"), &[]).unwrap();
|
|
sibling.start(&node("sibling"), &[]).unwrap();
|
|
|
|
cancelled.cancel();
|
|
assert!(matches!(cancelled.poll(), Poll::Pending));
|
|
assert!(
|
|
matches!(sibling.poll(), Poll::Progress { .. }),
|
|
"origin goodbye cancelled the sibling"
|
|
);
|
|
{
|
|
let hub = farm.hubs[0].lock().unwrap();
|
|
assert_eq!(hub.bye_calls, 1);
|
|
assert_eq!(hub.leases.len(), 1);
|
|
assert_eq!(
|
|
hub.requests.len(),
|
|
2,
|
|
"the sibling was not silently retried"
|
|
);
|
|
}
|
|
farm.hubs[0].lock().unwrap().complete_active();
|
|
assert!(matches!(sibling.poll(), Poll::Done(_)));
|
|
}
|
|
|
|
#[test]
|
|
fn local_use_retry_preserves_its_request_without_cancelling_a_sibling() {
|
|
let farm = Farm::default();
|
|
let mut retrying = executor(&farm);
|
|
let mut sibling = executor(&farm);
|
|
retrying
|
|
.start(&node("retry"), &[("prompt".into(), makepad_flow::PortValue::One(makepad_flow::Value::text("retry")))])
|
|
.unwrap();
|
|
sibling.start(&node("sibling"), &[]).unwrap();
|
|
|
|
let Poll::Progress { stage, .. } = retrying.poll() else {
|
|
panic!("no failover");
|
|
};
|
|
assert!(stage.contains("retrying on hub-b"), "{stage}");
|
|
assert!(
|
|
matches!(sibling.poll(), Poll::Progress { .. }),
|
|
"retry goodbye cancelled the sibling"
|
|
);
|
|
{
|
|
let original = farm.hubs[0].lock().unwrap();
|
|
let alternate = farm.hubs[1].lock().unwrap();
|
|
assert_eq!(original.bye_calls, 1);
|
|
assert_eq!(original.leases.len(), 1);
|
|
assert_eq!(original.requests.len(), 2);
|
|
assert_eq!(alternate.requests.len(), 1);
|
|
// Model, dimensions, seed and lease identity all survive failover.
|
|
assert_eq!(
|
|
original.requests[0].serialize_json(),
|
|
alternate.requests[0].serialize_json()
|
|
);
|
|
assert_eq!(alternate.requests[0].seed, Some(77));
|
|
assert_eq!(alternate.requests[0].model, "flux1-schnell");
|
|
}
|
|
for hub in &farm.hubs {
|
|
hub.lock().unwrap().complete_active();
|
|
}
|
|
assert!(matches!(retrying.poll(), Poll::Done(_)));
|
|
assert!(matches!(sibling.poll(), Poll::Done(_)));
|
|
}
|