Cooldown state rework

Cargo.toml

@@ -79,6 +79,7 @@ indicatif = "0.17.5"
 tokenizers = { version = "0.20.0", default-features = false, features = [
   "onig",
 ] }
+google-cloud-storage = "1.6.0"
 tch = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" }
 torch-sys = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" }
 pyo3-tch = { git = "https://github.com/jquesnelle/tch-rs.git", rev = "11d1ca2ef6dbd3f1e5b0986fab0a90fbb6734496" }

architectures/centralized/client/Cargo.toml

@@ -25,6 +25,8 @@ time.workspace = true
 bytemuck.workspace = true
 clap-markdown.workspace = true
 hex = "0.4.3"
+bytes.workspace = true
+google-cloud-storage.workspace = true
 psyche-python-extension-impl = { workspace = true, optional = true }
 
 [features]

architectures/centralized/client/src/app.rs

@@ -1,13 +1,13 @@
 use anyhow::{Error, Result};
 use bytemuck::Zeroable;
-use hf_hub::Repo;
+use google_cloud_storage::client::{Storage, StorageControl};
 use psyche_centralized_shared::{ClientId, ClientToServerMessage, ServerToClientMessage};
-use psyche_client::HubUploadInfo;
-use psyche_client::UploadInfo;
 use psyche_client::{
     Client, ClientTUI, ClientTUIState, NC, RunInitConfig, TrainArgs, read_identity_secret_key,
 };
-use psyche_coordinator::{Coordinator, HealthChecks, model};
+use psyche_client::{GcsUploadInfo, HubUploadInfo, UploadInfo};
+use psyche_coordinator::model::Checkpoint;
+use psyche_coordinator::{Coordinator, HealthChecks};
 use psyche_metrics::ClientMetrics;
 use psyche_network::{
     AuthenticatableIdentity, EndpointId, NetworkTUIState, NetworkTui, SecretKey, TcpClient,
@@ -31,7 +31,7 @@ pub type TabsData = <Tabs as CustomWidget>::Data;
 pub enum ToSend {
     Witness(Box<OpportunisticData>),
     HealthCheck(HealthChecks<ClientId>),
-    Checkpoint(model::Checkpoint),
+    Checkpoint(Checkpoint),
 }
 
 struct Backend {
@@ -69,7 +69,7 @@ impl WatcherBackend<ClientId> for Backend {
         Ok(())
     }
 
-    async fn send_checkpoint(&mut self, checkpoint: model::Checkpoint) -> Result<()> {
+    async fn send_checkpoint(&mut self, checkpoint: Checkpoint) -> Result<()> {
         self.tx.send(ToSend::Checkpoint(checkpoint))?;
         Ok(())
     }
@@ -84,13 +84,15 @@ pub struct App {
     server_conn: TcpClient<ClientId, ClientToServerMessage, ServerToClientMessage>,
 
     metrics: Arc<ClientMetrics>,
+    skip_upload_check: bool,
 }
 
 pub async fn build_app(
     cancel: CancellationToken,
     server_addr: String,
     tx_tui_state: Option<Sender<TabsData>>,
     p: TrainArgs,
+    is_test: bool,
 ) -> Result<(
     App,
     allowlist::AllowDynamic,
@@ -162,6 +164,7 @@ pub async fn build_app(
         server_conn,
         run_id: p.run_id,
         metrics,
+        skip_upload_check: is_test,
     };
     Ok((app, allowlist, p2p, state_options))
 }
@@ -173,22 +176,59 @@ impl App {
         p2p: NC,
         state_options: RunInitConfig<ClientId, ClientId>,
     ) -> Result<()> {
-        // sanity checks
-        if let Some(checkpoint_config) = &state_options.checkpoint_config {
-            if let Some(UploadInfo::Hub(HubUploadInfo {
-                hub_repo,
-                hub_token,
-            })) = &checkpoint_config.upload_info
-            {
-                let api = hf_hub::api::tokio::ApiBuilder::new()
-                    .with_token(Some(hub_token.clone()))
-                    .build()?;
-                let repo_api = api.repo(Repo::new(hub_repo.clone(), hf_hub::RepoType::Model));
-                if !repo_api.is_writable().await {
+        // Sanity checks using the checkpoint config from state_options, not the zeroed coordinator state.
+        // The coordinator_state is only populated after receiving the first ServerToClientMessage::Coordinator.
+        if !self.skip_upload_check {
+            let upload_info = match &state_options.checkpoint_config {
+                config if config.skip_upload => Some(UploadInfo::Dummy()),
+                config => {
+                    // Use HF_TOKEN from checkpoint_config for Hub uploads
+                    if let Some(ref hub_token) = config.hub_token {
+                        Some(UploadInfo::Hub(HubUploadInfo {
+                            hub_repo: String::new(), // Will be validated when actual checkpoint is received
+                            hub_token: hub_token.clone(),
+                        }))
+                    } else {
+                        // Check if GCS credentials are available by attempting to create a client
+                        match Storage::builder().build().await {
+                            Ok(_) => Some(UploadInfo::Gcs(GcsUploadInfo {
+                                gcs_bucket: String::new(), // Will be validated when actual checkpoint is received
+                                gcs_prefix: None,
+                            })),
+                            Err(_) => None,
+                        }
+                    }
+                }
+            };
+
+            match upload_info {
+                Some(UploadInfo::Hub(HubUploadInfo {
+                    hub_repo: _,
+                    hub_token,
+                })) => {
+                    let _api = hf_hub::api::tokio::ApiBuilder::new()
+                        .with_token(Some(hub_token.clone()))
+                        .build()?;
+                }
+                Some(UploadInfo::Gcs(_gcs_info)) => {
+                    let _storage = Storage::builder()
+                        .build()
+                        .await
+                        .map_err(|e| anyhow::anyhow!("Failed to create GCS client: {}", e))?;
+
+                    let _storage_control =
+                        StorageControl::builder().build().await.map_err(|e| {
+                            anyhow::anyhow!("Failed to create GCS control client: {}", e)
+                        })?;
+                    // GCS credentials are valid - actual bucket writability will be checked during checkpoint
+                }
+                Some(UploadInfo::Dummy()) => {
+                    // In test mode or skip_upload mode, we skip upload checks
+                }
+                None => {
                     anyhow::bail!(
-                        "Checkpoint upload repo {} is not writable with the passed API key.",
-                        hub_repo
-                    )
+                        "No upload credentials found for checkpointing. Set HF_TOKEN for HuggingFace Hub or configure GCS credentials."
+                    );
                 }
             }
         }

architectures/centralized/client/src/main.rs

@@ -105,7 +105,7 @@ async fn async_main() -> Result<()> {
             )?;
 
             let (mut app, allowlist, p2p, state_options) =
-                build_app(cancel, server_addr, tx_tui_state, args)
+                build_app(cancel, server_addr, tx_tui_state, args, false)
                     .await
                     .unwrap();
 

architectures/centralized/server/src/app.rs

@@ -1,7 +1,7 @@
 use anyhow::{Result, anyhow, bail};
 use async_trait::async_trait;
 use psyche_centralized_shared::{ClientId, ClientToServerMessage, ServerToClientMessage};
-use psyche_coordinator::model::{self, Checkpoint, LLM, LLMTrainingDataLocation, Model};
+use psyche_coordinator::model::{Checkpoint, LLM, LLMTrainingDataLocation, Model};
 use psyche_coordinator::{
     Client, ClientState, Coordinator, CoordinatorError, HealthChecks, Round, RunState,
     SOLANA_MAX_NUM_CLIENTS, TickResult,
@@ -81,7 +81,7 @@ impl psyche_watcher::Backend<ClientId> for ChannelCoordinatorBackend {
         bail!("Server does not send health checks");
     }
 
-    async fn send_checkpoint(&mut self, _checkpoint: model::Checkpoint) -> Result<()> {
+    async fn send_checkpoint(&mut self, _checkpoint: Checkpoint) -> Result<()> {
         bail!("Server does not send checkpoints");
     }
 }
@@ -402,6 +402,9 @@ impl App {
                         Self::get_timestamp(),
                         rand::rng().next_u64(),
                     ),
+                    OpportunisticData::CooldownStep(witness) => {
+                        self.coordinator.cooldown_witness(&from, witness)
+                    }
                 } {
                     warn!("Error when processing witness: {error}");
                 };

architectures/centralized/testing/src/client.rs

@@ -33,6 +33,7 @@ impl Client {
             client_app_params.server_addr,
             None,
             client_app_params.train_args,
+            true,
         )
         .await
         .unwrap();
@@ -57,6 +58,7 @@ impl Client {
             client_app_params.server_addr,
             None,
             client_app_params.train_args,
+            true,
         )
         .await
         .unwrap();

architectures/centralized/testing/src/lib.rs

@@ -6,4 +6,4 @@ pub mod test_utils;
 pub const WARMUP_TIME: u64 = 60;
 pub const MAX_ROUND_TRAIN_TIME: u64 = 5;
 pub const ROUND_WITNESS_TIME: u64 = 2;
-pub const COOLDOWN_TIME: u64 = 3;
+pub const COOLDOWN_TIME: u64 = 5;

architectures/centralized/testing/src/test_utils.rs

@@ -127,6 +127,7 @@ pub fn dummy_client_app_params_with_training_delay(
     run_id: &str,
     training_delay_secs: u64,
 ) -> AppParams {
+    std::env::set_var("HF_TOKEN", "dummy_token");
     AppParams {
         cancel: CancellationToken::default(),
         server_addr: format!("localhost:{server_port}").to_string(),
@@ -141,6 +142,7 @@ pub fn dummy_client_app_params_with_training_delay(
             "--max-concurrent-parameter-requests", "10",
             "--hub-max-concurrent-downloads", "1",
             "--dummy-training-delay-secs", training_delay_secs.to_string().as_str(),
+            "--skip-checkpoint-upload",
         ])
         .train_args,
     }

architectures/decentralized/justfile

@@ -5,6 +5,8 @@ set working-directory := '../../'
 # In case a recipe is not found here, it will fallback to the root justfile.
 
 AUTHORIZER := env_var_or_default("AUTHORIZER", "11111111111111111111111111111111")
+HF_TOKEN := env_var_or_default("HF_TOKEN", "")
+GOOGLE_APPLICATION_CREDENTIALS := env_var_or_default("GOOGLE_APPLICATION_CREDENTIALS", "")
 
 set fallback := true
 
@@ -37,10 +39,16 @@ setup-solana-localnet-permissioned-light-test-run-treasurer run_id="test" *args=
     RUN_ID={{ run_id }} CONFIG_FILE=./config/solana-test/light-config.toml ./scripts/deploy-solana-test.sh --treasurer {{ args }}
 
 start-training-localnet-client run_id="test" *args='':
-    AUTHORIZER={{ AUTHORIZER }} RUN_ID={{ run_id }} ./scripts/train-solana-test.sh {{ args }}
+    HF_TOKEN={{ HF_TOKEN }} GOOGLE_APPLICATION_CREDENTIALS={{ GOOGLE_APPLICATION_CREDENTIALS }} AUTHORIZER={{ AUTHORIZER }} CHECKPOINT="false" RUN_ID={{ run_id }} ./scripts/train-solana-test.sh {{ args }}
 
 start-training-localnet-light-client run_id="test" *args='':
-    AUTHORIZER={{ AUTHORIZER }} RUN_ID={{ run_id }} BATCH_SIZE=1 DP=1 ./scripts/train-solana-test.sh {{ args }}
+    HF_TOKEN={{ HF_TOKEN }} GOOGLE_APPLICATION_CREDENTIALS={{ GOOGLE_APPLICATION_CREDENTIALS }} AUTHORIZER={{ AUTHORIZER }} CHECKPOINT="false" RUN_ID={{ run_id }} BATCH_SIZE=1 DP=1 ./scripts/train-solana-test.sh {{ args }}
+
+start-training-localnet-light-client-checkpoint run_id="test" *args='':
+    HF_TOKEN={{ HF_TOKEN }} GOOGLE_APPLICATION_CREDENTIALS={{ GOOGLE_APPLICATION_CREDENTIALS }} AUTHORIZER={{ AUTHORIZER }} CHECKPOINT="true" RUN_ID={{ run_id }} BATCH_SIZE=1 DP=1 ./scripts/train-solana-test.sh {{ args }}
+
+start-training-localnet-client-checkpoint run_id="test" *args='':
+    HF_TOKEN={{ HF_TOKEN }} GOOGLE_APPLICATION_CREDENTIALS={{ GOOGLE_APPLICATION_CREDENTIALS }} AUTHORIZER={{ AUTHORIZER }} CHECKPOINT="true" RUN_ID={{ run_id }} ./scripts/train-solana-test.sh {{ args }}
 
 OTLP_METRICS_URL := "http://localhost:4318/v1/metrics"
 OTLP_LOGS_URL := "http://localhost:4318/v1/logs"

architectures/decentralized/solana-client/Cargo.toml

@@ -29,6 +29,8 @@ time.workspace = true
 tokio.workspace = true
 tokio-util.workspace = true
 tracing.workspace = true
+google-cloud-storage.workspace = true
+hf-hub.workspace = true
 psyche-python-extension-impl = { workspace = true, optional = true }
 
 [features]

architectures/decentralized/solana-client/src/app.rs

@@ -1,4 +1,6 @@
 use crate::network_identity::NetworkIdentity;
+use google_cloud_storage::client::StorageControl;
+use hf_hub::Repo;
 use psyche_solana_rpc::SolanaBackend;
 
 use anchor_client::{
@@ -11,9 +13,13 @@ use anchor_client::{
 };
 use anyhow::{Result, anyhow};
 use psyche_client::{
-    Client, ClientTUI, ClientTUIState, NC, RunInitConfig, TrainArgs, read_identity_secret_key,
+    Client, ClientTUI, ClientTUIState, GcsUploadInfo, HubUploadInfo, NC, RunInitConfig, TrainArgs,
+    UploadInfo, read_identity_secret_key,
+};
+use psyche_coordinator::{
+    ClientState, Coordinator, CoordinatorError, RunState,
+    model::{self, GcsRepo, HubRepo, LLM, Model},
 };
-use psyche_coordinator::{ClientState, Coordinator, CoordinatorError, RunState};
 use psyche_core::sha256;
 use psyche_metrics::ClientMetrics;
 
@@ -226,6 +232,77 @@ impl App {
         let mut joined_run_this_epoch = None;
         let mut ever_joined_run = false;
 
+        // sanity checks
+        let Model::LLM(LLM { checkpoint, .. }) = start_coordinator_state.model;
+        let upload_info = match checkpoint {
+            model::Checkpoint::Hub(HubRepo { repo_id, revision })
+            | model::Checkpoint::P2P(HubRepo { repo_id, revision }) => {
+                Some(UploadInfo::Hub(HubUploadInfo {
+                    hub_repo: (&repo_id).into(),
+                    hub_token: (&revision.unwrap_or_default()).into(),
+                }))
+            }
+            model::Checkpoint::Gcs(GcsRepo { bucket, prefix })
+            | model::Checkpoint::P2PGcs(model::GcsRepo { bucket, prefix }) => {
+                Some(UploadInfo::Gcs(GcsUploadInfo {
+                    gcs_bucket: (&bucket).into(),
+                    gcs_prefix: Some((&prefix.unwrap_or_default()).into()),
+                }))
+            }
+            _ => None,
+        };
+        match upload_info {
+            Some(UploadInfo::Hub(hub_info)) => {
+                let api = hf_hub::api::tokio::ApiBuilder::new()
+                    .with_token(Some(hub_info.hub_token))
+                    .build()?;
+                let repo_api = api.repo(Repo::new(
+                    hub_info.hub_repo.clone(),
+                    hf_hub::RepoType::Model,
+                ));
+                if !repo_api.is_writable().await {
+                    anyhow::bail!(
+                        "Checkpoint upload repo {} is not writable with the passed API key.",
+                        hub_info.hub_repo
+                    )
+                }
+            }
+            Some(UploadInfo::Gcs(gcs_info)) => {
+                let client = StorageControl::builder().build().await?;
+
+                let permissions_to_test = vec![
+                    "storage.objects.list",
+                    "storage.objects.get",
+                    "storage.objects.create",
+                    "storage.objects.delete",
+                ];
+
+                let resource = format!("projects/_/buckets/{}", gcs_info.gcs_bucket);
+                let perms_vec: Vec<String> =
+                    permissions_to_test.iter().map(|s| s.to_string()).collect();
+                let response = client
+                    .test_iam_permissions()
+                    .set_resource(&resource)
+                    .set_permissions(perms_vec)
+                    .send()
+                    .await?;
+
+                let correct_permissions = permissions_to_test
+                    .into_iter()
+                    .all(|p| response.permissions.contains(&p.to_string()));
+                if !correct_permissions {
+                    anyhow::bail!(
+                        "GCS bucket {} does not have the required permissions for checkpoint upload make sure to set GOOGLE_APPLICATION_CREDENTIALS environment variable correctly and have the correct permissions to the bucket.",
+                        gcs_info.gcs_bucket
+                    )
+                }
+            }
+            Some(UploadInfo::Dummy()) => {
+                // In test mode, we skip upload checks
+            }
+            None => {}
+        }
+
         // if we're already in "WaitingForMembers" we won't get an update saying that
         // (subscription is on change), so check if it's in that state right at boot
         // and join the run if so

architectures/decentralized/solana-common/src/backend.rs

@@ -19,7 +19,7 @@ use anchor_client::{
 };
 use anyhow::{Context, Result, anyhow};
 use futures_util::StreamExt;
-use psyche_coordinator::model::{self, Checkpoint};
+use psyche_coordinator::model::Checkpoint;
 use psyche_coordinator::{CommitteeProof, Coordinator, HealthChecks};
 use psyche_core::IntegrationTestLogMarker;
 use psyche_watcher::{Backend as WatcherBackend, OpportunisticData};
@@ -309,6 +309,12 @@ impl SolanaBackend {
                 &user,
                 witness,
             ),
+            OpportunisticData::CooldownStep(witness) => instructions::coordinator_cooldown_witness(
+                &coordinator_instance,
+                &coordinator_account,
+                &user,
+                witness,
+            ),
         };
         self.spawn_scheduled_send("Witness", &[instruction], &[]);
     }
@@ -605,7 +611,7 @@ impl WatcherBackend<psyche_solana_coordinator::ClientId> for SolanaBackendRunner
         Ok(())
     }
 
-    async fn send_checkpoint(&mut self, checkpoint: model::Checkpoint) -> Result<()> {
+    async fn send_checkpoint(&mut self, checkpoint: Checkpoint) -> Result<()> {
         self.backend
             .send_checkpoint(self.instance, self.account, checkpoint);
         Ok(())