-
Notifications
You must be signed in to change notification settings - Fork 16
fix(prism): keep BYOK seal alive for full train wall #146
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,9 +1,11 @@ | ||
| //! Miner-funded Lium keys (BYOK): vault + live-client factory. | ||
| //! | ||
| //! Keys are held in process memory and optionally in a short-TTL | ||
| //! encrypted seal directory (`PRISM_PAYER_VAULT_DIR` + key file). They are | ||
| //! never written to the submission store and never logged. Master still SSHs | ||
| //! with the operator keypair; the miner key pays for rent/terminate only. | ||
| //! Keys are held in process memory and optionally in a TTL-bounded | ||
| //! encrypted seal directory (`PRISM_PAYER_VAULT_DIR` + key file). Default TTL | ||
| //! is ≥36h (train wall + eval + skew) and heartbeats re-seal so long GPU runs | ||
| //! survive control-plane restarts. Keys are never written to the submission | ||
| //! store and never logged. Master still SSHs with the operator keypair; the | ||
| //! miner key pays for rent/terminate only. | ||
|
|
||
| #![forbid(unsafe_code)] | ||
|
|
||
|
|
@@ -15,7 +17,10 @@ use std::sync::{Arc, Mutex}; | |
|
|
||
| use prism_lium::{EvalJobBackend, LiumClient, LiumError, LiumSshConfig, LIUM_API_BASE_URL}; | ||
|
|
||
| pub use sealed::{SealedVaultConfig, DEFAULT_TTL_SECS, DIR_ENV, KEY_ENV}; | ||
| pub use sealed::{ | ||
| expiry_secs, recommended_ttl_secs, SealedVaultConfig, DEFAULT_TTL_SECS, DIR_ENV, | ||
| EVAL_BUDGET_SECS, KEY_ENV, SEAL_SKEW_SECS, TRAIN_WALL_SECS, TTL_ENV, | ||
| }; | ||
|
|
||
| /// In-memory map `submission_id → Lium API key`, with optional sealed persist. | ||
| pub struct PayerKeyVault { | ||
|
|
@@ -103,6 +108,9 @@ impl PayerKeyVault { | |
| } | ||
|
|
||
| /// Clone of the stored key, if any (memory, then sealed file). | ||
| /// | ||
| /// Measure / resume paths rely on this hydrate-from-seal fallback when the | ||
| /// in-memory map is empty after a restart (as long as the seal is unexpired). | ||
| #[must_use] | ||
| pub fn get(&self, submission_id: &str) -> Option<String> { | ||
| if let Ok(g) = self.inner.lock() { | ||
|
|
@@ -118,6 +126,26 @@ impl PayerKeyVault { | |
| Some(key) | ||
| } | ||
|
|
||
| /// Re-persist the seal with a fresh TTL window (memory and/or disk). | ||
| /// | ||
| /// Call before measure and on heartbeats so full-budget trains cannot | ||
| /// outlive the on-disk seal. Returns `false` when no key is available. | ||
| pub fn refresh(&self, submission_id: &str) -> bool { | ||
| let Some(key) = self.get(submission_id) else { | ||
| return false; | ||
| }; | ||
| if let Some(cfg) = &self.sealed { | ||
| if let Err(e) = sealed::persist(cfg, submission_id, &key) { | ||
| tracing::warn!(error = %e, "payer seal refresh failed"); | ||
| return false; | ||
| } | ||
| } | ||
| if let Ok(mut g) = self.inner.lock() { | ||
| g.insert(submission_id.to_owned(), key); | ||
| } | ||
| true | ||
|
Comment on lines
+133
to
+146
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift Preserve and handle sealed-vault refresh failures.
📍 Affects 3 files
🤖 Prompt for AI Agents |
||
| } | ||
|
|
||
| /// Drop the key after the eval finishes (memory + seal). | ||
| pub fn remove(&self, submission_id: &str) { | ||
| if let Ok(mut g) = self.inner.lock() { | ||
|
|
@@ -226,6 +254,7 @@ pub fn require_miner_lium(backend_mode: &str) -> bool { | |
| #[cfg(test)] | ||
| mod tests { | ||
| use super::*; | ||
| use tempfile::tempdir; | ||
|
|
||
| #[test] | ||
| fn vault_roundtrip_and_redacted_debug() { | ||
|
|
@@ -244,4 +273,57 @@ mod tests { | |
| assert!(!require_miner_lium("sim")); | ||
| assert!(!require_miner_lium("sim/openrouter")); | ||
| } | ||
|
|
||
| #[test] | ||
| fn default_ttl_covers_train_wall() { | ||
| const { | ||
| assert!(DEFAULT_TTL_SECS >= TRAIN_WALL_SECS + EVAL_BUDGET_SECS + SEAL_SKEW_SECS); | ||
| } | ||
| assert!(recommended_ttl_secs() >= TRAIN_WALL_SECS); | ||
| assert!(recommended_ttl_secs() >= 36 * 3600); | ||
| } | ||
|
|
||
| #[test] | ||
| fn refresh_extends_seal_expiry() { | ||
| let dir = tempdir().unwrap(); | ||
| let cfg = SealedVaultConfig { | ||
| dir: dir.path().to_path_buf(), | ||
| key: [3u8; 32], | ||
| ttl_secs: 120, | ||
| }; | ||
| let v = PayerKeyVault::new().with_sealed(cfg.clone()); | ||
| v.insert("job1", "sk_miner"); | ||
| let first = sealed::expiry_secs(&cfg, "job1").expect("sealed"); | ||
| // Advance perceived lifetime by re-sealing after a short sleep so the | ||
| // unix expiry second is guaranteed to move forward under CI load. | ||
| std::thread::sleep(std::time::Duration::from_secs(1)); | ||
| assert!(v.refresh("job1")); | ||
| let second = sealed::expiry_secs(&cfg, "job1").expect("refreshed"); | ||
| assert!( | ||
| second > first, | ||
| "refresh must extend expiry (first={first} second={second})" | ||
| ); | ||
| assert_eq!(v.get("job1").as_deref(), Some("sk_miner")); | ||
| } | ||
|
|
||
| #[test] | ||
| fn measure_hydrates_from_seal_without_ram() { | ||
| let dir = tempdir().unwrap(); | ||
| let cfg = SealedVaultConfig { | ||
| dir: dir.path().to_path_buf(), | ||
| key: [5u8; 32], | ||
| ttl_secs: 3600, | ||
| }; | ||
| sealed::persist(&cfg, "onlyseal", "sk_from_disk").unwrap(); | ||
| // Fresh vault: empty RAM, sealed backend only — same as post-restart | ||
| // before hydrate_all, or when get() mid-run must reload from disk. | ||
| let v = PayerKeyVault::new().with_sealed(cfg); | ||
| assert_eq!(v.get("onlyseal").as_deref(), Some("sk_from_disk")); | ||
| assert!(v.refresh("onlyseal")); | ||
| // Clear RAM and prove seal-only hydrate still works. | ||
| if let Ok(mut g) = v.inner.lock() { | ||
| g.clear(); | ||
| } | ||
| assert_eq!(v.get("onlyseal").as_deref(), Some("sk_from_disk")); | ||
| } | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Align the retry guidance with the implemented credential paths.
post_retrydoes not readX-Miner-Hotkeyor a body hotkey. It reusesrow.miner_hotkeyfrom storage. It also accepts a sealed payer-vault entry instead of requiringX-Lium-Api-Key. Update this message to describe the stored miner hotkey and theX-Lium-Api-Keyor sealed-vault alternatives, or add validation for the documented inputs.Proposed wording
📝 Committable suggestion
🤖 Prompt for AI Agents