From 3cebe625c8b7b3fa93d092f82f4d28a3f1aedfe6 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 13:32:29 +0200 Subject: [PATCH 01/19] eval: fix: no substitutes in the evaluator --- backend/gradient-eval/src/nix_eval.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/backend/gradient-eval/src/nix_eval.rs b/backend/gradient-eval/src/nix_eval.rs index c2d41b10d..d78461383 100644 --- a/backend/gradient-eval/src/nix_eval.rs +++ b/backend/gradient-eval/src/nix_eval.rs @@ -89,6 +89,7 @@ impl NixEvaluator { let ctx = Arc::new(Context::new().context("nix context init")?); ctx.set_setting("show-trace", "true")?; ctx.set_setting("builders", "")?; + ctx.set_setting("substitute", "false")?; let store = Arc::new(Store::open(&ctx, None).context("nix store open")?); let flake_settings = Arc::new(FlakeSettings::new(&ctx)?); From 2e64d5e1ab2dc37b830381f35284e47dbadce294 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 14:30:19 +0200 Subject: [PATCH 02/19] wire: feat: network, secrets and exec capabilities --- .../gradient-db/src/metrics/worker_samples.rs | 6 +- backend/gradient-entity/src/team_worker.rs | 4 +- .../src/worker_registration.rs | 4 +- backend/gradient-fixtures/src/rows/worker.rs | 4 +- backend/gradient-migration/src/lib.rs | 2 + ...261018_000000_worker_network_capability.rs | 49 +++++++++++ backend/gradient-pool/src/aggregate.rs | 4 +- backend/gradient-pool/src/worker_pool.rs | 12 ++- backend/gradient-proto/src/handler/auth.rs | 47 +++++++--- .../src/handler/cache_consumer.rs | 4 +- .../src/handler/cache_session.rs | 4 +- backend/gradient-proto/src/handler/socket.rs | 2 +- backend/gradient-report/src/tables.rs | 16 ++-- .../gradient-scheduler/src/cases/shapes.rs | 3 +- backend/gradient-scheduler/src/jobs.rs | 12 ++- .../gradient-scheduler/src/waiting_state.rs | 5 +- backend/gradient-state/src/config.rs | 6 +- backend/gradient-state/src/export.rs | 12 ++- .../src/provisioning/entities/workers.rs | 83 ++++++++++++----- backend/gradient-web/src/endpoints/board.rs | 28 +++--- .../src/endpoints/board_metrics.rs | 4 +- .../src/endpoints/gradient_ci_connections.rs | 10 ++- .../src/endpoints/projects/workers.rs | 88 +++++++++++++++---- .../src/endpoints/teams/workers.rs | 22 +++-- backend/gradient-web/tests/team_workers.rs | 4 +- .../gradient-wire/examples/probe_handshake.rs | 4 +- backend/gradient-wire/schema/v31.txt | 56 ------------ backend/gradient-wire/schema/v32.txt | 58 ------------ backend/gradient-wire/schema/v33.txt | 58 ------------ .../gradient-wire/schema/{v34.txt => v35.txt} | 4 +- backend/gradient-wire/src/messages/client.rs | 2 +- backend/gradient-wire/src/messages/server.rs | 2 +- backend/gradient-wire/src/types.rs | 8 +- .../src/connection/handshake.rs | 8 +- backend/gradient-worker/options.json | 50 ++++++++--- backend/gradient-worker/src/config/mod.rs | 4 +- backend/gradient-worker/src/config/options.rs | 18 +++- nix/modules/gradient-state.nix | 18 +++- nix/standalone/default.nix | 2 +- 39 files changed, 416 insertions(+), 311 deletions(-) create mode 100644 backend/gradient-migration/src/m20261018_000000_worker_network_capability.rs delete mode 100644 backend/gradient-wire/schema/v31.txt delete mode 100644 backend/gradient-wire/schema/v32.txt delete mode 100644 backend/gradient-wire/schema/v33.txt rename backend/gradient-wire/schema/{v34.txt => v35.txt} (95%) diff --git a/backend/gradient-db/src/metrics/worker_samples.rs b/backend/gradient-db/src/metrics/worker_samples.rs index c3d810292..cd62d8d94 100644 --- a/backend/gradient-db/src/metrics/worker_samples.rs +++ b/backend/gradient-db/src/metrics/worker_samples.rs @@ -78,7 +78,7 @@ fn board_fleet_sql(window_hours: i64, project_filter: Option<&str>) -> String { count(DISTINCT worker_id) AS connected, \ count(DISTINCT worker_id) FILTER (WHERE state = {draining}) AS draining, \ count(DISTINCT worker_id) FILTER (WHERE (capabilities->>'eval')::boolean) AS ev, \ - count(DISTINCT worker_id) FILTER (WHERE (capabilities->>'fetch')::boolean) AS ft, \ + count(DISTINCT worker_id) FILTER (WHERE (capabilities->>'network')::boolean) AS nw, \ count(DISTINCT worker_id) FILTER (WHERE (capabilities->>'build')::boolean) AS bd \ FROM worker_sample \ WHERE at >= (now() AT TIME ZONE 'UTC') - interval '{window_hours} hours'", @@ -103,7 +103,7 @@ pub struct FleetCount { pub connected: i64, pub draining: i64, pub eval: i64, - pub fetch: i64, + pub network: i64, pub build: i64, } @@ -123,7 +123,7 @@ pub async fn fleet_per_hour( connected: r.try_get("", "connected").unwrap_or(0), draining: r.try_get("", "draining").unwrap_or(0), eval: r.try_get("", "ev").unwrap_or(0), - fetch: r.try_get("", "ft").unwrap_or(0), + network: r.try_get("", "nw").unwrap_or(0), build: r.try_get("", "bd").unwrap_or(0), }) .collect()) diff --git a/backend/gradient-entity/src/team_worker.rs b/backend/gradient-entity/src/team_worker.rs index f6890c5a0..a3ea4f520 100644 --- a/backend/gradient-entity/src/team_worker.rs +++ b/backend/gradient-entity/src/team_worker.rs @@ -24,7 +24,9 @@ pub struct Model { pub url: Option, pub display_name: String, pub gradient_ci: bool, - pub enable_fetch: bool, + pub enable_network: bool, + pub enable_secrets: bool, + pub enable_exec: bool, pub enable_eval: bool, pub enable_build: bool, pub active: bool, diff --git a/backend/gradient-entity/src/worker_registration.rs b/backend/gradient-entity/src/worker_registration.rs index e18a680f7..228cdfe56 100644 --- a/backend/gradient-entity/src/worker_registration.rs +++ b/backend/gradient-entity/src/worker_registration.rs @@ -25,7 +25,9 @@ pub struct Model { /// connection. pub url: Option, pub active: bool, - pub enable_fetch: bool, + pub enable_network: bool, + pub enable_secrets: bool, + pub enable_exec: bool, pub enable_eval: bool, pub enable_build: bool, pub display_name: String, diff --git a/backend/gradient-fixtures/src/rows/worker.rs b/backend/gradient-fixtures/src/rows/worker.rs index 9f3f8bbd3..9322e5a69 100644 --- a/backend/gradient-fixtures/src/rows/worker.rs +++ b/backend/gradient-fixtures/src/rows/worker.rs @@ -13,7 +13,9 @@ pub fn worker_registration(active: bool, enable_eval: bool) -> worker_registrati peer_id: ProjectId::nil(), worker_id: "00000000-0000-4000-8000-000000000001".into(), active, - enable_fetch: true, + enable_network: true, + enable_secrets: true, + enable_exec: true, enable_eval, enable_build: true, created_by: Some(UserId::nil()), diff --git a/backend/gradient-migration/src/lib.rs b/backend/gradient-migration/src/lib.rs index 0dc32dfbd..54fd4ce11 100644 --- a/backend/gradient-migration/src/lib.rs +++ b/backend/gradient-migration/src/lib.rs @@ -20,6 +20,7 @@ mod m20261015_000002_pull_request_summary; mod m20261016_000000_metric_input_pnames; mod m20261016_000001_metric_pname_without_version; mod m20261016_000002_team_webhooks; +mod m20261018_000000_worker_network_capability; pub struct UpgradeFloor { pub release: &'static str, @@ -57,6 +58,7 @@ impl MigratorTrait for Migrator { Box::new(m20261016_000000_metric_input_pnames::Migration), Box::new(m20261016_000001_metric_pname_without_version::Migration), Box::new(m20261016_000002_team_webhooks::Migration), + Box::new(m20261018_000000_worker_network_capability::Migration), ] } } diff --git a/backend/gradient-migration/src/m20261018_000000_worker_network_capability.rs b/backend/gradient-migration/src/m20261018_000000_worker_network_capability.rs new file mode 100644 index 000000000..fcdde5394 --- /dev/null +++ b/backend/gradient-migration/src/m20261018_000000_worker_network_capability.rs @@ -0,0 +1,49 @@ +/* + * 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.worker_registration + ADD COLUMN enable_network boolean DEFAULT true NOT NULL, + ADD COLUMN enable_secrets boolean DEFAULT true NOT NULL, + ADD COLUMN enable_exec boolean DEFAULT false NOT NULL; + +UPDATE public.worker_registration SET enable_network = enable_fetch, enable_secrets = enable_fetch; + +ALTER TABLE public.worker_registration DROP COLUMN enable_fetch; + +ALTER TABLE public.team_worker + ADD COLUMN enable_network boolean DEFAULT true NOT NULL, + ADD COLUMN enable_secrets boolean DEFAULT true NOT NULL, + ADD COLUMN enable_exec boolean DEFAULT false NOT NULL; + +UPDATE public.team_worker SET enable_network = enable_fetch, enable_secrets = enable_fetch; + +ALTER TABLE public.team_worker DROP COLUMN enable_fetch; + +UPDATE public.evaluation +SET waiting_reason = jsonb_set(waiting_reason, '{capability}', '"network"') +WHERE waiting_reason ->> 'capability' = 'fetch'; +"#; + +#[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, _: &SchemaManager) -> Result<(), DbErr> { + Err(DbErr::Migration( + "m20261018_000000_worker_network_capability is irreversible".into(), + )) + } +} diff --git a/backend/gradient-pool/src/aggregate.rs b/backend/gradient-pool/src/aggregate.rs index 618a88758..444f8bc24 100644 --- a/backend/gradient-pool/src/aggregate.rs +++ b/backend/gradient-pool/src/aggregate.rs @@ -99,7 +99,7 @@ mod tests { let mut pool = WorkerPool::new(); let eval = GradientCapabilities { eval: true, - fetch: true, + network: true, ..Default::default() }; register(&mut pool, "a", build(), &["x86_64-linux"], &["kvm"], 4); @@ -113,7 +113,7 @@ mod tests { ); let agg = pool.aggregate(); - assert!(agg.capabilities.build && agg.capabilities.eval && agg.capabilities.fetch); + assert!(agg.capabilities.build && agg.capabilities.eval && agg.capabilities.network); assert!(!agg.capabilities.cache); assert_eq!( agg.architectures diff --git a/backend/gradient-pool/src/worker_pool.rs b/backend/gradient-pool/src/worker_pool.rs index 342e73517..776f29822 100644 --- a/backend/gradient-pool/src/worker_pool.rs +++ b/backend/gradient-pool/src/worker_pool.rs @@ -142,7 +142,7 @@ impl WorkerPool { self.workers.get(id).map(|slot| { let s = slot.shared(); crate::WorkerCaps { - fetch: s.capabilities.fetch, + fetch: s.capabilities.network && s.capabilities.secrets, architectures: s.architectures.clone(), system_features: s.system_features.clone(), capabilities: s.capabilities.clone(), @@ -297,7 +297,9 @@ impl WorkerPool { pub fn has_idle_eval_only_worker(&self) -> bool { self.workers.values().any(|slot| match slot { WorkerSlot::Active(w) => { - w.capabilities.eval && !w.capabilities.fetch && w.assigned_jobs.is_empty() + w.capabilities.eval + && !(w.capabilities.network && w.capabilities.secrets) + && w.assigned_jobs.is_empty() } WorkerSlot::Draining(_) => false, }) @@ -477,7 +479,8 @@ mod tests { fn caps_ef(eval: bool, fetch: bool) -> GradientCapabilities { GradientCapabilities { eval, - fetch, + network: fetch, + secrets: fetch, ..GradientCapabilities::default() } } @@ -680,7 +683,8 @@ mod tests { pool.register( "w1".into(), GradientCapabilities { - fetch: true, + network: true, + secrets: true, eval: true, ..Default::default() }, diff --git a/backend/gradient-proto/src/handler/auth.rs b/backend/gradient-proto/src/handler/auth.rs index 499adc86a..1c716a623 100644 --- a/backend/gradient-proto/src/handler/auth.rs +++ b/backend/gradient-proto/src/handler/auth.rs @@ -197,7 +197,9 @@ pub(super) async fn resolve_authorized( /// clamping the advertised set with it. #[derive(Clone, Copy, Debug)] pub(super) struct EnabledCapsAggregate { - pub enable_fetch: bool, + pub enable_network: bool, + pub enable_secrets: bool, + pub enable_exec: bool, pub enable_eval: bool, pub enable_build: bool, } @@ -205,7 +207,9 @@ pub(super) struct EnabledCapsAggregate { impl EnabledCapsAggregate { pub fn all() -> Self { Self { - enable_fetch: true, + enable_network: true, + enable_secrets: true, + enable_exec: true, enable_eval: true, enable_build: true, } @@ -234,7 +238,9 @@ pub(super) async fn aggregate_enabled_caps(db: &WorkerDb, worker_id: &str) -> En gradient_db::teams::workers::active_team_worker(db, worker_id).await { return EnabledCapsAggregate { - enable_fetch: worker.enable_fetch, + enable_network: worker.enable_network, + enable_secrets: worker.enable_secrets, + enable_exec: worker.enable_exec, enable_eval: worker.enable_eval, enable_build: worker.enable_build, }; @@ -244,7 +250,9 @@ pub(super) async fn aggregate_enabled_caps(db: &WorkerDb, worker_id: &str) -> En } EnabledCapsAggregate { - enable_fetch: rows.iter().all(|r| r.enable_fetch), + enable_network: rows.iter().all(|r| r.enable_network), + enable_secrets: rows.iter().all(|r| r.enable_secrets), + enable_exec: rows.iter().all(|r| r.enable_exec), enable_eval: rows.iter().all(|r| r.enable_eval), enable_build: rows.iter().all(|r| r.enable_build), } @@ -276,9 +284,11 @@ pub(super) fn negotiate_capabilities( core: true, cache: true, federate: client.federate && config.proto.federate, - fetch: client.fetch && enabled.enable_fetch, eval: client.eval && enabled.enable_eval, build: client.build && enabled.enable_build, + network: client.network && enabled.enable_network, + secrets: client.secrets && enabled.enable_secrets, + exec: client.exec && enabled.enable_exec, } } @@ -295,9 +305,11 @@ mod tests { core: val, cache: val, federate: val, - fetch: val, eval: val, build: val, + network: val, + secrets: val, + exec: val, } } @@ -363,14 +375,15 @@ mod tests { core: false, cache: false, federate: false, - fetch: true, eval: true, build: true, + network: true, + secrets: true, + exec: true, }; let result = negotiate_capabilities(&state.config, client, EnabledCapsAggregate::all()); - assert!(result.fetch); - assert!(result.eval); - assert!(result.build); + assert!(result.eval && result.build); + assert!(result.network && result.secrets && result.exec); } #[test] @@ -378,16 +391,26 @@ mod tests { let state = make_state(false); let client = all_caps(true); let enabled = EnabledCapsAggregate { - enable_fetch: false, enable_eval: true, enable_build: false, + ..EnabledCapsAggregate::all() }; let result = negotiate_capabilities(&state.config, client, enabled); - assert!(!result.fetch); assert!(result.eval); assert!(!result.build); } + #[test] + fn negotiation_keeps_network_secrets_and_exec_only_when_enabled() { + let state = make_state(false); + let enabled = EnabledCapsAggregate { + enable_secrets: false, + ..EnabledCapsAggregate::all() + }; + let result = negotiate_capabilities(&state.config, all_caps(true), enabled); + assert!(result.network && !result.secrets && result.exec); + } + use gradient_entity::project::Model as ProjectModel; use gradient_entity::project_cache::{CacheSubscriptionMode, Model as ProjectCacheModel}; diff --git a/backend/gradient-proto/src/handler/cache_consumer.rs b/backend/gradient-proto/src/handler/cache_consumer.rs index b5312a7cb..b0d5be446 100644 --- a/backend/gradient-proto/src/handler/cache_consumer.rs +++ b/backend/gradient-proto/src/handler/cache_consumer.rs @@ -29,7 +29,9 @@ fn consumer_capabilities() -> GradientCapabilities { core: true, build: false, eval: false, - fetch: false, + network: false, + secrets: false, + exec: false, cache: true, federate: false, } diff --git a/backend/gradient-proto/src/handler/cache_session.rs b/backend/gradient-proto/src/handler/cache_session.rs index d56295dff..568fa3d96 100644 --- a/backend/gradient-proto/src/handler/cache_session.rs +++ b/backend/gradient-proto/src/handler/cache_session.rs @@ -37,7 +37,9 @@ fn readonly_capabilities() -> GradientCapabilities { core: false, build: false, eval: false, - fetch: false, + network: false, + secrets: false, + exec: false, cache: true, federate: false, } diff --git a/backend/gradient-proto/src/handler/socket.rs b/backend/gradient-proto/src/handler/socket.rs index cffe338fd..b76d2b5eb 100644 --- a/backend/gradient-proto/src/handler/socket.rs +++ b/backend/gradient-proto/src/handler/socket.rs @@ -51,7 +51,7 @@ pub(super) async fn send_credentials_for_job( use gradient_wire::types::{FlakeStep, Job}; let caps = scheduler.worker_gradient_caps(worker_id).await; - let worker_can_fetch = caps.as_ref().map(|c| c.fetch).unwrap_or(false); + let worker_can_fetch = caps.as_ref().is_some_and(|c| c.network && c.secrets); match job { Job::Flake(flake_job) => { diff --git a/backend/gradient-report/src/tables.rs b/backend/gradient-report/src/tables.rs index 23c16bda0..75ce0bffe 100644 --- a/backend/gradient-report/src/tables.rs +++ b/backend/gradient-report/src/tables.rs @@ -501,8 +501,8 @@ pub fn instance_tables() -> &'static [TableSpec] { const SPECS: &[TableSpec] = &[ spec!( "worker_registration", - "CREATE TABLE worker_registration (id TEXT, peer_id TEXT, worker_id TEXT, created_at TEXT, managed INTEGER, url TEXT, active INTEGER, display_name TEXT, created_by TEXT, enable_fetch INTEGER, enable_eval INTEGER, enable_build INTEGER)", - "SELECT id::text, peer_id::text, worker_id::text, created_at::text, managed::int::text, url::text, active::int::text, display_name::text, created_by::text, enable_fetch::int::text, enable_eval::int::text, enable_build::int::text FROM worker_registration WHERE $1 IS NOT NULL", + "CREATE TABLE worker_registration (id TEXT, peer_id TEXT, worker_id TEXT, created_at TEXT, managed INTEGER, url TEXT, active INTEGER, display_name TEXT, created_by TEXT, enable_network INTEGER, enable_secrets INTEGER, enable_exec INTEGER, enable_eval INTEGER, enable_build INTEGER)", + "SELECT id::text, peer_id::text, worker_id::text, created_at::text, managed::int::text, url::text, active::int::text, display_name::text, created_by::text, enable_network::int::text, enable_secrets::int::text, enable_exec::int::text, enable_eval::int::text, enable_build::int::text FROM worker_registration WHERE $1 IS NOT NULL", "the whole instance", [ "id", @@ -514,7 +514,9 @@ pub fn instance_tables() -> &'static [TableSpec] { "active", "display_name", "created_by", - "enable_fetch", + "enable_network", + "enable_secrets", + "enable_exec", "enable_eval", "enable_build" ] @@ -589,8 +591,8 @@ pub fn instance_tables() -> &'static [TableSpec] { ), spec!( "team_worker", - "CREATE TABLE team_worker (id TEXT, team TEXT, worker_id TEXT, display_name TEXT, url TEXT, active INTEGER, gradient_ci INTEGER, enable_fetch INTEGER, enable_eval INTEGER, enable_build INTEGER, created_by TEXT, created_at TEXT)", - "SELECT id::text, team::text, worker_id::text, display_name::text, url::text, active::int::text, gradient_ci::int::text, enable_fetch::int::text, enable_eval::int::text, enable_build::int::text, created_by::text, created_at::text FROM team_worker WHERE $1 IS NOT NULL", + "CREATE TABLE team_worker (id TEXT, team TEXT, worker_id TEXT, display_name TEXT, url TEXT, active INTEGER, gradient_ci INTEGER, enable_network INTEGER, enable_secrets INTEGER, enable_exec INTEGER, enable_eval INTEGER, enable_build INTEGER, created_by TEXT, created_at TEXT)", + "SELECT id::text, team::text, worker_id::text, display_name::text, url::text, active::int::text, gradient_ci::int::text, enable_network::int::text, enable_secrets::int::text, enable_exec::int::text, enable_eval::int::text, enable_build::int::text, created_by::text, created_at::text FROM team_worker WHERE $1 IS NOT NULL", "the whole instance", [ "id", @@ -600,7 +602,9 @@ pub fn instance_tables() -> &'static [TableSpec] { "url", "active", "gradient_ci", - "enable_fetch", + "enable_network", + "enable_secrets", + "enable_exec", "enable_eval", "enable_build", "created_by", diff --git a/backend/gradient-scheduler/src/cases/shapes.rs b/backend/gradient-scheduler/src/cases/shapes.rs index adb55130f..976a99eee 100644 --- a/backend/gradient-scheduler/src/cases/shapes.rs +++ b/backend/gradient-scheduler/src/cases/shapes.rs @@ -99,7 +99,8 @@ impl Profile { fn caps(fetch: bool, eval: bool, build: bool) -> GradientCapabilities { GradientCapabilities { - fetch, + network: fetch, + secrets: fetch, eval, build, ..GradientCapabilities::default() diff --git a/backend/gradient-scheduler/src/jobs.rs b/backend/gradient-scheduler/src/jobs.rs index f88d20bec..08a2bd801 100644 --- a/backend/gradient-scheduler/src/jobs.rs +++ b/backend/gradient-scheduler/src/jobs.rs @@ -2171,7 +2171,8 @@ mod tests { let fetch_only = WorkerCaps { fetch: true, capabilities: GradientCapabilities { - fetch: true, + network: true, + secrets: true, ..Default::default() }, ..Default::default() @@ -2187,7 +2188,8 @@ mod tests { let both = WorkerCaps { fetch: true, capabilities: GradientCapabilities { - fetch: true, + network: true, + secrets: true, eval: true, ..Default::default() }, @@ -2376,7 +2378,8 @@ mod tests { architectures: vec!["x86_64-linux".into()], system_features: vec![], capabilities: GradientCapabilities { - fetch: true, + network: true, + secrets: true, eval: true, ..Default::default() }, @@ -2434,7 +2437,8 @@ mod tests { architectures: vec!["x86_64-linux".into()], system_features: vec![], capabilities: GradientCapabilities { - fetch: true, + network: true, + secrets: true, build: true, ..Default::default() }, diff --git a/backend/gradient-scheduler/src/waiting_state.rs b/backend/gradient-scheduler/src/waiting_state.rs index 486292246..b4fedeb33 100644 --- a/backend/gradient-scheduler/src/waiting_state.rs +++ b/backend/gradient-scheduler/src/waiting_state.rs @@ -714,7 +714,7 @@ impl ProjectWorkerCounts { .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), + fetch: counts.fetch + usize::from(w.capabilities.network && w.capabilities.secrets), connected: counts.connected + 1, }) } @@ -842,7 +842,8 @@ mod tests { id: "00000000-0000-4000-8000-000000000001".into(), capabilities: gradient_wire::types::GradientCapabilities { eval, - fetch: eval, + network: eval, + secrets: eval, ..Default::default() }, architectures: vec![], diff --git a/backend/gradient-state/src/config.rs b/backend/gradient-state/src/config.rs index 834656254..c42d2a71c 100644 --- a/backend/gradient-state/src/config.rs +++ b/backend/gradient-state/src/config.rs @@ -242,7 +242,11 @@ pub struct StateWorker { #[serde(default)] pub created_by: Option, #[serde(default = "default_true")] - pub enable_fetch: bool, + pub enable_network: bool, + #[serde(default = "default_true")] + pub enable_secrets: bool, + #[serde(default)] + pub enable_exec: bool, #[serde(default = "default_true")] pub enable_eval: bool, #[serde(default = "default_true")] diff --git a/backend/gradient-state/src/export.rs b/backend/gradient-state/src/export.rs index 566b37d4c..6176c9e15 100644 --- a/backend/gradient-state/src/export.rs +++ b/backend/gradient-state/src/export.rs @@ -342,7 +342,9 @@ pub async fn export_state(db: &C) -> Result StateApplicator<'a> { reg.url = Set(url.clone()); reg.display_name = Set(state_worker.display_name.clone()); reg.active = Set(state_worker.enabled); - reg.enable_fetch = Set(state_worker.enable_fetch); + reg.enable_network = Set(state_worker.enable_network); + reg.enable_secrets = Set(state_worker.enable_secrets); + reg.enable_exec = Set(state_worker.enable_exec); reg.enable_eval = Set(state_worker.enable_eval); reg.enable_build = Set(state_worker.enable_build); reg.created_by = Set(created_by_id); @@ -102,7 +104,9 @@ impl<'a> StateApplicator<'a> { display_name: state_worker.display_name.clone(), gradient_ci: false, active: state_worker.enabled, - enable_fetch: state_worker.enable_fetch, + enable_network: state_worker.enable_network, + enable_secrets: state_worker.enable_secrets, + enable_exec: state_worker.enable_exec, enable_eval: state_worker.enable_eval, enable_build: state_worker.enable_build, created_by: created_by_id, @@ -165,7 +169,9 @@ async fn apply_team_worker( active.token_encrypted = Set(token_encrypted); active.url = Set(worker.url.clone()); active.display_name = Set(worker.display_name.clone()); - active.enable_fetch = Set(worker.enable_fetch); + active.enable_network = Set(worker.enable_network); + active.enable_secrets = Set(worker.enable_secrets); + active.enable_exec = Set(worker.enable_exec); active.enable_eval = Set(worker.enable_eval); active.enable_build = Set(worker.enable_build); active.active = Set(worker.enabled); @@ -184,7 +190,9 @@ async fn apply_team_worker( url: worker.url.clone(), display_name: worker.display_name.clone(), gradient_ci: false, - enable_fetch: worker.enable_fetch, + enable_network: worker.enable_network, + enable_secrets: worker.enable_secrets, + enable_exec: worker.enable_exec, enable_eval: worker.enable_eval, enable_build: worker.enable_build, active: worker.enabled, @@ -205,8 +213,25 @@ mod team_worker_tests { use gradient_fixtures::mock::{all_statements, mock_db}; use sea_orm::MockExecResult; - #[tokio::test] - async fn a_declared_team_worker_is_inserted_as_managed() { + fn state_worker(enable_secrets: bool) -> StateWorker { + StateWorker { + worker_id: "w1".into(), + url: None, + projects: vec![], + team: Some("platform".into()), + token_file: "/dev/null".into(), + display_name: "W".into(), + created_by: None, + enable_network: true, + enable_secrets, + enable_exec: true, + enable_eval: true, + enable_build: true, + enabled: true, + } + } + + async fn inserted_team_worker(worker: &StateWorker) -> sea_orm::Statement { let team = TeamId::now_v7(); let db = mock_db() .append_query_results([Vec::::new()]) @@ -220,28 +245,42 @@ mod team_worker_tests { rows_affected: 1, }]) .into_connection(); - let worker = StateWorker { - worker_id: "w1".into(), - url: None, - projects: vec![], - team: Some("platform".into()), - token_file: "/dev/null".into(), - display_name: "W".into(), - created_by: None, - enable_fetch: true, - enable_eval: true, - enable_build: true, - enabled: true, - }; - apply_team_worker(&db, &worker, team, None, "hash".into(), None) + apply_team_worker(&db, worker, team, None, "hash".into(), None) .await .unwrap(); - let inserted = all_statements(db.into_transaction_log()) + all_statements(db.into_transaction_log()) .into_iter() .find(|s| s.sql.starts_with("INSERT INTO \"team_worker\"")) - .expect("an INSERT INTO team_worker"); + .expect("an INSERT INTO team_worker") + } + + fn inserted_value(statement: &sea_orm::Statement, column: &str) -> sea_orm::Value { + let columns = statement.sql.split(['(', ')']).nth(1).expect("column list"); + let index = columns + .split(", ") + .position(|c| c == format!("\"{column}\"")) + .expect("column in the INSERT"); + statement.values.as_ref().expect("values").0[index].clone() + } + + #[tokio::test] + async fn a_declared_team_worker_is_inserted_as_managed() { + let inserted = inserted_team_worker(&state_worker(true)).await; assert!(inserted.sql.contains("\"managed\"")); } + + #[tokio::test] + async fn a_worker_without_secrets_keeps_network() { + let inserted = inserted_team_worker(&state_worker(false)).await; + assert_eq!( + inserted_value(&inserted, "enable_secrets"), + sea_orm::Value::Bool(Some(false)) + ); + assert_eq!( + inserted_value(&inserted, "enable_network"), + sea_orm::Value::Bool(Some(true)) + ); + } } diff --git a/backend/gradient-web/src/endpoints/board.rs b/backend/gradient-web/src/endpoints/board.rs index 9a2f46182..edba209f4 100644 --- a/backend/gradient-web/src/endpoints/board.rs +++ b/backend/gradient-web/src/endpoints/board.rs @@ -684,7 +684,7 @@ pub struct BoardWorker { pub assigned_jobs: i64, pub max_concurrent_builds: i64, pub eval: bool, - pub fetch: bool, + pub network: bool, pub build: bool, pub architectures: Vec, pub cpu_usage_pct: Option, @@ -748,7 +748,7 @@ pub async fn get_board_workers( assigned_jobs: w.assigned_job_count as i64, max_concurrent_builds: w.max_concurrent_builds as i64, eval: w.capabilities.eval, - fetch: w.capabilities.fetch, + network: w.capabilities.network, build: w.capabilities.build, architectures: if accessible { w.architectures } else { vec![] }, cpu_usage_pct: accessible.then_some(w.cpu_usage_pct).flatten(), @@ -822,8 +822,8 @@ fn aggregate_worker_load( if w.capabilities.eval { bump_capacity(&mut cap, "eval", slots); } - if w.capabilities.fetch { - bump_capacity(&mut cap, "fetch", slots); + if w.capabilities.network { + bump_capacity(&mut cap, "network", slots); } if w.capabilities.build { bump_capacity(&mut cap, "build", slots); @@ -844,7 +844,7 @@ fn aggregate_worker_load( bump_in_flight(&mut cap, "eval"); } if j.fetch_step { - bump_in_flight(&mut cap, "fetch"); + bump_in_flight(&mut cap, "network"); } } } @@ -858,7 +858,7 @@ fn aggregate_worker_load( } } - let by_capability = ["eval", "fetch", "build"] + let by_capability = ["eval", "network", "build"] .into_iter() .map(|k| LoadBucket::from((k.to_owned(), cap.get(k).copied().unwrap_or_default()))) .collect(); @@ -876,7 +876,7 @@ fn aggregate_worker_load( operation_id = "getBoardWorkerLoad", tag = "board", summary = "Fleet load by capability, architecture, and feature", - description = "Real per-dimension fleet load: for each capability (eval/fetch/build), architecture, and system feature, the in-flight job count against the summed slot capacity of the workers that can serve it. Answers whether the fleet is eval-, build-, or architecture-bound rather than showing one blended busy %. Scoped to the caller's projects (superusers see the whole fleet). Each bucket carries `{ key, in_flight, capacity, workers }`; busy % = in_flight / capacity.", + description = "Real per-dimension fleet load: for each capability (eval/network/build), architecture, and system feature, the in-flight job count against the summed slot capacity of the workers that can serve it. Answers whether the fleet is eval-, build-, or architecture-bound rather than showing one blended busy %. Scoped to the caller's projects (superusers see the whole fleet). Each bucket carries `{ key, in_flight, capacity, workers }`; busy % = in_flight / capacity.", responses((status = OK, description = "Fleet load breakdowns", body = WorkerLoad), WebError) )] pub async fn get_board_worker_load( @@ -1374,7 +1374,7 @@ mod tests { fn worker( eval: bool, - fetch: bool, + network: bool, build: bool, arch: &[&str], features: &[&str], @@ -1384,7 +1384,7 @@ mod tests { id: "w".into(), capabilities: GradientCapabilities { eval, - fetch, + network, build, ..Default::default() }, @@ -1514,7 +1514,7 @@ mod tests { let cap = &load.by_capability; assert_eq!( cap.iter().map(|b| b.key.as_str()).collect::>(), - ["eval", "fetch", "build"] + ["eval", "network", "build"] ); assert_eq!( *bucket(cap, "eval"), @@ -1526,9 +1526,9 @@ mod tests { } ); assert_eq!( - *bucket(cap, "fetch"), + *bucket(cap, "network"), LoadBucket { - key: "fetch".into(), + key: "network".into(), in_flight: 1, capacity: 8, workers: 1 @@ -1586,9 +1586,9 @@ mod tests { let load = aggregate_worker_load(&workers, &job_refs); assert!(load.by_architecture.iter().all(|b| b.key != BUILTIN_ARCH)); assert_eq!( - *bucket(&load.by_capability, "fetch"), + *bucket(&load.by_capability, "network"), LoadBucket { - key: "fetch".into(), + key: "network".into(), in_flight: 0, capacity: 0, workers: 0 diff --git a/backend/gradient-web/src/endpoints/board_metrics.rs b/backend/gradient-web/src/endpoints/board_metrics.rs index 2f29674c1..e7e16ea63 100644 --- a/backend/gradient-web/src/endpoints/board_metrics.rs +++ b/backend/gradient-web/src/endpoints/board_metrics.rs @@ -336,7 +336,7 @@ pub struct BoardFleetPoint { pub connected: i64, pub draining: i64, pub eval: i64, - pub fetch: i64, + pub network: i64, pub build: i64, } @@ -347,7 +347,7 @@ impl From for BoardFleetPoint { connected: fleet.connected, draining: fleet.draining, eval: fleet.eval, - fetch: fleet.fetch, + network: fleet.network, build: fleet.build, } } diff --git a/backend/gradient-web/src/endpoints/gradient_ci_connections.rs b/backend/gradient-web/src/endpoints/gradient_ci_connections.rs index 3bc083e6d..16dc2e329 100644 --- a/backend/gradient-web/src/endpoints/gradient_ci_connections.rs +++ b/backend/gradient-web/src/endpoints/gradient_ci_connections.rs @@ -345,7 +345,9 @@ fn registration_row( managed: false, url: Some(url), active: true, - enable_fetch: false, + enable_network: true, + enable_secrets: false, + enable_exec: false, enable_eval: true, enable_build: true, display_name: DISPLAY_NAME.into(), @@ -371,7 +373,9 @@ fn team_worker_row( url: Some(url), display_name: DISPLAY_NAME.into(), gradient_ci: true, - enable_fetch: false, + enable_network: true, + enable_secrets: false, + enable_exec: false, enable_eval: true, enable_build: true, active: true, @@ -449,7 +453,7 @@ mod tests { ); assert!(row.gradient_ci && row.active && row.enable_eval && row.enable_build); - assert!(!row.enable_fetch && !row.managed); + assert!(row.enable_network && !row.enable_secrets && !row.managed); assert_eq!(row.token_encrypted.as_deref(), Some("encrypted")); assert_eq!(row.url.as_deref(), Some("wss://servers.gradient.ci/proto")); } diff --git a/backend/gradient-web/src/endpoints/projects/workers.rs b/backend/gradient-web/src/endpoints/projects/workers.rs index 663eff0ed..e95f3e3ed 100644 --- a/backend/gradient-web/src/endpoints/projects/workers.rs +++ b/backend/gradient-web/src/endpoints/projects/workers.rs @@ -45,7 +45,11 @@ pub struct RegisterWorkerRequest { pub display_name: String, pub token: Option, #[serde(default = "default_true")] - pub enable_fetch: bool, + pub enable_network: bool, + #[serde(default = "default_true")] + pub enable_secrets: bool, + #[serde(default)] + pub enable_exec: bool, #[serde(default = "default_true")] pub enable_eval: bool, #[serde(default = "default_true")] @@ -68,7 +72,9 @@ pub struct ProjectWorkerEntry { #[serde(skip_serializing_if = "Option::is_none")] pub url: Option, pub created_by: Option, - pub enable_fetch: bool, + pub enable_network: bool, + pub enable_secrets: bool, + pub enable_exec: bool, pub enable_eval: bool, pub enable_build: bool, #[serde(skip_serializing_if = "Option::is_none")] @@ -93,7 +99,9 @@ impl ProjectWorkerEntry { active: registration.active, url: registration.url, created_by: registration.created_by, - enable_fetch: registration.enable_fetch, + enable_network: registration.enable_network, + enable_secrets: registration.enable_secrets, + enable_exec: registration.enable_exec, enable_eval: registration.enable_eval, enable_build: registration.enable_build, team: None, @@ -116,7 +124,9 @@ impl ProjectWorkerEntry { active: worker.active, url: worker.url, created_by: worker.created_by, - enable_fetch: worker.enable_fetch, + enable_network: worker.enable_network, + enable_secrets: worker.enable_secrets, + enable_exec: worker.enable_exec, enable_eval: worker.enable_eval, enable_build: worker.enable_build, team: Some(team), @@ -149,16 +159,27 @@ pub(crate) fn worker_connection( pub struct PatchWorkerRequest { pub active: Option, pub display_name: Option, - pub enable_fetch: Option, + pub enable_network: Option, + pub enable_secrets: Option, + pub enable_exec: Option, pub enable_eval: Option, pub enable_build: Option, } +pub(crate) fn patch_changes_capabilities(body: &PatchWorkerRequest) -> bool { + [ + body.enable_network, + body.enable_secrets, + body.enable_exec, + body.enable_eval, + body.enable_build, + ] + .iter() + .any(Option::is_some) +} + pub(crate) fn patch_edits_managed_fields(body: &PatchWorkerRequest) -> bool { - body.display_name.is_some() - || body.enable_fetch.is_some() - || body.enable_eval.is_some() - || body.enable_build.is_some() + body.display_name.is_some() || patch_changes_capabilities(body) } #[derive(Serialize, utoipa::ToSchema)] @@ -263,7 +284,9 @@ pub async fn post_project_worker( url: body.url, display_name: body.display_name.trim().to_string(), active: true, - enable_fetch: body.enable_fetch, + enable_network: body.enable_network, + enable_secrets: body.enable_secrets, + enable_exec: body.enable_exec, enable_eval: body.enable_eval, enable_build: body.enable_build, created_by: Some(user.id), @@ -600,10 +623,15 @@ pub async fn patch_project_worker( if let Some(ref name) = body.display_name { active_model.display_name = Set(name.trim().to_string()); } - let caps_changed = - body.enable_fetch.is_some() || body.enable_eval.is_some() || body.enable_build.is_some(); - if let Some(v) = body.enable_fetch { - active_model.enable_fetch = Set(v); + let caps_changed = patch_changes_capabilities(&body); + if let Some(v) = body.enable_network { + active_model.enable_network = Set(v); + } + if let Some(v) = body.enable_secrets { + active_model.enable_secrets = Set(v); + } + if let Some(v) = body.enable_exec { + active_model.enable_exec = Set(v); } if let Some(v) = body.enable_eval { active_model.enable_eval = Set(v); @@ -753,7 +781,9 @@ mod tests { active: true, url: None, created_by: None, - enable_fetch: true, + enable_network: true, + enable_secrets: true, + enable_exec: true, enable_eval: true, enable_build: true, team: None, @@ -791,7 +821,9 @@ mod tests { PatchWorkerRequest { active: None, display_name: None, - enable_fetch: None, + enable_network: None, + enable_secrets: None, + enable_exec: None, enable_eval: None, enable_build: None, } @@ -815,7 +847,15 @@ mod tests { ..empty_patch() }, PatchWorkerRequest { - enable_fetch: Some(false), + enable_network: Some(false), + ..empty_patch() + }, + PatchWorkerRequest { + enable_secrets: Some(false), + ..empty_patch() + }, + PatchWorkerRequest { + enable_exec: Some(false), ..empty_patch() }, PatchWorkerRequest { @@ -831,6 +871,20 @@ mod tests { } } + #[test] + fn patching_secrets_changes_the_capabilities() { + let body = PatchWorkerRequest { + enable_secrets: Some(false), + ..empty_patch() + }; + assert!(patch_changes_capabilities(&body)); + assert!(!patch_changes_capabilities(&PatchWorkerRequest { + active: Some(false), + display_name: Some("x".into()), + ..empty_patch() + })); + } + #[test] fn worker_test_result_covers_all_states() { let (ok, msg) = worker_test_result(false, false); diff --git a/backend/gradient-web/src/endpoints/teams/workers.rs b/backend/gradient-web/src/endpoints/teams/workers.rs index 72589feec..2d88c72ea 100644 --- a/backend/gradient-web/src/endpoints/teams/workers.rs +++ b/backend/gradient-web/src/endpoints/teams/workers.rs @@ -42,7 +42,9 @@ pub struct TeamWorkerEntry { #[serde(skip_serializing_if = "Option::is_none")] pub url: Option, pub gradient_ci: bool, - pub enable_fetch: bool, + pub enable_network: bool, + pub enable_secrets: bool, + pub enable_exec: bool, pub enable_eval: bool, pub enable_build: bool, #[serde(flatten)] @@ -65,7 +67,9 @@ impl TeamWorkerEntry { managed: worker.managed, url: worker.url, gradient_ci: worker.gradient_ci, - enable_fetch: worker.enable_fetch, + enable_network: worker.enable_network, + enable_secrets: worker.enable_secrets, + enable_exec: worker.enable_exec, enable_eval: worker.enable_eval, enable_build: worker.enable_build, connection, @@ -200,7 +204,9 @@ pub async fn post_team_worker( url: body.url, display_name: body.display_name.trim().to_string(), gradient_ci: false, - enable_fetch: body.enable_fetch, + enable_network: body.enable_network, + enable_secrets: body.enable_secrets, + enable_exec: body.enable_exec, enable_eval: body.enable_eval, enable_build: body.enable_build, active: true, @@ -266,8 +272,14 @@ pub async fn patch_team_worker( if let Some(name) = body.display_name { active.display_name = Set(name.trim().to_string()); } - if let Some(value) = body.enable_fetch { - active.enable_fetch = Set(value); + if let Some(value) = body.enable_network { + active.enable_network = Set(value); + } + if let Some(value) = body.enable_secrets { + active.enable_secrets = Set(value); + } + if let Some(value) = body.enable_exec { + active.enable_exec = Set(value); } if let Some(value) = body.enable_eval { active.enable_eval = Set(value); diff --git a/backend/gradient-web/tests/team_workers.rs b/backend/gradient-web/tests/team_workers.rs index 2f16174b5..7f3dd6864 100644 --- a/backend/gradient-web/tests/team_workers.rs +++ b/backend/gradient-web/tests/team_workers.rs @@ -27,7 +27,9 @@ fn team_worker_row() -> team_worker::Model { token_hash: "hash".into(), display_name: "Builder".into(), active: true, - enable_fetch: true, + enable_network: true, + enable_secrets: true, + enable_exec: true, enable_eval: true, enable_build: true, created_at: test_date(), diff --git a/backend/gradient-wire/examples/probe_handshake.rs b/backend/gradient-wire/examples/probe_handshake.rs index 1e9b8c198..884435a19 100644 --- a/backend/gradient-wire/examples/probe_handshake.rs +++ b/backend/gradient-wire/examples/probe_handshake.rs @@ -45,7 +45,9 @@ async fn main() { capabilities: GradientCapabilities { core: false, federate: false, - fetch: true, + network: true, + secrets: true, + exec: true, eval: true, build: true, cache: false, diff --git a/backend/gradient-wire/schema/v31.txt b/backend/gradient-wire/schema/v31.txt deleted file mode 100644 index 498304a0c..000000000 --- a/backend/gradient-wire/schema/v31.txt +++ /dev/null @@ -1,56 +0,0 @@ -client -0:{{bool,bool,bool,bool,bool,bool},str} -1:{seq<(str,str)>} -2:{} -3:{u16,str} -4:{seq,seq,u32,u32,u64,u32,opt,opt} -5:{f32,u64,opt,opt,opt} -6:{} -7:{seq<{str,u32,u64,bool}>,bool} -8:{str,bool,opt} -9:{str,str,[0:{},1:{opt},2:{},3:{},4:{seq<{str,str,seq<{str,str}>,seq,seq,str,seq,opt,opt,bool,bool,bool,opt}>,seq,seq},5:{str},6:{str,seq<{str,str,str,opt,opt,seq<{str,str,str,str,opt}>}>,opt<{opt,opt,opt,opt,opt,bool,opt,opt,opt,opt}>,bool},7:{},8:{{u64,u64,u64,u64,u64,u64,u64,u64,str,seq<{str,u64,u64,u64,u64}>,seq<{str,opt,str,str,bool,opt}>}},9:{str,seq<{str,opt,str}>},10:{seq},11:{[0:{},1:{},2:{}]}]} -10:{str,str,seq<{[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{},8:{},9:{},10:{},11:{},12:{},13:{},14:{},15:{},16:{},17:{}],u64,u64,opt,u32,u64}>,u64} -11:{str,str,str,[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{}],seq,seq<{[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{},8:{},9:{},10:{},11:{},12:{},13:{},14:{},15:{},16:{},17:{}],u64,u64,opt,u32,u64}>,u64,opt<{opt,opt,opt,opt,opt,bool,opt,opt,opt,opt}>} -12:{} -13:{str,str,str,[0:{},1:{},2:{}],u64,opt,u32,opt} -14:{str,str,[0:{seq<{str,[0:{},1:{},2:{},3:{}],u64,u64}>},1:{u64}]} -15:{str,u32,bytes} -16:{str,seq} -17:{str,str,u64,str} -18:{str,str} -19:{[0:{},1:{}]} -20:{str,opt<{str,u32}>,bytes} -21:{str,str,seq,[0:{},1:{},2:{}],seq>,bool} -22:{str,[0:{},1:{},2:{}],str,str} -23:{str,str,seq} -24:{str,u64,[0:{str},1:{str}],u64} -25:{u64,bytes,u64,bool} -26:{u64,[0:{{str,u64,u64,str,seq,opt,opt,opt<{str,seq}>}},1:{u64}]} -27:{u64} -server -0:{seq} -1:{{bool,bool,bool,bool,bool,bool},seq,seq<{str,str}>} -2:{seq,seq<{str,str}>} -3:{u16,str} -4:{u16,str} -5:{} -6:{seq<{str,seq<{str,opt<{u64,u64}>}>,seq,seq,opt<{str,seq}>}>,bool} -7:{seq<{str,seq<{str,opt<{u64,u64}>}>,seq,seq,opt<{str,seq}>}>} -8:{str,str,[0:{{seq<[0:{},1:{},2:{}]>,[0:{str,str},1:{str}],seq,opt,seq<{str,opt}>,opt<{str,seq,bool}>}},1:{{seq<{str,str,[0:{},1:{},2:{}],bool,seq<{str,str}>,opt,opt}>,{str,seq}}}],opt<{str,str,u32,u32}>} -9:{str,str} -10:{str,seq<{str,u32,str,opt,opt}>} -11:{str,{str,u32},bytes} -12:{str,str} -13:{[0:{}],bytes} -14:{str,str,bytes,u64,bool} -15:{str,str,str} -16:{str,str,str} -17:{str,str,u64,str} -18:{str,[0:{},1:{str},2:{u64,str}]} -19:{str,bytes,u64,bool} -20:{str,seq<{str,bool,opt,opt,opt,opt,opt,opt>,opt>,opt,opt}>} -21:{str,seq} -22:{str,str} -23:{u64,[0:{},1:{u64},2:{str},3:{{str,u64,seq}}]} -24:{u64,[0:{},1:{str},2:{str}]} -25:{str,seq<(str,str)>} diff --git a/backend/gradient-wire/schema/v32.txt b/backend/gradient-wire/schema/v32.txt deleted file mode 100644 index f89b8e384..000000000 --- a/backend/gradient-wire/schema/v32.txt +++ /dev/null @@ -1,58 +0,0 @@ -client -0:{{bool,bool,bool,bool,bool,bool},str} -1:{seq<(str,str)>} -2:{} -3:{u16,str} -4:{seq,seq,u32,u32,u64,u32,opt,opt} -5:{f32,u64,opt,opt,opt} -6:{} -7:{seq<{str,u32,u64,bool}>,bool} -8:{str,bool,opt} -9:{str,str,[0:{},1:{opt},2:{},3:{},4:{seq<{str,str,seq<{str,str}>,seq,seq,str,seq,opt,opt,bool,bool,bool,opt,bool}>,seq,seq},5:{str},6:{str,seq<{str,str,str,opt,opt,seq<{str,str,str,str,opt}>}>,opt<{opt,opt,opt,opt,opt,bool,opt,opt,opt,opt}>,bool},7:{},8:{{u64,u64,u64,u64,u64,u64,u64,u64,str,seq<{str,u64,u64,u64,u64}>,seq<{str,opt,str,str,bool,opt}>}},9:{str,seq<{str,opt,str}>},10:{seq},11:{[0:{},1:{},2:{}]}]} -10:{str,str,seq<{[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{},8:{},9:{},10:{},11:{},12:{},13:{},14:{},15:{},16:{},17:{}],u64,u64,opt,u32,u64}>,u64} -11:{str,str,str,[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{}],seq,seq<{[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{},8:{},9:{},10:{},11:{},12:{},13:{},14:{},15:{},16:{},17:{}],u64,u64,opt,u32,u64}>,u64,opt<{opt,opt,opt,opt,opt,bool,opt,opt,opt,opt}>} -12:{} -13:{str,str,str,[0:{},1:{},2:{}],u64,opt,u32,opt} -14:{str,str,[0:{seq<{str,[0:{},1:{},2:{},3:{}],u64,u64}>},1:{u64}]} -15:{str,u32,bytes} -16:{str,seq} -17:{str,str,u64,str} -18:{str,str} -19:{[0:{},1:{}]} -20:{str,opt<{str,u32}>,bytes} -21:{str,str,seq,[0:{},1:{},2:{}],seq>,bool} -22:{str,[0:{},1:{},2:{}],str,str} -23:{str,str,seq} -24:{str,u64,[0:{str},1:{str}],u64} -25:{u64,bytes,u64,bool} -26:{u64,[0:{{str,u64,u64,str,seq,opt,opt,opt<{str,seq}>}},1:{u64}]} -27:{u64} -28:{str,str,seq} -server -0:{seq} -1:{{bool,bool,bool,bool,bool,bool},seq,seq<{str,str}>} -2:{seq,seq<{str,str}>} -3:{u16,str} -4:{u16,str} -5:{} -6:{seq<{str,seq<{str,opt<{u64,u64}>}>,seq,seq,opt<{str,seq}>}>,bool} -7:{seq<{str,seq<{str,opt<{u64,u64}>}>,seq,seq,opt<{str,seq}>}>} -8:{str,str,[0:{{seq<[0:{},1:{},2:{}]>,[0:{str,str},1:{str}],seq,opt,seq<{str,opt}>,opt<{str,seq,bool}>}},1:{{seq<{str,str,[0:{},1:{},2:{}],bool,seq<{str,str}>,opt,opt}>,{str,seq}}}],opt<{str,str,u32,u32}>} -9:{str,str} -10:{str,seq<{str,u32,str,opt,opt}>} -11:{str,{str,u32},bytes} -12:{str,str} -13:{[0:{}],bytes} -14:{str,str,bytes,u64,bool} -15:{str,str,str} -16:{str,str,str} -17:{str,str,u64,str} -18:{str,[0:{},1:{str},2:{u64,str}]} -19:{str,bytes,u64,bool} -20:{str,seq<{str,bool,opt,opt,opt,opt,opt,opt>,opt>,opt,opt}>} -21:{str,seq} -22:{str,str} -23:{u64,[0:{},1:{u64},2:{str},3:{{str,u64,seq}}]} -24:{u64,[0:{},1:{str},2:{str}]} -25:{str,seq<(str,str)>} -26:{str,str,[0:{},1:{str,str,str},2:{str}]} diff --git a/backend/gradient-wire/schema/v33.txt b/backend/gradient-wire/schema/v33.txt deleted file mode 100644 index 7ae257127..000000000 --- a/backend/gradient-wire/schema/v33.txt +++ /dev/null @@ -1,58 +0,0 @@ -client -0:{{bool,bool,bool,bool,bool,bool},str} -1:{seq<(str,str)>} -2:{} -3:{u16,str} -4:{seq,seq,u32,u32,u64,u32,opt,opt} -5:{f32,u64,opt,opt,opt,opt>} -6:{} -7:{seq<{str,u32,u64,bool}>,bool} -8:{str,bool,opt} -9:{str,str,[0:{},1:{opt},2:{},3:{},4:{seq<{str,str,seq<{str,str}>,seq,seq,str,seq,opt,opt,bool,bool,bool,opt,bool}>,seq,seq},5:{str},6:{str,seq<{str,str,str,opt,opt,seq<{str,str,str,str,opt}>}>,opt<{opt,opt,opt,opt,opt,bool,opt,opt,opt,opt}>,bool},7:{},8:{{u64,u64,u64,u64,u64,u64,u64,u64,str,seq<{str,u64,u64,u64,u64}>,seq<{str,opt,str,str,bool,opt}>}},9:{str,seq<{str,opt,str}>},10:{seq},11:{[0:{},1:{},2:{}]}]} -10:{str,str,seq<{[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{},8:{},9:{},10:{},11:{},12:{},13:{},14:{},15:{},16:{},17:{}],u64,u64,opt,u32,u64}>,u64} -11:{str,str,str,[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{}],seq,seq<{[0:{},1:{},2:{},3:{},4:{},5:{},6:{},7:{},8:{},9:{},10:{},11:{},12:{},13:{},14:{},15:{},16:{},17:{}],u64,u64,opt,u32,u64}>,u64,opt<{opt,opt,opt,opt,opt,bool,opt,opt,opt,opt}>} -12:{} -13:{str,str,str,[0:{},1:{},2:{}],u64,opt,u32,opt} -14:{str,str,[0:{seq<{str,[0:{},1:{},2:{},3:{}],u64,u64}>},1:{u64}]} -15:{str,u32,bytes} -16:{str,seq} -17:{str,str,u64,str} -18:{str,str} -19:{[0:{},1:{}]} -20:{str,opt<{str,u32}>,bytes} -21:{str,str,seq,[0:{},1:{},2:{}],seq>,bool} -22:{str,[0:{},1:{},2:{}],str,str} -23:{str,str,seq} -24:{str,u64,[0:{str},1:{str}],u64} -25:{u64,bytes,u64,bool} -26:{u64,[0:{{str,u64,u64,str,seq,opt,opt,opt<{str,seq}>}},1:{u64}]} -27:{u64} -28:{str,str,seq} -server -0:{seq} -1:{{bool,bool,bool,bool,bool,bool},seq,seq<{str,str}>} -2:{seq,seq<{str,str}>} -3:{u16,str} -4:{u16,str} -5:{} -6:{seq<{str,seq<{str,opt<{u64,u64}>}>,seq,seq,opt<{str,seq}>}>,bool} -7:{seq<{str,seq<{str,opt<{u64,u64}>}>,seq,seq,opt<{str,seq}>}>} -8:{str,str,[0:{{seq<[0:{},1:{},2:{}]>,[0:{str,str},1:{str}],seq,opt,seq<{str,opt}>,opt<{str,seq,bool}>}},1:{{seq<{str,str,[0:{},1:{},2:{}],bool,seq<{str,str}>,opt,opt}>,{str,seq}}}],opt<{str,str,u32,u32}>} -9:{str,str} -10:{str,seq<{str,u32,str,opt,opt}>} -11:{str,{str,u32},bytes} -12:{str,str} -13:{[0:{}],bytes} -14:{str,str,bytes,u64,bool} -15:{str,str,str} -16:{str,str,str} -17:{str,str,u64,str} -18:{str,[0:{},1:{str},2:{u64,str}]} -19:{str,bytes,u64,bool} -20:{str,seq<{str,bool,opt,opt,opt,opt,opt,opt>,opt>,opt,opt}>} -21:{str,seq} -22:{str,str} -23:{u64,[0:{},1:{u64},2:{str},3:{{str,u64,seq}}]} -24:{u64,[0:{},1:{str},2:{str}]} -25:{str,seq<(str,str)>} -26:{str,str,[0:{},1:{str,str,str},2:{str}]} diff --git a/backend/gradient-wire/schema/v34.txt b/backend/gradient-wire/schema/v35.txt similarity index 95% rename from backend/gradient-wire/schema/v34.txt rename to backend/gradient-wire/schema/v35.txt index c91c63387..81232c875 100644 --- a/backend/gradient-wire/schema/v34.txt +++ b/backend/gradient-wire/schema/v35.txt @@ -1,5 +1,5 @@ client -0:{{bool,bool,bool,bool,bool,bool},str} +0:{{bool,bool,bool,bool,bool,bool,bool,bool},str} 1:{seq<(str,str)>} 2:{} 3:{u16,str} @@ -30,7 +30,7 @@ client 28:{str,str,seq} server 0:{seq} -1:{{bool,bool,bool,bool,bool,bool},seq,seq<{str,str}>} +1:{{bool,bool,bool,bool,bool,bool,bool,bool},seq,seq<{str,str}>} 2:{seq,seq<{str,str}>} 3:{u16,str} 4:{u16,str} diff --git a/backend/gradient-wire/src/messages/client.rs b/backend/gradient-wire/src/messages/client.rs index 2eee40d5f..e988ff254 100644 --- a/backend/gradient-wire/src/messages/client.rs +++ b/backend/gradient-wire/src/messages/client.rs @@ -14,7 +14,7 @@ use crate::types::{ }; #[derive(Proto, Debug, Clone, PartialEq)] -#[proto(oldest = 31)] +#[proto(oldest = 35)] pub enum ClientMessage { InitConnection { capabilities: GradientCapabilities, diff --git a/backend/gradient-wire/src/messages/server.rs b/backend/gradient-wire/src/messages/server.rs index ff869a374..47baa25c6 100644 --- a/backend/gradient-wire/src/messages/server.rs +++ b/backend/gradient-wire/src/messages/server.rs @@ -20,7 +20,7 @@ pub struct FailedPeer { } #[derive(Proto, Debug, Clone, PartialEq)] -#[proto(oldest = 31)] +#[proto(oldest = 35)] pub enum ServerMessage { AuthChallenge { peers: Vec, diff --git a/backend/gradient-wire/src/types.rs b/backend/gradient-wire/src/types.rs index 2764bb9e7..13515df54 100644 --- a/backend/gradient-wire/src/types.rs +++ b/backend/gradient-wire/src/types.rs @@ -13,20 +13,24 @@ use serde::{Deserialize as SerdeDeserialize, Serialize as SerdeSerialize}; pub struct GradientCapabilities { pub core: bool, pub federate: bool, - pub fetch: bool, pub eval: bool, pub build: bool, pub cache: bool, + pub network: bool, + pub secrets: bool, + pub exec: bool, } impl std::ops::BitOrAssign for GradientCapabilities { fn bitor_assign(&mut self, rhs: Self) { self.core |= rhs.core; self.federate |= rhs.federate; - self.fetch |= rhs.fetch; self.eval |= rhs.eval; self.build |= rhs.build; self.cache |= rhs.cache; + self.network |= rhs.network; + self.secrets |= rhs.secrets; + self.exec |= rhs.exec; } } diff --git a/backend/gradient-worker-client/src/connection/handshake.rs b/backend/gradient-worker-client/src/connection/handshake.rs index a6c25e3af..7422f408d 100644 --- a/backend/gradient-worker-client/src/connection/handshake.rs +++ b/backend/gradient-worker-client/src/connection/handshake.rs @@ -161,7 +161,9 @@ mod tests { GradientCapabilities { core: false, federate: false, - fetch: true, + network: true, + secrets: true, + exec: true, eval: true, build: true, cache: false, @@ -172,7 +174,9 @@ mod tests { GradientCapabilities { core: false, federate: false, - fetch: false, + network: false, + secrets: false, + exec: false, eval: false, build: false, cache: false, diff --git a/backend/gradient-worker/options.json b/backend/gradient-worker/options.json index aade28092..2522ee8be 100644 --- a/backend/gradient-worker/options.json +++ b/backend/gradient-worker/options.json @@ -200,11 +200,11 @@ }, { "group": "capabilities", - "field": "fetch", - "flag": "capabilities-fetch", - "env": "GRADIENT_WORKER_CAPABILITIES_FETCH", + "field": "eval", + "flag": "capabilities-eval", + "env": "GRADIENT_WORKER_CAPABILITIES_EVAL", "type": "bool", - "description": "Whether to prefetch flake inputs and sources.", + "description": "Whether to run Nix flake evaluations.", "default": "true", "values": [ "true", @@ -213,11 +213,11 @@ }, { "group": "capabilities", - "field": "eval", - "flag": "capabilities-eval", - "env": "GRADIENT_WORKER_CAPABILITIES_EVAL", + "field": "build", + "flag": "capabilities-build", + "env": "GRADIENT_WORKER_CAPABILITIES_BUILD", "type": "bool", - "description": "Whether to run Nix flake evaluations.", + "description": "Whether to run Nix builds.", "default": "true", "values": [ "true", @@ -226,17 +226,43 @@ }, { "group": "capabilities", - "field": "build", - "flag": "capabilities-build", - "env": "GRADIENT_WORKER_CAPABILITIES_BUILD", + "field": "network", + "flag": "capabilities-network", + "env": "GRADIENT_WORKER_CAPABILITIES_NETWORK", "type": "bool", - "description": "Whether to run Nix builds.", + "description": "Whether this worker may reach the internet. Cloning repositories, fetching flake inputs,\nsubstitutions, downloads and fixed-output builds run only on workers with it.", + "default": "true", + "values": [ + "true", + "false" + ] + }, + { + "group": "capabilities", + "field": "secrets", + "flag": "capabilities-secrets", + "env": "GRADIENT_WORKER_CAPABILITIES_SECRETS", + "type": "bool", + "description": "Whether the server may send project credentials such as SSH keys to this worker.", "default": "true", "values": [ "true", "false" ] }, + { + "group": "capabilities", + "field": "exec", + "flag": "capabilities-exec", + "env": "GRADIENT_WORKER_CAPABILITIES_EXEC", + "type": "bool", + "description": "Whether this worker advertises the `exec` capability. The server stores it but does not\nschedule by it yet.", + "default": "false", + "values": [ + "true", + "false" + ] + }, { "group": "system", "field": "architectures", diff --git a/backend/gradient-worker/src/config/mod.rs b/backend/gradient-worker/src/config/mod.rs index d3f86a9d8..a1068b57e 100644 --- a/backend/gradient-worker/src/config/mod.rs +++ b/backend/gradient-worker/src/config/mod.rs @@ -130,10 +130,12 @@ impl WorkerConfig { GradientCapabilities { core: false, federate: self.capabilities.federate, - fetch: self.capabilities.fetch, eval: self.capabilities.eval, build: self.capabilities.build, cache: false, + network: self.capabilities.network, + secrets: self.capabilities.secrets, + exec: self.capabilities.exec, } } } diff --git a/backend/gradient-worker/src/config/options.rs b/backend/gradient-worker/src/config/options.rs index 4db33e091..c46c219dc 100644 --- a/backend/gradient-worker/src/config/options.rs +++ b/backend/gradient-worker/src/config/options.rs @@ -166,10 +166,6 @@ pub struct CapabilitiesArgs { #[arg(default_value = "false")] pub federate: bool, - /// Whether to prefetch flake inputs and sources. - #[arg(default_value = "true")] - pub fetch: bool, - /// Whether to run Nix flake evaluations. #[arg(default_value = "true")] pub eval: bool, @@ -177,6 +173,20 @@ pub struct CapabilitiesArgs { /// Whether to run Nix builds. #[arg(default_value = "true")] pub build: bool, + + /// Whether this worker may reach the internet. Cloning repositories, fetching flake inputs, + /// substitutions, downloads and fixed-output builds run only on workers with it. + #[arg(default_value = "true")] + pub network: bool, + + /// Whether the server may send project credentials such as SSH keys to this worker. + #[arg(default_value = "true")] + pub secrets: bool, + + /// Whether this worker advertises the `exec` capability. The server stores it but does not + /// schedule by it yet. + #[arg(default_value = "false")] + pub exec: bool, } #[config_group("system", env = "worker_system")] diff --git a/nix/modules/gradient-state.nix b/nix/modules/gradient-state.nix index e30d79b7d..d0840d458 100644 --- a/nix/modules/gradient-state.nix +++ b/nix/modules/gradient-state.nix @@ -1026,12 +1026,22 @@ ''; }; - enable_fetch = mkOption { + enable_network = mkOption { type = types.bool; default = true; - description = '' - Whether to grant this registration the worker's `fetch` capability. - ''; + description = "Whether to grant this registration the worker's `network` capability."; + }; + + enable_secrets = mkOption { + type = types.bool; + default = true; + description = "Whether to grant this registration the worker's `secrets` capability."; + }; + + enable_exec = mkOption { + type = types.bool; + default = false; + description = "Whether to grant this registration the worker's `exec` capability."; }; enable_eval = mkOption { diff --git a/nix/standalone/default.nix b/nix/standalone/default.nix index 65aec8835..938e85ac0 100644 --- a/nix/standalone/default.nix +++ b/nix/standalone/default.nix @@ -27,7 +27,7 @@ in { worker = { enable = true; - capabilities = { fetch = true; eval = true; build = true; }; + capabilities = { network = true; secrets = true; eval = true; build = true; }; }; state.users.admin = { From d3a8cc60eb5d4bc4831e8c36cdfe650c45e70fe2 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 14:40:32 +0200 Subject: [PATCH 03/19] scheduler: feat: schedule by network and secrets --- backend/gradient-db/src/teams/demand.rs | 13 +- backend/gradient-pool/src/lib.rs | 2 +- backend/gradient-pool/src/score/policy.rs | 4 +- backend/gradient-pool/src/score/rule.rs | 2 +- .../gradient-pool/src/score/rules/builtin.rs | 8 +- backend/gradient-pool/src/worker_caps.rs | 53 ++++++-- backend/gradient-pool/src/worker_pool.rs | 30 +++-- backend/gradient-proto/src/handler/socket.rs | 60 +++++++-- .../gradient-scheduler/src/assessment_memo.rs | 18 +-- .../gradient-scheduler/src/buildability.rs | 60 ++++++--- .../src/cases/eval_queue.rs | 94 ++++++++++++-- .../gradient-scheduler/src/cases/reports.rs | 8 +- .../gradient-scheduler/src/cases/shapes.rs | 35 ++++-- .../gradient-scheduler/src/cases/workers.rs | 62 +++++++++- .../src/job_handlers/import.rs | 6 +- backend/gradient-scheduler/src/jobs.rs | 45 ++++--- .../gradient-scheduler/src/waiting_state.rs | 115 +++++++++++++----- backend/gradient-types/src/waiting_reason.rs | 8 +- backend/gradient-web/tests/team_demand.rs | 3 +- 19 files changed, 477 insertions(+), 149 deletions(-) diff --git a/backend/gradient-db/src/teams/demand.rs b/backend/gradient-db/src/teams/demand.rs index b7a7b144c..42c5fb505 100644 --- a/backend/gradient-db/src/teams/demand.rs +++ b/backend/gradient-db/src/teams/demand.rs @@ -15,7 +15,8 @@ use std::collections::BTreeMap; #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, utoipa::ToSchema)] pub struct TeamDemand { pub builds: Vec, - pub fetch: u32, + pub network: u32, + pub secrets: u32, pub eval: u32, } @@ -63,7 +64,8 @@ fn aggregate(reasons: impl IntoIterator) -> TeamDemand { } } WaitingReason::EvalWorkers { capability, .. } => match capability { - EvalCapability::Fetch => demand.fetch += 1, + EvalCapability::Network => demand.network += 1, + EvalCapability::Secrets => demand.secrets += 1, EvalCapability::Eval => demand.eval += 1, }, _ => {} @@ -147,13 +149,14 @@ mod tests { } #[tokio::test] - async fn counts_eval_and_fetch_gaps() { + async fn counts_eval_network_and_secrets_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::eval_workers(EvalCapability::Network, 0)), + waiting(WaitingReason::eval_workers(EvalCapability::Secrets, 1)), waiting(WaitingReason::Draining), ]]) .into_connection(); @@ -162,7 +165,7 @@ mod tests { .await .expect("demand"); - assert_eq!((demand.eval, demand.fetch), (2, 1)); + assert_eq!((demand.eval, demand.network, demand.secrets), (2, 1, 1)); assert!(demand.builds.is_empty()); } diff --git a/backend/gradient-pool/src/lib.rs b/backend/gradient-pool/src/lib.rs index 361a65a17..dfb8fb6d8 100644 --- a/backend/gradient-pool/src/lib.rs +++ b/backend/gradient-pool/src/lib.rs @@ -15,6 +15,6 @@ pub mod worker_state; pub use self::aggregate::{Aggregate, aggregate}; pub use self::peer_auth::PeerAuth; -pub use self::worker_caps::WorkerCaps; +pub use self::worker_caps::{WorkerCaps, build_needs_network}; pub use self::worker_pool::{WorkerInfo, WorkerPool, WorkerSlot}; pub use self::worker_state::{WorkerProfile, WorkerShared}; diff --git a/backend/gradient-pool/src/score/policy.rs b/backend/gradient-pool/src/score/policy.rs index 875ee97b0..8ca15fcb1 100644 --- a/backend/gradient-pool/src/score/policy.rs +++ b/backend/gradient-pool/src/score/policy.rs @@ -303,7 +303,7 @@ mod tests { let with_cores = |cpu_core_score| WorkerContext { architectures: &archs, system_features: &feats, - fetch: false, + fetches_flakes: false, metrics: Some(WorkerMetricsView { cpu_core_score, ..Default::default() @@ -349,7 +349,7 @@ mod tests { let with_cores = |cpu_core_score, ram_free_mb| WorkerContext { architectures: &archs, system_features: &feats, - fetch: false, + fetches_flakes: false, metrics: Some(WorkerMetricsView { cpu_core_score, ram_total_mb: 128_000, diff --git a/backend/gradient-pool/src/score/rule.rs b/backend/gradient-pool/src/score/rule.rs index 6b75fa678..a32cc6279 100644 --- a/backend/gradient-pool/src/score/rule.rs +++ b/backend/gradient-pool/src/score/rule.rs @@ -83,7 +83,7 @@ impl<'a> JobContext<'a> { pub struct WorkerContext<'a> { pub architectures: &'a [String], pub system_features: &'a [String], - pub fetch: bool, + pub fetches_flakes: bool, pub metrics: Option, } diff --git a/backend/gradient-pool/src/score/rules/builtin.rs b/backend/gradient-pool/src/score/rules/builtin.rs index 4effe415b..a319f09c2 100644 --- a/backend/gradient-pool/src/score/rules/builtin.rs +++ b/backend/gradient-pool/src/score/rules/builtin.rs @@ -74,7 +74,7 @@ impl ScoreRule for ReserveFetchWorkersRule { instance: &InstanceContext, ) -> f64 { let cached_eval = matches!(job.job.kind(), JobKindContext::Eval(e) if !e.fetch_flake); - if !(worker.fetch && cached_eval) { + if !(worker.fetches_flakes && cached_eval) { return 0.0; } @@ -88,7 +88,7 @@ impl ScoreRule for ReserveFetchWorkersRule { } fn description(&self) -> &'static str { - "Keeps fetch-capable workers free for flake fetches by penalizing them for cached evaluations, easing off as idle workers appear." + "Keeps workers that can fetch flakes free for flake fetches by penalizing them for cached evaluations, easing off as idle workers appear." } } @@ -159,10 +159,10 @@ mod tests { ScoredJob::new_eval("test", ProjectId::now_v7(), fetch_flake, &NO_ESTIMATE) } - fn worker<'a>(archs: &'a [String], fetch: bool) -> WorkerContext<'a> { + fn worker<'a>(archs: &'a [String], fetches_flakes: bool) -> WorkerContext<'a> { WorkerContext { architectures: archs, - fetch, + fetches_flakes, ..Default::default() } } diff --git a/backend/gradient-pool/src/worker_caps.rs b/backend/gradient-pool/src/worker_caps.rs index ea651188e..ba24b3dd3 100644 --- a/backend/gradient-pool/src/worker_caps.rs +++ b/backend/gradient-pool/src/worker_caps.rs @@ -4,13 +4,15 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -use gradient_wire::types::{FlakeJob, FlakeStep, JobKind}; +use gradient_types::input::check_repository_url_is_ssh; +use gradient_wire::types::{BuildJob, BuildSpecKind, FlakeJob, FlakeSource, FlakeStep, JobKind}; + +pub fn build_needs_network(kind: BuildSpecKind, is_fixed_output: bool) -> bool { + kind != BuildSpecKind::Build || is_fixed_output +} #[derive(Debug, Clone, Default)] pub struct WorkerCaps { - /// The server is sending SSH credentials only to fetch-capable workers. Any `FetchFlake` step - /// is therefore requiring this flag. - pub fetch: bool, pub architectures: Vec, pub system_features: Vec, pub capabilities: gradient_wire::types::GradientCapabilities, @@ -22,7 +24,7 @@ pub struct WorkerCaps { impl WorkerCaps { pub fn takes_jobs_of(&self, kind: &JobKind) -> bool { match kind { - JobKind::Flake => self.capabilities.eval || self.fetch, + JobKind::Flake => self.capabilities.eval || self.capabilities.network, JobKind::Build => self.capabilities.build, } } @@ -36,14 +38,30 @@ impl WorkerCaps { self.takes_jobs_of(&JobKind::Build) && arch_ok && features_ok } + pub fn can_run_build(&self, job: &BuildJob) -> bool { + let needs_network = job + .builds + .iter() + .any(|b| build_needs_network(b.kind, b.is_fixed_output)); + self.can_build( + &job.requirement.architecture, + &job.requirement.required_features, + ) && (!needs_network || self.capabilities.network) + } + pub fn can_eval(&self, job: &FlakeJob) -> bool { - let needs_fetch = job.steps.contains(&FlakeStep::FetchFlake); let needs_eval = job .steps .iter() .any(|t| matches!(t, FlakeStep::EvaluateFlake | FlakeStep::EvaluateDerivations)); - (!needs_fetch || self.fetch) && (!needs_eval || self.capabilities.eval) + (!job.steps.contains(&FlakeStep::FetchFlake) || self.can_fetch(&job.source)) + && (!needs_eval || self.capabilities.eval) + } + + pub fn can_fetch(&self, source: &FlakeSource) -> bool { + let needs_secrets = matches!(source, FlakeSource::Repository { url, .. } if check_repository_url_is_ssh(url)); + self.capabilities.network && (!needs_secrets || self.capabilities.secrets) } } @@ -74,4 +92,25 @@ mod tests { assert!(builder.can_build(gradient_types::BUILTIN_ARCH, &[])); assert!(builder.can_build("x86_64-linux", &[])); } + + fn repository(url: &str) -> FlakeSource { + FlakeSource::Repository { + url: url.into(), + commit: "abc123".into(), + } + } + + #[test] + fn an_ssh_repository_needs_secrets_and_a_public_one_only_network() { + let network_only = WorkerCaps { + capabilities: GradientCapabilities { + network: true, + ..Default::default() + }, + ..Default::default() + }; + assert!(network_only.can_fetch(&repository("https://example.com/repo"))); + assert!(!network_only.can_fetch(&repository("git+ssh://git@example.com/repo"))); + assert!(!WorkerCaps::default().can_fetch(&repository("https://example.com/repo"))); + } } diff --git a/backend/gradient-pool/src/worker_pool.rs b/backend/gradient-pool/src/worker_pool.rs index 776f29822..8a3e7e84f 100644 --- a/backend/gradient-pool/src/worker_pool.rs +++ b/backend/gradient-pool/src/worker_pool.rs @@ -142,7 +142,6 @@ impl WorkerPool { self.workers.get(id).map(|slot| { let s = slot.shared(); crate::WorkerCaps { - fetch: s.capabilities.network && s.capabilities.secrets, architectures: s.architectures.clone(), system_features: s.system_features.clone(), capabilities: s.capabilities.clone(), @@ -297,9 +296,7 @@ impl WorkerPool { pub fn has_idle_eval_only_worker(&self) -> bool { self.workers.values().any(|slot| match slot { WorkerSlot::Active(w) => { - w.capabilities.eval - && !(w.capabilities.network && w.capabilities.secrets) - && w.assigned_jobs.is_empty() + w.capabilities.eval && !w.capabilities.network && w.assigned_jobs.is_empty() } WorkerSlot::Draining(_) => false, }) @@ -476,11 +473,10 @@ mod tests { (Arc::new(tx), rx) } - fn caps_ef(eval: bool, fetch: bool) -> GradientCapabilities { + fn eval_caps(eval: bool, network: bool) -> GradientCapabilities { GradientCapabilities { eval, - network: fetch, - secrets: fetch, + network, ..GradientCapabilities::default() } } @@ -488,13 +484,18 @@ mod tests { #[test] fn idle_eval_only_worker_detected() { let mut pool = WorkerPool::new(); - pool.register("f1".into(), caps_ef(true, true), HashSet::new(), port().0); + pool.register("f1".into(), eval_caps(true, true), HashSet::new(), port().0); assert!( !pool.has_idle_eval_only_worker(), - "only a fetch worker present" + "only a network worker present" ); - pool.register("e1".into(), caps_ef(true, false), HashSet::new(), port().0); + pool.register( + "e1".into(), + eval_caps(true, false), + HashSet::new(), + port().0, + ); assert!( pool.has_idle_eval_only_worker(), "idle eval-only worker present" @@ -510,7 +511,12 @@ mod tests { #[test] fn draining_eval_only_worker_does_not_count() { let mut pool = WorkerPool::new(); - pool.register("e1".into(), caps_ef(true, false), HashSet::new(), port().0); + pool.register( + "e1".into(), + eval_caps(true, false), + HashSet::new(), + port().0, + ); pool.mark_draining("e1"); assert!( !pool.has_idle_eval_only_worker(), @@ -706,7 +712,7 @@ mod tests { pool.update_metrics("w1", 12.5, 9000, None, None, None); let caps = pool.worker_caps("w1").unwrap(); - assert!(caps.fetch); + assert!(caps.capabilities.network); assert!(caps.capabilities.eval); assert_eq!(caps.architectures, vec!["x86_64-linux"]); assert_eq!(caps.system_features, vec!["kvm"]); diff --git a/backend/gradient-proto/src/handler/socket.rs b/backend/gradient-proto/src/handler/socket.rs index b76d2b5eb..529ea0ed1 100644 --- a/backend/gradient-proto/src/handler/socket.rs +++ b/backend/gradient-proto/src/handler/socket.rs @@ -48,21 +48,22 @@ pub(super) async fn send_credentials_for_job( job: &gradient_wire::types::Job, project_id: ProjectId, ) { - use gradient_wire::types::{FlakeStep, Job}; - let caps = scheduler.worker_gradient_caps(worker_id).await; - let worker_can_fetch = caps.as_ref().is_some_and(|c| c.network && c.secrets); - - match job { - Job::Flake(flake_job) => { - if worker_can_fetch && flake_job.steps.contains(&FlakeStep::FetchFlake) { - send_ssh_key_credential(writer, db, project_id).await; - } - } - Job::Build(_) => {} + if sends_ssh_key(caps.as_ref(), job) { + send_ssh_key_credential(writer, db, project_id).await; } } +fn sends_ssh_key( + caps: Option<&gradient_wire::types::GradientCapabilities>, + job: &gradient_wire::types::Job, +) -> bool { + use gradient_wire::types::{FlakeStep, Job}; + + caps.is_some_and(|c| c.secrets) + && matches!(job, Job::Flake(flake_job) if flake_job.steps.contains(&FlakeStep::FetchFlake)) +} + async fn send_ssh_key_credential(writer: &ProtoWriter, db: &DbContext, project_id: ProjectId) { use gradient_wire::types::CredentialKind; @@ -93,3 +94,40 @@ async fn send_ssh_key_credential(writer: &ProtoWriter, db: &DbContext, project_i Err(e) => warn!(%project_id, error = %e, "failed to fetch project for SSH key"), } } + +#[cfg(test)] +mod tests { + use super::*; + use gradient_wire::types::{FlakeJob, FlakeSource, FlakeStep, GradientCapabilities, Job}; + + fn fetch_job() -> Job { + Job::Flake(FlakeJob { + steps: vec![FlakeStep::FetchFlake], + source: FlakeSource::Repository { + url: "git+ssh://git@example.com/repo".into(), + commit: "abc123".into(), + }, + wildcards: vec![], + timeout_secs: None, + input_overrides: vec![], + input_update: None, + }) + } + + #[test] + fn no_ssh_key_without_secrets() { + let network_only = GradientCapabilities { + network: true, + eval: true, + ..Default::default() + }; + assert!(!sends_ssh_key(Some(&network_only), &fetch_job())); + assert!(!sends_ssh_key(None, &fetch_job())); + + let with_secrets = GradientCapabilities { + secrets: true, + ..network_only + }; + assert!(sends_ssh_key(Some(&with_secrets), &fetch_job())); + } +} diff --git a/backend/gradient-scheduler/src/assessment_memo.rs b/backend/gradient-scheduler/src/assessment_memo.rs index 58aa1938b..e7a7e064d 100644 --- a/backend/gradient-scheduler/src/assessment_memo.rs +++ b/backend/gradient-scheduler/src/assessment_memo.rs @@ -17,7 +17,9 @@ use gradient_db::evaluations::counters::EvalCounters; use gradient_entity::evaluation::EvaluationStatus; use gradient_types::*; -pub(crate) type Caps = [(Vec, Vec)]; +use crate::buildability::BuildWorkerCaps; + +pub(crate) type Caps = [BuildWorkerCaps]; struct Entry { key: (EvalCounters, u64), @@ -73,14 +75,14 @@ impl AssessmentMemo { } pub(crate) fn caps_fingerprint(caps: &Caps) -> u64 { - let mut workers: Vec<(Vec<&str>, Vec<&str>)> = caps + let mut workers: Vec<(Vec<&str>, Vec<&str>, bool)> = caps .iter() - .map(|(archs, feats)| { + .map(|(archs, feats, network)| { let mut archs: Vec<&str> = archs.iter().map(String::as_str).collect(); let mut feats: Vec<&str> = feats.iter().map(String::as_str).collect(); archs.sort_unstable(); feats.sort_unstable(); - (archs, feats) + (archs, feats, *network) }) .collect(); workers.sort_unstable(); @@ -94,8 +96,8 @@ pub(crate) fn caps_fingerprint(caps: &Caps) -> u64 { mod tests { use super::*; - fn caps() -> Vec<(Vec, Vec)> { - vec![(vec!["x86_64-linux".to_owned()], vec![])] + fn caps() -> Vec { + vec![(vec!["x86_64-linux".to_owned()], vec![], true)] } #[test] @@ -133,8 +135,8 @@ mod tests { #[test] fn the_fingerprint_ignores_worker_order() { let a = vec![ - (vec!["a".to_owned()], vec!["kvm".to_owned()]), - (vec!["b".to_owned()], vec![]), + (vec!["a".to_owned()], vec!["kvm".to_owned()], true), + (vec!["b".to_owned()], vec![], false), ]; let b = vec![a[1].clone(), a[0].clone()]; assert_eq!(caps_fingerprint(&a), caps_fingerprint(&b)); diff --git a/backend/gradient-scheduler/src/buildability.rs b/backend/gradient-scheduler/src/buildability.rs index 483d433f7..5c24d2dc6 100644 --- a/backend/gradient-scheduler/src/buildability.rs +++ b/backend/gradient-scheduler/src/buildability.rs @@ -16,8 +16,11 @@ use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; use crate::assign_mode::decide_build_spec_kind; use crate::count::clamp_to_u32; +use gradient_pool::build_needs_network; use gradient_wire::types::BuildSpecKind; +pub(crate) type BuildWorkerCaps = (Vec, Vec, bool); + pub(crate) struct BuildabilityChecker { drv_by_id: HashMap, features_by_drv: HashMap>, @@ -69,7 +72,7 @@ impl BuildabilityChecker { pub(crate) fn any_buildable( &self, shared_builds: &[MDerivationBuild], - worker_caps: &[(Vec, Vec)], + worker_caps: &[BuildWorkerCaps], ) -> bool { shared_builds.iter().any(|a| { if a.status == BuildStatus::Building { @@ -88,17 +91,22 @@ impl BuildabilityChecker { &self, build: &MDerivationBuild, drv: &MDerivation, - worker_caps: &[(Vec, Vec)], + worker_caps: &[BuildWorkerCaps], ) -> bool { - match decide_build_spec_kind( + let kind = decide_build_spec_kind( build.cache_available, &drv.architecture, drv.is_fixed_output, - ) { - BuildSpecKind::Substitute | BuildSpecKind::Download => !worker_caps.is_empty(), + ); + let needs_network = build_needs_network(kind, drv.is_fixed_output); + let mut able = worker_caps + .iter() + .filter(|(_, _, network)| *network || !needs_network); + match kind { + BuildSpecKind::Substitute | BuildSpecKind::Download => able.next().is_some(), BuildSpecKind::Build => { let required = self.required_features_for(&build.derivation); - worker_caps.iter().any(|(arch, feats)| { + able.any(|(arch, feats, _)| { let arch_ok = drv.architecture == gradient_types::BUILTIN_ARCH || arch.iter().any(|a| a == &drv.architecture); let feats_ok = required.iter().all(|f| feats.iter().any(|sf| sf == f)); @@ -126,7 +134,7 @@ impl BuildabilityChecker { pub(crate) fn unbuildable_imports( &self, imports: &[MDerivationBuild], - worker_caps: &[(Vec, Vec)], + worker_caps: &[BuildWorkerCaps], ) -> Vec<(DerivationBuildId, UnmetRequirement)> { imports .iter() @@ -154,7 +162,7 @@ impl BuildabilityChecker { pub(crate) fn compute_waiting_reason( &self, shared_builds: &[MDerivationBuild], - worker_caps: &[(Vec, Vec)], + worker_caps: &[BuildWorkerCaps], ) -> WaitingReason { let mut grouped: BTreeMap<(String, Vec), u32> = BTreeMap::new(); for a in shared_builds { @@ -187,7 +195,7 @@ impl BuildabilityChecker { let mut available_architectures: Vec = worker_caps .iter() - .flat_map(|(archs, _)| archs.iter().cloned()) + .flat_map(|(archs, _, _)| archs.iter().cloned()) .collect(); available_architectures.sort_unstable(); available_architectures.dedup(); @@ -287,7 +295,7 @@ mod tests { let builds = vec![build_for(d_x86.id, eval_id), build_for(d_arm.id, eval_id)]; let checker = checker_with(vec![d_x86, d_arm], vec![]); - let caps: Vec<(Vec, Vec)> = vec![(vec!["x86_64-linux".into()], vec![])]; + let caps: Vec = vec![(vec!["x86_64-linux".into()], vec![], true)]; let reason = checker.compute_waiting_reason(&builds, &caps); let (unmet, connected_workers, available_architectures) = workers_view(&reason); @@ -307,7 +315,7 @@ mod tests { let builds = vec![build_for(drv_id, eval_id)]; let checker = checker_with(vec![d], vec![(drv_id, feat_id, "kvm")]); - let caps: Vec<(Vec, Vec)> = vec![(vec!["x86_64-linux".into()], vec![])]; + let caps: Vec = vec![(vec!["x86_64-linux".into()], vec![], true)]; let reason = checker.compute_waiting_reason(&builds, &caps); let (unmet, _, _) = workers_view(&reason); @@ -345,7 +353,7 @@ mod tests { let builds = vec![build_for(d.id, eval_id)]; let checker = checker_with(vec![d], vec![]); - let caps: Vec<(Vec, Vec)> = vec![(vec!["x86_64-linux".into()], vec![])]; + let caps: Vec = vec![(vec!["x86_64-linux".into()], vec![], true)]; let reason = checker.compute_waiting_reason(&builds, &caps); let (unmet, _, _) = workers_view(&reason); @@ -369,7 +377,7 @@ mod tests { let build = cache_available_build(d.id, eval_id); let checker = checker_with(vec![d], vec![]); - let caps: Vec<(Vec, Vec)> = vec![(vec!["x86_64-linux".into()], vec![])]; + let caps: Vec = vec![(vec!["x86_64-linux".into()], vec![], true)]; let builds = [build]; assert!(checker.any_buildable(&builds, &caps)); let reason = checker.compute_waiting_reason(&builds, &caps); @@ -405,7 +413,7 @@ mod tests { build.cache_available = false; let checker = checker_with(vec![d], vec![]); - let caps: Vec<(Vec, Vec)> = vec![(vec!["x86_64-linux".into()], vec![])]; + let caps: Vec = vec![(vec!["x86_64-linux".into()], vec![], true)]; let builds = [build]; assert!(!checker.any_buildable(&builds, &caps)); let reason = checker.compute_waiting_reason(&builds, &caps); @@ -414,6 +422,17 @@ mod tests { assert_eq!(unmet[0].architecture, "aarch64-linux"); } + #[test] + fn a_fixed_output_build_needs_a_network_worker() { + let mut d = drv(DerivationId::now_v7(), "x86_64-linux"); + d.is_fixed_output = true; + let builds = [build_for(d.id, EvaluationId::now_v7())]; + let checker = checker_with(vec![d], vec![]); + + assert!(!checker.any_buildable(&builds, &[(vec!["x86_64-linux".into()], vec![], false)])); + assert!(checker.any_buildable(&builds, &[(vec!["x86_64-linux".into()], vec![], true)])); + } + #[test] fn dependency_blocked_shared_build_is_not_buildable() { let eval_id = EvaluationId::now_v7(); @@ -421,7 +440,7 @@ mod tests { let mut b = build_for(d.id, eval_id); b.status = BuildStatus::Created; let checker = checker_with(vec![d], vec![]); - let caps = vec![(vec!["x86_64-linux".to_string()], vec![])]; + let caps = vec![(vec!["x86_64-linux".to_string()], vec![], true)]; assert!(!checker.any_buildable(&[b], &caps)); } @@ -432,12 +451,13 @@ mod tests { let b = build_for(d.id, eval_id); let checker = checker_with(vec![d], vec![]); let shared_builds = std::slice::from_ref(&b); - assert!( - checker.any_buildable(shared_builds, &[(vec!["x86_64-linux".to_string()], vec![])]) - ); + assert!(checker.any_buildable( + shared_builds, + &[(vec!["x86_64-linux".to_string()], vec![], true)] + )); assert!(!checker.any_buildable( shared_builds, - &[(vec!["aarch64-linux".to_string()], vec![])] + &[(vec!["aarch64-linux".to_string()], vec![], true)] )); } @@ -455,7 +475,7 @@ mod tests { pending(linux.id, BuildStatus::Queued), ); let checker = checker_with(vec![darwin, linux], vec![]); - let caps = vec![(vec!["x86_64-linux".to_string()], vec![])]; + let caps = vec![(vec!["x86_64-linux".to_string()], vec![], true)]; let unbuildable = checker.unbuildable_imports(&[stuck.clone(), building, buildable], &caps); diff --git a/backend/gradient-scheduler/src/cases/eval_queue.rs b/backend/gradient-scheduler/src/cases/eval_queue.rs index 096eb6e9b..93c8db0e0 100644 --- a/backend/gradient-scheduler/src/cases/eval_queue.rs +++ b/backend/gradient-scheduler/src/cases/eval_queue.rs @@ -22,9 +22,12 @@ use gradient_wire::types::{ use super::{ A, A_OUT, ARM, B, B_OUT, C, C_OUT, CaseShapes, REPOSITORY, S_OUT, SchedulerCase, X86, ago, - attr_of, drv_nar, eval_key, store_path, w_all, w_build, w_eval, w_fetch, w_fetch_eval, + attr_of, drv_nar, eval_key, store_path, w_all, w_build, w_eval, w_network, w_network_eval, + w_network_secrets, }; +const SSH_REPOSITORY: &str = "git+ssh://git@example.invalid/repo.git"; + const W1: &str = "w1"; const W2: &str = "w2"; @@ -163,7 +166,7 @@ async fn q04_an_idle_evaluation_only_worker_splits_the_pass_and_without_one_no_s EvaluationId::now_v7(), ); let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; - case.connect(W1, w_fetch_eval(), &[]).await; + case.connect(W1, w_network_eval(), &[]).await; case.connect(W2, w_eval(), &[]).await; case.assign_queued_evals().await; @@ -213,7 +216,7 @@ async fn q05_an_evaluation_only_worker_running_a_job_causes_no_split() { EvaluationId::now_v7(), ); let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; - case.connect(W1, w_fetch_eval(), &[]).await; + case.connect(W1, w_network_eval(), &[]).await; case.connect(W2, w_eval(), &[]).await; let key = eval_key(e1); @@ -327,10 +330,10 @@ async fn q07_an_input_update_evaluation_is_queued_on_its_base_commit_with_the_up } #[tokio::test] -async fn q08_a_fetch_step_needs_fetch_and_an_evaluation_step_needs_eval() { +async fn q08_a_fetch_step_needs_network_and_an_evaluation_step_needs_eval() { let (e1, e2) = (EvaluationId::now_v7(), EvaluationId::now_v7()); let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; - case.connect(W1, w_fetch(), &[]).await; + case.connect(W1, w_network(), &[]).await; case.connect(W2, w_eval(), &[]).await; let key = eval_key(e1); @@ -419,7 +422,7 @@ async fn q09_a_queued_evaluation_without_an_eval_worker_waits_until_one_connects } #[tokio::test] -async fn q10_a_fetching_evaluation_without_a_fetch_worker_waits_on_fetch() { +async fn q10_a_fetching_evaluation_without_a_network_worker_waits_on_network() { let e1 = EvaluationId::now_v7(); let case = SchedulerCase::start( shape() @@ -434,13 +437,13 @@ async fn q10_a_fetching_evaluation_without_a_fetch_worker_waits_on_fetch() { waiting_state(&case, e1).await, ( EvaluationStatus::Waiting, - Some(WaitingReason::eval_workers(EvalCapability::Fetch, 1)), + Some(WaitingReason::eval_workers(EvalCapability::Network, 1)), ), "after 1: evaluation(E1)" ); case.check("step 1").await; - case.connect(W1, w_fetch_eval(), &[]).await; + case.connect(W1, w_network_eval(), &[]).await; case.refresh_waiting_state().await; assert_eq!( waiting_state(&case, e1).await, @@ -883,3 +886,78 @@ async fn recoveries_of( .filter(|e| e.kind == EvaluationKind::DrvRecovery && e.previous == Some(evaluation)) .collect() } + +#[tokio::test] +async fn q18_a_fetch_step_of_an_ssh_repository_needs_secrets() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start(shape().evalq(e1).repository(SSH_REPOSITORY)).await; + case.connect(W1, w_network(), &[]).await; + case.connect(W2, w_network_secrets(), &[]).await; + let key = eval_key(e1); + + case.assign_queued_evals().await; + case.check("step 1").await; + + let without_secrets = case.offers_full(W1).await; + let with_secrets = case.offers_full(W2).await; + assert_eq!( + candidate_keys(&without_secrets), + BTreeSet::new(), + "after 2: offers(W1).candidates" + ); + assert_eq!( + candidate_keys(&with_secrets), + BTreeSet::from([key.clone()]), + "after 2: offers(W2).candidates" + ); + case.check("step 2").await; + + case.score_all(W2).await; + let fetch = case + .request_job(W2, JobKind::Flake) + .await + .expect("after 3: W2 gets the fetch job of E1"); + assert_eq!(fetch.job_id(), key, "after 3: the fetch job of E1"); + case.report_flake_source(W2, &key, &store_path(S_OUT, "s")) + .await; + case.complete(W2, &key, fetch.assignment_id(), Vec::new(), 1) + .await; + case.check("step 3").await; + + case.finish().await; +} + +#[tokio::test] +async fn q19_a_fetching_evaluation_of_an_ssh_repository_waits_on_secrets() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start( + shape() + .evalq(e1) + .repository(SSH_REPOSITORY) + .evaluation_status(EvaluationStatus::Fetching), + ) + .await; + case.connect(W1, w_network_eval(), &[]).await; + + case.refresh_waiting_state().await; + assert_eq!( + waiting_state(&case, e1).await, + ( + EvaluationStatus::Waiting, + Some(WaitingReason::eval_workers(EvalCapability::Secrets, 1)), + ), + "after 1: evaluation(E1)" + ); + case.check("step 1").await; + + case.connect(W2, w_network_secrets(), &[]).await; + case.refresh_waiting_state().await; + assert_eq!( + waiting_state(&case, e1).await, + (EvaluationStatus::Queued, None), + "after 2: evaluation(E1)" + ); + case.check("step 2").await; + + case.finish().await; +} diff --git a/backend/gradient-scheduler/src/cases/reports.rs b/backend/gradient-scheduler/src/cases/reports.rs index 66a828f58..8a6a512b1 100644 --- a/backend/gradient-scheduler/src/cases/reports.rs +++ b/backend/gradient-scheduler/src/cases/reports.rs @@ -23,8 +23,8 @@ use gradient_wire::types::{ use super::{ A, A_OUT, B, B_OUT, C, C_OUT, CaseShapes, SchedulerCase, ago, attr_of, build_output, drv, - drv_nar, eval_key, failure, is_recent, job_ids, output_nar, span, w_all, w_eval, w_fetch_eval, - with_attr, with_output, + drv_nar, eval_key, failure, is_recent, job_ids, output_nar, span, w_all, w_eval, + w_network_eval, with_attr, with_output, }; use crate::jobs::PendingJob; @@ -1077,7 +1077,7 @@ async fn r17_a_completed_fetch_job_enqueues_the_cached_follow_up_for_an_eval_wor let e1 = EvaluationId::now_v7(); let case = SchedulerCase::start(shape().evalq(e1)).await; let k = eval_key(e1); - case.connect(W1, w_fetch_eval(), &[]).await; + case.connect(W1, w_network_eval(), &[]).await; case.connect(W2, w_eval(), &[]).await; case.assign_queued_evals().await; @@ -1221,7 +1221,7 @@ async fn r18_a_completed_fetch_job_without_a_flake_source_fails_the_evaluation() let e1 = EvaluationId::now_v7(); let case = SchedulerCase::start(shape().evalq(e1)).await; let k = eval_key(e1); - case.connect(W1, w_fetch_eval(), &[]).await; + case.connect(W1, w_network_eval(), &[]).await; case.connect(W2, w_eval(), &[]).await; case.assign_queued_evals().await; diff --git a/backend/gradient-scheduler/src/cases/shapes.rs b/backend/gradient-scheduler/src/cases/shapes.rs index 976a99eee..e61b1fdb3 100644 --- a/backend/gradient-scheduler/src/cases/shapes.rs +++ b/backend/gradient-scheduler/src/cases/shapes.rs @@ -97,10 +97,10 @@ impl Profile { } } -fn caps(fetch: bool, eval: bool, build: bool) -> GradientCapabilities { +fn caps(network: bool, secrets: bool, eval: bool, build: bool) -> GradientCapabilities { GradientCapabilities { - network: fetch, - secrets: fetch, + network, + secrets, eval, build, ..GradientCapabilities::default() @@ -109,7 +109,7 @@ fn caps(fetch: bool, eval: bool, build: bool) -> GradientCapabilities { pub(crate) fn w_all() -> Profile { Profile { - caps: caps(true, true, true), + caps: caps(true, true, true, true), hardware: None, } .slots(1) @@ -117,7 +117,15 @@ pub(crate) fn w_all() -> Profile { pub(crate) fn w_build() -> Profile { Profile { - caps: caps(false, false, true), + caps: caps(false, false, false, true), + hardware: None, + } + .slots(1) +} + +pub(crate) fn w_network_build() -> Profile { + Profile { + caps: caps(true, false, false, true), hardware: None, } .slots(1) @@ -125,21 +133,28 @@ pub(crate) fn w_build() -> Profile { pub(crate) fn w_eval() -> Profile { Profile { - caps: caps(false, true, false), + caps: caps(false, false, true, false), + hardware: None, + } +} + +pub(crate) fn w_network_eval() -> Profile { + Profile { + caps: caps(true, false, true, false), hardware: None, } } -pub(crate) fn w_fetch_eval() -> Profile { +pub(crate) fn w_network() -> Profile { Profile { - caps: caps(true, true, false), + caps: caps(true, false, false, false), hardware: None, } } -pub(crate) fn w_fetch() -> Profile { +pub(crate) fn w_network_secrets() -> Profile { Profile { - caps: caps(true, false, false), + caps: caps(true, true, false, false), hardware: None, } } diff --git a/backend/gradient-scheduler/src/cases/workers.rs b/backend/gradient-scheduler/src/cases/workers.rs index 142dcec50..a7b2d49c3 100644 --- a/backend/gradient-scheduler/src/cases/workers.rs +++ b/backend/gradient-scheduler/src/cases/workers.rs @@ -21,7 +21,7 @@ use gradient_wire::types::{BuildFailureKind, FlakeStep, Job, JobKind}; use super::{ A, A_OUT, B, B_OUT, CaseShapes, SchedulerCase, SlotKind, Totals, ago, eval_key, failure, - is_recent, output_nar, w_all, w_build, w_eval, w_fetch_eval, + is_recent, output_nar, w_all, w_build, w_eval, w_network_build, w_network_eval, }; use crate::actor::Offer; @@ -658,7 +658,7 @@ async fn w11_a_worker_without_build_gets_no_download_job() { .system("builtin") .fixed_output(); let case = SchedulerCase::start(seed).await; - case.connect(W1, w_fetch_eval(), &[]).await; + case.connect(W1, w_network_eval(), &[]).await; let a = case.key(A).await; case.admit_startable_moves().await; @@ -1370,6 +1370,64 @@ async fn w22_a_reauth_result_without_a_granted_project_aborts_no_running_job() { case.finish().await; } +#[tokio::test] +async fn w23_a_fixed_output_build_needs_network() { + only_a_network_worker_gets(shape().one(EvaluationId::now_v7(), A, A_OUT).fixed_output()).await; +} + +#[tokio::test] +async fn w24_a_substitute_needs_network() { + only_a_network_worker_gets(shape().passthrough(EvaluationId::now_v7())).await; +} + +#[tokio::test] +async fn w25_a_download_needs_network() { + only_a_network_worker_gets( + shape() + .one(EvaluationId::now_v7(), A, A_OUT) + .system("builtin") + .fixed_output(), + ) + .await; +} + +async fn only_a_network_worker_gets(seed: Shape) { + let case = SchedulerCase::start(seed).await; + case.connect(W1, w_build(), &[]).await; + case.connect(W2, w_network_build(), &[]).await; + let a = case.key(A).await; + + case.admit_startable_moves().await; + let offline = case.offers_full(W1).await; + let online = case.offers_full(W2).await; + assert_eq!( + offered(&offline), + Vec::::new(), + "after 1: offers(W1).candidates" + ); + assert_eq!( + offered(&online), + vec![a.clone()], + "after 1: offers(W2).candidates" + ); + case.check("step 1").await; + + let nothing = case.request_job(W1, JobKind::Build).await; + assert!(nothing.is_none(), "after 2: W1 gets nothing"); + assert!(case.pending(&a).await, "after 2: pending(build:A)"); + case.check("step 2").await; + + case.score_all(W2).await; + let job = case + .request_job(W2, JobKind::Build) + .await + .expect("after 3: W2 gets build:A"); + assert_eq!(job.job_id(), a, "after 3: result.job_id"); + case.check("step 3").await; + + case.finish().await; +} + fn drain_signals(case: &SchedulerCase, workers: &[&str]) { for worker in workers { case.signals(worker); diff --git a/backend/gradient-scheduler/src/job_handlers/import.rs b/backend/gradient-scheduler/src/job_handlers/import.rs index f845897c8..149e2bd91 100644 --- a/backend/gradient-scheduler/src/job_handlers/import.rs +++ b/backend/gradient-scheduler/src/job_handlers/import.rs @@ -19,7 +19,7 @@ use tracing::{debug, warn}; use crate::Scheduler; use crate::actor::SchedulerMsg; -use crate::buildability::BuildabilityChecker; +use crate::buildability::{BuildWorkerCaps, BuildabilityChecker}; use crate::import_waits::{ImportAnswer, ImportWait, ImportedBuild}; pub(crate) const NO_ENTRY_POINT: &str = "no entry point recorded for the import"; @@ -210,11 +210,11 @@ impl Scheduler { return Ok(Vec::new()); } - let worker_caps: Vec<(Vec, Vec)> = self + let worker_caps: Vec = self .board_workers() .await .into_iter() - .map(|w| (w.architectures, w.system_features)) + .map(|w| (w.architectures, w.system_features, w.capabilities.network)) .collect(); let checker = BuildabilityChecker::load(&self.db.worker_db, &candidates).await?; let unbuildable: HashMap = checker diff --git a/backend/gradient-scheduler/src/jobs.rs b/backend/gradient-scheduler/src/jobs.rs index 08a2bd801..22c7b66aa 100644 --- a/backend/gradient-scheduler/src/jobs.rs +++ b/backend/gradient-scheduler/src/jobs.rs @@ -373,10 +373,7 @@ fn job_eligible_for_caps(job: &PendingJob, caps: Option<&WorkerCaps>) -> bool { match (job, caps) { (_, None) => true, (PendingJob::Eval(j), Some(c)) => c.can_eval(&j.job), - (PendingJob::Build(j), Some(c)) => c.can_build( - &j.job.requirement.architecture, - &j.job.requirement.required_features, - ), + (PendingJob::Build(j), Some(c)) => c.can_run_build(&j.job), } } @@ -459,13 +456,13 @@ fn worker_context_of(caps: Option<&WorkerCaps>) -> WorkerContext<'_> { Some(c) => WorkerContext { architectures: &c.architectures, system_features: &c.system_features, - fetch: c.fetch, + fetches_flakes: c.capabilities.network, metrics: c.metrics, }, None => WorkerContext { architectures: &[], system_features: &[], - fetch: false, + fetches_flakes: false, metrics: None, }, } @@ -2169,7 +2166,6 @@ mod tests { #[test] fn can_eval_requires_eval_for_evaluation_steps() { let fetch_only = WorkerCaps { - fetch: true, capabilities: GradientCapabilities { network: true, secrets: true, @@ -2177,8 +2173,14 @@ mod tests { }, ..Default::default() }; + let network_only = WorkerCaps { + capabilities: GradientCapabilities { + network: true, + ..Default::default() + }, + ..Default::default() + }; let eval_only = WorkerCaps { - fetch: false, capabilities: GradientCapabilities { eval: true, ..Default::default() @@ -2186,7 +2188,6 @@ mod tests { ..Default::default() }; let both = WorkerCaps { - fetch: true, capabilities: GradientCapabilities { network: true, secrets: true, @@ -2201,7 +2202,13 @@ mod tests { FlakeStep::EvaluateDerivations, ]); let cached_eval = flake_job(vec![FlakeStep::EvaluateDerivations]); - let fetch = flake_job(vec![FlakeStep::FetchFlake]); + let fetch = FlakeJob { + source: FlakeSource::Repository { + url: "git+ssh://git@example.com/repo".into(), + commit: "abc123".into(), + }, + ..flake_job(vec![FlakeStep::FetchFlake]) + }; assert!(!fetch_only.can_eval(&bundled), "fetch-only lacks eval"); assert!(!eval_only.can_eval(&bundled), "eval-only lacks fetch"); @@ -2209,7 +2216,11 @@ mod tests { assert!(eval_only.can_eval(&cached_eval)); assert!(!fetch_only.can_eval(&cached_eval), "cached eval needs eval"); assert!(fetch_only.can_eval(&fetch)); - assert!(!eval_only.can_eval(&fetch), "fetch step needs fetch"); + assert!(!eval_only.can_eval(&fetch), "fetch step needs network"); + assert!( + !network_only.can_eval(&fetch), + "fetch step of an ssh repository needs secrets" + ); } #[test] @@ -2352,13 +2363,12 @@ mod tests { #[test] fn fetch_flake_job_requires_fetch_capability() { // Regression guard for #252. A FetchFlake job is cloning over SSH and must only go to - // fetch-capable workers. Only those workers are receiving SSH credentials. + // workers with network and secrets. Only those workers are receiving SSH credentials. let mut tracker = JobTracker::new(); let peer = ProjectId::now_v7(); tracker.add_pending("j1".into(), fetch_eval_job(peer)); let no_fetch = WorkerCaps { - fetch: false, architectures: vec!["x86_64-linux".into()], system_features: vec![], ..Default::default() @@ -2369,12 +2379,11 @@ mod tests { tracker .take_best_of_kind("w1", None, Some(&no_fetch), &JobKind::Flake, &*p, &inst) .is_none(), - "worker without fetch must not receive a fetch flake job" + "worker without network and secrets must not receive a fetch flake job" ); assert_eq!(tracker.pending_count(), 1); let with_fetch = WorkerCaps { - fetch: true, architectures: vec!["x86_64-linux".into()], system_features: vec![], capabilities: GradientCapabilities { @@ -2389,7 +2398,7 @@ mod tests { tracker .take_best_of_kind("w2", None, Some(&with_fetch), &JobKind::Flake, &*p, &inst) .is_some(), - "fetch- and eval-capable worker must receive the bundled flake job" + "worker with network, secrets and eval must receive the bundled flake job" ); } @@ -2400,7 +2409,6 @@ mod tests { tracker.add_pending("j1".into(), eval_job(peer)); let eval_no_fetch = WorkerCaps { - fetch: false, architectures: vec![], system_features: vec![], capabilities: GradientCapabilities { @@ -2422,7 +2430,7 @@ mod tests { &inst ) .is_some(), - "cached eval job must run on an eval-capable worker without fetch" + "cached eval job must run on an eval-capable worker without network" ); } @@ -2433,7 +2441,6 @@ mod tests { tracker.add_pending("j1".into(), fetch_eval_job(peer)); let fetch_build = WorkerCaps { - fetch: true, architectures: vec!["x86_64-linux".into()], system_features: vec![], capabilities: GradientCapabilities { diff --git a/backend/gradient-scheduler/src/waiting_state.rs b/backend/gradient-scheduler/src/waiting_state.rs index b4fedeb33..109bedbee 100644 --- a/backend/gradient-scheduler/src/waiting_state.rs +++ b/backend/gradient-scheduler/src/waiting_state.rs @@ -18,6 +18,7 @@ use gradient_entity::evaluation::{EvaluationKind, EvaluationStatus}; use gradient_graph::Graph; use gradient_pool::WorkerInfo; use gradient_types::events::worker_demand::WorkerDemand; +use gradient_types::input::check_repository_url_is_ssh; use gradient_types::*; use sea_orm::{ColumnTrait, ConnectionTrait, EntityTrait, QueryFilter, QuerySelect, Value}; use tracing::{error, info, warn}; @@ -25,7 +26,7 @@ use tracing::{error, info, warn}; const DRV_RECOVERY_GRACE_SECS: i64 = 120; use crate::assessment_memo::AssessmentMemo; -use crate::buildability::BuildabilityChecker; +use crate::buildability::{BuildWorkerCaps, BuildabilityChecker}; use crate::unbuildable::{Unbuildable, unbuildable}; use gradient_db::evaluations::counters::EvalCounters; @@ -92,7 +93,7 @@ impl WaitingRefresh<'_> { } let tasks = EvalTasks::load(&self.db.worker_db, &evals).await?; - let mut caps_by_project: HashMap, Vec> = HashMap::new(); + 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(); @@ -148,9 +149,12 @@ impl WaitingRefresh<'_> { EvaluationStatus::Queued | EvaluationStatus::Fetching | EvaluationStatus::EvaluatingFlake - | EvaluationStatus::EvaluatingDerivation => { - decide_pre_build_target(eval.kind, eval.status, project_workers) - } + | EvaluationStatus::EvaluatingDerivation => decide_pre_build_target( + eval.kind, + eval.status, + check_repository_url_is_ssh(&eval.repository), + project_workers, + ), EvaluationStatus::Building => { match self .build_phase_decision(eval.id, eval_counters, worker_caps, reason.as_ref()) @@ -185,7 +189,7 @@ impl WaitingRefresh<'_> { to = ?target, workers = project_workers.connected, eval_workers = project_workers.eval, - fetch_workers = project_workers.fetch, + network_workers = project_workers.network, "refreshing evaluation waiting state" ); } @@ -309,7 +313,7 @@ impl WaitingRefresh<'_> { &self, evaluation_id: EvaluationId, counters: EvalCounters, - worker_caps: &[(Vec, Vec)], + worker_caps: &[BuildWorkerCaps], current: Option<&WaitingReason>, ) -> Result { let a = match assess_buildability( @@ -396,7 +400,7 @@ async fn assess_buildability( memo: Option<&Mutex>, evaluation_id: EvaluationId, counters: EvalCounters, - worker_caps: &[(Vec, Vec)], + worker_caps: &[BuildWorkerCaps], ) -> Result { if let Some(phase) = phase_from_counters(counters) { return Ok(phase); @@ -456,7 +460,7 @@ impl WaitingRefresh<'_> { async fn attempt_graph_unstick( &self, evaluation_id: EvaluationId, - worker_caps: &[(Vec, Vec)], + worker_caps: &[BuildWorkerCaps], ) -> Result { info!(%evaluation_id, "graph stuck: pool can build every pending shared build but none is dispatchable; self-healing"); @@ -669,13 +673,18 @@ async fn eval_blocking_shared_builds( .collect()) } -fn pre_build_capability(kind: EvaluationKind, status: EvaluationStatus) -> Option { +fn pre_build_capability( + kind: EvaluationKind, + status: EvaluationStatus, + ssh_repository: bool, +) -> Option { if kind == EvaluationKind::Ssh { return None; } match status { - EvaluationStatus::Fetching => Some(EvalCapability::Fetch), + EvaluationStatus::Fetching if ssh_repository => Some(EvalCapability::Secrets), + EvaluationStatus::Fetching => Some(EvalCapability::Network), EvaluationStatus::Queued | EvaluationStatus::EvaluatingFlake | EvaluationStatus::EvaluatingDerivation => Some(EvalCapability::Eval), @@ -683,8 +692,6 @@ fn pre_build_capability(kind: EvaluationKind, status: EvaluationStatus) -> Optio } } -type WorkerCaps = (Vec, Vec); - fn eligible_for(worker: &WorkerInfo, project: Option) -> bool { worker .authorized_peers @@ -692,18 +699,25 @@ fn eligible_for(worker: &WorkerInfo, project: Option) -> bool { .is_none_or(|peers| project.is_some_and(|p| peers.contains(&p))) } -fn project_worker_caps(workers: &[WorkerInfo], project: Option) -> Vec { +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())) + .map(|w| { + ( + w.architectures.clone(), + w.system_features.clone(), + w.capabilities.network, + ) + }) .collect() } #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] struct ProjectWorkerCounts { eval: usize, - fetch: usize, + network: usize, + secrets: usize, connected: u32, } @@ -714,14 +728,17 @@ impl ProjectWorkerCounts { .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.network && w.capabilities.secrets), + network: counts.network + usize::from(w.capabilities.network), + secrets: counts.secrets + + usize::from(w.capabilities.network && w.capabilities.secrets), connected: counts.connected + 1, }) } fn provide(self, capability: EvalCapability) -> bool { match capability { - EvalCapability::Fetch => self.fetch > 0, + EvalCapability::Network => self.network > 0, + EvalCapability::Secrets => self.secrets > 0, EvalCapability::Eval => self.eval > 0, } } @@ -730,9 +747,10 @@ impl ProjectWorkerCounts { fn decide_pre_build_target( kind: EvaluationKind, current: EvaluationStatus, + ssh_repository: bool, workers: ProjectWorkerCounts, ) -> Option<(EvaluationStatus, Option)> { - let capability = pre_build_capability(kind, current)?; + let capability = pre_build_capability(kind, current, ssh_repository)?; if workers.provide(capability) { return None; } @@ -829,10 +847,11 @@ mod tests { } } - fn counts(eval: usize, fetch: usize, connected: u32) -> ProjectWorkerCounts { + fn counts(eval: usize, network: usize, connected: u32) -> ProjectWorkerCounts { ProjectWorkerCounts { eval, - fetch, + network, + secrets: network, connected, } } @@ -922,7 +941,11 @@ mod tests { 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()])] + vec![( + vec!["x86_64-linux".to_string()], + vec!["kvm".to_string()], + true + )] ); } @@ -934,6 +957,7 @@ mod tests { let (target, reason) = decide_pre_build_target( EvaluationKind::Normal, EvaluationStatus::Queued, + false, ProjectWorkerCounts::of(&workers, Some(own)), ) .expect("a project without its own eval worker must stall"); @@ -966,6 +990,7 @@ mod tests { let (target, reason) = decide_pre_build_target( EvaluationKind::Normal, EvaluationStatus::Queued, + false, counts(0, 1, 3), ) .expect("stall must produce a transition"); @@ -976,19 +1001,49 @@ mod tests { } #[test] - fn pre_build_target_fetching_no_fetch_worker_stalls_to_fetch_waiting() { + fn pre_build_target_fetching_without_network_waits_for_network() { let (target, reason) = decide_pre_build_target( EvaluationKind::Normal, EvaluationStatus::Fetching, + false, counts(2, 0, 2), ) .expect("stall must produce a transition"); assert_eq!(target, EvaluationStatus::Waiting); let (cap, connected) = eval_workers_view(&reason.expect("stall carries a reason")); - assert_eq!(cap, EvalCapability::Fetch); + assert_eq!(cap, EvalCapability::Network); assert_eq!(connected, 2); } + #[test] + fn fetching_an_ssh_repository_waits_for_secrets() { + let workers = ProjectWorkerCounts { + secrets: 0, + ..counts(2, 2, 2) + }; + let (_, reason) = decide_pre_build_target( + EvaluationKind::Normal, + EvaluationStatus::Fetching, + true, + workers, + ) + .expect("an ssh repository needs a worker with secrets"); + assert_eq!( + eval_workers_view(&reason.expect("stall carries a reason")).0, + EvalCapability::Secrets + ); + assert!( + decide_pre_build_target( + EvaluationKind::Normal, + EvaluationStatus::Fetching, + false, + workers, + ) + .is_none(), + "a public repository needs only network" + ); + } + #[test] fn pre_build_target_active_pre_build_with_capability_left_alone() { for status in [ @@ -998,7 +1053,8 @@ mod tests { EvaluationStatus::Queued, ] { assert!( - decide_pre_build_target(EvaluationKind::Normal, status, counts(1, 1, 1)).is_none(), + decide_pre_build_target(EvaluationKind::Normal, status, false, counts(1, 1, 1)) + .is_none(), "{status:?} with capable workers must be left alone" ); } @@ -1010,6 +1066,7 @@ mod tests { decide_pre_build_target( EvaluationKind::Ssh, EvaluationStatus::EvaluatingDerivation, + false, counts(0, 0, 0) ) .is_none() @@ -1022,6 +1079,7 @@ mod tests { decide_pre_build_target( EvaluationKind::Normal, EvaluationStatus::Waiting, + false, counts(0, 0, 0) ) .is_none() @@ -1030,6 +1088,7 @@ mod tests { decide_pre_build_target( EvaluationKind::Normal, EvaluationStatus::Waiting, + false, counts(2, 2, 2) ) .is_none() @@ -1042,17 +1101,17 @@ mod tests { assert_eq!(target, EvaluationStatus::Queued); assert!(reason.is_none()); - let (target, reason) = decide_eval_recovery(EvalCapability::Fetch, counts(0, 1, 1)); + let (target, reason) = decide_eval_recovery(EvalCapability::Network, counts(0, 1, 1)); assert_eq!(target, EvaluationStatus::Queued); assert!(reason.is_none()); } #[test] fn eval_recovery_refreshes_reason_while_capability_absent() { - let (target, reason) = decide_eval_recovery(EvalCapability::Fetch, counts(5, 0, 5)); + let (target, reason) = decide_eval_recovery(EvalCapability::Network, counts(5, 0, 5)); assert_eq!(target, EvaluationStatus::Waiting); let (cap, connected) = eval_workers_view(&reason.expect("refresh carries a reason")); - assert_eq!(cap, EvalCapability::Fetch); + assert_eq!(cap, EvalCapability::Network); assert_eq!(connected, 5); } diff --git a/backend/gradient-types/src/waiting_reason.rs b/backend/gradient-types/src/waiting_reason.rs index 6c7c188a8..eaa4487c3 100644 --- a/backend/gradient-types/src/waiting_reason.rs +++ b/backend/gradient-types/src/waiting_reason.rs @@ -9,7 +9,8 @@ use serde::{Deserialize, Serialize}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, utoipa::ToSchema)] #[serde(rename_all = "snake_case")] pub enum EvalCapability { - Fetch, + Network, + Secrets, Eval, } @@ -22,7 +23,7 @@ pub enum WaitingReason { available_architectures: Vec, }, /// `connected_workers` is the whole connected pool. It can be above zero while only - /// build-only workers are online and an eval or fetch worker is missing. + /// build-only workers are online and an eval or network worker is missing. EvalWorkers { capability: EvalCapability, connected_workers: u32, @@ -114,7 +115,8 @@ impl WaitingReason { Some(gaps.join(";")) } Self::EvalWorkers { capability, .. } => Some(match capability { - EvalCapability::Fetch => "fetch".to_owned(), + EvalCapability::Network => "network".to_owned(), + EvalCapability::Secrets => "secrets".to_owned(), EvalCapability::Eval => "eval".to_owned(), }), _ => None, diff --git a/backend/gradient-web/tests/team_demand.rs b/backend/gradient-web/tests/team_demand.rs index 575f5bcdc..f2b7eb0ca 100644 --- a/backend/gradient-web/tests/team_demand.rs +++ b/backend/gradient-web/tests/team_demand.rs @@ -58,6 +58,7 @@ async fn admin_reads_demand() { 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"]["network"], 0); + assert_eq!(body["message"]["secrets"], 0); assert_eq!(body["message"]["builds"], serde_json::json!([])); } From 8b816f3de18d5f41de2bc422958489432f42476c Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 14:52:31 +0200 Subject: [PATCH 04/19] workers: refactor: replace fetch with network --- cli/connector/tests/workers_api.rs | 4 +- cli/src/commands/completion.rs | 3 +- cli/src/commands/worker.rs | 4 +- docs/generated/configuration-options.md | 4 +- docs/gradient-api.yaml | 72 ++++++++++++++----- docs/src/concepts/workers.md | 20 ++++-- docs/src/contributors/proto/connection.md | 6 +- docs/src/contributors/proto/federation.md | 4 +- docs/src/contributors/proto/transfer.md | 2 +- docs/src/contributors/scheduler/scoring.md | 2 +- .../scheduler/waiting-and-recovery.md | 8 +-- docs/src/guides/gradient-ci-servers.md | 6 +- docs/src/reference/events.md | 2 +- docs/src/reference/scheduler-policies.md | 2 +- docs/src/reference/state.md | 4 +- frontend/src/app/core/models/api.generated.ts | 29 +++++--- .../app/core/models/gradient-ci.model.spec.ts | 4 +- frontend/src/app/core/models/worker.model.ts | 21 +++++- .../src/app/core/services/workers.service.ts | 4 +- .../job-detail/job-detail.component.spec.ts | 4 +- .../board/job-detail/job-detail.component.ts | 9 ++- .../board/workers/workers.component.ts | 2 +- .../evaluation-log.component.spec.ts | 19 +++-- .../evaluation-log.component.ts | 12 +++- .../workers/workers.component.spec.ts | 20 ++++-- .../projects/workers/workers.component.ts | 7 +- .../styleguide/pages/elements.component.html | 2 +- .../styleguide/pages/elements.component.ts | 8 ++- .../team-workers.component.spec.ts | 8 ++- .../team-workers/team-workers.component.ts | 3 +- .../allowed-capabilities.component.spec.ts | 32 +++++++-- .../allowed-capabilities.component.ts | 4 +- nix/tests/gradient/api/default.nix | 2 +- nix/tests/gradient/api/test.py | 3 +- 34 files changed, 240 insertions(+), 96 deletions(-) diff --git a/cli/connector/tests/workers_api.rs b/cli/connector/tests/workers_api.rs index 5bfbfe358..9698d7022 100644 --- a/cli/connector/tests/workers_api.rs +++ b/cli/connector/tests/workers_api.rs @@ -53,7 +53,9 @@ async fn create_worker_returns_response() { token: None, enable_build: None, enable_eval: None, - enable_fetch: None, + enable_exec: None, + enable_network: None, + enable_secrets: None, }, ) .await diff --git a/cli/src/commands/completion.rs b/cli/src/commands/completion.rs index 18036e1e4..b5a794ff3 100644 --- a/cli/src/commands/completion.rs +++ b/cli/src/commands/completion.rs @@ -295,7 +295,8 @@ mod tests { "worker_id": "builder-1", "display_name": "Builder One", "registered_at": "now", "active": true, "connected": true, "gradient_ci": false, "url": null, - "created_by": null, "enable_fetch": true, "enable_eval": true, + "created_by": null, "enable_network": true, "enable_secrets": true, + "enable_exec": false, "enable_eval": true, "enable_build": true, "live": null } ] diff --git a/cli/src/commands/worker.rs b/cli/src/commands/worker.rs index 3287533be..349910a04 100644 --- a/cli/src/commands/worker.rs +++ b/cli/src/commands/worker.rs @@ -69,7 +69,9 @@ pub async fn handle(cmd: Commands, out: Output) { token, enable_build: None, enable_eval: None, - enable_fetch: None, + enable_exec: None, + enable_network: None, + enable_secrets: None, }, ) .await diff --git a/docs/generated/configuration-options.md b/docs/generated/configuration-options.md index 161d3a067..579307b87 100644 --- a/docs/generated/configuration-options.md +++ b/docs/generated/configuration-options.md @@ -334,8 +334,10 @@ |---|---|---|---|---| | `worker.capabilities.build` | bool | `true` | `GRADIENT_WORKER_CAPABILITIES_BUILD` | Whether to run Nix builds. | | `worker.capabilities.eval` | bool | `true` | `GRADIENT_WORKER_CAPABILITIES_EVAL` | Whether to run Nix flake evaluations. | +| `worker.capabilities.exec` | bool | `false` | `GRADIENT_WORKER_CAPABILITIES_EXEC` | Whether this worker advertises the `exec` capability. The server stores it but does not schedule by it yet. | | `worker.capabilities.federate` | bool | `false` | `GRADIENT_WORKER_CAPABILITIES_FEDERATE` | Whether to forward work and NARs between workers and servers (requires `discoverable`). | -| `worker.capabilities.fetch` | bool | `true` | `GRADIENT_WORKER_CAPABILITIES_FETCH` | Whether to prefetch flake inputs and sources. | +| `worker.capabilities.network` | bool | `true` | `GRADIENT_WORKER_CAPABILITIES_NETWORK` | Whether this worker may reach the internet. Cloning repositories, fetching flake inputs, substitutions, downloads and fixed-output builds run only on workers with it. | +| `worker.capabilities.secrets` | bool | `true` | `GRADIENT_WORKER_CAPABILITIES_SECRETS` | Whether the server may send project credentials such as SSH keys to this worker. | ## `worker.eval` diff --git a/docs/gradient-api.yaml b/docs/gradient-api.yaml index 19827b6f9..f88c11894 100644 --- a/docs/gradient-api.yaml +++ b/docs/gradient-api.yaml @@ -1831,7 +1831,7 @@ paths: tags: - board summary: Fleet load by capability, architecture, and feature - description: 'Real per-dimension fleet load: for each capability (eval/fetch/build), architecture, and system feature, the in-flight job count against the summed slot capacity of the workers that can serve it. Answers whether the fleet is eval-, build-, or architecture-bound rather than showing one blended busy %. Scoped to the caller''s projects (superusers see the whole fleet). Each bucket carries `{ key, in_flight, capacity, workers }`; busy % = in_flight / capacity.' + description: 'Real per-dimension fleet load: for each capability (eval/network/build), architecture, and system feature, the in-flight job count against the summed slot capacity of the workers that can serve it. Answers whether the fleet is eval-, build-, or architecture-bound rather than showing one blended busy %. Scoped to the caller''s projects (superusers see the whole fleet). Each bucket carries `{ key, in_flight, capacity, workers }`; busy % = in_flight / capacity.' operationId: getBoardWorkerLoad responses: '200': @@ -13168,7 +13168,7 @@ components: - connected - draining - eval - - fetch + - network - build properties: bucket_start: @@ -13182,7 +13182,7 @@ components: eval: type: integer format: int64 - fetch: + network: type: integer format: int64 build: @@ -13391,7 +13391,7 @@ components: - assigned_jobs - max_concurrent_builds - eval - - fetch + - network - build - architectures properties: @@ -13417,7 +13417,7 @@ components: format: int64 eval: type: boolean - fetch: + network: type: boolean build: type: boolean @@ -14973,7 +14973,8 @@ components: EvalCapability: type: string enum: - - fetch + - network + - secrets - eval EvaluateRequest: type: object @@ -15607,23 +15608,29 @@ components: required: - core - federate - - fetch - eval - build - cache + - network + - secrets + - exec properties: core: type: boolean federate: type: boolean - fetch: - type: boolean eval: type: boolean build: type: boolean cache: type: boolean + network: + type: boolean + secrets: + type: boolean + exec: + type: boolean GrantCacheRequest: type: object required: @@ -17352,7 +17359,15 @@ components: type: - string - 'null' - enable_fetch: + enable_network: + type: + - boolean + - 'null' + enable_secrets: + type: + - boolean + - 'null' + enable_exec: type: - boolean - 'null' @@ -17647,7 +17662,9 @@ components: - display_name - registered_at - active - - enable_fetch + - enable_network + - enable_secrets + - enable_exec - enable_eval - enable_build - gradient_ci @@ -17669,7 +17686,11 @@ components: oneOf: - $ref: '#/components/schemas/UserId' - type: 'null' - enable_fetch: + enable_network: + type: boolean + enable_secrets: + type: boolean + enable_exec: type: boolean enable_eval: type: boolean @@ -17825,7 +17846,11 @@ components: type: - string - 'null' - enable_fetch: + enable_network: + type: boolean + enable_secrets: + type: boolean + enable_exec: type: boolean enable_eval: type: boolean @@ -18672,14 +18697,19 @@ components: type: object required: - builds - - fetch + - network + - secrets - eval properties: builds: type: array items: $ref: '#/components/schemas/UnmetRequirement' - fetch: + network: + type: integer + format: int32 + minimum: 0 + secrets: type: integer format: int32 minimum: 0 @@ -18901,7 +18931,9 @@ components: - active - managed - gradient_ci - - enable_fetch + - enable_network + - enable_secrets + - enable_exec - enable_eval - enable_build properties: @@ -18922,7 +18954,11 @@ components: - 'null' gradient_ci: type: boolean - enable_fetch: + enable_network: + type: boolean + enable_secrets: + type: boolean + enable_exec: type: boolean enable_eval: type: boolean @@ -19397,7 +19433,7 @@ components: - type: object description: |- `connected_workers` is the whole connected pool. It can be above zero while only - build-only workers are online and an eval or fetch worker is missing. + build-only workers are online and an eval or network worker is missing. required: - capability - connected_workers diff --git a/docs/src/concepts/workers.md b/docs/src/concepts/workers.md index f815568a1..a3bd7a8e4 100644 --- a/docs/src/concepts/workers.md +++ b/docs/src/concepts/workers.md @@ -9,15 +9,21 @@ flowchart LR ## Capabilities -Each worker declares which job kinds it accepts. Three kinds exist. +Workers declare their capabilities at the handshake. Capabilities cover the jobs and the access of a worker. -| Capability | Job | -|---|---| -| Fetch | Clone the repository and download the flake inputs | -| Evaluate | Walk the flake and report derivations | -| Build | Build derivations and upload the outputs | +| Capability | Default | Meaning | +|---|---|---| +| `network` | on | Internet access. Needed to clone the repository, fetch flake inputs, substitute from upstream caches and run fixed-output builds | +| `secrets` | on | The server may send the project SSH key. Needed to clone a repository behind an SSH URL | +| `exec` | off | Declared and stored. No job checks it yet | +| `eval` | on | Walk the flake and report derivations | +| `build` | on | Build derivations and upload the outputs | + +- Public repositories need a worker with `network` for the fetch step. +- Workers without `secrets` never get an SSH key. +- A project or team registration can withhold each capability through `enable_network`, `enable_secrets`, `enable_exec`, `enable_eval` and `enable_build`. -A single machine can handle all three kinds. Splitting them across machines also works: one machine with a lot of memory can evaluate, while several smaller machines build. +A single machine can hold every capability. A machine with a lot of memory can evaluate while smaller machines build. ## Matching Builds diff --git a/docs/src/contributors/proto/connection.md b/docs/src/contributors/proto/connection.md index e17a93ca0..9e10f3312 100644 --- a/docs/src/contributors/proto/connection.md +++ b/docs/src/contributors/proto/connection.md @@ -90,7 +90,9 @@ Two sides of a connection negotiate capabilities during the handshake. A capabil |---|---| | `core` | The server side of the protocol. Always enabled on the server, disabled on workers | | `cache` | Acting as a binary cache. Always enabled on the server | -| `fetch` | Cloning repositories and prefetching flake inputs | +| `network` | Internet access for clones, flake inputs, substitutions and fixed-output builds | +| `secrets` | Receiving project credentials such as the SSH key | +| `exec` | Declared and stored, with no behavior attached yet | | `eval` | Running flake evaluations | | `build` | Running Nix builds | | `federate` | Reserved. Part of the handshake, with no behavior attached yet | @@ -160,7 +162,7 @@ In-flight work can be lost during a restart. Queued jobs survive. - The agreed version is the lower of both newest versions. Every later frame is for that version. - Two ranges without overlap can end the session with close code `1002` and `no shared protocol version: ours 30..=30, peer 31..=32`. - A first frame without `GRAD` came from protocol 26 or older. The session can end with `peer protocol is older than 30`. -- The oldest supported version is 30. Older workers need an upgrade. +- The oldest supported version is 35. Older workers need an upgrade. - Both sides are logging the reason at the `warn` level. - A server connecting to a worker can keep the reason as the worker's offline reason. - `gradient_wire::PROTO_VERSIONS` is the supported range, computed from the `#[proto]` annotations. diff --git a/docs/src/contributors/proto/federation.md b/docs/src/contributors/proto/federation.md index 419578f3e..5d4402917 100644 --- a/docs/src/contributors/proto/federation.md +++ b/docs/src/contributors/proto/federation.md @@ -22,7 +22,7 @@ Each tenant gets one regular worker session toward the customer server. The hand | Aspect | Behavior | |---|---| -| Capabilities | A fixed set, `GRADIENT_PROXY_UPSTREAM_CAPABILITIES` (default `fetch,eval,build`) | +| Capabilities | A fixed set, `GRADIENT_PROXY_UPSTREAM_CAPABILITIES` (default `network,secrets,eval,build`) | | Hardware | `gradient_pool::aggregate()` over the downstream workers, sent as `WorkerCapabilities` and `WorkerMetrics` | | Aggregation | Capabilities OR'd, architectures and features unioned, slots, CPUs and RAM summed, fastest core score | | Aggregation scope | Per tenant: only that tenant's workers count | @@ -82,7 +82,7 @@ Access control works at two levels. | Level | Controlled by | |---|---| -| Server -> proxy | Per project: the proxy's `worker_registration` rows, tokens and `enable_fetch` / `enable_eval` / `enable_build` | +| Server -> proxy | Per project: the proxy's `worker_registration` rows, tokens and the `enable_*` capability grants | | Proxy -> worker | Per worker: `authorized_peers` rows with tenant, token hash and allowed capabilities | The handshake negotiates the `federate` capability. The server exposes this as `proto.federate`, the worker as `capabilities.federate`. The flag has no effect in the code yet. diff --git a/docs/src/contributors/proto/transfer.md b/docs/src/contributors/proto/transfer.md index 8185e4fa9..6ba432766 100644 --- a/docs/src/contributors/proto/transfer.md +++ b/docs/src/contributors/proto/transfer.md @@ -112,7 +112,7 @@ Before any build work starts, workers fetch all inputs that the local store does - The project's SSH key serve as the sole credential. Clones of private repositories and inputs depend on it. - Right before the `AssignJob` message, the server decrypt the key and send `Credential { SshKey }`. -- Only a flake job with a fetch step on a `fetch`-capable worker receive the credential. +- Only a flake job with a fetch step on a worker with `secrets` can receive the credential. - A project without a key receive no credential. - Workers hold the key in locked memory and zero it on drop. - At job completion, workers clear the key. diff --git a/docs/src/contributors/scheduler/scoring.md b/docs/src/contributors/scheduler/scoring.md index 5208a026e..a22a99244 100644 --- a/docs/src/contributors/scheduler/scoring.md +++ b/docs/src/contributors/scheduler/scoring.md @@ -32,7 +32,7 @@ flowchart LR | Context | Fields | Filled by | |---|---|---| | `JobContext` | `ScoredJob` (kind, architecture, `prefer_local_build`, `is_fixed_output`, `pname`, closure size, history), `missing_count`, `missing_nar_size`, `outputs_present`, `substitute_outputs`, `dependency_count`, `queued_at`, `ready_at`, `project_work_share`, `prioritized`, `build_request`, `ifd`, `rescore_count`, `now` | `JobTracker::score_candidates` in `gradient-scheduler/src/jobs.rs` | -| `WorkerContext` | `architectures`, `system_features`, `fetch`, `metrics` | `worker_context_of` from the worker's `WorkerCaps` | +| `WorkerContext` | `architectures`, `system_features`, `fetches_flakes`, `metrics` | `worker_context_of` from the worker's `WorkerCaps` | | `WorkerMetricsView` | `cpu_count`, `cpu_core_score`, `ram_total_mb`, `ram_free_mb`, `cpu_usage_pct`, `disk_speed_mbps`, `upload_speed_mbps`, `download_speed_mbps`, `running_builds`, `ram_reserved_mb`, `ram_reserved_unused_mb`, `ram_reserved_for_waiting_build_mb` | `WorkerCapabilities` (static), the 10 s `WorkerMetrics` heartbeat, the build stages, the assigned jobs and the waiting build (live) | | `InstanceContext` | 13 `Windowed` averages, `active_builds`, `pending_builds`, `total_workers`, `idle_workers`, `cpu_core_score_mean`, `upload_speed_mean_mbps`, `download_speed_mean_mbps`, `storage_read_mbps`, `storage_write_mbps`, `compression_ratio`, `per_path_secs`, `substitute_cost`, `download_slots`, `upload_slots`, `downloads_in_flight`, `uploads_in_flight` | `instance_metrics_pass`, see below. The two in-flight counts come from the build stages at assignment time | diff --git a/docs/src/contributors/scheduler/waiting-and-recovery.md b/docs/src/contributors/scheduler/waiting-and-recovery.md index 7de0737b3..aa2cdcb70 100644 --- a/docs/src/contributors/scheduler/waiting-and-recovery.md +++ b/docs/src/contributors/scheduler/waiting-and-recovery.md @@ -12,7 +12,7 @@ flowchart LR ## Fetch and Eval Split -`assign_queued_evals` (`gradient-scheduler/src/loops/eval.rs`) can ask the pool once per pass for an idle evaluation-only worker (`WorkerPool::has_idle_eval_only_worker`: active, `eval` without `fetch`, no assigned job). +`assign_queued_evals` (`gradient-scheduler/src/loops/eval.rs`) can ask the pool once per pass for an idle evaluation-only worker (`WorkerPool::has_idle_eval_only_worker`). Such workers are active, with `eval`, without `network` and without a job. | Pool at assignment | Job steps | |---|---| @@ -21,10 +21,10 @@ flowchart LR - **Decision per pass:** one idle worker will split every evaluation queued in that pass. - **Source:** a repository evaluation must carry `FlakeSource::Repository` as source. A `/nix/store/...` source (`gradient build` upload) must carry `FlakeSource::Cached` as source. That source will still run `FetchFlake` to resolve inputs with the project SSH key (`flake_job_for_eval_source`). -- **Routing:** `WorkerCaps::can_eval` (`gradient-pool/src/worker_caps.rs`) will require `fetch` for a `FetchFlake` step and `eval` for an evaluation step. +- **Routing:** `WorkerCaps::can_eval` (`gradient-pool/src/worker_caps.rs`) will require `network` for a `FetchFlake` step and `eval` for an evaluation step. Repositories behind an SSH URL also need `secrets` (`WorkerCaps::can_fetch`). - **Follow-up:** `handle_job_completed` (`job_handlers/build_status.rs`) will react to `JobCompleted` of a fetch-only job (`is_fetch_only`). The handler must read `evaluation.flake_source` and enqueue `PendingEvalJob::cached_followup` under the same `eval:{id}` key. The follow-up will carry `FlakeSource::Cached` plus the source as a required path. A missing `flake_source` will fail the evaluation permanently. - **Worker side:** the evaluating worker will substitute the source from the cache and evaluate `path:` (`gradient-worker/src/executor/eval.rs`). -- **Steering:** `ReserveFetchWorkersRule` (`gradient-pool/src/score/rules/builtin.rs`) will score a fetch-capable worker `-300 * (1 - idle / total)` for a job without `FetchFlake`. The penalty will fade as workers go idle. The rule will never refuse the job. See [Scoring](scoring.md). +- **Steering:** `ReserveFetchWorkersRule` (`gradient-pool/src/score/rules/builtin.rs`) will score a worker with `network` at `-300 * (1 - idle / total)` for a job without `FetchFlake`. The penalty will fade as workers go idle. The rule will never refuse the job. See [Scoring](scoring.md). ## Waiting Reasons @@ -32,7 +32,7 @@ The waiting pass `refresh_waiting_state` (`gradient-scheduler/src/waiting_state. | `WaitingReason` | Parked from | Unparked | Owner | |---|---|---|---| -| `eval_workers { capability, connected_workers }` | `Fetching` with no `fetch` worker. `Queued`, `EvaluatingFlake`, `EvaluatingDerivation` with no `eval` worker | To `Queued` once the capability can connect | Waiting pass | +| `eval_workers { capability, connected_workers }` | `Fetching` with no `network` worker, or no worker with `network` and `secrets` for an SSH repository (capability `secrets`). `Queued`, `EvaluatingFlake`, `EvaluatingDerivation` with no `eval` worker | To `Queued` once the capability can connect | Waiting pass | | `workers { unmet, connected_workers, available_architectures }` | `Building` with no worker that can fit the `(architecture, required_features)` of any blocking shared build | To `Building` once one worker can fit. To `Aborted` after 300 s on the same `unmet` set, unless `wait_for_workers` can be set on the task | Waiting pass | | `graph_stuck { pending_shared_builds }` | `Building` with a fitting worker for every blocking shared build, yet none can be handed out | To `Building` once the heal can free a shared build | Waiting pass | | `draining` | Every in-flight evaluation during an instance drain | At the end of draining or at startup (`resume_drained_evals`): to `Building` when the walk had finished (`status_after_draining`), otherwise to `Queued` | Waiting pass | diff --git a/docs/src/guides/gradient-ci-servers.md b/docs/src/guides/gradient-ci-servers.md index 942dd3c9e..f78837857 100644 --- a/docs/src/guides/gradient-ci-servers.md +++ b/docs/src/guides/gradient-ci-servers.md @@ -51,12 +51,12 @@ Connection tokens start with `gci1_`. ## Allowed Capabilities -A new connection grants `eval` and `build`. The entry shows both as **Allowed**. Fetching of sources stays the job of the project's own workers. +A new connection grants `network`, `eval` and `build`. All three appear as **Allowed** on the entry. Only the project's own workers get the project SSH key. | Workers of the project | Allowed on Gradient.CI Servers | |---|---| -| An own worker allowed to fetch | `eval`, `build` (default) | -| Gradient.CI Servers only | `fetch`, `eval`, `build` | +| An own worker with `secrets`, or a public repository | `network`, `eval`, `build` (default) | +| Gradient.CI Servers only, repository behind an SSH URL | `network`, `secrets`, `eval`, `build` | - **Edit** on the entry changes the allowed [capabilities](../concepts/workers.md#capabilities). - **Deactivate** pauses the connection without a deletion. diff --git a/docs/src/reference/events.md b/docs/src/reference/events.md index ce8fb31ad..4bf3fd1f0 100644 --- a/docs/src/reference/events.md +++ b/docs/src/reference/events.md @@ -29,7 +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 | +| `worker.demand` | yes, team webhooks only | waiting evaluations that need a system, feature, network, secrets 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` | diff --git a/docs/src/reference/scheduler-policies.md b/docs/src/reference/scheduler-policies.md index f069ee3cb..6542c7300 100644 --- a/docs/src/reference/scheduler-policies.md +++ b/docs/src/reference/scheduler-policies.md @@ -27,7 +27,7 @@ Workers take the job with the highest total score. Only jobs with a total of at | `WaitTimeRule` | Bonus, up to 4000 | Raise the score of jobs that wait long, to prevent starvation. The wait count start at the moment all dependencies finish | | `QosRule` | Bonus, 5000, 1000 and 1000 | Put prioritized jobs above all others. Jobs of a build request ([`gradient build`](../guides/build-before-push.md) or [SSH](../guides/build-over-ssh.md)) receive another 1000. Builds of an imported derivation and their unfinished dependencies also receive 1000 while evaluations wait on the import | | `RescoreWaitRule` | Veto | Hold a build until a worker report the size of its missing data. The hold lift after 4 rounds | -| `ReserveFetchWorkersRule` | Penalty, up to 300 | Keep fetch-capable workers free for fetching while capacity run short | +| `ReserveFetchWorkersRule` | Penalty, up to 300 | Keep workers with `network` free for flake fetches while capacity run short | | `TransferLimitRule` | Penalty, 4000 | Hold large transfers while all transfer slots sit full, see [Transfer Limits](#transfer-limits) | The **Prioritize** entry in the task or evaluation menu set the `QosRule` flag on a build and its dependencies, or on a whole evaluation. A failed or aborted build or evaluation clear the flag. Only superusers can prioritize. diff --git a/docs/src/reference/state.md b/docs/src/reference/state.md index c3eeb12f9..8de1b77d8 100644 --- a/docs/src/reference/state.md +++ b/docs/src/reference/state.md @@ -184,7 +184,9 @@ Each worker registration has the following options. | `display_name` | string | attribute name | Name shown for the worker in the UI. | | `enable_build` | bool | `true` | Grant of the worker's `build` capability by the server. | | `enable_eval` | bool | `true` | Grant of the worker's `eval` capability by the server. | -| `enable_fetch` | bool | `true` | Grant of the worker's `fetch` capability by the server. | +| `enable_exec` | bool | `false` | Grant of the worker's `exec` capability by the server. | +| `enable_network` | bool | `true` | Grant of the worker's `network` capability by the server. | +| `enable_secrets` | bool | `true` | Grant of the worker's `secrets` capability by the server. | | `enabled` | bool | `true` | Activation state of the worker. UI changes persist only until the next server start. | | `projects` | list of string | `[ ]` | Projects under which the worker registered, one registration per project. An empty list marks a team worker. | | `team` | null or string | `null` | Owning team. A team worker can build for every project that grants the team's workers. Cannot be combined with `projects`. | diff --git a/frontend/src/app/core/models/api.generated.ts b/frontend/src/app/core/models/api.generated.ts index ad35a0594..f3877020a 100644 --- a/frontend/src/app/core/models/api.generated.ts +++ b/frontend/src/app/core/models/api.generated.ts @@ -170,7 +170,7 @@ export interface BoardFleetPoint { connected: number; draining: number; eval: number; - fetch: number; + network: number; build: number; } @@ -240,7 +240,7 @@ export interface BoardWorker { assigned_jobs: number; max_concurrent_builds: number; eval: boolean; - fetch: boolean; + network: boolean; build: boolean; architectures: string[]; cpu_usage_pct?: number | null; @@ -796,7 +796,7 @@ export interface ErrorResponseBody { message: string; } -export type EvalCapability = "fetch" | "eval"; +export type EvalCapability = "network" | "secrets" | "eval"; export interface EvaluateRequest { commit?: string | null; @@ -1009,10 +1009,12 @@ export interface GradientCacheInfo { export interface GradientCapabilities { core: boolean; federate: boolean; - fetch: boolean; eval: boolean; build: boolean; cache: boolean; + network: boolean; + secrets: boolean; + exec: boolean; } export interface GrantCacheRequest { @@ -1618,7 +1620,9 @@ export interface PatchUserSettingsRequest { export interface PatchWorkerRequest { active?: boolean | null; display_name?: string | null; - enable_fetch?: boolean | null; + enable_network?: boolean | null; + enable_secrets?: boolean | null; + enable_exec?: boolean | null; enable_eval?: boolean | null; enable_build?: boolean | null; } @@ -1729,7 +1733,9 @@ export type ProjectWorkerEntry = WorkerConnection & { active: boolean; url?: string | null; created_by?: UserId | null; - enable_fetch: boolean; + enable_network: boolean; + enable_secrets: boolean; + enable_exec: boolean; enable_eval: boolean; enable_build: boolean; team?: string | null; @@ -1788,7 +1794,9 @@ export interface RegisterWorkerRequest { url?: string | null; display_name: string; token?: string | null; - enable_fetch?: boolean; + enable_network?: boolean; + enable_secrets?: boolean; + enable_exec?: boolean; enable_eval?: boolean; enable_build?: boolean; } @@ -2112,7 +2120,8 @@ export interface TeamCacheGrant { export interface TeamDemand { builds: UnmetRequirement[]; - fetch: number; + network: number; + secrets: number; eval: number; } @@ -2203,7 +2212,9 @@ export type TeamWorkerEntry = WorkerConnection & { managed: boolean; url?: string | null; gradient_ci: boolean; - enable_fetch: boolean; + enable_network: boolean; + enable_secrets: boolean; + enable_exec: boolean; enable_eval: boolean; enable_build: boolean; live?: WorkerLiveInfo | null; diff --git a/frontend/src/app/core/models/gradient-ci.model.spec.ts b/frontend/src/app/core/models/gradient-ci.model.spec.ts index 6ec0f8102..5d96ff2c7 100644 --- a/frontend/src/app/core/models/gradient-ci.model.spec.ts +++ b/frontend/src/app/core/models/gradient-ci.model.spec.ts @@ -16,7 +16,9 @@ function worker(overrides: Partial): ProjectWorkerEntry { active: true, gradient_ci: false, connected: false, - enable_fetch: true, + enable_network: true, + enable_secrets: true, + enable_exec: false, enable_eval: true, enable_build: true, ...overrides, diff --git a/frontend/src/app/core/models/worker.model.ts b/frontend/src/app/core/models/worker.model.ts index f625fe1a1..53dbc3ede 100644 --- a/frontend/src/app/core/models/worker.model.ts +++ b/frontend/src/app/core/models/worker.model.ts @@ -8,14 +8,29 @@ import type { ProjectWorkerEntry } from './api.generated'; export type AllowedCapabilities = Pick< ProjectWorkerEntry, - 'enable_fetch' | 'enable_eval' | 'enable_build' + 'enable_network' | 'enable_secrets' | 'enable_exec' | 'enable_eval' | 'enable_build' >; +export const ALLOWED_CAPABILITY_KEYS: (keyof AllowedCapabilities)[] = [ + 'enable_network', + 'enable_secrets', + 'enable_exec', + 'enable_eval', + 'enable_build', +]; + +export const ALL_ALLOWED: AllowedCapabilities = { + enable_network: true, + enable_secrets: true, + enable_exec: false, + enable_eval: true, + enable_build: true, +}; + export function changedCapabilities( before: AllowedCapabilities, after: AllowedCapabilities, ): Partial { - const keys: (keyof AllowedCapabilities)[] = ['enable_fetch', 'enable_eval', 'enable_build']; - return Object.fromEntries(keys.filter((k) => before[k] !== after[k]).map((k) => [k, after[k]])); + return Object.fromEntries(ALLOWED_CAPABILITY_KEYS.filter((k) => before[k] !== after[k]).map((k) => [k, after[k]])); } diff --git a/frontend/src/app/core/services/workers.service.ts b/frontend/src/app/core/services/workers.service.ts index 4f2c5dbaf..d224e8119 100644 --- a/frontend/src/app/core/services/workers.service.ts +++ b/frontend/src/app/core/services/workers.service.ts @@ -43,7 +43,9 @@ export class WorkersService { display_name: displayName, url: url || undefined, token: token || undefined, - enable_fetch: caps?.enable_fetch ?? true, + enable_network: caps?.enable_network ?? true, + enable_secrets: caps?.enable_secrets ?? true, + enable_exec: caps?.enable_exec ?? false, enable_eval: caps?.enable_eval ?? true, enable_build: caps?.enable_build ?? true, }); 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 698c23b44..ebd647edd 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 @@ -45,7 +45,7 @@ const DETAIL: AssignedJobDetail = { worker_context: { architectures: ['x86_64-linux'], system_features: ['kvm'], - capabilities: { core: true, federate: false, fetch: true, eval: false, build: true, cache: true }, + capabilities: { core: true, federate: false, network: true, secrets: false, exec: false, eval: false, build: true, cache: true }, cpu_count: 8, cpu_core_score: 1.5, ram_total_mb: 32000, @@ -204,7 +204,7 @@ describe('BoardJobDetailComponent - structured context panels', () => { const caps = Array.from(el.querySelectorAll('.ctx.worker tr')).find((r) => r.textContent?.includes('Capabilities'), ) as HTMLElement; - expect(caps.textContent).toContain('fetch'); + expect(caps.textContent).toContain('network'); expect(caps.textContent).toContain('build'); expect(caps.textContent).not.toContain('core'); expect(caps.textContent).not.toContain('cache'); 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 e13ff17bb..0ac863b63 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 @@ -358,7 +358,14 @@ export class BoardJobDetailComponent implements OnInit { protected readonly formatMegabytes = formatMegabytes; protected readonly formatDuration = formatDuration; - private static readonly WORKER_CAPS: (keyof GradientCapabilities)[] = ['federate', 'fetch', 'eval', 'build']; + private static readonly WORKER_CAPS: (keyof GradientCapabilities)[] = [ + 'federate', + 'network', + 'secrets', + 'exec', + 'eval', + 'build', + ]; currentState = computed(() => { const j = this.job(); diff --git a/frontend/src/app/features/board/workers/workers.component.ts b/frontend/src/app/features/board/workers/workers.component.ts index a4b3d8add..1557b3078 100644 --- a/frontend/src/app/features/board/workers/workers.component.ts +++ b/frontend/src/app/features/board/workers/workers.component.ts @@ -132,7 +132,7 @@ export class BoardWorkersComponent implements OnInit { ]); capOverTime = computed(() => [ { name: 'eval', data: this.fleet().map((p) => p.eval) }, - { name: 'fetch', data: this.fleet().map((p) => p.fetch) }, + { name: 'network', data: this.fleet().map((p) => p.network) }, { name: 'build', data: this.fleet().map((p) => p.build) }, ]); diff --git a/frontend/src/app/features/evaluations/evaluation-log/evaluation-log.component.spec.ts b/frontend/src/app/features/evaluations/evaluation-log/evaluation-log.component.spec.ts index c941a38b6..92b08ce38 100644 --- a/frontend/src/app/features/evaluations/evaluation-log/evaluation-log.component.spec.ts +++ b/frontend/src/app/features/evaluations/evaluation-log/evaluation-log.component.spec.ts @@ -308,13 +308,22 @@ describe('EvaluationLogComponent', () => { }); // #381: pre-build evals park with an eval_workers reason naming the missing - // capability (fetch while Fetching, eval while Evaluating). + // capability (network or secrets while Fetching, eval while Evaluating). describe('eval_workers waiting reason', () => { - it('titles and explains a missing fetch worker', () => { + it('titles and explains a missing network worker', () => { const { cmp } = setup(); - const reason = { kind: 'eval_workers', capability: 'fetch', connected_workers: 0 } as const; - expect(cmp.waitingTitle(reason)).toBe('Waiting for a Fetch Worker'); - expect(cmp.formatWaitingReason(reason)).toContain('fetch the flake sources'); + const reason = { kind: 'eval_workers', capability: 'network', connected_workers: 0 } as const; + expect(cmp.waitingTitle(reason)).toBe('Waiting for a Network Worker'); + expect(cmp.formatWaitingReason(reason)).toBe('No workers are connected to fetch the flake sources.'); + }); + + it('titles and explains a missing worker with secrets', () => { + const { cmp } = setup(); + const reason = { kind: 'eval_workers', capability: 'secrets', connected_workers: 1 } as const; + expect(cmp.waitingTitle(reason)).toBe('Waiting for a Worker with Secrets'); + expect(cmp.formatWaitingReason(reason)).toBe( + '1 worker is connected, but none can fetch the flake sources with the project SSH key.', + ); }); it('titles and explains a missing eval worker with connected count', () => { diff --git a/frontend/src/app/features/evaluations/evaluation-log/evaluation-log.component.ts b/frontend/src/app/features/evaluations/evaluation-log/evaluation-log.component.ts index b7cc85012..daac4487e 100644 --- a/frontend/src/app/features/evaluations/evaluation-log/evaluation-log.component.ts +++ b/frontend/src/app/features/evaluations/evaluation-log/evaluation-log.component.ts @@ -1625,7 +1625,11 @@ export class EvaluationLogComponent implements OnInit, OnDestroy { return `${reason.connected_workers} connected ${workerWord} (${archList}) cannot satisfy:`; } case 'eval_workers': { - const action = reason.capability === 'fetch' ? 'fetch the flake sources' : 'run the evaluation'; + const action = { + network: 'fetch the flake sources', + secrets: 'fetch the flake sources with the project SSH key', + eval: 'run the evaluation', + }[reason.capability]; const workerWord = reason.connected_workers === 1 ? 'worker is' : 'workers are'; return reason.connected_workers === 0 ? `No workers are connected to ${action}.` @@ -1648,7 +1652,11 @@ export class EvaluationLogComponent implements OnInit, OnDestroy { waitingTitle(reason?: WaitingReason | null): string { switch (reason?.kind) { - case 'eval_workers': return reason.capability === 'fetch' ? 'Waiting for a Fetch Worker' : 'Waiting for an Eval Worker'; + case 'eval_workers': return { + network: 'Waiting for a Network Worker', + secrets: 'Waiting for a Worker with Secrets', + eval: 'Waiting for an Eval Worker', + }[reason.capability]; case 'approval': return 'Awaiting Approval'; case 'no_cache': return 'No Cache Configured'; case 'cache_storage_full': return 'Cache Storage Full'; diff --git a/frontend/src/app/features/projects/workers/workers.component.spec.ts b/frontend/src/app/features/projects/workers/workers.component.spec.ts index 34219129e..e0f6f55e8 100644 --- a/frontend/src/app/features/projects/workers/workers.component.spec.ts +++ b/frontend/src/app/features/projects/workers/workers.component.spec.ts @@ -38,7 +38,9 @@ const workerUnmanaged: ProjectWorkerEntry = { active: true, gradient_ci: false, connected: false, - enable_fetch: true, + enable_network: true, + enable_secrets: true, + enable_exec: false, enable_eval: true, enable_build: true, }; @@ -51,7 +53,9 @@ const workerOfTeam: ProjectWorkerEntry = { team: 'platform', gradient_ci: false, connected: false, - enable_fetch: true, + enable_network: true, + enable_secrets: true, + enable_exec: false, enable_eval: true, enable_build: true, }; @@ -253,7 +257,7 @@ describe('WorkersComponent - team workers', () => { describe('WorkersComponent - system features', () => { it('lists the architectures and keeps the system features behind an info button', async () => { const live = { - capabilities: { fetch: true, eval: true, build: true, federate: false }, + capabilities: { network: true, secrets: true, exec: false, eval: true, build: true, federate: false }, architectures: ['x86_64-linux'], system_features: ['kvm', 'nixos-test'], max_concurrent_builds: 1, assigned_job_count: 0, draining: false, } as unknown as ProjectWorkerEntry['live']; const fixture = setup({ @@ -274,7 +278,7 @@ describe('WorkersComponent - system features', () => { describe('WorkersComponent - connection badge', () => { const live = { - capabilities: { fetch: true, eval: true, build: true, federate: false }, + capabilities: { network: true, secrets: true, exec: false, eval: true, build: true, federate: false }, architectures: [], system_features: [], max_concurrent_builds: 1, assigned_job_count: 0, draining: false, } as unknown as ProjectWorkerEntry['live']; const badges = (fixture: ComponentFixture) => @@ -325,7 +329,9 @@ const gciRegistration: ProjectWorkerEntry = { direction: 'outbound', before_auth: false, }, - enable_fetch: false, + enable_network: true, + enable_secrets: false, + enable_exec: false, enable_eval: true, enable_build: true, }; @@ -379,11 +385,11 @@ describe('WorkersComponent - Gradient.CI Servers entry', () => { const dialog = document.querySelector('.gr-dialog')!; expect(dialog.querySelector('#rename-name')).toBeNull(); Array.from(dialog.querySelectorAll('gr-allowed-capabilities button')) - .find((b) => b.textContent?.includes('fetch'))! + .find((b) => b.textContent?.includes('secrets'))! .dispatchEvent(new Event('click')); buttonsLabelled(dialog, 'Save')[0].click(); - expect(workers.patchWorker).toHaveBeenCalledWith('demo', 'g1', { enable_fetch: true }); + expect(workers.patchWorker).toHaveBeenCalledWith('demo', 'g1', { enable_secrets: true }); }); it('links to the Gradient.CI site only while connected', async () => { diff --git a/frontend/src/app/features/projects/workers/workers.component.ts b/frontend/src/app/features/projects/workers/workers.component.ts index 6f73b7be5..55a9371a3 100644 --- a/frontend/src/app/features/projects/workers/workers.component.ts +++ b/frontend/src/app/features/projects/workers/workers.component.ts @@ -18,6 +18,7 @@ import { TeamsService } from '@core/services/teams.service'; import { AuthService } from '@core/services/auth.service'; import { AccessState, + ALL_ALLOWED, AllowedCapabilities, changedCapabilities, gradientCiEntry, @@ -56,8 +57,6 @@ import { AllowedCapabilitiesComponent, GradientCiConnectComponent, LabelHelpComp import { WritableDirective, ManagedDisableDirective, canOpenTeam } from '@shared/access'; import { TeamGrantsComponent } from '@features/teams/team-grants/team-grants.component'; -const ALL_ALLOWED: AllowedCapabilities = { enable_fetch: true, enable_eval: true, enable_build: true }; - @Component({ selector: 'app-workers', standalone: true, @@ -115,7 +114,9 @@ export class WorkersComponent implements OnInit { readonly capLabels: { key: keyof GradientCapabilities; label: string }[] = [ { key: 'federate', label: 'federate' }, - { key: 'fetch', label: 'fetch' }, + { key: 'network', label: 'network' }, + { key: 'secrets', label: 'secrets' }, + { key: 'exec', label: 'exec' }, { key: 'eval', label: 'eval' }, { key: 'build', label: 'build' }, ]; diff --git a/frontend/src/app/features/styleguide/pages/elements.component.html b/frontend/src/app/features/styleguide/pages/elements.component.html index dacebbdf7..01d935130 100644 --- a/frontend/src/app/features/styleguide/pages/elements.component.html +++ b/frontend/src/app/features/styleguide/pages/elements.component.html @@ -143,7 +143,7 @@

Star button

Allowed capabilities

What a worker may run. The editable variant is the toggle group of worker dialogs.

- +

Permission picker

diff --git a/frontend/src/app/features/styleguide/pages/elements.component.ts b/frontend/src/app/features/styleguide/pages/elements.component.ts index 932455fdf..adfc6d6fa 100644 --- a/frontend/src/app/features/styleguide/pages/elements.component.ts +++ b/frontend/src/app/features/styleguide/pages/elements.component.ts @@ -83,7 +83,13 @@ export class ElementsComponent { ] as const; statusPhases: StatusPhase[] = ['queued', 'waiting', 'running', 'success', 'failure', 'aborted']; statusPhase = signal('queued'); - allowed: AllowedCapabilities = { enable_fetch: true, enable_eval: true, enable_build: false }; + allowed: AllowedCapabilities = { + enable_network: true, + enable_secrets: true, + enable_exec: false, + enable_eval: true, + enable_build: false, + }; thunkSamples: { caption: string; progress: ThunkProgress }[] = [ { caption: 'Against the median thunks of the recent evaluations', progress: thunkProgress({ kind: 'evaluating', thunks: 128_032_032 }, 450_000_000)! }, { caption: 'Without a history, or past the last total, pulsing at full width', progress: thunkProgress({ kind: 'evaluating', thunks: 128_032_032 }, null)! }, diff --git a/frontend/src/app/features/teams/team-workers/team-workers.component.spec.ts b/frontend/src/app/features/teams/team-workers/team-workers.component.spec.ts index 89342dfee..29a170f97 100644 --- a/frontend/src/app/features/teams/team-workers/team-workers.component.spec.ts +++ b/frontend/src/app/features/teams/team-workers/team-workers.component.spec.ts @@ -25,7 +25,9 @@ const connection: TeamWorkerEntry = { active: true, managed: false, gradient_ci: true, - enable_fetch: false, + enable_network: true, + enable_secrets: false, + enable_exec: false, enable_eval: true, enable_build: true, connected: true, @@ -73,11 +75,11 @@ describe('TeamWorkersComponent', () => { fixture.detectChanges(); const dialog = document.querySelector('.gr-dialog')!; Array.from(dialog.querySelectorAll('gr-allowed-capabilities button')) - .find((b) => b.textContent?.includes('fetch'))! + .find((b) => b.textContent?.includes('secrets'))! .dispatchEvent(new Event('click')); button(dialog, 'Save')!.click(); - expect(updateWorker).toHaveBeenCalledWith('platform', 'g1', { enable_fetch: true }); + expect(updateWorker).toHaveBeenCalledWith('platform', 'g1', { enable_secrets: true }); }); it('links a Gradient.CI connection to its site', () => { diff --git a/frontend/src/app/features/teams/team-workers/team-workers.component.ts b/frontend/src/app/features/teams/team-workers/team-workers.component.ts index 9059b3a3e..0ff2cd248 100644 --- a/frontend/src/app/features/teams/team-workers/team-workers.component.ts +++ b/frontend/src/app/features/teams/team-workers/team-workers.component.ts @@ -12,6 +12,7 @@ import { TeamsService } from '@core/services/teams.service'; import { AuthService } from '@core/services/auth.service'; import { ConfigService } from '@core/services/config.service'; import { + ALL_ALLOWED, AllowedCapabilities, changedCapabilities, gradientCiHost, @@ -81,7 +82,7 @@ export class TeamWorkersComponent implements OnInit { issuedToken = signal(null); editing = signal(null); editName = ''; - editAllowed: AllowedCapabilities = { enable_fetch: true, enable_eval: true, enable_build: true }; + editAllowed: AllowedCapabilities = ALL_ALLOWED; form: RegisterWorkerRequest = { worker_id: '', display_name: '' }; isAdmin = computed(() => this.team()?.role === 'admin' || this.authService.user()?.superuser === true); diff --git a/frontend/src/app/shared/ui/allowed-capabilities/allowed-capabilities.component.spec.ts b/frontend/src/app/shared/ui/allowed-capabilities/allowed-capabilities.component.spec.ts index bdbf52701..0999a0627 100644 --- a/frontend/src/app/shared/ui/allowed-capabilities/allowed-capabilities.component.spec.ts +++ b/frontend/src/app/shared/ui/allowed-capabilities/allowed-capabilities.component.spec.ts @@ -10,7 +10,13 @@ import { AllowedCapabilitiesComponent } from './allowed-capabilities.component'; function render(editable: boolean) { TestBed.configureTestingModule({ imports: [AllowedCapabilitiesComponent] }); const fixture = TestBed.createComponent(AllowedCapabilitiesComponent); - fixture.componentRef.setInput('value', { enable_fetch: false, enable_eval: true, enable_build: true }); + fixture.componentRef.setInput('value', { + enable_network: true, + enable_secrets: false, + enable_exec: false, + enable_eval: true, + enable_build: true, + }); fixture.componentRef.setInput('editable', editable); fixture.detectChanges(); return fixture; @@ -23,19 +29,31 @@ describe('AllowedCapabilitiesComponent', () => { expect(root.textContent).toContain('Allowed:'); expect(root.querySelectorAll('button').length).toBe(0); const off = Array.from(root.querySelectorAll('[data-allowed="false"]')).map((el) => el.textContent?.trim()); - expect(off).toEqual(['fetch']); + expect(off).toEqual(['secrets', 'exec']); + }); + + it('lists the capabilities in the order network, secrets, exec, eval, build', () => { + const root = render(false).nativeElement as HTMLElement; + const labels = Array.from(root.querySelectorAll('[data-allowed]')).map((el) => el.textContent?.trim()); + expect(labels).toEqual(['network', 'secrets', 'exec', 'eval', 'build']); }); it('toggles only the clicked capability when editable', () => { const fixture = render(true); - const fetch = Array.from((fixture.nativeElement as HTMLElement).querySelectorAll('button')).find( - (b) => b.textContent?.includes('fetch'), + const secrets = Array.from((fixture.nativeElement as HTMLElement).querySelectorAll('button')).find( + (b) => b.textContent?.includes('secrets'), )!; - fetch.click(); + secrets.click(); fixture.detectChanges(); - expect(fixture.componentInstance.value()).toEqual({ enable_fetch: true, enable_eval: true, enable_build: true }); - expect(fetch.getAttribute('aria-pressed')).toBe('true'); + expect(fixture.componentInstance.value()).toEqual({ + enable_network: true, + enable_secrets: true, + enable_exec: false, + enable_eval: true, + enable_build: true, + }); + expect(secrets.getAttribute('aria-pressed')).toBe('true'); }); }); diff --git a/frontend/src/app/shared/ui/allowed-capabilities/allowed-capabilities.component.ts b/frontend/src/app/shared/ui/allowed-capabilities/allowed-capabilities.component.ts index 759ace597..680ed121e 100644 --- a/frontend/src/app/shared/ui/allowed-capabilities/allowed-capabilities.component.ts +++ b/frontend/src/app/shared/ui/allowed-capabilities/allowed-capabilities.component.ts @@ -45,7 +45,9 @@ export class AllowedCapabilitiesComponent { editable = input(false, { transform: booleanAttribute }); readonly capabilities: { key: keyof AllowedCapabilities; label: string }[] = [ - { key: 'enable_fetch', label: 'fetch' }, + { key: 'enable_network', label: 'network' }, + { key: 'enable_secrets', label: 'secrets' }, + { key: 'enable_exec', label: 'exec' }, { key: 'enable_eval', label: 'eval' }, { key: 'enable_build', label: 'build' }, ]; diff --git a/nix/tests/gradient/api/default.nix b/nix/tests/gradient/api/default.nix index 6fb99bb5f..bbd714db3 100644 --- a/nix/tests/gradient/api/default.nix +++ b/nix/tests/gradient/api/default.nix @@ -135,7 +135,7 @@ in { token_file = toString stateWorkerToken; display_name = "State Worker"; created_by = "stateadmin"; - enable_fetch = false; + enable_secrets = false; }; teams.stateteam = { diff --git a/nix/tests/gradient/api/test.py b/nix/tests/gradient/api/test.py index 00b34f7ed..5d030b499 100644 --- a/nix/tests/gradient/api/test.py +++ b/nix/tests/gradient/api/test.py @@ -576,7 +576,8 @@ def accept_invite(kind, scope, invitee_token): # Worker registration applied with its capability gates. worker = next(w for w in api("GET", "projects/stateproject/workers", token=sa_token) if w["worker_id"] == "a0000000-0000-0000-0000-0000000000aa") -assert worker["enable_eval"] is True and worker["enable_fetch"] is False, worker +assert worker["enable_eval"] is True and worker["enable_network"] is True, worker +assert worker["enable_secrets"] is False and worker["enable_exec"] is False, worker # team workers team_worker_id = "a0000000-0000-0000-0000-0000000000bb" From 383b99b00efef0d5ccb2de3a503f16232fea43b3 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 15:07:02 +0200 Subject: [PATCH 05/19] nix: feat: let the embedder fetch --- ...-libexpr-feat-let-the-embedder-fetch.patch | 314 ++++++++++++++++++ 1 file changed, 314 insertions(+) create mode 100644 nix/patches/nix/0023-libexpr-feat-let-the-embedder-fetch.patch diff --git a/nix/patches/nix/0023-libexpr-feat-let-the-embedder-fetch.patch b/nix/patches/nix/0023-libexpr-feat-let-the-embedder-fetch.patch new file mode 100644 index 000000000..5ab35b258 --- /dev/null +++ b/nix/patches/nix/0023-libexpr-feat-let-the-embedder-fetch.patch @@ -0,0 +1,314 @@ +From c46e4d6dfc4f07003799c7e87ea09b1e6e4f0c28 Mon Sep 17 00:00:00 2001 +From: Dennis Wuitz +Date: Sun, 18 Oct 2026 14:56:12 +0200 +Subject: [PATCH] libexpr: feat: let the embedder fetch + +--- + src/libexpr-c/nix_api_expr.cc | 20 +++++ + src/libexpr-c/nix_api_expr.h | 34 ++++++++ + src/libexpr/primops/fetchTree.cc | 47 +++++++++-- + src/libfetchers/fetchers.cc | 77 +++++++++++++++---- + .../include/nix/fetchers/fetch-settings.hh | 8 ++ + 5 files changed, 164 insertions(+), 22 deletions(-) + +diff --git a/src/libexpr-c/nix_api_expr.cc b/src/libexpr-c/nix_api_expr.cc +index 612cc3c84..fc4fe73e4 100644 +--- a/src/libexpr-c/nix_api_expr.cc ++++ b/src/libexpr-c/nix_api_expr.cc +@@ -221,6 +221,26 @@ nix_err nix_eval_state_builder_set_realise_hook( + NIXC_CATCH_ERRS + } + ++nix_err nix_eval_state_builder_set_fetch_hook( ++ nix_c_context * context, nix_eval_state_builder * builder, nix_fetch_hook hook, void * user_data) ++{ ++ if (context) ++ context->last_err_code = NIX_OK; ++ try { ++ builder->fetchSettings.fetchHook = [hook, user_data](const std::string & request) { ++ std::string response; ++ auto setResponse = [](const char * start, unsigned int n, void * data) { ++ static_cast(data)->assign(start, n); ++ }; ++ nix_c_context hookContext; ++ if (hook(user_data, &hookContext, request.c_str(), setResponse, &response) != NIX_OK) ++ throw nix::Error("%s", hookContext.last_err.value_or("fetch hook failed")); ++ return response; ++ }; ++ } ++ NIXC_CATCH_ERRS ++} ++ + EvalState * nix_eval_state_build(nix_c_context * context, nix_eval_state_builder * builder) + { + if (context) +diff --git a/src/libexpr-c/nix_api_expr.h b/src/libexpr-c/nix_api_expr.h +index 1578539ea..04a829c8c 100644 +--- a/src/libexpr-c/nix_api_expr.h ++++ b/src/libexpr-c/nix_api_expr.h +@@ -274,6 +274,40 @@ typedef nix_err (*nix_realise_hook)( + nix_err nix_eval_state_builder_set_realise_hook( + nix_c_context * context, nix_eval_state_builder * builder, nix_realise_hook hook, void * user_data); + ++/** ++ * @brief Callback that fetches an input missing from the store ++ * @ingroup libexpr_init ++ * ++ * @param[in] user_data The pointer passed to nix_eval_state_builder_set_fetch_hook(). ++ * @param[out] context Stores the error message when the fetch fails. ++ * @param[in] request_json `{"builtin": "fetchTree" | "fetchTarball" | "fetchurl", "args": {...}, "storePath"?: "..."}`. ++ * @param[in] callback Called with the response `{"storePath": "...", "attrs"?: {...}}`. ++ * @param[in] callback_user_data Passed to the callback unchanged. ++ * @return NIX_OK once the response's store path is valid in the store, an error code otherwise. ++ */ ++typedef nix_err (*nix_fetch_hook)( ++ void * user_data, ++ nix_c_context * context, ++ const char * request_json, ++ nix_get_string_callback callback, ++ void * callback_user_data); ++ ++/** ++ * @brief Let the embedder fetch inputs that are not in the store ++ * @ingroup libexpr_init ++ * ++ * The hook replaces the built-in fetchers for every input except `path`, and the download of ++ * `fetchurl` and `fetchTarball`. user_data must stay valid as long as the state built from this builder. ++ * ++ * @param[out] context Optional, stores error information ++ * @param[in] builder The builder to modify. ++ * @param[in] hook The callback that fetches the input. ++ * @param[in] user_data Passed to the hook unchanged. ++ * @return NIX_OK if successful, an error code otherwise. ++ */ ++nix_err nix_eval_state_builder_set_fetch_hook( ++ nix_c_context * context, nix_eval_state_builder * builder, nix_fetch_hook hook, void * user_data); ++ + /** + * @brief Create a new Nix language evaluator state + * @ingroup libexpr_init +diff --git a/src/libexpr/primops/fetchTree.cc b/src/libexpr/primops/fetchTree.cc +index 0e092927b..90f59640d 100644 +--- a/src/libexpr/primops/fetchTree.cc ++++ b/src/libexpr/primops/fetchTree.cc +@@ -6,6 +6,7 @@ + #include "nix/expr/fetch-tree.hh" + #include "nix/store/store-api.hh" + #include "nix/fetchers/fetchers.hh" ++#include "nix/fetchers/fetch-settings.hh" + #include "nix/store/filetransfer.hh" + #include "nix/fetchers/registry.hh" + #include "nix/fetchers/tarball.hh" +@@ -13,6 +14,7 @@ + #include "nix/expr/value-to-json.hh" + #include "nix/fetchers/fetch-to-store.hh" + #include "nix/fetchers/input-cache.hh" ++#include "nix/util/json-utils.hh" + + #include + +@@ -471,6 +473,35 @@ static RegisterPrimOp primop_fetchFinalTree({ + .internal = true, + }); + ++static StorePath fetchThroughHook( ++ EvalState & state, ++ const std::string & who, ++ const std::string & url, ++ const std::string & name, ++ const std::optional & expectedHash, ++ const std::optional & expectedPath) ++{ ++ auto args = nlohmann::json{{"url", url}, {"name", name}}; ++ if (expectedHash) ++ args["sha256"] = expectedHash->to_string(HashFormat::SRI, true); ++ auto request = nlohmann::json{{"builtin", who}, {"args", args}}; ++ if (expectedPath) ++ request["storePath"] = state.store->printStorePath(*expectedPath); ++ ++ auto response = nlohmann::json::parse(state.fetchSettings.fetchHook(request.dump())); ++ auto storePath = state.store->parseStorePath(getString(valueAt(getObject(response), "storePath"))); ++ if (expectedPath && storePath != *expectedPath) ++ throw Error( ++ "the fetch hook returned '%s' for '%s', expected '%s'", ++ state.store->printStorePath(storePath), ++ url, ++ state.store->printStorePath(*expectedPath)); ++ if (!state.store->isValidPath(storePath)) ++ throw Error( ++ "the fetch hook returned '%s', which is not valid in the store", state.store->printStorePath(storePath)); ++ return storePath; ++} ++ + static void fetch( + EvalState & state, + const PosIdx pos, +@@ -552,8 +583,9 @@ static void fetch( + .debugThrow(); + + // early exit if pinned and already in the store ++ std::optional expectedPath; + if (expectedHash && expectedHash->algo == HashAlgorithm::SHA256) { +- auto expectedPath = state.store->makeFixedOutputPath( ++ expectedPath = state.store->makeFixedOutputPath( + name, + FixedOutputInfo{ + .method = unpack ? FileIngestionMethod::NixArchive : FileIngestionMethod::Flat, +@@ -562,19 +594,24 @@ static void fetch( + + // Try to get the path from the local store or substituters + try { +- state.store->ensurePath(expectedPath); +- debug("using substituted/cached path '%s' for '%s'", state.store->printStorePath(expectedPath), *url); +- state.allowAndSetStorePathString(expectedPath, v); ++ state.store->ensurePath(*expectedPath); ++ debug("using substituted/cached path '%s' for '%s'", state.store->printStorePath(*expectedPath), *url); ++ state.allowAndSetStorePathString(*expectedPath, v); + return; + } catch (Error & e) { + debug( + "substitution of '%s' failed, will try to download: %s", +- state.store->printStorePath(expectedPath), ++ state.store->printStorePath(*expectedPath), + e.what()); + // Fall through to download + } + } + ++ if (state.fetchSettings.fetchHook) { ++ state.allowAndSetStorePathString(fetchThroughHook(state, who, *url, name, expectedHash, expectedPath), v); ++ return; ++ } ++ + if (unpack) { + auto attrs = fetchers::Attrs{ + {"type", "tarball"}, +diff --git a/src/libfetchers/fetchers.cc b/src/libfetchers/fetchers.cc +index 6d7266e09..977eaa3cf 100644 +--- a/src/libfetchers/fetchers.cc ++++ b/src/libfetchers/fetchers.cc +@@ -297,6 +297,62 @@ std::pair, Input> Input::getAccessor(const Settings & settin + } + } + ++static ref ++storeObjectAccessor(const Settings & settings, Store & store, const Input & input, const StorePath & storePath) ++{ ++ auto accessor = store.requireStoreObjectAccessor(storePath); ++ ++ accessor->fingerprint = input.getFingerprint(store); ++ ++ // Store a cache entry for the tree so later fetches ++ // can reuse the existing nar instead of copying the unpacked ++ // input back into the store on every evaluation. ++ if (accessor->fingerprint) { ++ settings.getCache()->upsert( ++ makeSourcePathToHashCacheKey( ++ *accessor->fingerprint, ContentAddressMethod::Raw::NixArchive, CanonPath::root), ++ {{"hash", store.queryPathInfo(storePath)->narHash.to_string(HashFormat::SRI, true)}}); ++ } ++ ++ accessor->setPathDisplay("«" + input.to_string() + "»"); ++ ++ return accessor; ++} ++ ++static std::pair, Input> ++fetchThroughHook(const Settings & settings, Store & store, const Input & input) ++{ ++ auto args = input.toAttrs(); ++ args.erase("__final"); ++ auto request = nlohmann::json{{"builtin", "fetchTree"}, {"args", attrsToJSON(args)}}; ++ std::optional expectedPath; ++ if (input.getNarHash()) { ++ expectedPath = input.computeStorePath(store); ++ request["storePath"] = store.printStorePath(*expectedPath); ++ } ++ ++ auto response = nlohmann::json::parse(settings.fetchHook(request.dump())); ++ auto storePath = store.parseStorePath(getString(valueAt(getObject(response), "storePath"))); ++ if (expectedPath && storePath != *expectedPath) ++ throw Error( ++ "the fetch hook returned '%s' for input '%s', expected '%s'", ++ store.printStorePath(storePath), ++ input.to_string(), ++ store.printStorePath(*expectedPath)); ++ store.addTempRoot(storePath); ++ if (!store.isValidPath(storePath)) ++ throw Error("the fetch hook returned '%s', which is not valid in the store", store.printStorePath(storePath)); ++ ++ auto result = input; ++ if (auto attrs = optionalValueAt(getObject(response), "attrs")) { ++ result = Input::fromAttrs(settings, jsonToAttrs(*attrs)); ++ if (input.isFinal()) ++ result.attrs.insert_or_assign("__final", Explicit(true)); ++ } ++ ++ return {storeObjectAccessor(settings, store, result, storePath), std::move(result)}; ++} ++ + std::pair, Input> Input::getAccessorUnchecked(const Settings & settings, Store & store) const + { + // FIXME: cache the accessor +@@ -325,28 +381,15 @@ std::pair, Input> Input::getAccessorUnchecked(const Settings + + debug("using substituted/cached input '%s' in '%s'", to_string(), store.printStorePath(storePath)); + +- auto accessor = store.requireStoreObjectAccessor(storePath); +- +- accessor->fingerprint = getFingerprint(store); +- +- // Store a cache entry for the substituted tree so later fetches +- // can reuse the existing nar instead of copying the unpacked +- // input back into the store on every evaluation. +- if (accessor->fingerprint) { +- settings.getCache()->upsert( +- makeSourcePathToHashCacheKey( +- *accessor->fingerprint, ContentAddressMethod::Raw::NixArchive, CanonPath::root), +- {{"hash", store.queryPathInfo(storePath)->narHash.to_string(HashFormat::SRI, true)}}); +- } +- +- accessor->setPathDisplay("«" + to_string() + "»"); +- +- return {accessor, *this}; ++ return {storeObjectAccessor(settings, store, *this, storePath), *this}; + } catch (Error & e) { + debug("substitution of input '%s' failed: %s", to_string(), e.what()); + } + } + ++ if (settings.fetchHook && scheme->schemeName() != "path") ++ return fetchThroughHook(settings, store, *this); ++ + /* Acquire a path lock on this input. Note that fetching the same input in parallel is supposed to be safe (it's up + * to the fetchers to guarantee this), so this is merely intended to avoid work duplication. Note that we don't need + * this when substituting the input. */ +diff --git a/src/libfetchers/include/nix/fetchers/fetch-settings.hh b/src/libfetchers/include/nix/fetchers/fetch-settings.hh +index bb9a67d06..8d64c6ef0 100644 +--- a/src/libfetchers/include/nix/fetchers/fetch-settings.hh ++++ b/src/libfetchers/include/nix/fetchers/fetch-settings.hh +@@ -6,6 +6,7 @@ + #include "nix/util/ref.hh" + #include "nix/util/sync.hh" + ++#include + #include + #include + +@@ -161,6 +162,13 @@ struct Settings : public Config + + const ref srcToStore = createSrcToStore(); + ++ /** ++ * Fetches an input that is not in the store instead of the built-in ++ * fetchers. Takes and returns a JSON request and response; on return ++ * the response's `storePath` is valid. Throws `Error` on failure. ++ */ ++ std::function fetchHook; ++ + + private: + void anchor() override; +-- +2.55.0 + From 552132a1120ac07fa7021a53bed8b4dae481bcc4 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 15:08:44 +0200 Subject: [PATCH 06/19] deps: chore: nix-bindings with fetch hook --- backend/Cargo.lock | 4 ++-- cli/Cargo.lock | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/backend/Cargo.lock b/backend/Cargo.lock index de8ec3e90..9f9f33ecd 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -4746,7 +4746,7 @@ dependencies = [ [[package]] name = "nix-bindings" version = "0.2352.0" -source = "git+https://github.com/DerDennisOP/nix-bindings?branch=feat%2Fnix-2.35.3#b8aadf8622140447fa65ef4905a90da05c0510cc" +source = "git+https://github.com/DerDennisOP/nix-bindings?branch=feat%2Fnix-2.35.3#109cb88a1f8d3189cf0e8b7399afe01ef4916be2" dependencies = [ "nix-bindings-sys", ] @@ -4754,7 +4754,7 @@ dependencies = [ [[package]] name = "nix-bindings-sys" version = "0.2352.0" -source = "git+https://github.com/DerDennisOP/nix-bindings?branch=feat%2Fnix-2.35.3#b8aadf8622140447fa65ef4905a90da05c0510cc" +source = "git+https://github.com/DerDennisOP/nix-bindings?branch=feat%2Fnix-2.35.3#109cb88a1f8d3189cf0e8b7399afe01ef4916be2" dependencies = [ "bindgen", "cc", diff --git a/cli/Cargo.lock b/cli/Cargo.lock index e50dda6d3..b7176657d 100644 --- a/cli/Cargo.lock +++ b/cli/Cargo.lock @@ -2228,7 +2228,7 @@ dependencies = [ [[package]] name = "nix-bindings" version = "0.2352.0" -source = "git+https://github.com/DerDennisOP/nix-bindings?branch=feat%2Fnix-2.35.3#b8aadf8622140447fa65ef4905a90da05c0510cc" +source = "git+https://github.com/DerDennisOP/nix-bindings?branch=feat%2Fnix-2.35.3#109cb88a1f8d3189cf0e8b7399afe01ef4916be2" dependencies = [ "nix-bindings-sys", ] @@ -2236,7 +2236,7 @@ dependencies = [ [[package]] name = "nix-bindings-sys" version = "0.2352.0" -source = "git+https://github.com/DerDennisOP/nix-bindings?branch=feat%2Fnix-2.35.3#b8aadf8622140447fa65ef4905a90da05c0510cc" +source = "git+https://github.com/DerDennisOP/nix-bindings?branch=feat%2Fnix-2.35.3#109cb88a1f8d3189cf0e8b7399afe01ef4916be2" dependencies = [ "bindgen", "cc", From 65a8e98f2f3f90f3858a90e5c34db6fd8ee12d85 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 15:18:26 +0200 Subject: [PATCH 07/19] eval: feat: ask the parent for missing inputs --- backend/gradient-eval/src/eval_worker.rs | 114 +++++++-- backend/gradient-eval/src/ipc.rs | 49 +++- backend/gradient-eval/src/nix_eval.rs | 33 ++- .../src/pool/{import.rs => dependencies.rs} | 14 +- backend/gradient-eval/src/pool/driver.rs | 15 +- backend/gradient-eval/src/pool/mod.rs | 87 ++----- backend/gradient-eval/src/pool/transport.rs | 242 ++++++++++++------ backend/gradient-eval/src/pool/workers.rs | 38 +-- backend/gradient-sources/src/evaluator.rs | 14 +- .../src/fakes/derivation_resolver.rs | 4 +- backend/gradient-worker/options.json | 13 + backend/gradient-worker/src/config/options.rs | 5 + backend/gradient-worker/src/executor/eval.rs | 4 +- .../gradient-worker/src/executor/import.rs | 8 +- backend/gradient-worker/src/main.rs | 6 +- .../src/worker_pool/resolver.rs | 55 ++-- cli/src/commands/base.rs | 2 +- cli/src/commands/eval.rs | 11 +- 18 files changed, 468 insertions(+), 246 deletions(-) rename backend/gradient-eval/src/pool/{import.rs => dependencies.rs} (65%) diff --git a/backend/gradient-eval/src/eval_worker.rs b/backend/gradient-eval/src/eval_worker.rs index bcb0d7536..4481213c9 100644 --- a/backend/gradient-eval/src/eval_worker.rs +++ b/backend/gradient-eval/src/eval_worker.rs @@ -14,7 +14,7 @@ use crate::frames::Frames; use crate::ipc::{ EVAL_IPC_VERSION, EvalRequest, EvalResponse, ResolvedItem, decode_request, read_frame, }; -use crate::nix_eval::{NixEvaluator, RealiseHook, StatsReader}; +use crate::nix_eval::{FetchHook, NixEvaluator, RealiseHook, StatsReader}; pub(crate) struct ParentPipe { reader: Mutex, @@ -34,36 +34,59 @@ impl ParentPipe { pub(crate) fn read_frame(&self) -> std::io::Result>> { if self.broken.load(Ordering::Acquire) { return Err(std::io::Error::other( - "eval pipe broke during a build request", + "eval pipe broke during a request to the parent", )); } read_frame(&mut *self.reader.lock().unwrap_or_else(PoisonError::into_inner)) } pub(crate) fn request_build(&self, derived_paths: &[String]) -> Result<(), String> { + let need = EvalResponse::NeedsBuild { + derived_paths: derived_paths.to_vec(), + }; + match self.ask(&need, "a build")? { + EvalRequest::BuildDone { error: None } => Ok(()), + EvalRequest::BuildDone { error: Some(error) } => Err(error), + other => Err(self.break_on(Some(other), "a build")), + } + } + + pub(crate) fn request_input(&self, request: &str) -> Result { + let need = EvalResponse::NeedsInput { + request: request.to_owned(), + }; + match self.ask(&need, "an input")? { + EvalRequest::InputDone { + response: Some(response), + error: None, + } => Ok(response), + EvalRequest::InputDone { + error: Some(error), .. + } => Err(error), + other => Err(self.break_on(Some(other), "an input")), + } + } + + fn ask(&self, need: &EvalResponse, what: &str) -> Result { self.frames - .send(&EvalResponse::NeedsBuild { - derived_paths: derived_paths.to_vec(), - }) - .map_err(|e| format!("eval pipe closed while asking for a build: {e}"))?; + .send(need) + .map_err(|e| format!("eval pipe closed while asking for {what}: {e}"))?; let answer = match self.read_frame() { Ok(Some(payload)) => decode_request(&payload).ok(), Ok(None) | Err(_) => None, }; - match answer { - Some(EvalRequest::BuildDone { error: None }) => Ok(()), - Some(EvalRequest::BuildDone { error: Some(error) }) => Err(error), - other => { - self.broken.store(true, Ordering::Release); - error!(?other, "eval worker: expected a build result"); - Err("eval pipe closed while waiting for a build".to_string()) - } - } + answer.ok_or_else(|| self.break_on(None, what)) + } + + fn break_on(&self, answer: Option, what: &str) -> String { + self.broken.store(true, Ordering::Release); + error!(?answer, "eval worker: expected the answer for {what}"); + format!("eval pipe closed while waiting for {what}") } } -pub fn run_eval_worker(collect_stats: bool) -> std::io::Result<()> { +pub fn run_eval_worker(collect_stats: bool, fetch_through_parent: bool) -> std::io::Result<()> { let mut stdout = std::io::stdout(); stdout.write_all(&[EVAL_IPC_VERSION])?; stdout.flush()?; @@ -72,7 +95,8 @@ pub fn run_eval_worker(collect_stats: bool) -> std::io::Result<()> { std::io::stdin(), Frames::new(stdout, nix_bindings::EvalStats::default()), )); - let evaluator = match NixEvaluator::new(Some(realise_through(Arc::clone(&pipe)))) { + let fetch_hook = fetch_through_parent.then(|| fetch_through(Arc::clone(&pipe))); + let evaluator = match NixEvaluator::new(Some(realise_through(Arc::clone(&pipe))), fetch_hook) { Ok(e) => Some(e), Err(e) => { error!( @@ -113,6 +137,14 @@ where Box::new(move |derived_paths| pipe.request_build(derived_paths)) } +fn fetch_through(pipe: Arc>) -> FetchHook +where + R: Read + Send + 'static, + W: Write + Send + 'static, +{ + Box::new(move |request| pipe.request_input(request)) +} + const STATS_TICK: std::time::Duration = std::time::Duration::from_secs(1); fn tick_stats( @@ -168,8 +200,8 @@ fn answer_requests( trace!("eval worker shutting down on request"); return Ok(()); } - EvalRequest::BuildDone { .. } => EvalResponse::Err { - message: "a build result arrived outside a build request".to_string(), + EvalRequest::BuildDone { .. } | EvalRequest::InputDone { .. } => EvalResponse::Err { + message: "an answer arrived outside a request to the parent".to_string(), }, EvalRequest::Plan { repository, @@ -218,12 +250,12 @@ fn answer_requests( ) }), EvalRequest::FetchInput { - locked, + request, git_ssh_command, } => with_evaluator(evaluator, |ev| { or_err( - ev.fetch_tree(&locked, git_ssh_command.as_deref()) - .map(|store_path| EvalResponse::FetchOk { store_path }), + ev.fetch_tree(&request, git_ssh_command.as_deref()) + .map(|response| EvalResponse::FetchOk { response }), ) }), EvalRequest::Fingerprint { @@ -336,11 +368,12 @@ fn response_kind(resp: &EvalResponse) -> String { format!("FingerprintOk({})", fingerprint.is_some()) } EvalResponse::CheckpointOk => "CheckpointOk".to_string(), - EvalResponse::FetchOk { store_path } => format!("FetchOk({store_path})"), + EvalResponse::FetchOk { response } => format!("FetchOk({response})"), EvalResponse::Stats { delta } => format!("Stats({} thunks)", delta.nr_thunks), EvalResponse::NeedsBuild { derived_paths } => { format!("NeedsBuild({} paths)", derived_paths.len()) } + EvalResponse::NeedsInput { request } => format!("NeedsInput({request})"), EvalResponse::Err { message } => format!("Err({message})"), } } @@ -528,6 +561,41 @@ mod tests { ); } + #[test] + fn request_input_returns_the_parent_answer() { + let request = r#"{"builtin":"fetchTree","args":{"type":"tarball","url":"https://x"}}"#; + let pipe = pipe_over(&[EvalRequest::InputDone { + response: Some(r#"{"storePath":"/nix/store/aaaa-source"}"#.into()), + error: None, + }]); + + let fetched = pipe.request_input(request); + + assert_eq!(fetched, Ok(r#"{"storePath":"/nix/store/aaaa-source"}"#.to_string())); + let sent = sent(pipe); + assert!( + matches!(&sent[..], [EvalResponse::NeedsInput { request: r }] if r == request), + "{sent:?}" + ); + } + + #[test] + fn a_wrong_answer_to_an_input_request_breaks_the_pipe() { + let pipe = pipe_over(&[ + EvalRequest::BuildDone { error: None }, + EvalRequest::Shutdown, + ]); + + assert_eq!( + pipe.request_input("{}"), + Err("eval pipe closed while waiting for an input".to_string()) + ); + assert!( + pipe.read_frame().is_err(), + "a broken pipe ends the request loop" + ); + } + #[test] fn parse_warnings_keeps_multiline_warning() { let captured = "warning: the following insecure packages:\n - foo-1.2.3\nKnown issues:\n - CVE-1234\ntrace: unrelated\n"; diff --git a/backend/gradient-eval/src/ipc.rs b/backend/gradient-eval/src/ipc.rs index 7b12227b2..1d51e7c5f 100644 --- a/backend/gradient-eval/src/ipc.rs +++ b/backend/gradient-eval/src/ipc.rs @@ -14,7 +14,7 @@ use crate::stats::StatsDelta; /// The version must be bumped whenever the frame layout or a type's rkyv shape changes. /// A mismatch can only happen when the binary is replaced mid-run. -pub const EVAL_IPC_VERSION: u8 = 10; +pub const EVAL_IPC_VERSION: u8 = 11; pub const MAX_FRAME_BYTES: u32 = 64 * 1024 * 1024; @@ -85,7 +85,7 @@ pub enum EvalRequest { input_overrides: Vec<(String, String)>, }, FetchInput { - locked: String, + request: String, #[serde(default)] git_ssh_command: Option, }, @@ -93,6 +93,12 @@ pub enum EvalRequest { #[serde(default)] error: Option, }, + InputDone { + #[serde(default)] + response: Option, + #[serde(default)] + error: Option, + }, Shutdown, } @@ -117,7 +123,7 @@ pub enum EvalResponse { }, CheckpointOk, FetchOk { - store_path: String, + response: String, }, Stats { delta: StatsDelta, @@ -125,6 +131,9 @@ pub enum EvalResponse { NeedsBuild { derived_paths: Vec, }, + NeedsInput { + request: String, + }, Err { message: String, }, @@ -231,14 +240,14 @@ mod tests { #[test] fn a_fetch_request_and_its_frames_round_trip() { let req = EvalRequest::FetchInput { - locked: r#"{"type":"github","owner":"o","repo":"r","rev":"abc","narHash":"sha256-x"}"# + request: r#"{"builtin":"fetchTree","args":{"type":"github","owner":"o","repo":"r"}}"# .into(), git_ssh_command: Some("ssh -i key".into()), }; let back = decode_request(&encode_request(&req).unwrap()).unwrap(); assert!( - matches!(back, EvalRequest::FetchInput { ref locked, ref git_ssh_command } - if locked.contains("github") && git_ssh_command.as_deref() == Some("ssh -i key")) + matches!(back, EvalRequest::FetchInput { ref request, ref git_ssh_command } + if request.contains("github") && git_ssh_command.as_deref() == Some("ssh -i key")) ); let stats = EvalResponse::Stats { @@ -276,6 +285,34 @@ mod tests { assert_eq!(json["kind"], "needs_build"); } + #[test] + fn needs_input_and_input_done_round_trip() { + let request = r#"{"builtin":"fetchurl","args":{"url":"https://example.org/x"}}"#; + let needs = EvalResponse::NeedsInput { + request: request.into(), + }; + let back = decode_response(&encode_response(&needs).unwrap()).unwrap(); + assert!(matches!(back, EvalResponse::NeedsInput { request: r } if r == request)); + + let done = EvalRequest::InputDone { + response: Some(r#"{"storePath":"/nix/store/aaaa-x"}"#.into()), + error: None, + }; + let back = decode_request(&encode_request(&done).unwrap()).unwrap(); + assert!(matches!( + back, + EvalRequest::InputDone { response: Some(r), error: None } if r.contains("aaaa-x") + )); + + let json = serde_json::to_value(EvalRequest::InputDone { + response: None, + error: Some("404".into()), + }) + .unwrap(); + assert_eq!(json["op"], "input_done"); + assert_eq!(serde_json::to_value(&needs).unwrap()["kind"], "needs_input"); + } + #[test] fn decode_survives_misaligned_input() { let bytes = encode_request(&EvalRequest::Shutdown).unwrap(); diff --git a/backend/gradient-eval/src/nix_eval.rs b/backend/gradient-eval/src/nix_eval.rs index d78461383..fcc17115a 100644 --- a/backend/gradient-eval/src/nix_eval.rs +++ b/backend/gradient-eval/src/nix_eval.rs @@ -10,7 +10,7 @@ use anyhow::{Context as _, Result}; use nix_bindings::flake::{FetchersSettings, FlakeSettings}; use nix_bindings::{Context, EvalState, EvalStateBuilder, Store}; -pub use nix_bindings::RealiseHook; +pub use nix_bindings::{FetchHook, RealiseHook}; pub struct NixEvaluator { ctx: Arc, @@ -20,6 +20,19 @@ pub struct NixEvaluator { state: EvalState, } +const FETCH_REQUEST: &str = r#"json: + let + request = builtins.fromJSON json; + fetch = builtins.getAttr request.builtin { inherit (builtins) fetchTree fetchTarball fetchurl; }; + fetched = fetch request.args; + locked = builtins.intersectAttrs { narHash = null; rev = null; revCount = null; lastModified = null; } fetched; + in + builtins.toJSON ( + if request.builtin == "fetchTree" + then { storePath = fetched.outPath; attrs = request.args // locked; } + else { storePath = fetched; } + )"#; + pub struct StatsReader<'ev>(&'ev EvalState); // SAFETY: `stats_with` reads only atomic counters and `GC_get_heap_usage_safe`, so a reader may run beside the evaluating thread. @@ -84,7 +97,7 @@ fn ensure_store_writable() {} impl NixEvaluator { #[allow(clippy::arc_with_non_send_sync)] - pub fn new(realise_hook: Option) -> Result { + pub fn new(realise_hook: Option, fetch_hook: Option) -> Result { ensure_store_writable(); let ctx = Arc::new(Context::new().context("nix context init")?); ctx.set_setting("show-trace", "true")?; @@ -102,6 +115,9 @@ impl NixEvaluator { if let Some(hook) = realise_hook { builder = builder.set_realise_hook(hook)?; } + if let Some(hook) = fetch_hook { + builder = builder.set_fetch_hook(hook)?; + } let state = builder.build().context("nix eval state build")?; Ok(NixEvaluator { @@ -123,20 +139,17 @@ impl NixEvaluator { StatsReader(&self.state) } - pub fn fetch_tree(&self, locked: &str, git_ssh_command: Option<&str>) -> Result { + pub fn fetch_tree(&self, request: &str, git_ssh_command: Option<&str>) -> Result { self.ctx.set_log_format("internal-json")?; let _restore = SshCommand::set(git_ssh_command); - let fetched = self.fetch_tree_out_path(locked); + let fetched = self.answer_fetch(request); self.ctx.set_log_format("raw-with-logs")?; fetched } - fn fetch_tree_out_path(&self, locked: &str) -> Result { - let fetch = self.state.eval_from_string( - "json: (builtins.fetchTree (builtins.fromJSON json)).outPath", - "/", - )?; - let json = self.state.make_string(locked)?; + fn answer_fetch(&self, request: &str) -> Result { + let fetch = self.state.eval_from_string(FETCH_REQUEST, "/")?; + let json = self.state.make_string(request)?; Ok(fetch.call(&json)?.as_string()?) } diff --git a/backend/gradient-eval/src/pool/import.rs b/backend/gradient-eval/src/pool/dependencies.rs similarity index 65% rename from backend/gradient-eval/src/pool/import.rs rename to backend/gradient-eval/src/pool/dependencies.rs index 3bf7ad4bf..c5bfafef2 100644 --- a/backend/gradient-eval/src/pool/import.rs +++ b/backend/gradient-eval/src/pool/dependencies.rs @@ -7,20 +7,26 @@ use async_trait::async_trait; #[async_trait] -pub trait ImportBuilder: Send + Sync { +pub trait EvalDependencies: Send + Sync { async fn build_imports(&self, derived_paths: Vec) -> Result<(), String>; + + async fn fetch_input(&self, request: String) -> Result; } -pub struct RefuseImports(pub &'static str); +pub struct RefuseDependencies(pub &'static str); #[async_trait] -impl ImportBuilder for RefuseImports { +impl EvalDependencies for RefuseDependencies { async fn build_imports(&self, _derived_paths: Vec) -> Result<(), String> { Err(format!( "import from derivation is not available during {}", self.0 )) } + + async fn fetch_input(&self, _request: String) -> Result { + Err(format!("fetching inputs is not available during {}", self.0)) + } } #[cfg(test)] @@ -29,7 +35,7 @@ mod tests { #[tokio::test] async fn refusing_imports_names_the_caller() { - let refused = RefuseImports("the local evaluation") + let refused = RefuseDependencies("the local evaluation") .build_imports(vec!["/nix/store/aaaa-src.drv^out".into()]) .await; diff --git a/backend/gradient-eval/src/pool/driver.rs b/backend/gradient-eval/src/pool/driver.rs index 80d12f3b2..574c4e325 100644 --- a/backend/gradient-eval/src/pool/driver.rs +++ b/backend/gradient-eval/src/pool/driver.rs @@ -14,7 +14,7 @@ use std::sync::{Arc, Mutex}; use crate::ipc::EvalRequest; -use super::import::RefuseImports; +use super::dependencies::RefuseDependencies; use super::transport::{EvalWorker, SpawnConfig}; pub async fn run_eval_driver( @@ -29,9 +29,10 @@ pub async fn run_eval_driver( let live = Arc::new(Mutex::new(HashSet::new())); let mut worker = EvalWorker::spawn( &SpawnConfig { - cache_dir: eval_cache_dir, + cache_dir: eval_cache_dir.to_owned(), trace_dir: None, metrics: eval_metrics, + fetch_through_parent: false, }, live, Arc::default(), @@ -39,7 +40,7 @@ pub async fn run_eval_driver( .await .context("spawn eval worker for driver")?; - let imports = RefuseImports("the eval driver"); + let dependencies = RefuseDependencies("the eval driver"); for line in text.lines().filter(|l| !l.trim().is_empty()) { let req: EvalRequest = serde_json::from_str(line).with_context(|| format!("parse request: {line}"))?; @@ -54,7 +55,7 @@ pub async fn run_eval_driver( wildcards, input_overrides, } => worker - .plan(repository, wildcards, input_overrides, &imports) + .plan(repository, wildcards, input_overrides, &dependencies) .await .map(|(shards, errors)| json!({"kind": "plan_ok", "shards": shards, "errors": errors})), EvalRequest::List { @@ -63,7 +64,7 @@ pub async fn run_eval_driver( only, input_overrides, } => worker - .list(repository, wildcards, only, input_overrides, &imports) + .list(repository, wildcards, only, input_overrides, &dependencies) .await .map(|l| { json!({ @@ -91,8 +92,8 @@ pub async fn run_eval_driver( EvalRequest::FetchInput { .. } => { Err(anyhow::anyhow!("the eval driver does not fetch inputs")) } - EvalRequest::BuildDone { .. } => Err(anyhow::anyhow!( - "the eval driver answers build requests itself" + EvalRequest::BuildDone { .. } | EvalRequest::InputDone { .. } => Err(anyhow::anyhow!( + "the eval driver answers requests to the parent itself" )), }; diff --git a/backend/gradient-eval/src/pool/mod.rs b/backend/gradient-eval/src/pool/mod.rs index c1befb5a6..a021943dc 100644 --- a/backend/gradient-eval/src/pool/mod.rs +++ b/backend/gradient-eval/src/pool/mod.rs @@ -4,10 +4,10 @@ * SPDX-License-Identifier: AGPL-3.0-only */ +pub mod dependencies; pub mod download; mod driver; pub mod fan_out; -pub mod import; mod live_thunks; pub mod memory; mod transport; @@ -24,21 +24,23 @@ use self::fan_out::{ DiscoveryCall, MAX_CRASH_ATTEMPTS, crashed_listing, discovery_calls, entry_point_of, pooled_fan_out, }; -use self::transport::{EvalErrorResponse, Listing}; +use self::transport::{Listing, SpawnConfig}; use self::workers::EvalWorkerPool; pub use self::download::{DownloadSink, forward_nix_log}; pub use self::driver::run_eval_driver; pub use self::fan_out::{CallStats, Discovery, FoundItems}; -pub use self::import::{ImportBuilder, RefuseImports}; +pub use self::dependencies::{EvalDependencies, RefuseDependencies}; pub use self::live_thunks::LiveThunks; pub use self::memory::{MemoryGuard, memory_guard_bytes, min_free_ram_guard, total_memory_bytes}; +pub use self::transport::EvalWorker; pub use self::workers::default_workers; pub(crate) fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { mutex.lock().unwrap_or_else(PoisonError::into_inner) } +#[derive(Debug, Clone)] pub struct EvalPoolConfig { pub workers: usize, pub max_rss: u64, @@ -46,6 +48,7 @@ pub struct EvalPoolConfig { pub cache_dir: String, pub trace_dir: Option, pub metrics: bool, + pub fetch_through_parent: bool, } #[derive(Debug)] @@ -62,9 +65,12 @@ impl EvalPool { config.workers, config.max_rss, config.memory, - config.cache_dir.clone(), - config.trace_dir, - config.metrics, + SpawnConfig { + cache_dir: config.cache_dir.clone(), + trace_dir: config.trace_dir, + metrics: config.metrics, + fetch_through_parent: config.fetch_through_parent, + }, thunks, )); if tokio::runtime::Handle::try_current().is_ok() { @@ -110,7 +116,7 @@ impl EvalPool { repository: String, attrs: Vec, overrides: &[(String, String)], - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, found: FoundItems<'_>, ) -> Result { let requested = attrs.clone(); @@ -127,7 +133,7 @@ impl EvalPool { }) .collect(); let mut discovery = self - .run_calls(repository, calls, &[], overrides, imports, &recording) + .run_calls(repository, calls, &[], overrides, dependencies, &recording) .await?; let listed = listed.into_inner().unwrap_or_default(); @@ -150,7 +156,7 @@ impl EvalPool { repository: String, wildcards: Vec, overrides: &[(String, String)], - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, found: FoundItems<'_>, ) -> Result { *lock(&self.patterns) = wildcards @@ -166,7 +172,7 @@ impl EvalPool { repository.clone(), wildcards.clone(), overrides.to_vec(), - imports, + dependencies, ) .await { @@ -191,7 +197,7 @@ impl EvalPool { ); let mut discovery = self - .run_calls(repository, calls, &excludes, overrides, imports, found) + .run_calls(repository, calls, &excludes, overrides, dependencies, found) .await?; discovery.errors.extend(plan_errors); discovery.errors.sort_unstable(); @@ -205,7 +211,7 @@ impl EvalPool { calls: Vec, excludes: &[String], overrides: &[(String, String)], - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, found: FoundItems<'_>, ) -> Result { let listed_attrs = Mutex::new(HashSet::::new()); @@ -217,7 +223,7 @@ impl EvalPool { let (listed_attrs, warnings, errors, calls_seen) = (&listed_attrs, &warnings, &errors, &calls_seen); pooled_fan_out(self.workers.max(), calls, |call| async move { - let (listing, call_stats) = self.list_shard(repo, call, overrides, imports).await?; + let (listing, call_stats) = self.list_shard(repo, call, overrides, dependencies).await?; let fresh: Vec = { let mut listed = lock(listed_attrs); listing @@ -259,12 +265,12 @@ impl EvalPool { repository: &str, call: DiscoveryCall, overrides: &[(String, String)], - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, ) -> Result<(Listing, Option)> { let mut attempt = 0; loop { match self - .list_once(repository, call.clone(), overrides, imports) + .list_once(repository, call.clone(), overrides, dependencies) .await { Ok((listing, stats)) => return Ok((listing, Some(stats))), @@ -286,7 +292,7 @@ impl EvalPool { repository: &str, call: DiscoveryCall, overrides: &[(String, String)], - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, ) -> Result<(Listing, CallStats)> { let bucket = entry_point_of( call.wildcards @@ -302,7 +308,7 @@ impl EvalPool { call.wildcards, call.only, overrides.to_vec(), - imports, + dependencies, ) .await { @@ -363,23 +369,6 @@ impl EvalPool { } } } - - pub async fn fetch_input( - &self, - locked: String, - git_ssh_command: Option, - sink: Arc, - ) -> Result { - let mut worker = self.workers.acquire().await?; - let fetched = worker.fetch_input(locked, git_ssh_command, sink).await; - if fetched - .as_ref() - .is_err_and(|e| !e.is::()) - { - worker.mark_dead(); - } - fetched - } } #[cfg(test)] @@ -396,6 +385,7 @@ mod tests { cache_dir: String::new(), trace_dir: None, metrics: false, + fetch_through_parent: false, }) } @@ -425,7 +415,7 @@ mod tests { "repo".into(), vec!["packages.x86_64-linux.hello".into()], &[], - &RefuseImports("the test"), + &RefuseDependencies("the test"), &move |items| lock(found_ref).extend(items), ) .await @@ -454,7 +444,7 @@ mod tests { "repo".into(), vec!["packages.x86_64-linux.nope".into()], &[], - &RefuseImports("the test"), + &RefuseDependencies("the test"), &|_| {}, ) .await @@ -463,29 +453,4 @@ mod tests { assert_eq!(discovery.errors.len(), 1, "{:?}", discovery.errors); assert_eq!(discovery.errors[0].attr, "packages.x86_64-linux.nope"); } - - struct NullSink; - - impl DownloadSink for NullSink { - fn transfer(&self, _id: u64, _done: u64, _expected: u64) {} - } - - #[tokio::test] - async fn a_failed_fetch_keeps_the_worker_unless_its_pipe_broke() { - let pool = pool(1); - pool.workers.push_for_test(replying_worker( - &EvalResponse::Err { - message: "404".into(), - }, - "fetch-404", - )); - - let err = pool - .fetch_input("{}".into(), None, Arc::new(NullSink)) - .await - .expect_err("a refused fetch is an error"); - - assert_eq!(err.to_string(), "eval worker: 404"); - assert_eq!(pool.workers.idle_count(), 1, "an answered error keeps it"); - } } diff --git a/backend/gradient-eval/src/pool/transport.rs b/backend/gradient-eval/src/pool/transport.rs index c57e0655c..873aa823f 100644 --- a/backend/gradient-eval/src/pool/transport.rs +++ b/backend/gradient-eval/src/pool/transport.rs @@ -21,7 +21,8 @@ use crate::ipc::{ use crate::stats::StatsDelta; use super::download::{DownloadSink, DownloadSlot, forward_nix_log}; -use super::import::{ImportBuilder, RefuseImports}; +use super::EvalPoolConfig; +use super::dependencies::{EvalDependencies, RefuseDependencies}; use super::live_thunks::LiveThunks; use super::lock; @@ -37,22 +38,27 @@ const SHUTDOWN_GRACE: Duration = Duration::from_secs(5); const EXIT_PROBE_TIMEOUT: Duration = Duration::from_secs(2); -pub(super) struct SpawnConfig<'a> { - pub(super) cache_dir: &'a str, - pub(super) trace_dir: Option<&'a str>, +#[derive(Debug)] +pub(super) struct SpawnConfig { + pub(super) cache_dir: String, + pub(super) trace_dir: Option, pub(super) metrics: bool, + pub(super) fetch_through_parent: bool, } /// The subprocess is this binary re-executed, so the flags and the cache home form the contract /// with `run_eval_worker`. -fn eval_subprocess_command(exe: &Path, config: &SpawnConfig<'_>) -> Command { +fn eval_subprocess_command(exe: &Path, config: &SpawnConfig) -> Command { let mut command = Command::new(exe); command.arg("--eval-subprocess"); if config.metrics { command.arg("--eval-metrics"); } - command.env("NIX_CACHE_HOME", config.cache_dir); - if let Some(dir) = config.trace_dir { + if config.fetch_through_parent { + command.arg("--eval-fetch-through-parent"); + } + command.env("NIX_CACHE_HOME", &config.cache_dir); + if let Some(dir) = &config.trace_dir { command.env("GRADIENT_WORKER_LOG_TRACE_DIR", dir); } for &(k, v) in eval_worker_stats_env(config.metrics) { @@ -94,7 +100,7 @@ pub(super) struct Listing { pub(super) stats: Option, } -pub(super) struct EvalWorker { +pub struct EvalWorker { child: Child, stdin: ChildStdin, stdout: BufReader, @@ -120,8 +126,18 @@ impl Drop for PidGuard { } impl EvalWorker { + pub async fn spawn_for_fetch(config: &EvalPoolConfig) -> Result { + let spawn = SpawnConfig { + cache_dir: config.cache_dir.clone(), + trace_dir: config.trace_dir.clone(), + metrics: false, + fetch_through_parent: false, + }; + Self::spawn(&spawn, Arc::default(), Arc::default()).await + } + pub(super) async fn spawn( - config: &SpawnConfig<'_>, + config: &SpawnConfig, live: Arc>>, thunks: Arc, ) -> Result { @@ -236,7 +252,7 @@ impl EvalWorker { self.stdin.flush().await.context("flush eval worker stdin") } - async fn recv(&mut self, imports: &dyn ImportBuilder) -> Result { + async fn recv(&mut self, dependencies: &dyn EvalDependencies) -> Result { loop { match self.recv_frame().await? { EvalResponse::Stats { delta } => { @@ -244,14 +260,23 @@ impl EvalWorker { } EvalResponse::NeedsBuild { derived_paths } => { let error = tokio::select! { - built = imports.build_imports(derived_paths) => built.err(), - status = self.child.wait() => anyhow::bail!( - "eval worker exited during an import: {}", - status.map_or_else(|e| e.to_string(), |s| s.to_string()) - ), + built = dependencies.build_imports(derived_paths) => built.err(), + status = self.child.wait() => anyhow::bail!(exited_during("an import", status)), }; self.send(&EvalRequest::BuildDone { error }).await?; } + EvalResponse::NeedsInput { request } => { + let fetched = tokio::select! { + fetched = dependencies.fetch_input(request) => fetched, + status = self.child.wait() => anyhow::bail!(exited_during("an input fetch", status)), + }; + let (response, error) = match fetched { + Ok(response) => (Some(response), None), + Err(error) => (None, Some(error)), + }; + self.send(&EvalRequest::InputDone { response, error }) + .await?; + } other => return Ok(other), } } @@ -313,11 +338,11 @@ impl EvalWorker { &mut self, req: EvalRequest, what: &'static str, - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, extract: impl FnOnce(EvalResponse) -> std::result::Result>, ) -> Result { self.send(&req).await?; - let resp = self.recv(imports).await?; + let resp = self.recv(dependencies).await?; self.in_flight = false; match extract(resp) { Ok(v) => Ok(v), @@ -333,7 +358,7 @@ impl EvalWorker { repository: String, wildcards: Vec, input_overrides: Vec<(String, String)>, - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, ) -> Result<(Vec, Vec)> { self.call( EvalRequest::Plan { @@ -342,7 +367,7 @@ impl EvalWorker { input_overrides, }, "Plan", - imports, + dependencies, |resp| match resp { EvalResponse::PlanOk { shards, errors } => Ok((shards, errors)), other => Err(Box::new(other)), @@ -357,7 +382,7 @@ impl EvalWorker { wildcards: Vec, only: Option>, input_overrides: Vec<(String, String)>, - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, ) -> Result { self.call( EvalRequest::List { @@ -367,7 +392,7 @@ impl EvalWorker { input_overrides, }, "List", - imports, + dependencies, |resp| match resp { EvalResponse::ListOk { items, @@ -399,7 +424,7 @@ impl EvalWorker { input_overrides, }, "Fingerprint", - &RefuseImports("the fingerprint"), + &RefuseDependencies("the fingerprint"), |resp| match resp { EvalResponse::FingerprintOk { fingerprint } => Ok(fingerprint), other => Err(Box::new(other)), @@ -419,7 +444,7 @@ impl EvalWorker { input_overrides, }, "Checkpoint", - &RefuseImports("the eval-cache checkpoint"), + &RefuseDependencies("the eval-cache checkpoint"), |resp| match resp { EvalResponse::CheckpointOk => Ok(()), other => Err(Box::new(other)), @@ -428,9 +453,9 @@ impl EvalWorker { .await } - pub(super) async fn fetch_input( + pub async fn fetch( &mut self, - locked: String, + request: String, git_ssh_command: Option, sink: Arc, ) -> Result { @@ -438,13 +463,13 @@ impl EvalWorker { let fetched = self .call( EvalRequest::FetchInput { - locked, + request, git_ssh_command, }, "FetchInput", - &RefuseImports("an input fetch"), + &RefuseDependencies("an input fetch"), |resp| match resp { - EvalResponse::FetchOk { store_path } => Ok(store_path), + EvalResponse::FetchOk { response } => Ok(response), other => Err(Box::new(other)), }, ) @@ -453,7 +478,7 @@ impl EvalWorker { fetched } - pub(super) async fn shutdown(mut self) { + pub async fn shutdown(mut self) { let pid = self.child.id(); trace!(pid, "sending Shutdown to eval worker"); if let Err(e) = self.send(&EvalRequest::Shutdown).await { @@ -485,6 +510,13 @@ impl EvalWorker { } } +fn exited_during(what: &str, status: std::io::Result) -> String { + format!( + "eval worker exited during {what}: {}", + status.map_or_else(|e| e.to_string(), |s| s.to_string()) + ) +} + fn forward_stderr(stderr: ChildStderr, downloads: DownloadSlot) { #[expect( clippy::disallowed_methods, @@ -508,14 +540,48 @@ mod tests { use super::*; use crate::ipc::{decode_request, encode_response}; - struct RecordingImports(Mutex>>); + #[derive(Default)] + struct RecordingDependencies { + builds: Mutex>>, + inputs: Mutex>, + } #[async_trait::async_trait] - impl ImportBuilder for RecordingImports { + impl EvalDependencies for RecordingDependencies { async fn build_imports(&self, derived_paths: Vec) -> Result<(), String> { - lock(&self.0).push(derived_paths); + lock(&self.builds).push(derived_paths); Err("no builder here".into()) } + + async fn fetch_input(&self, request: String) -> Result { + lock(&self.inputs).push(request); + Ok(r#"{"storePath":"/nix/store/aaaa-source"}"#.into()) + } + } + + fn asking_worker( + tag: &str, + request: &EvalRequest, + need: &EvalResponse, + answer_len: usize, + done: &EvalResponse, + ) -> (EvalWorker, std::path::PathBuf) { + let request_len = framed(&encode_request(request).unwrap()).len(); + let need = temp_file(&format!("{tag}-need"), &framed(&encode_response(need).unwrap())); + let done = temp_file(&format!("{tag}-done"), &framed(&encode_response(done).unwrap())); + let answer = std::env::temp_dir().join(format!( + "gradient-transport-test-{}-{tag}-answer", + std::process::id() + )); + let mut cmd = Command::new("sh"); + cmd.arg("-c").arg(format!( + "head -c {request_len} >/dev/null; cat '{}'; head -c {answer_len} >'{}'; cat '{}'; cat >/dev/null", + need.display(), + answer.display(), + done.display() + )); + let worker = EvalWorker::from_command(cmd, Arc::default()).expect("spawn sh"); + (worker, answer) } fn framed(payload: &[u8]) -> Vec { @@ -536,60 +602,45 @@ mod tests { path } - #[tokio::test] - async fn recv_answers_a_build_request_and_returns_the_final_frame() { - let plan = EvalRequest::Plan { + fn plan_request() -> EvalRequest { + EvalRequest::Plan { repository: "repo".into(), wildcards: vec![], input_overrides: vec![], - }; + } + } + + fn plan_ok() -> EvalResponse { + EvalResponse::PlanOk { + shards: vec![], + errors: vec![], + } + } + + #[tokio::test] + async fn recv_answers_a_build_request_and_returns_the_final_frame() { let expected_answer = EvalRequest::BuildDone { error: Some("no builder here".into()), }; - let plan_len = framed(&encode_request(&plan).unwrap()).len(); - let answer_len = framed(&encode_request(&expected_answer).unwrap()).len(); - let needs = temp_file( - "needs", - &framed( - &encode_response(&EvalResponse::NeedsBuild { - derived_paths: vec!["/nix/store/aaaa-src.drv^out".into()], - }) - .unwrap(), - ), - ); - let done = temp_file( - "done", - &framed( - &encode_response(&EvalResponse::PlanOk { - shards: vec![], - errors: vec![], - }) - .unwrap(), - ), + let (mut worker, answer) = asking_worker( + "build", + &plan_request(), + &EvalResponse::NeedsBuild { + derived_paths: vec!["/nix/store/aaaa-src.drv^out".into()], + }, + framed(&encode_request(&expected_answer).unwrap()).len(), + &plan_ok(), ); - let answer = std::env::temp_dir().join(format!( - "gradient-transport-test-{}-answer", - std::process::id() - )); - - let mut cmd = Command::new("sh"); - cmd.arg("-c").arg(format!( - "head -c {plan_len} >/dev/null; cat '{}'; head -c {answer_len} >'{}'; cat '{}'; cat >/dev/null", - needs.display(), - answer.display(), - done.display() - )); - let mut worker = EvalWorker::from_command(cmd, Arc::default()).expect("spawn sh"); - let imports = RecordingImports(Mutex::new(Vec::new())); + let dependencies = RecordingDependencies::default(); let (shards, errors) = worker - .plan("repo".into(), vec![], vec![], &imports) + .plan("repo".into(), vec![], vec![], &dependencies) .await .expect("the final frame after the build request"); assert!(shards.is_empty() && errors.is_empty()); assert_eq!( - imports.0.into_inner().unwrap_or_default(), + dependencies.builds.into_inner().unwrap_or_default(), vec![vec!["/nix/store/aaaa-src.drv^out".to_string()]] ); let written = std::fs::read(&answer).expect("the answer frame"); @@ -601,14 +652,51 @@ mod tests { assert!(!worker.in_flight(), "the call ended in lockstep"); } + #[tokio::test] + async fn recv_answers_an_input_request_through_the_dependencies() { + let request = r#"{"builtin":"fetchTree","args":{"type":"tarball","url":"https://x"}}"#; + let expected_answer = EvalRequest::InputDone { + response: Some(r#"{"storePath":"/nix/store/aaaa-source"}"#.into()), + error: None, + }; + let (mut worker, answer) = asking_worker( + "input", + &plan_request(), + &EvalResponse::NeedsInput { + request: request.into(), + }, + framed(&encode_request(&expected_answer).unwrap()).len(), + &plan_ok(), + ); + let dependencies = RecordingDependencies::default(); + + worker + .plan("repo".into(), vec![], vec![], &dependencies) + .await + .expect("the final frame after the input request"); + + assert_eq!( + dependencies.inputs.into_inner().unwrap_or_default(), + vec![request.to_string()] + ); + let written = std::fs::read(&answer).expect("the answer frame"); + let back = decode_request(&written[4..]).expect("a request frame"); + assert!( + matches!(&back, EvalRequest::InputDone { response: Some(r), error: None } if r.contains("aaaa-source")), + "{back:?}" + ); + assert!(!worker.in_flight(), "the call ended in lockstep"); + } + #[test] fn the_spawned_command_asks_for_the_eval_subprocess() { let plain = eval_subprocess_command( Path::new("/usr/bin/gradient"), &SpawnConfig { - cache_dir: "/tmp/cache", + cache_dir: "/tmp/cache".into(), trace_dir: None, metrics: false, + fetch_through_parent: false, }, ); let args: Vec = plain @@ -632,9 +720,10 @@ mod tests { let metrics = eval_subprocess_command( Path::new("/usr/bin/gradient"), &SpawnConfig { - cache_dir: "/tmp/cache", - trace_dir: Some("/tmp/trace"), + cache_dir: "/tmp/cache".into(), + trace_dir: Some("/tmp/trace".into()), metrics: true, + fetch_through_parent: true, }, ); let args: Vec = metrics @@ -646,7 +735,8 @@ mod tests { args, vec![ "--eval-subprocess".to_string(), - "--eval-metrics".to_string() + "--eval-metrics".to_string(), + "--eval-fetch-through-parent".to_string() ] ); assert!( diff --git a/backend/gradient-eval/src/pool/workers.rs b/backend/gradient-eval/src/pool/workers.rs index 1775446ca..a2b3ded6b 100644 --- a/backend/gradient-eval/src/pool/workers.rs +++ b/backend/gradient-eval/src/pool/workers.rs @@ -31,28 +31,19 @@ pub(super) struct EvalWorkerPool { max: usize, max_rss: u64, guard: MemoryGuard, - settings: WorkerSettings, + spawn: SpawnConfig, shutting_down: Arc, live: Arc>>, under_pressure: Arc, thunks: Arc, } -#[derive(Debug)] -struct WorkerSettings { - cache_dir: String, - trace_dir: Option, - metrics: bool, -} - impl EvalWorkerPool { pub(super) fn new( workers: usize, max_rss: u64, guard: MemoryGuard, - cache_dir: String, - trace_dir: Option, - metrics: bool, + spawn: SpawnConfig, thunks: Arc, ) -> Self { let max = workers.max(1); @@ -62,11 +53,7 @@ impl EvalWorkerPool { max, max_rss, guard, - settings: WorkerSettings { - cache_dir, - trace_dir, - metrics, - }, + spawn, shutting_down: Arc::new(AtomicBool::new(false)), live: Arc::new(Mutex::new(HashSet::new())), under_pressure: Arc::new(AtomicBool::new(false)), @@ -123,11 +110,7 @@ impl EvalWorkerPool { } None => { break EvalWorker::spawn( - &SpawnConfig { - cache_dir: &self.settings.cache_dir, - trace_dir: self.settings.trace_dir.as_deref(), - metrics: self.settings.metrics, - }, + &self.spawn, Arc::clone(&self.live), Arc::clone(&self.thunks), ) @@ -360,9 +343,12 @@ pub(super) mod tests { workers, 2 * GIB, guard, - String::new(), - None, - true, + SpawnConfig { + cache_dir: String::new(), + trace_dir: None, + metrics: true, + fetch_through_parent: false, + }, Arc::default(), ) } @@ -413,7 +399,7 @@ pub(super) mod tests { "repo".into(), vec![], vec![], - &super::super::RefuseImports("a test"), + &super::super::RefuseDependencies("a test"), ); tokio::time::timeout(Duration::from_millis(200), call) .await @@ -447,7 +433,7 @@ pub(super) mod tests { "repo".into(), vec![], vec![], - &super::super::RefuseImports("a test"), + &super::super::RefuseDependencies("a test"), ) .await .expect("plan"); diff --git a/backend/gradient-sources/src/evaluator.rs b/backend/gradient-sources/src/evaluator.rs index 1262c77a2..416b0247a 100644 --- a/backend/gradient-sources/src/evaluator.rs +++ b/backend/gradient-sources/src/evaluator.rs @@ -24,20 +24,26 @@ pub struct FlakeDiscovery { } #[async_trait] -pub trait ImportBuilder: Send + Sync { +pub trait EvalDependencies: Send + Sync { async fn build_imports(&self, derived_paths: Vec) -> Result<(), String>; + + async fn fetch_input(&self, request: String) -> Result; } -pub struct RefuseImports(pub &'static str); +pub struct RefuseDependencies(pub &'static str); #[async_trait] -impl ImportBuilder for RefuseImports { +impl EvalDependencies for RefuseDependencies { async fn build_imports(&self, _derived_paths: Vec) -> Result<(), String> { Err(format!( "import from derivation is not available during {}", self.0 )) } + + async fn fetch_input(&self, _request: String) -> Result { + Err(format!("fetching inputs is not available during {}", self.0)) + } } /// Production impls must run inside `tokio::task::spawn_blocking`. The embedded Nix C API with @@ -49,7 +55,7 @@ pub trait DerivationResolver: Send + Sync + std::fmt::Debug + 'static { repository: String, wildcards: Vec, overrides: &[(String, String)], - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, found: FoundDerivations<'_>, ) -> Result; diff --git a/backend/gradient-test-support/src/fakes/derivation_resolver.rs b/backend/gradient-test-support/src/fakes/derivation_resolver.rs index d0b86697c..341e46f83 100644 --- a/backend/gradient-test-support/src/fakes/derivation_resolver.rs +++ b/backend/gradient-test-support/src/fakes/derivation_resolver.rs @@ -8,7 +8,7 @@ use anyhow::{Result, anyhow}; use async_trait::async_trait; use gradient_derivation::Derivation; use gradient_sources::{ - AttrError, DerivationResolver, FlakeDiscovery, FoundDerivations, ImportBuilder, + AttrError, DerivationResolver, FlakeDiscovery, FoundDerivations, EvalDependencies, ResolvedDerivation, }; use std::collections::HashMap; @@ -88,7 +88,7 @@ impl DerivationResolver for FakeDerivationResolver { repository: String, _wildcards: Vec, _overrides: &[(String, String)], - _imports: &dyn ImportBuilder, + _dependencies: &dyn EvalDependencies, found: FoundDerivations<'_>, ) -> Result { let attrs = self diff --git a/backend/gradient-worker/options.json b/backend/gradient-worker/options.json index 2522ee8be..0e8b2dd44 100644 --- a/backend/gradient-worker/options.json +++ b/backend/gradient-worker/options.json @@ -175,6 +175,19 @@ "false" ] }, + { + "group": "", + "field": "eval_fetch_through_parent", + "flag": "eval-fetch-through-parent", + "env": "GRADIENT_WORKER_EVAL_FETCH_THROUGH_PARENT", + "type": "bool", + "description": "Let the eval subprocess ask its parent for missing inputs (internal - do not set manually).", + "default": null, + "values": [ + "true", + "false" + ] + }, { "group": "", "field": "eval_driver", diff --git a/backend/gradient-worker/src/config/options.rs b/backend/gradient-worker/src/config/options.rs index c46c219dc..b09762c83 100644 --- a/backend/gradient-worker/src/config/options.rs +++ b/backend/gradient-worker/src/config/options.rs @@ -106,6 +106,10 @@ pub struct WorkerConfig { #[arg(hide = true)] pub eval_subprocess: bool, + /// Let the eval subprocess ask its parent for missing inputs (internal - do not set manually). + #[arg(hide = true)] + pub eval_fetch_through_parent: bool, + /// Drive one eval subprocess over the production rkyv transport from a JSONL request file. /// This internal harness serves the NixOS VM integration test. #[arg(hide = true)] @@ -147,6 +151,7 @@ impl Default for WorkerConfig { idle_timeout: None, max_lifetime: None, eval_subprocess: false, + eval_fetch_through_parent: false, eval_driver: None, capabilities: CapabilitiesArgs::default(), system: SystemArgs::default(), diff --git a/backend/gradient-worker/src/executor/eval.rs b/backend/gradient-worker/src/executor/eval.rs index 8fc159dec..7cf7d21b8 100644 --- a/backend/gradient-worker/src/executor/eval.rs +++ b/backend/gradient-worker/src/executor/eval.rs @@ -1552,7 +1552,7 @@ mod tests { _: String, _: Vec, _: &[(String, String)], - _: &dyn gradient_sources::ImportBuilder, + _: &dyn gradient_sources::EvalDependencies, _: gradient_sources::FoundDerivations<'_>, ) -> Result { std::future::pending().await @@ -1737,7 +1737,7 @@ mod tests { _: String, _: Vec, _: &[(String, String)], - _: &dyn gradient_sources::ImportBuilder, + _: &dyn gradient_sources::EvalDependencies, found: gradient_sources::FoundDerivations<'_>, ) -> Result { found(vec![("hello".into(), Ok((self.drv_path.clone(), vec![])))]); diff --git a/backend/gradient-worker/src/executor/import.rs b/backend/gradient-worker/src/executor/import.rs index 5ec63c53c..6a8bb4a45 100644 --- a/backend/gradient-worker/src/executor/import.rs +++ b/backend/gradient-worker/src/executor/import.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use async_trait::async_trait; use futures::future::{BoxFuture, FutureExt as _, Shared, join_all}; use gradient_derivation::{Derivation, parse_drv}; -use gradient_sources::ImportBuilder; +use gradient_sources::EvalDependencies; use gradient_util::store_path::strip_store_prefix; use gradient_util::sync::Mutex; use gradient_wire::types::ImportOutcome; @@ -81,7 +81,7 @@ impl<'a> EvalImportBuilder<'a> { } #[async_trait] -impl ImportBuilder for EvalImportBuilder<'_> { +impl EvalDependencies for EvalImportBuilder<'_> { async fn build_imports(&self, derived_paths: Vec) -> Result<(), String> { let mut drv_paths: Vec = derived_paths .iter() @@ -97,6 +97,10 @@ impl ImportBuilder for EvalImportBuilder<'_> { let imports: Vec<_> = drv_paths.into_iter().map(|drv| self.import(drv)).collect(); join_all(imports).await.into_iter().collect() } + + async fn fetch_input(&self, _request: String) -> Result { + Err("this worker's evaluator fetches its inputs itself".to_owned()) + } } async fn import_derivation( diff --git a/backend/gradient-worker/src/main.rs b/backend/gradient-worker/src/main.rs index 1b4168863..2735a2186 100644 --- a/backend/gradient-worker/src/main.rs +++ b/backend/gradient-worker/src/main.rs @@ -59,7 +59,11 @@ fn main() -> Result<()> { }); if config.eval_subprocess { - return nix::eval_worker::run_eval_worker(config.eval.metrics).map_err(anyhow::Error::from); + return nix::eval_worker::run_eval_worker( + config.eval.metrics, + config.eval_fetch_through_parent, + ) + .map_err(anyhow::Error::from); } if let Some(requests_path) = config.eval_driver.clone() { diff --git a/backend/gradient-worker/src/worker_pool/resolver.rs b/backend/gradient-worker/src/worker_pool/resolver.rs index 95dd5b175..04ce1db8e 100644 --- a/backend/gradient-worker/src/worker_pool/resolver.rs +++ b/backend/gradient-worker/src/worker_pool/resolver.rs @@ -8,9 +8,11 @@ use anyhow::{Context, Result}; use async_trait::async_trait; use gradient_derivation::{Derivation, parse_drv}; use gradient_eval::ipc::ResolvedItem; -use gradient_eval::pool::{CallStats, EvalPool, EvalPoolConfig, LiveThunks, MemoryGuard}; +use gradient_eval::pool::{ + CallStats, EvalPool, EvalPoolConfig, EvalWorker, LiveThunks, MemoryGuard, +}; use gradient_sources::{ - AttrError, DerivationResolver, FlakeDiscovery, FoundDerivations, ImportBuilder, + AttrError, DerivationResolver, EvalDependencies, FlakeDiscovery, FoundDerivations, ResolvedDerivation, }; use gradient_util::store_path::nix_store_path; @@ -23,16 +25,21 @@ use super::input_fetch::{DownloadTarget, InputFetcher}; #[derive(Debug)] pub struct WorkerPoolResolver { pool: Arc, + config: EvalPoolConfig, stats: Arc>, } -struct SourceImports<'a>(&'a dyn ImportBuilder); +struct SourceDependencies<'a>(&'a dyn EvalDependencies); #[async_trait] -impl gradient_eval::pool::ImportBuilder for SourceImports<'_> { +impl gradient_eval::pool::EvalDependencies for SourceDependencies<'_> { async fn build_imports(&self, derived_paths: Vec) -> Result<(), String> { self.0.build_imports(derived_paths).await } + + async fn fetch_input(&self, request: String) -> Result { + self.0.fetch_input(request).await + } } fn item_to_resolved(item: ResolvedItem) -> ResolvedDerivation { @@ -61,15 +68,18 @@ impl WorkerPoolResolver { trace_dir: Option, eval_metrics: bool, ) -> Self { + let config = EvalPoolConfig { + workers: pool_size, + max_rss, + memory, + cache_dir: eval_cache_dir, + trace_dir, + metrics: eval_metrics, + fetch_through_parent: false, + }; Self { - pool: EvalPool::new(EvalPoolConfig { - workers: pool_size, - max_rss, - memory, - cache_dir: eval_cache_dir, - trace_dir, - metrics: eval_metrics, - }), + pool: EvalPool::new(config.clone()), + config, stats: Arc::new(Mutex::new(EvalStatsAccumulator::default())), } } @@ -126,9 +136,20 @@ impl InputFetcher for WorkerPoolResolver { git_ssh_command: Option, target: DownloadTarget, ) -> Result { - self.pool - .fetch_input(locked, git_ssh_command, Arc::new(target)) - .await + let args: serde_json::Value = + serde_json::from_str(&locked).context("parse locked input attributes")?; + let request = serde_json::json!({ "builtin": "fetchTree", "args": args }).to_string(); + let mut worker = EvalWorker::spawn_for_fetch(&self.config).await?; + let fetched = worker + .fetch(request, git_ssh_command, Arc::new(target)) + .await; + worker.shutdown().await; + let response: serde_json::Value = + serde_json::from_str(&fetched?).context("parse fetch response")?; + response["storePath"] + .as_str() + .map(str::to_owned) + .context("fetch response without a store path") } } @@ -139,7 +160,7 @@ impl DerivationResolver for WorkerPoolResolver { repository: String, wildcards: Vec, overrides: &[(String, String)], - imports: &dyn ImportBuilder, + dependencies: &dyn EvalDependencies, found: FoundDerivations<'_>, ) -> Result { let found_items = |items: Vec| { @@ -151,7 +172,7 @@ impl DerivationResolver for WorkerPoolResolver { repository, wildcards, overrides, - &SourceImports(imports), + &SourceDependencies(dependencies), &found_items, ) .await?; diff --git a/cli/src/commands/base.rs b/cli/src/commands/base.rs index 15998a7f5..a2621ff4f 100644 --- a/cli/src/commands/base.rs +++ b/cli/src/commands/base.rs @@ -172,7 +172,7 @@ pub fn run() -> std::io::Result<()> { #[cfg(feature = "eval")] if let Some(metrics) = eval::subprocess_flags(&std::env::args().collect::>()) { - return gradient_eval::eval_worker::run_eval_worker(metrics); + return gradient_eval::eval_worker::run_eval_worker(metrics, false); } let cli = Cli::parse(); diff --git a/cli/src/commands/eval.rs b/cli/src/commands/eval.rs index e1b96757e..a02e4a220 100644 --- a/cli/src/commands/eval.rs +++ b/cli/src/commands/eval.rs @@ -8,7 +8,9 @@ use crate::commands::attr_spec; use clap::Args; use etcetera::base_strategy::{BaseStrategy, choose_base_strategy}; use gradient_eval::ipc::ResolvedItem; -use gradient_eval::pool::{EvalPool, EvalPoolConfig, MemoryGuard, RefuseImports, default_workers}; +use gradient_eval::pool::{ + EvalPool, EvalPoolConfig, MemoryGuard, RefuseDependencies, default_workers, +}; use serde::Serialize; use std::io::Write; use std::num::{NonZeroU64, NonZeroUsize}; @@ -101,6 +103,7 @@ async fn evaluate( cache_dir: eval_cache_dir(), trace_dir: None, metrics: false, + fetch_through_parent: false, }); let out = std::sync::Mutex::new(std::io::stdout()); @@ -111,14 +114,14 @@ async fn evaluate( } }; - let imports = RefuseImports("the local evaluation"); + let dependencies = RefuseDependencies("the local evaluation"); let found = |items: Vec| items.into_iter().map(Job::from_item).for_each(print); let evaluated = if all_concrete(&wildcards) { - pool.list_attrs(flake_ref, wildcards, &[], &imports, &found) + pool.list_attrs(flake_ref, wildcards, &[], &dependencies, &found) .await } else { - pool.list_shards(flake_ref, wildcards, &[], &imports, &found) + pool.list_shards(flake_ref, wildcards, &[], &dependencies, &found) .await }; From f132d52c109d21247279e85031758cb22a682601 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 15:25:46 +0200 Subject: [PATCH 08/19] wire: feat: input requests and input fetches --- backend/gradient-eval/src/eval_worker.rs | 5 +- .../gradient-eval/src/pool/dependencies.rs | 5 +- backend/gradient-eval/src/pool/mod.rs | 5 +- backend/gradient-eval/src/pool/transport.rs | 12 +- backend/gradient-proto/src/handler/inbound.rs | 6 + backend/gradient-sources/src/evaluator.rs | 5 +- .../src/fakes/derivation_resolver.rs | 2 +- .../src/fakes/job_reporter.rs | 51 ++++++- backend/gradient-wire/schema/v35.txt | 6 + backend/gradient-wire/src/messages/client.rs | 30 +++- backend/gradient-wire/src/messages/server.rs | 41 +++++- backend/gradient-wire/src/traits.rs | 10 ++ backend/gradient-wire/src/types.rs | 6 + .../gradient-worker-client/src/correlation.rs | 129 +++++++++++++++++- .../gradient-worker/src/executor/compress.rs | 1 + backend/gradient-worker/src/proto/job.rs | 94 ++++++++++++- .../src/worker/message_loop.rs | 39 +++++- docs/src/contributors/proto/jobs.md | 1 + docs/src/contributors/proto/messages.md | 6 + 19 files changed, 431 insertions(+), 23 deletions(-) diff --git a/backend/gradient-eval/src/eval_worker.rs b/backend/gradient-eval/src/eval_worker.rs index 4481213c9..1218eddd5 100644 --- a/backend/gradient-eval/src/eval_worker.rs +++ b/backend/gradient-eval/src/eval_worker.rs @@ -571,7 +571,10 @@ mod tests { let fetched = pipe.request_input(request); - assert_eq!(fetched, Ok(r#"{"storePath":"/nix/store/aaaa-source"}"#.to_string())); + assert_eq!( + fetched, + Ok(r#"{"storePath":"/nix/store/aaaa-source"}"#.to_string()) + ); let sent = sent(pipe); assert!( matches!(&sent[..], [EvalResponse::NeedsInput { request: r }] if r == request), diff --git a/backend/gradient-eval/src/pool/dependencies.rs b/backend/gradient-eval/src/pool/dependencies.rs index c5bfafef2..f55de019b 100644 --- a/backend/gradient-eval/src/pool/dependencies.rs +++ b/backend/gradient-eval/src/pool/dependencies.rs @@ -25,7 +25,10 @@ impl EvalDependencies for RefuseDependencies { } async fn fetch_input(&self, _request: String) -> Result { - Err(format!("fetching inputs is not available during {}", self.0)) + Err(format!( + "fetching inputs is not available during {}", + self.0 + )) } } diff --git a/backend/gradient-eval/src/pool/mod.rs b/backend/gradient-eval/src/pool/mod.rs index a021943dc..8b5afe8eb 100644 --- a/backend/gradient-eval/src/pool/mod.rs +++ b/backend/gradient-eval/src/pool/mod.rs @@ -27,10 +27,10 @@ use self::fan_out::{ use self::transport::{Listing, SpawnConfig}; use self::workers::EvalWorkerPool; +pub use self::dependencies::{EvalDependencies, RefuseDependencies}; pub use self::download::{DownloadSink, forward_nix_log}; pub use self::driver::run_eval_driver; pub use self::fan_out::{CallStats, Discovery, FoundItems}; -pub use self::dependencies::{EvalDependencies, RefuseDependencies}; pub use self::live_thunks::LiveThunks; pub use self::memory::{MemoryGuard, memory_guard_bytes, min_free_ram_guard, total_memory_bytes}; pub use self::transport::EvalWorker; @@ -223,7 +223,8 @@ impl EvalPool { let (listed_attrs, warnings, errors, calls_seen) = (&listed_attrs, &warnings, &errors, &calls_seen); pooled_fan_out(self.workers.max(), calls, |call| async move { - let (listing, call_stats) = self.list_shard(repo, call, overrides, dependencies).await?; + let (listing, call_stats) = + self.list_shard(repo, call, overrides, dependencies).await?; let fresh: Vec = { let mut listed = lock(listed_attrs); listing diff --git a/backend/gradient-eval/src/pool/transport.rs b/backend/gradient-eval/src/pool/transport.rs index 873aa823f..fdf6c5bb7 100644 --- a/backend/gradient-eval/src/pool/transport.rs +++ b/backend/gradient-eval/src/pool/transport.rs @@ -20,9 +20,9 @@ use crate::ipc::{ }; use crate::stats::StatsDelta; -use super::download::{DownloadSink, DownloadSlot, forward_nix_log}; use super::EvalPoolConfig; use super::dependencies::{EvalDependencies, RefuseDependencies}; +use super::download::{DownloadSink, DownloadSlot, forward_nix_log}; use super::live_thunks::LiveThunks; use super::lock; @@ -567,8 +567,14 @@ mod tests { done: &EvalResponse, ) -> (EvalWorker, std::path::PathBuf) { let request_len = framed(&encode_request(request).unwrap()).len(); - let need = temp_file(&format!("{tag}-need"), &framed(&encode_response(need).unwrap())); - let done = temp_file(&format!("{tag}-done"), &framed(&encode_response(done).unwrap())); + let need = temp_file( + &format!("{tag}-need"), + &framed(&encode_response(need).unwrap()), + ); + let done = temp_file( + &format!("{tag}-done"), + &framed(&encode_response(done).unwrap()), + ); let answer = std::env::temp_dir().join(format!( "gradient-transport-test-{}-{tag}-answer", std::process::id() diff --git a/backend/gradient-proto/src/handler/inbound.rs b/backend/gradient-proto/src/handler/inbound.rs index ca37a5fbb..d7b052ed0 100644 --- a/backend/gradient-proto/src/handler/inbound.rs +++ b/backend/gradient-proto/src/handler/inbound.rs @@ -368,6 +368,12 @@ impl<'a> InboundContext<'a> { } true } + msg @ (ClientMessage::InputRequest { .. } + | ClientMessage::InputFetched { .. } + | ClientMessage::InputFetchProgress { .. }) => { + warn!(peer_id = %self.peer_id, variant = msg.variant_name(), "input fetches are not served yet"); + true + } ClientMessage::EvalMessage { job_id, level, diff --git a/backend/gradient-sources/src/evaluator.rs b/backend/gradient-sources/src/evaluator.rs index 416b0247a..8de1a9bfc 100644 --- a/backend/gradient-sources/src/evaluator.rs +++ b/backend/gradient-sources/src/evaluator.rs @@ -42,7 +42,10 @@ impl EvalDependencies for RefuseDependencies { } async fn fetch_input(&self, _request: String) -> Result { - Err(format!("fetching inputs is not available during {}", self.0)) + Err(format!( + "fetching inputs is not available during {}", + self.0 + )) } } diff --git a/backend/gradient-test-support/src/fakes/derivation_resolver.rs b/backend/gradient-test-support/src/fakes/derivation_resolver.rs index 341e46f83..e45196283 100644 --- a/backend/gradient-test-support/src/fakes/derivation_resolver.rs +++ b/backend/gradient-test-support/src/fakes/derivation_resolver.rs @@ -8,7 +8,7 @@ use anyhow::{Result, anyhow}; use async_trait::async_trait; use gradient_derivation::Derivation; use gradient_sources::{ - AttrError, DerivationResolver, FlakeDiscovery, FoundDerivations, EvalDependencies, + AttrError, DerivationResolver, EvalDependencies, FlakeDiscovery, FoundDerivations, ResolvedDerivation, }; use std::collections::HashMap; diff --git a/backend/gradient-test-support/src/fakes/job_reporter.rs b/backend/gradient-test-support/src/fakes/job_reporter.rs index 6a7b2fbcb..c081a2743 100644 --- a/backend/gradient-test-support/src/fakes/job_reporter.rs +++ b/backend/gradient-test-support/src/fakes/job_reporter.rs @@ -10,8 +10,8 @@ use gradient_util::sync::Mutex; use gradient_wire::messages::{ BuildMetrics, BuildOutput, CachedPath, DiscoveredDerivation, EvalMessageLevel, QueryMode, }; -use gradient_wire::traits::{EvalProgressSink, JobReporter}; -use gradient_wire::types::{EvalProgress, ImportOutcome}; +use gradient_wire::traits::{EvalProgressSink, InputProgressSink, JobReporter}; +use gradient_wire::types::{EvalProgress, ImportOutcome, InputOutcome}; use std::sync::Arc; #[derive(Debug, Clone)] @@ -36,6 +36,9 @@ pub enum ReportedEvent { ImportRequested { drv_paths: Vec, }, + InputRequested { + request: String, + }, Building { build_id: String, }, @@ -64,6 +67,8 @@ pub struct RecordingJobReporter { pub known_drv_paths: Vec, pub upstream: std::collections::HashMap, pub import_outcomes: std::collections::HashMap, + pub input_outcomes: Vec<(String, InputOutcome)>, + pub input_progress: Vec<(u64, Option)>, } impl RecordingJobReporter { @@ -91,6 +96,26 @@ impl RecordingJobReporter { self } + pub fn with_input_outcome(mut self, request_part: &str, outcome: InputOutcome) -> Self { + self.input_outcomes.push((request_part.to_owned(), outcome)); + self + } + + pub fn with_input_progress(mut self, bytes_done: u64, bytes_total: Option) -> Self { + self.input_progress.push((bytes_done, bytes_total)); + self + } + + pub fn input_requests(&self) -> Vec { + self.events() + .into_iter() + .filter_map(|e| match e { + ReportedEvent::InputRequested { request } => Some(request), + _ => None, + }) + .collect() + } + pub fn events(&self) -> Vec { self.events.lock().clone() } @@ -263,6 +288,28 @@ impl JobReporter for RecordingJobReporter { Ok(outcome) } + async fn request_input( + &self, + request: String, + progress: Arc, + ) -> Result { + let outcome = self + .input_outcomes + .iter() + .find(|(part, _)| request.contains(part.as_str())) + .map_or_else( + || InputOutcome::Failed { + message: "no input fetch configured".to_owned(), + }, + |(_, outcome)| outcome.clone(), + ); + for &(bytes_done, bytes_total) in &self.input_progress { + progress.transfer(bytes_done, bytes_total); + } + self.record(ReportedEvent::InputRequested { request }); + Ok(outcome) + } + async fn report_building(&mut self, build_id: String) -> Result<()> { self.record(ReportedEvent::Building { build_id }); Ok(()) diff --git a/backend/gradient-wire/schema/v35.txt b/backend/gradient-wire/schema/v35.txt index 81232c875..f04239de1 100644 --- a/backend/gradient-wire/schema/v35.txt +++ b/backend/gradient-wire/schema/v35.txt @@ -28,6 +28,9 @@ client 26:{u64,[0:{{str,u64,u64,str,seq,opt,opt,opt<{str,seq}>}},1:{u64}]} 27:{u64} 28:{str,str,seq} +29:{str,str,str} +30:{str,[0:{str},1:{str}]} +31:{str,u64,opt} server 0:{seq} 1:{{bool,bool,bool,bool,bool,bool,bool,bool},seq,seq<{str,str}>} @@ -56,3 +59,6 @@ server 24:{u64,[0:{},1:{str},2:{str}]} 25:{str,seq<(str,str)>} 26:{str,str,[0:{},1:{str,str,str},2:{str}]} +27:{str,str,[0:{str},1:{str}]} +28:{str,str,str,opt} +29:{str,str,u64,opt} diff --git a/backend/gradient-wire/src/messages/client.rs b/backend/gradient-wire/src/messages/client.rs index e988ff254..0f8c4273f 100644 --- a/backend/gradient-wire/src/messages/client.rs +++ b/backend/gradient-wire/src/messages/client.rs @@ -9,8 +9,8 @@ use bytes::Bytes; use crate::codec::Proto; use crate::types::{ BuildFailureKind, BuildMetrics, BuildProgressPhase, CandidateScore, ClusterAddress, - EvalMessageLevel, EvalProgress, GradientCapabilities, JobKind, JobPhaseSpan, JobUpdateKind, - QueryMode, UploadMetadata, UploadObject, + EvalMessageLevel, EvalProgress, GradientCapabilities, InputOutcome, JobKind, JobPhaseSpan, + JobUpdateKind, QueryMode, UploadMetadata, UploadObject, }; #[derive(Proto, Debug, Clone, PartialEq)] @@ -197,6 +197,26 @@ pub enum ClientMessage { request_id: String, drv_paths: Vec, }, + + #[proto(35)] + InputRequest { + job_id: String, + request_id: String, + request: String, + }, + + #[proto(35)] + InputFetched { + fetch_id: String, + outcome: InputOutcome, + }, + + #[proto(35)] + InputFetchProgress { + fetch_id: String, + bytes_done: u64, + bytes_total: Option, + }, } impl ClientMessage { @@ -216,7 +236,8 @@ impl ClientMessage { | ClientMessage::EvalMessage { job_id, .. } | ClientMessage::QueryKnownDerivations { job_id, .. } | ClientMessage::UploadRequest { job_id, .. } - | ClientMessage::ImportRequest { job_id, .. } => Some(job_id), + | ClientMessage::ImportRequest { job_id, .. } + | ClientMessage::InputRequest { job_id, .. } => Some(job_id), _ => None, } } @@ -252,6 +273,9 @@ impl ClientMessage { ClientMessage::UploadFinished { .. } => "UploadFinished", ClientMessage::UploadCancel { .. } => "UploadCancel", ClientMessage::ImportRequest { .. } => "ImportRequest", + ClientMessage::InputRequest { .. } => "InputRequest", + ClientMessage::InputFetched { .. } => "InputFetched", + ClientMessage::InputFetchProgress { .. } => "InputFetchProgress", } } diff --git a/backend/gradient-wire/src/messages/server.rs b/backend/gradient-wire/src/messages/server.rs index 47baa25c6..be8eb487e 100644 --- a/backend/gradient-wire/src/messages/server.rs +++ b/backend/gradient-wire/src/messages/server.rs @@ -9,8 +9,8 @@ use bytes::Bytes; use crate::codec::Proto; use crate::types::{ CachedPath, ClusterAddress, ClusterMembership, ClusterPeer, CredentialKind, - EvalCachePullOutcome, GradientCapabilities, GrantTarget, ImportOutcome, Job, JobCandidate, - UploadOutcome, + EvalCachePullOutcome, GradientCapabilities, GrantTarget, ImportOutcome, InputOutcome, Job, + JobCandidate, UploadOutcome, }; #[derive(Proto, Debug, Clone, PartialEq)] @@ -168,6 +168,29 @@ pub enum ServerMessage { request_id: String, outcome: ImportOutcome, }, + + #[proto(35)] + InputResult { + job_id: String, + request_id: String, + outcome: InputOutcome, + }, + + #[proto(35)] + FetchInput { + fetch_id: String, + project: String, + request: String, + ssh_key: Option, + }, + + #[proto(35)] + InputProgress { + job_id: String, + request_id: String, + bytes_done: u64, + bytes_total: Option, + }, } impl ServerMessage { @@ -181,7 +204,9 @@ impl ServerMessage { | ServerMessage::NarStreamHeader { job_id, .. } | ServerMessage::EvalCachePullResult { job_id, .. } | ServerMessage::EvalCacheChunk { job_id, .. } - | ServerMessage::ImportResult { job_id, .. } => Some(job_id), + | ServerMessage::ImportResult { job_id, .. } + | ServerMessage::InputResult { job_id, .. } + | ServerMessage::InputProgress { job_id, .. } => Some(job_id), _ => None, } } @@ -215,13 +240,21 @@ impl ServerMessage { ServerMessage::UploadCommitted { .. } => "UploadCommitted", ServerMessage::Authenticate { .. } => "Authenticate", ServerMessage::ImportResult { .. } => "ImportResult", + ServerMessage::InputResult { .. } => "InputResult", + ServerMessage::FetchInput { .. } => "FetchInput", + ServerMessage::InputProgress { .. } => "InputProgress", } } pub fn carries_secret(&self) -> bool { matches!( self, - ServerMessage::Authenticate { .. } | ServerMessage::Credential { .. } + ServerMessage::Authenticate { .. } + | ServerMessage::Credential { .. } + | ServerMessage::FetchInput { + ssh_key: Some(_), + .. + } ) } } diff --git a/backend/gradient-wire/src/traits.rs b/backend/gradient-wire/src/traits.rs index 7752b1e8c..8a2a1affe 100644 --- a/backend/gradient-wire/src/traits.rs +++ b/backend/gradient-wire/src/traits.rs @@ -66,6 +66,12 @@ pub trait JobReporter: Send + Sync { async fn request_import(&self, drv_paths: Vec) -> Result; + async fn request_input( + &self, + request: String, + progress: std::sync::Arc, + ) -> Result; + async fn report_building(&mut self, build_id: String) -> Result<()>; async fn report_build_output( &mut self, @@ -90,6 +96,10 @@ pub trait EvalProgressSink: Send + Sync { async fn report(&self, progress: crate::types::EvalProgress); } +pub trait InputProgressSink: Send + Sync { + fn transfer(&self, bytes_done: u64, bytes_total: Option); +} + #[async_trait] pub trait PeerIdentity: Send + Sync { fn peer_id(&self) -> String; diff --git a/backend/gradient-wire/src/types.rs b/backend/gradient-wire/src/types.rs index 13515df54..2255a0ca2 100644 --- a/backend/gradient-wire/src/types.rs +++ b/backend/gradient-wire/src/types.rs @@ -322,6 +322,12 @@ pub enum ImportOutcome { }, } +#[derive(Proto, Debug, Clone, PartialEq, Eq)] +pub enum InputOutcome { + Fetched { response: String }, + Failed { message: String }, +} + #[derive(Proto, Debug, Clone, PartialEq)] pub struct DerivationOutput { pub name: String, diff --git a/backend/gradient-worker-client/src/correlation.rs b/backend/gradient-worker-client/src/correlation.rs index eb6c0de8d..d6a695a6e 100644 --- a/backend/gradient-worker-client/src/correlation.rs +++ b/backend/gradient-worker-client/src/correlation.rs @@ -14,7 +14,8 @@ use gradient_wire::messages::{ CACHE_QUERY_MAX_PATHS, CACHE_QUERY_TIMEOUT, CACHE_QUERY_WINDOW, CachedPath, ClientMessage, QueryMode, }; -use gradient_wire::types::ImportOutcome; +use gradient_wire::traits::InputProgressSink; +use gradient_wire::types::{ImportOutcome, InputOutcome}; use tokio::sync::oneshot; use crate::connection::ProtoWriter; @@ -154,6 +155,88 @@ pub async fn request_import( }) } +pub struct InputWaiter { + job_id: String, + reply: oneshot::Sender, + progress: Arc, +} + +pub type InputWaiters = Arc>>; + +pub fn register_input_waiter( + waiters: &InputWaiters, + request_id: String, + job_id: String, + progress: Arc, +) -> oneshot::Receiver { + let (reply, rx) = oneshot::channel(); + waiters.lock().insert( + request_id, + InputWaiter { + job_id, + reply, + progress, + }, + ); + rx +} + +pub fn deliver_input_result( + waiters: &InputWaiters, + request_id: &str, + outcome: InputOutcome, +) -> bool { + match waiters.lock().remove(request_id) { + Some(w) => { + let _ = w.reply.send(outcome); + true + } + None => false, + } +} + +pub fn deliver_input_progress( + waiters: &InputWaiters, + request_id: &str, + bytes_done: u64, + bytes_total: Option, +) -> bool { + let Some(progress) = waiters + .lock() + .get(request_id) + .map(|w| Arc::clone(&w.progress)) + else { + return false; + }; + progress.transfer(bytes_done, bytes_total); + true +} + +pub fn forget_input_waiters_for_job(waiters: &InputWaiters, job_id: &str) { + waiters.lock().retain(|_, w| w.job_id != job_id); +} + +pub async fn request_input( + job_id: &str, + writer: &ProtoWriter, + waiters: &InputWaiters, + request: String, + progress: Arc, +) -> Result { + let request_id = uuid::Uuid::now_v7().to_string(); + let rx = register_input_waiter(waiters, request_id.clone(), job_id.to_owned(), progress); + writer + .send(ClientMessage::InputRequest { + job_id: job_id.to_owned(), + request_id, + request, + }) + .await?; + rx.await.map_err(|_| { + anyhow::anyhow!("input waiter dropped - the job ended or the connection closed") + }) +} + #[derive(Clone, Debug)] pub struct AssignmentHandle(Arc>); @@ -371,4 +454,48 @@ mod tests { assert_eq!(rx_a.try_recv(), Err(TryRecvError::Closed)); assert!(waiters.lock().is_empty()); } + + #[derive(Default)] + struct RecordedProgress(Mutex)>>); + + impl InputProgressSink for RecordedProgress { + fn transfer(&self, bytes_done: u64, bytes_total: Option) { + self.0.lock().push((bytes_done, bytes_total)); + } + } + + #[test] + fn an_input_result_reaches_its_waiter() { + use tokio::sync::oneshot::error::TryRecvError; + let waiters: InputWaiters = Arc::new(Mutex::new(HashMap::new())); + let none = Arc::new(RecordedProgress::default()); + let mut rx_a = register_input_waiter(&waiters, "r1".into(), "job-A".into(), none.clone()); + let mut rx_b = register_input_waiter(&waiters, "r2".into(), "job-B".into(), none); + let fetched = InputOutcome::Fetched { + response: r#"{"storePath":"/nix/store/aaaa-source"}"#.into(), + }; + + assert!(deliver_input_result(&waiters, "r2", fetched.clone())); + assert_eq!(rx_b.try_recv(), Ok(fetched.clone())); + assert!(!deliver_input_result(&waiters, "r2", fetched)); + + forget_input_waiters_for_job(&waiters, "job-A"); + assert_eq!(rx_a.try_recv(), Err(TryRecvError::Closed)); + assert!(waiters.lock().is_empty()); + } + + #[test] + fn input_progress_reaches_the_sink_of_its_request() { + let waiters: InputWaiters = Arc::new(Mutex::new(HashMap::new())); + let first = Arc::new(RecordedProgress::default()); + let second = Arc::new(RecordedProgress::default()); + let _rx1 = register_input_waiter(&waiters, "r1".into(), "job".into(), first.clone()); + let _rx2 = register_input_waiter(&waiters, "r2".into(), "job".into(), second.clone()); + + assert!(deliver_input_progress(&waiters, "r2", 512, Some(1024))); + assert!(!deliver_input_progress(&waiters, "gone", 1, None)); + + assert!(first.0.lock().is_empty()); + assert_eq!(*second.0.lock(), vec![(512, Some(1024))]); + } } diff --git a/backend/gradient-worker/src/executor/compress.rs b/backend/gradient-worker/src/executor/compress.rs index 7f0bc819e..097f1dea0 100644 --- a/backend/gradient-worker/src/executor/compress.rs +++ b/backend/gradient-worker/src/executor/compress.rs @@ -225,6 +225,7 @@ mod tests { cache_waiters, Arc::new(Mutex::new(HashMap::new())), Arc::new(Mutex::new(HashMap::new())), + Arc::new(Mutex::new(HashMap::new())), NarReceiver::new(), EvalCacheReceiver::new(), None, diff --git a/backend/gradient-worker/src/proto/job.rs b/backend/gradient-worker/src/proto/job.rs index a5605783c..82642b12d 100644 --- a/backend/gradient-worker/src/proto/job.rs +++ b/backend/gradient-worker/src/proto/job.rs @@ -14,7 +14,7 @@ use gradient_wire::messages::{ EvalCachePullOutcome, EvalMessageLevel, EvalStatsReport, JobPhase, JobUpdateKind, QueryMode, }; use gradient_worker_client::correlation::{ - AssignmentHandle, CacheWaiters, ImportWaiters, KnownDerivationWaiters, + AssignmentHandle, CacheWaiters, ImportWaiters, InputWaiters, KnownDerivationWaiters, cache_query_with_timeout, known_derivations_with_timeout, }; use tracing::debug; @@ -27,9 +27,10 @@ use crate::proto::prefetch::MissingInputs; use crate::proto::progress::{ BuildProgressSink, EvalProgressSender, Progress, Tally, count_delivery, }; -use gradient_wire::traits::{EvalProgressSink, JobReporter}; +use gradient_wire::traits::{EvalProgressSink, InputProgressSink, JobReporter}; use gradient_wire::types::{ - BuildProgressPhase, BuildStage, GrantTarget, ImportOutcome, UploadMetadata, UploadObject, + BuildProgressPhase, BuildStage, GrantTarget, ImportOutcome, InputOutcome, UploadMetadata, + UploadObject, }; use gradient_worker_client::connection::ProtoWriter; use gradient_worker_client::nar_recv::{NarPayload, NarReceiver, NarUnavailable}; @@ -42,6 +43,7 @@ pub struct JobUpdater { pub(crate) cache_waiters: CacheWaiters, pub(crate) known_derivation_waiters: KnownDerivationWaiters, pub(crate) import_waiters: ImportWaiters, + pub(crate) input_waiters: InputWaiters, pub(crate) nar_recv: NarReceiver, pub(crate) eval_cache_recv: EvalCacheReceiver, pub(crate) store: Option>, @@ -93,6 +95,7 @@ impl JobUpdater { cache_waiters: CacheWaiters, known_derivation_waiters: KnownDerivationWaiters, import_waiters: ImportWaiters, + input_waiters: InputWaiters, nar_recv: NarReceiver, eval_cache_recv: EvalCacheReceiver, store: Option>, @@ -106,6 +109,7 @@ impl JobUpdater { cache_waiters, known_derivation_waiters, import_waiters, + input_waiters, nar_recv, eval_cache_recv, store, @@ -126,6 +130,7 @@ impl JobUpdater { cache_waiters: Arc::clone(&self.cache_waiters), known_derivation_waiters: Arc::clone(&self.known_derivation_waiters), import_waiters: Arc::clone(&self.import_waiters), + input_waiters: Arc::clone(&self.input_waiters), nar_recv: self.nar_recv.clone(), eval_cache_recv: self.eval_cache_recv.clone(), store: self.store.clone(), @@ -561,6 +566,21 @@ impl JobReporter for JobUpdater { .await } + async fn request_input( + &self, + request: String, + progress: Arc, + ) -> Result { + gradient_worker_client::correlation::request_input( + &self.job_id, + &self.writer, + &self.input_waiters, + request, + progress, + ) + .await + } + async fn report_building(&mut self, build_id: String) -> Result<()> { self.send_update(JobUpdateKind::Building { build_id }).await } @@ -649,6 +669,7 @@ mod tests { let cache_waiters = Arc::new(Mutex::new(HashMap::new())); let known_derivation_waiters = Arc::new(Mutex::new(HashMap::new())); let import_waiters = Arc::new(Mutex::new(HashMap::new())); + let input_waiters = Arc::new(Mutex::new(HashMap::new())); let nar_recv = NarReceiver::new(); let eval_cache_recv = EvalCacheReceiver::new(); let updater = JobUpdater::new( @@ -658,6 +679,7 @@ mod tests { cache_waiters, known_derivation_waiters, import_waiters, + input_waiters, nar_recv, eval_cache_recv, None, @@ -970,4 +992,70 @@ mod tests { ); pump.abort(); } + + struct NoProgress; + + impl InputProgressSink for NoProgress { + fn transfer(&self, _bytes_done: u64, _bytes_total: Option) {} + } + + #[tokio::test] + async fn request_input_sends_an_input_request_and_returns_the_result() { + use gradient_wire::messages::ServerMessage; + const REQUEST: &str = + r#"{"builtin":"fetchTree","args":{"type":"tarball","url":"https://x"}}"#; + let (conn, server_task, job_id) = server_then_client!("job-input", |sc| { + let ClientMessage::InputRequest { + job_id, + request_id, + request, + } = sc.recv().await.unwrap() + else { + panic!("expected an InputRequest"); + }; + assert_eq!((job_id.as_str(), request.as_str()), ("job-input", REQUEST)); + sc.send(ServerMessage::InputResult { + job_id, + request_id, + outcome: InputOutcome::Fetched { + response: r#"{"storePath":"/nix/store/aaaa-source"}"#.into(), + }, + }) + .await + .unwrap(); + }); + + let (updater, mut reader) = make_updater(job_id, conn); + let waiters = Arc::clone(&updater.input_waiters); + let pump = tokio::spawn(async move { + while let Some(msg) = reader.recv().await { + if let ServerMessage::InputResult { + request_id, + outcome, + .. + } = msg + { + gradient_worker_client::correlation::deliver_input_result( + &waiters, + &request_id, + outcome, + ); + } + } + }); + + let outcome = updater + .request_input(REQUEST.into(), Arc::new(NoProgress)) + .await + .unwrap(); + + assert_eq!( + outcome, + InputOutcome::Fetched { + response: r#"{"storePath":"/nix/store/aaaa-source"}"#.into() + } + ); + server_task.await.unwrap(); + pump.abort(); + } } diff --git a/backend/gradient-worker/src/worker/message_loop.rs b/backend/gradient-worker/src/worker/message_loop.rs index 00f90bc3e..8b7981c85 100644 --- a/backend/gradient-worker/src/worker/message_loop.rs +++ b/backend/gradient-worker/src/worker/message_loop.rs @@ -28,7 +28,7 @@ use crate::proto::job::JobUpdater; use crate::proto::scorer::JobScorer; use gradient_worker_client::connection::{ProtoReader, ProtoWriter}; use gradient_worker_client::correlation::{ - AssignmentHandle, CacheWaiters, ImportWaiters, KnownDerivationWaiters, + AssignmentHandle, CacheWaiters, ImportWaiters, InputWaiters, KnownDerivationWaiters, }; use super::cluster::{ClusterChannels, ClusterHolds, HeldJob}; @@ -236,6 +236,7 @@ pub(super) struct MessageLoopState { cache_waiters: CacheWaiters, known_derivation_waiters: KnownDerivationWaiters, import_waiters: ImportWaiters, + input_waiters: InputWaiters, nar_recv: gradient_worker_client::nar_recv::NarReceiver, eval_cache_recv: crate::proto::eval_cache_recv::EvalCacheReceiver, uploads: gradient_worker_client::upload::UploadClient, @@ -284,6 +285,7 @@ impl MessageLoopState { cache_waiters: Arc::new(Mutex::new(HashMap::new())), known_derivation_waiters: Arc::new(Mutex::new(HashMap::new())), import_waiters: Arc::new(Mutex::new(HashMap::new())), + input_waiters: Arc::new(Mutex::new(HashMap::new())), nar_recv, eval_cache_recv: crate::proto::eval_cache_recv::EvalCacheReceiver::new(), jobs: JobRegistry { @@ -443,6 +445,35 @@ impl MessageLoopState { } => { self.on_import_result(job_id, request_id, outcome); } + ServerMessage::InputResult { + job_id, + request_id, + outcome, + } => { + if !gradient_worker_client::correlation::deliver_input_result( + &self.input_waiters, + &request_id, + outcome, + ) { + debug!(%job_id, %request_id, "InputResult arrived after its waiter cleared"); + } + } + ServerMessage::InputProgress { + request_id, + bytes_done, + bytes_total, + .. + } => { + gradient_worker_client::correlation::deliver_input_progress( + &self.input_waiters, + &request_id, + bytes_done, + bytes_total, + ); + } + ServerMessage::FetchInput { fetch_id, .. } => { + warn!(%fetch_id, "this worker does not run input fetches yet"); + } } Ok(()) } @@ -467,6 +498,10 @@ impl MessageLoopState { &self.import_waiters, &job_id, ); + gradient_worker_client::correlation::forget_input_waiters_for_job( + &self.input_waiters, + &job_id, + ); self.nar_recv.forget_job(&job_id); self.eval_cache_recv.forget_job(&job_id); self.uploads.forget_job(&job_id); @@ -753,6 +788,7 @@ impl MessageLoopState { let job_cache_waiters = Arc::clone(&self.cache_waiters); let job_known_derivation_waiters = Arc::clone(&self.known_derivation_waiters); let job_import_waiters = Arc::clone(&self.import_waiters); + let job_input_waiters = Arc::clone(&self.input_waiters); let job_nar_recv = self.nar_recv.clone(); let job_eval_cache_recv = self.eval_cache_recv.clone(); let job_uploads = self.uploads.clone(); @@ -771,6 +807,7 @@ impl MessageLoopState { job_cache_waiters, job_known_derivation_waiters, job_import_waiters, + job_input_waiters, job_nar_recv, job_eval_cache_recv, Some(job_store), diff --git a/docs/src/contributors/proto/jobs.md b/docs/src/contributors/proto/jobs.md index b84a79e6e..bf9d757f1 100644 --- a/docs/src/contributors/proto/jobs.md +++ b/docs/src/contributors/proto/jobs.md @@ -52,6 +52,7 @@ Flake jobs fetch and evaluate a flake. The server can pick the steps at queue ti - **Eval Cache:** Workers pull the evaluation cache before discovery and push the cache after a walk with derivations. Empty or schemaless caches stay local. - **Attribute Failures:** Workers send them as `EvalMessage { level: Error, source: "nix-eval:" }` before the last `EvalResult`. - **Import From Derivation:** Workers send `ImportRequest { drv_paths }` during discovery and wait for `ImportResult`. +- **Missing Inputs:** Workers send `InputRequest { request }` for an input in neither the store nor the cache, and wait for `InputResult`. The server can send the download as `FetchInput` to a worker with `network`. That worker can push the input under the `fetch_id` and answer with `InputFetched`. - **Walk:** Workers walk an entry point as soon as discovery returned it. - **Waves:** Waves hold up to 256 of the most recently found derivations. Leaf derivations then come early in the walk, and their builds can start mid-evaluation. diff --git a/docs/src/contributors/proto/messages.md b/docs/src/contributors/proto/messages.md index bb0c11400..726a21494 100644 --- a/docs/src/contributors/proto/messages.md +++ b/docs/src/contributors/proto/messages.md @@ -33,6 +33,9 @@ The message definitions live in `backend/gradient-wire/src/messages` and cover e | `UploadGrant` | Upload admission: skip, passthrough (with resume offset), presigned PUT or multipart | `request_id`, `target` | | `UploadCommitted` | Upload outcome: ok, retry or rejected | `request_id`, `outcome` | | `ImportResult` | Reply to `ImportRequest`, sent after the imported derivations finished: completed, failed with the build, or unknown | `job_id`, `request_id`, `outcome` | +| `InputResult` | Reply to `InputRequest`: the fetch answer with the store path, or the download error | `job_id`, `request_id`, `outcome` | +| `InputProgress` | Download bytes of the input fetch an `InputRequest` waits for | `job_id`, `request_id`, `bytes_done`, `bytes_total` | +| `FetchInput` | Download one input for a project and push the result to the cache. Carries the project SSH key only for a worker with `secrets` | `fetch_id`, `project`, `request`, `ssh_key` | ## Worker -> Server @@ -67,6 +70,9 @@ The message definitions live in `backend/gradient-wire/src/messages` and cover e | `UploadFinished` | Marks an upload done and carries NAR metadata | `request_id`, `metadata` | | `UploadCancel` | Cancels an upload in progress | `request_id` | | `ImportRequest` | Builds the given derivations now, because the evaluation needs them | `job_id`, `request_id`, `drv_paths` | +| `InputRequest` | Asks for an input missing from the store and the cache, because the evaluation needs the input | `job_id`, `request_id`, `request` | +| `InputFetchProgress` | Reports download bytes of a `FetchInput` | `fetch_id`, `bytes_done`, `bytes_total` | +| `InputFetched` | Reply to `FetchInput` after the push: the fetch answer or the download error | `fetch_id`, `outcome` | ## Cache Query Modes From 5abab4d91d6a78d81de50f529036098585d13bdb Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 15:33:14 +0200 Subject: [PATCH 09/19] scheduler: feat: dispatch input fetches --- backend/gradient-pool/Cargo.toml | 8 +- backend/gradient-pool/src/input_request.rs | 86 ++++ backend/gradient-pool/src/lib.rs | 2 + backend/gradient-pool/src/session_port.rs | 19 +- backend/gradient-pool/src/worker_caps.rs | 4 + backend/gradient-proto/src/handler/inbound.rs | 73 +++- .../gradient-proto/src/handler/job_events.rs | 38 +- .../src/handler/session_actor.rs | 68 ++- backend/gradient-proto/src/handler/socket.rs | 113 ++++- .../gradient-proto/src/handler/upload/mod.rs | 6 +- .../gradient-scheduler/src/cases/inputs.rs | 390 ++++++++++++++++++ backend/gradient-scheduler/src/cases/mod.rs | 1 + .../gradient-scheduler/src/cases/readers.rs | 4 + .../gradient-scheduler/src/cases/session.rs | 30 +- backend/gradient-scheduler/src/input_waits.rs | 183 ++++++++ .../src/job_handlers/abort.rs | 2 + .../src/job_handlers/input.rs | 223 ++++++++++ .../src/job_handlers/mod.rs | 1 + backend/gradient-scheduler/src/lib.rs | 3 + backend/gradient-scheduler/src/loops/build.rs | 3 + .../gradient-scheduler/src/waiting_state.rs | 37 +- .../src/worker_lifecycle.rs | 3 +- 22 files changed, 1257 insertions(+), 40 deletions(-) create mode 100644 backend/gradient-pool/src/input_request.rs create mode 100644 backend/gradient-scheduler/src/cases/inputs.rs create mode 100644 backend/gradient-scheduler/src/input_waits.rs create mode 100644 backend/gradient-scheduler/src/job_handlers/input.rs diff --git a/backend/gradient-pool/Cargo.toml b/backend/gradient-pool/Cargo.toml index 0b0ac30d1..117e24434 100644 --- a/backend/gradient-pool/Cargo.toml +++ b/backend/gradient-pool/Cargo.toml @@ -14,12 +14,12 @@ workspace = true gradient-types = { workspace = true } gradient-wire = { workspace = true } -chrono = { workspace = true } -serde = { workspace = true } -tokio = { workspace = true, features = ["sync"] } +chrono = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true, features = ["sync"] } tracing = { workspace = true } utoipa = { workspace = true } [dev-dependencies] tokio = { workspace = true, features = ["sync", "macros", "rt"] } -serde_json = { workspace = true } diff --git a/backend/gradient-pool/src/input_request.rs b/backend/gradient-pool/src/input_request.rs new file mode 100644 index 000000000..2d686afaf --- /dev/null +++ b/backend/gradient-pool/src/input_request.rs @@ -0,0 +1,86 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use gradient_types::input::check_repository_url_is_ssh; +use serde_json::Value; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InputRequest { + pub fetch_key: String, + pub store_path: Option, + pub needs_secrets: bool, +} + +impl InputRequest { + pub fn parse(request: &str) -> Result { + let json: Value = + serde_json::from_str(request).map_err(|e| format!("malformed input request: {e}"))?; + let builtin = json["builtin"] + .as_str() + .ok_or("an input request without a builtin")?; + let args = &json["args"]; + if !args.is_object() { + return Err("an input request without arguments".to_owned()); + } + + let store_path = json["storePath"].as_str().map(str::to_owned); + let fetch_key = store_path.clone().unwrap_or_else(|| { + sorted(&serde_json::json!({ "builtin": builtin, "args": args })).to_string() + }); + Ok(Self { + fetch_key, + store_path, + needs_secrets: args["url"] + .as_str() + .is_some_and(check_repository_url_is_ssh), + }) + } +} + +fn sorted(value: &Value) -> Value { + match value { + Value::Object(map) => { + let mut entries: Vec<_> = map.iter().map(|(k, v)| (k.clone(), sorted(v))).collect(); + entries.sort_by(|a, b| a.0.cmp(&b.0)); + Value::Object(entries.into_iter().collect()) + } + Value::Array(items) => Value::Array(items.iter().map(sorted).collect()), + other => other.clone(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn requests_without_a_store_path_share_a_key_regardless_of_attribute_order() { + let first = InputRequest::parse( + r#"{"builtin":"fetchTree","args":{"type":"git","url":"git+ssh://git@example.com/r","rev":"abc"}}"#, + ) + .unwrap(); + let second = InputRequest::parse( + r#"{"args":{"rev":"abc","url":"git+ssh://git@example.com/r","type":"git"},"builtin":"fetchTree"}"#, + ) + .unwrap(); + + assert_eq!(first.fetch_key, second.fetch_key); + assert!(first.needs_secrets); + assert_eq!(first.store_path, None); + } + + #[test] + fn a_known_store_path_is_the_key_and_https_needs_no_secrets() { + let request = InputRequest::parse( + r#"{"builtin":"fetchTree","args":{"type":"tarball","url":"https://example.com/x.tar.gz"},"storePath":"/nix/store/aaaa-source"}"#, + ) + .unwrap(); + + assert_eq!(request.fetch_key, "/nix/store/aaaa-source"); + assert!(!request.needs_secrets); + assert!(InputRequest::parse(r#"{"builtin":"fetchTree"}"#).is_err()); + } +} diff --git a/backend/gradient-pool/src/lib.rs b/backend/gradient-pool/src/lib.rs index dfb8fb6d8..0214de9ed 100644 --- a/backend/gradient-pool/src/lib.rs +++ b/backend/gradient-pool/src/lib.rs @@ -6,6 +6,7 @@ pub mod aggregate; pub mod estimate; +pub mod input_request; pub mod peer_auth; pub mod score; pub mod session_port; @@ -14,6 +15,7 @@ pub mod worker_pool; pub mod worker_state; pub use self::aggregate::{Aggregate, aggregate}; +pub use self::input_request::InputRequest; pub use self::peer_auth::PeerAuth; pub use self::worker_caps::{WorkerCaps, build_needs_network}; pub use self::worker_pool::{WorkerInfo, WorkerPool, WorkerSlot}; diff --git a/backend/gradient-pool/src/session_port.rs b/backend/gradient-pool/src/session_port.rs index e1f81eacd..7017a8064 100644 --- a/backend/gradient-pool/src/session_port.rs +++ b/backend/gradient-pool/src/session_port.rs @@ -4,7 +4,8 @@ * SPDX-License-Identifier: AGPL-3.0-only */ -use gradient_wire::types::{ClusterAddress, ClusterPeer, ImportOutcome}; +use gradient_types::ids::ProjectId; +use gradient_wire::types::{ClusterAddress, ClusterPeer, ImportOutcome, InputOutcome}; #[derive(Debug, Clone, PartialEq, Eq)] pub enum SessionSignal { @@ -40,6 +41,22 @@ pub enum SessionSignal { request_id: String, outcome: ImportOutcome, }, + InputResult { + job_id: String, + request_id: String, + outcome: InputOutcome, + }, + InputProgress { + job_id: String, + request_id: String, + bytes_done: u64, + bytes_total: Option, + }, + FetchInput { + fetch_id: String, + project: ProjectId, + request: String, + }, } pub trait SessionPort: Send + Sync + 'static { diff --git a/backend/gradient-pool/src/worker_caps.rs b/backend/gradient-pool/src/worker_caps.rs index ba24b3dd3..551ac6e73 100644 --- a/backend/gradient-pool/src/worker_caps.rs +++ b/backend/gradient-pool/src/worker_caps.rs @@ -61,6 +61,10 @@ impl WorkerCaps { pub fn can_fetch(&self, source: &FlakeSource) -> bool { let needs_secrets = matches!(source, FlakeSource::Repository { url, .. } if check_repository_url_is_ssh(url)); + self.can_fetch_input(needs_secrets) + } + + pub fn can_fetch_input(&self, needs_secrets: bool) -> bool { self.capabilities.network && (!needs_secrets || self.capabilities.secrets) } } diff --git a/backend/gradient-proto/src/handler/inbound.rs b/backend/gradient-proto/src/handler/inbound.rs index d7b052ed0..f59957a53 100644 --- a/backend/gradient-proto/src/handler/inbound.rs +++ b/backend/gradient-proto/src/handler/inbound.rs @@ -48,9 +48,24 @@ pub(crate) struct ActiveJob { } #[derive(Clone, Default)] -pub(crate) struct ActiveJobs(Arc>>); +pub(crate) struct ActiveJobs( + Arc>>, + Arc>>, +); impl ActiveJobs { + pub(crate) fn insert_fetch(&self, fetch_id: String, project: ProjectId) { + self.1.lock().insert(fetch_id, project); + } + + pub(crate) fn fetch_project(&self, fetch_id: &str) -> Option { + self.1.lock().get(fetch_id).copied() + } + + pub(crate) fn remove_fetch(&self, fetch_id: &str) -> Option { + self.1.lock().remove(fetch_id) + } + pub(crate) fn insert(&self, job_id: String, job: ActiveJob) { self.0.lock().insert(job_id, job); } @@ -117,7 +132,10 @@ impl ActiveJobs { impl From> for ActiveJobs { fn from(jobs: HashMap) -> Self { - Self(Arc::new(gradient_util::sync::Mutex::new(jobs))) + Self( + Arc::new(gradient_util::sync::Mutex::new(jobs)), + Arc::default(), + ) } } @@ -368,10 +386,53 @@ impl<'a> InboundContext<'a> { } true } - msg @ (ClientMessage::InputRequest { .. } - | ClientMessage::InputFetched { .. } - | ClientMessage::InputFetchProgress { .. }) => { - warn!(peer_id = %self.peer_id, variant = msg.variant_name(), "input fetches are not served yet"); + ClientMessage::InputRequest { + job_id, + request_id, + request, + } => { + if self.owns_handed_dispatch(&job_id).await && self.active.eval(&job_id).is_some() { + self.job_events + .push(JobEvent::InputRequest { + job_id, + request_id, + request, + }) + .await; + } else { + warn!(peer_id = %self.peer_id, %job_id, "input request for an evaluation this session does not run"); + } + true + } + ClientMessage::InputFetched { fetch_id, outcome } => { + if self.active.remove_fetch(&fetch_id).is_some() { + self.job_events + .push(JobEvent::InputFetched { fetch_id, outcome }) + .await; + } else { + warn!(peer_id = %self.peer_id, %fetch_id, "answer for an input fetch this session does not run"); + } + true + } + ClientMessage::InputFetchProgress { + fetch_id, + bytes_done, + bytes_total, + } => { + if self.active.fetch_project(&fetch_id).is_some() { + let scheduler = Arc::clone(self.scheduler); + let peer_id = self.peer_id.to_owned(); + self.handles.db.shutdown.spawn(async move { + scheduler + .handle_input_fetch_progress( + &peer_id, + &fetch_id, + bytes_done, + bytes_total, + ) + .await; + }); + } true } ClientMessage::EvalMessage { diff --git a/backend/gradient-proto/src/handler/job_events.rs b/backend/gradient-proto/src/handler/job_events.rs index 981c8bd04..47725f2be 100644 --- a/backend/gradient-proto/src/handler/job_events.rs +++ b/backend/gradient-proto/src/handler/job_events.rs @@ -15,7 +15,7 @@ use gradient_graph::writer::RECORD_ROW_BUDGET; use gradient_scheduler::{ReportedTimeline, Scheduler}; use gradient_types::ids::{DispatchedJobId, EvaluationId}; use gradient_util::shutdown::Shutdown; -use gradient_wire::types::{BuildFailureKind, BuildMetrics}; +use gradient_wire::types::{BuildFailureKind, BuildMetrics, InputOutcome}; use tokio::sync::{mpsc, oneshot}; use tokio::task::JoinHandle; use tracing::{Instrument as _, debug, debug_span, error, info, warn}; @@ -52,6 +52,15 @@ pub(super) enum JobEvent { request_id: String, drv_paths: Vec, }, + InputRequest { + job_id: String, + request_id: String, + request: String, + }, + InputFetched { + fetch_id: String, + outcome: InputOutcome, + }, } impl JobEvent { @@ -60,7 +69,11 @@ impl JobEvent { JobEvent::Update { job_id, .. } | JobEvent::Completed { job_id, .. } | JobEvent::Failed { job_id, .. } - | JobEvent::ImportRequest { job_id, .. } => job_id, + | JobEvent::ImportRequest { job_id, .. } + | JobEvent::InputRequest { job_id, .. } + | JobEvent::InputFetched { + fetch_id: job_id, .. + } => job_id, } } @@ -74,6 +87,8 @@ impl JobEvent { JobEvent::Completed { .. } => "completed", JobEvent::Failed { .. } => "failed", JobEvent::ImportRequest { .. } => "import_request", + JobEvent::InputRequest { .. } => "input_request", + JobEvent::InputFetched { .. } => "input_fetched", } } @@ -243,6 +258,25 @@ impl ApplyJobEvent for SchedulerJobEvents { self.import_request(job_id, evaluation, request_id, drv_paths) .await } + JobEvent::InputRequest { + job_id, + request_id, + request, + } => { + let peer_id = self.peer_id.as_str(); + if let Err(e) = self + .scheduler + .handle_input_request(peer_id, &job_id, request_id, request) + .await + { + error!(%peer_id, %job_id, error = %e, "handle_input_request failed"); + } + } + JobEvent::InputFetched { fetch_id, outcome } => { + self.scheduler + .handle_input_fetched(&self.peer_id, &fetch_id, outcome) + .await; + } } } } diff --git a/backend/gradient-proto/src/handler/session_actor.rs b/backend/gradient-proto/src/handler/session_actor.rs index 89e09c333..77f37f262 100644 --- a/backend/gradient-proto/src/handler/session_actor.rs +++ b/backend/gradient-proto/src/handler/session_actor.rs @@ -30,7 +30,10 @@ use super::inbound::{ActiveJobs, InboundContext, RpcContext}; use super::job_events::{JobEvents, SchedulerJobEvents}; use super::log_lane::LogLane; use super::session::on_reauth_notify; -use super::socket::{ProtoSocket, ProtoWriter, job_offer_chunks, recv_client_msg, send_server_msg}; +use super::socket::{ + InputFetch, ProtoSocket, ProtoWriter, job_offer_chunks, recv_client_msg, send_input_fetch, + send_server_msg, +}; use super::upload::{UploadSession, UploadTable, abandon_transfer}; use gradient_wire::Rejection; use gradient_wire::messages::{ClientMessage, GradientCapabilities, ServerMessage}; @@ -283,6 +286,69 @@ impl Actor for SessionActor { myself.stop(Some("write failed".into())); } } + SessionMsg::Signal(SessionSignal::InputResult { + job_id, + request_id, + outcome, + }) => { + if !st.active.contains(&job_id) { + debug!(peer_id = %st.peer_id, %job_id, "input result for a job this session no longer runs"); + return Ok(()); + } + + let msg = ServerMessage::InputResult { + job_id, + request_id, + outcome, + }; + if send_server_msg(&st.writer, &msg).await.is_err() { + myself.stop(Some("write failed".into())); + } + } + SessionMsg::Signal(SessionSignal::InputProgress { + job_id, + request_id, + bytes_done, + bytes_total, + }) => { + if !st.active.contains(&job_id) { + return Ok(()); + } + + let msg = ServerMessage::InputProgress { + job_id, + request_id, + bytes_done, + bytes_total, + }; + if send_server_msg(&st.writer, &msg).await.is_err() { + myself.stop(Some("write failed".into())); + } + } + SessionMsg::Signal(SessionSignal::FetchInput { + fetch_id, + project, + request, + }) => { + st.active.insert_fetch(fetch_id.clone(), project); + let fetch = InputFetch { + fetch_id, + project, + request, + }; + if send_input_fetch( + &st.writer, + &st.handles.db, + &st.scheduler, + &st.peer_id, + fetch, + ) + .await + .is_err() + { + myself.stop(Some("write failed".into())); + } + } SessionMsg::Signal(SessionSignal::StartCluster { attempt, roster }) => { let msg = ServerMessage::StartCluster { attempt, roster }; if send_server_msg(&st.writer, &msg).await.is_err() { diff --git a/backend/gradient-proto/src/handler/socket.rs b/backend/gradient-proto/src/handler/socket.rs index 529ea0ed1..adfab8862 100644 --- a/backend/gradient-proto/src/handler/socket.rs +++ b/backend/gradient-proto/src/handler/socket.rs @@ -64,34 +64,82 @@ fn sends_ssh_key( && matches!(job, Job::Flake(flake_job) if flake_job.steps.contains(&FlakeStep::FetchFlake)) } +pub(super) async fn send_input_fetch( + writer: &ProtoWriter, + db: &DbContext, + scheduler: &gradient_scheduler::Scheduler, + worker_id: &str, + fetch: InputFetch, +) -> Result<(), gradient_wire::session::frame::SendError> { + let caps = scheduler.worker_gradient_caps(worker_id).await; + let ssh_key = if sends_input_ssh_key(caps.as_ref(), &fetch.request) { + project_ssh_key(db, fetch.project).await + } else { + None + }; + send_server_msg( + writer, + &ServerMessage::FetchInput { + fetch_id: fetch.fetch_id, + project: fetch.project.to_string(), + request: fetch.request, + ssh_key, + }, + ) + .await +} + +pub(super) struct InputFetch { + pub fetch_id: String, + pub project: ProjectId, + pub request: String, +} + +fn sends_input_ssh_key( + caps: Option<&gradient_wire::types::GradientCapabilities>, + request: &str, +) -> bool { + caps.is_some_and(|c| c.secrets) + && gradient_pool::InputRequest::parse(request).is_ok_and(|r| r.needs_secrets) +} + async fn send_ssh_key_credential(writer: &ProtoWriter, db: &DbContext, project_id: ProjectId) { use gradient_wire::types::CredentialKind; + if let Some(private_key) = project_ssh_key(db, project_id).await { + let _ = send_server_msg( + writer, + &ServerMessage::Credential { + kind: CredentialKind::SshKey, + data: private_key.into_bytes().into(), + }, + ) + .await; + debug!(%project_id, "SSH key credential sent"); + } +} + +async fn project_ssh_key(db: &DbContext, project_id: ProjectId) -> Option { match by_id::(&db.worker_db, project_id).await { - Ok(Some(project)) => { - match gradient_sources::ssh_key::decrypt_ssh_private_key( - &db.config.secrets.crypt_file, - project, - &db.config.server.serve_url, - ) { - Ok((private_key, _public_key)) => { - let _ = send_server_msg( - writer, - &ServerMessage::Credential { - kind: CredentialKind::SshKey, - data: private_key.into_bytes().into(), - }, - ) - .await; - debug!(%project_id, "SSH key credential sent"); - } - Err(e) => { - debug!(%project_id, error = %e, "no SSH key for project (may be HTTPS repo)"); - } + Ok(Some(project)) => match gradient_sources::ssh_key::decrypt_ssh_private_key( + &db.config.secrets.crypt_file, + project, + &db.config.server.serve_url, + ) { + Ok((private_key, _public_key)) => Some(private_key), + Err(e) => { + debug!(%project_id, error = %e, "no SSH key for project (may be HTTPS repo)"); + None } + }, + Ok(None) => { + warn!(%project_id, "project not found for SSH key lookup"); + None + } + Err(e) => { + warn!(%project_id, error = %e, "failed to fetch project for SSH key"); + None } - Ok(None) => warn!(%project_id, "project not found for SSH key lookup"), - Err(e) => warn!(%project_id, error = %e, "failed to fetch project for SSH key"), } } @@ -130,4 +178,25 @@ mod tests { }; assert!(sends_ssh_key(Some(&with_secrets), &fetch_job())); } + + #[test] + fn no_input_ssh_key_without_negotiated_secrets() { + let ssh = + r#"{"builtin":"fetchTree","args":{"type":"git","url":"git+ssh://git@example.com/r"}}"#; + let https = + r#"{"builtin":"fetchTree","args":{"type":"git","url":"https://example.com/r"}}"#; + let network_only = GradientCapabilities { + network: true, + ..Default::default() + }; + let with_secrets = GradientCapabilities { + secrets: true, + ..network_only + }; + + assert!(!sends_input_ssh_key(Some(&network_only), ssh)); + assert!(!sends_input_ssh_key(None, ssh)); + assert!(!sends_input_ssh_key(Some(&with_secrets), https)); + assert!(sends_input_ssh_key(Some(&with_secrets), ssh)); + } } diff --git a/backend/gradient-proto/src/handler/upload/mod.rs b/backend/gradient-proto/src/handler/upload/mod.rs index 7a6b1b79d..3fec4019b 100644 --- a/backend/gradient-proto/src/handler/upload/mod.rs +++ b/backend/gradient-proto/src/handler/upload/mod.rs @@ -65,7 +65,7 @@ impl InboundContext<'_> { ) .await; }; - if !self.active.contains(&job_id) { + if !self.active.contains(&job_id) && self.active.fetch_project(&job_id).is_none() { return self .settle( request_id, @@ -315,7 +315,11 @@ impl InboundContext<'_> { let db = self.handles.db.worker_db.clone(); let peer_id = self.peer_id.to_owned(); let job_id = granted.job_id.clone(); + let fetch_project = self.active.fetch_project(&job_id); Box::pin(async move { + if fetch_project.is_some() { + return fetch_project; + } match scheduler.project_for_job(&job_id).await { Some(id) => Some(id), None => super::nar::project_for_dispatched_job(&db, &peer_id, &job_id).await, diff --git a/backend/gradient-scheduler/src/cases/inputs.rs b/backend/gradient-scheduler/src/cases/inputs.rs new file mode 100644 index 000000000..407411631 --- /dev/null +++ b/backend/gradient-scheduler/src/cases/inputs.rs @@ -0,0 +1,390 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use gradient_entity::evaluation::EvaluationStatus; +use gradient_entity::ids::{EvaluationId, ProjectId}; +use gradient_fixtures::shape::shape; +use gradient_pool::session_port::SessionSignal; +use gradient_types::waiting_reason::{EvalCapability, WaitingReason}; +use gradient_wire::types::{InputOutcome, JobKind}; + +use super::{ + A, A_OUT, CaseShapes, SchedulerCase, eval_key, w_all, w_network_eval, w_network_secrets, + without_offers, +}; + +const W1: &str = "w1"; +const W2: &str = "w2"; +const W3: &str = "w3"; +const W4: &str = "w4"; + +const SSH_REQUEST: &str = r#"{"builtin":"fetchTree","args":{"type":"git","url":"git+ssh://git@example.invalid/input.git","rev":"abc"},"storePath":"/nix/store/aaaa-input"}"#; +const FETCHED: &str = r#"{"storePath":"/nix/store/aaaa-input"}"#; + +#[derive(Debug, Clone, PartialEq, Eq)] +struct Fetch { + worker: String, + fetch_id: String, + project: ProjectId, + request: String, +} + +#[tokio::test] +async fn i01_an_input_request_goes_to_a_network_and_secrets_worker() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start(shape().evalq(e1)).await; + evaluating(&case, W1, e1).await; + case.connect(W2, super::w_network(), &[]).await; + case.connect(W3, w_network_secrets(), &[]).await; + + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + let fetches = fetches_of(&case, &[W1, W2, W3]); + assert_eq!( + fetches, + vec![Fetch { + worker: W3.into(), + fetch_id: fetches[0].fetch_id.clone(), + project: case.db.project_of(e1).await, + request: SSH_REQUEST.into(), + }], + "after 1: only W3 gets FetchInput" + ); + assert_eq!(case.input_waits(), 1, "after 1: one input wait"); + case.check("step 1").await; + + case.finish().await; +} + +#[tokio::test] +async fn i02_two_evaluations_share_one_fetch() { + let (e1, e2) = (EvaluationId::now_v7(), EvaluationId::now_v7()); + let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; + evaluating(&case, W1, e1).await; + evaluating(&case, W2, e2).await; + case.connect(W3, w_network_secrets(), &[]).await; + + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + case.input_request(W2, &eval_key(e2), "r2", SSH_REQUEST) + .await; + let fetches = fetches_of(&case, &[W1, W2, W3]); + assert_eq!(fetches.len(), 1, "after 1: one FetchInput, got {fetches:?}"); + assert_eq!(case.input_waits(), 2, "after 1: two waiting requests"); + case.check("step 1").await; + + case.input_fetched(W3, &fetches[0].fetch_id, fetched()) + .await; + assert_eq!( + without_offers(case.signals(W1)), + vec![result(e1, "r1", fetched())], + "after 2: signals(W1)" + ); + assert_eq!( + without_offers(case.signals(W2)), + vec![result(e2, "r2", fetched())], + "after 2: signals(W2)" + ); + assert_eq!(case.input_waits(), 0, "after 2: no input wait left"); + case.check("step 2").await; + + case.finish().await; +} + +#[tokio::test] +async fn i03_a_failed_fetch_answers_every_waiter() { + let (e1, e2) = (EvaluationId::now_v7(), EvaluationId::now_v7()); + let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; + evaluating(&case, W1, e1).await; + evaluating(&case, W2, e2).await; + case.connect(W3, w_network_secrets(), &[]).await; + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + case.input_request(W2, &eval_key(e2), "r2", SSH_REQUEST) + .await; + let fetch = single_fetch(&case, &[W1, W2, W3]); + case.check("step 1").await; + + let failed = InputOutcome::Failed { + message: "HTTP error 404".into(), + }; + case.input_fetched(W3, &fetch.fetch_id, failed.clone()) + .await; + assert_eq!( + without_offers(case.signals(W1)), + vec![result(e1, "r1", failed.clone())], + "after 2: signals(W1)" + ); + assert_eq!( + without_offers(case.signals(W2)), + vec![result(e2, "r2", failed)], + "after 2: signals(W2)" + ); + case.check("step 2").await; + + case.finish().await; +} + +#[tokio::test] +async fn i04_a_lost_fetch_worker_hands_the_fetch_to_another() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start(shape().evalq(e1)).await; + evaluating(&case, W1, e1).await; + case.connect(W2, w_network_secrets(), &[]).await; + case.connect(W3, w_network_secrets(), &[]).await; + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + let first = single_fetch(&case, &[W1, W2, W3]); + assert_eq!( + first.worker, W2, + "after 1: the lower worker id wins the tie" + ); + case.check("step 1").await; + + case.disconnect(W2).await; + let second = single_fetch(&case, &[W1, W2, W3]); + assert_eq!(second.worker, W3, "after 2: W3 takes over"); + assert_ne!(second.fetch_id, first.fetch_id, "after 2: a new fetch id"); + case.check("step 2").await; + + case.input_fetched(W2, &first.fetch_id, fetched()).await; + assert_eq!( + without_offers(case.signals(W1)), + vec![], + "after 3: the late answer of W2 is ignored" + ); + case.check("step 3").await; + + case.input_fetched(W3, &second.fetch_id, fetched()).await; + assert_eq!( + without_offers(case.signals(W1)), + vec![result(e1, "r1", fetched())], + "after 4: signals(W1)" + ); + case.check("step 4").await; + + case.finish().await; +} + +#[tokio::test] +async fn i05_an_abort_drops_the_input_wait() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start(shape().evalq(e1)).await; + evaluating(&case, W1, e1).await; + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + assert_eq!( + fetches_of(&case, &[W1]), + vec![], + "after 1: W1 lacks secrets" + ); + case.check("step 1").await; + + case.abort_evaluation(e1).await; + case.signals(W1); + assert_eq!(case.input_waits(), 0, "after 2: the wait is gone"); + case.check("step 2").await; + + case.connect(W2, w_network_secrets(), &[]).await; + case.build_assign_pass(false, false).await; + assert_eq!( + fetches_of(&case, &[W1, W2]), + vec![], + "after 3: no FetchInput for an aborted evaluation" + ); + assert_eq!( + without_offers(case.signals(W1)), + vec![], + "after 3: no InputResult" + ); + case.check("step 3").await; + + case.finish().await; +} + +#[tokio::test] +async fn i06_without_a_fetch_worker_the_evaluation_waits_for_network() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start(shape().evalq(e1)).await; + evaluating(&case, W1, e1).await; + + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + case.refresh_waiting_state().await; + let row = case.evaluation(e1).await; + assert_eq!( + ( + row.status, + row.waiting_reason + .as_ref() + .and_then(WaitingReason::from_json) + ), + ( + EvaluationStatus::EvaluatingDerivation, + Some(WaitingReason::eval_workers(EvalCapability::Network, 1)) + ), + "after 1: evaluation(E1) keeps evaluating and names the gap" + ); + case.check("step 1").await; + + case.connect(W2, w_network_secrets(), &[]).await; + case.build_assign_pass(false, false).await; + let fetch = single_fetch(&case, &[W1, W2]); + assert_eq!(fetch.worker, W2, "after 2: the new worker gets the fetch"); + let row = case.evaluation(e1).await; + assert_eq!( + (row.status, row.waiting_reason), + (EvaluationStatus::EvaluatingDerivation, None), + "after 2: evaluation(E1) has no waiting reason" + ); + case.check("step 2").await; + + case.finish().await; +} + +#[tokio::test] +async fn i07_an_input_request_from_a_build_job_is_dropped() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start(shape().one(e1, A, A_OUT)).await; + case.connect(W1, w_all(), &[]).await; + case.connect(W2, w_network_secrets(), &[]).await; + case.admit_startable_moves().await; + case.offers_full(W1).await; + case.score_all(W1).await; + let build = case + .request_job(W1, JobKind::Build) + .await + .expect("after 1: W1 gets build:A"); + case.signals(W1); + case.check("step 1").await; + + case.input_request(W1, build.job_id(), "r1", SSH_REQUEST) + .await; + assert_eq!(fetches_of(&case, &[W1, W2]), vec![], "after 2: no fetch"); + assert_eq!(case.input_waits(), 0, "after 2: no wait opened"); + case.check("step 2").await; + + case.finish().await; +} + +#[tokio::test] +async fn i08_fetch_progress_reaches_every_waiting_evaluation() { + let (e1, e2) = (EvaluationId::now_v7(), EvaluationId::now_v7()); + let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; + evaluating(&case, W1, e1).await; + evaluating(&case, W2, e2).await; + case.connect(W3, w_network_secrets(), &[]).await; + case.connect(W4, w_network_secrets(), &[]).await; + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + case.input_request(W2, &eval_key(e2), "r2", SSH_REQUEST) + .await; + let first = single_fetch(&case, &[W1, W2, W3, W4]); + case.check("step 1").await; + + case.input_fetch_progress(&first.worker, &first.fetch_id, 512, Some(1024)) + .await; + assert_eq!( + without_offers(case.signals(W1)), + vec![progress(e1, "r1", 512, Some(1024))], + "after 2: signals(W1)" + ); + assert_eq!( + without_offers(case.signals(W2)), + vec![progress(e2, "r2", 512, Some(1024))], + "after 2: signals(W2)" + ); + case.check("step 2").await; + + case.disconnect(&first.worker).await; + let second = single_fetch(&case, &[W1, W2, W3, W4]); + case.input_fetch_progress(&first.worker, &first.fetch_id, 1024, Some(1024)) + .await; + for worker in [W1, W2] { + assert_eq!( + without_offers(case.signals(worker)), + vec![], + "after 3: progress of the replaced fetch reaches no {worker}" + ); + } + case.input_fetch_progress(&second.worker, &second.fetch_id, 64, None) + .await; + assert_eq!( + without_offers(case.signals(W1)), + vec![progress(e1, "r1", 64, None)], + "after 3: signals(W1)" + ); + case.signals(W2); + case.check("step 3").await; + + case.finish().await; +} + +async fn evaluating(case: &SchedulerCase, worker: &str, evaluation: EvaluationId) { + let key = eval_key(evaluation); + case.connect(worker, w_network_eval(), &[]).await; + case.evalrun(worker, evaluation).await; + case.advance_evaluation(worker, &key, EvaluationStatus::EvaluatingDerivation) + .await; + case.signals(worker); + case.check("step 0").await; +} + +fn fetches_of(case: &SchedulerCase, workers: &[&str]) -> Vec { + let mut fetches = Vec::new(); + for worker in workers { + for signal in without_offers(case.signals(worker)) { + match signal { + SessionSignal::FetchInput { + fetch_id, + project, + request, + } => fetches.push(Fetch { + worker: (*worker).to_owned(), + fetch_id, + project, + request, + }), + other => panic!("{worker} got {other:?} besides FetchInput"), + } + } + } + fetches +} + +fn single_fetch(case: &SchedulerCase, workers: &[&str]) -> Fetch { + let mut fetches = fetches_of(case, workers); + assert_eq!(fetches.len(), 1, "one FetchInput, got {fetches:?}"); + fetches.remove(0) +} + +fn fetched() -> InputOutcome { + InputOutcome::Fetched { + response: FETCHED.into(), + } +} + +fn result(evaluation: EvaluationId, request_id: &str, outcome: InputOutcome) -> SessionSignal { + SessionSignal::InputResult { + job_id: eval_key(evaluation), + request_id: request_id.into(), + outcome, + } +} + +fn progress( + evaluation: EvaluationId, + request_id: &str, + bytes_done: u64, + bytes_total: Option, +) -> SessionSignal { + SessionSignal::InputProgress { + job_id: eval_key(evaluation), + request_id: request_id.into(), + bytes_done, + bytes_total, + } +} diff --git a/backend/gradient-scheduler/src/cases/mod.rs b/backend/gradient-scheduler/src/cases/mod.rs index 90e4ade44..d90d3c16a 100644 --- a/backend/gradient-scheduler/src/cases/mod.rs +++ b/backend/gradient-scheduler/src/cases/mod.rs @@ -13,6 +13,7 @@ mod control; mod estimates; mod eval_queue; mod imports; +mod inputs; mod invariants; mod outside; mod probe; diff --git a/backend/gradient-scheduler/src/cases/readers.rs b/backend/gradient-scheduler/src/cases/readers.rs index 993435181..582cf2b62 100644 --- a/backend/gradient-scheduler/src/cases/readers.rs +++ b/backend/gradient-scheduler/src/cases/readers.rs @@ -349,6 +349,10 @@ impl SchedulerCase { } } + pub fn input_waits(&self) -> usize { + self.scheduler.input_waits.waiting_requests() + } + pub fn import_waits(&self) -> usize { self.scheduler.import_waits.waiting_requests() } diff --git a/backend/gradient-scheduler/src/cases/session.rs b/backend/gradient-scheduler/src/cases/session.rs index 469ad2bff..db4e3163f 100644 --- a/backend/gradient-scheduler/src/cases/session.rs +++ b/backend/gradient-scheduler/src/cases/session.rs @@ -12,7 +12,7 @@ use gradient_entity::evaluation::EvaluationStatus; use gradient_entity::ids::{DispatchedJobId, EvaluationId, ProjectId}; use gradient_wire::types::{ BuildFailureKind, BuildMetrics, BuildOutput, BumpedInputWire, CandidateScore, - DiscoveredDerivation, EvalMessageLevel, EvalStatsReport, JobKind, JobPhaseSpan, + DiscoveredDerivation, EvalMessageLevel, EvalStatsReport, InputOutcome, JobKind, JobPhaseSpan, }; use super::aliases; @@ -527,6 +527,34 @@ impl SchedulerCase { .await; } + pub async fn input_request(&self, worker: &str, key: &str, request_id: &str, request: &str) { + if !self.owns_remembered(worker, key).await { + return; + } + let _ = self + .scheduler + .handle_input_request(worker, key, request_id.to_owned(), request.to_owned()) + .await; + } + + pub async fn input_fetched(&self, worker: &str, fetch_id: &str, outcome: InputOutcome) { + self.scheduler + .handle_input_fetched(worker, fetch_id, outcome) + .await; + } + + pub async fn input_fetch_progress( + &self, + worker: &str, + fetch_id: &str, + bytes_done: u64, + bytes_total: Option, + ) { + self.scheduler + .handle_input_fetch_progress(worker, fetch_id, bytes_done, bytes_total) + .await; + } + pub async fn respawn_core(&self) { let old = lock(&self.core).clone(); old.stop_and_wait(None, None) diff --git a/backend/gradient-scheduler/src/input_waits.rs b/backend/gradient-scheduler/src/input_waits.rs new file mode 100644 index 000000000..4aa6af49c --- /dev/null +++ b/backend/gradient-scheduler/src/input_waits.rs @@ -0,0 +1,183 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; + +use gradient_types::ids::{EvaluationId, ProjectId}; +use gradient_util::sync::Mutex; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct InputWaiter { + pub worker: String, + pub job_id: String, + pub request_id: String, + pub evaluation: EvaluationId, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AssignedFetch { + pub fetch_id: String, + pub worker: String, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UnassignedFetch { + pub fetch_key: String, + pub project: ProjectId, + pub request: String, + pub needs_secrets: bool, +} + +#[derive(Debug)] +struct InputWait { + request: String, + project: ProjectId, + needs_secrets: bool, + waiters: Vec, + fetch: Option, +} + +#[derive(Debug, Clone)] +pub struct NewInputWait { + pub fetch_key: String, + pub request: String, + pub project: ProjectId, + pub needs_secrets: bool, + pub waiter: InputWaiter, +} + +#[derive(Clone, Debug, Default)] +pub struct InputWaits(Arc>>); + +impl InputWaits { + pub fn register(&self, new: NewInputWait) { + self.0 + .lock() + .entry(new.fetch_key) + .or_insert_with(|| InputWait { + request: new.request, + project: new.project, + needs_secrets: new.needs_secrets, + waiters: Vec::new(), + fetch: None, + }) + .waiters + .push(new.waiter); + } + + pub fn unassigned(&self) -> Vec { + let mut unassigned: Vec = self + .0 + .lock() + .iter() + .filter(|(_, wait)| wait.fetch.is_none()) + .map(|(key, wait)| UnassignedFetch { + fetch_key: key.clone(), + project: wait.project, + request: wait.request.clone(), + needs_secrets: wait.needs_secrets, + }) + .collect(); + unassigned.sort_by(|a, b| a.fetch_key.cmp(&b.fetch_key)); + unassigned + } + + pub fn assigned_per_worker(&self) -> HashMap { + let mut counts = HashMap::new(); + for fetch in self.0.lock().values().filter_map(|w| w.fetch.as_ref()) { + *counts.entry(fetch.worker.clone()).or_default() += 1; + } + counts + } + + pub fn assign(&self, fetch_key: &str, fetch: AssignedFetch) -> bool { + match self.0.lock().get_mut(fetch_key) { + Some(wait) if wait.fetch.is_none() => { + wait.fetch = Some(fetch); + true + } + _ => false, + } + } + + pub fn take_fetched(&self, worker: &str, fetch_id: &str) -> Vec { + let mut waits = self.0.lock(); + let Some(key) = key_of_fetch(&waits, worker, fetch_id) else { + return Vec::new(); + }; + waits.remove(&key).map(|w| w.waiters).unwrap_or_default() + } + + pub fn waiters_of_fetch(&self, worker: &str, fetch_id: &str) -> Vec { + let waits = self.0.lock(); + key_of_fetch(&waits, worker, fetch_id) + .and_then(|key| waits.get(&key)) + .map(|w| w.waiters.clone()) + .unwrap_or_default() + } + + pub fn release_worker(&self, worker: &str) { + let mut waits = self.0.lock(); + for wait in waits.values_mut() { + wait.waiters.retain(|w| w.worker != worker); + if wait.fetch.as_ref().is_some_and(|f| f.worker == worker) { + wait.fetch = None; + } + } + waits.retain(|_, wait| !wait.waiters.is_empty()); + } + + pub fn drop_evaluations(&self, evaluations: &HashSet) -> usize { + let mut waits = self.0.lock(); + let mut dropped = 0; + for wait in waits.values_mut() { + let before = wait.waiters.len(); + wait.waiters + .retain(|w| !evaluations.contains(&w.evaluation)); + dropped += before - wait.waiters.len(); + } + waits.retain(|_, wait| !wait.waiters.is_empty()); + dropped + } + + pub fn waiting_evaluations(&self) -> HashSet { + self.0 + .lock() + .values() + .flat_map(|w| w.waiters.iter().map(|waiter| waiter.evaluation)) + .collect() + } + + pub fn unassigned_evaluations(&self) -> HashSet { + self.0 + .lock() + .values() + .filter(|w| w.fetch.is_none()) + .flat_map(|w| w.waiters.iter().map(|waiter| waiter.evaluation)) + .collect() + } + + #[cfg(test)] + pub fn waiting_requests(&self) -> usize { + self.0.lock().values().map(|w| w.waiters.len()).sum() + } +} + +fn key_of_fetch( + waits: &HashMap, + worker: &str, + fetch_id: &str, +) -> Option { + waits + .iter() + .find(|(_, wait)| { + wait.fetch + .as_ref() + .is_some_and(|f| f.fetch_id == fetch_id && f.worker == worker) + }) + .map(|(key, _)| key.clone()) +} diff --git a/backend/gradient-scheduler/src/job_handlers/abort.rs b/backend/gradient-scheduler/src/job_handlers/abort.rs index 3cee4e0bb..cbfbf1446 100644 --- a/backend/gradient-scheduler/src/job_handlers/abort.rs +++ b/backend/gradient-scheduler/src/job_handlers/abort.rs @@ -58,6 +58,7 @@ impl Scheduler { pub async fn stop_aborted_evaluation(&self, evaluation: EvaluationId, hard_abort: bool) { self.drop_import_waits(evaluation).await; + self.drop_input_waits(evaluation); let shared_builds = self .abort_evaluation_shared_builds(evaluation, hard_abort) .await; @@ -208,6 +209,7 @@ impl Scheduler { aborted_shared_builds: Vec, ) -> Vec<(String, String)> { self.drop_import_waits(evaluation_id).await; + self.drop_input_waits(evaluation_id); self.call(|reply| SchedulerMsg::AbortEvaluation { evaluation_id, aborted_shared_builds, diff --git a/backend/gradient-scheduler/src/job_handlers/input.rs b/backend/gradient-scheduler/src/job_handlers/input.rs new file mode 100644 index 000000000..fad45b928 --- /dev/null +++ b/backend/gradient-scheduler/src/job_handlers/input.rs @@ -0,0 +1,223 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use std::collections::{HashMap, HashSet}; + +use anyhow::{Context, Result}; +use gradient_entity::evaluation::EvaluationStatus; +use gradient_pool::session_port::SessionSignal; +use gradient_pool::{InputRequest, WorkerCaps, WorkerInfo}; +use gradient_types::*; +use gradient_wire::types::InputOutcome; +use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; +use tracing::debug; + +use crate::Scheduler; +use crate::input_waits::{AssignedFetch, InputWaiter, NewInputWait, UnassignedFetch}; +use crate::jobs::PendingJob; + +impl Scheduler { + pub async fn handle_input_request( + &self, + worker_id: &str, + job_id: &str, + request_id: String, + request: String, + ) -> Result<()> { + let Some(PendingJob::Eval(job)) = self.active_job(job_id).await else { + debug!(%worker_id, %job_id, "input request from a job that is no running evaluation; dropped"); + return Ok(()); + }; + let waiter = InputWaiter { + worker: worker_id.to_owned(), + job_id: job_id.to_owned(), + request_id, + evaluation: job.evaluation_id, + }; + let parsed = match InputRequest::parse(&request) { + Ok(parsed) => parsed, + Err(message) => { + self.answer_inputs(vec![waiter], InputOutcome::Failed { message }) + .await; + return Ok(()); + } + }; + + self.input_waits.register(NewInputWait { + fetch_key: parsed.fetch_key, + request, + project: job.project_id, + needs_secrets: parsed.needs_secrets, + waiter, + }); + self.dispatch_input_fetches().await; + Ok(()) + } + + pub async fn handle_input_fetched( + &self, + worker_id: &str, + fetch_id: &str, + outcome: InputOutcome, + ) { + let waiters = self.input_waits.take_fetched(worker_id, fetch_id); + if waiters.is_empty() { + debug!(%worker_id, %fetch_id, "input fetch answered without waiters; dropped"); + return; + } + + self.answer_inputs(waiters, outcome).await; + } + + pub async fn handle_input_fetch_progress( + &self, + worker_id: &str, + fetch_id: &str, + bytes_done: u64, + bytes_total: Option, + ) { + let signals = self + .input_waits + .waiters_of_fetch(worker_id, fetch_id) + .into_iter() + .map(|w| { + ( + w.worker, + SessionSignal::InputProgress { + job_id: w.job_id, + request_id: w.request_id, + bytes_done, + bytes_total, + }, + ) + }) + .collect::>(); + if !signals.is_empty() { + self.signal_workers(signals).await; + } + } + + pub(crate) async fn reassign_input_fetches_of(&self, worker_id: &str) { + self.input_waits.release_worker(worker_id); + self.dispatch_input_fetches().await; + } + + pub(crate) fn drop_input_waits(&self, evaluation: EvaluationId) { + let dropped = self + .input_waits + .drop_evaluations(&HashSet::from([evaluation])); + if dropped > 0 { + debug!(%evaluation, dropped, "dropped the input waits of an aborted evaluation"); + } + } + + pub(crate) async fn settle_input_waits(&self) -> Result<()> { + let waiting: Vec = + self.input_waits.waiting_evaluations().into_iter().collect(); + if waiting.is_empty() { + return Ok(()); + } + + let db = &self.db.worker_db; + let finished: HashSet = + gradient_db::fetch_in_chunks(&waiting, |chunk| async move { + EEvaluation::find() + .filter(CEvaluation::Id.is_in(chunk)) + .filter(CEvaluation::Status.is_in(EvaluationStatus::TERMINAL.to_vec())) + .all(db) + .await + }) + .await + .context("read the evaluations waiting on inputs")? + .into_iter() + .map(|e| e.id) + .collect(); + self.input_waits.drop_evaluations(&finished); + self.dispatch_input_fetches().await; + Ok(()) + } + + pub(crate) fn input_waiting_evaluations(&self) -> HashSet { + self.input_waits.unassigned_evaluations() + } + + async fn dispatch_input_fetches(&self) { + let unassigned = self.input_waits.unassigned(); + if unassigned.is_empty() { + return; + } + + let workers = self.board_workers().await; + let mut load = self.input_waits.assigned_per_worker(); + let mut signals = Vec::new(); + for fetch in unassigned { + let Some(worker) = fetch_worker(&workers, &load, &fetch) else { + debug!(fetch_key = %fetch.fetch_key, "no connected worker can fetch the input yet"); + continue; + }; + let fetch_id = uuid::Uuid::now_v7().to_string(); + let assigned = AssignedFetch { + fetch_id: fetch_id.clone(), + worker: worker.clone(), + }; + if self.input_waits.assign(&fetch.fetch_key, assigned) { + *load.entry(worker.clone()).or_default() += 1; + signals.push(( + worker, + SessionSignal::FetchInput { + fetch_id, + project: fetch.project, + request: fetch.request, + }, + )); + } + } + if !signals.is_empty() { + self.signal_workers(signals).await; + } + } + + async fn answer_inputs(&self, waiters: Vec, outcome: InputOutcome) { + let signals = waiters + .into_iter() + .map(|w| { + ( + w.worker, + SessionSignal::InputResult { + job_id: w.job_id, + request_id: w.request_id, + outcome: outcome.clone(), + }, + ) + }) + .collect(); + self.signal_workers(signals).await; + } +} + +fn fetch_worker( + workers: &[WorkerInfo], + load: &HashMap, + fetch: &UnassignedFetch, +) -> Option { + workers + .iter() + .filter(|w| !w.draining) + .filter(|w| { + w.authorized_peers + .as_ref() + .is_none_or(|peers| peers.contains(&fetch.project)) + }) + .filter(|w| { + WorkerCaps { + capabilities: w.capabilities.clone(), + ..WorkerCaps::default() + } + .can_fetch_input(fetch.needs_secrets) + }) + .min_by_key(|w| (load.get(&w.id).copied().unwrap_or_default(), w.id.clone())) + .map(|w| w.id.clone()) +} diff --git a/backend/gradient-scheduler/src/job_handlers/mod.rs b/backend/gradient-scheduler/src/job_handlers/mod.rs index f4dc3e23d..358ebef2b 100644 --- a/backend/gradient-scheduler/src/job_handlers/mod.rs +++ b/backend/gradient-scheduler/src/job_handlers/mod.rs @@ -10,6 +10,7 @@ mod build_status; mod cluster; mod eval_status; pub(crate) mod import; +mod input; mod priority; mod queue; pub(crate) mod timeline; diff --git a/backend/gradient-scheduler/src/lib.rs b/backend/gradient-scheduler/src/lib.rs index 7680067d4..a02ada893 100644 --- a/backend/gradient-scheduler/src/lib.rs +++ b/backend/gradient-scheduler/src/lib.rs @@ -14,6 +14,7 @@ pub mod connection_failures; pub mod draining; pub mod eval; pub mod import_waits; +pub mod input_waits; pub mod instance; pub mod jobs; pub mod log_substitution; @@ -73,6 +74,7 @@ pub struct Scheduler { pub draining: draining::Draining, pub connection_failures: Arc, pub(crate) import_waits: import_waits::ImportWaits, + pub(crate) input_waits: input_waits::InputWaits, pub(crate) workers: worker_memory::WorkerMemory, } @@ -97,6 +99,7 @@ impl Scheduler { draining: draining::Draining::default(), connection_failures: Arc::default(), import_waits: import_waits::ImportWaits::default(), + input_waits: input_waits::InputWaits::default(), workers: worker_memory::WorkerMemory::default(), } } diff --git a/backend/gradient-scheduler/src/loops/build.rs b/backend/gradient-scheduler/src/loops/build.rs index 081052a9b..fceb488b2 100644 --- a/backend/gradient-scheduler/src/loops/build.rs +++ b/backend/gradient-scheduler/src/loops/build.rs @@ -51,6 +51,9 @@ pub(crate) async fn build_assign_pass(scheduler: &Scheduler, timer_tick: bool, r if scheduler.scoring.policy.uses_history() { reestimate_pending(scheduler).await; } + if let Err(e) = scheduler.settle_input_waits().await { + error!(error = %e, "settling input waits in dispatch loop failed"); + } if let Err(e) = scheduler.refresh_waiting_state().await { error!(error = %e, "refresh_waiting_state in dispatch loop failed"); } diff --git a/backend/gradient-scheduler/src/waiting_state.rs b/backend/gradient-scheduler/src/waiting_state.rs index 109bedbee..f08075060 100644 --- a/backend/gradient-scheduler/src/waiting_state.rs +++ b/backend/gradient-scheduler/src/waiting_state.rs @@ -40,6 +40,7 @@ impl WaitingRefresh<'_> { pub(crate) async fn run( &self, workers: &[WorkerInfo], + input_waiting: &HashSet, draining: bool, ) -> Result> { gradient_db::evaluations::counters::fold_shared_build_deltas(&self.db.worker_db) @@ -154,7 +155,15 @@ impl WaitingRefresh<'_> { eval.status, check_repository_url_is_ssh(&eval.repository), project_workers, - ), + ) + .or_else(|| { + decide_input_wait( + eval.status, + reason.as_ref(), + input_waiting.contains(&eval.id), + project_workers.connected, + ) + }), EvaluationStatus::Building => { match self .build_phase_decision(eval.id, eval_counters, worker_caps, reason.as_ref()) @@ -775,6 +784,32 @@ fn decide_eval_recovery( } } +fn decide_input_wait( + status: EvaluationStatus, + current: Option<&WaitingReason>, + waits_for_input: bool, + connected: u32, +) -> Option<(EvaluationStatus, Option)> { + let waits_for_network = matches!( + current, + Some(WaitingReason::EvalWorkers { + capability: EvalCapability::Network, + .. + }) + ); + match (waits_for_input, waits_for_network) { + (true, _) => Some(( + status, + Some(WaitingReason::eval_workers( + EvalCapability::Network, + connected, + )), + )), + (false, true) => Some((status, None)), + (false, false) => None, + } +} + pub(crate) async fn persist_waiting_reason( db: &WorkerDb, evaluation_id: EvaluationId, diff --git a/backend/gradient-scheduler/src/worker_lifecycle.rs b/backend/gradient-scheduler/src/worker_lifecycle.rs index dec3b2f0e..ef1bd245b 100644 --- a/backend/gradient-scheduler/src/worker_lifecycle.rs +++ b/backend/gradient-scheduler/src/worker_lifecycle.rs @@ -314,6 +314,7 @@ impl Scheduler { warn!(error = %e, %worker_id, key = %lost.key, "settling a lost cluster member failed"); } } + self.reassign_input_fetches_of(worker_id).await; self.db.events.publish(worker::Disconnected { worker_id: worker_id.to_owned(), }); @@ -328,7 +329,7 @@ impl Scheduler { graph: &self.graph, memo: &self.scoring.assessments, } - .run(&workers, draining) + .run(&workers, &self.input_waiting_evaluations(), draining) .await?; for unbuildable in unbuildables { self.abort_unbuildable_evaluation(unbuildable).await; From 94e407ffd4ef2b0e38a42086a1f6b8f2bc47e6de Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 15:36:55 +0200 Subject: [PATCH 10/19] nix: fix: removal notes for fetch options --- nix/modules/gradient.nix | 1 + nix/modules/renamed-options.nix | 31 +++++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+) create mode 100644 nix/modules/renamed-options.nix diff --git a/nix/modules/gradient.nix b/nix/modules/gradient.nix index c3977b7f1..d06e71a28 100644 --- a/nix/modules/gradient.nix +++ b/nix/modules/gradient.nix @@ -126,6 +126,7 @@ in { imports = [ ./gradient-state.nix ./gradient-worker.nix + ./renamed-options.nix { options.services.gradient = serverOptions.options; } ]; diff --git a/nix/modules/renamed-options.nix b/nix/modules/renamed-options.nix new file mode 100644 index 000000000..c74120419 --- /dev/null +++ b/nix/modules/renamed-options.nix @@ -0,0 +1,31 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +{ lib, options, ... }: let + option = path: [ "services" "gradient" ] ++ path; + removed = path: lib.mkRemovedOptionModule (option path); + + removedStateOptions = [ + { group = "workers"; name = "enable_fetch"; reason = "replaced by enable_network and enable_secrets"; } + ]; + + stateEntries = group: value: + let entries = if lib.isAttrs value then value.${group} or { } else { }; in + if lib.isAttrs entries then entries else { }; +in { + imports = [ + (removed [ "worker" "capabilities" "fetch" ] "replaced by services.gradient.worker.capabilities.network and services.gradient.worker.capabilities.secrets") + ]; + + config.assertions = lib.concatMap (def: + lib.concatMap (r: + lib.mapAttrsToList (entry: _: { + assertion = false; + message = "The option `services.gradient.state.${r.group}.${entry}.${r.name}' defined in ${def.file} has been removed: ${r.reason}."; + }) (lib.filterAttrs (_: v: lib.isAttrs v && v ? ${r.name}) (stateEntries r.group def.value)) + ) removedStateOptions + ) options.services.gradient.state.definitionsWithLocations; +} From 45e790d7ea1ff653db1361f251397f2512805a1c Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 15:38:01 +0200 Subject: [PATCH 11/19] worker: feat: run input fetches --- backend/gradient-proto/src/handler/inbound.rs | 6 +- backend/gradient-worker/src/executor/fetch.rs | 2 +- .../src/executor/input_fetch.rs | 330 ++++++++++++++++++ backend/gradient-worker/src/executor/mod.rs | 1 + .../src/worker/message_loop.rs | 164 ++++++++- .../src/worker_pool/resolver.rs | 4 + 6 files changed, 503 insertions(+), 4 deletions(-) create mode 100644 backend/gradient-worker/src/executor/input_fetch.rs diff --git a/backend/gradient-proto/src/handler/inbound.rs b/backend/gradient-proto/src/handler/inbound.rs index f59957a53..b0c9dbbcd 100644 --- a/backend/gradient-proto/src/handler/inbound.rs +++ b/backend/gradient-proto/src/handler/inbound.rs @@ -941,7 +941,11 @@ impl RpcContext { ) { debug!(peer_id = %self.peer_id, %job_id, %query_id, count = paths.len(), ?mode, external, "CacheQuery"); let answer = async { - let project_id = match self.active.project(&job_id) { + let project_id = match self + .active + .project(&job_id) + .or_else(|| self.active.fetch_project(&job_id)) + { Some(project_id) => Some(project_id), None => self.scheduler.project_for_job(&job_id).await, }; diff --git a/backend/gradient-worker/src/executor/fetch.rs b/backend/gradient-worker/src/executor/fetch.rs index f7aa9fe60..f13e73f24 100644 --- a/backend/gradient-worker/src/executor/fetch.rs +++ b/backend/gradient-worker/src/executor/fetch.rs @@ -475,7 +475,7 @@ async fn run_input_update( updater.report_input_update(candidate, bumped).await } -async fn ssh_key_env( +pub(super) async fn ssh_key_env( ssh_key: Option<&str>, binpath_ssh: &str, ) -> Result> { diff --git a/backend/gradient-worker/src/executor/input_fetch.rs b/backend/gradient-worker/src/executor/input_fetch.rs new file mode 100644 index 000000000..e6a7c13b7 --- /dev/null +++ b/backend/gradient-worker/src/executor/input_fetch.rs @@ -0,0 +1,330 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ + +use std::collections::HashMap; +use std::convert::Infallible; +use std::sync::Arc; + +use anyhow::{Context, Result}; +use async_trait::async_trait; +use gradient_eval::pool::{DownloadSink, EvalPoolConfig, EvalWorker}; +use gradient_util::sync::Mutex; +use gradient_wire::messages::{ClientMessage, PROGRESS_INTERVAL}; +use gradient_wire::traits::JobReporter; +use gradient_wire::types::InputOutcome; +use tokio::sync::watch; +use tokio::time::MissedTickBehavior; +use tracing::{info, warn}; + +use super::fetch::ssh_key_env; +use crate::proto::job::JobUpdater; + +#[derive(Debug, Clone)] +pub struct InputFetchJob { + pub fetch_id: String, + pub project: String, + pub request: String, + pub ssh_key: Option, +} + +#[async_trait] +pub(crate) trait InputDownloader: Send + Sync { + async fn download( + &self, + request: String, + git_ssh_command: Option, + sink: Arc, + ) -> Result; +} + +#[async_trait] +pub(crate) trait FetchReport: Send + Sync { + async fn progress(&self, bytes_done: u64, bytes_total: Option); + + async fn push(&self, store_path: String) -> Result<()>; +} + +pub(crate) struct SubprocessDownloader(pub EvalPoolConfig); + +#[async_trait] +impl InputDownloader for SubprocessDownloader { + async fn download( + &self, + request: String, + git_ssh_command: Option, + sink: Arc, + ) -> Result { + let mut worker = EvalWorker::spawn_for_fetch(&self.0).await?; + let fetched = worker.fetch(request, git_ssh_command, sink).await; + worker.shutdown().await; + fetched + } +} + +pub(crate) struct UpdaterFetchReport(pub JobUpdater); + +#[async_trait] +impl FetchReport for UpdaterFetchReport { + async fn progress(&self, bytes_done: u64, bytes_total: Option) { + let sent = self + .0 + .writer + .send(ClientMessage::InputFetchProgress { + fetch_id: self.0.job_id.clone(), + bytes_done, + bytes_total, + }) + .await; + if let Err(e) = sent { + warn!(fetch_id = %self.0.job_id, error = %e, "input fetch progress not sent"); + } + } + + async fn push(&self, store_path: String) -> Result<()> { + JobReporter::push_paths(&self.0, &[(store_path, None)]).await + } +} + +pub(crate) async fn fetch_input( + job: InputFetchJob, + downloader: &dyn InputDownloader, + report: &dyn FetchReport, + binpath_ssh: &str, +) -> InputOutcome { + info!(fetch_id = %job.fetch_id, project = %job.project, "fetching an input"); + match fetch_and_push(&job, downloader, report, binpath_ssh).await { + Ok(response) => InputOutcome::Fetched { response }, + Err(e) => InputOutcome::Failed { + message: format!("{e:#}"), + }, + } +} + +async fn fetch_and_push( + job: &InputFetchJob, + downloader: &dyn InputDownloader, + report: &dyn FetchReport, + binpath_ssh: &str, +) -> Result { + let key_env = ssh_key_env(job.ssh_key.as_deref(), binpath_ssh).await?; + let git_ssh_command = key_env.as_ref().map(|(_, command)| command.clone()); + let (transfers, totals) = Transfers::new(); + let download = downloader.download(job.request.clone(), git_ssh_command, transfers); + let response = tokio::select! { + downloaded = download => downloaded?, + never = report_progress(totals, report) => match never {}, + }; + + let store_path = checked_store_path(&job.request, &response)?; + report + .push(store_path.clone()) + .await + .with_context(|| format!("push {store_path} to the cache"))?; + Ok(response) +} + +fn checked_store_path(request: &str, response: &str) -> Result { + let request: serde_json::Value = serde_json::from_str(request).context("parse the request")?; + let response: serde_json::Value = + serde_json::from_str(response).context("parse the fetch answer")?; + let fetched = response["storePath"] + .as_str() + .context("a fetch answer without a store path")?; + if let Some(expected) = request["storePath"].as_str() + && expected != fetched + { + anyhow::bail!("the input was fetched as {fetched}, expected {expected}"); + } + Ok(fetched.to_owned()) +} + +type Totals = (u64, Option); + +struct Transfers { + sizes: Mutex>, + totals: watch::Sender, +} + +impl Transfers { + fn new() -> (Arc, watch::Receiver) { + let (totals, rx) = watch::channel((0, None)); + let transfers = Arc::new(Self { + sizes: Mutex::default(), + totals, + }); + (transfers, rx) + } +} + +impl DownloadSink for Transfers { + fn transfer(&self, id: u64, done: u64, expected: u64) { + let mut sizes = self.sizes.lock(); + sizes.insert(id, (done, expected)); + let done = sizes.values().map(|(done, _)| done).sum(); + let total = sizes + .values() + .all(|(_, expected)| *expected > 0) + .then(|| sizes.values().map(|(_, expected)| expected).sum()); + self.totals.send_replace((done, total)); + } +} + +async fn report_progress( + mut totals: watch::Receiver, + report: &dyn FetchReport, +) -> Infallible { + let mut ticks = tokio::time::interval(PROGRESS_INTERVAL); + ticks.set_missed_tick_behavior(MissedTickBehavior::Delay); + loop { + ticks.tick().await; + if totals.has_changed().unwrap_or(false) { + let (done, total) = *totals.borrow_and_update(); + report.progress(done, total).await; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + const REQUEST: &str = r#"{"builtin":"fetchTree","args":{"type":"git","url":"git+ssh://git@example.com/r"},"storePath":"/nix/store/aaaa-source"}"#; + + #[derive(Default)] + struct FakeDownloader { + answer: String, + transfers: Vec<(u64, u64, u64)>, + seen: Mutex, Option)>>, + } + + impl FakeDownloader { + fn answering(answer: &str) -> Self { + Self { + answer: answer.to_owned(), + ..Self::default() + } + } + } + + #[async_trait] + impl InputDownloader for FakeDownloader { + async fn download( + &self, + _request: String, + git_ssh_command: Option, + sink: Arc, + ) -> Result { + let key = git_ssh_command.as_deref().and_then(|command| { + let path = command + .split_whitespace() + .skip_while(|w| *w != "-i") + .nth(1)?; + std::fs::read_to_string(path).ok() + }); + self.seen.lock().push((git_ssh_command, key)); + for &(id, done, expected) in &self.transfers { + sink.transfer(id, done, expected); + } + tokio::time::sleep(PROGRESS_INTERVAL * 3).await; + Ok(self.answer.clone()) + } + } + + #[derive(Default)] + struct RecordedReport { + progress: Mutex>, + pushed: Mutex>, + } + + #[async_trait] + impl FetchReport for RecordedReport { + async fn progress(&self, bytes_done: u64, bytes_total: Option) { + self.progress.lock().push((bytes_done, bytes_total)); + } + + async fn push(&self, store_path: String) -> Result<()> { + self.pushed.lock().push(store_path); + Ok(()) + } + } + + fn job(ssh_key: Option<&str>) -> InputFetchJob { + InputFetchJob { + fetch_id: "f1".into(), + project: "p1".into(), + request: REQUEST.into(), + ssh_key: ssh_key.map(str::to_owned), + } + } + + #[tokio::test(start_paused = true)] + async fn a_fetched_input_is_pushed_and_answered() { + let downloader = FakeDownloader::answering(r#"{"storePath":"/nix/store/aaaa-source"}"#); + let report = RecordedReport::default(); + + let outcome = fetch_input(job(None), &downloader, &report, "ssh").await; + + assert_eq!( + outcome, + InputOutcome::Fetched { + response: r#"{"storePath":"/nix/store/aaaa-source"}"#.into() + } + ); + assert_eq!( + *report.pushed.lock(), + vec!["/nix/store/aaaa-source".to_owned()] + ); + } + + #[tokio::test(start_paused = true)] + async fn a_response_with_another_store_path_fails() { + let downloader = FakeDownloader::answering(r#"{"storePath":"/nix/store/bbbb-source"}"#); + let report = RecordedReport::default(); + + let outcome = fetch_input(job(None), &downloader, &report, "ssh").await; + + assert_eq!( + outcome, + InputOutcome::Failed { + message: "the input was fetched as /nix/store/bbbb-source, expected /nix/store/aaaa-source".into() + } + ); + assert!(report.pushed.lock().is_empty(), "nothing pushed"); + } + + #[tokio::test(start_paused = true)] + async fn the_ssh_key_reaches_the_fetch_as_git_ssh_command() { + let downloader = FakeDownloader::answering(r#"{"storePath":"/nix/store/aaaa-source"}"#); + let report = RecordedReport::default(); + + fetch_input(job(Some("PRIVATE KEY")), &downloader, &report, "/bin/ssh").await; + fetch_input(job(None), &downloader, &report, "/bin/ssh").await; + + let seen = downloader.seen.lock(); + let (command, key) = &seen[0]; + assert!( + command + .as_deref() + .is_some_and(|c| c.starts_with("/bin/ssh -i ")), + "{command:?}" + ); + assert_eq!(key.as_deref(), Some("PRIVATE KEY")); + assert_eq!(seen[1], (None, None), "no key, no ssh command"); + } + + #[tokio::test(start_paused = true)] + async fn download_progress_is_sent_as_input_fetch_progress() { + let downloader = FakeDownloader { + transfers: vec![(1, 100, 400), (2, 50, 100)], + ..FakeDownloader::answering(r#"{"storePath":"/nix/store/aaaa-source"}"#) + }; + let report = RecordedReport::default(); + + fetch_input(job(None), &downloader, &report, "ssh").await; + + assert_eq!(*report.progress.lock(), vec![(150, Some(500))]); + } +} diff --git a/backend/gradient-worker/src/executor/mod.rs b/backend/gradient-worker/src/executor/mod.rs index 135f0f8aa..298ee919a 100644 --- a/backend/gradient-worker/src/executor/mod.rs +++ b/backend/gradient-worker/src/executor/mod.rs @@ -13,6 +13,7 @@ pub mod eval; pub(crate) mod failure; pub mod fetch; mod import; +pub(crate) mod input_fetch; pub mod log_limit; mod progress_report; mod source; diff --git a/backend/gradient-worker/src/worker/message_loop.rs b/backend/gradient-worker/src/worker/message_loop.rs index 8b7981c85..dcc4f2b75 100644 --- a/backend/gradient-worker/src/worker/message_loop.rs +++ b/backend/gradient-worker/src/worker/message_loop.rs @@ -4,6 +4,9 @@ * SPDX-License-Identifier: AGPL-3.0-only */ +use crate::executor::input_fetch::{ + InputDownloader, InputFetchJob, SubprocessDownloader, UpdaterFetchReport, fetch_input, +}; use gradient_util::sync::Mutex; use std::collections::HashMap; use std::sync::Arc; @@ -249,6 +252,8 @@ pub(super) struct MessageLoopState { credentials: CredentialStore, scorer: JobScorer, executor: JobExecutor, + downloader: Arc, + fetches: tokio::task::JoinSet<()>, config: WorkerConfig, draining: bool, stopping: bool, @@ -299,6 +304,10 @@ impl MessageLoopState { max_build: config.build.max_concurrent.get(), credentials, scorer, + downloader: Arc::new(SubprocessDownloader( + executor.evaluator.resolver().config().clone(), + )), + fetches: tokio::task::JoinSet::new(), executor, config, draining: false, @@ -325,6 +334,7 @@ impl MessageLoopState { async fn begin_stop(&mut self) -> Result<()> { self.stopping = true; self.begin_drain().await?; + self.fetches.abort_all(); let aborted = self.jobs.abort_all(); info!( aborted, @@ -471,8 +481,18 @@ impl MessageLoopState { bytes_total, ); } - ServerMessage::FetchInput { fetch_id, .. } => { - warn!(%fetch_id, "this worker does not run input fetches yet"); + ServerMessage::FetchInput { + fetch_id, + project, + request, + ssh_key, + } => { + self.on_fetch_input(InputFetchJob { + fetch_id, + project, + request, + ssh_key, + }); } } Ok(()) @@ -819,6 +839,39 @@ impl MessageLoopState { }); } + fn on_fetch_input(&mut self, fetch: InputFetchJob) { + while self.fetches.try_join_next().is_some() {} + let updater = JobUpdater::new( + fetch.fetch_id.clone(), + AssignmentHandle::new(String::new()), + self.writer.clone(), + Arc::clone(&self.cache_waiters), + Arc::clone(&self.known_derivation_waiters), + Arc::clone(&self.import_waiters), + Arc::clone(&self.input_waiters), + self.nar_recv.clone(), + self.eval_cache_recv.clone(), + Some(Arc::clone(&self.executor.store)), + crate::executor::timeline::JobTimeline::new(), + self.uploads.clone(), + ); + let downloader = Arc::clone(&self.downloader); + let writer = self.writer.clone(); + let binpath_ssh = self.executor.binpath_ssh.clone(); + self.fetches.spawn(async move { + let fetch_id = fetch.fetch_id.clone(); + let report = UpdaterFetchReport(updater); + let outcome = fetch_input(fetch, &*downloader, &report, &binpath_ssh).await; + report.0.uploads.forget_job(&fetch_id); + if let Err(e) = writer + .send(ClientMessage::InputFetched { fetch_id, outcome }) + .await + { + warn!(error = %e, "input fetch answer not sent"); + } + }); + } + async fn on_start_cluster(&mut self, attempt: String, roster: Vec) -> Result<()> { let Some((membership, held)) = self.holds.start(&attempt) else { warn!(%attempt, "StartCluster for an attempt this worker does not hold - ignoring"); @@ -1201,4 +1254,111 @@ mod tests { ); assert!(idle.is_aborted(), "a job between checkpoints sees it too"); } + + struct OfflineDownloader; + + #[async_trait::async_trait] + impl InputDownloader for OfflineDownloader { + async fn download( + &self, + _request: String, + _git_ssh_command: Option, + _sink: Arc, + ) -> Result { + Err(anyhow::anyhow!("offline")) + } + } + + fn executor(gcroots: &std::path::Path) -> JobExecutor { + let store = + crate::nix::store::LocalNixStore::connect_at("/var/empty/gradient-nonexistent.sock", 1) + .unwrap(); + let config = WorkerConfig::default(); + JobExecutor::new( + store.clone(), + crate::executor::WorkerEvaluator::new( + std::num::NonZeroUsize::MIN, + 0, + config.memory_guard().unwrap(), + String::new(), + false, + false, + ), + crate::nix::gcroots::GcRootKeeper::new(gcroots.to_str().unwrap(), Arc::new(store)), + "ssh".into(), + crate::executor::log_limit::LogRateLimits { + burst_bytes_per_min: 0, + sustained_bytes_per_hour: 0, + }, + false, + crate::executor::BuildHost { + build_cores: None, + cpu_core_score: 0, + }, + ) + } + + #[tokio::test] + #[expect( + clippy::disallowed_methods, + reason = "the test stands in for the server by hand" + )] + async fn an_input_fetch_runs_while_the_eval_pool_is_full() { + let server = gradient_test_support::prelude::MockProtoServer::bind().await; + let url = server.url().to_owned(); + let answered = tokio::spawn(async move { + let mut sc = server.accept().await; + sc.recv().await + }); + let conn = gradient_worker_client::connection::ProtoConnection::open(&url) + .await + .unwrap(); + let (writer, _reader, _flush) = conn.split(); + let gcroots = tempfile::tempdir().unwrap(); + let mut config = WorkerConfig::default(); + config.eval.max_concurrent = std::num::NonZeroU32::MIN; + let mut state = MessageLoopState::new( + writer, + config, + executor(gcroots.path()), + JobScorer::new(), + CredentialStore::new(), + ); + state.downloader = Arc::new(OfflineDownloader); + let _running = state.jobs.register( + "eval:e1".to_owned(), + JobKind::Flake, + "dispatch-1".to_owned(), + None, + ); + assert_eq!( + state.occupied(JobKind::Flake), + state.max_for(JobKind::Flake) + ); + + state + .route(ServerMessage::FetchInput { + fetch_id: "f1".into(), + project: "p1".into(), + request: "{}".into(), + ssh_key: None, + }) + .await + .unwrap(); + + let sent = tokio::time::timeout(std::time::Duration::from_secs(5), answered) + .await + .expect("the fetch answers while every eval slot is taken") + .unwrap() + .unwrap(); + assert_eq!( + sent, + ClientMessage::InputFetched { + fetch_id: "f1".into(), + outcome: gradient_wire::types::InputOutcome::Failed { + message: "offline".into() + }, + } + ); + } } diff --git a/backend/gradient-worker/src/worker_pool/resolver.rs b/backend/gradient-worker/src/worker_pool/resolver.rs index 04ce1db8e..78023e01a 100644 --- a/backend/gradient-worker/src/worker_pool/resolver.rs +++ b/backend/gradient-worker/src/worker_pool/resolver.rs @@ -92,6 +92,10 @@ impl WorkerPoolResolver { self.pool.cache_dir() } + pub(crate) fn config(&self) -> &EvalPoolConfig { + &self.config + } + pub async fn shutdown(&self) { self.pool.shutdown().await; } From 4145f6a84068f80408f020ea5eea78c6ce07e76a Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 15:42:17 +0200 Subject: [PATCH 12/19] worker: feat: inputs through the server --- backend/gradient-eval/src/eval_worker.rs | 5 +- backend/gradient-eval/src/ipc.rs | 2 +- backend/gradient-eval/src/pool/transport.rs | 11 +- backend/gradient-pool/src/input_request.rs | 13 +- backend/gradient-worker/options.json | 13 - backend/gradient-worker/src/config/options.rs | 5 - .../executor/{import.rs => dependencies.rs} | 188 +++++- backend/gradient-worker/src/executor/eval.rs | 4 +- backend/gradient-worker/src/executor/fetch.rs | 551 ++++++------------ .../src/executor/input_fetch.rs | 25 +- backend/gradient-worker/src/executor/mod.rs | 4 +- backend/gradient-worker/src/main.rs | 6 +- .../src/worker_pool/input_fetch.rs | 129 ---- .../gradient-worker/src/worker_pool/mod.rs | 2 - .../src/worker_pool/resolver.rs | 32 +- cli/src/commands/base.rs | 2 +- 16 files changed, 416 insertions(+), 576 deletions(-) rename backend/gradient-worker/src/executor/{import.rs => dependencies.rs} (60%) delete mode 100644 backend/gradient-worker/src/worker_pool/input_fetch.rs diff --git a/backend/gradient-eval/src/eval_worker.rs b/backend/gradient-eval/src/eval_worker.rs index 1218eddd5..1e4a72c7f 100644 --- a/backend/gradient-eval/src/eval_worker.rs +++ b/backend/gradient-eval/src/eval_worker.rs @@ -86,10 +86,13 @@ impl ParentPipe { } } -pub fn run_eval_worker(collect_stats: bool, fetch_through_parent: bool) -> std::io::Result<()> { +pub fn run_eval_worker(collect_stats: bool) -> std::io::Result<()> { let mut stdout = std::io::stdout(); stdout.write_all(&[EVAL_IPC_VERSION])?; stdout.flush()?; + let mut fetch_through_parent = [0u8; 1]; + std::io::stdin().read_exact(&mut fetch_through_parent)?; + let fetch_through_parent = fetch_through_parent[0] == 1; let pipe = Arc::new(ParentPipe::new( std::io::stdin(), diff --git a/backend/gradient-eval/src/ipc.rs b/backend/gradient-eval/src/ipc.rs index 1d51e7c5f..94ad9eef5 100644 --- a/backend/gradient-eval/src/ipc.rs +++ b/backend/gradient-eval/src/ipc.rs @@ -14,7 +14,7 @@ use crate::stats::StatsDelta; /// The version must be bumped whenever the frame layout or a type's rkyv shape changes. /// A mismatch can only happen when the binary is replaced mid-run. -pub const EVAL_IPC_VERSION: u8 = 11; +pub const EVAL_IPC_VERSION: u8 = 12; pub const MAX_FRAME_BYTES: u32 = 64 * 1024 * 1024; diff --git a/backend/gradient-eval/src/pool/transport.rs b/backend/gradient-eval/src/pool/transport.rs index fdf6c5bb7..ee564f62c 100644 --- a/backend/gradient-eval/src/pool/transport.rs +++ b/backend/gradient-eval/src/pool/transport.rs @@ -54,9 +54,6 @@ fn eval_subprocess_command(exe: &Path, config: &SpawnConfig) -> Command { if config.metrics { command.arg("--eval-metrics"); } - if config.fetch_through_parent { - command.arg("--eval-fetch-through-parent"); - } command.env("NIX_CACHE_HOME", &config.cache_dir); if let Some(dir) = &config.trace_dir { command.env("GRADIENT_WORKER_LOG_TRACE_DIR", dir); @@ -172,6 +169,11 @@ impl EvalWorker { } worker.expect_handshake().await?; + worker + .stdin + .write_all(&[u8::from(config.fetch_through_parent)]) + .await + .context("tell the eval worker how to fetch inputs")?; // Register the pid after the handshake: a kill mid-handshake is a spawn error, not an // attribute error. @@ -741,8 +743,7 @@ mod tests { args, vec![ "--eval-subprocess".to_string(), - "--eval-metrics".to_string(), - "--eval-fetch-through-parent".to_string() + "--eval-metrics".to_string() ] ); assert!( diff --git a/backend/gradient-pool/src/input_request.rs b/backend/gradient-pool/src/input_request.rs index 2d686afaf..8cf0310bb 100644 --- a/backend/gradient-pool/src/input_request.rs +++ b/backend/gradient-pool/src/input_request.rs @@ -10,7 +10,6 @@ use serde_json::Value; #[derive(Debug, Clone, PartialEq, Eq)] pub struct InputRequest { pub fetch_key: String, - pub store_path: Option, pub needs_secrets: bool, } @@ -26,13 +25,14 @@ impl InputRequest { return Err("an input request without arguments".to_owned()); } - let store_path = json["storePath"].as_str().map(str::to_owned); - let fetch_key = store_path.clone().unwrap_or_else(|| { - sorted(&serde_json::json!({ "builtin": builtin, "args": args })).to_string() - }); + let fetch_key = json["storePath"] + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| { + sorted(&serde_json::json!({ "builtin": builtin, "args": args })).to_string() + }); Ok(Self { fetch_key, - store_path, needs_secrets: args["url"] .as_str() .is_some_and(check_repository_url_is_ssh), @@ -69,7 +69,6 @@ mod tests { assert_eq!(first.fetch_key, second.fetch_key); assert!(first.needs_secrets); - assert_eq!(first.store_path, None); } #[test] diff --git a/backend/gradient-worker/options.json b/backend/gradient-worker/options.json index 0e8b2dd44..2522ee8be 100644 --- a/backend/gradient-worker/options.json +++ b/backend/gradient-worker/options.json @@ -175,19 +175,6 @@ "false" ] }, - { - "group": "", - "field": "eval_fetch_through_parent", - "flag": "eval-fetch-through-parent", - "env": "GRADIENT_WORKER_EVAL_FETCH_THROUGH_PARENT", - "type": "bool", - "description": "Let the eval subprocess ask its parent for missing inputs (internal - do not set manually).", - "default": null, - "values": [ - "true", - "false" - ] - }, { "group": "", "field": "eval_driver", diff --git a/backend/gradient-worker/src/config/options.rs b/backend/gradient-worker/src/config/options.rs index b09762c83..c46c219dc 100644 --- a/backend/gradient-worker/src/config/options.rs +++ b/backend/gradient-worker/src/config/options.rs @@ -106,10 +106,6 @@ pub struct WorkerConfig { #[arg(hide = true)] pub eval_subprocess: bool, - /// Let the eval subprocess ask its parent for missing inputs (internal - do not set manually). - #[arg(hide = true)] - pub eval_fetch_through_parent: bool, - /// Drive one eval subprocess over the production rkyv transport from a JSONL request file. /// This internal harness serves the NixOS VM integration test. #[arg(hide = true)] @@ -151,7 +147,6 @@ impl Default for WorkerConfig { idle_timeout: None, max_lifetime: None, eval_subprocess: false, - eval_fetch_through_parent: false, eval_driver: None, capabilities: CapabilitiesArgs::default(), system: SystemArgs::default(), diff --git a/backend/gradient-worker/src/executor/import.rs b/backend/gradient-worker/src/executor/dependencies.rs similarity index 60% rename from backend/gradient-worker/src/executor/import.rs rename to backend/gradient-worker/src/executor/dependencies.rs index 6a8bb4a45..968dc1ccc 100644 --- a/backend/gradient-worker/src/executor/import.rs +++ b/backend/gradient-worker/src/executor/dependencies.rs @@ -13,23 +13,27 @@ use gradient_derivation::{Derivation, parse_drv}; use gradient_sources::EvalDependencies; use gradient_util::store_path::strip_store_prefix; use gradient_util::sync::Mutex; -use gradient_wire::types::ImportOutcome; +use gradient_wire::traits::InputProgressSink; +use gradient_wire::types::{ImportOutcome, InputOutcome}; use super::AbortSignal; use super::eval::{Root, walk_and_publish}; use crate::traits::{DrvReader, JobReporter}; +use gradient_wire::messages::QueryMode; type ImportFuture<'a> = Shared>>; +type InputFuture<'a> = Shared>>; -pub(super) struct EvalImportBuilder<'a> { +pub(super) struct EvalDependencySource<'a> { reporter: &'a dyn JobReporter, drv_reader: &'a dyn DrvReader, abort: AbortSignal, names: Arc>>, in_flight: Mutex>>, + inputs: InputSource<'a>, } -impl<'a> EvalImportBuilder<'a> { +impl<'a> EvalDependencySource<'a> { pub(super) fn new( reporter: &'a dyn JobReporter, drv_reader: &'a dyn DrvReader, @@ -41,6 +45,7 @@ impl<'a> EvalImportBuilder<'a> { abort, names: Arc::default(), in_flight: Mutex::default(), + inputs: InputSource::new(reporter), } } @@ -81,7 +86,7 @@ impl<'a> EvalImportBuilder<'a> { } #[async_trait] -impl EvalDependencies for EvalImportBuilder<'_> { +impl EvalDependencies for EvalDependencySource<'_> { async fn build_imports(&self, derived_paths: Vec) -> Result<(), String> { let mut drv_paths: Vec = derived_paths .iter() @@ -98,11 +103,95 @@ impl EvalDependencies for EvalImportBuilder<'_> { join_all(imports).await.into_iter().collect() } - async fn fetch_input(&self, _request: String) -> Result { - Err("this worker's evaluator fetches its inputs itself".to_owned()) + async fn fetch_input(&self, request: String) -> Result { + self.inputs.fetch(request, Arc::new(NoProgress)).await } } +struct NoProgress; + +impl InputProgressSink for NoProgress { + fn transfer(&self, _bytes_done: u64, _bytes_total: Option) {} +} + +pub(super) struct InputSource<'a> { + reporter: &'a dyn JobReporter, + in_flight: Mutex>>, +} + +impl<'a> InputSource<'a> { + pub(super) fn new(reporter: &'a dyn JobReporter) -> Self { + Self { + reporter, + in_flight: Mutex::default(), + } + } + + pub(super) async fn fetch( + &self, + request: String, + progress: Arc, + ) -> Result { + let store_path = store_path_of(&request)?; + let key = store_path.clone().unwrap_or_else(|| request.clone()); + let fetch = self + .in_flight + .lock() + .entry(key) + .or_insert_with(|| { + fetch_input(self.reporter, request, store_path, progress) + .boxed() + .shared() + }) + .clone(); + fetch.await + } +} + +async fn fetch_input( + reporter: &dyn JobReporter, + request: String, + store_path: Option, + progress: Arc, +) -> Result { + if let Some(path) = store_path + && in_cache(reporter, &path).await + && reporter.pull_paths(vec![path.clone()]).await.is_ok() + { + return Ok(serde_json::json!({ "storePath": path }).to_string()); + } + + let outcome = reporter + .request_input(request, progress) + .await + .map_err(|e| format!("input request failed: {e:#}"))?; + match outcome { + InputOutcome::Fetched { response } => { + let fetched = store_path_of(&response)? + .ok_or_else(|| "the fetch answer names no store path".to_owned())?; + reporter + .pull_paths(vec![fetched.clone()]) + .await + .map_err(|e| format!("pulling the fetched input {fetched} failed: {e:#}"))?; + Ok(response) + } + InputOutcome::Failed { message } => Err(message), + } +} + +async fn in_cache(reporter: &dyn JobReporter, path: &str) -> bool { + reporter + .query_cache(vec![path.to_owned()], QueryMode::Pull) + .await + .is_ok_and(|cached| cached.iter().any(|c| c.cached && c.path == path)) +} + +fn store_path_of(json: &str) -> Result, String> { + let value: serde_json::Value = + serde_json::from_str(json).map_err(|e| format!("malformed input request: {e}"))?; + Ok(value["storePath"].as_str().map(str::to_owned)) +} + async fn import_derivation( reporter: &dyn JobReporter, drv_reader: &dyn DrvReader, @@ -111,7 +200,7 @@ async fn import_derivation( drv_path: String, ) -> Result<(), String> { let drv = read_drv(drv_reader, &drv_path).await?; - let attr = EvalImportBuilder::import_attr(&mut names.lock(), &drv.system, &drv_path); + let attr = EvalDependencySource::import_attr(&mut names.lock(), &drv.system, &drv_path); let root = Root { drv_path: drv_path.clone(), attr, @@ -216,7 +305,7 @@ mod tests { async fn concurrent_requests_for_a_path_share_the_result() { let reporter = RecordingJobReporter::new().with_import_outcome(SOURCE_A, failed(SOURCE_A)); let reader = reader(); - let builder = EvalImportBuilder::new(&reporter, &reader, AbortSignal::never()); + let builder = EvalDependencySource::new(&reporter, &reader, AbortSignal::never()); let (first, second) = tokio::join!( builder.build_imports(vec![format!("{SOURCE_A}^out")]), @@ -232,7 +321,7 @@ mod tests { async fn a_completed_import_pulls_its_outputs() { let reporter = RecordingJobReporter::new(); let reader = reader(); - let builder = EvalImportBuilder::new(&reporter, &reader, AbortSignal::never()); + let builder = EvalDependencySource::new(&reporter, &reader, AbortSignal::never()); builder .build_imports(vec![format!("{SOURCE_A}^out")]) @@ -266,7 +355,7 @@ mod tests { async fn a_failed_import_names_the_build() { let reporter = RecordingJobReporter::new().with_import_outcome(SOURCE_A, failed(SOURCE_A)); let reader = reader(); - let builder = EvalImportBuilder::new(&reporter, &reader, AbortSignal::never()); + let builder = EvalDependencySource::new(&reporter, &reader, AbortSignal::never()); let built = builder.build_imports(vec![format!("{SOURCE_A}^out")]).await; @@ -287,7 +376,7 @@ mod tests { async fn a_repeated_name_gets_a_hash_suffix() { let reporter = RecordingJobReporter::new(); let reader = reader(); - let builder = EvalImportBuilder::new(&reporter, &reader, AbortSignal::never()); + let builder = EvalDependencySource::new(&reporter, &reader, AbortSignal::never()); builder .build_imports(vec![format!("{SOURCE_A}^out"), format!("{SOURCE_B}^out")]) @@ -316,13 +405,86 @@ mod tests { let drv = "/nix/store/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa-hello-2.12.drv"; assert_eq!( - EvalImportBuilder::import_attr(&mut names, "aarch64-linux", drv), + EvalDependencySource::import_attr(&mut names, "aarch64-linux", drv), "other.aarch64-linux.hello-2.12" ); assert_eq!( - EvalImportBuilder::import_attr(&mut names, "aarch64-linux", drv), + EvalDependencySource::import_attr(&mut names, "aarch64-linux", drv), "other.aarch64-linux.hello-2.12", "the same derivation keeps its attribute" ); } + + const INPUT: &str = "/nix/store/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb-source"; + + fn request() -> String { + format!( + r#"{{"builtin":"fetchTree","args":{{"type":"tarball","url":"https://example.com/x.tar.gz"}},"storePath":"{INPUT}"}}"# + ) + } + + fn fetched() -> InputOutcome { + InputOutcome::Fetched { + response: format!(r#"{{"storePath":"{INPUT}","attrs":{{"type":"tarball"}}}}"#), + } + } + + fn pulled(reporter: &RecordingJobReporter) -> Vec> { + reporter + .events() + .into_iter() + .filter_map(|e| match e { + ReportedEvent::PathsPulled { paths } => Some(paths), + _ => None, + }) + .collect() + } + + #[tokio::test] + async fn an_input_in_the_cache_is_pulled_without_a_fetch() { + let reporter = RecordingJobReporter::new().with_cached_paths(vec![INPUT.into()]); + let reader = reader(); + let source = EvalDependencySource::new(&reporter, &reader, AbortSignal::never()); + + let answer = source.fetch_input(request()).await; + + assert_eq!(answer, Ok(format!(r#"{{"storePath":"{INPUT}"}}"#))); + assert_eq!(pulled(&reporter), vec![vec![INPUT.to_owned()]]); + assert!(reporter.input_requests().is_empty(), "no InputRequest"); + } + + #[tokio::test] + async fn a_missing_input_is_fetched_through_the_server_and_pulled() { + let reporter = RecordingJobReporter::new().with_input_outcome(INPUT, fetched()); + let reader = reader(); + let source = EvalDependencySource::new(&reporter, &reader, AbortSignal::never()); + + let answer = source.fetch_input(request()).await; + + let InputOutcome::Fetched { response } = fetched() else { + unreachable!() + }; + assert_eq!(answer, Ok(response)); + assert_eq!(reporter.input_requests(), vec![request()]); + assert_eq!(pulled(&reporter), vec![vec![INPUT.to_owned()]]); + } + + #[tokio::test] + async fn a_failed_input_is_not_requested_twice() { + let reporter = RecordingJobReporter::new().with_input_outcome( + INPUT, + InputOutcome::Failed { + message: "HTTP error 404".into(), + }, + ); + let reader = reader(); + let source = EvalDependencySource::new(&reporter, &reader, AbortSignal::never()); + + let first = source.fetch_input(request()).await; + let second = source.fetch_input(request()).await; + + assert_eq!(first, Err("HTTP error 404".to_owned())); + assert_eq!(second, first); + assert_eq!(reporter.input_requests().len(), 1); + } } diff --git a/backend/gradient-worker/src/executor/eval.rs b/backend/gradient-worker/src/executor/eval.rs index 7cf7d21b8..9f0c962e4 100644 --- a/backend/gradient-worker/src/executor/eval.rs +++ b/backend/gradient-worker/src/executor/eval.rs @@ -9,7 +9,7 @@ use std::num::NonZeroUsize; use std::sync::Arc; use std::time::Instant; -use super::import::EvalImportBuilder; +use super::dependencies::EvalDependencySource; use super::progress_report::{ChangeReporter, thunk_progress}; use crate::worker_pool::WorkerPoolResolver; use anyhow::{Context, Result}; @@ -928,7 +928,7 @@ async fn discover( let eval_overrides = eval_input_overrides(job, local_flake_path); debug!(repo = %repo, "listing flake derivations"); - let imports = EvalImportBuilder::new(reporter, drv_reader, abort.clone()); + let imports = EvalDependencySource::new(reporter, drv_reader, abort.clone()); let send = move |resolved: Vec| { for derivation in resolved { let _ = found.send(derivation); diff --git a/backend/gradient-worker/src/executor/fetch.rs b/backend/gradient-worker/src/executor/fetch.rs index f13e73f24..9ab03b27e 100644 --- a/backend/gradient-worker/src/executor/fetch.rs +++ b/backend/gradient-worker/src/executor/fetch.rs @@ -8,17 +8,17 @@ use std::collections::{HashMap, HashSet}; use std::sync::Arc; use anyhow::{Context, Result}; +use gradient_util::sync::Mutex; use gradient_wire::messages::{FlakeJob, FlakeSource}; -use gradient_wire::traits::{EvalProgressSink, JobReporter, WorkerStore}; -use gradient_wire::types::{EvalProgress, InputFetchState}; -use tempfile::NamedTempFile; +use gradient_wire::traits::{EvalProgressSink, InputProgressSink, JobReporter, WorkerStore}; +use gradient_wire::types::{EvalProgress, InputFetch, InputFetchState}; use tracing::{debug, info}; use super::AbortSignal; +use super::dependencies::InputSource; use super::progress_report::ChangeReporter; use crate::error::WorkerError; use crate::proto::credentials::CredentialStore; -use crate::worker_pool::{DownloadTarget, InputBoard, InputFetcher}; pub struct FetchOutcome { pub source_path: String, @@ -39,8 +39,6 @@ pub async fn fetch_repository( updater: &mut dyn JobReporter, credentials: &CredentialStore, store: &dyn WorkerStore, - fetcher: &dyn InputFetcher, - binpath_ssh: &str, mut abort: AbortSignal, ) -> Result { if abort.is_aborted() { @@ -128,11 +126,11 @@ pub async fn fetch_repository( .as_ref() .map(|lock| locked_inputs(lock, &overridden)) .unwrap_or_default(); - let key_env = ssh_key_env(ssh_key.as_deref(), binpath_ssh).await?; - let git_ssh_command = key_env.as_ref().map(|(_, command)| command.clone()); - pull_from_cache(updater, store, &inputs, &mut abort).await?; let sink = updater.eval_progress_sink(); - let fetched = fetch_inputs(inputs, store, fetcher, &*sink, git_ssh_command, &mut abort).await?; + let fetched = { + let source = InputSource::new(&*updater); + fetch_inputs(inputs, store, &source, &*sink, &mut abort).await? + }; for msg in &fetched.warnings { updater .send_eval_message( @@ -156,9 +154,8 @@ pub async fn fetch_repository( #[derive(Debug, Clone)] struct LockedInput { name: String, - domain: String, store_path: String, - locked: String, + request: String, } async fn read_lock(flake_root: &str) -> Result> { @@ -207,9 +204,13 @@ fn locked_inputs(lock: &serde_json::Value, overridden: &HashSet) -> Vec< }; inputs.push(LockedInput { name: root_names.get(key).copied().unwrap_or(key).to_owned(), - domain: input_domain(locked), + request: serde_json::json!({ + "builtin": "fetchTree", + "args": fetcher_attrs(locked), + "storePath": store_path, + }) + .to_string(), store_path, - locked: fetcher_attrs(locked).to_string(), }); } inputs.sort_by(|a, b| a.name.cmp(&b.name)); @@ -236,75 +237,76 @@ fn source_store_path(nar_hash: &str) -> Option { Some(store_dir.display(&path).to_string()) } -fn input_domain(locked: &serde_json::Value) -> String { - let default_host = match locked["type"].as_str() { - Some("github") => Some("github.com"), - Some("gitlab") => Some("gitlab.com"), - Some("sourcehut") => Some("git.sr.ht"), - _ => None, - }; - let host = match default_host { - Some(default) => locked["host"].as_str().unwrap_or(default), - None => locked["url"].as_str().map(url_host).unwrap_or_default(), - }; - second_level(host) +#[derive(Debug)] +struct InputRow { + name: String, + state: InputFetchState, + downloaded_bytes: u64, + expected_bytes: Option, } -fn url_host(url: &str) -> &str { - let rest = url.split_once("://").map_or(url, |(_, rest)| rest); - let authority = rest.split('/').next().unwrap_or_default(); - let host = authority - .rsplit_once('@') - .map_or(authority, |(_, host)| host); - host.split(':').next().unwrap_or_default() +#[derive(Debug)] +struct InputBoard { + rows: Mutex>, } -fn second_level(host: &str) -> String { - if host.parse::().is_ok() { - return host.to_owned(); +impl InputBoard { + fn new(rows: Vec<(String, InputFetchState)>) -> Arc { + let rows = rows + .into_iter() + .map(|(name, state)| InputRow { + name, + state, + downloaded_bytes: 0, + expected_bytes: None, + }) + .collect(); + Arc::new(Self { + rows: Mutex::new(rows), + }) } - let labels: Vec<&str> = host.rsplitn(3, '.').collect(); - match labels.as_slice() { - [tld, sld, ..] => format!("{sld}.{tld}"), - _ => host.to_owned(), + + fn set_state(&self, index: usize, state: InputFetchState) { + self.rows.lock()[index].state = state; } -} -#[tracing::instrument(level = "debug", skip_all, fields(inputs = inputs.len()))] -async fn pull_from_cache( - updater: &mut dyn JobReporter, - store: &dyn WorkerStore, - inputs: &[LockedInput], - abort: &mut AbortSignal, -) -> Result<()> { - let mut paths: Vec = inputs.iter().map(|i| i.store_path.clone()).collect(); - paths.sort(); - paths.dedup(); - let missing = missing_paths(store, &paths).await?; - if missing.is_empty() { - return Ok(()); + fn transfer(&self, index: usize, bytes_done: u64, bytes_total: Option) { + let row = &mut self.rows.lock()[index]; + row.downloaded_bytes = bytes_done; + row.expected_bytes = bytes_total; } - let pulled = tokio::select! { - biased; - () = abort.aborted() => anyhow::bail!(WorkerError::FlakeInputFetchAborted), - pulled = updater.pull_paths(missing) => pulled, - }; - if let Err(e) = pulled { - debug!( - error = format!("{e:#}"), - "flake inputs not pulled from the cache, fetching them from their origin" - ); + + fn snapshot(&self) -> Vec { + self.rows + .lock() + .iter() + .map(|row| InputFetch { + name: row.name.clone(), + state: row.state, + downloaded_bytes: row.downloaded_bytes, + expected_bytes: row.expected_bytes.unwrap_or_default(), + }) + .collect() + } +} + +struct BoardRow { + board: Arc, + index: usize, +} + +impl InputProgressSink for BoardRow { + fn transfer(&self, bytes_done: u64, bytes_total: Option) { + self.board.transfer(self.index, bytes_done, bytes_total); } - Ok(()) } #[tracing::instrument(level = "debug", skip_all, fields(inputs = inputs.len()))] async fn fetch_inputs( inputs: Vec, store: &dyn WorkerStore, - fetcher: &dyn InputFetcher, + source: &InputSource<'_>, sink: &dyn EvalProgressSink, - git_ssh_command: Option, abort: &mut AbortSignal, ) -> Result { let paths: Vec = inputs.iter().map(|i| i.store_path.clone()).collect(); @@ -323,46 +325,29 @@ async fn fetch_inputs( .collect(), ); - let mut by_domain: HashMap<&str, Vec> = HashMap::new(); - for (index, input) in inputs - .iter() - .enumerate() - .filter(|(_, i)| missing.contains(&i.store_path)) - { - by_domain - .entry(input.domain.as_str()) - .or_default() - .push(index); - } - let fetch_domain = |indexes: Vec| { - let (board, inputs, git_ssh_command) = (&board, &inputs, &git_ssh_command); + let fetch_one = |index: usize| { + let (board, inputs) = (&board, &inputs); async move { - let mut done = Vec::new(); - for index in indexes { - board.set_state(index, InputFetchState::Fetching); - let target = DownloadTarget { - board: Arc::clone(board), - index, - }; - let fetched = fetcher - .fetch_input( - inputs[index].locked.clone(), - git_ssh_command.clone(), - target, - ) - .await; - let state = if fetched.is_ok() { - InputFetchState::Done - } else { - InputFetchState::Failed - }; - board.set_state(index, state); - done.push((index, fetched)); - } - done + board.set_state(index, InputFetchState::Fetching); + let progress = Arc::new(BoardRow { + board: Arc::clone(board), + index, + }); + let fetched = source.fetch(inputs[index].request.clone(), progress).await; + let state = if fetched.is_ok() { + InputFetchState::Done + } else { + InputFetchState::Failed + }; + board.set_state(index, state); + (index, fetched) } }; - let fetches = futures::future::join_all(by_domain.into_values().map(fetch_domain)); + let fetches = futures::future::join_all( + (0..inputs.len()) + .filter(|&i| missing.contains(&inputs[i].store_path)) + .map(fetch_one), + ); let snapshot = || { (!missing.is_empty()).then(|| EvalProgress::Fetching { @@ -380,13 +365,13 @@ async fn fetch_inputs( let mut all: Vec = paths.into_iter().filter(|p| !missing.contains(p)).collect(); let mut warnings = Vec::new(); - for (index, result) in fetched.into_iter().flatten() { + for (index, result) in fetched { match result { - Ok(path) => all.push(path), - Err(e) => warnings.push(format!( + Ok(_) => all.push(inputs[index].store_path.clone()), + Err(message) => warnings.push(format!( "skipping flake input '{}': {}", inputs[index].name, - format!("{e:#}").trim() + message.trim() )), } } @@ -475,30 +460,6 @@ async fn run_input_update( updater.report_input_update(candidate, bumped).await } -pub(super) async fn ssh_key_env( - ssh_key: Option<&str>, - binpath_ssh: &str, -) -> Result> { - use std::os::unix::fs::PermissionsExt; - - let Some(key) = ssh_key else { - return Ok(None); - }; - let kf = NamedTempFile::with_suffix(".key").context("create SSH key temp file")?; - tokio::fs::set_permissions(kf.path(), std::fs::Permissions::from_mode(0o600)) - .await - .context("chmod SSH key file")?; - tokio::fs::write(kf.path(), key.as_bytes()) - .await - .context("write SSH key file")?; - let ssh_command = format!( - "{} -i {} -o IdentitiesOnly=yes -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null", - binpath_ssh, - kf.path().display() - ); - Ok(Some((kf, ssh_command))) -} - async fn require_present(store: &dyn WorkerStore, paths: &[String]) -> Result<()> { let missing = missing_paths(store, paths).await?; if !missing.is_empty() { @@ -665,8 +626,6 @@ mod tests { &mut reporter, &credentials, &FakeWorkerStore::new(), - &FakeFetcher::new(""), - "ssh", no_abort(), ) .await; @@ -695,8 +654,6 @@ mod tests { &mut reporter, &credentials, &FakeWorkerStore::new(), - &FakeFetcher::new(""), - "ssh", no_abort(), ) .await; @@ -774,17 +731,9 @@ mod tests { let mut reporter = RecordingJobReporter::new(); let store = FakeWorkerStore::new(); - let outcome = fetch_repository( - &job, - &mut reporter, - &credentials, - &store, - &FakeFetcher::new(""), - "ssh", - no_abort(), - ) - .await - .unwrap(); + let outcome = fetch_repository(&job, &mut reporter, &credentials, &store, no_abort()) + .await + .unwrap(); assert!(outcome.source_path.starts_with("/nix/store/")); assert!(outcome.source_path.ends_with("-source")); @@ -961,10 +910,7 @@ mod tests { assert!(warnings[0].contains("missing")); } - use crate::worker_pool::{DownloadTarget, InputFetcher}; - use gradient_util::sync::Mutex; - use gradient_wire::types::{EvalProgress, InputFetchState}; - use std::sync::atomic::{AtomicBool, Ordering}; + use gradient_wire::types::InputOutcome; fn lock_with(nodes: serde_json::Value) -> serde_json::Value { serde_json::json!({ "root": "root", "version": 7, "nodes": nodes }) @@ -993,18 +939,8 @@ mod tests { })); let overridden = ["home".to_owned()].into_iter().collect(); let inputs = locked_inputs(&lock, &overridden); - let names: Vec<_> = inputs - .iter() - .map(|i| (i.name.as_str(), i.domain.as_str())) - .collect(); - assert_eq!( - names, - vec![ - ("nixpkgs", "github.com"), - ("systems", "github.com"), - ("utils", "github.com") - ] - ); + let names: Vec<_> = inputs.iter().map(|i| i.name.as_str()).collect(); + assert_eq!(names, vec!["nixpkgs", "systems", "utils"]); assert!( inputs .iter() @@ -1021,264 +957,163 @@ mod tests { "rev": "r", "dir": "lib", "narHash": nar_hash('A') } }, })); let inputs = locked_inputs(&lock, &HashSet::new()); - let locked: serde_json::Value = serde_json::from_str(&inputs[0].locked).unwrap(); - assert_eq!( - locked, - serde_json::json!({ "type": "github", "owner": "NixOS", "repo": "nixpkgs", - "rev": "r", "narHash": nar_hash('A') }) - ); - } - - #[test] - fn domains_reduce_to_the_second_level() { - let d = |v: serde_json::Value| input_domain(&v); - assert_eq!(d(serde_json::json!({"type":"github"})), "github.com"); - assert_eq!( - d(serde_json::json!({"type":"github","host":"git.corp.example.org"})), - "example.org" - ); - assert_eq!(d(serde_json::json!({"type":"sourcehut"})), "sr.ht"); + let request: serde_json::Value = serde_json::from_str(&inputs[0].request).unwrap(); assert_eq!( - d(serde_json::json!({"type":"git","url":"ssh://git@gitlab.com/a/b"})), - "gitlab.com" + request, + serde_json::json!({ "builtin": "fetchTree", "storePath": inputs[0].store_path, + "args": { "type": "github", "owner": "NixOS", "repo": "nixpkgs", + "rev": "r", "narHash": nar_hash('A') } }) ); - assert_eq!( - d(serde_json::json!({"type":"tarball","url":"https://releases.nixos.org/x.tar.xz"})), - "nixos.org" - ); - assert_eq!(d(serde_json::json!({"type":"path","path":"/x"})), ""); - } - - struct FakeFetcher { - in_flight: Mutex>, - overlapped: AtomicBool, - fail: &'static str, - } - - impl FakeFetcher { - fn new(fail: &'static str) -> Self { - Self { - in_flight: Mutex::default(), - overlapped: AtomicBool::default(), - fail, - } - } } - #[async_trait::async_trait] - impl InputFetcher for FakeFetcher { - async fn fetch_input( - &self, - locked: String, - _: Option, - target: DownloadTarget, - ) -> anyhow::Result { - let v: serde_json::Value = serde_json::from_str(&locked)?; - let domain = input_domain(&v); - { - let mut in_flight = self.in_flight.lock(); - let n = in_flight.entry(domain.clone()).or_default(); - *n += 1; - assert_eq!(*n, 1, "two downloads in flight for {domain}"); - if in_flight.values().filter(|n| **n > 0).count() > 1 { - self.overlapped.store(true, Ordering::SeqCst); - } - } - target.board.transfer(target.index, 1, 10, 0); - tokio::time::sleep(std::time::Duration::from_millis(20)).await; - *self.in_flight.lock().get_mut(&domain).unwrap() -= 1; - if v["repo"] == self.fail { - anyhow::bail!("404"); - } - Ok(format!("/nix/store/{}-source", v["repo"].as_str().unwrap())) + fn input(name: &str, repo: &str) -> LockedInput { + let store_path = format!("/nix/store/{repo}-source"); + LockedInput { + name: name.into(), + request: serde_json::json!({ "builtin": "fetchTree", + "args": { "type": "github", "repo": repo }, "storePath": store_path }) + .to_string(), + store_path, } } - fn input(name: &str, domain: &str, repo: &str) -> LockedInput { - LockedInput { - name: name.into(), - domain: domain.into(), - store_path: format!("/nix/store/{repo}-source"), - locked: serde_json::json!({ "type": "github", "host": domain, "repo": repo }) - .to_string(), + fn fetched(repo: &str) -> InputOutcome { + InputOutcome::Fetched { + response: format!(r#"{{"storePath":"/nix/store/{repo}-source"}}"#), } } - #[tokio::test] - async fn domains_overlap_while_each_domain_stays_serial_and_a_failure_is_a_warning() { - let fetcher = FakeFetcher::new("bad"); - let reporter = RecordingJobReporter::new(); - let store = FakeWorkerStore::default(); - let inputs = vec![ - input("a", "github.com", "a"), - input("b", "github.com", "bad"), - input("c", "gitlab.com", "c"), - ]; - let FetchedInputs { - paths, warnings, .. - } = fetch_inputs( + async fn fetch_with( + reporter: &RecordingJobReporter, + store: &FakeWorkerStore, + inputs: Vec, + abort: &mut AbortSignal, + ) -> Result { + let source = InputSource::new(reporter); + fetch_inputs( inputs, - &store, - &fetcher, + store, + &source, &*reporter.eval_progress_sink(), - None, - &mut no_abort(), + abort, ) .await - .unwrap(); - assert!(fetcher.overlapped.load(Ordering::SeqCst)); - assert_eq!(paths.len(), 2); - assert_eq!(warnings, vec!["skipping flake input 'b': 404".to_owned()]); - let Some(ReportedEvent::EvalProgress(EvalProgress::Fetching { inputs })) = - reporter.events().last().cloned() - else { - panic!("no final fetch snapshot"); + } + + fn rows(progress: Option) -> Vec<(String, InputFetchState, u64, u64)> { + let Some(EvalProgress::Fetching { inputs }) = progress else { + panic!("no fetch snapshot"); }; - let states: Vec<_> = inputs - .iter() - .map(|i| (i.name.as_str(), i.state, i.downloaded_bytes)) - .collect(); - assert_eq!( - states, - vec![ - ("a", InputFetchState::Done, 10), - ("b", InputFetchState::Failed, 10), - ("c", InputFetchState::Done, 10), - ] - ); + inputs + .into_iter() + .map(|i| (i.name, i.state, i.downloaded_bytes, i.expected_bytes)) + .collect() } #[tokio::test] - async fn inputs_sharing_a_source_return_it_once() { - let fetcher = FakeFetcher::new(""); - let reporter = RecordingJobReporter::new(); - let FetchedInputs { paths, .. } = fetch_inputs( - vec![ - input("systems", "github.com", "systems"), - input("systems_2", "github.com", "systems"), - ], + async fn locked_inputs_are_requested_up_front_and_failures_become_warnings() { + let reporter = RecordingJobReporter::new() + .with_input_outcome("/nix/store/a-source", fetched("a")) + .with_input_outcome( + "/nix/store/bad-source", + InputOutcome::Failed { + message: "404 ".into(), + }, + ); + let fetched = fetch_with( + &reporter, &FakeWorkerStore::default(), - &fetcher, - &*reporter.eval_progress_sink(), - None, + vec![input("a", "a"), input("b", "bad")], &mut no_abort(), ) .await .unwrap(); - assert_eq!(paths, vec!["/nix/store/systems-source".to_owned()]); - } - #[tokio::test] - async fn only_inputs_missing_from_the_store_are_pulled_from_the_cache_once() { - let mut reporter = RecordingJobReporter::new(); - let store = FakeWorkerStore::new().with_present_path("/nix/store/a-source"); - pull_from_cache( - &mut reporter, - &store, - &[ - input("a", "github.com", "a"), - input("b", "github.com", "b"), - input("b_2", "github.com", "b"), - ], - &mut no_abort(), - ) - .await - .unwrap(); - let pulled: Vec<_> = reporter - .events() + assert_eq!(fetched.paths, vec!["/nix/store/a-source".to_owned()]); + assert_eq!( + fetched.warnings, + vec!["skipping flake input 'b': 404".to_owned()] + ); + assert_eq!(reporter.input_requests().len(), 2); + let states: Vec<_> = rows(fetched.progress) .into_iter() - .filter_map(|e| match e { - ReportedEvent::PathsPulled { paths } => Some(paths), - _ => None, - }) + .map(|(name, state, ..)| (name, state)) .collect(); - assert_eq!(pulled, vec![vec!["/nix/store/b-source".to_owned()]]); + assert_eq!( + states, + vec![ + ("a".to_owned(), InputFetchState::Done), + ("b".to_owned(), InputFetchState::Failed) + ] + ); } #[tokio::test] - async fn present_inputs_send_no_progress() { - let fetcher = FakeFetcher::new(""); + async fn an_input_in_the_store_is_answered_without_the_server() { let reporter = RecordingJobReporter::new(); let store = FakeWorkerStore::new().with_present_path("/nix/store/a-source"); - let FetchedInputs { paths, .. } = fetch_inputs( - vec![input("a", "github.com", "a")], - &store, - &fetcher, - &*reporter.eval_progress_sink(), - None, - &mut no_abort(), - ) - .await - .unwrap(); - assert_eq!(paths, vec!["/nix/store/a-source".to_owned()]); - assert!(reporter.events().is_empty()); + let fetched = fetch_with(&reporter, &store, vec![input("a", "a")], &mut no_abort()) + .await + .unwrap(); + + assert_eq!(fetched.paths, vec!["/nix/store/a-source".to_owned()]); + assert!(reporter.events().is_empty(), "{:?}", reporter.events()); } #[tokio::test] - async fn a_present_input_is_a_done_row_next_to_a_missing_one() { - let fetcher = FakeFetcher::new(""); - let reporter = RecordingJobReporter::new(); + async fn server_progress_moves_the_input_board() { + let reporter = RecordingJobReporter::new() + .with_input_outcome("/nix/store/b-source", fetched("b")) + .with_input_progress(40, Some(90)); let store = FakeWorkerStore::new().with_present_path("/nix/store/a-source"); - let fetched = fetch_inputs( - vec![input("a", "github.com", "a"), input("b", "github.com", "b")], + let fetched = fetch_with( + &reporter, &store, - &fetcher, - &*reporter.eval_progress_sink(), - None, + vec![input("a", "a"), input("b", "b")], &mut no_abort(), ) .await .unwrap(); - let Some(EvalProgress::Fetching { inputs }) = fetched.progress else { - panic!("no final fetch snapshot"); - }; - let rows: Vec<_> = inputs - .iter() - .map(|i| (i.name.as_str(), i.state, i.downloaded_bytes)) - .collect(); + assert_eq!( - rows, + rows(fetched.progress), vec![ - ("a", InputFetchState::Done, 0), - ("b", InputFetchState::Done, 10), + ("a".to_owned(), InputFetchState::Done, 0, 0), + ("b".to_owned(), InputFetchState::Done, 40, 90), ] ); let Some(ReportedEvent::EvalProgress(sent)) = reporter.events().last().cloned() else { panic!("the final snapshot was not sent"); }; - assert_eq!(sent, EvalProgress::Fetching { inputs }); + assert!(matches!(sent, EvalProgress::Fetching { .. })); } #[tokio::test] - async fn no_rows_sends_no_progress() { - let fetcher = FakeFetcher::new(""); - let reporter = RecordingJobReporter::new(); - let FetchedInputs { paths, .. } = fetch_inputs( - vec![], + async fn inputs_sharing_a_source_are_requested_once() { + let reporter = RecordingJobReporter::new() + .with_input_outcome("/nix/store/systems-source", fetched("systems")); + let fetched = fetch_with( + &reporter, &FakeWorkerStore::default(), - &fetcher, - &*reporter.eval_progress_sink(), - None, + vec![input("systems", "systems"), input("systems_2", "systems")], &mut no_abort(), ) .await .unwrap(); - assert!(paths.is_empty()); - assert!(reporter.events().is_empty()); + + assert_eq!(fetched.paths, vec!["/nix/store/systems-source".to_owned()]); + assert_eq!(reporter.input_requests().len(), 1); } #[tokio::test] async fn an_abort_stops_the_fetch() { - let fetcher = FakeFetcher::new(""); let reporter = RecordingJobReporter::new(); let (tx, mut rx) = AbortSignal::channel(); tx.send(true).unwrap(); - let err = fetch_inputs( - vec![input("a", "github.com", "a")], + let err = fetch_with( + &reporter, &FakeWorkerStore::default(), - &fetcher, - &*reporter.eval_progress_sink(), - None, + vec![input("a", "a")], &mut rx, ) .await diff --git a/backend/gradient-worker/src/executor/input_fetch.rs b/backend/gradient-worker/src/executor/input_fetch.rs index e6a7c13b7..c0b310555 100644 --- a/backend/gradient-worker/src/executor/input_fetch.rs +++ b/backend/gradient-worker/src/executor/input_fetch.rs @@ -19,7 +19,6 @@ use tokio::sync::watch; use tokio::time::MissedTickBehavior; use tracing::{info, warn}; -use super::fetch::ssh_key_env; use crate::proto::job::JobUpdater; #[derive(Debug, Clone)] @@ -141,6 +140,30 @@ fn checked_store_path(request: &str, response: &str) -> Result { Ok(fetched.to_owned()) } +async fn ssh_key_env( + ssh_key: Option<&str>, + binpath_ssh: &str, +) -> Result> { + use std::os::unix::fs::PermissionsExt; + + let Some(key) = ssh_key else { + return Ok(None); + }; + let kf = tempfile::NamedTempFile::with_suffix(".key").context("create SSH key temp file")?; + tokio::fs::set_permissions(kf.path(), std::fs::Permissions::from_mode(0o600)) + .await + .context("chmod SSH key file")?; + tokio::fs::write(kf.path(), key.as_bytes()) + .await + .context("write SSH key file")?; + let ssh_command = format!( + "{} -i {} -o IdentitiesOnly=yes -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null", + binpath_ssh, + kf.path().display() + ); + Ok(Some((kf, ssh_command))) +} + type Totals = (u64, Option); struct Transfers { diff --git a/backend/gradient-worker/src/executor/mod.rs b/backend/gradient-worker/src/executor/mod.rs index 298ee919a..b19c26609 100644 --- a/backend/gradient-worker/src/executor/mod.rs +++ b/backend/gradient-worker/src/executor/mod.rs @@ -8,11 +8,11 @@ pub mod abort; pub mod build; mod build_metrics; pub mod compress; +mod dependencies; mod download; pub mod eval; pub(crate) mod failure; pub mod fetch; -mod import; pub(crate) mod input_fetch; pub mod log_limit; mod progress_report; @@ -397,8 +397,6 @@ impl JobExecutor { updater as &mut dyn gradient_wire::traits::JobReporter, credentials, &*self.store, - &**self.evaluator.resolver(), - &self.binpath_ssh, abort.clone(), ) .await diff --git a/backend/gradient-worker/src/main.rs b/backend/gradient-worker/src/main.rs index 2735a2186..1b4168863 100644 --- a/backend/gradient-worker/src/main.rs +++ b/backend/gradient-worker/src/main.rs @@ -59,11 +59,7 @@ fn main() -> Result<()> { }); if config.eval_subprocess { - return nix::eval_worker::run_eval_worker( - config.eval.metrics, - config.eval_fetch_through_parent, - ) - .map_err(anyhow::Error::from); + return nix::eval_worker::run_eval_worker(config.eval.metrics).map_err(anyhow::Error::from); } if let Some(requests_path) = config.eval_driver.clone() { diff --git a/backend/gradient-worker/src/worker_pool/input_fetch.rs b/backend/gradient-worker/src/worker_pool/input_fetch.rs deleted file mode 100644 index 607313fef..000000000 --- a/backend/gradient-worker/src/worker_pool/input_fetch.rs +++ /dev/null @@ -1,129 +0,0 @@ -/* - * SPDX-FileCopyrightText: 2026 Wavelens GmbH - * - * SPDX-License-Identifier: AGPL-3.0-only - */ - -use std::collections::HashMap; -use std::sync::Arc; - -use async_trait::async_trait; -use gradient_util::sync::Mutex; -use gradient_wire::types::{InputFetch, InputFetchState}; - -#[derive(Debug)] -struct InputRow { - name: String, - state: InputFetchState, - transfers: HashMap, -} - -#[derive(Debug)] -pub(crate) struct InputBoard { - rows: Mutex>, -} - -impl InputBoard { - pub(crate) fn new(rows: Vec<(String, InputFetchState)>) -> Arc { - let rows = rows - .into_iter() - .map(|(name, state)| InputRow { - name, - state, - transfers: HashMap::new(), - }) - .collect(); - Arc::new(Self { - rows: Mutex::new(rows), - }) - } - - pub(crate) fn set_state(&self, index: usize, state: InputFetchState) { - self.rows.lock()[index].state = state; - } - - pub(crate) fn transfer(&self, index: usize, id: u64, done: u64, expected: u64) { - self.rows.lock()[index] - .transfers - .insert(id, (done, expected)); - } - - pub(crate) fn snapshot(&self) -> Vec { - self.rows - .lock() - .iter() - .map(|row| { - let sizes = row.transfers.values(); - let known = sizes.clone().all(|(_, expected)| *expected > 0); - InputFetch { - name: row.name.clone(), - state: row.state, - downloaded_bytes: sizes.clone().map(|(done, _)| done).sum(), - expected_bytes: if known { - sizes.map(|(_, expected)| expected).sum() - } else { - 0 - }, - } - }) - .collect() - } -} - -#[derive(Debug, Clone)] -pub(crate) struct DownloadTarget { - pub board: Arc, - pub index: usize, -} - -impl gradient_eval::pool::DownloadSink for DownloadTarget { - fn transfer(&self, id: u64, done: u64, expected: u64) { - self.board.transfer(self.index, id, done, expected); - } -} - -#[async_trait] -pub(crate) trait InputFetcher: Send + Sync { - async fn fetch_input( - &self, - locked: String, - git_ssh_command: Option, - target: DownloadTarget, - ) -> anyhow::Result; -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn a_row_sums_its_transfers_and_drops_an_unknown_size() { - let board = InputBoard::new(vec![("nixpkgs".into(), InputFetchState::Queued)]); - board.set_state(0, InputFetchState::Fetching); - board.transfer(0, 1, 100, 400); - board.transfer(0, 2, 50, 100); - assert_eq!(board.snapshot()[0].downloaded_bytes, 150); - assert_eq!(board.snapshot()[0].expected_bytes, 500); - board.transfer(0, 3, 10, 0); - assert_eq!(board.snapshot()[0].expected_bytes, 0); - assert_eq!(board.snapshot()[0].state, InputFetchState::Fetching); - } - - #[test] - fn the_sink_draws_transfers_on_its_own_row() { - let board = InputBoard::new(vec![ - ("nixpkgs".into(), InputFetchState::Fetching), - ("utils".into(), InputFetchState::Fetching), - ]); - let second = DownloadTarget { - board: Arc::clone(&board), - index: 1, - }; - - gradient_eval::pool::DownloadSink::transfer(&second, 7, 40, 90); - - assert_eq!(board.snapshot()[0].downloaded_bytes, 0); - assert_eq!(board.snapshot()[1].downloaded_bytes, 40); - assert_eq!(board.snapshot()[1].expected_bytes, 90); - } -} diff --git a/backend/gradient-worker/src/worker_pool/mod.rs b/backend/gradient-worker/src/worker_pool/mod.rs index f077dc163..48fa0b640 100644 --- a/backend/gradient-worker/src/worker_pool/mod.rs +++ b/backend/gradient-worker/src/worker_pool/mod.rs @@ -5,8 +5,6 @@ */ pub mod eval_stats; -mod input_fetch; mod resolver; -pub(crate) use self::input_fetch::{DownloadTarget, InputBoard, InputFetcher}; pub use self::resolver::WorkerPoolResolver; diff --git a/backend/gradient-worker/src/worker_pool/resolver.rs b/backend/gradient-worker/src/worker_pool/resolver.rs index 78023e01a..9bc5d729a 100644 --- a/backend/gradient-worker/src/worker_pool/resolver.rs +++ b/backend/gradient-worker/src/worker_pool/resolver.rs @@ -8,9 +8,7 @@ use anyhow::{Context, Result}; use async_trait::async_trait; use gradient_derivation::{Derivation, parse_drv}; use gradient_eval::ipc::ResolvedItem; -use gradient_eval::pool::{ - CallStats, EvalPool, EvalPoolConfig, EvalWorker, LiveThunks, MemoryGuard, -}; +use gradient_eval::pool::{CallStats, EvalPool, EvalPoolConfig, LiveThunks, MemoryGuard}; use gradient_sources::{ AttrError, DerivationResolver, EvalDependencies, FlakeDiscovery, FoundDerivations, ResolvedDerivation, @@ -20,7 +18,6 @@ use gradient_util::sync::Mutex; use std::sync::Arc; use super::eval_stats::{EvalStatsAccumulator, EvalStatsTotals}; -use super::input_fetch::{DownloadTarget, InputFetcher}; #[derive(Debug)] pub struct WorkerPoolResolver { @@ -75,7 +72,7 @@ impl WorkerPoolResolver { cache_dir: eval_cache_dir, trace_dir, metrics: eval_metrics, - fetch_through_parent: false, + fetch_through_parent: true, }; Self { pool: EvalPool::new(config.clone()), @@ -132,31 +129,6 @@ impl WorkerPoolResolver { } } -#[async_trait] -impl InputFetcher for WorkerPoolResolver { - async fn fetch_input( - &self, - locked: String, - git_ssh_command: Option, - target: DownloadTarget, - ) -> Result { - let args: serde_json::Value = - serde_json::from_str(&locked).context("parse locked input attributes")?; - let request = serde_json::json!({ "builtin": "fetchTree", "args": args }).to_string(); - let mut worker = EvalWorker::spawn_for_fetch(&self.config).await?; - let fetched = worker - .fetch(request, git_ssh_command, Arc::new(target)) - .await; - worker.shutdown().await; - let response: serde_json::Value = - serde_json::from_str(&fetched?).context("parse fetch response")?; - response["storePath"] - .as_str() - .map(str::to_owned) - .context("fetch response without a store path") - } -} - #[async_trait] impl DerivationResolver for WorkerPoolResolver { async fn list_flake_derivations( diff --git a/cli/src/commands/base.rs b/cli/src/commands/base.rs index a2621ff4f..15998a7f5 100644 --- a/cli/src/commands/base.rs +++ b/cli/src/commands/base.rs @@ -172,7 +172,7 @@ pub fn run() -> std::io::Result<()> { #[cfg(feature = "eval")] if let Some(metrics) = eval::subprocess_flags(&std::env::args().collect::>()) { - return gradient_eval::eval_worker::run_eval_worker(metrics, false); + return gradient_eval::eval_worker::run_eval_worker(metrics); } let cli = Cli::parse(); From 28f7f877c47f79d4c9895b357f633ac07345becf Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 16:02:05 +0200 Subject: [PATCH 13/19] scheduler: feat: input fetches take a build slot --- backend/gradient-scheduler/src/actor.rs | 34 +++++ .../gradient-scheduler/src/cases/inputs.rs | 126 ++++++++++++++++-- backend/gradient-scheduler/src/input_waits.rs | 4 + .../src/job_handlers/input.rs | 89 +++++++++---- .../src/worker/message_loop.rs | 86 +++++++++++- 5 files changed, 303 insertions(+), 36 deletions(-) diff --git a/backend/gradient-scheduler/src/actor.rs b/backend/gradient-scheduler/src/actor.rs index 4108233a9..e48f1f21a 100644 --- a/backend/gradient-scheduler/src/actor.rs +++ b/backend/gradient-scheduler/src/actor.rs @@ -166,6 +166,17 @@ pub enum SchedulerMsg { signals: Vec<(String, SessionSignal)>, reply: RpcReplyPort<()>, }, + ClaimInputFetch { + worker: String, + fetch_id: String, + signal: SessionSignal, + reply: RpcReplyPort, + }, + ReleaseInputFetch { + worker: String, + fetch_id: String, + reply: RpcReplyPort<()>, + }, DropCluster { seats: Vec<(String, String)>, reply: RpcReplyPort<()>, @@ -939,6 +950,29 @@ impl Actor for CoreActor { } let _ = reply.send(()); } + SchedulerMsg::ClaimInputFetch { + worker, + fetch_id, + signal, + reply, + } => { + let claimed = core.pool.has_capacity(&worker, &JobKind::Build); + if claimed { + let key = crate::input_waits::input_fetch_key(&fetch_id); + core.pool.assign_job(&worker, &key, JobKind::Build, 0); + core.pool.signal(&worker, signal); + } + let _ = reply.send(claimed); + } + SchedulerMsg::ReleaseInputFetch { + worker, + fetch_id, + reply, + } => { + core.pool + .release_job(&worker, &crate::input_waits::input_fetch_key(&fetch_id)); + let _ = reply.send(()); + } SchedulerMsg::Candidates { worker, only_new, diff --git a/backend/gradient-scheduler/src/cases/inputs.rs b/backend/gradient-scheduler/src/cases/inputs.rs index 407411631..829a75a44 100644 --- a/backend/gradient-scheduler/src/cases/inputs.rs +++ b/backend/gradient-scheduler/src/cases/inputs.rs @@ -12,8 +12,8 @@ use gradient_types::waiting_reason::{EvalCapability, WaitingReason}; use gradient_wire::types::{InputOutcome, JobKind}; use super::{ - A, A_OUT, CaseShapes, SchedulerCase, eval_key, w_all, w_network_eval, w_network_secrets, - without_offers, + A, A_OUT, CaseShapes, Profile, SchedulerCase, eval_key, output_nar, w_all, w_network_eval, + w_network_secrets, without_offers, }; const W1: &str = "w1"; @@ -38,7 +38,7 @@ async fn i01_an_input_request_goes_to_a_network_and_secrets_worker() { let case = SchedulerCase::start(shape().evalq(e1)).await; evaluating(&case, W1, e1).await; case.connect(W2, super::w_network(), &[]).await; - case.connect(W3, w_network_secrets(), &[]).await; + case.connect(W3, w_fetch(), &[]).await; case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) .await; @@ -65,7 +65,7 @@ async fn i02_two_evaluations_share_one_fetch() { let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; evaluating(&case, W1, e1).await; evaluating(&case, W2, e2).await; - case.connect(W3, w_network_secrets(), &[]).await; + case.connect(W3, w_fetch(), &[]).await; case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) .await; @@ -100,7 +100,7 @@ async fn i03_a_failed_fetch_answers_every_waiter() { let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; evaluating(&case, W1, e1).await; evaluating(&case, W2, e2).await; - case.connect(W3, w_network_secrets(), &[]).await; + case.connect(W3, w_fetch(), &[]).await; case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) .await; case.input_request(W2, &eval_key(e2), "r2", SSH_REQUEST) @@ -133,8 +133,8 @@ async fn i04_a_lost_fetch_worker_hands_the_fetch_to_another() { let e1 = EvaluationId::now_v7(); let case = SchedulerCase::start(shape().evalq(e1)).await; evaluating(&case, W1, e1).await; - case.connect(W2, w_network_secrets(), &[]).await; - case.connect(W3, w_network_secrets(), &[]).await; + case.connect(W2, w_fetch(), &[]).await; + case.connect(W3, w_fetch(), &[]).await; case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) .await; let first = single_fetch(&case, &[W1, W2, W3]); @@ -188,7 +188,7 @@ async fn i05_an_abort_drops_the_input_wait() { assert_eq!(case.input_waits(), 0, "after 2: the wait is gone"); case.check("step 2").await; - case.connect(W2, w_network_secrets(), &[]).await; + case.connect(W2, w_fetch(), &[]).await; case.build_assign_pass(false, false).await; assert_eq!( fetches_of(&case, &[W1, W2]), @@ -230,7 +230,7 @@ async fn i06_without_a_fetch_worker_the_evaluation_waits_for_network() { ); case.check("step 1").await; - case.connect(W2, w_network_secrets(), &[]).await; + case.connect(W2, w_fetch(), &[]).await; case.build_assign_pass(false, false).await; let fetch = single_fetch(&case, &[W1, W2]); assert_eq!(fetch.worker, W2, "after 2: the new worker gets the fetch"); @@ -250,7 +250,7 @@ async fn i07_an_input_request_from_a_build_job_is_dropped() { let e1 = EvaluationId::now_v7(); let case = SchedulerCase::start(shape().one(e1, A, A_OUT)).await; case.connect(W1, w_all(), &[]).await; - case.connect(W2, w_network_secrets(), &[]).await; + case.connect(W2, w_fetch(), &[]).await; case.admit_startable_moves().await; case.offers_full(W1).await; case.score_all(W1).await; @@ -276,8 +276,8 @@ async fn i08_fetch_progress_reaches_every_waiting_evaluation() { let case = SchedulerCase::start(shape().evalq(e1).evalq(e2)).await; evaluating(&case, W1, e1).await; evaluating(&case, W2, e2).await; - case.connect(W3, w_network_secrets(), &[]).await; - case.connect(W4, w_network_secrets(), &[]).await; + case.connect(W3, w_fetch(), &[]).await; + case.connect(W4, w_fetch(), &[]).await; case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) .await; case.input_request(W2, &eval_key(e2), "r2", SSH_REQUEST) @@ -323,6 +323,108 @@ async fn i08_fetch_progress_reaches_every_waiting_evaluation() { case.finish().await; } +fn w_fetch() -> Profile { + w_network_secrets().slots(1) +} + +#[tokio::test] +async fn i09_an_input_fetch_waits_for_a_free_build_slot() { + let (e1, e2) = (EvaluationId::now_v7(), EvaluationId::now_v7()); + let case = SchedulerCase::start(shape().evalq(e1).one(e2, A, A_OUT)).await; + evaluating(&case, W1, e1).await; + case.connect(W2, w_all(), &[]).await; + let (a, a2) = building_on_w2(&case).await; + + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + assert_eq!( + fetch_signals(&case, W2), + vec![], + "after 2: W2 has no free build slot" + ); + assert_eq!(case.input_waits(), 1, "after 2: the request waits"); + case.check("step 2").await; + + case.commit_nar(output_nar(A, A_OUT)).await; + case.complete(W2, &a, a2, Vec::new(), 5).await; + case.build_assign_pass(false, false).await; + assert_eq!( + fetch_signals(&case, W2).len(), + 1, + "after 3: the freed slot takes the fetch" + ); + case.signals(W1); + case.check("step 3").await; + + case.finish().await; +} + +#[tokio::test] +async fn i10_a_running_fetch_takes_a_build_slot() { + let (e1, e2) = (EvaluationId::now_v7(), EvaluationId::now_v7()); + let case = SchedulerCase::start(shape().evalq(e1).one(e2, A, A_OUT)).await; + evaluating(&case, W1, e1).await; + case.connect(W2, w_all(), &[]).await; + + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + let fetches = fetch_signals(&case, W2); + let [SessionSignal::FetchInput { fetch_id, .. }] = fetches.as_slice() else { + panic!("after 1: W2 gets one FetchInput, got {fetches:?}"); + }; + case.check("step 1").await; + + case.admit_startable_moves().await; + case.offers_full(W2).await; + case.score_all(W2).await; + let taken = case.request_job(W2, JobKind::Build).await; + assert!( + taken.is_none(), + "after 2: the fetch holds the only slot, got {:?}", + taken.map(|t| t.job_id().to_owned()) + ); + case.check("step 2").await; + + case.input_fetched(W2, fetch_id, fetched()).await; + assert_eq!( + without_offers(case.signals(W1)), + vec![result(e1, "r1", fetched())], + "after 3: signals(W1)" + ); + let build = case + .request_job(W2, JobKind::Build) + .await + .expect("after 3: the released slot takes build:A"); + assert_eq!(build.job_id(), case.key(A).await, "after 3: build:A"); + case.signals(W2); + case.check("step 3").await; + + case.finish().await; +} + +async fn building_on_w2(case: &SchedulerCase) -> (String, gradient_entity::ids::DispatchedJobId) { + let a = case.key(A).await; + case.admit_startable_moves().await; + case.offers_full(W2).await; + case.score_all(W2).await; + let assignment = case + .request_job(W2, JobKind::Build) + .await + .expect("after 1: W2 gets build:A"); + assert_eq!(assignment.job_id(), a, "after 1: build:A"); + case.report_building(W2, &a).await; + case.signals(W2); + case.check("step 1").await; + (a, assignment.assignment_id()) +} + +fn fetch_signals(case: &SchedulerCase, worker: &str) -> Vec { + case.signals(worker) + .into_iter() + .filter(|s| matches!(s, SessionSignal::FetchInput { .. })) + .collect() +} + async fn evaluating(case: &SchedulerCase, worker: &str, evaluation: EvaluationId) { let key = eval_key(evaluation); case.connect(worker, w_network_eval(), &[]).await; diff --git a/backend/gradient-scheduler/src/input_waits.rs b/backend/gradient-scheduler/src/input_waits.rs index 4aa6af49c..f6a27b562 100644 --- a/backend/gradient-scheduler/src/input_waits.rs +++ b/backend/gradient-scheduler/src/input_waits.rs @@ -167,6 +167,10 @@ impl InputWaits { } } +pub(crate) fn input_fetch_key(fetch_id: &str) -> String { + format!("input:{fetch_id}") +} + fn key_of_fetch( waits: &HashMap, worker: &str, diff --git a/backend/gradient-scheduler/src/job_handlers/input.rs b/backend/gradient-scheduler/src/job_handlers/input.rs index fad45b928..528512eed 100644 --- a/backend/gradient-scheduler/src/job_handlers/input.rs +++ b/backend/gradient-scheduler/src/job_handlers/input.rs @@ -13,9 +13,10 @@ use gradient_pool::{InputRequest, WorkerCaps, WorkerInfo}; use gradient_types::*; use gradient_wire::types::InputOutcome; use sea_orm::{ColumnTrait, EntityTrait, QueryFilter}; -use tracing::debug; +use tracing::{debug, warn}; use crate::Scheduler; +use crate::actor::SchedulerMsg; use crate::input_waits::{AssignedFetch, InputWaiter, NewInputWait, UnassignedFetch}; use crate::jobs::PendingJob; @@ -63,6 +64,7 @@ impl Scheduler { fetch_id: &str, outcome: InputOutcome, ) { + self.release_input_fetch(worker_id, fetch_id).await; let waiters = self.input_waits.take_fetched(worker_id, fetch_id); if waiters.is_empty() { debug!(%worker_id, %fetch_id, "input fetch answered without waiters; dropped"); @@ -150,34 +152,73 @@ impl Scheduler { return; } - let workers = self.board_workers().await; + let mut workers = self.board_workers().await; let mut load = self.input_waits.assigned_per_worker(); - let mut signals = Vec::new(); for fetch in unassigned { - let Some(worker) = fetch_worker(&workers, &load, &fetch) else { + if !self + .place_input_fetch(&fetch, &mut workers, &mut load) + .await + { debug!(fetch_key = %fetch.fetch_key, "no connected worker can fetch the input yet"); - continue; - }; + } + } + } + + async fn place_input_fetch( + &self, + fetch: &UnassignedFetch, + workers: &mut [WorkerInfo], + load: &mut HashMap, + ) -> bool { + for worker in fetch_workers(workers, load, fetch) { let fetch_id = uuid::Uuid::now_v7().to_string(); + let signal = SessionSignal::FetchInput { + fetch_id: fetch_id.clone(), + project: fetch.project, + request: fetch.request.clone(), + }; + let claimed = self + .call(|reply| SchedulerMsg::ClaimInputFetch { + worker: worker.clone(), + fetch_id: fetch_id.clone(), + signal, + reply, + }) + .await + .unwrap_or(false); + if !claimed { + continue; + } + let assigned = AssignedFetch { fetch_id: fetch_id.clone(), worker: worker.clone(), }; - if self.input_waits.assign(&fetch.fetch_key, assigned) { - *load.entry(worker.clone()).or_default() += 1; - signals.push(( - worker, - SessionSignal::FetchInput { - fetch_id, - project: fetch.project, - request: fetch.request, - }, - )); + if !self.input_waits.assign(&fetch.fetch_key, assigned) { + self.release_input_fetch(&worker, &fetch_id).await; + return false; } + *load.entry(worker.clone()).or_default() += 1; + if let Some(info) = workers.iter_mut().find(|w| w.id == worker) { + info.assigned_build_count += 1; + } + return true; } - if !signals.is_empty() { - self.signal_workers(signals).await; + false + } + + async fn release_input_fetch(&self, worker: &str, fetch_id: &str) { + let released = self + .call(|reply| SchedulerMsg::ReleaseInputFetch { + worker: worker.to_owned(), + fetch_id: fetch_id.to_owned(), + reply, + }) + .await; + if let Err(e) = released { + warn!(error = %e, %worker, %fetch_id, "the build slot of an input fetch was not released"); } + self.build_kick.kick(); } async fn answer_inputs(&self, waiters: Vec, outcome: InputOutcome) { @@ -198,14 +239,15 @@ impl Scheduler { } } -fn fetch_worker( +fn fetch_workers( workers: &[WorkerInfo], load: &HashMap, fetch: &UnassignedFetch, -) -> Option { - workers +) -> Vec { + let mut eligible: Vec<&WorkerInfo> = workers .iter() .filter(|w| !w.draining) + .filter(|w| w.assigned_build_count < w.max_concurrent_builds as usize) .filter(|w| { w.authorized_peers .as_ref() @@ -218,6 +260,7 @@ fn fetch_worker( } .can_fetch_input(fetch.needs_secrets) }) - .min_by_key(|w| (load.get(&w.id).copied().unwrap_or_default(), w.id.clone())) - .map(|w| w.id.clone()) + .collect(); + eligible.sort_by_key(|w| (load.get(&w.id).copied().unwrap_or_default(), w.id.clone())); + eligible.into_iter().map(|w| w.id.clone()).collect() } diff --git a/backend/gradient-worker/src/worker/message_loop.rs b/backend/gradient-worker/src/worker/message_loop.rs index dcc4f2b75..76ad051b5 100644 --- a/backend/gradient-worker/src/worker/message_loop.rs +++ b/backend/gradient-worker/src/worker/message_loop.rs @@ -254,6 +254,7 @@ pub(super) struct MessageLoopState { executor: JobExecutor, downloader: Arc, fetches: tokio::task::JoinSet<()>, + running_fetches: Arc, config: WorkerConfig, draining: bool, stopping: bool, @@ -308,6 +309,7 @@ impl MessageLoopState { executor.evaluator.resolver().config().clone(), )), fetches: tokio::task::JoinSet::new(), + running_fetches: Arc::default(), executor, config, draining: false, @@ -344,7 +346,13 @@ impl MessageLoopState { } fn occupied(&self, kind: JobKind) -> u32 { - self.jobs.active(kind.clone()) + self.holds.held(&kind) + let fetches = match kind { + JobKind::Build => self + .running_fetches + .load(std::sync::atomic::Ordering::Acquire), + JobKind::Flake => 0, + }; + self.jobs.active(kind.clone()) + self.holds.held(&kind) + fetches } fn max_for(&self, kind: JobKind) -> u32 { @@ -858,7 +866,9 @@ impl MessageLoopState { let downloader = Arc::clone(&self.downloader); let writer = self.writer.clone(); let binpath_ssh = self.executor.binpath_ssh.clone(); + let running = RunningFetch::start(Arc::clone(&self.running_fetches)); self.fetches.spawn(async move { + let _running = running; let fetch_id = fetch.fetch_id.clone(); let report = UpdaterFetchReport(updater); let outcome = fetch_input(fetch, &*downloader, &report, &binpath_ssh).await; @@ -1034,6 +1044,21 @@ impl MessageLoopState { } } +struct RunningFetch(Arc); + +impl RunningFetch { + fn start(count: Arc) -> Self { + count.fetch_add(1, std::sync::atomic::Ordering::AcqRel); + Self(count) + } +} + +impl Drop for RunningFetch { + fn drop(&mut self) { + self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel); + } +} + fn send_live_metrics(writer: &ProtoWriter, daemon_cgroup: Option) { let writer = writer.clone(); #[expect( @@ -1255,6 +1280,20 @@ mod tests { assert!(idle.is_aborted(), "a job between checkpoints sees it too"); } + struct StalledDownloader; + + #[async_trait::async_trait] + impl InputDownloader for StalledDownloader { + async fn download( + &self, + _request: String, + _git_ssh_command: Option, + _sink: Arc, + ) -> Result { + std::future::pending().await + } + } + struct OfflineDownloader; #[async_trait::async_trait] @@ -1361,4 +1400,49 @@ mod tests { } ); } + + #[tokio::test] + async fn a_running_input_fetch_occupies_a_build_slot() { + let server = gradient_test_support::prelude::MockProtoServer::bind().await; + let url = server.url().to_owned(); + let conn = gradient_worker_client::connection::ProtoConnection::open(&url) + .await + .unwrap(); + let (writer, _reader, _flush) = conn.split(); + let gcroots = tempfile::tempdir().unwrap(); + let mut state = MessageLoopState::new( + writer, + WorkerConfig::default(), + executor(gcroots.path()), + JobScorer::new(), + CredentialStore::new(), + ); + state.downloader = Arc::new(StalledDownloader); + let _running = state.jobs.register( + "build:b1".to_owned(), + JobKind::Build, + "dispatch-1".to_owned(), + None, + ); + + state + .route(ServerMessage::FetchInput { + fetch_id: "f1".into(), + project: "p1".into(), + request: "{}".into(), + ssh_key: None, + }) + .await + .unwrap(); + + assert_eq!(state.occupied(JobKind::Build), 2); + assert_eq!(state.occupied(JobKind::Flake), 0); + state.fetches.abort_all(); + while state.fetches.join_next().await.is_some() {} + assert_eq!( + state.occupied(JobKind::Build), + 1, + "an ended fetch frees its slot" + ); + } } From 586f4f1042695ee64489f60c5d917243d3471ca6 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 16:02:19 +0200 Subject: [PATCH 14/19] e2e: test: input fetched by another worker --- docs/src/concepts/evaluations-and-builds.md | 2 +- docs/src/concepts/workers.md | 3 +- docs/src/contributors/eval-worker.md | 29 ++++++++-- .../gradient/e2e/flake_repository_input.lock | 23 ++++++++ .../gradient/e2e/flake_repository_input.nix | 20 +++++++ nix/tests/gradient/e2e/mk.nix | 57 +++++++++++++++++++ 6 files changed, 127 insertions(+), 7 deletions(-) create mode 100644 nix/tests/gradient/e2e/flake_repository_input.lock create mode 100644 nix/tests/gradient/e2e/flake_repository_input.nix diff --git a/docs/src/concepts/evaluations-and-builds.md b/docs/src/concepts/evaluations-and-builds.md index f5044d563..339240d20 100644 --- a/docs/src/concepts/evaluations-and-builds.md +++ b/docs/src/concepts/evaluations-and-builds.md @@ -15,7 +15,7 @@ flowchart LR | Status | Meaning | |---|---| | Queued | The evaluation awaits an evaluating worker | -| Fetching | A worker clones the repository and its flake inputs | +| Fetching | A worker clones the repository and gathers its flake inputs from the store, the cache or a download by a worker with `network` | | Evaluating | A worker walks the flake and reports derivations in batches | | Building | All derivations are known. Builds are running | | Waiting | No connected worker can make progress, e.g. no worker with the needed system or features. The evaluation resumes on its own. Builds needing a system or features absent from every connected worker abort the evaluation after 5 minutes with a warning. A task can opt into waiting for workers instead | diff --git a/docs/src/concepts/workers.md b/docs/src/concepts/workers.md index a3bd7a8e4..1382f240d 100644 --- a/docs/src/concepts/workers.md +++ b/docs/src/concepts/workers.md @@ -13,7 +13,7 @@ Workers declare their capabilities at the handshake. Capabilities cover the jobs | Capability | Default | Meaning | |---|---|---| -| `network` | on | Internet access. Needed to clone the repository, fetch flake inputs, substitute from upstream caches and run fixed-output builds | +| `network` | on | Internet access. Needed to clone the repository, download missing inputs, substitute from upstream caches and run fixed-output builds | | `secrets` | on | The server may send the project SSH key. Needed to clone a repository behind an SSH URL | | `exec` | off | Declared and stored. No job checks it yet | | `eval` | on | Walk the flake and report derivations | @@ -21,6 +21,7 @@ Workers declare their capabilities at the handshake. Capabilities cover the jobs - Public repositories need a worker with `network` for the fetch step. - Workers without `secrets` never get an SSH key. +- An evaluating worker never downloads an input itself. The server can hand each missing input to a worker with `network`, plus `secrets` for an SSH URL. The evaluation can then pull the input from the cache. - A project or team registration can withhold each capability through `enable_network`, `enable_secrets`, `enable_exec`, `enable_eval` and `enable_build`. A single machine can hold every capability. A machine with a lot of memory can evaluate while smaller machines build. diff --git a/docs/src/contributors/eval-worker.md b/docs/src/contributors/eval-worker.md index d345c12bb..21ce01e98 100644 --- a/docs/src/contributors/eval-worker.md +++ b/docs/src/contributors/eval-worker.md @@ -14,7 +14,8 @@ flowchart LR ## Subprocess IPC - **Frames:** each frame start with a `u32` little-endian length prefix, followed by an rkyv payload (`gradient-eval/src/ipc.rs`). -- **Version byte:** before the first frame, subprocesses write `EVAL_IPC_VERSION` (currently 10). A binary swapped mid-evaluation then fail the handshake instead of sending undecodable frames. +- **Version byte:** before the first frame, subprocesses write `EVAL_IPC_VERSION` (currently 12). A binary swapped mid-evaluation then fail the handshake instead of sending undecodable frames. +- **Fetch byte:** parents reply with one byte. Subprocesses given `1` ask the parent for [missing inputs](#missing-inputs). Subprocesses given `0` let Nix download them. - **Responses:** each request pair with one response. `Stats` ticks and [build requests](#import-from-derivation) can arrive before that response. `List` responses carry the derivation path of every attribute found, resolved in the subprocess that forced the attribute. - **Shards:** `Plan` can split each include at its first wildcard. - With more segments after the wildcard, each child receive a sub-pattern (`packages.**.hello` -> `packages.x86_64-linux.hello`). @@ -41,7 +42,7 @@ An import from a derivation (IFD) requires a build in the middle of an evaluatio - A different request during the wait counts as a protocol error. Subprocesses then stop reading and exit. - Nix can raise the `error` of `BuildDone` as the error of the import. Later attributes with the same import receive the same error. -### Import Builder (`gradient-worker/src/executor/import.rs`) +### Import Builder (`gradient-worker/src/executor/dependencies.rs`) 1. Record the closure, with the imported derivation as the entry point `other..` and `ifd` set. Repeated names within a job end with `-`. 2. Send [`ImportRequest`](proto/messages.md) and wait for `ImportResult`. @@ -51,6 +52,25 @@ An import from a derivation (IFD) requires a build in the middle of an evaluatio - Requests for the same `.drv` from different subprocesses share the wait and the answer. - Fingerprint, checkpoint, input fetch and the `--eval-driver` harness refuse imports with `import from derivation is not available during `. +## Missing Inputs + +Evaluation subprocesses of the worker never download an input. The patched Nix of Gradient hands every missing non-`path` input to a fetch hook instead. + +| Frame | Direction | Content | +|---|---|---| +| `NeedsInput` | subprocess -> parent | `request`: JSON with `builtin` (`fetchTree`, `fetchTarball` or `fetchurl`), `args` and the expected `storePath` where known | +| `InputDone` | parent -> subprocess | `response`: JSON with `storePath` and, for `fetchTree`, the locked `attrs`. Or `error` | + +Workers answer from the first source holding the input. + +1. The Gradient cache, for a request with a known `storePath`. +2. An [`InputRequest`](proto/messages.md) to the server, followed by a pull of the fetched path. + +- The server can hand the download to a worker with `network`, plus `secrets` for an SSH URL. That worker can start one subprocess with fetch byte `0`, outside the eval pool. Downloads hold one build slot of that worker until `InputFetched`. +- Requests of one job for the same input share the wait and the answer. Failed inputs keep their error for the rest of the evaluation. +- The fetch step can request every locked input of `flake.lock` up front. Failures there remain warnings. The evaluation can retry such an input through the hook. +- `gradient eval` and the `--eval-driver` harness use fetch byte `0`. + ## Parent Side The pool code sits in `gradient-eval/src/pool/`. Both `gradient eval` and the worker drive it. @@ -61,9 +81,9 @@ The pool code sits in `gradient-eval/src/pool/`. Both `gradient eval` and the wo | `workers.rs` | Checkout and return with test-on-borrow | | `memory.rs` | The memory guard, its victim choice and its reaper | | `fan_out.rs` | Shard batching and crash isolation | -| `mod.rs` | The `EvalPool` facade: plan, list, fingerprint, checkpoint, input fetch | +| `mod.rs` | The `EvalPool` facade: plan, list, fingerprint, checkpoint | | `live_thunks.rs` | Live thunk counters | -| `import.rs` | The import-from-derivation seam, with `RefuseImports` as the refusal | +| `dependencies.rs` | The seam for imports and missing inputs, with `RefuseDependencies` as the refusal | | `download.rs` | The `@nix` log parser and the download progress seam | | `driver.rs` | The hidden `--eval-driver ` harness: JSONL requests through the real transport, JSON responses out. The NixOS VM test can drive both sides of the wire through the harness from Python | @@ -72,7 +92,6 @@ The worker adapter sits in `gradient-worker/src/worker_pool/`. | File | Role | |---|---| | `resolver.rs` | `WorkerPoolResolver`, the `gradient_sources::DerivationResolver` over `EvalPool` | -| `input_fetch.rs` | The download board, as a `DownloadSink` on `InputBoard` | | `eval_stats.rs` | Per-attribute evaluation statistics | ## Memory Safety diff --git a/nix/tests/gradient/e2e/flake_repository_input.lock b/nix/tests/gradient/e2e/flake_repository_input.lock new file mode 100644 index 000000000..155c86135 --- /dev/null +++ b/nix/tests/gradient/e2e/flake_repository_input.lock @@ -0,0 +1,23 @@ +{ + "nodes": { + "extra": { + "flake": false, + "locked": { + "narHash": "[hash]", + "type": "tarball", + "url": "http://server/upstream/input.tar.gz" + }, + "original": { + "type": "tarball", + "url": "http://server/upstream/input.tar.gz" + } + }, + "root": { + "inputs": { + "extra": "extra" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/nix/tests/gradient/e2e/flake_repository_input.nix b/nix/tests/gradient/e2e/flake_repository_input.nix new file mode 100644 index 000000000..b293d0bf2 --- /dev/null +++ b/nix/tests/gradient/e2e/flake_repository_input.nix @@ -0,0 +1,20 @@ +/* + * SPDX-FileCopyrightText: 2026 Wavelens GmbH + * + * SPDX-License-Identifier: AGPL-3.0-only + */ +{ + description = "Test Repository for Gradient Build Server"; + inputs.extra = { + url = "http://server/upstream/input.tar.gz"; + flake = false; + }; + outputs = { self, extra, ... }: { + packages.x86_64-linux.default = derivation { + name = "input-${builtins.readFile "${extra}/value"}"; + system = "x86_64-linux"; + builder = "/bin/sh"; + args = [ "-c" "echo ok > $out" ]; + }; + }; +} diff --git a/nix/tests/gradient/e2e/mk.nix b/nix/tests/gradient/e2e/mk.nix index 6a273a3b6..793f2fc22 100644 --- a/nix/tests/gradient/e2e/mk.nix +++ b/nix/tests/gradient/e2e/mk.nix @@ -277,6 +277,8 @@ pkgs.testers.runNixOSTest ({ pkgs, lib, ... }: { "L+ /var/lib/git/flake-busywrap.nix 0755 git git - ${./flake_repository_busywrap.nix}" "L+ /var/lib/git/flake-spin.nix 0755 git git - ${./flake_repository_spin.nix}" "L+ /var/lib/git/flake-slow.nix 0755 git git - ${./flake_repository_slow.nix}" + "L+ /var/lib/git/flake-input.nix 0755 git git - ${./flake_repository_input.nix}" + "L+ /var/lib/git/flake-input.lock 0755 git git - ${./flake_repository_input.lock}" ]; }; @@ -2442,6 +2444,61 @@ pkgs.testers.runNixOSTest ({ pkgs, lib, ... }: { f"the unguarded arm did not drift, so the keys are not what the others prove:\n{out}" ) + if requires("Phase 10n", "distinct-upstream-workers"): + banner("Phase 10n: an input in no cache arrives through the worker with network") + + def restart_with_caps(node, caps): + since = node.succeed("date +%s").strip() + dropin = "/run/systemd/system/gradient-worker.service.d" + if caps: + lines = "".join(f"Environment=GRADIENT_WORKER_CAPABILITIES_{k}={v}\n" for k, v in caps.items()) + node.succeed(f"mkdir -p {dropin} && cat > {dropin}/caps.conf <<'CAPS'\n[Service]\n{lines}CAPS") + else: + node.succeed(f"rm -f {dropin}/caps.conf") + node.succeed("systemctl daemon-reload && systemctl restart gradient-worker.service") + node.wait_until_succeeds( + f"journalctl -u gradient-worker --no-pager --since=@{since} | grep -q 'handshake successful'", + timeout=180, + ) + + server.succeed( + "mkdir -p /srv/upstream/input-src && printf fetched > /srv/upstream/input-src/value" + " && tar -czf /srv/upstream/input.tar.gz -C /srv/upstream input-src" + ) + input_hash = server.succeed( + f"{NIX} --extra-experimental-features nix-command hash path --sri /srv/upstream/input-src" + ).strip() + input_base32 = server.succeed( + f"{NIX} --extra-experimental-features nix-command hash convert --to nix32 {input_hash}" + ).strip() + input_path = server.succeed( + f"{NIX}-store --print-fixed-path --recursive sha256 {input_base32} source" + ).strip() + input_narinfo = input_path.removeprefix("/nix/store/").split("-", 1)[0] + ".narinfo" + server.fail(f"{CURL} -sf {CACHE}/{input_narinfo}") + + fetch_since = builder2.succeed("date +%s").strip() + restart_with_caps(builder, {"EVAL": "false"}) + restart_with_caps(builder2, {"NETWORK": "false", "SECRETS": "false"}) + + known = sql("SELECT string_agg(id::text, ',') FROM evaluation;") + input_evals = f"FROM evaluation WHERE NOT (id = ANY(string_to_array('{known}', ',')::uuid[]))" + server.succeed("cp /var/lib/git/flake-input.nix /var/lib/git/test/flake.nix") + server.succeed("cp /var/lib/git/flake-input.lock /var/lib/git/test/flake.lock") + server.succeed(f"sed -i 's#\\[hash\\]#{input_hash}#g' /var/lib/git/test/flake.lock") + server.succeed(f"{GIT} -C /var/lib/git/test commit -am 'input'") + server.succeed("chown git:git -R /var/lib/git/test") + + poll(f"SELECT count(*) {input_evals};", "2", "both tasks did not pick up the input commit") + poll(f"SELECT string_agg(DISTINCT status::text, ',') {input_evals};", "5", + "the evaluations of the input flake did not complete", timeout=600) + server.succeed(f"{CURL} -sf {CACHE}/{input_narinfo}") + builder.succeed(f"journalctl -u gradient-worker --no-pager --since=@{fetch_since} | grep -q 'fetching an input'") + builder2.fail(f"journalctl -u gradient-worker --no-pager --since=@{fetch_since} | grep -q 'fetching an input'") + + restart_with_caps(builder, {}) + restart_with_caps(builder2, {}) + banner("Phase 11: every supervised loop is running; SIGTERM stops the server") health = json.loads(api_get(token, "board/health"))["message"] names = sorted(l["name"] for l in health["supervised"]) From 88a76f82d1a2521c64352b5bfe49c47dafb3b26d Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 16:21:21 +0200 Subject: [PATCH 15/19] scheduler: feat: retry unplaced input fetches --- .../gradient-scheduler/src/cases/inputs.rs | 78 +++++++++++++++++++ backend/gradient-scheduler/src/input_waits.rs | 18 +++++ .../src/job_handlers/input.rs | 8 ++ backend/gradient-wire/schema/v35.txt | 4 +- backend/gradient-wire/src/types.rs | 1 + .../src/executor/dependencies.rs | 2 +- docs/src/contributors/proto/messages.md | 2 +- 7 files changed, 109 insertions(+), 4 deletions(-) diff --git a/backend/gradient-scheduler/src/cases/inputs.rs b/backend/gradient-scheduler/src/cases/inputs.rs index 829a75a44..8c28a0c20 100644 --- a/backend/gradient-scheduler/src/cases/inputs.rs +++ b/backend/gradient-scheduler/src/cases/inputs.rs @@ -402,6 +402,84 @@ async fn i10_a_running_fetch_takes_a_build_slot() { case.finish().await; } +#[tokio::test] +async fn i11_an_unplaced_fetch_goes_to_another_worker() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start(shape().evalq(e1)).await; + evaluating(&case, W1, e1).await; + case.connect(W2, w_fetch(), &[]).await; + case.connect(W3, w_fetch(), &[]).await; + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + let first = single_fetch(&case, &[W1, W2, W3]); + assert_eq!( + first.worker, W2, + "after 1: the lower worker id wins the tie" + ); + case.check("step 1").await; + + case.input_fetched(W2, &first.fetch_id, unplaced()).await; + let second = single_fetch(&case, &[W1, W2, W3]); + assert_eq!(second.worker, W3, "after 2: the fetch moves to W3"); + assert_ne!(second.fetch_id, first.fetch_id, "after 2: a new fetch id"); + case.check("step 2").await; + + case.input_fetched(W3, &second.fetch_id, fetched()).await; + assert_eq!( + without_offers(case.signals(W1)), + vec![result(e1, "r1", fetched())], + "after 3: the evaluation gets the result of W3" + ); + case.check("step 3").await; + + case.finish().await; +} + +#[tokio::test] +async fn i12_an_unplaced_fetch_waits_when_no_other_worker_qualifies() { + let e1 = EvaluationId::now_v7(); + let case = SchedulerCase::start(shape().evalq(e1)).await; + evaluating(&case, W1, e1).await; + case.connect(W2, w_fetch(), &[]).await; + case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) + .await; + let first = single_fetch(&case, &[W1, W2]); + case.check("step 1").await; + + case.input_fetched(W2, &first.fetch_id, unplaced()).await; + case.build_assign_pass(false, false).await; + assert_eq!( + fetches_of(&case, &[W1, W2]), + vec![], + "after 2: W2 is not asked again" + ); + assert_eq!(case.input_waits(), 1, "after 2: the request still waits"); + let reason = case.evaluation(e1).await.waiting_reason; + assert_eq!( + reason.as_ref().and_then(WaitingReason::from_json), + Some(WaitingReason::eval_workers(EvalCapability::Network, 2)), + "after 2: the evaluation waits for network" + ); + case.check("step 2").await; + + case.connect(W3, w_fetch(), &[]).await; + case.build_assign_pass(false, false).await; + assert_eq!( + single_fetch(&case, &[W1, W2, W3]).worker, + W3, + "after 3: W3 takes the fetch" + ); + case.check("step 3").await; + + case.finish().await; +} + +fn unplaced() -> InputOutcome { + InputOutcome::Unplaced { + message: "no worker behind the shared worker can fetch the input".into(), + } +} + async fn building_on_w2(case: &SchedulerCase) -> (String, gradient_entity::ids::DispatchedJobId) { let a = case.key(A).await; case.admit_startable_moves().await; diff --git a/backend/gradient-scheduler/src/input_waits.rs b/backend/gradient-scheduler/src/input_waits.rs index f6a27b562..c8936764e 100644 --- a/backend/gradient-scheduler/src/input_waits.rs +++ b/backend/gradient-scheduler/src/input_waits.rs @@ -30,6 +30,7 @@ pub struct UnassignedFetch { pub project: ProjectId, pub request: String, pub needs_secrets: bool, + pub excluded: HashSet, } #[derive(Debug)] @@ -39,6 +40,7 @@ struct InputWait { needs_secrets: bool, waiters: Vec, fetch: Option, + excluded: HashSet, } #[derive(Debug, Clone)] @@ -64,6 +66,7 @@ impl InputWaits { needs_secrets: new.needs_secrets, waiters: Vec::new(), fetch: None, + excluded: HashSet::new(), }) .waiters .push(new.waiter); @@ -80,6 +83,7 @@ impl InputWaits { project: wait.project, request: wait.request.clone(), needs_secrets: wait.needs_secrets, + excluded: wait.excluded.clone(), }) .collect(); unassigned.sort_by(|a, b| a.fetch_key.cmp(&b.fetch_key)); @@ -112,6 +116,19 @@ impl InputWaits { waits.remove(&key).map(|w| w.waiters).unwrap_or_default() } + pub fn unplace(&self, worker: &str, fetch_id: &str) -> bool { + let mut waits = self.0.lock(); + let Some(key) = key_of_fetch(&waits, worker, fetch_id) else { + return false; + }; + let Some(wait) = waits.get_mut(&key) else { + return false; + }; + wait.fetch = None; + wait.excluded.insert(worker.to_owned()); + true + } + pub fn waiters_of_fetch(&self, worker: &str, fetch_id: &str) -> Vec { let waits = self.0.lock(); key_of_fetch(&waits, worker, fetch_id) @@ -124,6 +141,7 @@ impl InputWaits { let mut waits = self.0.lock(); for wait in waits.values_mut() { wait.waiters.retain(|w| w.worker != worker); + wait.excluded.remove(worker); if wait.fetch.as_ref().is_some_and(|f| f.worker == worker) { wait.fetch = None; } diff --git a/backend/gradient-scheduler/src/job_handlers/input.rs b/backend/gradient-scheduler/src/job_handlers/input.rs index 528512eed..a713efca9 100644 --- a/backend/gradient-scheduler/src/job_handlers/input.rs +++ b/backend/gradient-scheduler/src/job_handlers/input.rs @@ -65,6 +65,13 @@ impl Scheduler { outcome: InputOutcome, ) { self.release_input_fetch(worker_id, fetch_id).await; + if let InputOutcome::Unplaced { message } = &outcome { + if self.input_waits.unplace(worker_id, fetch_id) { + debug!(%worker_id, %fetch_id, %message, "input fetch unplaced; trying another worker"); + self.dispatch_input_fetches().await; + } + return; + } let waiters = self.input_waits.take_fetched(worker_id, fetch_id); if waiters.is_empty() { debug!(%worker_id, %fetch_id, "input fetch answered without waiters; dropped"); @@ -247,6 +254,7 @@ fn fetch_workers( let mut eligible: Vec<&WorkerInfo> = workers .iter() .filter(|w| !w.draining) + .filter(|w| !fetch.excluded.contains(&w.id)) .filter(|w| w.assigned_build_count < w.max_concurrent_builds as usize) .filter(|w| { w.authorized_peers diff --git a/backend/gradient-wire/schema/v35.txt b/backend/gradient-wire/schema/v35.txt index f04239de1..364f9788d 100644 --- a/backend/gradient-wire/schema/v35.txt +++ b/backend/gradient-wire/schema/v35.txt @@ -29,7 +29,7 @@ client 27:{u64} 28:{str,str,seq} 29:{str,str,str} -30:{str,[0:{str},1:{str}]} +30:{str,[0:{str},1:{str},2:{str}]} 31:{str,u64,opt} server 0:{seq} @@ -59,6 +59,6 @@ server 24:{u64,[0:{},1:{str},2:{str}]} 25:{str,seq<(str,str)>} 26:{str,str,[0:{},1:{str,str,str},2:{str}]} -27:{str,str,[0:{str},1:{str}]} +27:{str,str,[0:{str},1:{str},2:{str}]} 28:{str,str,str,opt} 29:{str,str,u64,opt} diff --git a/backend/gradient-wire/src/types.rs b/backend/gradient-wire/src/types.rs index 2255a0ca2..04e938749 100644 --- a/backend/gradient-wire/src/types.rs +++ b/backend/gradient-wire/src/types.rs @@ -326,6 +326,7 @@ pub enum ImportOutcome { pub enum InputOutcome { Fetched { response: String }, Failed { message: String }, + Unplaced { message: String }, } #[derive(Proto, Debug, Clone, PartialEq)] diff --git a/backend/gradient-worker/src/executor/dependencies.rs b/backend/gradient-worker/src/executor/dependencies.rs index 968dc1ccc..d7bb6ed65 100644 --- a/backend/gradient-worker/src/executor/dependencies.rs +++ b/backend/gradient-worker/src/executor/dependencies.rs @@ -175,7 +175,7 @@ async fn fetch_input( .map_err(|e| format!("pulling the fetched input {fetched} failed: {e:#}"))?; Ok(response) } - InputOutcome::Failed { message } => Err(message), + InputOutcome::Failed { message } | InputOutcome::Unplaced { message } => Err(message), } } diff --git a/docs/src/contributors/proto/messages.md b/docs/src/contributors/proto/messages.md index 726a21494..9a22dfdfb 100644 --- a/docs/src/contributors/proto/messages.md +++ b/docs/src/contributors/proto/messages.md @@ -72,7 +72,7 @@ The message definitions live in `backend/gradient-wire/src/messages` and cover e | `ImportRequest` | Builds the given derivations now, because the evaluation needs them | `job_id`, `request_id`, `drv_paths` | | `InputRequest` | Asks for an input missing from the store and the cache, because the evaluation needs the input | `job_id`, `request_id`, `request` | | `InputFetchProgress` | Reports download bytes of a `FetchInput` | `fetch_id`, `bytes_done`, `bytes_total` | -| `InputFetched` | Reply to `FetchInput` after the push: the fetch answer or the download error | `fetch_id`, `outcome` | +| `InputFetched` | Reply to `FetchInput` after the push: the fetch answer or the download error. `Unplaced` answers return the fetch to the server for another worker | `fetch_id`, `outcome` | ## Cache Query Modes From 445a3e7c25c66f7593edf36af1a7e96e0f86762d Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 16:44:43 +0200 Subject: [PATCH 16/19] worker: test: accept the mock connection --- backend/gradient-worker/src/worker/message_loop.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/backend/gradient-worker/src/worker/message_loop.rs b/backend/gradient-worker/src/worker/message_loop.rs index 76ad051b5..26c88ca1d 100644 --- a/backend/gradient-worker/src/worker/message_loop.rs +++ b/backend/gradient-worker/src/worker/message_loop.rs @@ -1402,9 +1402,17 @@ mod tests { } #[tokio::test] + #[expect( + clippy::disallowed_methods, + reason = "the test stands in for the server by hand" + )] async fn a_running_input_fetch_occupies_a_build_slot() { let server = gradient_test_support::prelude::MockProtoServer::bind().await; let url = server.url().to_owned(); + let _accepted = tokio::spawn(async move { + let _sc = server.accept().await; + std::future::pending::<()>().await + }); let conn = gradient_worker_client::connection::ProtoConnection::open(&url) .await .unwrap(); From ea93618074ad02b7f79afbb9b19b3e1dd028111b Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 16:44:43 +0200 Subject: [PATCH 17/19] scheduler: test: q18 with an eval-only worker --- backend/gradient-scheduler/src/cases/eval_queue.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/backend/gradient-scheduler/src/cases/eval_queue.rs b/backend/gradient-scheduler/src/cases/eval_queue.rs index 93c8db0e0..1eb92a160 100644 --- a/backend/gradient-scheduler/src/cases/eval_queue.rs +++ b/backend/gradient-scheduler/src/cases/eval_queue.rs @@ -30,6 +30,7 @@ const SSH_REPOSITORY: &str = "git+ssh://git@example.invalid/repo.git"; const W1: &str = "w1"; const W2: &str = "w2"; +const W3: &str = "w3"; #[tokio::test] async fn q01_a_queued_evaluation_becomes_one_pending_job_with_wildcards_and_overrides() { @@ -893,6 +894,7 @@ async fn q18_a_fetch_step_of_an_ssh_repository_needs_secrets() { let case = SchedulerCase::start(shape().evalq(e1).repository(SSH_REPOSITORY)).await; case.connect(W1, w_network(), &[]).await; case.connect(W2, w_network_secrets(), &[]).await; + case.connect(W3, w_eval(), &[]).await; let key = eval_key(e1); case.assign_queued_evals().await; @@ -900,6 +902,11 @@ async fn q18_a_fetch_step_of_an_ssh_repository_needs_secrets() { let without_secrets = case.offers_full(W1).await; let with_secrets = case.offers_full(W2).await; + assert_eq!( + candidate_keys(&case.offers_full(W3).await), + BTreeSet::new(), + "after 2: offers(W3).candidates" + ); assert_eq!( candidate_keys(&without_secrets), BTreeSet::new(), From 3fa3eaa91a4d0d90a8e12dffb89f67c15e7ef5d5 Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 16:44:43 +0200 Subject: [PATCH 18/19] scheduler: fix: no waiting reason for input waits --- .../gradient-scheduler/src/cases/inputs.rs | 46 +++++++----- backend/gradient-scheduler/src/input_waits.rs | 22 ++++-- .../src/job_handlers/input.rs | 62 ++++++++++----- .../gradient-scheduler/src/waiting_state.rs | 75 +++++-------------- .../src/worker_lifecycle.rs | 2 +- 5 files changed, 108 insertions(+), 99 deletions(-) diff --git a/backend/gradient-scheduler/src/cases/inputs.rs b/backend/gradient-scheduler/src/cases/inputs.rs index 8c28a0c20..b46813851 100644 --- a/backend/gradient-scheduler/src/cases/inputs.rs +++ b/backend/gradient-scheduler/src/cases/inputs.rs @@ -8,7 +8,6 @@ use gradient_entity::evaluation::EvaluationStatus; use gradient_entity::ids::{EvaluationId, ProjectId}; use gradient_fixtures::shape::shape; use gradient_pool::session_port::SessionSignal; -use gradient_types::waiting_reason::{EvalCapability, WaitingReason}; use gradient_wire::types::{InputOutcome, JobKind}; use super::{ @@ -145,6 +144,7 @@ async fn i04_a_lost_fetch_worker_hands_the_fetch_to_another() { case.check("step 1").await; case.disconnect(W2).await; + case.signals(W2); let second = single_fetch(&case, &[W1, W2, W3]); assert_eq!(second.worker, W3, "after 2: W3 takes over"); assert_ne!(second.fetch_id, first.fetch_id, "after 2: a new fetch id"); @@ -206,7 +206,7 @@ async fn i05_an_abort_drops_the_input_wait() { } #[tokio::test] -async fn i06_without_a_fetch_worker_the_evaluation_waits_for_network() { +async fn i06_without_a_fetch_worker_the_evaluation_says_it_waits() { let e1 = EvaluationId::now_v7(); let case = SchedulerCase::start(shape().evalq(e1)).await; evaluating(&case, W1, e1).await; @@ -214,19 +214,17 @@ async fn i06_without_a_fetch_worker_the_evaluation_waits_for_network() { case.input_request(W1, &eval_key(e1), "r1", SSH_REQUEST) .await; case.refresh_waiting_state().await; + case.build_assign_pass(false, false).await; let row = case.evaluation(e1).await; assert_eq!( - ( - row.status, - row.waiting_reason - .as_ref() - .and_then(WaitingReason::from_json) - ), - ( - EvaluationStatus::EvaluatingDerivation, - Some(WaitingReason::eval_workers(EvalCapability::Network, 1)) - ), - "after 1: evaluation(E1) keeps evaluating and names the gap" + (row.status, row.waiting_reason), + (EvaluationStatus::EvaluatingDerivation, None), + "after 1: evaluation(E1) keeps evaluating without a waiting reason" + ); + assert_eq!( + wait_messages(&case, e1).await, + vec!["waiting for a worker with network and secrets to fetch an input".to_owned()], + "after 1: one message names the missing worker, also after a second pass" ); case.check("step 1").await; @@ -300,6 +298,7 @@ async fn i08_fetch_progress_reaches_every_waiting_evaluation() { case.check("step 2").await; case.disconnect(&first.worker).await; + case.signals(&first.worker); let second = single_fetch(&case, &[W1, W2, W3, W4]); case.input_fetch_progress(&first.worker, &first.fetch_id, 1024, Some(1024)) .await; @@ -454,11 +453,15 @@ async fn i12_an_unplaced_fetch_waits_when_no_other_worker_qualifies() { "after 2: W2 is not asked again" ); assert_eq!(case.input_waits(), 1, "after 2: the request still waits"); - let reason = case.evaluation(e1).await.waiting_reason; assert_eq!( - reason.as_ref().and_then(WaitingReason::from_json), - Some(WaitingReason::eval_workers(EvalCapability::Network, 2)), - "after 2: the evaluation waits for network" + case.evaluation(e1).await.waiting_reason, + None, + "after 2: no waiting reason while evaluating" + ); + assert_eq!( + wait_messages(&case, e1).await, + vec!["waiting for a worker with network and secrets to fetch an input".to_owned()], + "after 2: the excluded W2 counts as no fetch worker" ); case.check("step 2").await; @@ -474,6 +477,15 @@ async fn i12_an_unplaced_fetch_waits_when_no_other_worker_qualifies() { case.finish().await; } +async fn wait_messages(case: &SchedulerCase, evaluation: EvaluationId) -> Vec { + case.messages(evaluation) + .await + .into_iter() + .filter(|m| m.source.as_deref() == Some("fetch")) + .map(|m| m.message) + .collect() +} + fn unplaced() -> InputOutcome { InputOutcome::Unplaced { message: "no worker behind the shared worker can fetch the input".into(), diff --git a/backend/gradient-scheduler/src/input_waits.rs b/backend/gradient-scheduler/src/input_waits.rs index c8936764e..06532fc9b 100644 --- a/backend/gradient-scheduler/src/input_waits.rs +++ b/backend/gradient-scheduler/src/input_waits.rs @@ -41,6 +41,7 @@ struct InputWait { waiters: Vec, fetch: Option, excluded: HashSet, + announced: HashSet, } #[derive(Debug, Clone)] @@ -67,6 +68,7 @@ impl InputWaits { waiters: Vec::new(), fetch: None, excluded: HashSet::new(), + announced: HashSet::new(), }) .waiters .push(new.waiter); @@ -170,13 +172,19 @@ impl InputWaits { .collect() } - pub fn unassigned_evaluations(&self) -> HashSet { - self.0 - .lock() - .values() - .filter(|w| w.fetch.is_none()) - .flat_map(|w| w.waiters.iter().map(|waiter| waiter.evaluation)) - .collect() + pub fn announce_unplaced(&self, fetch_key: &str) -> Vec { + let mut waits = self.0.lock(); + let Some(wait) = waits.get_mut(fetch_key).filter(|w| w.fetch.is_none()) else { + return Vec::new(); + }; + let mut fresh: Vec = wait + .waiters + .iter() + .map(|waiter| waiter.evaluation) + .filter(|evaluation| wait.announced.insert(*evaluation)) + .collect(); + fresh.dedup(); + fresh } #[cfg(test)] diff --git a/backend/gradient-scheduler/src/job_handlers/input.rs b/backend/gradient-scheduler/src/job_handlers/input.rs index a713efca9..2990fe3f8 100644 --- a/backend/gradient-scheduler/src/job_handlers/input.rs +++ b/backend/gradient-scheduler/src/job_handlers/input.rs @@ -19,6 +19,7 @@ use crate::Scheduler; use crate::actor::SchedulerMsg; use crate::input_waits::{AssignedFetch, InputWaiter, NewInputWait, UnassignedFetch}; use crate::jobs::PendingJob; +use crate::waiting_state::announce_worker_gap; impl Scheduler { pub async fn handle_input_request( @@ -149,10 +150,6 @@ impl Scheduler { Ok(()) } - pub(crate) fn input_waiting_evaluations(&self) -> HashSet { - self.input_waits.unassigned_evaluations() - } - async fn dispatch_input_fetches(&self) { let unassigned = self.input_waits.unassigned(); if unassigned.is_empty() { @@ -167,6 +164,10 @@ impl Scheduler { .await { debug!(fetch_key = %fetch.fetch_key, "no connected worker can fetch the input yet"); + if !workers.iter().any(|w| can_ever_fetch(w, &fetch)) { + self.announce_missing_fetch_worker(&fetch, workers.len()) + .await; + } } } } @@ -214,6 +215,30 @@ impl Scheduler { false } + async fn announce_missing_fetch_worker(&self, fetch: &UnassignedFetch, connected: usize) { + let (capability, wanted) = if fetch.needs_secrets { + (EvalCapability::Secrets, "network and secrets") + } else { + (EvalCapability::Network, "network") + }; + let reason = + WaitingReason::eval_workers(capability, connected.try_into().unwrap_or(u32::MAX)); + for evaluation in self.input_waits.announce_unplaced(&fetch.fetch_key) { + let inserted = gradient_db::status::insert_worker_message( + &self.db.worker_db, + evaluation, + gradient_entity::evaluation_message::MessageLevel::Notice, + format!("waiting for a worker with {wanted} to fetch an input"), + "fetch".to_owned(), + ) + .await; + if let Err(e) = inserted { + warn!(error = %e, %evaluation, "input wait message not recorded"); + } + announce_worker_gap(&self.db, evaluation, fetch.project, &reason).await; + } + } + async fn release_input_fetch(&self, worker: &str, fetch_id: &str) { let released = self .call(|reply| SchedulerMsg::ReleaseInputFetch { @@ -246,6 +271,20 @@ impl Scheduler { } } +fn can_ever_fetch(worker: &WorkerInfo, fetch: &UnassignedFetch) -> bool { + !worker.draining + && !fetch.excluded.contains(&worker.id) + && worker + .authorized_peers + .as_ref() + .is_none_or(|peers| peers.contains(&fetch.project)) + && WorkerCaps { + capabilities: worker.capabilities.clone(), + ..WorkerCaps::default() + } + .can_fetch_input(fetch.needs_secrets) +} + fn fetch_workers( workers: &[WorkerInfo], load: &HashMap, @@ -253,21 +292,8 @@ fn fetch_workers( ) -> Vec { let mut eligible: Vec<&WorkerInfo> = workers .iter() - .filter(|w| !w.draining) - .filter(|w| !fetch.excluded.contains(&w.id)) + .filter(|w| can_ever_fetch(w, fetch)) .filter(|w| w.assigned_build_count < w.max_concurrent_builds as usize) - .filter(|w| { - w.authorized_peers - .as_ref() - .is_none_or(|peers| peers.contains(&fetch.project)) - }) - .filter(|w| { - WorkerCaps { - capabilities: w.capabilities.clone(), - ..WorkerCaps::default() - } - .can_fetch_input(fetch.needs_secrets) - }) .collect(); eligible.sort_by_key(|w| (load.get(&w.id).copied().unwrap_or_default(), w.id.clone())); eligible.into_iter().map(|w| w.id.clone()).collect() diff --git a/backend/gradient-scheduler/src/waiting_state.rs b/backend/gradient-scheduler/src/waiting_state.rs index f08075060..c89371334 100644 --- a/backend/gradient-scheduler/src/waiting_state.rs +++ b/backend/gradient-scheduler/src/waiting_state.rs @@ -40,7 +40,6 @@ impl WaitingRefresh<'_> { pub(crate) async fn run( &self, workers: &[WorkerInfo], - input_waiting: &HashSet, draining: bool, ) -> Result> { gradient_db::evaluations::counters::fold_shared_build_deltas(&self.db.worker_db) @@ -155,15 +154,7 @@ impl WaitingRefresh<'_> { eval.status, check_repository_url_is_ssh(&eval.repository), project_workers, - ) - .or_else(|| { - decide_input_wait( - eval.status, - reason.as_ref(), - input_waiting.contains(&eval.id), - project_workers.connected, - ) - }), + ), EvaluationStatus::Building => { match self .build_phase_decision(eval.id, eval_counters, worker_caps, reason.as_ref()) @@ -206,7 +197,7 @@ 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; + announce_worker_gap(self.db, eval.id, project, gap).await; } persist_waiting_reason( @@ -224,26 +215,24 @@ 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; - } +pub(crate) async fn announce_worker_gap( + db: &DbContext, + evaluation: EvaluationId, + project: ProjectId, + reason: &WaitingReason, +) { + let teams = + match gradient_db::teams::workers::teams_granting_workers(&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) { + db.record(event).await; } } @@ -784,32 +773,6 @@ fn decide_eval_recovery( } } -fn decide_input_wait( - status: EvaluationStatus, - current: Option<&WaitingReason>, - waits_for_input: bool, - connected: u32, -) -> Option<(EvaluationStatus, Option)> { - let waits_for_network = matches!( - current, - Some(WaitingReason::EvalWorkers { - capability: EvalCapability::Network, - .. - }) - ); - match (waits_for_input, waits_for_network) { - (true, _) => Some(( - status, - Some(WaitingReason::eval_workers( - EvalCapability::Network, - connected, - )), - )), - (false, true) => Some((status, None)), - (false, false) => None, - } -} - pub(crate) async fn persist_waiting_reason( db: &WorkerDb, evaluation_id: EvaluationId, diff --git a/backend/gradient-scheduler/src/worker_lifecycle.rs b/backend/gradient-scheduler/src/worker_lifecycle.rs index ef1bd245b..02edf0a4c 100644 --- a/backend/gradient-scheduler/src/worker_lifecycle.rs +++ b/backend/gradient-scheduler/src/worker_lifecycle.rs @@ -329,7 +329,7 @@ impl Scheduler { graph: &self.graph, memo: &self.scoring.assessments, } - .run(&workers, &self.input_waiting_evaluations(), draining) + .run(&workers, draining) .await?; for unbuildable in unbuildables { self.abort_unbuildable_evaluation(unbuildable).await; From 31d34db5a9b6d89926fe2f1ca1fa96f005941a8f Mon Sep 17 00:00:00 2001 From: Dennis Wuitz Date: Sun, 18 Oct 2026 17:43:20 +0200 Subject: [PATCH 19/19] e2e: fix: start the input evaluations by API --- nix/tests/gradient/e2e/mk.nix | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/nix/tests/gradient/e2e/mk.nix b/nix/tests/gradient/e2e/mk.nix index 793f2fc22..40adf0b4c 100644 --- a/nix/tests/gradient/e2e/mk.nix +++ b/nix/tests/gradient/e2e/mk.nix @@ -2488,8 +2488,13 @@ pkgs.testers.runNixOSTest ({ pkgs, lib, ... }: { server.succeed(f"sed -i 's#\\[hash\\]#{input_hash}#g' /var/lib/git/test/flake.lock") server.succeed(f"{GIT} -C /var/lib/git/test commit -am 'input'") server.succeed("chown git:git -R /var/lib/git/test") + for task in ("task", "task2"): + server.succeed( + f'{CURL} -sf -X POST -H "Authorization: Bearer {token}" ' + f'-H "Content-Type: application/json" -d \'{{}}\' {API}/tasks/project/{task}/evaluate' + ) - poll(f"SELECT count(*) {input_evals};", "2", "both tasks did not pick up the input commit") + poll(f"SELECT count(*) {input_evals};", "2", "both tasks did not start an evaluation of the input commit") poll(f"SELECT string_agg(DISTINCT status::text, ',') {input_evals};", "5", "the evaluations of the input flake did not complete", timeout=600) server.succeed(f"{CURL} -sf {CACHE}/{input_narinfo}")