diff --git a/backend/gradient-ci/src/webhooks/routing.rs b/backend/gradient-ci/src/webhooks/routing.rs index 8eb168b21..d0672e79a 100644 --- a/backend/gradient-ci/src/webhooks/routing.rs +++ b/backend/gradient-ci/src/webhooks/routing.rs @@ -20,6 +20,7 @@ pub fn routes_to(hook: &MWebhook, owner: &EventOwner, name: &str, personal: bool _ if personal => false, WebhookScope::Project => hook.project.is_some() && hook.project == owner.project, WebhookScope::Cache => hook.cache.is_some() && hook.cache == owner.cache, + WebhookScope::Team => hook.team.is_some() && hook.team == owner.team, } } @@ -34,6 +35,9 @@ pub async fn candidates( if let Some(cache) = owner.cache { scope = scope.add(CWebhook::Cache.eq(cache)); } + if let Some(team) = owner.team { + scope = scope.add(CWebhook::Team.eq(team)); + } EWebhook::find() .filter(CWebhook::Active.eq(true)) .filter(scope) @@ -116,6 +120,33 @@ mod tests { )); } + fn team_hook(team: TeamId) -> MWebhook { + MWebhook { + team: Some(team), + ..hook(WebhookScope::Team, None, None, &[]) + } + } + + #[test] + fn team_hook_receives_its_team_event() { + let a = TeamId::new(Uuid::from_u128(4)); + assert!(routes_to( + &team_hook(a), + &EventOwner::team(a), + "worker.demand", + false + )); + } + + #[test] + fn team_hook_ignores_other_team_and_project_events() { + let h = team_hook(TeamId::new(Uuid::from_u128(4))); + let other_team = EventOwner::team(TeamId::new(Uuid::from_u128(5))); + let project = EventOwner::project(ProjectId::new(Uuid::from_u128(1))); + assert!(!routes_to(&h, &other_team, "worker.demand", false)); + assert!(!routes_to(&h, &project, "build.failed", false)); + } + #[test] fn an_empty_event_list_subscribes_to_everything() { let h = hook(WebhookScope::Instance, None, None, &[]); diff --git a/backend/gradient-db/src/graph/runtime_can_start.rs b/backend/gradient-db/src/graph/runtime_can_start.rs index 2cc72aafc..2749c1c79 100644 --- a/backend/gradient-db/src/graph/runtime_can_start.rs +++ b/backend/gradient-db/src/graph/runtime_can_start.rs @@ -7,6 +7,7 @@ //! The seed cannot read the transition off the row it is writing. //! The commit making an output present wrote `cached_path` one statement earlier. //! `freshly_present` is the caller's endpoint for that transition. +//! A freshly present dependency still counts as missing for a referrer in the same seed, since its own ripple subtracts it. //! A re-push of an already backed path must not count as fresh. use std::collections::BTreeMap; @@ -21,14 +22,15 @@ use crate::sql::StatusWrite; fn seed_sql() -> String { format!( - "UPDATE derivation_build db SET missing_runtime_deps = x.n \ + "WITH s(derivation, fresh) AS (SELECT * FROM unnest($1::uuid[], $2::bool[])) \ + UPDATE derivation_build db SET missing_runtime_deps = x.n \ FROM (SELECT b.id, ({was_complete} AND NOT s.fresh) AS was_complete, \ (SELECT count(*) FROM derivation_dependency e \ JOIN derivation_build dep ON dep.derivation = e.dependency \ + LEFT JOIN s landed ON landed.derivation = e.dependency AND landed.fresh \ WHERE e.derivation = b.derivation AND e.kind IN (1, 2) \ - AND NOT {dep_complete})::int AS n \ - FROM unnest($1::uuid[], $2::bool[]) AS s(derivation, fresh) \ - JOIN derivation_build b ON b.derivation = s.derivation) x \ + AND NOT ({dep_complete} AND landed.derivation IS NULL))::int AS n \ + FROM s JOIN derivation_build b ON b.derivation = s.derivation) x \ WHERE db.id = x.id \ RETURNING db.derivation, x.was_complete, {complete} AS complete", was_complete = shared_build_complete_predicate("b"), diff --git a/backend/gradient-db/src/teams/demand.rs b/backend/gradient-db/src/teams/demand.rs new file mode 100644 index 000000000..b7a7b144c --- /dev/null +++ b/backend/gradient-db/src/teams/demand.rs @@ -0,0 +1,213 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use super::workers::projects_granted_with_workers; +use gradient_entity::evaluation::EvaluationStatus; +use gradient_types::*; +use sea_orm::sea_query::Query; +use sea_orm::{ColumnTrait, ConnectionTrait, DbErr, EntityTrait, QueryFilter}; +use serde::Serialize; +use std::collections::BTreeMap; + +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, utoipa::ToSchema)] +pub struct TeamDemand { + pub builds: Vec, + pub fetch: u32, + pub eval: u32, +} + +pub async fn team_demand( + db: &impl ConnectionTrait, + team: TeamId, + projects: Option<&[ProjectId]>, +) -> Result { + let granting: Vec = projects_granted_with_workers(db, team) + .await? + .into_iter() + .filter(|p| projects.is_none_or(|only| only.contains(p))) + .collect(); + if granting.is_empty() { + return Ok(TeamDemand::default()); + } + + let tasks = Query::select() + .column(CTask::Id) + .from(ETask::default()) + .and_where(CTask::Project.is_in(granting)) + .to_owned(); + let reasons = EEvaluation::find() + .filter(CEvaluation::Status.eq(EvaluationStatus::Waiting)) + .filter(CEvaluation::Task.in_subquery(tasks)) + .all(db) + .await? + .into_iter() + .filter_map(|e| e.waiting_reason.as_ref().and_then(WaitingReason::from_json)); + Ok(aggregate(reasons)) +} + +fn aggregate(reasons: impl IntoIterator) -> TeamDemand { + let mut demand = TeamDemand::default(); + let mut builds: BTreeMap<(String, Vec), u32> = BTreeMap::new(); + for reason in reasons { + match reason { + WaitingReason::Workers { unmet, .. } => { + for requirement in unmet { + let mut features = requirement.required_features; + features.sort(); + *builds + .entry((requirement.architecture, features)) + .or_default() += requirement.build_count; + } + } + WaitingReason::EvalWorkers { capability, .. } => match capability { + EvalCapability::Fetch => demand.fetch += 1, + EvalCapability::Eval => demand.eval += 1, + }, + _ => {} + } + } + demand.builds = builds + .into_iter() + .map( + |((architecture, required_features), build_count)| UnmetRequirement { + architecture, + required_features, + build_count, + }, + ) + .collect(); + demand +} + +#[cfg(test)] +mod tests { + use super::*; + use gradient_fixtures::mock::{all_statements, mock_db}; + + fn grant(project: ProjectId) -> MTeamProject { + MTeamProject { + project, + includes_workers: true, + ..Default::default() + } + } + + fn waiting(reason: WaitingReason) -> MEvaluation { + MEvaluation { + status: EvaluationStatus::Waiting, + waiting_reason: Some(reason.to_json()), + ..Default::default() + } + } + + fn unmet(architecture: &str, features: &[&str], build_count: u32) -> UnmetRequirement { + UnmetRequirement { + architecture: architecture.into(), + required_features: features.iter().map(|f| f.to_string()).collect(), + build_count, + } + } + + #[tokio::test] + async fn merges_unmet_builds_by_system_and_features() { + let project = ProjectId::now_v7(); + let db = mock_db() + .append_query_results([vec![grant(project)]]) + .append_query_results([vec![ + waiting(WaitingReason::workers( + vec![unmet("x86_64-linux", &["kvm", "big"], 2)], + 1, + vec![], + )), + waiting(WaitingReason::workers( + vec![ + unmet("x86_64-linux", &["big", "kvm"], 3), + unmet("aarch64-darwin", &[], 1), + ], + 1, + vec![], + )), + ]]) + .into_connection(); + + let demand = team_demand(&db, TeamId::now_v7(), None) + .await + .expect("demand"); + + assert_eq!( + demand.builds, + vec![ + unmet("aarch64-darwin", &[], 1), + unmet("x86_64-linux", &["big", "kvm"], 5), + ] + ); + } + + #[tokio::test] + async fn counts_eval_and_fetch_gaps() { + let db = mock_db() + .append_query_results([vec![grant(ProjectId::now_v7())]]) + .append_query_results([vec![ + waiting(WaitingReason::eval_workers(EvalCapability::Eval, 0)), + waiting(WaitingReason::eval_workers(EvalCapability::Eval, 2)), + waiting(WaitingReason::eval_workers(EvalCapability::Fetch, 0)), + waiting(WaitingReason::Draining), + ]]) + .into_connection(); + + let demand = team_demand(&db, TeamId::now_v7(), None) + .await + .expect("demand"); + + assert_eq!((demand.eval, demand.fetch), (2, 1)); + assert!(demand.builds.is_empty()); + } + + #[tokio::test] + async fn ignores_projects_without_worker_grant() { + let db = mock_db() + .append_query_results([Vec::::new()]) + .into_connection(); + + let demand = team_demand(&db, TeamId::now_v7(), None) + .await + .expect("demand"); + + assert_eq!(demand, TeamDemand::default()); + let statements = all_statements(db.into_transaction_log()); + assert!( + statements + .iter() + .all(|s| !s.sql.contains("FROM \"evaluation\"")) + ); + assert!( + statements + .iter() + .any(|s| s.sql.contains("\"includes_workers\" = $")) + ); + } + + #[tokio::test] + async fn project_filter_limits_aggregation() { + let (kept, dropped) = (ProjectId::now_v7(), ProjectId::now_v7()); + let db = mock_db() + .append_query_results([vec![grant(kept), grant(dropped)]]) + .append_query_results([Vec::::new()]) + .into_connection(); + + team_demand(&db, TeamId::now_v7(), Some(&[kept])) + .await + .expect("demand"); + + let evaluations = all_statements(db.into_transaction_log()) + .into_iter() + .find(|s| s.sql.contains("FROM \"evaluation\"")) + .expect("an evaluation query"); + let values = format!("{:?}", evaluations.values); + assert!(values.contains(&kept.to_string())); + assert!(!values.contains(&dropped.to_string())); + } +} diff --git a/backend/gradient-db/src/teams/mod.rs b/backend/gradient-db/src/teams/mod.rs index c0dd2f22b..5f7014fff 100644 --- a/backend/gradient-db/src/teams/mod.rs +++ b/backend/gradient-db/src/teams/mod.rs @@ -4,6 +4,7 @@ * SPDX-License-Identifier: AGPL-3.0-only */ +pub mod demand; pub mod evaluations; pub mod grants; pub mod invitations; diff --git a/backend/gradient-db/src/teams/workers.rs b/backend/gradient-db/src/teams/workers.rs index 8a86ba373..84915e420 100644 --- a/backend/gradient-db/src/teams/workers.rs +++ b/backend/gradient-db/src/teams/workers.rs @@ -44,18 +44,25 @@ pub async fn projects_granted_with_workers( .collect()) } -pub async fn team_workers_for_project( +pub async fn teams_granting_workers( db: &C, project: ProjectId, -) -> Result, DbErr> { - let teams: Vec = ETeamProject::find() +) -> Result, DbErr> { + Ok(ETeamProject::find() .filter(CTeamProject::Project.eq(project)) .filter(CTeamProject::IncludesWorkers.eq(true)) .all(db) .await? .into_iter() .map(|grant| grant.team) - .collect(); + .collect()) +} + +pub async fn team_workers_for_project( + db: &C, + project: ProjectId, +) -> Result, DbErr> { + let teams = teams_granting_workers(db, project).await?; if teams.is_empty() { return Ok(Vec::new()); } diff --git a/backend/gradient-db/src/webhooks.rs b/backend/gradient-db/src/webhooks.rs index 5e3d9bf7a..bcd9013b7 100644 --- a/backend/gradient-db/src/webhooks.rs +++ b/backend/gradient-db/src/webhooks.rs @@ -16,6 +16,7 @@ pub struct Owner { pub scope: WebhookScope, pub project: Option, pub cache: Option, + pub team: Option, } impl Owner { @@ -30,6 +31,10 @@ impl Owner { Some(c) => CWebhook::Cache.eq(c), None => CWebhook::Cache.is_null(), }) + .add(match self.team { + Some(t) => CWebhook::Team.eq(t), + None => CWebhook::Team.is_null(), + }) } } diff --git a/backend/gradient-entity/src/webhook.rs b/backend/gradient-entity/src/webhook.rs index a4e0fccba..c172511f8 100644 --- a/backend/gradient-entity/src/webhook.rs +++ b/backend/gradient-entity/src/webhook.rs @@ -9,7 +9,7 @@ use num_enum::{IntoPrimitive, TryFromPrimitive}; use sea_orm::entity::prelude::*; use serde::{Deserialize, Serialize}; -use crate::ids::{CacheId, ProjectId, UserId, WebhookId}; +use crate::ids::{CacheId, ProjectId, TeamId, UserId, WebhookId}; #[repr(i16)] #[derive( @@ -37,6 +37,8 @@ pub enum WebhookScope { Cache = 1, #[sea_orm(num_value = 2)] Instance = 2, + #[sea_orm(num_value = 3)] + Team = 3, } #[derive(Clone, Debug, Default, PartialEq, DeriveEntityModel, Deserialize, Serialize)] @@ -47,6 +49,7 @@ pub struct Model { pub scope: WebhookScope, pub project: Option, pub cache: Option, + pub team: Option, pub name: String, pub url: String, pub secret: String, diff --git a/backend/gradient-graph/src/cases/nar.rs b/backend/gradient-graph/src/cases/nar.rs index acc8e66ad..860485904 100644 --- a/backend/gradient-graph/src/cases/nar.rs +++ b/backend/gradient-graph/src/cases/nar.rs @@ -607,6 +607,67 @@ async fn n12_a_path_nobody_produces_is_stored_without_edges_and_marks_no_output( case.finish().await; } +#[tokio::test] +async fn n13_a_dependency_landing_in_the_same_flush_as_its_referrer_counts_for_it_once() { + let e1 = EvaluationId::now_v7(); + let case = Case::start_with_unstored_outputs( + shape() + .evaluation(e1) + .derivation(X) + .walked() + .entry_point("pkg.x") + .job() + .probed() + .drv_cached() + .output(X_OUT) + .depends_on(P) + .derivation(P) + .walked() + .job() + .probed() + .drv_cached() + .output(P_OUT) + .build_status(Completed) + .depends_on(D) + .depends_on(M) + .derivation(D) + .walked() + .job() + .probed() + .drv_cached() + .output(D_OUT) + .build_status(Completed) + .derivation(M) + .walked() + .job() + .probed() + .drv_cached() + .output(M_OUT) + .build_status(Completed), + 3, + ) + .await; + + let referrer = commit_of(P_OUT, "p", &[(P_OUT, "p"), (D_OUT, "d"), (M_OUT, "m")]); + let dependency = commit_of(D_OUT, "d", &[(D_OUT, "d")]); + let (p, d) = tokio::join!( + case.commit_nar_result(referrer), + case.commit_nar_result(dependency) + ); + p.expect("after 1: P_OUT is committed"); + d.expect("after 1: D_OUT is committed"); + + let d = case.closure_state(D).await; + assert_eq!(d, (true, 0), "after 1: sb(D) fetchable, missing 0"); + let p = case.closure_state(P).await; + assert_eq!(p, (false, 1), "after 1: sb(P) waits for M"); + let x = case.queue_state(X).await; + assert_eq!(x, (Created, 1), "after 1: sb(X)"); + assert_only_unstored_outputs(case.drift().await, 1, "step 1"); + let end = case.finish_with_drift().await; + assert_only_unstored_outputs(end, 1, "the end"); +} + #[tokio::test] async fn the_cached_path_row_is_read_under_a_row_lock() { let e1 = EvaluationId::now_v7(); diff --git a/backend/gradient-migration/src/lib.rs b/backend/gradient-migration/src/lib.rs index 7740774ad..0dc32dfbd 100644 --- a/backend/gradient-migration/src/lib.rs +++ b/backend/gradient-migration/src/lib.rs @@ -18,6 +18,8 @@ mod m20261015_000000_evaluation_metric_per_evaluation; mod m20261015_000001_output_cache_by_hash; mod m20261015_000002_pull_request_summary; mod m20261016_000000_metric_input_pnames; +mod m20261016_000001_metric_pname_without_version; +mod m20261016_000002_team_webhooks; pub struct UpgradeFloor { pub release: &'static str, @@ -53,6 +55,8 @@ impl MigratorTrait for Migrator { Box::new(m20261015_000001_output_cache_by_hash::Migration), Box::new(m20261015_000002_pull_request_summary::Migration), Box::new(m20261016_000000_metric_input_pnames::Migration), + Box::new(m20261016_000001_metric_pname_without_version::Migration), + Box::new(m20261016_000002_team_webhooks::Migration), ] } } diff --git a/backend/gradient-migration/src/m20261016_000001_metric_pname_without_version.rs b/backend/gradient-migration/src/m20261016_000001_metric_pname_without_version.rs new file mode 100644 index 000000000..21d1fe5dd --- /dev/null +++ b/backend/gradient-migration/src/m20261016_000001_metric_pname_without_version.rs @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use sea_orm_migration::prelude::*; +use sea_orm_migration::sea_orm::ConnectionTrait; + +const UP: &str = r#" +UPDATE public.derivation_metric m +SET pname = coalesce(nullif(d.pname, ''), regexp_replace(d.name, '-[0-9][^-]*$', '')) +FROM public.derivation d +WHERE d.id = m.derivation + AND m.pname IS DISTINCT FROM coalesce(nullif(d.pname, ''), regexp_replace(d.name, '-[0-9][^-]*$', '')); +"#; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager.get_connection().execute_unprepared(UP).await?; + Ok(()) + } + + async fn down(&self, _manager: &SchemaManager) -> Result<(), DbErr> { + Ok(()) + } +} diff --git a/backend/gradient-migration/src/m20261016_000002_team_webhooks.rs b/backend/gradient-migration/src/m20261016_000002_team_webhooks.rs new file mode 100644 index 000000000..51b14f2d5 --- /dev/null +++ b/backend/gradient-migration/src/m20261016_000002_team_webhooks.rs @@ -0,0 +1,56 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use sea_orm_migration::prelude::*; +use sea_orm_migration::sea_orm::ConnectionTrait; + +const UP: &str = r#" +ALTER TABLE public.webhook + ADD COLUMN team uuid NULL REFERENCES public.team(id) ON DELETE CASCADE; + +ALTER TABLE public.webhook DROP CONSTRAINT "chk-webhook-scope"; +ALTER TABLE public.webhook ADD CONSTRAINT "chk-webhook-scope" CHECK ( + ((scope = 0) AND (project IS NOT NULL) AND (cache IS NULL) AND (team IS NULL)) + OR ((scope = 1) AND (cache IS NOT NULL) AND (project IS NULL) AND (team IS NULL)) + OR ((scope = 2) AND (project IS NULL) AND (cache IS NULL) AND (team IS NULL)) + OR ((scope = 3) AND (team IS NOT NULL) AND (project IS NULL) AND (cache IS NULL)) +); + +DROP INDEX public."idx-webhook-name"; +CREATE UNIQUE INDEX "idx-webhook-name" ON public.webhook USING btree (scope, COALESCE(project, cache, team), name); + +CREATE INDEX "idx-webhook-team" ON public.webhook USING btree (team) WHERE (team IS NOT NULL); +"#; + +const DOWN: &str = r#" +DELETE FROM public.webhook WHERE scope = 3; + +DROP INDEX public."idx-webhook-team"; + +DROP INDEX public."idx-webhook-name"; +CREATE UNIQUE INDEX "idx-webhook-name" ON public.webhook USING btree (scope, COALESCE(project, cache), name); + +ALTER TABLE public.webhook DROP CONSTRAINT "chk-webhook-scope"; +ALTER TABLE public.webhook ADD CONSTRAINT "chk-webhook-scope" CHECK ((((scope = 0) AND (project IS NOT NULL) AND (cache IS NULL)) OR ((scope = 1) AND (cache IS NOT NULL) AND (project IS NULL)) OR ((scope = 2) AND (project IS NULL) AND (cache IS NULL)))); + +ALTER TABLE public.webhook DROP COLUMN team; +"#; + +#[derive(DeriveMigrationName)] +pub struct Migration; + +#[async_trait::async_trait] +impl MigrationTrait for Migration { + async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager.get_connection().execute_unprepared(UP).await?; + Ok(()) + } + + async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> { + manager.get_connection().execute_unprepared(DOWN).await?; + Ok(()) + } +} diff --git a/backend/gradient-pool/src/estimate/mod.rs b/backend/gradient-pool/src/estimate/mod.rs index 47207578e..aa3431f40 100644 --- a/backend/gradient-pool/src/estimate/mod.rs +++ b/backend/gradient-pool/src/estimate/mod.rs @@ -54,7 +54,7 @@ pub struct BuildEstimate { pub own_samples: u32, pub ram_floor: Vec<(u32, u64)>, pub max_cores: Option, - pub top_features: Vec<(String, f64)>, + pub similar_builds: Vec<(String, f64)>, } pub static NO_ESTIMATE: BuildEstimate = BuildEstimate { @@ -64,7 +64,7 @@ pub static NO_ESTIMATE: BuildEstimate = BuildEstimate { own_samples: 0, ram_floor: Vec::new(), max_cores: None, - top_features: Vec::new(), + similar_builds: Vec::new(), }; impl BuildEstimate { diff --git a/backend/gradient-pool/src/estimate/model.rs b/backend/gradient-pool/src/estimate/model.rs index 46922d2b4..21f8ab37a 100644 --- a/backend/gradient-pool/src/estimate/model.rs +++ b/backend/gradient-pool/src/estimate/model.rs @@ -10,9 +10,8 @@ use gradient_types::ids::DerivationId; use super::{BuildEstimate, LogCurve, log_cores}; use crate::score::weights::{ - BUILD_TIME_CORE_SLOPE_PRIOR, CORE_FEATURE_PENALTY, CORE_SLOPE_PENALTY, FEATURE_PENALTY, - FEW_SAMPLES_SPREAD, FIT_EPOCHS, MIN_SPREAD, OWN_SAMPLE_PRIOR_WEIGHT, RAM_CORE_SLOPE_PRIOR, - REFERENCE_CORES, TOP_FEATURES, + BUILD_TIME_CORE_SLOPE, FEW_SAMPLES_SPREAD, MIN_INPUT_OVERLAP, MIN_SPREAD, + OWN_SAMPLE_PRIOR_WEIGHT, RAM_CORE_SLOPE, REFERENCE_CORES, SIMILAR_BUILDS, }; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -25,11 +24,11 @@ enum Target { impl Target { const ALL: [Self; 3] = [Self::Ram, Self::BuildTime, Self::OutputNar]; - const fn core_slope_prior(self) -> Option { + const fn core_slope(self) -> f64 { match self { - Self::Ram => Some(RAM_CORE_SLOPE_PRIOR), - Self::BuildTime => Some(BUILD_TIME_CORE_SLOPE_PRIOR), - Self::OutputNar => None, + Self::Ram => RAM_CORE_SLOPE, + Self::BuildTime => BUILD_TIME_CORE_SLOPE, + Self::OutputNar => 0.0, } } } @@ -63,14 +62,6 @@ pub fn input_overlap(a: &[String], b: &[String]) -> f64 { a.intersection(&b).count() as f64 / union as f64 } -impl GroupKey { - fn features(&self) -> Vec { - std::iter::once(format!("pname:{}", self.pname)) - .chain(self.input_pnames.iter().map(|i| format!("input:{i}"))) - .collect() - } -} - impl Sample { pub fn group(&self) -> GroupKey { GroupKey { @@ -94,48 +85,114 @@ impl Sample { target == Target::Ram && self.lower_bound } - fn log_cores(&self, target: Target) -> f64 { - match (target.core_slope_prior(), self.cores) { - (Some(_), Some(cores)) => log_cores(cores), - _ => 0.0, - } + fn at_reference(&self, target: Target) -> Option { + let cores = self.cores.map_or(0.0, log_cores); + Some(self.value(target)? - target.core_slope() * cores) } } +/// The newest measured build of every pname on one architecture, and the pnames per input. #[derive(Clone, Debug, Default)] -struct LinearFit { - bias: f64, - core_slope: f64, - weights: HashMap, - core_weights: HashMap, - error: f64, +struct Builds { + newest: HashMap, + by_input: HashMap>, + overall: HashMap, } -impl LinearFit { - fn curve(&self, features: &[String]) -> LogCurve { - let sum = - |map: &HashMap| features.iter().filter_map(|f| map.get(f)).sum::(); - LogCurve { - at_reference: self.bias + sum(&self.weights), - per_log_core: self.core_slope + sum(&self.core_weights), - spread: self.error, +impl Builds { + fn new(samples: Vec<&Sample>) -> Self { + let mut newest = HashMap::new(); + for sample in samples.into_iter().filter(|s| !s.lower_bound) { + newest.insert(sample.pname.clone(), sample.clone()); + } + let mut by_input: HashMap> = HashMap::new(); + for sample in newest.values() { + for input in &sample.input_pnames { + by_input + .entry(input.clone()) + .or_default() + .push(sample.pname.clone()); + } + } + let overall = Target::ALL + .into_iter() + .filter_map(|target| { + let mut values: Vec = newest + .values() + .filter_map(|s| s.at_reference(target)) + .collect(); + Some((target, median_curve(&mut values, target)?)) + }) + .collect(); + Self { + newest, + by_input, + overall, + } + } + + fn similar(&self, key: &GroupKey) -> Vec<(f64, &Sample)> { + let mut shared: HashMap<&str, usize> = HashMap::new(); + let inputs: HashSet<&String> = key.input_pnames.iter().collect(); + for input in &inputs { + for pname in self.by_input.get(*input).into_iter().flatten() { + *shared.entry(pname).or_default() += 1; + } } + let mut similar: Vec<(f64, &Sample)> = shared + .into_iter() + .map(|(pname, both)| { + let sample = &self.newest[pname]; + let union = inputs.len() + sample.input_pnames.len() - both; + (both as f64 / union as f64, sample) + }) + .filter(|(overlap, _)| *overlap >= MIN_INPUT_OVERLAP) + .collect(); + similar.sort_by(|a, b| b.0.total_cmp(&a.0).then_with(|| a.1.pname.cmp(&b.1.pname))); + similar.truncate(SIMILAR_BUILDS); + similar } - fn top_features(&self, features: &[String]) -> Vec<(String, f64)> { - let mut weighed: Vec<(String, f64)> = features + fn curve(&self, similar: &[(f64, &Sample)], target: Target) -> Option { + let points: Vec<(f64, f64)> = similar .iter() - .filter_map(|f| self.weights.get(f).map(|w| (f.clone(), *w))) + .filter_map(|(overlap, s)| Some((overlap * overlap, s.at_reference(target)?))) .collect(); - weighed.sort_by(|a, b| b.1.abs().total_cmp(&a.1.abs())); - weighed.truncate(TOP_FEATURES); - weighed + let total: f64 = points.iter().map(|(weight, _)| weight).sum(); + if total <= 0.0 { + return self.overall.get(&target).copied(); + } + let mean = points.iter().map(|(w, v)| w * v).sum::() / total; + let variance = points + .iter() + .map(|(w, v)| w * (v - mean).powi(2)) + .sum::() + / total; + Some(LogCurve { + at_reference: mean, + per_log_core: target.core_slope(), + spread: variance.sqrt().max(FEW_SAMPLES_SPREAD), + }) } } +fn median_curve(values: &mut [f64], target: Target) -> Option { + if values.is_empty() { + return None; + } + values.sort_by(f64::total_cmp); + let median = values[values.len() / 2]; + let variance = values.iter().map(|v| (v - median).powi(2)).sum::() / values.len() as f64; + Some(LogCurve { + at_reference: median, + per_log_core: target.core_slope(), + spread: variance.sqrt().max(FEW_SAMPLES_SPREAD), + }) +} + #[derive(Clone, Debug, Default)] pub struct Model { - fits: HashMap<(String, Target), LinearFit>, + architectures: HashMap, } impl Model { @@ -147,15 +204,12 @@ impl Model { .or_default() .push(sample); } - let mut fits = HashMap::new(); - for (architecture, samples) in by_architecture { - for target in Target::ALL { - if let Some(fit) = Fitter::new(&samples, target).map(Fitter::fit) { - fits.insert((architecture.to_owned(), target), fit); - } - } + Self { + architectures: by_architecture + .into_iter() + .map(|(architecture, samples)| (architecture.to_owned(), Builds::new(samples))) + .collect(), } - Self { fits } } #[expect( @@ -164,9 +218,11 @@ impl Model { reason = "rounded and positive" )] pub fn estimate(&self, key: &GroupKey, own: &[(f64, &Sample)]) -> BuildEstimate { - let features = key.features(); - let fit = |target: Target| self.fits.get(&(key.architecture.clone(), target)); - let curve = |target: Target| Some(blend(fit(target)?.curve(&features), own, target)); + let Some(builds) = self.architectures.get(&key.architecture) else { + return BuildEstimate::default(); + }; + let similar = builds.similar(key); + let curve = |target: Target| Some(blend(builds.curve(&similar, target)?, own, target)); BuildEstimate { ram: curve(Target::Ram), build_time: curve(Target::BuildTime), @@ -180,15 +236,16 @@ impl Model { .filter_map(|(_, s)| Some((s.cores?, s.peak_ram_mb?))) .collect(), max_cores: None, - top_features: fit(Target::Ram) - .map(|fit| fit.top_features(&features)) - .unwrap_or_default(), + similar_builds: similar + .iter() + .map(|(overlap, s)| (s.pname.clone(), *overlap)) + .collect(), } } } -/// Own builds count at the cores they ran with. The core term only shapes the model part, so -/// it fades as own builds add up. +/// Own builds count at the cores they ran with. The core term only shapes the estimate from +/// similar builds, so it fades as own builds add up. fn blend(model: LogCurve, own: &[(f64, &Sample)], target: Target) -> LogCurve { let measured: Vec<(f64, f64)> = own .iter() @@ -218,204 +275,6 @@ fn blend(model: LogCurve, own: &[(f64, &Sample)], target: Target) -> LogCurve { } } -struct Row { - y: f64, - l: f64, - censored: bool, - features: Vec, -} - -/// Ridge regression in log space, solved by coordinate descent. A censored row (a lower bound) -/// only pulls the fit up, so it is active only while the prediction is below it. -struct Fitter { - names: Vec, - rows: Vec, - by_feature: Vec>, - prior: Option, - bias: f64, - slope: f64, - weights: Vec, - core_weights: Vec, - residual: Vec, - active: Vec, -} - -impl Fitter { - fn new(samples: &[&Sample], target: Target) -> Option { - let mut names: Vec = Vec::new(); - let mut index: HashMap = HashMap::new(); - let mut rows = Vec::new(); - for sample in samples { - let Some(y) = sample.value(target) else { - continue; - }; - let features = sample - .group() - .features() - .into_iter() - .map(|name| { - *index.entry(name.clone()).or_insert_with(|| { - names.push(name); - names.len() - 1 - }) - }) - .collect(); - rows.push(Row { - y, - l: sample.log_cores(target), - censored: sample.censored(target), - features, - }); - } - let measured: Vec = rows.iter().filter(|r| !r.censored).map(|r| r.y).collect(); - if measured.is_empty() { - return None; - } - - let mut by_feature = vec![Vec::new(); names.len()]; - for (i, row) in rows.iter().enumerate() { - for &f in &row.features { - by_feature[f].push(i); - } - } - let prior = target.core_slope_prior(); - let bias = measured.iter().sum::() / measured.len() as f64; - let slope = prior.unwrap_or(0.0); - let residual = rows.iter().map(|r| r.y - bias - slope * r.l).collect(); - Some(Self { - weights: vec![0.0; names.len()], - core_weights: vec![0.0; names.len()], - active: vec![true; rows.len()], - names, - rows, - by_feature, - prior, - bias, - slope, - residual, - }) - } - - fn fit(mut self) -> LinearFit { - for _ in 0..FIT_EPOCHS { - self.active = self - .rows - .iter() - .zip(&self.residual) - .map(|(row, e)| !row.censored || *e > 0.0) - .collect(); - self.fit_bias(); - self.fit_slope(); - for feature in 0..self.names.len() { - self.fit_weight(feature); - self.fit_core_weight(feature); - } - } - self.finish() - } - - fn active_rows(&self) -> impl Iterator + '_ { - (0..self.rows.len()).filter(|&i| self.active[i]) - } - - fn shift(&mut self, rows: &[usize], by: impl Fn(&Row) -> f64) { - for &i in rows { - self.residual[i] -= by(&self.rows[i]); - } - } - - fn fit_bias(&mut self) { - let (sum, count) = self - .active_rows() - .fold((0.0, 0.0), |(s, c), i| (s + self.residual[i], c + 1.0)); - if count == 0.0 { - return; - } - let step = sum / count; - self.bias += step; - let all: Vec = (0..self.rows.len()).collect(); - self.shift(&all, |_| step); - } - - fn fit_slope(&mut self) { - let Some(prior) = self.prior else { - return; - }; - let (num, den) = self.active_rows().fold((0.0, 0.0), |(n, d), i| { - let l = self.rows[i].l; - (n + l * (self.residual[i] + self.slope * l), d + l * l) - }); - let next = (num + CORE_SLOPE_PENALTY * prior) / (den + CORE_SLOPE_PENALTY); - let step = next - self.slope; - let all: Vec = (0..self.rows.len()).collect(); - self.shift(&all, |row| step * row.l); - self.slope = next; - } - - fn fit_weight(&mut self, feature: usize) { - let members = std::mem::take(&mut self.by_feature[feature]); - let current = self.weights[feature]; - let (num, den) = members - .iter() - .filter(|&&i| self.active[i]) - .fold((0.0, 0.0), |(n, d), &i| { - (n + self.residual[i] + current, d + 1.0) - }); - let next = num / (den + FEATURE_PENALTY); - self.shift(&members, |_| next - current); - self.weights[feature] = next; - self.by_feature[feature] = members; - } - - fn fit_core_weight(&mut self, feature: usize) { - if self.prior.is_none() { - return; - } - let members = std::mem::take(&mut self.by_feature[feature]); - let current = self.core_weights[feature]; - let (num, den) = - members - .iter() - .filter(|&&i| self.active[i]) - .fold((0.0, 0.0), |(n, d), &i| { - let l = self.rows[i].l; - (n + l * (self.residual[i] + current * l), d + l * l) - }); - let next = num / (den + CORE_FEATURE_PENALTY); - self.shift(&members, |row| (next - current) * row.l); - self.core_weights[feature] = next; - self.by_feature[feature] = members; - } - - fn finish(self) -> LinearFit { - let squared: Vec = self - .rows - .iter() - .zip(&self.residual) - .filter(|(row, _)| !row.censored) - .map(|(_, e)| e * e) - .collect(); - let error = (squared.iter().sum::() / squared.len() as f64) - .sqrt() - .max(MIN_SPREAD); - let named = |values: &[f64]| { - self.names - .iter() - .cloned() - .zip(values.iter().copied()) - .filter(|(_, v)| *v != 0.0) - .collect() - }; - LinearFit { - bias: self.bias, - core_slope: self.slope, - weights: named(&self.weights), - core_weights: named(&self.core_weights), - error, - } - } -} - #[cfg(test)] mod tests { use super::*; @@ -468,6 +327,26 @@ mod tests { samples } + fn python_packages() -> Vec { + vec![ + sample( + "llama-index-llms-openai", + &["python3", "hatchling", "llama-index-core", "openai"], + 16, + 126, + 7_000.0, + ), + sample( + "cuda-kernels", + &["cmake", "cudatoolkit", "ninja"], + 16, + 50_000, + 900_000.0, + ), + sample("unit-file", &["bash"], 16, 5, 100.0), + ] + } + #[test] fn a_shared_pname_with_other_inputs_gets_its_own_estimate() { let samples = torch_and_torch_bin(); @@ -489,18 +368,57 @@ mod tests { } #[test] - fn a_build_never_seen_is_estimated_from_its_inputs() { - let samples = torch_and_torch_bin(); + fn a_build_never_seen_takes_the_builds_with_the_most_similar_inputs() { + let model = Model::fit(&python_packages()); + let unseen = key( + "llama-index-embeddings-openai", + &[ + "python3", + "hatchling", + "llama-index-core", + "openai", + "numpy", + ], + ); + + let estimate = model.estimate(&unseen, &[]); + + let median = estimate.ram.unwrap().median(16); + assert!(median > 60.0 && median < 250.0, "{estimate:?}"); + assert_eq!(estimate.similar_builds[0].0, "llama-index-llms-openai"); + assert_eq!(estimate.own_samples, 0); + } + + #[test] + fn a_build_without_similar_inputs_gets_the_median_of_all_builds() { + let model = Model::fit(&python_packages()); + + let estimate = model.estimate(&key("never-built", &["also-unknown"]), &[]); + + let median = estimate.ram.unwrap().median(16); + assert!((median - 126.0).abs() < 1.0, "{estimate:?}"); + assert!(estimate.similar_builds.is_empty()); + assert!(estimate.build_ms(16).unwrap() > 0.0); + } + + #[test] + fn an_out_of_memory_kill_is_no_similar_build() { + let mut samples = python_packages(); + let mut killed = sample( + "llama-index-core", + &["python3", "hatchling", "openai"], + 16, + 90_000, + 1.0, + ); + killed.lower_bound = true; + killed.build_ms = None; + samples.push(killed); let model = Model::fit(&samples); - let bumped = key("torch", &["cmake", "cudatoolkit", "magma", "nccl"]); - let unknown = key("never-built", &["also-unknown"]); - let bumped = model.estimate(&bumped, &[]); - let unknown = model.estimate(&unknown, &[]); + let estimate = model.estimate(&key("other", &["python3", "hatchling", "openai"]), &[]); - assert_eq!(bumped.own_samples, 0); - assert!(bumped.ram_mb(64).unwrap() > 20_000, "{bumped:?}"); - assert!(unknown.ram_mb(16).unwrap() > 0 && unknown.build_ms(16).unwrap() > 0.0); + assert!(estimate.ram_mb(16).unwrap() < 1_000, "{estimate:?}"); } #[test] @@ -554,20 +472,14 @@ mod tests { } #[test] - fn a_build_without_own_builds_follows_the_core_term_of_the_model() { - let mut samples = Vec::new(); - for _ in 0..6 { - samples.push(sample("par", &["cc"], 8, 4_000, 20_000.0)); - samples.push(sample("par", &["cc"], 32, 16_000, 6_000.0)); - } - let model = Model::fit(&samples); - let estimate = model.estimate(&key("other", &["cc"]), &[]); + fn a_similar_build_counts_at_the_cores_it_ran_with() { + let model = Model::fit(&[sample("par", &["cc", "make"], 64, 16_000, 6_000.0)]); - assert!( - estimate.ram_mb(32).unwrap() > estimate.ram_mb(8).unwrap(), - "{estimate:?}" - ); - assert!(estimate.build_ms(32).unwrap() < estimate.build_ms(8).unwrap()); + let estimate = model.estimate(&key("other", &["cc", "make"]), &[]); + + assert!(estimate.ram.unwrap().median(16) < 16_000.0, "{estimate:?}"); + assert!(estimate.ram_mb(64).unwrap() > estimate.ram_mb(16).unwrap()); + assert!(estimate.build_ms(64).unwrap() < estimate.build_ms(16).unwrap()); } #[test] diff --git a/backend/gradient-pool/src/score/weights.rs b/backend/gradient-pool/src/score/weights.rs index a49cdcfc3..ebfcd7b6e 100644 --- a/backend/gradient-pool/src/score/weights.rs +++ b/backend/gradient-pool/src/score/weights.rs @@ -51,11 +51,7 @@ pub const MIN_INPUT_OVERLAP: f64 = 0.3; pub const OWN_SAMPLE_WINDOW: usize = 20; pub const FEW_SAMPLES_SPREAD: f64 = 0.3; pub const MIN_SPREAD: f64 = 0.05; -pub const RAM_CORE_SLOPE_PRIOR: f64 = 0.5; -pub const BUILD_TIME_CORE_SLOPE_PRIOR: f64 = -0.5; -pub const FEATURE_PENALTY: f64 = 1.0; -pub const CORE_FEATURE_PENALTY: f64 = 20.0; -pub const CORE_SLOPE_PENALTY: f64 = 50.0; -pub const FIT_EPOCHS: usize = 25; +pub const RAM_CORE_SLOPE: f64 = 0.5; +pub const BUILD_TIME_CORE_SLOPE: f64 = -0.5; +pub const SIMILAR_BUILDS: usize = 10; pub const RAM_FLOOR_MARGIN: f64 = 1.1; -pub const TOP_FEATURES: usize = 5; diff --git a/backend/gradient-scheduler/src/views.rs b/backend/gradient-scheduler/src/views.rs index 8af5e3dde..1f2fd6920 100644 --- a/backend/gradient-scheduler/src/views.rs +++ b/backend/gradient-scheduler/src/views.rs @@ -60,7 +60,7 @@ pub struct EstimateView { pub ram_mb: Option, pub build_ms: Option, pub output_nar_bytes: Option, - pub top_features: Vec<(String, f64)>, + pub similar_builds: Vec<(String, f64)>, } impl EstimateView { @@ -77,7 +77,7 @@ impl EstimateView { ram_mb: estimate.ram_mb(cores), build_ms: estimate.build_ms(cores).map(|ms| ms.round() as u64), output_nar_bytes: estimate.output_nar_bytes, - top_features: estimate.top_features.clone(), + similar_builds: estimate.similar_builds.clone(), } } } diff --git a/backend/gradient-scheduler/src/waiting_state.rs b/backend/gradient-scheduler/src/waiting_state.rs index 20b4c5982..486292246 100644 --- a/backend/gradient-scheduler/src/waiting_state.rs +++ b/backend/gradient-scheduler/src/waiting_state.rs @@ -17,6 +17,7 @@ use gradient_entity::build::BuildStatus; use gradient_entity::evaluation::{EvaluationKind, EvaluationStatus}; use gradient_graph::Graph; use gradient_pool::WorkerInfo; +use gradient_types::events::worker_demand::WorkerDemand; use gradient_types::*; use sea_orm::{ColumnTrait, ConnectionTrait, EntityTrait, QueryFilter, QuerySelect, Value}; use tracing::{error, info, warn}; @@ -91,11 +92,7 @@ impl WaitingRefresh<'_> { } let tasks = EvalTasks::load(&self.db.worker_db, &evals).await?; - let worker_caps: Vec<(Vec, Vec)> = workers - .iter() - .map(|w| (w.architectures.clone(), w.system_features.clone())) - .collect(); - let worker_caps = worker_caps.as_slice(); + let mut caps_by_project: HashMap, Vec> = HashMap::new(); let now = gradient_types::now(); let mut unbuildables = Vec::new(); let ids: Vec = evals.iter().map(|e| e.id).collect(); @@ -111,7 +108,12 @@ impl WaitingRefresh<'_> { .waiting_reason .as_ref() .and_then(WaitingReason::from_json); - let project_workers = ProjectWorkerCounts::of(workers, tasks.project_of(&eval)); + let project = tasks.project_of(&eval); + let project_workers = ProjectWorkerCounts::of(workers, project); + let worker_caps = caps_by_project + .entry(project) + .or_insert_with(|| project_worker_caps(workers, project)) + .as_slice(); if eval.status == EvaluationStatus::Waiting && reason.as_ref().is_some_and(|r| { @@ -188,6 +190,12 @@ impl WaitingRefresh<'_> { ); } + if let (Some(project), Some(gap)) = + (project, opened_gap(reason.as_ref(), new_reason.as_ref())) + { + self.announce_worker_gap(eval.id, project, gap).await; + } + persist_waiting_reason( &self.db.worker_db, eval.id, @@ -203,6 +211,53 @@ impl WaitingRefresh<'_> { Ok(unbuildables) } + + async fn announce_worker_gap( + &self, + evaluation: EvaluationId, + project: ProjectId, + reason: &WaitingReason, + ) { + let teams = + match gradient_db::teams::workers::teams_granting_workers(&self.db.worker_db, project) + .await + { + Ok(teams) => teams, + Err(e) => { + warn!(error = %e, %evaluation, "failed to read teams granting workers"); + return; + } + }; + for event in worker_demand_events(reason, evaluation, project, &teams) { + self.db.record(event).await; + } + } +} + +fn opened_gap<'a>( + current: Option<&WaitingReason>, + new: Option<&'a WaitingReason>, +) -> Option<&'a WaitingReason> { + let new = new?; + let gap = new.worker_gap()?; + (current.and_then(WaitingReason::worker_gap) != Some(gap)).then_some(new) +} + +fn worker_demand_events( + reason: &WaitingReason, + evaluation: EvaluationId, + project: ProjectId, + teams: &[TeamId], +) -> Vec { + teams + .iter() + .map(|&team| WorkerDemand { + team, + project, + evaluation, + reason: reason.clone(), + }) + .collect() } #[derive(Default)] @@ -628,6 +683,23 @@ fn pre_build_capability(kind: EvaluationKind, status: EvaluationStatus) -> Optio } } +type WorkerCaps = (Vec, Vec); + +fn eligible_for(worker: &WorkerInfo, project: Option) -> bool { + worker + .authorized_peers + .as_ref() + .is_none_or(|peers| project.is_some_and(|p| peers.contains(&p))) +} + +fn project_worker_caps(workers: &[WorkerInfo], project: Option) -> Vec { + workers + .iter() + .filter(|w| eligible_for(w, project)) + .map(|w| (w.architectures.clone(), w.system_features.clone())) + .collect() +} + #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] struct ProjectWorkerCounts { eval: usize, @@ -639,11 +711,7 @@ impl ProjectWorkerCounts { fn of(workers: &[WorkerInfo], project: Option) -> Self { workers .iter() - .filter(|w| { - w.authorized_peers - .as_ref() - .is_none_or(|peers| project.is_some_and(|p| peers.contains(&p))) - }) + .filter(|w| eligible_for(w, project)) .fold(Self::default(), |counts, w| Self { eval: counts.eval + usize::from(w.capabilities.eval), fetch: counts.fetch + usize::from(w.capabilities.fetch), @@ -793,6 +861,70 @@ mod tests { } } + fn capable_worker( + architectures: &[&str], + features: &[&str], + authorized_peers: Option>, + ) -> WorkerInfo { + WorkerInfo { + architectures: architectures.iter().map(|a| a.to_string()).collect(), + system_features: features.iter().map(|f| f.to_string()).collect(), + ..worker(true, authorized_peers) + } + } + + fn x86_gap(connected_workers: u32) -> WaitingReason { + WaitingReason::workers( + vec![UnmetRequirement { + architecture: "x86_64-linux".into(), + required_features: vec!["kvm".into()], + build_count: connected_workers + 1, + }], + connected_workers, + vec![], + ) + } + + #[test] + fn reason_change_records_one_event_per_granting_team() { + let new = x86_gap(0); + let gap = + opened_gap(Some(&WaitingReason::Draining), Some(&new)).expect("a new worker gap opens"); + let teams = [TeamId::now_v7(), TeamId::now_v7()]; + let events = worker_demand_events(gap, EvaluationId::now_v7(), ProjectId::now_v7(), &teams); + assert_eq!( + events.iter().map(|e| e.team).collect::>(), + teams.to_vec() + ); + assert_ne!( + gradient_types::events::EventKind::key(&events[0]), + gradient_types::events::EventKind::key(&events[1]) + ); + } + + #[test] + fn unchanged_reason_records_nothing() { + assert!(opened_gap(Some(&x86_gap(1)), Some(&x86_gap(3))).is_none()); + assert!(opened_gap(None, None).is_none()); + assert!(opened_gap(None, Some(&WaitingReason::graph_stuck(2))).is_none()); + } + + #[test] + fn worker_restricted_to_other_project_gives_no_caps() { + let (own, other) = (ProjectId::now_v7(), ProjectId::now_v7()); + let w = capable_worker(&["aarch64-darwin"], &[], Some(HashSet::from([other]))); + assert!(project_worker_caps(&[w], Some(own)).is_empty()); + } + + #[test] + fn open_worker_gives_caps_to_every_project() { + let w = capable_worker(&["x86_64-linux"], &["kvm"], None); + assert_eq!( + project_worker_caps(&[w], Some(ProjectId::now_v7())), + vec![(vec!["x86_64-linux".to_string()], vec!["kvm".to_string()])] + ); + } + #[test] fn a_queued_evaluation_waits_when_only_other_projects_have_eval_workers() { let (own, other) = (ProjectId::now_v7(), ProjectId::now_v7()); diff --git a/backend/gradient-types/src/events/audit.rs b/backend/gradient-types/src/events/audit.rs index 7fdaa7502..20774438c 100644 --- a/backend/gradient-types/src/events/audit.rs +++ b/backend/gradient-types/src/events/audit.rs @@ -120,6 +120,9 @@ actions! { TeamWorkerCreate => "team.worker.create", TeamWorkerUpdate => "team.worker.update", TeamWorkerDelete => "team.worker.delete", + TeamWebhookCreate => "team.webhook.create", + TeamWebhookUpdate => "team.webhook.update", + TeamWebhookDelete => "team.webhook.delete", } impl Action { diff --git a/backend/gradient-types/src/events/build.rs b/backend/gradient-types/src/events/build.rs index 1f9513b31..81dfb7858 100644 --- a/backend/gradient-types/src/events/build.rs +++ b/backend/gradient-types/src/events/build.rs @@ -94,7 +94,7 @@ impl EventKind for Reported { EventOwner { project: self.project, task: self.task, - cache: None, + ..EventOwner::default() } } diff --git a/backend/gradient-types/src/events/evaluation.rs b/backend/gradient-types/src/events/evaluation.rs index 0d16d62dc..e206001f2 100644 --- a/backend/gradient-types/src/events/evaluation.rs +++ b/backend/gradient-types/src/events/evaluation.rs @@ -126,7 +126,7 @@ impl EventKind for Reported { EventOwner { project: self.project, task: self.task, - cache: None, + ..EventOwner::default() } } diff --git a/backend/gradient-types/src/events/mod.rs b/backend/gradient-types/src/events/mod.rs index 0a3344ebe..26e4997be 100644 --- a/backend/gradient-types/src/events/mod.rs +++ b/backend/gradient-types/src/events/mod.rs @@ -20,12 +20,13 @@ pub mod names; pub mod proto; pub mod webhook; pub mod worker; +pub mod worker_demand; pub use bus::{EVENT_BUS_CAPACITY, EventBus, EventRx}; pub use envelope::Envelope; pub use filter::EventFilter; -use crate::ids::{CacheId, ProjectId, TaskId}; +use crate::ids::{CacheId, ProjectId, TaskId, TeamId}; use serde::{Deserialize, Serialize}; use std::borrow::Cow; @@ -34,6 +35,7 @@ pub struct EventOwner { pub project: Option, pub task: Option, pub cache: Option, + pub team: Option, } impl EventOwner { @@ -66,6 +68,13 @@ impl EventOwner { ..Self::default() } } + + pub fn team(team: TeamId) -> Self { + Self { + team: Some(team), + ..Self::default() + } + } } pub trait EventKind: Serialize { @@ -176,6 +185,7 @@ events! { WorkerDisconnected(worker::Disconnected), WorkerJobDispatched(worker::JobDispatched), WorkerQueueDepth(worker::QueueDepth), + WorkerDemand(worker_demand::WorkerDemand), Proto(proto::Message), CacheChanged(cache::Changed), CacheNarFetched(cache::NarFetched), diff --git a/backend/gradient-types/src/events/names.rs b/backend/gradient-types/src/events/names.rs index 785f06e04..9acd94c1c 100644 --- a/backend/gradient-types/src/events/names.rs +++ b/backend/gradient-types/src/events/names.rs @@ -35,6 +35,7 @@ pub const WORKER_CONNECTED: &str = "worker.connected"; pub const WORKER_DISCONNECTED: &str = "worker.disconnected"; pub const WORKER_JOB_DISPATCHED: &str = "worker.job_dispatched"; pub const WORKER_QUEUE_DEPTH: &str = "worker.queue_depth"; +pub const WORKER_DEMAND: &str = "worker.demand"; pub const GRAPH_NAR_COMMITTED: &str = "graph.nar_committed"; pub const GRAPH_TRANSITIONED: &str = "graph.transitioned"; diff --git a/backend/gradient-types/src/events/worker_demand.rs b/backend/gradient-types/src/events/worker_demand.rs new file mode 100644 index 000000000..eef0ecc07 --- /dev/null +++ b/backend/gradient-types/src/events/worker_demand.rs @@ -0,0 +1,63 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use super::{EventKind, EventOwner, names}; +use crate::ids::{EvaluationId, ProjectId, TeamId}; +use crate::waiting_reason::WaitingReason; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct WorkerDemand { + pub team: TeamId, + pub project: ProjectId, + pub evaluation: EvaluationId, + pub reason: WaitingReason, +} + +impl EventKind for WorkerDemand { + const NAME: &'static str = names::WORKER_DEMAND; + const DURABLE: bool = true; + + fn owner(&self) -> EventOwner { + EventOwner::team(self.team) + } + + fn key(&self) -> Option { + Some(format!( + "{}:{}:{}", + names::WORKER_DEMAND, + self.team, + self.reason.worker_gap().unwrap_or_default() + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::waiting_reason::EvalCapability; + + fn wanted(team: TeamId, evaluation: EvaluationId) -> WorkerDemand { + WorkerDemand { + team, + project: ProjectId::now_v7(), + evaluation, + reason: WaitingReason::eval_workers(EvalCapability::Eval, 0), + } + } + + #[test] + fn one_pending_delivery_per_team_and_gap() { + let team = TeamId::now_v7(); + let first = wanted(team, EvaluationId::now_v7()); + let second = wanted(team, EvaluationId::now_v7()); + assert_eq!(first.key(), second.key()); + assert_ne!( + first.key(), + wanted(TeamId::now_v7(), first.evaluation).key() + ); + } +} diff --git a/backend/gradient-types/src/waiting_reason.rs b/backend/gradient-types/src/waiting_reason.rs index 9cf267f4b..6c7c188a8 100644 --- a/backend/gradient-types/src/waiting_reason.rs +++ b/backend/gradient-types/src/waiting_reason.rs @@ -97,6 +97,29 @@ impl WaitingReason { pending_shared_builds, } } + + pub fn worker_gap(&self) -> Option { + match self { + Self::Workers { unmet, .. } if !unmet.is_empty() => { + let mut gaps: Vec = unmet + .iter() + .map(|u| { + let mut features = u.required_features.clone(); + features.sort(); + format!("{}[{}]", u.architecture, features.join(",")) + }) + .collect(); + gaps.sort(); + gaps.dedup(); + Some(gaps.join(";")) + } + Self::EvalWorkers { capability, .. } => Some(match capability { + EvalCapability::Fetch => "fetch".to_owned(), + EvalCapability::Eval => "eval".to_owned(), + }), + _ => None, + } + } } #[cfg(test)] @@ -113,4 +136,32 @@ mod tests { let decoded = WaitingReason::from_json(&reason.to_json()).expect("row decodes"); assert_eq!(decoded, reason); } + + #[test] + fn worker_gap_ignores_counts_and_order() { + let unmet = |arch: &str, features: &[&str], build_count| UnmetRequirement { + architecture: arch.to_owned(), + required_features: features.iter().map(|f| f.to_string()).collect(), + build_count, + }; + let a = WaitingReason::workers( + vec![ + unmet("x86_64-linux", &["kvm", "big"], 1), + unmet("aarch64-linux", &[], 2), + ], + 1, + vec![], + ); + let b = WaitingReason::workers( + vec![ + unmet("aarch64-linux", &[], 7), + unmet("x86_64-linux", &["big", "kvm"], 3), + ], + 4, + vec![], + ); + assert_eq!(a.worker_gap(), b.worker_gap()); + assert_eq!(WaitingReason::workers(vec![], 1, vec![]).worker_gap(), None); + assert_eq!(WaitingReason::Draining.worker_gap(), None); + } } diff --git a/backend/gradient-web/src/endpoints/projects/workers.rs b/backend/gradient-web/src/endpoints/projects/workers.rs index 7c84ecfd3..663eff0ed 100644 --- a/backend/gradient-web/src/endpoints/projects/workers.rs +++ b/backend/gradient-web/src/endpoints/projects/workers.rs @@ -169,6 +169,8 @@ pub struct WorkerLiveInfo { pub max_concurrent_builds: u32, pub assigned_job_count: usize, pub draining: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub ram_total_mb: Option, } impl From<&WorkerInfo> for WorkerLiveInfo { @@ -180,6 +182,7 @@ impl From<&WorkerInfo> for WorkerLiveInfo { max_concurrent_builds: w.max_concurrent_builds, assigned_job_count: w.assigned_job_count, draining: w.draining, + ram_total_mb: (w.ram_total_mb > 0).then_some(w.ram_total_mb), } } } diff --git a/backend/gradient-web/src/endpoints/teams/demand.rs b/backend/gradient-web/src/endpoints/teams/demand.rs new file mode 100644 index 000000000..72388c767 --- /dev/null +++ b/backend/gradient-web/src/endpoints/teams/demand.rs @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use crate::api_spec::TeamPath; +use crate::error::{WebError, WebResult}; +use crate::helpers::ok_json; +use crate::scope::{Admin, TeamScope}; +use axum::Json; +use axum::extract::Query; +use gradient_db::WebDb; +use gradient_db::teams::demand::TeamDemand; +use gradient_types::*; +use serde::Deserialize; +use uuid::Uuid; + +#[derive(Deserialize, Debug, Default, utoipa::IntoParams)] +#[into_params(parameter_in = Query)] +pub struct DemandQuery { + pub projects: Option, +} + +#[utoipa::path( + get, + path = "/teams/{team}/workers/demand", + operation_id = "getTeamWorkerDemand", + tag = "teams", + summary = "Read missing worker capabilities", + description = "Needs Admin in the team. Sums the waiting evaluations of projects granting the team its workers that no connected worker eligible for the project can take: builds per system and features, plus fetch and evaluation gaps. `projects` takes comma-separated project names or IDs and limits the sum to those granting projects.", + params(TeamPath, DemandQuery), + responses((status = OK, description = "OK", body = TeamDemand), WebError) +)] +pub async fn get_team_worker_demand( + TeamScope { team, db, .. }: TeamScope, + Query(query): Query, +) -> WebResult>> { + let projects = match query.projects.as_deref() { + Some(list) => Some(resolve_projects(&db.web_db, list).await?), + None => None, + }; + let demand = + gradient_db::teams::demand::team_demand(&db.web_db, team.id, projects.as_deref()).await?; + Ok(ok_json(demand)) +} + +async fn resolve_projects(db: &WebDb, list: &str) -> Result, WebError> { + let mut ids = Vec::new(); + for entry in list.split(',').map(str::trim).filter(|e| !e.is_empty()) { + if let Ok(id) = entry.parse::() { + ids.push(ProjectId::new(id)); + } else if let Some(project) = gradient_db::projects::by_name(db, entry).await? { + ids.push(project.id); + } + } + Ok(ids) +} diff --git a/backend/gradient-web/src/endpoints/teams/mod.rs b/backend/gradient-web/src/endpoints/teams/mod.rs index 7fe607718..9d0f8c376 100644 --- a/backend/gradient-web/src/endpoints/teams/mod.rs +++ b/backend/gradient-web/src/endpoints/teams/mod.rs @@ -4,6 +4,7 @@ * SPDX-License-Identifier: AGPL-3.0-only */ +pub mod demand; pub mod grants; pub mod invitations; pub mod management; diff --git a/backend/gradient-web/src/endpoints/teams/workers.rs b/backend/gradient-web/src/endpoints/teams/workers.rs index 0d8879829..72589feec 100644 --- a/backend/gradient-web/src/endpoints/teams/workers.rs +++ b/backend/gradient-web/src/endpoints/teams/workers.rs @@ -8,8 +8,8 @@ use super::refusal::TeamRefusal; use crate::api_spec::TeamPath; use crate::audit::Audit; use crate::endpoints::projects::workers::{ - PatchWorkerRequest, RegisterWorkerRequest, WorkerConnection, encrypt_for_dialing, live_workers, - patch_edits_managed_fields, worker_connection, + PatchWorkerRequest, RegisterWorkerRequest, WorkerConnection, WorkerLiveInfo, + encrypt_for_dialing, live_workers, patch_edits_managed_fields, worker_connection, }; use crate::error::{WebError, WebResult}; use crate::helpers::{OptionExt, ok_json}; @@ -21,13 +21,14 @@ use axum::{Extension, Json}; use chrono::NaiveDateTime; use gradient_core::ServerState; use gradient_db::WebDb; +use gradient_pool::WorkerInfo; use gradient_scheduler::Scheduler; use gradient_types::events::EventOwner; use gradient_types::events::audit::Action; use gradient_types::*; use sea_orm::ActiveValue::Set; use serde::Serialize; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use uuid::Uuid; @@ -46,10 +47,16 @@ pub struct TeamWorkerEntry { pub enable_build: bool, #[serde(flatten)] pub connection: WorkerConnection, + #[serde(skip_serializing_if = "Option::is_none")] + pub live: Option, } impl TeamWorkerEntry { - fn new(worker: MTeamWorker, connection: WorkerConnection) -> Self { + fn new( + worker: MTeamWorker, + connection: WorkerConnection, + live: Option, + ) -> Self { Self { worker_id: worker.worker_id, display_name: worker.display_name, @@ -62,10 +69,15 @@ impl TeamWorkerEntry { enable_eval: worker.enable_eval, enable_build: worker.enable_build, connection, + live, } } } +fn team_worker_live(live: &HashMap, worker_id: &str) -> Option { + live.get(worker_id).map(WorkerLiveInfo::from) +} + #[derive(Serialize, utoipa::ToSchema)] pub struct RegisteredTeamWorker { pub team: TeamId, @@ -135,7 +147,8 @@ pub async fn get_team_workers( let mut entries = Vec::with_capacity(rows.len()); for worker in rows { let connection = worker_connection(&live_workers, failures, &worker.worker_id); - entries.push(TeamWorkerEntry::new(worker, connection)); + let live = team_worker_live(&live_workers, &worker.worker_id); + entries.push(TeamWorkerEntry::new(worker, connection, live)); } Ok(ok_json(entries)) @@ -209,7 +222,7 @@ pub async fn post_team_worker( audit .write( Action::TeamWorkerCreate, - EventOwner::default(), + EventOwner::team(team.id), serde_json::json!({ "team_id": team.id.to_string(), "worker_id": worker_id }), ) .await; @@ -279,7 +292,7 @@ pub async fn patch_team_worker( audit .write( Action::TeamWorkerUpdate, - EventOwner::default(), + EventOwner::team(team.id), serde_json::json!({ "team_id": team.id.to_string(), "worker_id": worker_id }), ) .await; @@ -319,10 +332,52 @@ pub async fn delete_team_worker( audit .write( Action::TeamWorkerDelete, - EventOwner::default(), + EventOwner::team(team.id), serde_json::json!({ "team_id": team.id.to_string(), "worker_id": worker_id }), ) .await; Ok(ok_json(format!("worker '{}' unregistered", worker_id))) } + +#[cfg(test)] +mod tests { + use super::*; + + fn connected(id: &str) -> WorkerInfo { + WorkerInfo { + id: id.into(), + capabilities: Default::default(), + architectures: vec!["x86_64-linux".into()], + system_features: vec!["kvm".into()], + max_concurrent_builds: 4, + assigned_job_count: 0, + assigned_build_count: 0, + draining: false, + authorized_peers: None, + cpu_usage_pct: None, + ram_free_mb: None, + ram_total_mb: 16384, + disk_speed_mbps: None, + upload_speed_mbps: None, + download_speed_mbps: None, + } + } + + #[test] + fn connected_team_worker_lists_live_capabilities() { + let live = live_workers(vec![connected("w1")]); + let info = team_worker_live(&live, "w1").expect("connected worker has live details"); + let json = serde_json::to_value(&info).expect("serializes"); + assert_eq!(json["architectures"], serde_json::json!(["x86_64-linux"])); + assert_eq!(json["system_features"], serde_json::json!(["kvm"])); + assert_eq!(json["max_concurrent_builds"], 4); + assert_eq!(json["ram_total_mb"], 16384); + } + + #[test] + fn disconnected_team_worker_has_no_live() { + let live = live_workers(vec![connected("w1")]); + assert!(team_worker_live(&live, "w2").is_none()); + } +} diff --git a/backend/gradient-web/src/endpoints/webhooks/handlers.rs b/backend/gradient-web/src/endpoints/webhooks/handlers.rs index f33781d1b..ea3a4885d 100644 --- a/backend/gradient-web/src/endpoints/webhooks/handlers.rs +++ b/backend/gradient-web/src/endpoints/webhooks/handlers.rs @@ -143,6 +143,7 @@ pub async fn create_webhook( scope: owner.scope, project: owner.project, cache: owner.cache, + team: owner.team, name: body.name, url: body.url, secret: encrypted, diff --git a/backend/gradient-web/src/endpoints/webhooks/mod.rs b/backend/gradient-web/src/endpoints/webhooks/mod.rs index 1bbffece5..05076c451 100644 --- a/backend/gradient-web/src/endpoints/webhooks/mod.rs +++ b/backend/gradient-web/src/endpoints/webhooks/mod.rs @@ -12,7 +12,7 @@ mod scope; pub use responses::{ActionDeliveryItem, DeliveryDetail, DeliveryListItem, WebhookDeliveryItem}; pub use scope::{Verb, WebhookOwner}; -use crate::api_spec::{CachePath, ProjectPath, operations_mut}; +use crate::api_spec::{CachePath, ProjectPath, TeamPath, operations_mut}; use gradient_core::ServerState; use std::sync::Arc; use utoipa::IntoParams; @@ -25,6 +25,7 @@ use utoipa_axum::routes; pub enum WebhookMount { Project, Cache, + Team, Instance, } @@ -33,6 +34,7 @@ impl WebhookMount { match self { Self::Project => "project", Self::Cache => "cache", + Self::Team => "team", Self::Instance => "instance", } } @@ -41,6 +43,7 @@ impl WebhookMount { match self { Self::Project => "`manageWebhooks` on the project", Self::Cache => "`manageCacheWebhooks` on the cache", + Self::Team => "Admin role in the team", Self::Instance => "superuser", } } @@ -50,6 +53,7 @@ impl WebhookMount { match self { Self::Project => ProjectPath::into_params(path), Self::Cache => CachePath::into_params(path), + Self::Team => TeamPath::into_params(path), Self::Instance => Vec::new(), } } diff --git a/backend/gradient-web/src/endpoints/webhooks/scope.rs b/backend/gradient-web/src/endpoints/webhooks/scope.rs index 249c12daa..a7818cee4 100644 --- a/backend/gradient-web/src/endpoints/webhooks/scope.rs +++ b/backend/gradient-web/src/endpoints/webhooks/scope.rs @@ -5,7 +5,9 @@ */ use super::refusal::WebhookRefusal; -use crate::access::{CacheAccess, Caller, ProjectAccess, load_cache, load_project}; +use crate::access::{ + CacheAccess, Caller, ProjectAccess, TeamAccess, load_cache, load_project, load_team, +}; use crate::authorization::MaybeApiKey; use crate::error::{WebError, require_superuser}; use crate::permissions::{CachePermission, Permission}; @@ -24,6 +26,7 @@ pub struct WebhookOwner { pub scope: WebhookScope, pub project: Option, pub cache: Option, + pub team: Option, pub user: MUser, pub params: PathParams, } @@ -41,6 +44,7 @@ impl WebhookOwner { project: self.project, task: None, cache: self.cache, + team: self.team, } } @@ -55,6 +59,9 @@ impl WebhookOwner { (WebhookScope::Instance, Verb::Create) => Action::InstanceWebhookCreate, (WebhookScope::Instance, Verb::Update) => Action::InstanceWebhookUpdate, (WebhookScope::Instance, Verb::Delete) => Action::InstanceWebhookDelete, + (WebhookScope::Team, Verb::Create) => Action::TeamWebhookCreate, + (WebhookScope::Team, Verb::Update) => Action::TeamWebhookUpdate, + (WebhookScope::Team, Verb::Delete) => Action::TeamWebhookDelete, } } @@ -69,6 +76,7 @@ impl From<&WebhookOwner> for gradient_db::webhooks::Owner { scope: owner.scope, project: owner.project, cache: owner.cache, + team: owner.team, } } } @@ -108,6 +116,7 @@ impl FromRequestParts> for WebhookOwner { scope: WebhookScope::Project, project: Some(project.id), cache: None, + team: None, user, params, }); @@ -128,6 +137,27 @@ impl FromRequestParts> for WebhookOwner { scope: WebhookScope::Cache, project: None, cache: Some(cache.id), + team: None, + user, + params, + }); + } + if params.contains("team") { + let (team, _) = load_team( + &db.web_db, + &user, + api_key.as_ref(), + params.param("team")?, + TeamAccess::Admin { + reject_managed: false, + }, + ) + .await?; + return Ok(Self { + scope: WebhookScope::Team, + project: None, + cache: None, + team: Some(team.id), user, params, }); @@ -137,6 +167,7 @@ impl FromRequestParts> for WebhookOwner { scope: WebhookScope::Instance, project: None, cache: None, + team: None, user, params, }) diff --git a/backend/gradient-web/src/routes.rs b/backend/gradient-web/src/routes.rs index 7a0fc3ed6..9a068939a 100644 --- a/backend/gradient-web/src/routes.rs +++ b/backend/gradient-web/src/routes.rs @@ -105,6 +105,7 @@ pub(crate) fn auth_api(limits: &BodyLimits) -> OpenApiRouter> { teams::workers::get_team_workers, teams::workers::post_team_worker )) + .routes(routes!(teams::demand::get_team_worker_demand)) .routes(routes!( teams::workers::patch_team_worker, teams::workers::delete_team_worker @@ -294,6 +295,10 @@ pub(crate) fn auth_api(limits: &BodyLimits) -> OpenApiRouter> { "/caches/{cache}/webhooks", webhooks::router(WebhookMount::Cache), ) + .nest( + "/teams/{team}/webhooks", + webhooks::router(WebhookMount::Team), + ) .nest("/admin", admin::admin_router()), ), BEARER, diff --git a/backend/gradient-web/tests/team_demand.rs b/backend/gradient-web/tests/team_demand.rs new file mode 100644 index 000000000..575f5bcdc --- /dev/null +++ b/backend/gradient-web/tests/team_demand.rs @@ -0,0 +1,63 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use axum::http::StatusCode; +use gradient_entity::evaluation::EvaluationStatus; +use gradient_entity::team_user::TeamRole; +use gradient_entity::{evaluation, team_project}; +use gradient_fixtures::ids::project_id; +use gradient_fixtures::rows::{team, team_member}; +use gradient_test_support::web::{make_test_server, signed_in}; +use gradient_types::{EvalCapability, WaitingReason}; +use serde_json::Value; + +const DEMAND_URL: &str = "/api/v1/teams/platform/workers/demand"; + +#[tokio::test] +async fn team_member_cannot_read_demand() { + let (db, auth) = signed_in(); + let db = db + .append_query_results([vec![team()]]) + .append_query_results([vec![team_member(TeamRole::Member)]]); + let server = make_test_server(db.into_connection()); + + let res = server + .get(DEMAND_URL) + .add_header("authorization", auth) + .await; + + res.assert_status(StatusCode::FORBIDDEN); +} + +#[tokio::test] +async fn admin_reads_demand() { + let (db, auth) = signed_in(); + let db = db + .append_query_results([vec![team()]]) + .append_query_results([vec![team_member(TeamRole::Admin)]]) + .append_query_results([vec![team_project::Model { + project: project_id(), + includes_workers: true, + ..Default::default() + }]]) + .append_query_results([vec![evaluation::Model { + status: EvaluationStatus::Waiting, + waiting_reason: Some(WaitingReason::eval_workers(EvalCapability::Eval, 0).to_json()), + ..Default::default() + }]]); + let server = make_test_server(db.into_connection()); + + let res = server + .get(DEMAND_URL) + .add_header("authorization", auth) + .await; + + res.assert_status_ok(); + let body: Value = res.json(); + assert_eq!(body["message"]["eval"], 1); + assert_eq!(body["message"]["fetch"], 0); + assert_eq!(body["message"]["builds"], serde_json::json!([])); +} diff --git a/backend/gradient-web/tests/webhooks.rs b/backend/gradient-web/tests/webhooks.rs index dc1fa812f..0e53d2d25 100644 --- a/backend/gradient-web/tests/webhooks.rs +++ b/backend/gradient-web/tests/webhooks.rs @@ -9,11 +9,12 @@ reason = "test scaffolding: a fixture helper that cannot build its value should fail the test loudly" )] +use gradient_entity::team_user::TeamRole; use gradient_entity::webhook::WebhookScope; use gradient_entity::{ids::*, role, webhook}; -use gradient_fixtures::ids::{project_id, user_id}; -use gradient_fixtures::mock::mock_db; -use gradient_fixtures::rows::{project, project_member}; +use gradient_fixtures::ids::{project_id, team_id, user_id}; +use gradient_fixtures::mock::{all_statements, mock_db}; +use gradient_fixtures::rows::{project, project_member, team, team_member}; use gradient_fixtures::test_date; use gradient_test_support::state::TestState; use gradient_test_support::web::{bearer, make_test_server}; @@ -24,6 +25,7 @@ use serde_json::{Value, json}; use uuid::Uuid; const PROJECT_URL: &str = "/api/v1/projects/test-project/webhooks"; +const TEAM_URL: &str = "/api/v1/teams/platform/webhooks"; fn webhook_id() -> WebhookId { WebhookId::new(Uuid::from_u128(0xb1)) @@ -35,6 +37,7 @@ fn webhook_row() -> webhook::Model { scope: WebhookScope::Project, project: Some(project_id()), cache: None, + team: None, name: "ci".into(), url: "https://example.com/hook".into(), secret: "ENCRYPTED".into(), @@ -47,6 +50,22 @@ fn webhook_row() -> webhook::Model { } } +fn team_webhook_row() -> webhook::Model { + webhook::Model { + scope: WebhookScope::Team, + project: None, + team: Some(team_id()), + events: json!(["worker.demand"]), + ..webhook_row() + } +} + +fn as_team(session_id: SessionId, role: TeamRole) -> MockDatabase { + with_auth(mock_db(), session_id) + .append_query_results([vec![team()]]) + .append_query_results([vec![team_member(role)]]) +} + fn role_row(id: RoleId, permission: i64) -> role::Model { role::Model { id, @@ -180,3 +199,68 @@ async fn instance_webhooks_require_a_superuser() { res.assert_status(axum::http::StatusCode::FORBIDDEN); } + +#[tokio::test] +async fn team_admin_creates_webhook_and_gets_secret_once() { + let session_id = SessionId::now_v7(); + let db = as_team(session_id, TeamRole::Admin) + .append_query_results([Vec::::new()]) + .append_query_results([vec![team_webhook_row()]]); + let server = TestState::new() + .web_db(db.into_connection()) + .cli(|cli| cli.secrets.crypt_file = temp_crypt_secret_file()) + .server(); + + let res = server + .post(TEAM_URL) + .add_header("authorization", bearer(session_id)) + .json(&json!({ "name": "ci", "url": "https://example.com/hook", "events": ["worker.demand"] })) + .await; + + res.assert_status_ok(); + let body: Value = res.json(); + assert!( + body["message"]["secret"] + .as_str() + .unwrap() + .starts_with("whs_") + ); + assert_eq!(body["message"]["webhook"]["scope"], "team"); +} + +#[tokio::test] +async fn team_member_cannot_create_webhook() { + let session_id = SessionId::now_v7(); + let server = make_test_server(as_team(session_id, TeamRole::Member).into_connection()); + + let res = server + .post(TEAM_URL) + .add_header("authorization", bearer(session_id)) + .json(&json!({ "name": "ci", "url": "https://example.com/hook", "events": [] })) + .await; + + res.assert_status(axum::http::StatusCode::FORBIDDEN); +} + +#[tokio::test] +async fn team_webhook_list_hides_other_team() { + let session_id = SessionId::now_v7(); + let db = as_team(session_id, TeamRole::Admin) + .append_query_results([vec![team_webhook_row()]]) + .into_connection(); + let server = make_test_server(db.clone()); + + let res = server + .get(TEAM_URL) + .add_header("authorization", bearer(session_id)) + .await; + + res.assert_status_ok(); + drop(server); + let list = all_statements(db.into_transaction_log()) + .into_iter() + .find(|s| s.sql.contains("FROM \"webhook\"")) + .unwrap(); + assert!(list.sql.contains("\"webhook\".\"team\" = $")); + assert!(format!("{:?}", list.values).contains(&team_id().to_string())); +} diff --git a/backend/gradient-worker/options.json b/backend/gradient-worker/options.json index 6242bbb24..aade28092 100644 --- a/backend/gradient-worker/options.json +++ b/backend/gradient-worker/options.json @@ -142,6 +142,26 @@ "default": null, "values": [] }, + { + "group": "", + "field": "idle_timeout", + "flag": "idle-timeout", + "env": "GRADIENT_WORKER_IDLE_TIMEOUT", + "type": "Option", + "description": "Seconds without an assigned job after which the worker drains and exits. `null` keeps it\nrunning.", + "default": null, + "values": [] + }, + { + "group": "", + "field": "max_lifetime", + "flag": "max-lifetime", + "env": "GRADIENT_WORKER_MAX_LIFETIME", + "type": "Option", + "description": "Seconds after start after which the worker drains, finishes its running jobs and exits.\n`null` keeps it running.", + "default": null, + "values": [] + }, { "group": "", "field": "eval_subprocess", diff --git a/backend/gradient-worker/src/config/mod.rs b/backend/gradient-worker/src/config/mod.rs index 71d4aa7a1..d3f86a9d8 100644 --- a/backend/gradient-worker/src/config/mod.rs +++ b/backend/gradient-worker/src/config/mod.rs @@ -9,6 +9,7 @@ mod options; pub use options::*; use std::num::{NonZeroU32, NonZeroUsize}; +use std::time::Duration; use gradient_eval::pool::MemoryGuard; use gradient_wire::types::GradientCapabilities; @@ -118,6 +119,13 @@ impl WorkerConfig { .map_or_else(crate::metrics::cpu_core_score, NonZeroU32::get) } + pub fn exit_limits(&self) -> crate::exit_clock::ExitLimits { + crate::exit_clock::ExitLimits { + idle_timeout: self.idle_timeout.map(Duration::from_secs), + max_lifetime: self.max_lifetime.map(Duration::from_secs), + } + } + pub fn capabilities(&self) -> GradientCapabilities { GradientCapabilities { core: false, diff --git a/backend/gradient-worker/src/config/options.rs b/backend/gradient-worker/src/config/options.rs index 4e3143624..4db33e091 100644 --- a/backend/gradient-worker/src/config/options.rs +++ b/backend/gradient-worker/src/config/options.rs @@ -94,6 +94,14 @@ pub struct WorkerConfig { /// Set it to `null` to accept every server; avoid that in a production setup. pub accepted_server_tokens_file: Option, + /// Seconds without an assigned job after which the worker drains and exits. `null` keeps it + /// running. + pub idle_timeout: Option, + + /// Seconds after start after which the worker drains, finishes its running jobs and exits. + /// `null` keeps it running. + pub max_lifetime: Option, + /// Re-exec as a Nix evaluator subprocess (internal - do not set manually). #[arg(hide = true)] pub eval_subprocess: bool, @@ -136,6 +144,8 @@ impl Default for WorkerConfig { listen_addr: "127.0.0.1".to_owned(), port: 3100, accepted_server_tokens_file: None, + idle_timeout: None, + max_lifetime: None, eval_subprocess: false, eval_driver: None, capabilities: CapabilitiesArgs::default(), diff --git a/backend/gradient-worker/src/exit_clock.rs b/backend/gradient-worker/src/exit_clock.rs new file mode 100644 index 000000000..0b11f85fa --- /dev/null +++ b/backend/gradient-worker/src/exit_clock.rs @@ -0,0 +1,208 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use std::collections::HashMap; +use std::sync::LazyLock; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +use gradient_util::sync::Mutex; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ExitLimits { + pub idle_timeout: Option, + pub max_lifetime: Option, +} + +impl ExitLimits { + pub fn is_set(&self) -> bool { + self.idle_timeout.is_some() || self.max_lifetime.is_some() + } +} + +// One clock per process: every session reports its busy slots here, so an exit limit judges the +// whole worker and never one session alone. +pub struct ExitClock { + started: Instant, + state: Mutex, +} + +struct ClockState { + busy: HashMap, + idle_since: Instant, + exiting: bool, +} + +static CLOCK: LazyLock = LazyLock::new(|| ExitClock::new(Instant::now())); +static NEXT_SESSION: AtomicU64 = AtomicU64::new(0); + +pub fn clock() -> &'static ExitClock { + &CLOCK +} + +pub fn new_session() -> u64 { + NEXT_SESSION.fetch_add(1, Ordering::Relaxed) +} + +impl ExitClock { + fn new(started: Instant) -> Self { + Self { + started, + state: Mutex::new(ClockState { + busy: HashMap::new(), + idle_since: started, + exiting: false, + }), + } + } + + pub fn report(&self, session: u64, busy: usize, now: Instant) { + let mut state = self.state.lock(); + let was_idle = state.busy.is_empty(); + if busy == 0 { + state.busy.remove(&session); + } else { + state.busy.insert(session, busy); + } + let is_idle = state.busy.is_empty(); + if is_idle && !was_idle { + state.idle_since = now; + } + } + + pub fn is_idle(&self) -> bool { + self.state.lock().busy.is_empty() + } + + pub fn exit_due(&self, limits: &ExitLimits, now: Instant) -> bool { + let mut state = self.state.lock(); + let idle_since = state.busy.is_empty().then_some(state.idle_since); + state.exiting = state.exiting || exit_due(idle_since, self.started, now, limits); + state.exiting + } + + pub fn deadline(&self, limits: &ExitLimits) -> Option { + let state = self.state.lock(); + let lifetime = limits.max_lifetime.map(|l| self.started + l); + let idle = limits + .idle_timeout + .filter(|_| state.busy.is_empty()) + .map(|t| state.idle_since + t); + lifetime.into_iter().chain(idle).min() + } +} + +fn exit_due( + idle_since: Option, + started: Instant, + now: Instant, + limits: &ExitLimits, +) -> bool { + let elapsed = |since: Instant| now.saturating_duration_since(since); + let idle_over = limits + .idle_timeout + .zip(idle_since) + .is_some_and(|(timeout, since)| elapsed(since) >= timeout); + let lifetime_over = limits + .max_lifetime + .is_some_and(|lifetime| elapsed(started) >= lifetime); + idle_over || lifetime_over +} + +#[cfg(test)] +mod tests { + use super::*; + + fn limits(idle: Option, lifetime: Option) -> ExitLimits { + ExitLimits { + idle_timeout: idle.map(Duration::from_secs), + max_lifetime: lifetime.map(Duration::from_secs), + } + } + + fn secs(n: u64) -> Duration { + Duration::from_secs(n) + } + + #[test] + fn idle_past_timeout_exits() { + let started = Instant::now(); + let clock = ExitClock::new(started); + assert!(clock.exit_due(&limits(Some(60), None), started + secs(120))); + } + + #[test] + fn busy_never_exits_on_idle() { + let started = Instant::now(); + let clock = ExitClock::new(started); + clock.report(1, 1, started); + assert!(!clock.exit_due(&limits(Some(60), None), started + secs(3600))); + } + + #[test] + fn a_held_cluster_slot_counts_as_busy() { + let started = Instant::now(); + let clock = ExitClock::new(started); + let held_slots = 1; + clock.report(1, held_slots, started + secs(10)); + assert!(!clock.exit_due(&limits(Some(60), None), started + secs(600))); + } + + #[test] + fn job_done_resets_idle() { + let started = Instant::now(); + let clock = ExitClock::new(started); + clock.report(1, 1, started + secs(10)); + clock.report(1, 0, started + secs(100)); + let idle = limits(Some(60), None); + assert!(!clock.exit_due(&idle, started + secs(130))); + assert!(clock.exit_due(&idle, started + secs(160))); + } + + #[test] + fn one_idle_session_does_not_end_a_busy_one() { + let started = Instant::now(); + let clock = ExitClock::new(started); + clock.report(1, 0, started); + clock.report(2, 1, started); + assert!(!clock.is_idle()); + assert!(!clock.exit_due(&limits(Some(60), None), started + secs(600))); + } + + #[test] + fn lifetime_reached_exits_even_when_busy() { + let started = Instant::now(); + let lifetime = limits(None, Some(600)); + let clock = ExitClock::new(started); + clock.report(1, 2, started); + assert!(!clock.exit_due(&lifetime, started + secs(599))); + assert!(clock.exit_due(&lifetime, started + secs(600))); + assert!(!clock.is_idle()); + } + + #[test] + fn none_limits_never_exit() { + let started = Instant::now(); + let clock = ExitClock::new(started); + assert!(!clock.exit_due(&limits(None, None), started + secs(1_000_000))); + assert_eq!(clock.deadline(&limits(None, None)), None); + } + + #[test] + fn disconnected_worker_exits_at_the_earliest_deadline() { + let started = Instant::now(); + let clock = ExitClock::new(started); + assert_eq!( + clock.deadline(&limits(Some(60), Some(600))), + Some(started + secs(60)) + ); + clock.report(1, 1, started); + assert_eq!( + clock.deadline(&limits(Some(60), Some(600))), + Some(started + secs(600)) + ); + } +} diff --git a/backend/gradient-worker/src/main.rs b/backend/gradient-worker/src/main.rs index a59d62794..1b4168863 100644 --- a/backend/gradient-worker/src/main.rs +++ b/backend/gradient-worker/src/main.rs @@ -9,6 +9,7 @@ mod connection_state; mod count; mod error; mod executor; +mod exit_clock; mod listener; mod metrics; mod nix; @@ -72,6 +73,7 @@ fn main() -> Result<()> { } config.memory_guard()?; + exit_clock::clock(); // The crypto provider, before the first TLS handshake, for rustls 0.23 (#232). gradient_util::http::init_crypto_provider(); @@ -147,9 +149,11 @@ fn main() -> Result<()> { sd_notify::notify(&[sd_notify::NotifyState::Ready])?; let mut backoff = INITIAL_BACKOFF; + let limits = config.exit_limits(); let initial = tokio::select! { _ = stop.cancelled() => None, + () = exit_deadline(limits) => None, w = async { loop { match Worker::connect(config.clone()).await { @@ -168,7 +172,7 @@ fn main() -> Result<()> { } => Some(w), }; let Some(mut worker) = initial else { - info!("shutdown requested during initial connect"); + info!("shutdown requested or exit limit reached during initial connect"); wait_for_inbound_sessions(&sessions).await; return Ok(()); }; @@ -180,7 +184,7 @@ fn main() -> Result<()> { let (disconnected, outcome) = worker.run(stop.clone()).await; if stop.is_cancelled() { - info!("worker stopped; tearing down"); + info!("worker stopped or reached an exit limit; tearing down"); drop(disconnected); wait_for_inbound_sessions(&sessions).await; executor_handle.shutdown().await; @@ -212,6 +216,7 @@ fn main() -> Result<()> { let reconnected = tokio::select! { _ = stop.cancelled() => None, + () = exit_deadline(limits) => None, conn = retry_reconnect( async || disconnected.reconnect().await, |delay| tokio::time::sleep(delay), @@ -227,7 +232,7 @@ fn main() -> Result<()> { backoff_after_session(backoff, session_end, INITIAL_BACKOFF, MAX_BACKOFF); } None => { - info!("shutdown requested during reconnect"); + info!("shutdown requested or exit limit reached during reconnect"); wait_for_inbound_sessions(&sessions).await; executor_handle.shutdown().await; return Ok(()); @@ -249,6 +254,13 @@ fn install_signal_handler(stop: CancellationToken) { }); } +async fn exit_deadline(limits: exit_clock::ExitLimits) { + match exit_clock::clock().deadline(&limits) { + Some(deadline) => tokio::time::sleep_until(deadline.into()).await, + None => std::future::pending().await, + } +} + async fn wait_for_inbound_sessions(sessions: &TaskTracker) { sessions.close(); if !sessions.is_empty() { diff --git a/backend/gradient-worker/src/worker/message_loop.rs b/backend/gradient-worker/src/worker/message_loop.rs index d7d8cae82..00f90bc3e 100644 --- a/backend/gradient-worker/src/worker/message_loop.rs +++ b/backend/gradient-worker/src/worker/message_loop.rs @@ -7,7 +7,7 @@ use gradient_util::sync::Mutex; use std::collections::HashMap; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; use anyhow::{Context as _, Result}; use gradient_wire::messages::{ @@ -22,6 +22,7 @@ use crate::config::WorkerConfig; use crate::count::clamp_to_u32; use crate::executor::timeline::{JobTimeline, TimelineSnapshot}; use crate::executor::{AbortSignal, JobExecutor}; +use crate::exit_clock; use crate::proto::credentials::CredentialStore; use crate::proto::job::JobUpdater; use crate::proto::scorer::JobScorer; @@ -56,6 +57,9 @@ pub(super) async fn run_message_loop( let mut stop_deadline = None; + let limits = state.config.exit_limits(); + let mut exit_check = tokio::time::interval(Duration::from_secs(1)); + info!("entering dispatch loop"); loop { @@ -83,6 +87,23 @@ pub(super) async fn run_message_loop( } } + _ = exit_check.tick(), if limits.is_set() => { + state.report_busy(); + if exit_clock::clock().exit_due(&limits, Instant::now()) { + if !state.exiting { + info!(?limits, "exit limit reached; draining before exit"); + state.exiting = true; + state.begin_drain().await?; + state.report_busy(); + } + if exit_clock::clock().is_idle() { + info!("exit limit reached and every job has finished; exiting"); + stop.cancel(); + break; + } + } + } + _ = heartbeat.tick() => { state.on_heartbeat().await?; } @@ -105,8 +126,10 @@ pub(super) async fn run_message_loop( result?; } } + state.report_busy(); } + exit_clock::clock().report(state.session, 0, Instant::now()); Ok(LoopEnd { draining: state.draining && !state.stopping, refused: state.refused, @@ -229,6 +252,8 @@ pub(super) struct MessageLoopState { draining: bool, stopping: bool, refused: bool, + exiting: bool, + session: u64, } impl MessageLoopState { @@ -277,9 +302,17 @@ impl MessageLoopState { draining: false, stopping: false, refused: false, + exiting: false, + session: exit_clock::new_session(), } } + fn report_busy(&self) { + let held = self.holds.held(&JobKind::Flake) + self.holds.held(&JobKind::Build); + let busy = self.jobs.len() + usize::try_from(held).unwrap_or(usize::MAX); + exit_clock::clock().report(self.session, busy, Instant::now()); + } + async fn begin_drain(&mut self) -> Result<()> { self.draining = true; let held = self.holds.release_all(); diff --git a/docs/generated/configuration-options.md b/docs/generated/configuration-options.md index f219702d1..161d3a067 100644 --- a/docs/generated/configuration-options.md +++ b/docs/generated/configuration-options.md @@ -309,7 +309,9 @@ | `worker.environmentVariables` | lazy attribute set of (null or bool or int or string or path) | `{ }` | - | Environment variables of the Gradient worker. The [configuration reference](https://wavelens.github.io/gradient/reference/configuration) lists every variable and its default. Set a variable to `null` to unset it. | | `worker.gcrootsDir` | string | `"/nix/var/nix/gcroots/gradient"` | `GRADIENT_WORKER_GCROOTS_DIR` | Directory for the indirect GC roots pinning each running build's inputs and outputs against a concurrent `nix-collect-garbage`. An empty string disables pinning. | | `worker.id` | null or string | `null` | `GRADIENT_WORKER_ID` | Worker UUID. `null` generates one on first start and stores it in `/worker-id`. Set it when the UUID must be known in advance, for example to register the worker in `state.workers`. | +| `worker.idleTimeout` | null or (unsigned int) | `null` | `GRADIENT_WORKER_IDLE_TIMEOUT` | Seconds without an assigned job after which the worker drains and exits. `null` keeps it running. | | `worker.listenAddr` | string | `"127.0.0.1"` | `GRADIENT_WORKER_LISTEN_ADDR` | IP address the worker listens on for incoming server connections. | +| `worker.maxLifetime` | null or (unsigned int) | `null` | `GRADIENT_WORKER_MAX_LIFETIME` | Seconds after start after which the worker drains, finishes its running jobs and exits. `null` keeps it running. | | `worker.peersFile` | null or path | `null` | `GRADIENT_WORKER_PEERS_FILE` | File of peer tokens for challenge-response authentication with the server, with `peer_id:token` per line and `#` starting a comment. ``` : *: ``` Use peer ID `*` to match any UUID in the server's challenge. Tokens are 48 byte random secrets, for example from `openssl rand -base64 48`. Register them via the API through `POST /api/v1/projects/{project}/workers`. Set it to `null` for open mode without authentication; avoid that in a production setup. | | `worker.port` | port | `3100` | `GRADIENT_WORKER_PORT` | Port the worker listens on for incoming server connections. | | `worker.serverUrl` | null or string | `null` | `GRADIENT_WORKER_SERVER_URL` | WebSocket URL of the Gradient server's `/proto` endpoint. `null` keeps the worker waiting for the server to connect. | diff --git a/docs/gradient-api.yaml b/docs/gradient-api.yaml index 6d697ca3c..19827b6f9 100644 --- a/docs/gradient-api.yaml +++ b/docs/gradient-api.yaml @@ -10364,6 +10364,357 @@ paths: $ref: '#/components/responses/Error' security: - bearerAuth: [] + /api/v1/teams/{team}/webhooks: + get: + tags: + - webhooks + summary: List team webhooks + description: Requires Admin role in the team. + operationId: listTeamWebhooks + parameters: + - name: team + in: path + required: true + schema: + type: string + responses: + '200': + description: Webhooks + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + type: array + items: + $ref: '#/components/schemas/WebhookResponse' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] + post: + tags: + - webhooks + summary: Create a team webhook + description: |- + Requires Admin role in the team. The response carries the signing `secret` once; it is + never returned again (use `rotate-secret` to replace it). + operationId: createTeamWebhook + parameters: + - name: team + in: path + required: true + schema: + type: string + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/CreateWebhookRequest' + required: true + responses: + '200': + description: Created webhook and its secret + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + $ref: '#/components/schemas/CreateWebhookResponse' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] + /api/v1/teams/{team}/webhooks/{id}: + get: + tags: + - webhooks + summary: Read a team webhook + operationId: readTeamWebhook + parameters: + - name: id + in: path + required: true + schema: + $ref: '#/components/schemas/WebhookId' + - name: team + in: path + required: true + schema: + type: string + responses: + '200': + description: Webhook + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + $ref: '#/components/schemas/WebhookResponse' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] + delete: + tags: + - webhooks + summary: Delete a team webhook + operationId: deleteTeamWebhook + parameters: + - name: id + in: path + required: true + schema: + $ref: '#/components/schemas/WebhookId' + - name: team + in: path + required: true + schema: + type: string + responses: + '200': + description: Deleted + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + $ref: '#/components/schemas/DeletedResponse' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] + patch: + tags: + - webhooks + summary: Update a team webhook + operationId: updateTeamWebhook + parameters: + - name: id + in: path + required: true + schema: + $ref: '#/components/schemas/WebhookId' + - name: team + in: path + required: true + schema: + type: string + requestBody: + content: + application/json: + schema: + $ref: '#/components/schemas/UpdateWebhookRequest' + required: true + responses: + '200': + description: Updated webhook + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + $ref: '#/components/schemas/WebhookResponse' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] + /api/v1/teams/{team}/webhooks/{id}/deliveries: + get: + tags: + - webhooks + summary: List a team webhook's deliveries + operationId: listTeamWebhookDeliveries + parameters: + - name: limit + in: query + required: false + schema: + type: integer + format: int64 + minimum: 0 + - name: offset + in: query + required: false + schema: + type: integer + format: int64 + minimum: 0 + - name: id + in: path + required: true + schema: + $ref: '#/components/schemas/WebhookId' + - name: team + in: path + required: true + schema: + type: string + responses: + '200': + description: Deliveries, newest first + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + type: array + items: + $ref: '#/components/schemas/DeliveryListItem' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] + /api/v1/teams/{team}/webhooks/{id}/deliveries/{delivery_id}: + get: + tags: + - webhooks + summary: Read one team webhook delivery with bodies + operationId: getTeamWebhookDelivery + parameters: + - name: id + in: path + required: true + schema: + $ref: '#/components/schemas/WebhookId' + - name: delivery_id + in: path + required: true + schema: + $ref: '#/components/schemas/WebhookDeliveryId' + - name: team + in: path + required: true + schema: + type: string + responses: + '200': + description: Delivery + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + $ref: '#/components/schemas/DeliveryDetail_WebhookDeliveryId' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] + /api/v1/teams/{team}/webhooks/{id}/rotate-secret: + post: + tags: + - webhooks + summary: Replace a team webhook's signing secret + operationId: rotateTeamWebhookSecret + parameters: + - name: id + in: path + required: true + schema: + $ref: '#/components/schemas/WebhookId' + - name: team + in: path + required: true + schema: + type: string + responses: + '200': + description: The new secret, shown once + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + $ref: '#/components/schemas/RotatedSecret' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] + /api/v1/teams/{team}/webhooks/{id}/test: + post: + tags: + - webhooks + summary: Send a `webhook.ping` to a team webhook + description: Delivers immediately (no retry) and returns the recorded delivery. + operationId: testTeamWebhook + parameters: + - name: id + in: path + required: true + schema: + $ref: '#/components/schemas/WebhookId' + - name: team + in: path + required: true + schema: + type: string + responses: + '200': + description: Recorded delivery + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + $ref: '#/components/schemas/DeliveryListItem_WebhookDeliveryId' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] /api/v1/teams/{team}/workers: get: tags: @@ -10435,6 +10786,43 @@ paths: $ref: '#/components/responses/Error' security: - bearerAuth: [] + /api/v1/teams/{team}/workers/demand: + get: + tags: + - teams + summary: Read missing worker capabilities + description: 'Needs Admin in the team. Sums the waiting evaluations of projects granting the team its workers that no connected worker eligible for the project can take: builds per system and features, plus fetch and evaluation gaps. `projects` takes comma-separated project names or IDs and limits the sum to those granting projects.' + operationId: getTeamWorkerDemand + parameters: + - name: team + in: path + required: true + schema: + type: string + - name: projects + in: query + required: false + schema: + type: string + responses: + '200': + description: OK + content: + application/json: + schema: + type: object + required: + - error + - message + properties: + error: + type: boolean + message: + $ref: '#/components/schemas/TeamDemand' + default: + $ref: '#/components/responses/Error' + security: + - bearerAuth: [] /api/v1/teams/{team}/workers/{worker_id}: delete: tags: @@ -18280,6 +18668,25 @@ components: type: - string - 'null' + TeamDemand: + type: object + required: + - builds + - fetch + - eval + properties: + builds: + type: array + items: + $ref: '#/components/schemas/UnmetRequirement' + fetch: + type: integer + format: int32 + minimum: 0 + eval: + type: integer + format: int32 + minimum: 0 TeamEvaluation: type: object required: @@ -18521,6 +18928,10 @@ components: type: boolean enable_build: type: boolean + live: + oneOf: + - $ref: '#/components/schemas/WorkerLiveInfo' + - type: 'null' Tier: type: string enum: @@ -19112,6 +19523,7 @@ components: - project - cache - instance + - team WorkerConnection: type: object required: @@ -19204,6 +19616,12 @@ components: minimum: 0 draining: type: boolean + ram_total_mb: + type: + - integer + - 'null' + format: int64 + minimum: 0 WorkerLoad: type: object required: diff --git a/docs/src/contributors/scheduler/scoring.md b/docs/src/contributors/scheduler/scoring.md index 7210ffe89..5208a026e 100644 --- a/docs/src/contributors/scheduler/scoring.md +++ b/docs/src/contributors/scheduler/scoring.md @@ -90,19 +90,21 @@ A released job drops out of these counts together with its slot. ```mermaid flowchart LR rows["derivation_metric rows"] --> store["sample store per instance"] - store --> fit["fit per architecture"] + store --> similar["similar builds per architecture"] store --> own["own past builds"] - fit --> estimate["BuildEstimate"] + similar --> estimate["BuildEstimate"] own --> estimate ``` - **Samples:** build-assign passes read the rows added since the last read. Rows count once per id, also across instances. - **Lower bounds:** out-of-memory rows and rows without a build time set a lower bound for RAM only. -- **Features:** the pname, each input pname, and each of those times `log(cores / 16)`. +- **Similar builds:** the 10 pnames whose newest build shares the most input pnames, at least 30 %. The weight of each is the square of that overlap. Lower-bound rows stay out. +- **No similar build:** the median over the newest build of every pname. +- **Cores:** values move to 16 cores by `log(cores / 16)` times `+0.5` for RAM and `-0.5` for build time. - **Own builds:** the last 20 rows of the pname and architecture with at least 30 % input overlap, weighted by that overlap. torch and torch-bin share no inputs and stay apart. - **Targets:** peak RAM, build time at `REFERENCE_CORE_SCORE` without concurrent builds, and output NAR size, all in log space. -- **Fit:** `BuildHistoryRefit` can refit a ridge regression every 300 s on each instance. The shared core term can start at `+0.5` for RAM and `-0.5` for build time. -- **Blend:** the estimate is `(W * own + 0.1 * model) / (W + 0.1)`, with `W` the summed weight of the own samples. Own samples keep the core count of their run. Never-measured derivations still get the model's estimate. +- **Refresh:** `BuildHistoryRefit` can rebuild the similar builds every 300 s on each instance. +- **Blend:** the estimate is `(W * own + 0.1 * similar) / (W + 0.1)`, with `W` the summed weight of the own samples. Own samples keep the core count of their run. Never-measured derivations still get the estimate of their similar builds. - **RAM reserve:** `BuildEstimate::ram_mb(cores)` can give the 90th percentile at that core count. Core counts at or above an own out-of-memory kill keep at least 1.1 times its killed peak. - **No samples:** architectures without rows give no estimate. Their builds count as `RamNeed::Unmeasured`. diff --git a/docs/src/reference/events.md b/docs/src/reference/events.md index 25bb65ceb..ce8fb31ad 100644 --- a/docs/src/reference/events.md +++ b/docs/src/reference/events.md @@ -29,6 +29,7 @@ All three payloads use one JSON structure: websocket messages, webhook bodies, a | `evaluation.` | yes | `evaluation.queued`, `evaluation.building`, `evaluation.failed`, `evaluation.approval_granted` | | `project.*`, `task.*`, `cache.*` | yes | `project.create`, `task.star`, `task.action.update`, `cache.nar.upload`, `cache.member.create` | | `gc.*` | yes | `gc.swept`, `gc.deep_finished` | +| `worker.demand` | yes, team webhooks only | waiting evaluations that need a system, feature, fetch or eval worker missing from the connected workers. Delivered once per team granting its workers to the project | | account activity | yes, instance webhooks only | `login.success`, `api_key.create`, `session.revoke` | | `build.status_changed`, `build.progress`, `evaluation.progress`, `evaluation.activity` | no | each build job transition, build transfer progress, fetch and eval progress | | `graph.*` | no | `graph.recorded`, `graph.demoted`, `graph.collected` | @@ -40,12 +41,13 @@ Every `build.*` event with a per-evaluation `build_id` also includes `derivation ## Webhooks -Manage webhooks under **Webhooks** in the project settings or on the cache page. Superusers manage instance webhooks under **Job Board -> System Health -> Instance Webhooks**. [Task actions](../guides/actions.md) receive the same envelope for single tasks. +Manage webhooks under **Webhooks** in the project settings, on the cache page or in the team settings. Superusers manage instance webhooks under **Job Board -> System Health -> Instance Webhooks**. [Task actions](../guides/actions.md) receive the same envelope for single tasks. | Scope | API | Who may manage | Receives | |---|---|---|---| | Project | `/api/v1/projects/{project}/webhooks` | `manageWebhooks` | events of the project, its tasks and evaluations | | Cache | `/api/v1/caches/{cache}/webhooks` | `manageCacheWebhooks` | events of the cache | +| Team | `/api/v1/teams/{team}/webhooks` | team Admin | team events (`worker.demand`, `team.worker.*`, `team.webhook.*`), never events of granting projects | | Instance | `/api/v1/admin/webhooks` | superuser | every durable event | - `events` is a list of globs (`build.*`, `task.star`). With an empty list, the webhook receives everything in scope. diff --git a/flake.nix b/flake.nix index 868e74b69..05d47238b 100644 --- a/flake.nix +++ b/flake.nix @@ -170,6 +170,7 @@ nix = final: prev: { gradient-nix = ((prev.nixVersions.nixComponents_2_35.appendPatches ( + # We would love someone getting this upstream! map (patch: ./nix/patches/nix + "/${patch}") (builtins.attrNames (builtins.readDir ./nix/patches/nix)) )).overrideScope (finalScope: prevScope: { withAWS = false; })).nix-everything; }; diff --git a/frontend/src/app/app.routes.ts b/frontend/src/app/app.routes.ts index ab0a5cf8c..ee2286c80 100644 --- a/frontend/src/app/app.routes.ts +++ b/frontend/src/app/app.routes.ts @@ -472,6 +472,13 @@ export const routes: Routes = [ loadComponent: () => import('./features/teams/team-workers/team-workers.component').then((m) => m.TeamWorkersComponent), }, + { + path: 'team/:team/webhooks', + title: 'Team Webhooks', + data: { webhookScope: 'team' }, + loadComponent: () => + import('./features/webhooks/webhooks.component').then((m) => m.WebhooksComponent), + }, // Integrations { diff --git a/frontend/src/app/core/models/api.generated.ts b/frontend/src/app/core/models/api.generated.ts index b0a564c58..ad35a0594 100644 --- a/frontend/src/app/core/models/api.generated.ts +++ b/frontend/src/app/core/models/api.generated.ts @@ -2110,6 +2110,12 @@ export interface TeamCacheGrant { role?: string | null; } +export interface TeamDemand { + builds: UnmetRequirement[]; + fetch: number; + eval: number; +} + export interface TeamEvaluation { id: EvaluationId; project: string; @@ -2200,6 +2206,7 @@ export type TeamWorkerEntry = WorkerConnection & { enable_fetch: boolean; enable_eval: boolean; enable_build: boolean; + live?: WorkerLiveInfo | null; }; export type Tier = "starred_active" | "active" | "starred" | "team" | "member"; @@ -2405,7 +2412,7 @@ export interface WebhookResponse { updated_at: string; } -export type WebhookScope = "project" | "cache" | "instance"; +export type WebhookScope = "project" | "cache" | "instance" | "team"; export interface WorkerConnection { connected: boolean; @@ -2435,6 +2442,7 @@ export interface WorkerLiveInfo { max_concurrent_builds: number; assigned_job_count: number; draining: boolean; + ram_total_mb?: number | null; } export interface WorkerLoad { diff --git a/frontend/src/app/core/models/webhook.model.ts b/frontend/src/app/core/models/webhook.model.ts index 92b3489db..020fce9f5 100644 --- a/frontend/src/app/core/models/webhook.model.ts +++ b/frontend/src/app/core/models/webhook.model.ts @@ -7,6 +7,7 @@ export type WebhookScopeRef = | { kind: 'project'; name: string } | { kind: 'cache'; name: string } + | { kind: 'team'; name: string } | { kind: 'instance' }; export function webhooksBase(scope: WebhookScopeRef): string { @@ -15,6 +16,8 @@ export function webhooksBase(scope: WebhookScopeRef): string { return `projects/${scope.name}/webhooks`; case 'cache': return `caches/${scope.name}/webhooks`; + case 'team': + return `teams/${scope.name}/webhooks`; case 'instance': return 'admin/webhooks'; } diff --git a/frontend/src/app/core/services/board.service.ts b/frontend/src/app/core/services/board.service.ts index 76a91205f..2eec93209 100644 --- a/frontend/src/app/core/services/board.service.ts +++ b/frontend/src/app/core/services/board.service.ts @@ -56,7 +56,7 @@ export interface DerivationRef { build_id: string; drv_path: string; pname: stri export interface JobEstimateView { own_samples: number; cores: number; ram_mb: number | null; build_ms: number | null; - output_nar_bytes: number | null; top_features: [string, number][]; + output_nar_bytes: number | null; similar_builds: [string, number][]; } export interface JobContextView { diff --git a/frontend/src/app/features/board/job-detail/job-detail.component.spec.ts b/frontend/src/app/features/board/job-detail/job-detail.component.spec.ts index 2259676de..698c23b44 100644 --- a/frontend/src/app/features/board/job-detail/job-detail.component.spec.ts +++ b/frontend/src/app/features/board/job-detail/job-detail.component.spec.ts @@ -76,7 +76,7 @@ const DETAIL: AssignedJobDetail = { ram_mb: 777, build_ms: 4200, output_nar_bytes: null, - top_features: [['input:cudatoolkit', 2.1], ['input:python3-wheel', -0.5]], + similar_builds: [['python3.14-torchaudio', 0.75], ['magma', 0.4]], }, derivations: [{ build_id: 'b1', drv_path: '/nix/store/xxx-foo', pname: 'foo' }], }, @@ -272,14 +272,14 @@ describe('BoardJobDetailComponent - structured context panels', () => { expect(job.textContent).toContain('Architecture'); }); - it('shows the chosen cores and the inputs that weigh most in the estimate', () => { + it('shows the chosen cores and the builds the estimate is taken from', () => { const el = setup().nativeElement as HTMLElement; const text = el.textContent ?? ''; expect(text).toContain('Estimate'); expect(text).toContain('43'); - expect(text).toContain('input:cudatoolkit'); - expect(text).toContain('+2.10'); - expect(text).toContain('-0.50'); + expect(text).toContain('Like python3.14-torchaudio'); + expect(text).toContain('75% same inputs'); + expect(text).toContain('40% same inputs'); }); it('hides the architecture for eval jobs', () => { diff --git a/frontend/src/app/features/board/job-detail/job-detail.component.ts b/frontend/src/app/features/board/job-detail/job-detail.component.ts index e0fabba90..e13ff17bb 100644 --- a/frontend/src/app/features/board/job-detail/job-detail.component.ts +++ b/frontend/src/app/features/board/job-detail/job-detail.component.ts @@ -203,8 +203,8 @@ interface RuleRow { Build time{{ formatDuration(e.build_ms) }} Output size{{ formatBytes(e.output_nar_bytes) }} Own builds{{ e.own_samples }} - @for (f of e.top_features; track f[0]) { - {{ f[0] }}{{ f[1] >= 0 ? '+' : '' }}{{ f[1] | number: '1.2-2' }} + @for (b of e.similar_builds; track b[0]) { + Like {{ b[0] }}{{ b[1] | percent }} same inputs } diff --git a/frontend/src/app/features/teams/team-settings/team-settings.component.html b/frontend/src/app/features/teams/team-settings/team-settings.component.html index 1a7815244..cc57e4cff 100644 --- a/frontend/src/app/features/teams/team-settings/team-settings.component.html +++ b/frontend/src/app/features/teams/team-settings/team-settings.component.html @@ -54,6 +54,17 @@ [appManagedDisable]="access()"> + @if (team.role === 'admin' || isSuperuser()) { + + + + Webhooks + Signed HTTP requests for this team's events + + + + } + @if (access().canEdit) { diff --git a/frontend/src/app/features/teams/team-settings/team-settings.component.ts b/frontend/src/app/features/teams/team-settings/team-settings.component.ts index a48e3371f..df4e1ebc2 100644 --- a/frontend/src/app/features/teams/team-settings/team-settings.component.ts +++ b/frontend/src/app/features/teams/team-settings/team-settings.component.ts @@ -22,6 +22,7 @@ import { MessageBannerComponent, PageLayoutComponent, RowComponent, + RowListComponent, SelectComponent, SettingsSectionComponent, } from '@gradient/ui/ui'; @@ -52,6 +53,7 @@ interface SettingsForm { MessageBannerComponent, PageLayoutComponent, RowComponent, + RowListComponent, SelectComponent, SettingsSectionComponent, ], diff --git a/frontend/src/app/features/webhooks/webhooks.component.spec.ts b/frontend/src/app/features/webhooks/webhooks.component.spec.ts index 6940c92a1..d6214432f 100644 --- a/frontend/src/app/features/webhooks/webhooks.component.spec.ts +++ b/frontend/src/app/features/webhooks/webhooks.component.spec.ts @@ -11,8 +11,10 @@ import { WebhooksComponent } from './webhooks.component'; import { WebhooksService } from '@core/services/webhooks.service'; import { CachesService } from '@core/services/caches.service'; import { ProjectsService } from '@core/services/projects.service'; +import { TeamsService } from '@core/services/teams.service'; import { EventsService } from '@core/services/events.service'; import type { WebhookResponse } from '@core/models'; +import { webhooksBase } from '@core/models/webhook.model'; const hook: WebhookResponse = { id: 'w1', @@ -44,6 +46,7 @@ function setup(data: Record, params: Record) { { provide: WebhooksService, useValue: service }, { provide: ProjectsService, useValue: { getProject: () => of({ display_name: 'Acme' }) } }, { provide: CachesService, useValue: { getCache: () => of({ display_name: 'Main' }) } }, + { provide: TeamsService, useValue: { get: () => of({ display_name: 'Platform' }) } }, { provide: EventsService, useValue: { catalog$: of([]) } }, { provide: ActivatedRoute, useValue: { snapshot: { data, paramMap: convertToParamMap(params) } } }, ], @@ -57,6 +60,7 @@ describe('WebhooksComponent', () => { it.each([ { scope: 'project', trail: ['Projects', 'Acme', 'Settings', 'Webhooks'] }, { scope: 'cache', trail: ['Caches', 'Main', 'Settings', 'Webhooks'] }, + { scope: 'team', trail: ['Teams', 'Platform', 'Settings', 'Webhooks'] }, { scope: 'instance', trail: ['Job Board', 'Webhooks'] }, ])('leads the $scope scope back to the page its webhooks are reached from', ({ scope, trail }) => { const { root } = setup({ webhookScope: scope }, { [scope]: 'main' }); @@ -70,6 +74,13 @@ describe('WebhooksComponent', () => { expect(service.list).toHaveBeenCalledWith({ kind: 'cache', name: 'main' }); }); + it('team scope builds /teams/{team}/webhooks base', () => { + const { component, service } = setup({ webhookScope: 'team' }, { team: 'platform' }); + expect(component.scope).toEqual({ kind: 'team', name: 'platform' }); + expect(webhooksBase(component.scope)).toBe('teams/platform/webhooks'); + expect(service.list).toHaveBeenCalledWith({ kind: 'team', name: 'platform' }); + }); + it('shows the secret once after create', () => { const { component } = setup({ webhookScope: 'project' }, { project: 'p' }); component.name.set('ci'); diff --git a/frontend/src/app/features/webhooks/webhooks.component.ts b/frontend/src/app/features/webhooks/webhooks.component.ts index 4bf1e8366..2aa12c6c1 100644 --- a/frontend/src/app/features/webhooks/webhooks.component.ts +++ b/frontend/src/app/features/webhooks/webhooks.component.ts @@ -94,6 +94,10 @@ export class WebhooksComponent implements OnInit { this.scope = { kind: 'cache', name: paramMap.get('cache') ?? '' }; this.title = 'Cache Webhooks'; break; + case 'team': + this.scope = { kind: 'team', name: paramMap.get('team') ?? '' }; + this.title = 'Team Webhooks'; + break; default: this.scope = { kind: 'instance' }; this.title = 'Instance Webhooks'; @@ -108,6 +112,8 @@ export class WebhooksComponent implements OnInit { return this.crumbs.projectSettings(this.scope.name, page); case 'cache': return this.crumbs.cacheSettings(this.scope.name, page); + case 'team': + return this.crumbs.team(this.scope.name, { label: 'Settings', link: ['/team', this.scope.name, 'settings'] }, page); case 'instance': return [{ label: 'Job Board', link: ['/board', 'health'] }, page]; }