Nvidia/BossAI
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Integrate ComfyUI agents and video workflows

Browse Source
This commit is contained in:
boos
2026-05-22 04:46:15 +08:00
parent 3c12738138
commit c1074ddf71
35 changed files with 2527 additions and 213 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/services/builtin_keys.rs
+1 -1
View File
@@ -166,7 +166,7 @@ mod tests {
.fetch_optional(db)
.await
.unwrap()
}
}
async fn read_row(db: &SqlitePool) -> Option<(i64, String, i64)> {
sqlx::query_as(
src/services/comfyui_pool.rs
+25 -15
View File
@@ -230,12 +230,7 @@ impl ComfyUIPool {
let map = self.containers.read().await;
let live = map
.values()
.filter(|e| {
matches!(
e.state,
ContainerState::Starting | ContainerState::Ready
)
})
.filter(|e| matches!(e.state, ContainerState::Starting | ContainerState::Ready))
.count();
if live >= self.config.max_concurrent {
return Err(PoolError::AtCapacity(live));
@@ -292,11 +287,7 @@ impl ComfyUIPool {
}
/// Wait for the container to answer `/system_stats`, then mark Ready.
async fn probe_until_ready(
&self,
user_id: i64,
entry: PoolEntry,
) -> PoolResult<PoolEntry> {
async fn probe_until_ready(&self, user_id: i64, entry: PoolEntry) -> PoolResult<PoolEntry> {
let url = format!("http://127.0.0.1:{}/system_stats", entry.host_port);
let deadline = Instant::now() + READINESS_TIMEOUT;
@@ -464,7 +455,10 @@ impl ComfyUIPool {
"BOOSAPI_BUILTIN_KEY".into(),
self.config.builtin_key.clone(),
);
env.insert("BOOSAPI_BASE".into(), self.config.boosapi_base.clone());
env.insert(
"BOOSAPI_BASE".into(),
container_reachable_boosapi_base(&self.config.boosapi_base),
);
env.insert("BOOSAPI_KEY".into(), self.config.builtin_key.clone());
ContainerSpec {
@@ -488,6 +482,24 @@ impl ComfyUIPool {
}
}
fn container_reachable_boosapi_base(base: &str) -> String {
if let Ok(override_base) = std::env::var("BOOSAPI_COMFY_CONTAINER_BASE") {
let trimmed = override_base.trim();
if !trimmed.is_empty() {
return trimmed.trim_end_matches('/').to_string();
}
}
let trimmed = base.trim().trim_end_matches('/');
for local_host in ["127.0.0.1", "localhost"] {
let needle = format!("//{}:", local_host);
if let Some((prefix, rest)) = trimmed.split_once(&needle) {
return format!("{}//host.docker.internal:{}", prefix, rest);
}
}
trimmed.to_string()
}
/// Spawn the idle-reaper task. Lives for the life of the process.
pub fn spawn_reaper(pool: Arc<ComfyUIPool>) {
tokio::spawn(async move {
@@ -529,9 +541,7 @@ mod tests {
use async_trait::async_trait;
use std::sync::atomic::{AtomicUsize, Ordering};
use crate::services::container_runtime::{
ContainerHandle, ContainerListing, RuntimeResult,
};
use crate::services::container_runtime::{ContainerHandle, ContainerListing, RuntimeResult};
/// In-memory runtime that succeeds and pretends `/system_stats` is up.
/// Real readiness probe is tested separately via the live Docker setup.
src/services/docker_runtime.rs
+13 -10
View File
@@ -11,8 +11,7 @@ use tokio::process::Command;
use tracing::{debug, warn};
use super::container_runtime::{
ContainerHandle, ContainerListing, ContainerRuntime, ContainerSpec, RuntimeError,
RuntimeResult,
ContainerHandle, ContainerListing, ContainerRuntime, ContainerSpec, RuntimeError, RuntimeResult,
};
/// Internal port the ComfyUI process listens on inside the container.
@@ -74,6 +73,8 @@ impl ContainerRuntime for DockerRuntime {
"127.0.0.1:{}:{}",
spec.host_port, CONTAINER_INTERNAL_PORT
))
.arg("--add-host")
.arg("host.docker.internal:host-gateway")
// Resource limits (§design Decision 5)
.arg(format!("--cpus={}", spec.limits.cpus))
.arg(format!("--memory={}", spec.limits.memory_bytes))
@@ -173,12 +174,7 @@ impl ContainerRuntime for DockerRuntime {
async fn is_alive(&self, name: &str) -> RuntimeResult<bool> {
let output = self
.cmd()
.args([
"inspect",
"-f",
"{{.State.Running}}",
name,
])
.args(["inspect", "-f", "{{.State.Running}}", name])
.output()
.await
.map_err(io_to_runtime)?;
@@ -230,11 +226,18 @@ impl ContainerRuntime for DockerRuntime {
let host_port = match parse_host_port(ports) {
Some(p) => p,
None => {
warn!(name, ports, "could not parse host port from docker ps; skipping");
warn!(
name,
ports, "could not parse host port from docker ps; skipping"
);
continue;
}
};
out.push(ContainerListing { id, name, host_port });
out.push(ContainerListing {
id,
name,
host_port,
});
}
Ok(out)
}
src/services/route_rebuild.rs
+31 -13
View File
@@ -124,8 +124,7 @@ pub(crate) async fn rebuild_routes_inner(
sqlx::query_as("SELECT model_pattern, id FROM token_routes")
.fetch_all(&mut *tx)
.await?;
let mut route_id_by_pattern: HashMap<String, i64> =
existing_routes.into_iter().collect();
let mut route_id_by_pattern: HashMap<String, i64> = existing_routes.into_iter().collect();
// 4. For each model, ensure a token_routes row, then reconcile channels.
for (model_name, account_ids) in &by_model {
@@ -169,7 +168,10 @@ pub(crate) async fn rebuild_routes_inner(
let mut auto_by_key: HashMap<AutoKey, i64> = HashMap::new();
let mut manual_keys: HashSet<AutoKey> = HashSet::new();
for (ch_id, account_id, source_model, manual) in existing_channels {
let key = (account_id, source_model.unwrap_or_else(|| model_name.clone()));
let key = (
account_id,
source_model.unwrap_or_else(|| model_name.clone()),
);
if manual.unwrap_or(false) {
manual_keys.insert(key);
summary.skipped_manual_channels += 1;
@@ -294,10 +296,24 @@ pub(crate) async fn rebuild_routes_inner(
// has no channels (channels will be re-inserted on next refresh).
continue;
}
let touched_by_operator = display_name.as_deref().map(str::trim).unwrap_or("").is_empty()
let touched_by_operator = display_name
.as_deref()
.map(str::trim)
.unwrap_or("")
.is_empty()
== false
|| display_icon.as_deref().map(str::trim).unwrap_or("").is_empty() == false
|| model_mapping.as_deref().map(str::trim).unwrap_or("").is_empty() == false;
|| display_icon
.as_deref()
.map(str::trim)
.unwrap_or("")
.is_empty()
== false
|| model_mapping
.as_deref()
.map(str::trim)
.unwrap_or("")
.is_empty()
== false;
if touched_by_operator {
continue;
}
@@ -493,12 +509,11 @@ mod tests {
"route is pinned by a manual_override channel and must survive"
);
let manual_count: i64 = sqlx::query_scalar(
"SELECT COUNT(*) FROM route_channels WHERE manual_override = 1",
)
.fetch_one(&pool)
.await
.unwrap();
let manual_count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM route_channels WHERE manual_override = 1")
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(manual_count, 1);
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
.fetch_one(&pool)
@@ -539,7 +554,10 @@ mod tests {
route_count, 0,
"all phantom routes must be swept when no account advertises any model"
);
assert_eq!(summary.deleted_routes, 3, "summary reports the three sweeps");
assert_eq!(
summary.deleted_routes, 3,
"summary reports the three sweeps"
);
assert_eq!(summary.created_routes, 0);
assert_eq!(summary.created_channels, 0);
}