diff --git a/.github/workflows/solana-build-anchor-programs.yml b/.github/workflows/solana-build-anchor-programs.yml index 0bafeaedb..1662a847b 100644 --- a/.github/workflows/solana-build-anchor-programs.yml +++ b/.github/workflows/solana-build-anchor-programs.yml @@ -19,7 +19,7 @@ jobs: path: validator-image.tar.gz # FIXME: this key will make the cache succeed even if some other nix files which affect it change. # we should *really* build this with nix instead of a manual hashFiles / cache here - key: validator-image-${{ runner.os }}-${{ hashFiles('shared/coordinator/src/coordinator.rs', 'architectures/decentralized/solana-coordinator/**/*.rs', 'architectures/decentralized/solana-coordinator/**/*.toml', 'architectures/decentralized/solana-coordinator/Cargo.lock', 'architectures/decentralized/solana-authorizer/**/*.rs', 'architectures/decentralized/solana-authorizer/**/*.toml', 'architectures/decentralized/solana-authorizer/Cargo.lock', 'docker/test/psyche_solana_validator_entrypoint.sh', 'nix/docker.nix', 'flake.lock') }} + key: validator-image-${{ runner.os }}-${{ hashFiles('shared/coordinator/src/coordinator.rs', 'shared/coordinator/src/model.rs', 'architectures/decentralized/solana-coordinator/**/*.rs', 'architectures/decentralized/solana-coordinator/**/*.toml', 'architectures/decentralized/solana-coordinator/Cargo.lock', 'architectures/decentralized/solana-authorizer/**/*.rs', 'architectures/decentralized/solana-authorizer/**/*.toml', 'architectures/decentralized/solana-authorizer/Cargo.lock', 'docker/test/psyche_solana_validator_entrypoint.sh', 'nix/docker.nix', 'flake.lock') }} lookup-only: true # Build validator image if cache fails @@ -89,7 +89,7 @@ jobs: uses: actions/cache/save@v4 with: path: validator-image.tar.gz - key: validator-image-${{ runner.os }}-${{ hashFiles('shared/coordinator/src/coordinator.rs', 'architectures/decentralized/solana-coordinator/**/*.rs', 'architectures/decentralized/solana-coordinator/**/*.toml', 'architectures/decentralized/solana-coordinator/Cargo.lock', 'architectures/decentralized/solana-authorizer/**/*.rs', 'architectures/decentralized/solana-authorizer/**/*.toml', 'architectures/decentralized/solana-authorizer/Cargo.lock', 'docker/test/psyche_solana_validator_entrypoint.sh', 'nix/docker.nix', 'flake.lock') }} + key: validator-image-${{ runner.os }}-${{ hashFiles('shared/coordinator/src/coordinator.rs', 'shared/coordinator/src/model.rs', 'architectures/decentralized/solana-coordinator/**/*.rs', 'architectures/decentralized/solana-coordinator/**/*.toml', 'architectures/decentralized/solana-coordinator/Cargo.lock', 'architectures/decentralized/solana-authorizer/**/*.rs', 'architectures/decentralized/solana-authorizer/**/*.toml', 'architectures/decentralized/solana-authorizer/Cargo.lock', 'docker/test/psyche_solana_validator_entrypoint.sh', 'nix/docker.nix', 'flake.lock') }} - name: Build complete if: steps.cache-validator.outputs.cache-hit != 'true' diff --git a/.github/workflows/solana-integration-test-base.yml b/.github/workflows/solana-integration-test-base.yml index 9284fda23..63270a336 100644 --- a/.github/workflows/solana-integration-test-base.yml +++ b/.github/workflows/solana-integration-test-base.yml @@ -46,7 +46,7 @@ jobs: uses: actions/cache/restore@v4 with: path: validator-image.tar.gz - key: validator-image-${{ runner.os }}-${{ hashFiles('shared/coordinator/src/coordinator.rs', 'architectures/decentralized/solana-coordinator/**/*.rs', 'architectures/decentralized/solana-coordinator/**/*.toml', 'architectures/decentralized/solana-coordinator/Cargo.lock', 'architectures/decentralized/solana-authorizer/**/*.rs', 'architectures/decentralized/solana-authorizer/**/*.toml', 'architectures/decentralized/solana-authorizer/Cargo.lock', 'docker/test/psyche_solana_validator_entrypoint.sh', 'nix/docker.nix', 'flake.lock') }} + key: validator-image-${{ runner.os }}-${{ hashFiles('shared/coordinator/src/coordinator.rs', 'shared/coordinator/src/model.rs', 'architectures/decentralized/solana-coordinator/**/*.rs', 'architectures/decentralized/solana-coordinator/**/*.toml', 'architectures/decentralized/solana-coordinator/Cargo.lock', 'architectures/decentralized/solana-authorizer/**/*.rs', 'architectures/decentralized/solana-authorizer/**/*.toml', 'architectures/decentralized/solana-authorizer/Cargo.lock', 'docker/test/psyche_solana_validator_entrypoint.sh', 'nix/docker.nix', 'flake.lock') }} fail-on-cache-miss: true - name: Load Validator Image diff --git a/Cargo.lock b/Cargo.lock index 405b53793..d11277d3e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -36,7 +36,7 @@ checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" dependencies = [ "cfg-if", "cipher", - "cpufeatures", + "cpufeatures 0.2.17", ] [[package]] @@ -753,28 +753,6 @@ dependencies = [ "pin-project-lite", ] -[[package]] -name = "async-stream" -version = "0.3.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" -dependencies = [ - "async-stream-impl", - "futures-core", - "pin-project-lite", -] - -[[package]] -name = "async-stream-impl" -version = "0.3.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" -dependencies = [ - "proc-macro2", - "quote", - "syn 2.0.115", -] - [[package]] name = "async-trait" version = "0.1.89" @@ -850,6 +828,28 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "aws-lc-rs" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9a7b350e3bb1767102698302bc37256cbd48422809984b98d292c40e2579aa9" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.37.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b092fe214090261288111db7a2b2c2118e5a7f30dc2569f1732c4069a6840549" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", +] + [[package]] name = "axum" version = "0.7.9" @@ -1058,7 +1058,7 @@ dependencies = [ "cc", "cfg-if", "constant_time_eq 0.4.2", - "cpufeatures", + "cpufeatures 0.2.17", "digest 0.10.7", ] @@ -1467,6 +1467,17 @@ dependencies = [ "syn 2.0.115", ] +[[package]] +name = "chacha20" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.0", +] + [[package]] name = "chrono" version = "0.4.43" @@ -1603,6 +1614,15 @@ dependencies = [ "roff", ] +[[package]] +name = "cmake" +version = "0.1.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75443c44cd6b379beb8c5b45d85d0773baf31cce901fe7bb252f4eff3008ef7d" +dependencies = [ + "cc", +] + [[package]] name = "cobs" version = "0.3.0" @@ -1723,12 +1743,6 @@ dependencies = [ "web-sys", ] -[[package]] -name = "const-oid" -version = "0.9.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" - [[package]] name = "const-oid" version = "0.10.2" @@ -1883,6 +1897,24 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32c" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a47af21622d091a8f0fb295b88bc886ac74efcc613efc19f5d0b21de5c89e47" +dependencies = [ + "rustc_version", +] + [[package]] name = "crc32fast" version = "1.5.0" @@ -2023,7 +2055,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "curve25519-dalek-derive", "digest 0.10.7", "fiat-crypto 0.2.9", @@ -2041,7 +2073,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6f9200d1d13637f15a6acb71e758f64624048d85b31a5fdbfd8eca1e2687d0b7" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "curve25519-dalek-derive", "digest 0.11.0-rc.10", "fiat-crypto 0.3.0", @@ -2187,25 +2219,14 @@ version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" -[[package]] -name = "der" -version = "0.7.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" -dependencies = [ - "const-oid 0.9.6", - "pem-rfc7468 0.7.0", - "zeroize", -] - [[package]] name = "der" version = "0.8.0-rc.12" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c0182be35043efdd2df327a443bb600606e350cfb090cccb233e9451e76f5a3" dependencies = [ - "const-oid 0.10.2", - "pem-rfc7468 1.0.0", + "const-oid", + "pem-rfc7468", "zeroize", ] @@ -2363,7 +2384,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "afa94b64bfc6549e6e4b5a3216f22593224174083da7a90db47e951c4fb31725" dependencies = [ "block-buffer 0.11.0", - "const-oid 0.10.2", + "const-oid", "crypto-common 0.2.0", ] @@ -2488,6 +2509,12 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1435fa1053d8b2fbbe9be7e97eca7f33d37b28409959813daefc1446a14247f1" +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + [[package]] name = "dwrote" version = "0.11.5" @@ -2527,7 +2554,7 @@ version = "3.0.0-rc.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6e914c7c52decb085cea910552e24c63ac019e3ab8bf001ff736da9a9d9d890" dependencies = [ - "pkcs8 0.11.0-rc.11", + "pkcs8", "serde", "signature 3.0.0-rc.10", ] @@ -3035,6 +3062,12 @@ dependencies = [ "pkg-config", ] +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + [[package]] name = "funty" version = "2.0.0" @@ -3328,6 +3361,7 @@ dependencies = [ "cfg-if", "libc", "r-efi", + "rand_core 0.10.0", "wasip2", "wasip3", ] @@ -3389,9 +3423,9 @@ dependencies = [ [[package]] name = "google-cloud-auth" -version = "0.17.2" +version = "0.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e57a13fbacc5e9c41ded3ad8d0373175a6b7a6ad430d99e89d314ac121b7ab06" +checksum = "1112c453c2e155b3e683204ffff52bcc6d6495d04b68d9e90cd24161270c5058" dependencies = [ "async-trait", "base64 0.21.7", @@ -3409,6 +3443,134 @@ dependencies = [ "urlencoding", ] +[[package]] +name = "google-cloud-auth" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36ad774d41426ab205eeec577540f209a5485c366814dd5c89a7e3018fe84e7c" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "google-cloud-gax", + "http 1.4.0", + "reqwest 0.13.2", + "rustc_version", + "rustls 0.23.36", + "rustls-pki-types", + "serde", + "serde_json", + "thiserror 2.0.18", + "time", + "tokio", +] + +[[package]] +name = "google-cloud-gax" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2973715fe664ecb0d883926c8b5f66cb9d52a44add1d0be1cad1907d832bf0af" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures", + "google-cloud-rpc", + "google-cloud-wkt", + "http 1.4.0", + "pin-project", + "rand 0.10.0", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", +] + +[[package]] +name = "google-cloud-gax-internal" +version = "0.7.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "598e5ffec2c1c9b43e83847b2badc0f128a03e541b75f1ecd8acf0a2605c40cf" +dependencies = [ + "bytes", + "futures", + "google-cloud-auth 1.6.0", + "google-cloud-gax", + "google-cloud-rpc", + "google-cloud-wkt", + "http 1.4.0", + "http-body 1.0.1", + "http-body-util", + "hyper 1.8.1", + "opentelemetry-semantic-conventions", + "percent-encoding", + "pin-project", + "prost 0.14.3", + "prost-types", + "reqwest 0.13.2", + "rustc_version", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tokio-stream", + "tonic 0.14.5", + "tonic-prost", + "tower 0.5.3", + "tracing", +] + +[[package]] +name = "google-cloud-iam-v1" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b41025359d1f52c24966c24d5ecb7f91198ded13c9cbf57c2d2e735dae93c814" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-type", + "google-cloud-wkt", + "lazy_static", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-longrunning" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daafe6c2859976ca51571db6c7bbf96bdce0122e81e6854a6c62554c4f63bae8" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-rpc", + "google-cloud-wkt", + "lazy_static", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-lro" +version = "1.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef3969c85cf6b163c655f7ddcdee5bb5c3241f680ab648951239ef9cd4ffb0f" +dependencies = [ + "google-cloud-gax", + "google-cloud-longrunning", + "google-cloud-rpc", + "google-cloud-wkt", + "serde", + "tokio", +] + [[package]] name = "google-cloud-metadata" version = "0.5.1" @@ -3420,37 +3582,60 @@ dependencies = [ "tokio", ] +[[package]] +name = "google-cloud-rpc" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd10e97751ca894f9dad6be69fcef1cb72f5bc187329e0254817778fc8235030" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + [[package]] name = "google-cloud-storage" -version = "0.24.0" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34a73d9e94d35665909050f02e035d8bdc82e419241b1b027ebf1ea51dc8a470" +checksum = "3ce225eccfece751251e07ac787199cfc53bc2b5d60ac0570f58404507ceb309" dependencies = [ - "anyhow", - "async-stream", "async-trait", - "base64 0.21.7", + "base64 0.22.1", "bytes", - "futures-util", - "google-cloud-auth", - "google-cloud-metadata", - "google-cloud-token", + "chrono", + "crc32c", + "futures", + "google-cloud-auth 1.6.0", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-iam-v1", + "google-cloud-longrunning", + "google-cloud-lro", + "google-cloud-rpc", + "google-cloud-type", + "google-cloud-wkt", "hex", - "once_cell", + "http 1.4.0", + "http-body 1.0.1", + "hyper 1.8.1", + "lazy_static", + "md5", "percent-encoding", - "pkcs8 0.10.2", - "regex", - "reqwest 0.12.28", - "reqwest-middleware 0.4.2", - "ring", + "pin-project", + "prost 0.14.3", + "prost-types", "serde", "serde_json", + "serde_with", "sha2 0.10.9", - "thiserror 1.0.69", - "time", + "thiserror 2.0.18", "tokio", + "tokio-stream", "tracing", "url", + "uuid", ] [[package]] @@ -3462,6 +3647,35 @@ dependencies = [ "async-trait", ] +[[package]] +name = "google-cloud-type" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9390ac2f3f9882ff42956b25ea65b9f546c8dd44c131726d75a96bf744ec75f6" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-wkt" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0ade65b0e4fa9cb4b6f147c8e726803bff453e3190910a53cbd3b0c019f5c2a" +dependencies = [ + "base64 0.22.1", + "bytes", + "serde", + "serde_json", + "serde_with", + "thiserror 2.0.18", + "time", + "url", +] + [[package]] name = "governor" version = "0.6.3" @@ -4033,6 +4247,19 @@ dependencies = [ "webpki-roots 1.0.6", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper 1.8.1", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-tls" version = "0.6.0" @@ -4466,7 +4693,7 @@ dependencies = [ "papaya", "pin-project", "pkarr", - "pkcs8 0.11.0-rc.11", + "pkcs8", "portmapper", "rand 0.9.2", "reqwest 0.12.28", @@ -5021,7 +5248,7 @@ version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ecc2af9a1119c51f12a14607e783cb977bde58bc069ff0c3da1095e635d70654" dependencies = [ - "cpufeatures", + "cpufeatures 0.2.17", ] [[package]] @@ -5422,6 +5649,12 @@ dependencies = [ "syn 2.0.115", ] +[[package]] +name = "md5" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae960838283323069879657ca3de837e9f7bbb4c7bf6ea7f1b290d5e9476d2e0" + [[package]] name = "memchr" version = "2.8.0" @@ -6353,7 +6586,7 @@ dependencies = [ "opentelemetry-http", "opentelemetry-proto", "opentelemetry_sdk", - "prost", + "prost 0.13.5", "reqwest 0.12.28", "thiserror 2.0.18", "tracing", @@ -6367,10 +6600,16 @@ checksum = "56f8870d3024727e99212eb3bb1762ec16e255e3e6f58eeb3dc8db1aa226746d" dependencies = [ "opentelemetry 0.28.0", "opentelemetry_sdk", - "prost", - "tonic", + "prost 0.13.5", + "tonic 0.12.3", ] +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.31.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e62e29dfe041afb8ed2a6c9737ab57db4907285d999ef8ad3a59092a36bdc846" + [[package]] name = "opentelemetry_sdk" version = "0.28.0" @@ -6543,15 +6782,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "pem-rfc7468" -version = "0.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412" -dependencies = [ - "base64ct", -] - [[package]] name = "pem-rfc7468" version = "1.0.0" @@ -6649,24 +6879,14 @@ dependencies = [ "wasm-bindgen-futures", ] -[[package]] -name = "pkcs8" -version = "0.10.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" -dependencies = [ - "der 0.7.10", - "spki 0.7.3", -] - [[package]] name = "pkcs8" version = "0.11.0-rc.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "12922b6296c06eb741b02d7b5161e3aaa22864af38dfa025a1a3ba3f68c84577" dependencies = [ - "der 0.8.0-rc.12", - "spki 0.8.0-rc.4", + "der", + "spki", ] [[package]] @@ -6754,7 +6974,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "opaque-debug", "universal-hash", ] @@ -7016,7 +7236,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" dependencies = [ "bytes", - "prost-derive", + "prost-derive 0.13.5", +] + +[[package]] +name = "prost" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" +dependencies = [ + "bytes", + "prost-derive 0.14.3", ] [[package]] @@ -7032,6 +7262,28 @@ dependencies = [ "syn 2.0.115", ] +[[package]] +name = "prost-derive" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn 2.0.115", +] + +[[package]] +name = "prost-types" +version = "0.14.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" +dependencies = [ + "prost 0.14.3", +] + [[package]] name = "psyche-centralized-client" version = "0.2.0" @@ -7042,7 +7294,6 @@ dependencies = [ "clap", "clap-markdown", "hex", - "hf-hub", "psyche-centralized-shared", "psyche-client", "psyche-coordinator", @@ -7055,6 +7306,7 @@ dependencies = [ "psyche-tui", "psyche-watcher", "rand 0.9.2", + "rustls 0.23.36", "tikv-jemallocator", "time", "tokio", @@ -7156,6 +7408,7 @@ dependencies = [ "anyhow", "clap", "futures", + "google-cloud-storage", "hex", "hf-hub", "iroh-blobs", @@ -7191,6 +7444,7 @@ name = "psyche-coordinator" version = "0.2.0" dependencies = [ "anchor-lang", + "anyhow", "async-trait", "bytemuck", "cfg_eval", @@ -7228,10 +7482,14 @@ dependencies = [ "anchor-lang", "anyhow", "async-trait", + "bs58", "bytemuck", + "bytes", "chrono", "clap", "futures", + "google-cloud-auth 0.16.0", + "google-cloud-gax", "google-cloud-storage", "hf-hub", "memmap2 0.9.9", @@ -7258,6 +7516,7 @@ dependencies = [ "tokio-util 0.7.18", "tracing", "ts-rs", + "urlencoding", ] [[package]] @@ -7353,6 +7612,7 @@ dependencies = [ "psyche-metrics", "psyche-network", "pyo3", + "reqwest 0.12.28", "serde", "serde_json", "tikv-jemallocator", @@ -7522,6 +7782,7 @@ dependencies = [ "psyche-watcher", "rand 0.9.2", "rand_chacha 0.9.0", + "rustls 0.23.36", "serde", "serde_json", "tikv-jemallocator", @@ -7796,6 +8057,7 @@ version = "0.11.13" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1906b49b0c3bc04b5fe5d86a77925ae6524a19b816ae38ce1e426255f1d8a31" dependencies = [ + "aws-lc-rs", "bytes", "fastbloom", "getrandom 0.3.4", @@ -7882,6 +8144,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc266eb313df6c5c09c1c7b1fbe2510961e5bcd3add930c1e31f7ed9da0feff8" +dependencies = [ + "chacha20", + "getrandom 0.4.1", + "rand_core 0.10.0", +] + [[package]] name = "rand_chacha" version = "0.2.2" @@ -7939,6 +8212,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c8d0fd677905edcbeedbf2edb6494d676f0e98d54d5cf9bda0b061cb8fb8aba" + [[package]] name = "rand_hc" version = "0.2.0" @@ -8252,7 +8531,6 @@ dependencies = [ "js-sys", "log", "mime", - "mime_guess", "native-tls", "percent-encoding", "pin-project-lite", @@ -8274,39 +8552,66 @@ dependencies = [ "url", "wasm-bindgen", "wasm-bindgen-futures", - "wasm-streams", + "wasm-streams 0.4.2", "web-sys", "webpki-roots 1.0.6", ] [[package]] -name = "reqwest-middleware" -version = "0.2.5" +name = "reqwest" +version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a735987236a8e238bf0296c7e351b999c188ccc11477f311b82b55c93984216" +checksum = "ab3f43e3283ab1488b624b44b0e988d0acea0b3214e694730a055cb6b2efa801" dependencies = [ - "anyhow", - "async-trait", - "http 0.2.12", - "reqwest 0.11.27", + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "http 1.4.0", + "http-body 1.0.1", + "http-body-util", + "hyper 1.8.1", + "hyper-rustls 0.27.7", + "hyper-util", + "js-sys", + "log", + "mime_guess", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls 0.23.36", + "rustls-pki-types", + "rustls-platform-verifier", "serde", - "task-local-extensions", - "thiserror 1.0.69", + "serde_json", + "serde_urlencoded", + "sync_wrapper 1.0.2", + "tokio", + "tokio-rustls 0.26.4", + "tokio-util 0.7.18", + "tower 0.5.3", + "tower-http 0.6.8", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams 0.5.0", + "web-sys", ] [[package]] name = "reqwest-middleware" -version = "0.4.2" +version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57f17d28a6e6acfe1733fe24bcd30774d13bffa4b8a22535b4c8c98423088d4e" +checksum = "5a735987236a8e238bf0296c7e351b999c188ccc11477f311b82b55c93984216" dependencies = [ "anyhow", "async-trait", - "http 1.4.0", - "reqwest 0.12.28", + "http 0.2.12", + "reqwest 0.11.27", "serde", + "task-local-extensions", "thiserror 1.0.69", - "tower-service", ] [[package]] @@ -8476,6 +8781,7 @@ version = "0.23.36" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c665f33d38cea657d9614f766881e4d510e0eda4239891eea56b4cadcf01801b" dependencies = [ + "aws-lc-rs", "log", "once_cell", "ring", @@ -8568,6 +8874,7 @@ version = "0.103.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7df23109aa6c1567d1c575b9952556388da57401e4ace1d15f79eedad0d8f53" dependencies = [ + "aws-lc-rs", "ring", "rustls-pki-types", "untrusted", @@ -8966,7 +9273,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest 0.10.7", ] @@ -8984,7 +9291,7 @@ checksum = "4d58a1e1bf39749807d89cf2d98ac2dfa0ff1cb3faa38fbb64dd88ac8013d800" dependencies = [ "block-buffer 0.9.0", "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest 0.9.0", "opaque-debug", ] @@ -8996,7 +9303,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest 0.10.7", ] @@ -9007,7 +9314,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d1e3878ab0f98e35b2df35fe53201d088299b41a6bb63e3e34dada2ac4abd924" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.2.17", "digest 0.11.0-rc.10", ] @@ -10293,7 +10600,7 @@ dependencies = [ "indicatif", "log", "reqwest 0.11.27", - "reqwest-middleware 0.2.5", + "reqwest-middleware", "semver", "serde", "serde_derive", @@ -10318,7 +10625,7 @@ dependencies = [ "bs58", "jsonrpc-core", "reqwest 0.11.27", - "reqwest-middleware 0.2.5", + "reqwest-middleware", "semver", "serde", "serde_derive", @@ -11173,16 +11480,6 @@ dependencies = [ "lock_api", ] -[[package]] -name = "spki" -version = "0.7.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" -dependencies = [ - "base64ct", - "der 0.7.10", -] - [[package]] name = "spki" version = "0.8.0-rc.4" @@ -11190,7 +11487,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8baeff88f34ed0691978ec34440140e1572b68c7dd4a495fd14a3dc1944daa80" dependencies = [ "base64ct", - "der 0.8.0-rc.12", + "der", ] [[package]] @@ -12563,13 +12860,51 @@ dependencies = [ "http-body-util", "percent-encoding", "pin-project", - "prost", + "prost 0.13.5", "tokio-stream", "tower-layer", "tower-service", "tracing", ] +[[package]] +name = "tonic" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" +dependencies = [ + "base64 0.22.1", + "bytes", + "http 1.4.0", + "http-body 1.0.1", + "http-body-util", + "hyper 1.8.1", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "rustls-native-certs", + "sync_wrapper 1.0.2", + "tokio", + "tokio-rustls 0.26.4", + "tokio-stream", + "tower 0.5.3", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" +dependencies = [ + "bytes", + "prost 0.14.3", + "tonic 0.14.5", +] + [[package]] name = "torch-sys" version = "0.22.0" @@ -12600,9 +12935,12 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap 2.13.0", "pin-project-lite", + "slab", "sync_wrapper 1.0.2", "tokio", + "tokio-util 0.7.18", "tower-layer", "tower-service", "tracing", @@ -13294,6 +13632,19 @@ dependencies = [ "web-sys", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" diff --git a/Cargo.toml b/Cargo.toml index a63967dbe..6864cc200 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -61,6 +61,7 @@ serde_json = "1.0.28" tracing = { version = "0.1.40", features = ["std"] } bytes = "1.7.1" rand = "0.9.2" +rustls = { version = "0.23", default-features = false, features = ["ring"] } rand_chacha = "0.9.0" postcard = { version = "1.0.10", features = ["use-std"] } sha2 = "0.10.8" @@ -82,6 +83,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/NousResearch/tch-rs.git", rev = "dda507e05a776547a112b6854d1e611684f8c729" } torch-sys = { git = "https://github.com/NousResearch/tch-rs.git", rev = "dda507e05a776547a112b6854d1e611684f8c729" } pyo3-tch = { git = "https://github.com/NousResearch/tch-rs.git", rev = "dda507e05a776547a112b6854d1e611684f8c729" } diff --git a/architectures/centralized/client/Cargo.toml b/architectures/centralized/client/Cargo.toml index 70b5fc075..065e22595 100644 --- a/architectures/centralized/client/Cargo.toml +++ b/architectures/centralized/client/Cargo.toml @@ -21,13 +21,13 @@ clap.workspace = true tokio.workspace = true tracing.workspace = true tokio-util.workspace = true -hf-hub.workspace = true time.workspace = true bytemuck.workspace = true clap-markdown.workspace = true hex = "0.4.3" psyche-python-extension-impl = { workspace = true, optional = true } tikv-jemallocator.workspace = true +rustls.workspace = true [features] parallelism = ["psyche-client/parallelism"] diff --git a/architectures/centralized/client/src/app.rs b/architectures/centralized/client/src/app.rs index cd9ed42db..630627cf8 100644 --- a/architectures/centralized/client/src/app.rs +++ b/architectures/centralized/client/src/app.rs @@ -1,13 +1,12 @@ use anyhow::{Error, Result}; use bytemuck::Zeroable; -use hf_hub::Repo; use psyche_centralized_shared::{ClientToServerMessage, ServerToClientMessage}; -use psyche_client::HubUploadInfo; -use psyche_client::UploadInfo; use psyche_client::{ - Client, ClientTUI, ClientTUIState, NC, RunInitConfig, TrainArgs, read_identity_secret_key, + CheckpointUploader, Client, ClientTUI, ClientTUIState, NC, RunInitConfig, TrainArgs, + read_identity_secret_key, }; -use psyche_coordinator::{Coordinator, HealthChecks, model}; +use psyche_coordinator::model::{self, Checkpoint}; +use psyche_coordinator::{Coordinator, HealthChecks}; use psyche_core::NodeIdentity; use psyche_metrics::ClientMetrics; use psyche_network::{EndpointId, NetworkTUIState, NetworkTui, SecretKey, TcpClient, allowlist}; @@ -29,7 +28,7 @@ pub type TabsData = ::Data; pub enum ToSend { Witness(Box), HealthCheck(HealthChecks), - Checkpoint(model::Checkpoint), + Checkpoint(Checkpoint), } struct Backend { @@ -67,7 +66,7 @@ impl WatcherBackend 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(()) } @@ -167,33 +166,55 @@ impl App { p2p: NC, state_options: RunInitConfig, ) -> 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 { - anyhow::bail!( - "Checkpoint upload repo {} is not writable with the passed API key.", - hub_repo - ) - } - } - } - self.server_conn .send(ClientToServerMessage::Join { run_id: self.run_id.clone(), }) .await?; + // Wait for the first coordinator state to validate upload credentials with repo info. + let first_coordinator_state = select! { + _ = self.cancel.cancelled() => { + return Ok(()); + } + message = self.server_conn.receive() => { + match message? { + ServerToClientMessage::Coordinator(state) => { + self.coordinator_state = *state; + *state + } + } + } + }; + + // Validate upload credentials now that we have the coordinator state with checkpoint info. + if !state_options.checkpoint_config.skip_upload { + let model::Model::LLM(model::LLM { checkpoint, .. }) = first_coordinator_state.model; + match checkpoint { + Checkpoint::Hub(ref hub_repo) | Checkpoint::P2P(ref hub_repo) => { + let token = state_options.checkpoint_config.hub_token.as_ref() + .ok_or_else(|| anyhow::anyhow!( + "No HF_TOKEN found for checkpointing to Hub repo {}. Set HF_TOKEN environment variable.", + hub_repo.repo_id + ))?; + // Validate write permissions — the uploader is dropped after validation, + // it will be re-created during cooldown with the current coordinator state. + CheckpointUploader::new_hub(hub_repo.repo_id.to_string(), token.clone()) + .await?; + } + Checkpoint::Gcs(_) | Checkpoint::P2PGcs(_) => { + // GCS uploads use run-down signed URLs; auth is validated at request time. + if state_options.checkpoint_config.run_down_client.is_none() { + anyhow::bail!("RunDownClient not configured for GCS checkpoint upload"); + } + } + _ => {} + } + } + let (tx_from_server_message, rx_from_server_message) = mpsc::unbounded_channel(); + // Forward the first coordinator state we already received. + tx_from_server_message.send(first_coordinator_state)?; let (tx_to_server_message, mut rx_to_server_message) = mpsc::unbounded_channel(); let mut client = Client::new( Backend { diff --git a/architectures/centralized/client/src/main.rs b/architectures/centralized/client/src/main.rs index bfc381c2d..1e7cecdf1 100644 --- a/architectures/centralized/client/src/main.rs +++ b/architectures/centralized/client/src/main.rs @@ -129,6 +129,10 @@ async fn async_main() -> Result<()> { } fn main() -> Result<()> { + rustls::crypto::ring::default_provider() + .install_default() + .expect("Failed to install default CryptoProvider"); + #[cfg(feature = "python")] psyche_python_extension_impl::init_embedded_python()?; diff --git a/architectures/centralized/server/src/app.rs b/architectures/centralized/server/src/app.rs index 770f62944..b4897b000 100644 --- a/architectures/centralized/server/src/app.rs +++ b/architectures/centralized/server/src/app.rs @@ -1,7 +1,7 @@ use anyhow::{Result, anyhow, bail}; use async_trait::async_trait; use psyche_centralized_shared::{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 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"); } } @@ -403,6 +403,10 @@ impl App { Self::get_timestamp(), rand::rng().next_u64(), ), + OpportunisticData::CooldownStep(witness) => self.coordinator.cooldown_witness( + &NodeIdentity::from_single_key(*from.as_bytes()), + witness, + ), } { warn!("Error when processing witness: {error}"); }; diff --git a/architectures/centralized/testing/src/lib.rs b/architectures/centralized/testing/src/lib.rs index dfa92e11a..6c067c639 100644 --- a/architectures/centralized/testing/src/lib.rs +++ b/architectures/centralized/testing/src/lib.rs @@ -6,4 +6,4 @@ pub mod test_utils; pub const WARMUP_TIME: u64 = 20; 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; diff --git a/architectures/centralized/testing/src/test_utils.rs b/architectures/centralized/testing/src/test_utils.rs index 1ec700b4e..ae835da36 100644 --- a/architectures/centralized/testing/src/test_utils.rs +++ b/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, } diff --git a/architectures/decentralized/justfile b/architectures/decentralized/justfile index e1eddf893..ad174cf39 100644 --- a/architectures/decentralized/justfile +++ b/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/setup-and-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" diff --git a/architectures/decentralized/solana-client/Cargo.toml b/architectures/decentralized/solana-client/Cargo.toml index 6e19d5cf3..940dc2e5f 100644 --- a/architectures/decentralized/solana-client/Cargo.toml +++ b/architectures/decentralized/solana-client/Cargo.toml @@ -31,6 +31,7 @@ tokio-util.workspace = true tracing.workspace = true psyche-python-extension-impl = { workspace = true, optional = true } tikv-jemallocator.workspace = true +rustls.workspace = true [features] parallelism = ["psyche-client/parallelism"] diff --git a/architectures/decentralized/solana-client/src/app.rs b/architectures/decentralized/solana-client/src/app.rs index c541d5e74..cf46c6d38 100644 --- a/architectures/decentralized/solana-client/src/app.rs +++ b/architectures/decentralized/solana-client/src/app.rs @@ -10,10 +10,16 @@ use anchor_client::{ }; use anyhow::{Result, anyhow}; use psyche_client::{ - Client, ClientTUI, ClientTUIState, NC, RunInitConfig, TrainArgs, read_identity_secret_key, + CheckpointUploader, Client, ClientTUI, ClientTUIState, NC, RunInitConfig, TrainArgs, + read_identity_secret_key, }; -use psyche_coordinator::{ClientState, Coordinator, CoordinatorError, RunState}; +use psyche_coordinator::{ + ClientState, Coordinator, CoordinatorError, RunState, + model::{self, LLM, Model}, +}; + use psyche_core::sha256; +use psyche_data_provider::RunDownClient; use psyche_metrics::ClientMetrics; use psyche_network::{DiscoveryMode, NetworkTUIState, NetworkTui, SecretKey, allowlist}; @@ -91,7 +97,17 @@ pub async fn build_app( let eval_tasks = p.eval_tasks()?; let hub_read_token = std::env::var("HF_TOKEN").ok(); - let checkpoint_config = p.checkpoint_config()?; + let mut checkpoint_config = p.checkpoint_config()?; + + // Construct RunDownClient using the wallet keypair for signing. + // This enables GCS checkpoint upload/download via run-down signed URLs. + let run_down_keypair = wallet_keypair.clone(); + let run_down_client = Arc::new(RunDownClient::new( + p.run_id.clone(), + wallet_keypair.pubkey().to_string(), + move |msg| run_down_keypair.sign_message(msg).as_ref().to_vec(), + )); + checkpoint_config.run_down_client = Some(run_down_client); let solana_pubkey = wallet_keypair.pubkey(); let wandb_info = p.wandb_info(format!("{}-{solana_pubkey}", p.run_id))?; @@ -230,6 +246,34 @@ impl App { let mut joined_run_this_epoch = None; let mut ever_joined_run = false; + // sanity checks — skip credential validation when checkpoint upload is disabled + if !self.state_options.checkpoint_config.skip_upload { + let Model::LLM(LLM { checkpoint, .. }) = start_coordinator_state.model; + match checkpoint { + model::Checkpoint::Hub(ref hub_repo) | model::Checkpoint::P2P(ref hub_repo) => { + let token = self.state_options.checkpoint_config.hub_token.as_ref() + .ok_or_else(|| anyhow!( + "No HF_TOKEN found for checkpointing to Hub repo {}. Set HF_TOKEN environment variable.", + hub_repo.repo_id + ))?; + CheckpointUploader::new_hub(hub_repo.repo_id.to_string(), token.clone()) + .await?; + } + model::Checkpoint::Gcs(_) | model::Checkpoint::P2PGcs(_) => { + // GCS uploads use run-down signed URLs; auth is validated at request time. + if self + .state_options + .checkpoint_config + .run_down_client + .is_none() + { + anyhow::bail!("RunDownClient not configured for GCS checkpoint upload"); + } + } + _ => {} + } + } + // 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 diff --git a/architectures/decentralized/solana-client/src/main.rs b/architectures/decentralized/solana-client/src/main.rs index 73eb7f019..352427111 100644 --- a/architectures/decentralized/solana-client/src/main.rs +++ b/architectures/decentralized/solana-client/src/main.rs @@ -344,6 +344,10 @@ async fn async_main() -> Result<()> { } fn main() -> Result<()> { + rustls::crypto::ring::default_provider() + .install_default() + .expect("Failed to install default CryptoProvider"); + #[cfg(feature = "python")] psyche_python_extension_impl::init_embedded_python()?; diff --git a/architectures/decentralized/solana-common/src/backend.rs b/architectures/decentralized/solana-common/src/backend.rs index bd66d06b7..89f60e2f7 100644 --- a/architectures/decentralized/solana-common/src/backend.rs +++ b/architectures/decentralized/solana-common/src/backend.rs @@ -20,7 +20,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}; @@ -315,6 +315,12 @@ impl SolanaBackend { &user, witness, ), + OpportunisticData::CooldownStep(witness) => instructions::coordinator_cooldown_witness( + &coordinator_instance, + &coordinator_account, + &user, + witness, + ), }; self.spawn_scheduled_send("Witness", &[instruction], &[]); } @@ -650,7 +656,7 @@ impl WatcherBackend 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(()) diff --git a/architectures/decentralized/solana-common/src/instructions.rs b/architectures/decentralized/solana-common/src/instructions.rs index c6c0e28cc..54a8b4a12 100644 --- a/architectures/decentralized/solana-common/src/instructions.rs +++ b/architectures/decentralized/solana-common/src/instructions.rs @@ -179,6 +179,28 @@ pub fn coordinator_warmup_witness( ) } +pub fn coordinator_cooldown_witness( + coordinator_instance: &Pubkey, + coordinator_account: &Pubkey, + user: &Pubkey, + witness: psyche_coordinator::Witness, +) -> Instruction { + anchor_instruction( + psyche_solana_coordinator::ID, + psyche_solana_coordinator::accounts::PermissionlessCoordinatorAccounts { + user: *user, + coordinator_instance: *coordinator_instance, + coordinator_account: *coordinator_account, + }, + psyche_solana_coordinator::instruction::CooldownWitness { + proof: witness.proof, + participant_bloom: witness.participant_bloom, + broadcast_bloom: witness.broadcast_bloom, + broadcast_merkle: witness.broadcast_merkle, + }, + ) +} + pub fn coordinator_health_check( coordinator_instance: &Pubkey, coordinator_account: &Pubkey, diff --git a/architectures/decentralized/solana-coordinator/Cargo.lock b/architectures/decentralized/solana-coordinator/Cargo.lock index b63749b7b..5c316632e 100644 --- a/architectures/decentralized/solana-coordinator/Cargo.lock +++ b/architectures/decentralized/solana-coordinator/Cargo.lock @@ -1605,6 +1605,7 @@ name = "psyche-coordinator" version = "0.2.0" dependencies = [ "anchor-lang", + "anyhow", "async-trait", "bytemuck", "cfg_eval", diff --git a/architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/instance_state.rs b/architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/instance_state.rs index c4010acb7..d092d9353 100644 --- a/architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/instance_state.rs +++ b/architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/instance_state.rs @@ -236,6 +236,20 @@ impl CoordinatorInstanceState { self.tick() } + pub fn cooldown_witness( + &mut self, + payer: &Pubkey, + witness: Witness, + ) -> Result<()> { + let id = self.clients_state.find_signer(payer)?; + + self.coordinator + .cooldown_witness(&id, witness) + .map_err(|err| anchor_lang::error!(ProgramError::from(err)))?; + + self.tick() + } + pub fn warmup_witness( &mut self, payer: &Pubkey, diff --git a/architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/lib.rs b/architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/lib.rs index 9cb58b948..3d18a7066 100644 --- a/architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/lib.rs +++ b/architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/lib.rs @@ -268,7 +268,6 @@ pub mod psyche_solana_coordinator { ) } - #[allow(unused_variables)] // for the metadata field. adding a _ prefix results in anchor's IDL not matching the actual types. lol. pub fn warmup_witness( ctx: Context, proof: WitnessProof, @@ -289,6 +288,26 @@ pub mod psyche_solana_coordinator { ) } + pub fn cooldown_witness( + ctx: Context, + proof: WitnessProof, + participant_bloom: WitnessBloom, + broadcast_bloom: WitnessBloom, + broadcast_merkle: MerkleRoot, + ) -> Result<()> { + let mut account = ctx.accounts.coordinator_account.load_mut()?; + account.increment_nonce(); + account.state.cooldown_witness( + ctx.accounts.user.key, + Witness { + proof, + participant_bloom, + broadcast_bloom, + broadcast_merkle, + }, + ) + } + pub fn health_check( ctx: Context, id: NodeIdentity, diff --git a/architectures/decentralized/solana-treasurer/Cargo.lock b/architectures/decentralized/solana-treasurer/Cargo.lock index 9be05b852..58d9afc91 100644 --- a/architectures/decentralized/solana-treasurer/Cargo.lock +++ b/architectures/decentralized/solana-treasurer/Cargo.lock @@ -1605,6 +1605,7 @@ name = "psyche-coordinator" version = "0.2.0" dependencies = [ "anchor-lang", + "anyhow", "async-trait", "bytemuck", "cfg_eval", diff --git a/architectures/decentralized/testing/src/docker_setup.rs b/architectures/decentralized/testing/src/docker_setup.rs index ad1c46981..f656e7acd 100644 --- a/architectures/decentralized/testing/src/docker_setup.rs +++ b/architectures/decentralized/testing/src/docker_setup.rs @@ -203,6 +203,9 @@ pub async fn spawn_new_client( }) .collect(); + // Skip checkpoint upload validation in tests + env_vars.push("SKIP_CHECKPOINT_UPLOAD=true".to_string()); + // Add extra env vars for keypair if provided env_vars.extend(extra_env); diff --git a/config/llama2-20m-dolma-noverify-no-checkpointer/state.toml b/config/llama2-20m-dolma-noverify-no-checkpointer/state.toml index 11df63854..692440638 100644 --- a/config/llama2-20m-dolma-noverify-no-checkpointer/state.toml +++ b/config/llama2-20m-dolma-noverify-no-checkpointer/state.toml @@ -24,7 +24,7 @@ max_seq_len = 2048 cold_start_warmup_steps = 0 [model.LLM.data_location] Server = "127.0.0.1:20001" -[model.LLM.checkpoint.Hub] +[model.LLM.checkpoint.Dummy] repo_id = "emozilla/llama2-20m-init" [model.LLM.lr_schedule.Cosine] base_lr = 4.0e-4 diff --git a/config/solana-test/nano-config.toml b/config/solana-test/nano-config.toml index 204b2f58f..20bb85139 100644 --- a/config/solana-test/nano-config.toml +++ b/config/solana-test/nano-config.toml @@ -1,6 +1,6 @@ [config] warmup_time = 100 -cooldown_time = 5 +cooldown_time = 30 epoch_time = 120 max_round_train_time = 30 round_witness_time = 5 diff --git a/docker/test/docker-compose.yml b/docker/test/docker-compose.yml index 7c209a12a..7997802c6 100644 --- a/docker/test/docker-compose.yml +++ b/docker/test/docker-compose.yml @@ -12,6 +12,7 @@ services: - '../../config/client/.env.local' environment: - HF_TOKEN=${HF_TOKEN:-} + - SKIP_CHECKPOINT_UPLOAD=true networks: - psyche-test-network diff --git a/psyche-book/src/enduser/join-run.md b/psyche-book/src/enduser/join-run.md index 968b09980..d17fcd909 100644 --- a/psyche-book/src/enduser/join-run.md +++ b/psyche-book/src/enduser/join-run.md @@ -68,6 +68,11 @@ WS_RPC=wss://your-primary-rpc-provider.com # If not set, the client will automatically discover and join an available run RUN_ID=your_run_id_here +# Required: access token to write model states in the storage. +# Depending on the config this will be a HuggingFace token or Google Cloud Storage credentials file path, only one of them should be needed +HF_TOKEN=HuggingFace_write_token_for_repo +GOOGLE_CREDENTIALS_FILE_PATH=/path/to/application_default_credentials.json + # Recommended: Fallback RPC Endpoints (for reliability) RPC_2=https://your-backup-rpc-provider.com WS_RPC_2=wss://your-backup-rpc-provider.com diff --git a/psyche-book/src/enduser/run-config.md b/psyche-book/src/enduser/run-config.md index 8440dd0a7..4974e3ac3 100644 --- a/psyche-book/src/enduser/run-config.md +++ b/psyche-book/src/enduser/run-config.md @@ -16,7 +16,7 @@ Here's a sample config with some of its options documented. # maximum time, in seconds, to let nodes download the model from a checkpoint / other nodes warmup_time = 30 -# time, in seconds, to let nodes bring the model from the GPU to disk, and to opt to join the next round. +# time, in seconds, to let nodes bring the model from the GPU to disk, upload the model to the remote storage and to opt to join the next round. cooldown_time = 30 # time, in seconds, that an epoch will last. @@ -70,10 +70,16 @@ architecture = "HfLlama" data_type = "Pretraining" max_seq_len = 2048 +# Use one of the following checkpoint options (Hub or Gcs, not both): + [model.LLM.checkpoint.Hub] -# Repo where the model is located in HugggingFace, will be used to download the model at the beginning of training. +# Repo where the model is located in HuggingFace, will be used to download the model at the beginning of training. repo_id = "emozilla/llama2-20m-init" +# Or, use Google Cloud Storage instead: +# [model.LLM.checkpoint.Gcs] +# bucket = "bucket_name" + [model.LLM.data_location.Http] # Token size in bytes, can be "TwoBytes" or "FourBytes" token_size_in_bytes = "TwoBytes" diff --git a/psyche-book/src/explain/gcs-checkpoints.md b/psyche-book/src/explain/gcs-checkpoints.md index e1f21e53f..81782aa16 100644 --- a/psyche-book/src/explain/gcs-checkpoints.md +++ b/psyche-book/src/explain/gcs-checkpoints.md @@ -92,6 +92,14 @@ or gcloud auth application-default login --scopes="https://www.googleapis.com/auth/cloud-platform" ``` +This creates a credentials file at `~/.config/gcloud/application_default_credentials.json`. To use these credentials with the run-manager, set the `GOOGLE_CREDENTIALS_FILE_PATH` variable in your env file: + +```bash +GOOGLE_CREDENTIALS_FILE_PATH=/home/youruser/.config/gcloud/application_default_credentials.json +``` + +The run-manager will mount this file into the Docker container and configure the `GOOGLE_APPLICATION_CREDENTIALS` environment variable automatically. + ### Useful Links - [Google Cloud Console](https://console.cloud.google.com) diff --git a/psyche-book/src/explain/general-workflow.md b/psyche-book/src/explain/general-workflow.md index 4ac8bfa63..f18c207c9 100644 --- a/psyche-book/src/explain/general-workflow.md +++ b/psyche-book/src/explain/general-workflow.md @@ -95,18 +95,27 @@ Any clients that have failed [health checks](#health-checks) will also be remove ### Cooldown phase (state: Cooldown) -The _Cooldown_ phase is the last phase of an epoch, during which the Coordinator waits the _Cooldown_ period to elapse. At this point the clients will begin to do a new checkpoint of the model, this is saving the state of the model at that time to a external storage, such as a Hugging Face. +The **Cooldown** phase is the last phase of an epoch. At this point, clients begin creating a new checkpoint of the model. This means saving the current state of the model to external storage, such as Hugging Face or a bucket in Google Cloud Storage (GCS). -When the _Cooldown_ phase begins, the Coordinator also resets the current model checkpoint state to `Checkpoint::P2P`, indicating that new joiners should download the latest copy of the model from the other participants and not from the usual checkpoint. +At the beginning of this state, the run elects a subset of clients that will be designated as **checkpointers**. All clients are potential checkpointers: one third of the total clients in the run will be elected pseudo-randomly at this stage. If a client is elected, it will start uploading the model state to the storage declared in the run configuration by the run owner. -Upon exiting the _Cooldown_ phase, the Coordinator transitions to the next epoch, saving the previous epoch state, and moving back to the _WaitingForMembers_ phase. All the clients that were participating in the previous epoch automatically join to the new epoch unless they exit manually. +The client that finishes uploading the model sends a transaction to the coordinator, called the **opportunistic cooldown**, indicating that the entire model was uploaded successfully. + +There are two ways the coordinator can transition from this state to the next one: + +- As soon as the first opportunistic cooldown transaction arrives, the coordinator moves to the next state and cancels all upload tasks from the remaining clients, since it already knows that at least one checkpointer has uploaded the complete model correctly. +- If no transaction is received, there is a maximum cooldown time defined in the run configuration. If this time is reached, the coordinator will move to the next state even if no new checkpoint was produced. + +When the _Cooldown_ phase begins, the coordinator also resets the current model checkpoint state to `Checkpoint::P2P`, indicating that new joiners should download the latest copy of the model from other participants rather than from the usual checkpoint storage. + +Upon exiting the _Cooldown_ phase, the coordinator transitions to the next epoch, saving the previous epoch state and moving back to the _WaitingForMembers_ phase. All clients that participated in the previous epoch automatically join the new epoch unless they exit manually. ### It all comes together Here's is an overview of how the state of the run can change depending on the situation: ```mermaid -%%{init: {'theme':'base', 'themeVariables': { 'fontSize':'35px'}}}%% +%%{init: {'theme':'base', 'themeVariables': { 'fontSize':'45px'}}}%% flowchart LR WFM((Waiting For Members)) W((Warmup)) @@ -119,6 +128,8 @@ flowchart LR d{Witness quorum reached} e{Max training time passed} f{End of the epoch reached} + g{Client checkpoints} + h{Max cooldown time passed} WFM --> a a -->|Yes| W @@ -135,7 +146,11 @@ flowchart LR WI --> f f -->|Yes| CD f -->|No| T - CD --> WFM + CD -->g + g -->|Yes| WFM + g -->|No|h + h -->|Yes| WFM + h -->|No| CD ``` And this is how it fits with real the real clients and how they interact in each of the stages. The committee in this case is the structure that contains all the witness data for the round. diff --git a/psyche-book/src/explain/index.md b/psyche-book/src/explain/index.md index 16d6eb6aa..a2e6baa27 100644 --- a/psyche-book/src/explain/index.md +++ b/psyche-book/src/explain/index.md @@ -62,7 +62,7 @@ These three phases constitute a **round** of training and will be looping until At the start of an **epoch**, all clients have a window of time to join the run by requesting to be added by coordinator, and then connecting to the other participating clients. This state will be known as the _Waiting for Members_ phase. -Once a minimum threshold of clients has been met, the run will transition to the _Warmup_ phase and begin a countdown to allow connected clients to update their copy of the model. To obtain a copy of the model, the Coordinator will either direct clients to a checkpoint uploaded somewhere like HuggingFace and they will have to download it from there or direct clients to [download the model from other clients](./model-sharing.md) via the p2p network. In the first epoch, all clients will download the model from HuggingFace and after that every new epoch, clients will download the model from other clients via the p2p network. +Once a minimum threshold of clients has been met, the run will transition to the _Warmup_ phase and begin a countdown to allow connected clients to update their copy of the model. To obtain a copy of the model, the Coordinator will either direct clients to a checkpoint uploaded somewhere like HuggingFace or Google Cloud Storage and they will have to download it from there or direct clients to [download the model from other clients](./model-sharing.md) via the p2p network. In the first epoch, all clients will download the model from the external storage and after that every new epoch, clients will download the model from other clients via the p2p network. After the _Warmup_ phase ends, it will enter the _Training_ phase. @@ -84,7 +84,7 @@ At the start of each round, one or more clients are randomly selected as witness These bloom filters are sent to the coordinator, which then combines them into a provable consensus of which results to apply to the model. -Once a witness quorum is reached, the coordinator advances to the _Training_ phase to allow all clients a brief window to download every training result of the previous round, clients are assigned new data, and the process repeats. After a fixed amount of time, a _Cooldown_ round occurs, marking the end of an **epoch**. This time is configurable in the run creation process that we'll explore in the other sections. +Once a witness quorum is reached, the coordinator advances to the _Training_ phase to allow all clients a brief window to download every training result of the previous round, clients are assigned new data, and the process repeats. After a fixed amount of time, a _Cooldown_ round occurs, marking the end of an **epoch**. At this state, one third of the clients are randomly selected as checkpointers and all of them starts uploading the state of the model to an external storage. There's a maximum amount of time for staying in this state, this time is configurable in the run creation process that we'll explore in the other sections. ## The witness/train loop visualized diff --git a/psyche-book/src/explain/model-sharing.md b/psyche-book/src/explain/model-sharing.md index ba61ebb12..5bc0bb8cd 100644 --- a/psyche-book/src/explain/model-sharing.md +++ b/psyche-book/src/explain/model-sharing.md @@ -6,15 +6,19 @@ At the beginning of a run, all clients must download the model parameters, token Each client will then modify their copy of the model by receiving new training results from other clients and applying them. This keeps everyone's copy of model identical within an **epoch** without an additional full synchronization step. -When a new client joins a run that has already progressed past its first epoch, it would not be correct for the client to download the original model from HuggingFace, as the model parameters would have already been updated during training. Instead, the new client must acquire a copy of the model from the peers who have been actively training it. +When a new client joins a run that has already progressed past its first epoch, it would not be correct for the client to download the original model from the external storage, as the model parameters would have already been updated during training. Instead, the new client must acquire a copy of the model from the peers who have been actively training it. This synchronization process occurs during the _Warmup_ phase, while the coordinator waits to begin the next _Training_ phase. -To address this, we **checkpoint** the model at the end of an **epoch**, where clients save and share the entire model for new peers to join. There are two checkpointing variants: HuggingFace based and P2P based. +To address this, we **checkpoint** the model at the end of an **epoch**, where clients save and share the entire model for new peers to join. There are three checkpointing variants: HuggingFace based, Google Cloud Storage based and P2P based. ## HuggingFace checkpoint -In this approach, a client or a set of clients can optionally run as **checkpointers** if they declare a checkpoint URL when joining the run. These clients upload their copy of updated model to HuggingFace after each epoch, and send the URL for this checkpoint to the coordinator. When a new client joins the run, it retrieves the checkpoint URL from the coordinator, and connects to HuggingFace to download the latest copy of the model parameters and configuration files. +In this approach, a client or a set of clients will be elected randomly as **checkpointers**. These clients upload their copy of updated model to HuggingFace at Cooldown state at the end of the epoch. The model will be uploaded to the HuggingFace repository that is declared in the run configuration by the run owner. When a new client joins the run it connects to HuggingFace to download the latest copy of the model parameters and configuration files. + +## Google Cloud Storage checkpoint + +Very similar to the previous approach but based on Google Cloud Storage bucket. Every elected checkpointer will upload the model at the end of an epoch. The bucket name is declared by the run owner in the initial configuration. If a client joins the run, it connects to the GCS and download the model parameters and configuration files. ## P2P checkpoint diff --git a/scripts/train-solana-test.sh b/scripts/train-solana-test.sh index 55699e591..6c6bdf4fa 100755 --- a/scripts/train-solana-test.sh +++ b/scripts/train-solana-test.sh @@ -20,15 +20,23 @@ elif [[ -z "${WALLET_FILE:-}" ]]; then trap "echo 'Cleaning up ephemeral wallet file...'; rm -f '${WALLET_FILE}'" EXIT fi -RPC=${RPC:-"http://127.0.0.1:8899"} -WS_RPC=${WS_RPC:-"ws://127.0.0.1:8900"} +RPC=${RPC:-"http://localhost:8899"} +WS_RPC=${WS_RPC:-"ws://localhost:8900"} RUN_ID=${RUN_ID:-"test"} AUTHORIZER=${AUTHORIZER:-"11111111111111111111111111111111"} +if [[ "$CHECKPOINT" == true ]]; then + echo -e "\n[+] Starting Solana training with checkpointing enabled..." +else + echo -e "\n[+] Starting Solana training without checkpointing..." +fi + # presets for a DGX or an HGX DP=${DP:-"8"} TP=${TP:-"1"} BATCH_SIZE=${BATCH_SIZE:-"1"} +export HF_TOKEN=${HF_TOKEN:-""} +export GOOGLE_APPLICATION_CREDENTIALS=${GOOGLE_APPLICATION_CREDENTIALS:-""} # fine if this fails solana airdrop 10 "$(solana-keygen pubkey ${WALLET_FILE})" --url "${RPC}" || true @@ -47,6 +55,7 @@ if [[ "$OTLP_METRICS_URL" == "" ]]; then --micro-batch-size ${BATCH_SIZE} \ --authorizer ${AUTHORIZER} \ --logs "console" \ + $( [[ "$CHECKPOINT" != "true" ]] && echo "--skip-checkpoint-upload" ) \ "$@" else cargo run --release --bin psyche-solana-client -- \ @@ -62,5 +71,6 @@ else --authorizer ${AUTHORIZER} \ --oltp-metrics-url "http://localhost:4318/v1/metrics" \ --oltp-logs-url "http://localhost:4318/v1/logs" \ + $( [[ "$CHECKPOINT" != "true" ]] && echo "--skip-checkpoint-upload" ) \ "$@" fi diff --git a/shared/client/Cargo.toml b/shared/client/Cargo.toml index d46726040..7ba5de8f0 100644 --- a/shared/client/Cargo.toml +++ b/shared/client/Cargo.toml @@ -35,6 +35,7 @@ hf-hub.workspace = true clap.workspace = true sysinfo = "0.32.0" iroh-blobs.workspace = true +google-cloud-storage.workspace = true [features] parallelism = ["psyche-modeling/parallelism"] diff --git a/shared/client/src/cli.rs b/shared/client/src/cli.rs index e48c9c1f4..b3871695b 100644 --- a/shared/client/src/cli.rs +++ b/shared/client/src/cli.rs @@ -1,14 +1,13 @@ use crate::{CheckpointConfig, WandBInfo}; -use crate::UploadInfo; use anyhow::{Result, anyhow, bail}; use clap::Args; -use psyche_data_provider::{GcsUploadInfo, HubUploadInfo}; use psyche_eval::tasktype_from_name; use psyche_modeling::Devices; use psyche_network::{DiscoveryMode, RelayKind, SecretKey}; use psyche_tui::LogOutput; use std::{path::PathBuf, time::Duration}; +use tracing::info; pub fn read_identity_secret_key( identity_secret_key_path: Option<&PathBuf>, @@ -141,20 +140,8 @@ pub struct TrainArgs { pub prompt_task: bool, /// If provided, every model parameters update will be save in this directory after each epoch. - #[clap(long, env)] - pub checkpoint_dir: Option, - - /// Path to the Hugging Face repository containing model data and configuration. - #[clap(long, env)] - pub hub_repo: Option, - - /// Name of the GCS bucket containing model data and configuration. - #[clap(long, env)] - pub gcs_bucket: Option, - - /// Prefix within the GCS bucket for model data and configuration. - #[clap(long, env)] - pub gcs_prefix: Option, + #[clap(long, env, default_value_os_t = default_checkpoint_dir())] + pub checkpoint_dir: PathBuf, #[clap(long, env, default_value_t = 3)] pub hub_max_concurrent_downloads: usize, @@ -204,6 +191,10 @@ pub struct TrainArgs { #[clap(long, default_value_t = 3, env)] pub keep_steps: u32, + + /// Skip saving and uploading checkpoints (for testing). + #[clap(long, default_value_t = false, env, hide = true)] + pub skip_checkpoint_upload: bool, } impl TrainArgs { @@ -232,74 +223,24 @@ impl TrainArgs { Ok(wandb_info) } - pub fn checkpoint_config(&self) -> Result> { - let hub_read_token = std::env::var("HF_TOKEN").ok(); - - if self.hub_repo.is_some() && self.gcs_bucket.is_some() { - bail!("Use either GCS or HF hub for checkpoint uploads, not both."); - } - - let checkpoint_dir = match &self.checkpoint_dir { - Some(dir) => dir, - None => { - if self.hub_repo.is_some() || self.gcs_bucket.is_some() { - bail!( - "--hub-repo or --gcs-bucket was set, but no --checkpoint-dir was passed!" - ); - } - return Ok(None); - } - }; - - let upload_info = self.build_upload_info(&hub_read_token)?; + pub fn checkpoint_config(&self) -> Result { + let hub_token = std::env::var("HF_TOKEN").ok(); - if upload_info.is_some() && self.keep_steps == 0 { + if self.keep_steps == 0 { bail!( "keep_steps must be >= 1 for checkpoint uploads (got {})", self.keep_steps ); } - Ok(Some(CheckpointConfig { - checkpoint_dir: checkpoint_dir.clone(), - upload_info, + Ok(CheckpointConfig { + checkpoint_dir: self.checkpoint_dir.clone(), delete_old_steps: self.delete_old_steps, keep_steps: self.keep_steps, - })) - } - - fn build_upload_info(&self, hub_token: &Option) -> Result> { - if let Some(repo) = &self.hub_repo { - return self.build_hub_upload_info(repo, hub_token); - } - - if let Some(bucket) = &self.gcs_bucket { - return self.build_gcs_upload_info(bucket); - } - - Ok(None) - } - - fn build_hub_upload_info( - &self, - repo: &str, - token: &Option, - ) -> Result> { - let token = token.as_ref().ok_or_else(|| { - anyhow::anyhow!("hub-repo and checkpoint-dir set, but no HF_TOKEN env variable.") - })?; - - Ok(Some(UploadInfo::Hub(HubUploadInfo { - hub_repo: repo.to_string(), - hub_token: token.to_string(), - }))) - } - - fn build_gcs_upload_info(&self, bucket: &str) -> Result> { - Ok(Some(UploadInfo::Gcs(GcsUploadInfo { - gcs_bucket: bucket.to_string(), - gcs_prefix: self.gcs_prefix.clone(), - }))) + hub_token, + skip_upload: self.skip_checkpoint_upload, + run_down_client: None, // Set by app.rs after construction + }) } pub fn eval_tasks(&self) -> Result> { @@ -329,6 +270,17 @@ impl TrainArgs { } } +fn default_checkpoint_dir() -> PathBuf { + let final_dir = if std::path::Path::new("/scratch").exists() { + PathBuf::from("/scratch/checkpoints") + } else { + let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string()); + PathBuf::from(home).join(".cache/psyche/local_checkpoints") + }; + info!("Default checkpoint directory set to {:?}", final_dir); + final_dir +} + pub fn prepare_environment() { psyche_modeling::set_suggested_env_vars(); diff --git a/shared/client/src/client.rs b/shared/client/src/client.rs index 80f4795d9..7f2957798 100644 --- a/shared/client/src/client.rs +++ b/shared/client/src/client.rs @@ -80,7 +80,6 @@ impl Client { // From Run let (tx_witness, mut rx_witness) = mpsc::unbounded_channel(); let (tx_health_check, mut rx_health_check) = mpsc::unbounded_channel(); - let (tx_checkpoint, mut rx_checkpoint) = mpsc::unbounded_channel(); let (tx_model, mut rx_model) = mpsc::unbounded_channel(); let (tx_distro_result, mut rx_distro_result) = mpsc::unbounded_channel(); let (tx_request_download, mut rx_request_download) = mpsc::unbounded_channel(); @@ -103,7 +102,6 @@ impl Client { metrics: metrics.clone(), tx_witness, tx_health_check, - tx_checkpoint, tx_model, tx_parameters_req, tx_config, @@ -130,7 +128,6 @@ impl Client { let mut retry_check_interval = interval(DOWNLOAD_RETRY_CHECK_INTERVAL); let mut opportunistic_witness_interval = interval(OPPROTUNISTIC_WITNESS_INTERVAL); let mut check_connection_interval = interval(CHECK_CONNECTION_INTERVAL); - let mut wait_for_checkpoint = false; let mut last_gossip_connection_time = SystemTime::now(); debug!("Starting client loop"); @@ -138,9 +135,6 @@ impl Client { select! { _ = cancel.cancelled() => { info!("Got request to cancel main client loop"); - if run.doing_checkpoint() { - wait_for_checkpoint = true; - } break; } @@ -522,9 +516,6 @@ impl Client { Some(health_check) = rx_health_check.recv() => { watcher.backend_mut().send_health_check(health_check).await?; } - Some(checkpoint) = rx_checkpoint.recv() => { - watcher.backend_mut().send_checkpoint(checkpoint).await?; - } Some(model) = rx_model.recv() => { sharable_model.update_parameters(model)?; }, @@ -646,30 +637,6 @@ impl Client { let p2p_shutdown = p2p.shutdown(); - if wait_for_checkpoint { - info!("Waiting for all pending checkpoints to finish"); - - // Keep waiting for checkpoints while there are uploads pending - let mut checkpoint_check_interval = interval(Duration::from_secs(10)); - while run.doing_checkpoint() { - tokio::select! { - checkpoint = rx_checkpoint.recv() => { - if let Some(checkpoint) = checkpoint { - info!("Checkpoint upload completed, sending to Solana"); - watcher.backend_mut().send_checkpoint(checkpoint).await?; - } else { - // Channel closed, no more checkpoints coming - break; - } - } - _ = checkpoint_check_interval.tick() => { - } - } - } - - info!("All checkpoints finished, exiting main client loop"); - } - p2p_shutdown .await .map_err(|e| anyhow!("Error shutting down p2p: {e}")) diff --git a/shared/client/src/lib.rs b/shared/client/src/lib.rs index bdad43e30..b1e6bed99 100644 --- a/shared/client/src/lib.rs +++ b/shared/client/src/lib.rs @@ -9,8 +9,8 @@ pub use cli::{TrainArgs, prepare_environment, print_identity_keys, read_identity pub use client::Client; pub use protocol::{Broadcast, BroadcastType, Finished, NC, TrainingResult}; pub use state::{ - CheckpointConfig, GcsUploadInfo, HubUploadInfo, InitRunError, RoundState, RunInitConfig, - RunInitConfigAndIO, UploadInfo, + CheckpointConfig, CheckpointUploader, InitRunError, RoundState, RunInitConfig, + RunInitConfigAndIO, }; pub use tui::{ClientTUI, ClientTUIState}; diff --git a/shared/client/src/state/cooldown.rs b/shared/client/src/state/cooldown.rs index 8ffcb62f5..0eb0426bc 100644 --- a/shared/client/src/state/cooldown.rs +++ b/shared/client/src/state/cooldown.rs @@ -1,9 +1,9 @@ -use crate::UploadInfo; +use crate::CheckpointUploader; use psyche_coordinator::{ - Coordinator, - model::{self}, + CheckpointerSelection, Coordinator, + model::{self, HubRepo, LLM, Model}, }; -use psyche_data_provider::{GcsManifestMetadata, UploadError, upload_to_gcs, upload_to_hub}; +use psyche_data_provider::{GcsManifestMetadata, UploadError, upload_to_gcs_signed, upload_to_hub}; #[cfg(feature = "python")] use psyche_modeling::CausalLM; use psyche_modeling::{ @@ -13,7 +13,10 @@ use std::{ cmp::Reverse, collections::{BinaryHeap, HashMap}, path::PathBuf, - sync::Arc, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, }; use tch::Tensor; use thiserror::Error; @@ -21,6 +24,7 @@ use tokio::{ sync::{Mutex, mpsc}, task::JoinHandle, }; +use tracing::error; use tracing::{Instrument, info, info_span, warn}; use super::{ @@ -38,12 +42,14 @@ pub enum CooldownError { #[error("error while checkpointing: {0}")] Checkpoint(#[from] CheckpointError), + + #[error("error in cooldown step: {0}")] + CoordinatorError(#[from] psyche_coordinator::CoordinatorError), } pub struct CooldownStepMetadata { - tx_checkpoint: mpsc::UnboundedSender, tx_model: mpsc::UnboundedSender>, - checkpoint_info: Option, + checkpoint_info: CheckpointConfig, checkpoint_extra_files: Vec, model_task_runner: ModelTaskRunner, @@ -58,14 +64,12 @@ pub struct CooldownStepMetadata { impl CooldownStepMetadata { pub fn new( - tx_checkpoint: mpsc::UnboundedSender, tx_model: mpsc::UnboundedSender>, - checkpoint_info: Option, + checkpoint_info: CheckpointConfig, checkpoint_extra_files: Vec, model_task_runner: ModelTaskRunner, ) -> Self { Self { - tx_checkpoint, tx_model, checkpoint_info, checkpoint_extra_files, @@ -128,6 +132,7 @@ impl CooldownStepMetadata { pub fn start( &self, mut trainers: Vec, + client_index: u64, state: &Coordinator, ) -> Result { let Some(mut trainer) = trainers.pop() else { @@ -139,77 +144,129 @@ impl CooldownStepMetadata { let epoch = state.progress.epoch as u32; let checkpoint_extra_files = self.checkpoint_extra_files.clone(); let checkpoint_info = self.checkpoint_info.clone(); - let tx_checkpoint = self.tx_checkpoint.clone(); + let Model::LLM(LLM { checkpoint, .. }) = state.model; let tx_model = self.tx_model.clone(); let model_task_runner = self.model_task_runner.clone(); let delete_queue = self.delete_queue.clone(); + let checkpointer_selection = CheckpointerSelection::from_coordinator(state, 0)?; + let is_checkpointer = checkpointer_selection + .is_checkpointer(client_index, state.epoch_state.clients.len() as u64); + let cancellation_token = tokio_util::sync::CancellationToken::new(); + let checkpoint_completed = Arc::new(AtomicBool::new(false)); + + let checkpointing_and_evals: JoinHandle> = + tokio::task::spawn({ + let cancellation_token = cancellation_token.clone(); + let checkpoint_completed = checkpoint_completed.clone(); + async move { + info!("Extracting full model..."); + let (variables, trainer) = + tokio::task::spawn_blocking::<_, Result<_, CheckpointError>>(|| { + let variables = trainer.extract()?; + info!("Model extracted; {} parameters", variables.len()); + Ok((variables, trainer)) + }) + .await + .map_err(|_| CheckpointError::ExtractThreadCrashed)??; + + let variables_clone: HashMap = variables + .iter() + .map(|(name, var)| (name.clone(), var.shallow_clone())) + .collect(); + + // for p2p model sharing we use the native trainer shape + tx_model + .send(variables_clone) + .map_err(|_| CheckpointError::SendCheckpoint)?; + + // convert from internal shape to serialized shape (e.g. torchtitan to hf) + let (variables, trainer) = match trainer { + #[cfg(feature = "python")] + Trainer::PythonDistributed(_) => { + info!("Converting distributed trainer variables for checkpointing..."); + tokio::task::spawn_blocking(|| (trainer.convert(Some(variables)), trainer)) + .await + .map_err(|_| CheckpointError::ExtractThreadCrashed)? + } + _ => (variables, trainer), + }; + + trainers.push(trainer); + let evals = model_task_runner.start(trainers); + if !is_checkpointer { + info!("Skipping checkpoint upload as this node is not a checkpointer for this epoch"); + return Ok(evals); + } - let checkpointing_and_evals: CheckpointAndEvalsHandle = tokio::task::spawn( - async move { - info!("Extracting full model..."); - let (variables, trainer) = - tokio::task::spawn_blocking::<_, Result<_, CheckpointError>>(|| { - let variables = trainer.extract()?; - info!("Model extracted; {} parameters", variables.len()); - Ok((variables, trainer)) - }) - .await - .map_err(|_| CheckpointError::ExtractThreadCrashed)??; - - let variables_clone: HashMap = variables - .iter() - .map(|(name, var)| (name.clone(), var.shallow_clone())) - .collect(); - - // for p2p model sharing we use the native trainer shape - tx_model - .send(variables_clone) - .map_err(|_| CheckpointError::SendCheckpoint)?; - - // convert from internal shape to serialized shape (e.g. torchtitan to hf) - let (variables, trainer) = match trainer { - #[cfg(feature = "python")] - Trainer::PythonDistributed(_) => { - info!("Converting distributed trainer variables for checkpointing..."); - tokio::task::spawn_blocking(|| (trainer.convert(Some(variables)), trainer)) - .await - .map_err(|_| CheckpointError::ExtractThreadCrashed)? + let CheckpointConfig { + checkpoint_dir, + delete_old_steps, + keep_steps, + hub_token, + skip_upload, + run_down_client, + } = checkpoint_info; + + // When skip_upload is true (testing), skip all checkpoint saving + if skip_upload { + info!("Skipping checkpoint save and upload (skip_upload flag is set)"); + checkpoint_completed.store(true, Ordering::Release); + return Ok(evals); } - _ => (variables, trainer), - }; - - trainers.push(trainer); - let evals = model_task_runner.start(trainers); - - let Some(CheckpointConfig { - upload_info, - checkpoint_dir, - delete_old_steps, - keep_steps, - }) = checkpoint_info - else { - return Ok((evals, None)); - }; - - let upload_handle = tokio::task::spawn(async move { + + let uploader = match checkpoint { + model::Checkpoint::Hub(HubRepo { + repo_id, + revision: _, + }) + | model::Checkpoint::P2P(HubRepo { + repo_id, + revision: _, + }) => { + if let Some(token) = hub_token { + match CheckpointUploader::new_hub(repo_id.to_string(), token).await { + Ok(uploader) => Some(uploader), + Err(err) => { + error!("Failed to create HF uploader: {}", err); + None + } + } + } else { + warn!("HF_TOKEN env not provided, skipping upload to HuggingFace Hub"); + None + } + } + model::Checkpoint::Gcs(_) | model::Checkpoint::P2PGcs(_) => { + if let Some(client) = run_down_client { + Some(CheckpointUploader::new_gcs(client)) + } else { + warn!("RunDownClient not configured, skipping GCS checkpoint upload"); + None + } + } + _ => None, + }; + let path = checkpoint_dir.join(format!("{run_id}-step{step}")); let local = save_checkpoint_locally(path, variables, checkpoint_extra_files).await?; - if let Some(upload_info) = upload_info { + let upload_succeeded = if let Some(uploader) = uploader { let manifest_metadata = GcsManifestMetadata { epoch, run_id: run_id.clone(), }; - upload_checkpoint( - upload_info, - manifest_metadata, - local.clone(), - step as u64, - tx_checkpoint, - ) - .await?; - } + match upload_checkpoint(uploader, manifest_metadata, local.clone(), step as u64, cancellation_token.clone()).await { + Ok(()) => true, + Err(err) => { + error!("Error uploading checkpoint: {}", err); + false + } + } + } else { + true + }; + checkpoint_completed.store(upload_succeeded, Ordering::Release); cleanup_dirs( delete_queue, @@ -221,16 +278,16 @@ impl CooldownStepMetadata { ) .await; - Ok(()) - }); - - Ok((evals, Some(upload_handle))) - } - .instrument(info_span!("checkpointing")), - ); + Ok(evals) + } + .instrument(info_span!("checkpointing")) + }); Ok(CooldownStep { checkpointing_and_evals, + cancellation_token, + checkpoint_completed, + sent_witness: false, }) } } @@ -260,54 +317,58 @@ async fn save_checkpoint_locally( } async fn upload_checkpoint( - upload_info: UploadInfo, + uploader: CheckpointUploader, manifest_metadata: GcsManifestMetadata, local: Vec, step: u64, - tx_checkpoint: mpsc::UnboundedSender, + cancellation_token: tokio_util::sync::CancellationToken, ) -> Result<(), CheckpointError> { - match upload_info { - UploadInfo::Gcs(gcs_info) => { - upload_to_gcs(gcs_info, manifest_metadata, local, step, tx_checkpoint) + match uploader { + CheckpointUploader::Gcs(run_down) => upload_to_gcs_signed( + &run_down, + manifest_metadata, + local, + step, + cancellation_token, + ) + .await + .map_err(CheckpointError::UploadError), + CheckpointUploader::Hub(hub_info) => { + upload_to_hub(hub_info, local, step, cancellation_token) .await .map_err(CheckpointError::UploadError) } - UploadInfo::Hub(hub_info) => upload_to_hub(hub_info, local, step, tx_checkpoint) - .await - .map_err(CheckpointError::UploadError), + CheckpointUploader::Dummy => { + info!("Dummy upload info provided; skipping upload"); + Ok(()) + } } } -type CheckpointAndEvalsHandle = JoinHandle< - Result< - ( - RunningEvals, - Option>>, - ), - CheckpointError, - >, ->; - #[derive(Debug)] pub struct CooldownStep { - checkpointing_and_evals: CheckpointAndEvalsHandle, + checkpointing_and_evals: JoinHandle>, + cancellation_token: tokio_util::sync::CancellationToken, + checkpoint_completed: Arc, + pub sent_witness: bool, } impl CooldownStep { - pub async fn finish( - self, - ) -> Result< - ( - RunningEvals, - Option>>, - ), - CooldownError, - > { - let (running_evals, upload_handle) = self + pub async fn finish(self) -> Result { + self.cancellation_token.cancel(); + let running_evals = self .checkpointing_and_evals .await .map_err(|_| CooldownError::CheckpointThreadCrashed)??; - Ok((running_evals, upload_handle)) + Ok(running_evals) + } + + pub fn is_finished(&self) -> bool { + self.checkpointing_and_evals.is_finished() + } + + pub fn checkpoint_complete(&self) -> bool { + self.checkpoint_completed.load(Ordering::Acquire) } } diff --git a/shared/client/src/state/init.rs b/shared/client/src/state/init.rs index 0fb4394b9..a8a4b780d 100644 --- a/shared/client/src/state/init.rs +++ b/shared/client/src/state/init.rs @@ -9,7 +9,7 @@ use psyche_core::{ use psyche_data_provider::{ DataProvider, DataProviderTcpClient, DownloadError, DummyDataProvider, PreprocessedDataProvider, Split, WeightedDataProvider, download_dataset_repo_async, - download_model_from_gcs_async, download_model_repo_async, + download_model_from_gcs_signed_async, download_model_repo_async, http::{FileURLs, HttpDataProvider}, }; use psyche_metrics::ClientMetrics; @@ -68,7 +68,7 @@ pub struct RunInitConfig { pub write_gradients_dir: Option, // checkpointing - pub checkpoint_config: Option, + pub checkpoint_config: CheckpointConfig, // configurable dummy training time (in seconds) for this client - relevant just for testing pub dummy_training_delay_secs: Option, @@ -157,7 +157,6 @@ pub struct RunInitConfigAndIO { pub tx_health_check: UnboundedSender, pub tx_witness: UnboundedSender, - pub tx_checkpoint: UnboundedSender, pub tx_model: UnboundedSender>, pub tx_parameters_req: UnboundedSender<(Vec, OneshotModelParameterSender)>, pub tx_config: UnboundedSender<(String, String)>, @@ -176,7 +175,6 @@ impl RunInitConfigAndIO { init_config, tx_witness, tx_health_check, - tx_checkpoint, tx_model, tx_config, tx_parameters_req, @@ -266,6 +264,8 @@ impl RunInitConfigAndIO { Ok(data_provider) }; + let run_down_client = init_config.checkpoint_config.run_down_client.clone(); + let model_future: JoinHandle> = match &llm.architecture { model::LLMArchitecture::HfLlama @@ -315,8 +315,8 @@ impl RunInitConfigAndIO { }), model::Checkpoint::Hub(_) | model::Checkpoint::P2P(_) - | model::Checkpoint::P2PGcs(_) - | model::Checkpoint::Gcs(_) => { + | model::Checkpoint::Gcs(_) + | model::Checkpoint::P2PGcs(_) => { let checkpoint = llm.checkpoint; tokio::spawn(async move { let (source, tokenizer, checkpoint_extra_files) = match checkpoint { @@ -438,19 +438,16 @@ impl RunInitConfigAndIO { vec![], ) } - model::Checkpoint::Gcs(gcs_repo) => { - let bucket: String = (&gcs_repo.bucket).into(); - let prefix: Option = gcs_repo.prefix.map(|p| (&p).into()); - + model::Checkpoint::Gcs(_) => { + let run_down = run_down_client.as_ref().ok_or_else(|| { + DownloadError::RunDown("RunDownClient not configured".into()) + })?; info!( - "Downloading model from gs://{}/{}", - bucket, - prefix.as_deref().unwrap_or("") + "Downloading model via run-down signed URLs for run {}", + run_down.run_id() ); - let repo_files = - download_model_from_gcs_async(&bucket, prefix.as_deref()) - .await?; + download_model_from_gcs_signed_async(run_down).await?; let checkpoint_extra_files = repo_files .iter() @@ -868,7 +865,6 @@ impl RunInitConfigAndIO { }; let cooldown = CooldownStepMetadata::new( - tx_checkpoint, tx_model, init_config.checkpoint_config, checkpoint_extra_files, diff --git a/shared/client/src/state/mod.rs b/shared/client/src/state/mod.rs index 78e6cd1eb..c63aa1d19 100644 --- a/shared/client/src/state/mod.rs +++ b/shared/client/src/state/mod.rs @@ -14,7 +14,8 @@ mod warmup; mod witness; pub use init::{InitRunError, RunInitConfig, RunInitConfigAndIO}; -pub use psyche_data_provider::{GcsUploadInfo, HubUploadInfo}; pub use round_state::RoundState; pub use steps::{ApplyMessageOutcome, RunManager}; -pub use types::{CheckpointConfig, DistroBroadcastAndPayload, FinishedBroadcast, UploadInfo}; +pub use types::{ + CheckpointConfig, CheckpointUploader, DistroBroadcastAndPayload, FinishedBroadcast, +}; diff --git a/shared/client/src/state/steps.rs b/shared/client/src/state/steps.rs index c584d2613..0661fc574 100644 --- a/shared/client/src/state/steps.rs +++ b/shared/client/src/state/steps.rs @@ -4,6 +4,7 @@ use crate::{ }; use iroh_blobs::api::Tag; +use psyche_coordinator::CheckpointerSelection; use psyche_coordinator::{Committee, Coordinator, RunState, Witness, WitnessProof}; use psyche_core::{IntegrationTestLogMarker, MerkleRoot, MerkleTree, NodeIdentity, sha256}; use psyche_modeling::{DistroResult, Trainer}; @@ -54,14 +55,8 @@ pub struct StepStateMachine { current_round: RoundState, previous_round: RoundState, step_finish_time: Option, - sent_warmup_finished: bool, - sent_warmup_witness: bool, coordinator_state: Coordinator, - - // Handles for HuggingFace uploads running in background - pending_upload_handles: - Vec>>, } #[derive(Error, Debug)] @@ -161,182 +156,244 @@ impl StepStateMachine { coordinator_state, step_finish_time: None, - sent_warmup_finished: false, - sent_warmup_witness: false, - - pending_upload_handles: Vec::new(), } } pub fn try_send_opportunistic_witness(&mut self) -> Result<(), OpportunisticWitnessError> { - if let Some(committee_info) = &self.current_round.committee_info { - // trace!("Checking for opprotunistic witness with committee info"); - if let ActiveStep::Training(step) = &self.active_step { - let all_prev_round_batches_are_trained = self - .previous_round - .batch_ids_not_yet_trained_on - .lock() - .unwrap() - .is_none(); - - if step.finished() && all_prev_round_batches_are_trained { - // Finished training and finished downloading the previous round's results - // (or we're on the first or last which has nothing to download) - - // check that all batches from the previous round are done deserializing - { - let prev_round_downloads = self.previous_round.downloads.lock().unwrap(); - for batch in &*prev_round_downloads { - match batch.1 { - // this batch is done deserializing, we can witness on it now. - PayloadState::Deserializing(thread) if thread.is_finished() => (), - // we're still downloading or deserializing this batch, so we're not ready to send an opportunistic witness. - // this function will get called again when a deserialize finishes. - _ => return Ok(()), - } - } - } + match self.coordinator_state.run_state { + RunState::Warmup => self.try_send_warmup_witness(), + RunState::Cooldown => self.try_send_cooldown_witness(), + _ => self.try_send_training_witness(), + } + } - if !self.current_round.sent_finished { - // okay, we're all done. we've trained and downloaded everything. - // send our early "finished message" + fn all_clients_finished(&self) -> bool { + self.coordinator_state + .epoch_state + .clients + .iter() + .all(|client| self.current_round.clients_finished.contains_key(&client.id)) + } - let merkle = MerkleTree::new(&self.previous_round.broadcasts) - .get_root() - .cloned() - .unwrap_or(MerkleRoot::default()); + fn get_unfinished_clients(&self) -> Vec { + self.coordinator_state + .epoch_state + .clients + .iter() + .filter(|client| !self.current_round.clients_finished.contains_key(&client.id)) + .map(|client| client.id) + .collect() + } - self.tx_broadcast_finished - .send(FinishedBroadcast { - step: self.current_round.step, - commitment_data_hash: sha256(&merkle.inner), - merkle, - proof: committee_info.0, - warmup: false, - }) - .map_err(|_| OpportunisticWitnessError::Finished)?; + fn get_client_index(&self) -> Option { + self.coordinator_state + .epoch_state + .clients + .iter() + .position(|x| x.id == self.identity) + .map(|i| i as u64) + } - self.current_round.sent_finished = true; + fn get_merkle_root(&self, broadcasts: &[[u8; 32]]) -> MerkleRoot { + MerkleTree::new(broadcasts) + .get_root() + .cloned() + .unwrap_or_default() + } - return Ok(()); - } + fn try_send_training_witness(&mut self) -> Result<(), OpportunisticWitnessError> { + let Some(committee_info) = self.current_round.committee_info.as_ref() else { + return Ok(()); + }; + let committee_proof = committee_info.0; - // if we get here we've sent our own finished message. - // now we just need to wait until we've received everyone else's finished - let unfinished_clients: Vec<_> = self - .coordinator_state - .epoch_state - .clients - .iter() - .filter_map(|client| { - if self.current_round.clients_finished.contains_key(&client.id) { - None - } else { - Some(client.id) - } - }) - .collect(); - if !unfinished_clients.is_empty() { - return Ok(()); - } + let ActiveStep::Training(step) = &self.active_step else { + return Ok(()); + }; - if let Some(witness) = WitnessStep::get_witness_to_send( - &mut self.previous_round, - &mut self.current_round, - ) { - info!(target: "witness", id = %self.identity, merkle=witness.broadcast_merkle.fmt_short(), "Sending opportunistic witness"); - - let metadata = self - .stats_logger - .lock() - .map_err(|_| OpportunisticWitnessError::StatsLoggerMutex)? - .get_witness_metadata(&self.coordinator_state); - self.tx_opportunistic_data - .send(OpportunisticData::WitnessStep(witness, metadata)) - .map_err(|_| OpportunisticWitnessError::Send)?; - } + if !step.finished() { + return Ok(()); + } + + let all_prev_round_batches_trained = self + .previous_round + .batch_ids_not_yet_trained_on + .lock() + .unwrap() + .is_none(); + + if !all_prev_round_batches_trained { + return Ok(()); + } + + // Check that all batches from the previous round are done deserializing + { + let prev_round_downloads = self.previous_round.downloads.lock().unwrap(); + for batch in &*prev_round_downloads { + match batch.1 { + PayloadState::Deserializing(thread) if thread.is_finished() => (), + // Still downloading or deserializing - will be called again when done + _ => return Ok(()), } } - } else if self.coordinator_state.run_state == RunState::Warmup { - if !self.sent_warmup_finished { - let merkle = MerkleTree::new(&self.current_round.broadcasts) - .get_root() - .cloned() - .unwrap_or(MerkleRoot::default()); - - info!(name: "send_warmup_broadcast", epoch = self.coordinator_state.progress.epoch, "Sending warmup ready broadcast"); - self.tx_broadcast_finished - .send(FinishedBroadcast { - step: 0, - commitment_data_hash: sha256(&merkle.inner), - merkle, - proof: Default::default(), - warmup: true, - }) - .map_err(|_| OpportunisticWitnessError::Finished)?; - - self.sent_warmup_finished = true; + } - return Ok(()); - } + // Send our early "finished message" if we haven't yet + if !self.current_round.sent_finished { + let merkle = self.get_merkle_root(&self.previous_round.broadcasts); + + self.tx_broadcast_finished + .send(FinishedBroadcast { + step: self.current_round.step, + commitment_data_hash: sha256(&merkle.inner), + merkle, + proof: committee_proof, + warmup: false, + }) + .map_err(|_| OpportunisticWitnessError::Finished)?; - let unfinished_clients: Vec<_> = self - .coordinator_state - .epoch_state - .clients - .iter() - .filter_map(|client| { - if self.current_round.clients_finished.contains_key(&client.id) { - None - } else { - Some(client.id) - } + self.current_round.sent_finished = true; + return Ok(()); + } + + // Wait until we've received everyone else's finished message + if !self.all_clients_finished() { + return Ok(()); + } + + // Send witness if available + if let Some(witness) = + WitnessStep::get_witness_to_send(&mut self.previous_round, &mut self.current_round) + { + info!(target: "witness", id = %self.identity, merkle=witness.broadcast_merkle.fmt_short(), "Sending opportunistic witness"); + + let metadata = self + .stats_logger + .lock() + .map_err(|_| OpportunisticWitnessError::StatsLoggerMutex)? + .get_witness_metadata(&self.coordinator_state); + + self.tx_opportunistic_data + .send(OpportunisticData::WitnessStep(witness, metadata)) + .map_err(|_| OpportunisticWitnessError::Send)?; + } + + Ok(()) + } + + fn try_send_warmup_witness(&mut self) -> Result<(), OpportunisticWitnessError> { + let ActiveStep::Warmup(ref warmup) = self.active_step else { + return Ok(()); + }; + + // Send warmup finished broadcast if we haven't yet + if !warmup.sent_finished { + let merkle = self.get_merkle_root(&self.current_round.broadcasts); + + info!(name: "send_warmup_broadcast", epoch = self.coordinator_state.progress.epoch, "Sending warmup ready broadcast"); + self.tx_broadcast_finished + .send(FinishedBroadcast { + step: 0, + commitment_data_hash: sha256(&merkle.inner), + merkle, + proof: Default::default(), + warmup: true, }) - .collect(); - if !unfinished_clients.is_empty() { - trace!( - unfinished_clients = ?unfinished_clients, - "Still waiting on {} warmup finish broadcasts", - unfinished_clients.len() - ); - return Ok(()); + .map_err(|_| OpportunisticWitnessError::Finished)?; + + if let ActiveStep::Warmup(ref mut warmup) = self.active_step { + warmup.sent_finished = true; } + return Ok(()); + } - if !self.sent_warmup_witness { - info!(name: "send_warmup_witness", epoch = self.coordinator_state.progress.epoch, "Sending warmup witness"); + // Wait for all clients to finish + let unfinished_clients = self.get_unfinished_clients(); + if !unfinished_clients.is_empty() { + trace!( + unfinished_clients = ?unfinished_clients, + "Still waiting on {} warmup finish broadcasts", + unfinished_clients.len() + ); + return Ok(()); + } - let merkle = MerkleTree::new(&self.current_round.broadcasts) - .get_root() - .cloned() - .unwrap_or(MerkleRoot::default()); + // Send warmup witness if we haven't yet + if warmup.sent_witness { + return Ok(()); + } - if let Some(index) = self - .coordinator_state - .epoch_state - .clients - .iter() - .position(|x| x.id == self.identity) - { - // coordinator needs to check the index for duplicate detection - let index = index as u64; - let witness = Witness { - proof: WitnessProof { - position: index, - index, - witness: Default::default(), - }, - participant_bloom: Default::default(), - broadcast_bloom: Default::default(), - broadcast_merkle: merkle, - }; - self.tx_opportunistic_data - .send(OpportunisticData::WarmupStep(witness)) - .map_err(|_| OpportunisticWitnessError::Send)?; - }; + info!(name: "send_warmup_witness", epoch = self.coordinator_state.progress.epoch, "Sending warmup witness"); + + let merkle = self.get_merkle_root(&self.current_round.broadcasts); + let Some(index) = self.get_client_index() else { + return Ok(()); + }; + + let witness = Witness { + proof: WitnessProof { + position: index, + index, + witness: Default::default(), + }, + participant_bloom: Default::default(), + broadcast_bloom: Default::default(), + broadcast_merkle: merkle, + }; - self.sent_warmup_witness = true; + self.tx_opportunistic_data + .send(OpportunisticData::WarmupStep(witness)) + .map_err(|_| OpportunisticWitnessError::Send)?; + + if let ActiveStep::Warmup(ref mut warmup) = self.active_step { + warmup.sent_witness = true; + } + Ok(()) + } + + fn try_send_cooldown_witness(&mut self) -> Result<(), OpportunisticWitnessError> { + if let ActiveStep::Cooldown(ref active_step) = self.active_step { + if !active_step.checkpoint_complete() || active_step.sent_witness { + return Ok(()); } } + + // Panic if client not found (should never happen) + let client_index = self + .get_client_index() + .expect("client should be in epoch_state.clients"); + + let checkpointer_selection = + CheckpointerSelection::from_coordinator(&self.coordinator_state, 0) + .map_err(|_| OpportunisticWitnessError::Send)?; + + if !checkpointer_selection.is_checkpointer( + client_index, + self.coordinator_state.epoch_state.clients.len() as u64, + ) { + return Ok(()); + } + + let merkle = self.get_merkle_root(&self.current_round.broadcasts); + + let witness = Witness { + proof: WitnessProof { + position: client_index, + index: client_index, + witness: Default::default(), + }, + participant_bloom: Default::default(), + broadcast_bloom: Default::default(), + broadcast_merkle: merkle, + }; + + self.tx_opportunistic_data + .send(OpportunisticData::CooldownStep(witness)) + .map_err(|_| OpportunisticWitnessError::Send)?; + + if let ActiveStep::Cooldown(ref mut active_step) = self.active_step { + active_step.sent_witness = true; + } Ok(()) } @@ -719,10 +776,7 @@ impl StepStateMachine { trace!( "since we're not a member of this step, killing cooldown step and returning to warmup to wait." ); - let (trainers, upload_handle) = cooldown.finish().await?; - if let Some(handle) = upload_handle { - self.pending_upload_handles.push(handle); - } + let trainers = cooldown.finish().await?; ActiveStep::Warmup(self.warmup.start( trainers, &mut self.previous_round, @@ -761,8 +815,6 @@ impl StepStateMachine { (ActiveStep::Warmup(warmup), RunState::RoundTrain) => { let trainers = warmup.finish().stop_evals().await?; self.step_finish_time = None; - self.sent_warmup_finished = false; - self.sent_warmup_witness = false; self.stats_logger .lock() .map_err(|_| StepError::StatsLoggerMutex)? @@ -834,19 +886,16 @@ impl StepStateMachine { // the epoch ended & we're transitioning to cooldown (ActiveStep::Witness(witnessing), RunState::Cooldown) => { let trainers = witnessing.finish().await?.stop_evals().await?; - // check here - self.cleanup_completed_uploads(); - ActiveStep::Cooldown(self.cooldown.start(trainers, &state)?) + ActiveStep::Cooldown(self.cooldown.start(trainers, client_index, &state)?) } - // cooldown is done, we consider waiting for members and warmup to be basically the same + // Cooldown is done — at least one checkpointer has successfully uploaded + // (the coordinator only transitions out of Cooldown after receiving a witness). + // Safe to cancel our own upload if still in progress. (ActiveStep::Cooldown(cooldown), RunState::WaitingForMembers) | (ActiveStep::Cooldown(cooldown), RunState::Warmup) | (ActiveStep::Cooldown(cooldown), RunState::Paused) => { - let (trainers, upload_handle) = cooldown.finish().await?; - if let Some(handle) = upload_handle { - self.pending_upload_handles.push(handle); - } + let trainers = cooldown.finish().await?; ActiveStep::Warmup(self.warmup.start( trainers, &mut self.previous_round, @@ -880,11 +929,6 @@ impl StepStateMachine { .endpoint_info = endpoint_info; Ok(()) } - - fn cleanup_completed_uploads(&mut self) { - self.pending_upload_handles - .retain(|handle| !handle.is_finished()); - } } #[derive(Default, Debug)] @@ -1088,20 +1132,6 @@ impl RunManager { } Ok(()) } - - pub fn doing_checkpoint(&self) -> bool { - match &self.0 { - InitStage::Running(step_state_machine) => { - let has_pending_uploads = step_state_machine - .pending_upload_handles - .iter() - .any(|handle| !handle.is_finished()); - - has_pending_uploads - } - _ => false, - } - } } impl From<&RunManager> for ClientTUIState { diff --git a/shared/client/src/state/train.rs b/shared/client/src/state/train.rs index 0993608e8..c349bcb91 100644 --- a/shared/client/src/state/train.rs +++ b/shared/client/src/state/train.rs @@ -506,10 +506,7 @@ impl TrainingStepMetadata { let (cold_start_warmup_steps, checkpoint_is_p2p) = match &state.model { model::Model::LLM(llm) => ( llm.cold_start_warmup_steps, - matches!( - llm.checkpoint, - model::Checkpoint::P2P(_) | model::Checkpoint::P2PGcs(_) - ), + matches!(llm.checkpoint, model::Checkpoint::P2P(_)), ), }; let warmup_lr_between = state.get_cold_start_warmup_bounds(); diff --git a/shared/client/src/state/types.rs b/shared/client/src/state/types.rs index 6385139c1..f17ef811c 100644 --- a/shared/client/src/state/types.rs +++ b/shared/client/src/state/types.rs @@ -1,26 +1,71 @@ use std::path::PathBuf; +use std::sync::Arc; use psyche_coordinator::CommitteeProof; use psyche_core::{BatchId, MerkleRoot, NodeIdentity}; -use psyche_data_provider::{GcsUploadInfo, HubUploadInfo}; +use psyche_data_provider::{HubUploadInfo, RunDownClient}; use psyche_modeling::DistroResult; use psyche_network::{BlobTicket, TransmittableDistroResult}; use tch::TchError; use thiserror::Error; use tokio::task::JoinHandle; +/// Validated checkpoint uploader. Can only be constructed via async methods +/// that validate credentials and permissions on creation. #[derive(Debug, Clone)] -pub enum UploadInfo { +pub enum CheckpointUploader { Hub(HubUploadInfo), - Gcs(GcsUploadInfo), + Gcs(Arc), + Dummy, +} + +impl CheckpointUploader { + /// Creates a new HF Hub uploader after validating write permissions to the repo. + pub async fn new_hub(repo: String, token: String) -> anyhow::Result { + let api = hf_hub::api::tokio::ApiBuilder::new() + .with_token(Some(token.clone())) + .build()?; + let api_repo = api.repo(hf_hub::Repo::model(repo.clone())); + if !api_repo.is_writable().await { + anyhow::bail!( + "Checkpoint upload repo {} is not writable with the provided HF token.", + repo + ); + } + Ok(Self::Hub(HubUploadInfo { + hub_repo: repo, + hub_token: token, + })) + } + + /// Creates a new GCS uploader using run-down signed URLs. + /// Auth is validated at upload time via signed URLs. + pub fn new_gcs(run_down_client: Arc) -> Self { + Self::Gcs(run_down_client) + } } #[derive(Debug, Clone)] pub struct CheckpointConfig { - pub upload_info: Option, pub checkpoint_dir: PathBuf, pub delete_old_steps: bool, pub keep_steps: u32, + pub hub_token: Option, + pub skip_upload: bool, + pub run_down_client: Option>, +} + +impl CheckpointConfig { + pub fn dummy() -> Self { + Self { + checkpoint_dir: PathBuf::from("./checkpoints"), + delete_old_steps: false, + keep_steps: 1, + hub_token: None, + skip_upload: false, + run_down_client: None, + } + } } #[derive(Debug)] diff --git a/shared/client/src/state/warmup.rs b/shared/client/src/state/warmup.rs index 14b123cbf..e5052cf44 100644 --- a/shared/client/src/state/warmup.rs +++ b/shared/client/src/state/warmup.rs @@ -21,13 +21,19 @@ impl WarmupStepMetadata { let evals = self .model_task_runner .start_if_not_running(evals_or_trainers.into()); - WarmupStep { evals } + WarmupStep { + evals, + sent_finished: false, + sent_witness: false, + } } } #[derive(Debug)] pub struct WarmupStep { evals: RunningEvals, + pub sent_finished: bool, + pub sent_witness: bool, } impl WarmupStep { diff --git a/shared/coordinator/Cargo.toml b/shared/coordinator/Cargo.toml index f7cdecc81..024696555 100644 --- a/shared/coordinator/Cargo.toml +++ b/shared/coordinator/Cargo.toml @@ -9,6 +9,7 @@ async-trait.workspace = true anchor-lang.workspace = true bytemuck.workspace = true serde_with.workspace = true +anyhow.workspace = true serde.workspace = true cfg_eval = "0.1.2" ts-rs.workspace = true diff --git a/shared/coordinator/src/checkpointer_selection.rs b/shared/coordinator/src/checkpointer_selection.rs new file mode 100644 index 000000000..64193d12a --- /dev/null +++ b/shared/coordinator/src/checkpointer_selection.rs @@ -0,0 +1,63 @@ +use std::cmp::max; + +use crate::{Coordinator, CoordinatorError, coordinator::SOLANA_MAX_NUM_CHECKPOINTERS}; +use psyche_core::{compute_shuffled_index, sha256, sha256v}; + +use super::types::salts; + +#[derive(Clone)] +pub struct CheckpointerSelection { + checkpointers: u64, + seed: [u8; 32], +} + +impl CheckpointerSelection { + pub fn new(checkpointers: u64, seed: [u8; 32]) -> Self { + Self { + checkpointers, + seed, + } + } + + pub fn from_coordinator( + coordinator: &Coordinator, + offset: isize, + ) -> Result { + let round = get_round_by_offset(coordinator, offset)?; + let seed = sha256(&round.random_seed.to_le_bytes()); + + let checkpointers = max( + (coordinator.epoch_state.clients.len() / 3).min(SOLANA_MAX_NUM_CHECKPOINTERS), + 1, + ) as u64; + Ok(Self { + checkpointers, + seed, + }) + } + + pub fn is_checkpointer(&self, client_index: u64, total_clients: u64) -> bool { + let final_seed = compute_salted_seed(&self.seed, salts::COOLDOWN); + let index = compute_shuffled_index(client_index, total_clients, &final_seed); + index < self.checkpointers + } +} + +pub(crate) fn get_round_by_offset( + coordinator: &Coordinator, + offset: isize, +) -> Result<&crate::Round, CoordinatorError> { + match offset { + -2 => coordinator.previous_previous_round(), + -1 => coordinator.previous_round(), + 0 => coordinator.current_round(), + _ => return Err(CoordinatorError::NoActiveRound), + } + .ok_or(CoordinatorError::NoActiveRound) +} + +pub(crate) fn compute_salted_seed(seed: &[u8; 32], salt: &str) -> [u8; 32] { + let mut result = [0u8; 32]; + result.copy_from_slice(&sha256v(&[&sha256(seed), salt.as_bytes()])); + result +} diff --git a/shared/coordinator/src/committee_selection.rs b/shared/coordinator/src/committee_selection.rs index 98e817c9a..ca1c350ee 100644 --- a/shared/coordinator/src/committee_selection.rs +++ b/shared/coordinator/src/committee_selection.rs @@ -1,84 +1,16 @@ use crate::{Client, Coordinator, CoordinatorError, SOLANA_MAX_NUM_WITNESSES}; +use psyche_core::{NodeIdentity, compute_shuffled_index, sha256}; -use anchor_lang::{AnchorDeserialize, AnchorSerialize, InitSpace, prelude::borsh}; -use bytemuck::Zeroable; -use psyche_core::{NodeIdentity, SmallBoolean, compute_shuffled_index, sha256, sha256v}; -use serde::{Deserialize, Serialize}; -use ts_rs::TS; - -pub const COMMITTEE_SALT: &str = "committee"; -pub const WITNESS_SALT: &str = "witness"; - -#[derive( - Clone, - Copy, - Debug, - Default, - PartialEq, - Zeroable, - AnchorDeserialize, - AnchorSerialize, - Serialize, - Deserialize, -)] -#[repr(C)] -pub enum Committee { - #[default] - TieBreaker, - Verifier, - Trainer, -} +use super::checkpointer_selection::get_round_by_offset; +use super::types::{Committee, CommitteeProof, WitnessProof, salts}; #[derive(Clone)] pub struct CommitteeSelection { - tie_breaker_nodes: u64, - verifier_nodes: u64, - total_nodes: u64, - witness_nodes: u64, - seed: [u8; 32], -} - -#[derive( - Clone, - Copy, - Debug, - Default, - PartialEq, - Zeroable, - AnchorDeserialize, - AnchorSerialize, - Serialize, - Deserialize, -)] -#[repr(C)] -pub struct CommitteeProof { - pub committee: Committee, - pub position: u64, - pub index: u64, -} - -#[derive( - Clone, - Copy, - Debug, - PartialEq, - Zeroable, - Default, - AnchorDeserialize, - AnchorSerialize, - Serialize, - Deserialize, - InitSpace, - TS, -)] -#[repr(C)] -pub struct WitnessProof { - // position in virtual shuffle, as determined by seed - pub position: u64, - // index into epoch_state.clients of sender - pub index: u64, - // assertion of witness membership or non-membership - pub witness: SmallBoolean, + pub(crate) tie_breaker_nodes: u64, + pub(crate) verifier_nodes: u64, + pub(crate) total_nodes: u64, + pub(crate) witness_nodes: u64, + pub(crate) seed: [u8; 32], } impl CommitteeSelection { @@ -89,25 +21,15 @@ impl CommitteeSelection { total_nodes: usize, seed: u64, ) -> Result { - if total_nodes >= u64::MAX as usize { - return Err(CoordinatorError::InvalidCommitteeSelection); - } - - if total_nodes < tie_breaker_nodes { - return Err(CoordinatorError::InvalidCommitteeSelection); - } - - if witness_nodes != 0 && total_nodes < witness_nodes { - return Err(CoordinatorError::InvalidCommitteeSelection); - } - - if verification_percent > 100 { - return Err(CoordinatorError::InvalidCommitteeSelection); - } + Self::validate_params( + tie_breaker_nodes, + witness_nodes, + verification_percent, + total_nodes, + )?; let free_nodes = total_nodes - tie_breaker_nodes; let verifier_nodes = (free_nodes * verification_percent as usize) / 100; - let seed = sha256(&seed.to_le_bytes()); Ok(Self { @@ -119,19 +41,32 @@ impl CommitteeSelection { }) } + fn validate_params( + tie_breaker_nodes: usize, + witness_nodes: usize, + verification_percent: u8, + total_nodes: usize, + ) -> Result<(), CoordinatorError> { + if total_nodes >= u64::MAX as usize { + return Err(CoordinatorError::InvalidCommitteeSelection); + } + if total_nodes < tie_breaker_nodes { + return Err(CoordinatorError::InvalidCommitteeSelection); + } + if witness_nodes != 0 && total_nodes < witness_nodes { + return Err(CoordinatorError::InvalidCommitteeSelection); + } + if verification_percent > 100 { + return Err(CoordinatorError::InvalidCommitteeSelection); + } + Ok(()) + } + pub fn from_coordinator( coordinator: &Coordinator, offset: isize, ) -> Result { - let round = match offset { - -2 => coordinator.previous_previous_round(), - -1 => coordinator.previous_round(), - 0 => coordinator.current_round(), - _ => { - return Err(CoordinatorError::NoActiveRound); - } - } - .ok_or(CoordinatorError::NoActiveRound)?; + let round = get_round_by_offset(coordinator, offset)?; Self::new( round.tie_breaker_tasks as usize, coordinator.config.witness_nodes as usize, @@ -142,8 +77,8 @@ impl CommitteeSelection { } pub fn get_witness(&self, index: u64) -> WitnessProof { - let position = self.compute_shuffled_index(index, WITNESS_SALT); - let witness = self.get_witness_from_position(position); + let position = self.compute_shuffled_index(index, salts::WITNESS); + let witness = self.is_witness_at_position(position); WitnessProof { witness: witness.into(), position, @@ -151,8 +86,29 @@ impl CommitteeSelection { } } + pub fn verify_witness(&self, proof: &WitnessProof) -> bool { + let position = self.compute_shuffled_index(proof.index, salts::WITNESS); + proof.position == position && proof.witness == self.is_witness_at_position(position).into() + } + + pub fn verify_witness_for_client( + &self, + client_id: &NodeIdentity, + proof: &WitnessProof, + clients: &[Client], + ) -> bool { + Self::verify_client(client_id, proof.index, clients) && self.verify_witness(proof) + } + + fn is_witness_at_position(&self, position: u64) -> bool { + match self.witness_nodes { + 0 => position < SOLANA_MAX_NUM_WITNESSES as u64, + witness_nodes => position < witness_nodes, + } + } + pub fn get_committee(&self, index: u64) -> CommitteeProof { - let position = self.compute_shuffled_index(index, COMMITTEE_SALT); + let position = self.compute_shuffled_index(index, salts::COMMITTEE); let committee = self.get_committee_from_position(position); CommitteeProof { committee, @@ -161,21 +117,19 @@ impl CommitteeSelection { } } - pub fn get_committee_from_position(&self, committee_position: u64) -> Committee { - if committee_position < self.tie_breaker_nodes { + pub fn get_committee_from_position(&self, position: u64) -> Committee { + if position < self.tie_breaker_nodes { Committee::TieBreaker - } else if committee_position < self.tie_breaker_nodes + self.verifier_nodes { + } else if position < self.tie_breaker_nodes + self.verifier_nodes { Committee::Verifier } else { Committee::Trainer } } - fn get_witness_from_position(&self, witness_position: u64) -> bool { - match self.witness_nodes { - 0 => witness_position < SOLANA_MAX_NUM_WITNESSES as u64, - witness_nodes => witness_position < witness_nodes, - } + pub fn verify_committee(&self, proof: &CommitteeProof) -> bool { + let position = self.compute_shuffled_index(proof.index, salts::COMMITTEE); + proof.position == position && proof.committee == self.get_committee_from_position(position) } pub fn verify_committee_for_client( @@ -187,34 +141,13 @@ impl CommitteeSelection { Self::verify_client(client_id, proof.index, clients) && self.verify_committee(proof) } - pub fn verify_witness_for_client( - &self, - client_id: &NodeIdentity, - proof: &WitnessProof, - clients: &[Client], - ) -> bool { - Self::verify_client(client_id, proof.index, clients) && self.verify_witness(proof) - } - fn verify_client(client_id: &NodeIdentity, index: u64, clients: &[Client]) -> bool { clients.get(index as usize).map(|c| &c.id) == Some(client_id) } - fn verify_committee(&self, proof: &CommitteeProof) -> bool { - let position = self.compute_shuffled_index(proof.index, COMMITTEE_SALT); - proof.position == position && proof.committee == self.get_committee_from_position(position) - } - - fn verify_witness(&self, proof: &WitnessProof) -> bool { - let position = self.compute_shuffled_index(proof.index, WITNESS_SALT); - proof.position == position - && proof.witness == self.get_witness_from_position(position).into() - } - fn compute_shuffled_index(&self, index: u64, salt: &str) -> u64 { let mut seed = [0u8; 32]; - seed.copy_from_slice(&sha256v(&[&self.seed, salt.as_bytes()])); - + seed.copy_from_slice(&psyche_core::sha256v(&[&self.seed, salt.as_bytes()])); compute_shuffled_index(index, self.total_nodes, &seed) } @@ -235,19 +168,10 @@ impl CommitteeSelection { } } -impl std::fmt::Display for Committee { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Committee::TieBreaker => write!(f, "Tie breaker"), - Committee::Verifier => write!(f, "Verifier"), - Committee::Trainer => write!(f, "Trainer"), - } - } -} - #[cfg(test)] mod tests { use super::*; + use crate::types::{Committee, CommitteeProof, WitnessProof}; #[test] fn test_new_committee_selection() { @@ -385,7 +309,7 @@ mod tests { } #[test] - fn test_invalid_committee_selections() { + fn test_invalid_comittee_selections() { // verification_percent > 100 assert!(CommitteeSelection::new(10, 5, 101, 100, 12345).is_err()); // total_nodes < tie_breaker_nodes diff --git a/shared/coordinator/src/coordinator.rs b/shared/coordinator/src/coordinator.rs index 57dde7710..aa9833a6a 100644 --- a/shared/coordinator/src/coordinator.rs +++ b/shared/coordinator/src/coordinator.rs @@ -1,5 +1,5 @@ use crate::{ - Commitment, Committee, CommitteeProof, CommitteeSelection, WitnessProof, + CheckpointerSelection, Commitment, Committee, CommitteeProof, CommitteeSelection, WitnessProof, model::{Checkpoint, Model}, }; @@ -17,6 +17,7 @@ pub const SOLANA_MAX_STRING_LEN: usize = 64; pub const SOLANA_MAX_URL_STRING_LEN: usize = 192; pub const SOLANA_MAX_NUM_CLIENTS: usize = 256; pub const SOLANA_MAX_NUM_WITNESSES: usize = 32; +pub const SOLANA_MAX_NUM_CHECKPOINTERS: usize = 16; // run_id must be at most 32 bytes because of PDA constraints pub const SOLANA_RUN_ID_MAX_LEN: usize = 32; @@ -282,6 +283,7 @@ pub struct CoordinatorEpochState { pub start_timestamp: u64, pub first_round: SmallBoolean, pub cold_start_epoch: SmallBoolean, + pub checkpointed: bool, } #[derive( @@ -409,6 +411,7 @@ impl Default for CoordinatorEpochState { start_step: Default::default(), last_step: Default::default(), start_timestamp: Default::default(), + checkpointed: false, } } } @@ -501,6 +504,37 @@ impl Coordinator { Ok(()) } + pub fn cooldown_witness( + &mut self, + from: &NodeIdentity, + witness: Witness, + ) -> std::result::Result<(), CoordinatorError> { + if self.halted() { + return Err(CoordinatorError::Halted); + } + + if !matches!(self.run_state, RunState::Cooldown) { + return Ok(()); + } + + // Verify the sender matches the witness index to prevent spoofing + let index = witness.proof.index as usize; + if index >= self.epoch_state.clients.len() || self.epoch_state.clients[index].id != *from { + return Err(CoordinatorError::InvalidWitness); + } + + let checkpointer_selection = CheckpointerSelection::from_coordinator(self, 0)?; + if !checkpointer_selection + .is_checkpointer(witness.proof.index, self.epoch_state.clients.len() as u64) + { + return Err(CoordinatorError::InvalidWitness); + } + + self.epoch_state.checkpointed = true; + + Ok(()) + } + pub fn witness( &mut self, from: &NodeIdentity, @@ -597,36 +631,40 @@ impl Coordinator { return Err(CoordinatorError::InvalidCommitteeProof); } - // TODO: In the case of more than one checkpointer, this will overwrite the checkpoint - // with the last checkpointed one. We could instead have a vector of checkpoints to have - // more download options. + if self.halted() { + return Err(CoordinatorError::Halted); + } + + if !matches!(self.run_state, RunState::Cooldown) { + return Err(CoordinatorError::InvalidRunState); + } + + let checkpointer_selection = CheckpointerSelection::from_coordinator(self, 0)?; + if !checkpointer_selection + .is_checkpointer(index as u64, self.epoch_state.clients.len() as u64) + { + return Err(CoordinatorError::InvalidWitness); + } + let Model::LLM(llm) = &mut self.model; match (&llm.checkpoint, checkpoint_repo) { - // If current is P2P, wrap the new checkpoint in P2P - (Checkpoint::P2P(_), Checkpoint::Hub(hub_repo)) => { - llm.checkpoint = Checkpoint::P2P(hub_repo); - } - (Checkpoint::P2PGcs(_), Checkpoint::Gcs(gcs_repo)) => { - llm.checkpoint = Checkpoint::P2PGcs(gcs_repo); - } - // If current is Hub, only accept Hub updates - (Checkpoint::Hub(_), Checkpoint::Hub(hub_repo)) => { - llm.checkpoint = Checkpoint::Hub(hub_repo); - } - // If current is Gcs, only accept Gcs updates - (Checkpoint::Gcs(_), Checkpoint::Gcs(gcs_repo)) => { - llm.checkpoint = Checkpoint::Gcs(gcs_repo); - } - (Checkpoint::P2PGcs(_), Checkpoint::Hub(hub_repo)) => { - llm.checkpoint = Checkpoint::P2P(hub_repo); + // If current is P2P, wrap the new checkpoint in its P2P equivalent + (Checkpoint::P2P(_) | Checkpoint::P2PGcs(_), new) => { + llm.checkpoint = new.to_p2p(); } - (Checkpoint::P2P(_), Checkpoint::Gcs(gcs_repo)) => { - llm.checkpoint = Checkpoint::P2PGcs(gcs_repo); + // If current is hosted (Hub/Gcs), accept hosted updates directly + ( + Checkpoint::Hub(_) | Checkpoint::Gcs(_), + new @ (Checkpoint::Hub(_) | Checkpoint::Gcs(_)), + ) => { + llm.checkpoint = new; } // Ignore other combinations _ => {} } + self.epoch_state.checkpointed = true; + Ok(()) } @@ -952,11 +990,7 @@ impl Coordinator { .any(|client| pending_clients_unordered.contains(&client.id)); if all_prev_clients_disconnected { let Model::LLM(llm) = &mut self.model; - match llm.checkpoint { - Checkpoint::P2P(hub_repo) => llm.checkpoint = Checkpoint::Hub(hub_repo), - Checkpoint::P2PGcs(gcs_repo) => llm.checkpoint = Checkpoint::Gcs(gcs_repo), - _ => {} - } + llm.checkpoint = llm.checkpoint.to_hosted(); } let cold_start_epoch = self.epoch_state.cold_start_epoch; @@ -1071,7 +1105,9 @@ impl Coordinator { &mut self, unix_timestamp: u64, ) -> std::result::Result { - if self.check_timeout(unix_timestamp, self.config.cooldown_time) { + if self.check_timeout(unix_timestamp, self.config.cooldown_time) + || self.epoch_state.checkpointed + { let last_round_batch_size = self.get_target_global_batch_size(self.current_round()); self.progress.epoch_start_data_index = self.current_round_unchecked().data_index + last_round_batch_size as u64; @@ -1083,13 +1119,7 @@ impl Coordinator { // we've completed an epoch, switch to P2P from now on let Model::LLM(llm) = &mut self.model; - match llm.checkpoint { - Checkpoint::Hub(hub_repo) | Checkpoint::Dummy(hub_repo) => { - llm.checkpoint = Checkpoint::P2P(hub_repo) - } - Checkpoint::Gcs(gcs_repo) => llm.checkpoint = Checkpoint::P2PGcs(gcs_repo), - _ => {} - } + llm.checkpoint = llm.checkpoint.to_p2p(); if self.pending_pause.is_true() { self.withdraw_all(); diff --git a/shared/coordinator/src/lib.rs b/shared/coordinator/src/lib.rs index bef26863e..e09bccef5 100644 --- a/shared/coordinator/src/lib.rs +++ b/shared/coordinator/src/lib.rs @@ -1,15 +1,16 @@ #![allow(unexpected_cfgs)] +mod checkpointer_selection; mod commitment; mod committee_selection; mod coordinator; mod data_selection; pub mod model; +mod types; +pub use checkpointer_selection::CheckpointerSelection; pub use commitment::Commitment; -pub use committee_selection::{ - COMMITTEE_SALT, Committee, CommitteeProof, CommitteeSelection, WITNESS_SALT, WitnessProof, -}; +pub use committee_selection::CommitteeSelection; pub use coordinator::{ BLOOM_FALSE_RATE, Client, ClientState, Coordinator, CoordinatorConfig, CoordinatorEpochState, CoordinatorError, CoordinatorProgress, HealthChecks, MAX_TOKENS_TO_SEND, NUM_STORED_ROUNDS, @@ -20,3 +21,4 @@ pub use coordinator::{ pub use data_selection::{ assign_data_for_state, get_batch_ids_for_node, get_batch_ids_for_round, get_data_index_for_step, }; +pub use types::{Committee, CommitteeProof, WitnessProof, salts}; diff --git a/shared/coordinator/src/model.rs b/shared/coordinator/src/model.rs index 3176f276e..edb0140da 100644 --- a/shared/coordinator/src/model.rs +++ b/shared/coordinator/src/model.rs @@ -287,18 +287,63 @@ pub enum Checkpoint { P2PGcs(GcsRepo), } +impl Checkpoint { + /// Returns the HubRepo if this is a Hub or P2P checkpoint. + pub fn hub_repo(&self) -> Option<&HubRepo> { + match self { + Checkpoint::Hub(repo) | Checkpoint::P2P(repo) | Checkpoint::Dummy(repo) => Some(repo), + _ => None, + } + } + + /// Returns the GcsRepo if this is a Gcs or P2PGcs checkpoint. + pub fn gcs_repo(&self) -> Option<&GcsRepo> { + match self { + Checkpoint::Gcs(repo) | Checkpoint::P2PGcs(repo) => Some(repo), + _ => None, + } + } + + /// Returns true if this checkpoint uses P2P model sharing. + pub fn is_p2p(&self) -> bool { + matches!(self, Checkpoint::P2P(_) | Checkpoint::P2PGcs(_)) + } + + /// Converts a hosted checkpoint to its P2P variant. + pub fn to_p2p(self) -> Self { + match self { + Checkpoint::Hub(repo) | Checkpoint::Dummy(repo) => Checkpoint::P2P(repo), + Checkpoint::Gcs(repo) => Checkpoint::P2PGcs(repo), + other => other, + } + } + + /// Converts a P2P checkpoint back to its hosted variant. + pub fn to_hosted(self) -> Self { + match self { + Checkpoint::P2P(repo) => Checkpoint::Hub(repo), + Checkpoint::P2PGcs(repo) => Checkpoint::Gcs(repo), + other => other, + } + } +} + impl std::fmt::Display for Checkpoint { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { - Checkpoint::Dummy(_hub_repo) => write!(f, "Dummy"), + Checkpoint::Dummy(_) => write!(f, "Dummy"), Checkpoint::Ephemeral => write!(f, "Ephemeral"), Checkpoint::Hub(hub_repo) => write!(f, "{}", &hub_repo.repo_id), + Checkpoint::Gcs(gcs_repo) => match &gcs_repo.prefix { + Some(prefix) => write!(f, "gs://{}/{}", &gcs_repo.bucket, prefix), + None => write!(f, "gs://{}", &gcs_repo.bucket), + }, Checkpoint::P2P(hub_repo) => { write!(f, "P2P - Hub repo: {}", &hub_repo.repo_id) } - Checkpoint::Gcs(gcs_repo) | Checkpoint::P2PGcs(gcs_repo) => match &gcs_repo.prefix { - Some(prefix) => write!(f, "gs://{}/{}", &gcs_repo.bucket, prefix), - None => write!(f, "gs://{}", &gcs_repo.bucket), + Checkpoint::P2PGcs(gcs_repo) => match &gcs_repo.prefix { + Some(prefix) => write!(f, "P2P - gs://{}/{}", &gcs_repo.bucket, prefix), + None => write!(f, "P2P - gs://{}", &gcs_repo.bucket), }, } } @@ -336,11 +381,11 @@ impl Model { return false; } let bad_checkpoint = match llm.checkpoint { - Checkpoint::Dummy(_hub_repo) => false, - Checkpoint::Ephemeral => true, - Checkpoint::Hub(hub_repo) => hub_repo.repo_id.is_empty(), - Checkpoint::P2P(hub_repo) => hub_repo.repo_id.is_empty(), - Checkpoint::Gcs(gcs_repo) | Checkpoint::P2PGcs(gcs_repo) => { + Checkpoint::Dummy(_) | Checkpoint::Ephemeral => false, + Checkpoint::Hub(ref hub_repo) | Checkpoint::P2P(ref hub_repo) => { + hub_repo.repo_id.is_empty() + } + Checkpoint::Gcs(ref gcs_repo) | Checkpoint::P2PGcs(ref gcs_repo) => { gcs_repo.bucket.is_empty() } }; diff --git a/shared/coordinator/src/types.rs b/shared/coordinator/src/types.rs new file mode 100644 index 000000000..ac49815d7 --- /dev/null +++ b/shared/coordinator/src/types.rs @@ -0,0 +1,85 @@ +use anchor_lang::{AnchorDeserialize, AnchorSerialize, InitSpace, prelude::borsh}; +use bytemuck::Zeroable; +use psyche_core::SmallBoolean; +use serde::{Deserialize, Serialize}; +use ts_rs::TS; + +/// Salt constants for deterministic shuffling +pub mod salts { + pub const COMMITTEE: &str = "committee"; + pub const WITNESS: &str = "witness"; + pub const COOLDOWN: &str = "cooldown"; +} + +#[derive( + Clone, + Copy, + Debug, + Default, + PartialEq, + Zeroable, + AnchorDeserialize, + AnchorSerialize, + Serialize, + Deserialize, +)] +#[repr(C)] +pub enum Committee { + #[default] + TieBreaker, + Verifier, + Trainer, +} + +impl std::fmt::Display for Committee { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Committee::TieBreaker => write!(f, "Tie breaker"), + Committee::Verifier => write!(f, "Verifier"), + Committee::Trainer => write!(f, "Trainer"), + } + } +} + +#[derive( + Clone, + Copy, + Debug, + Default, + PartialEq, + Zeroable, + AnchorDeserialize, + AnchorSerialize, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct CommitteeProof { + pub committee: Committee, + pub position: u64, + pub index: u64, +} + +#[derive( + Clone, + Copy, + Debug, + Default, + PartialEq, + Zeroable, + AnchorDeserialize, + AnchorSerialize, + Serialize, + Deserialize, + InitSpace, + TS, +)] +#[repr(C)] +pub struct WitnessProof { + /// Position in virtual shuffle, as determined by seed + pub position: u64, + /// Index into epoch_state.clients of sender + pub index: u64, + /// Assertion of witness membership or non-membership + pub witness: SmallBoolean, +} diff --git a/shared/data-provider/Cargo.toml b/shared/data-provider/Cargo.toml index 7ad69e616..5637bd1e8 100644 --- a/shared/data-provider/Cargo.toml +++ b/shared/data-provider/Cargo.toml @@ -25,8 +25,13 @@ serde.workspace = true thiserror.workspace = true postcard.workspace = true bytemuck.workspace = true +google-cloud-storage.workspace = true reqwest = "0.12.12" -google-cloud-storage = "0.24.0" +bytes = "1" +bs58 = "0.5" +google-cloud-auth = "0.16" +google-cloud-gax = "1.4.0" +urlencoding = "2.1.3" chrono = { version = "0.4", features = ["serde"] } serde_json.workspace = true ts-rs.workspace = true diff --git a/shared/data-provider/examples/tcp.rs b/shared/data-provider/examples/tcp.rs index 904cefa60..d8ffd354e 100644 --- a/shared/data-provider/examples/tcp.rs +++ b/shared/data-provider/examples/tcp.rs @@ -2,7 +2,8 @@ use anyhow::{Result, bail}; use async_trait::async_trait; use bytemuck::Zeroable; use futures::future::try_join_all; -use psyche_coordinator::{Coordinator, HealthChecks, model}; +use parquet::data_type::AsBytes; +use psyche_coordinator::{Coordinator, HealthChecks, model::Checkpoint}; use psyche_core::BatchId; use psyche_data_provider::{ DataProviderTcpClient, DataProviderTcpServer, LengthKnownDataProvider, TokenizedData, @@ -32,7 +33,7 @@ impl WatcherBackend for DummyBackend { bail!("Data provider does not send health check"); } - async fn send_checkpoint(&mut self, _checkpoint: model::Checkpoint) -> anyhow::Result<()> { + async fn send_checkpoint(&mut self, _checkpoint: Checkpoint) -> anyhow::Result<()> { bail!("Data provider does not send checkpoints"); } } diff --git a/shared/data-provider/src/errors.rs b/shared/data-provider/src/errors.rs index b84bc5f9a..a12a357a8 100644 --- a/shared/data-provider/src/errors.rs +++ b/shared/data-provider/src/errors.rs @@ -1,3 +1,4 @@ +use hf_hub::api::tokio::CommitError; use std::path::PathBuf; use thiserror::Error; @@ -9,45 +10,42 @@ pub enum UploadError { #[error("file {0} doesn't have a valid utf-8 representation")] InvalidFilename(PathBuf), - #[error("failed to send checkpoint notification")] - SendCheckpoint, - - // Hub-specific errors - #[error("failed to connect to HF hub: {0}")] - HfHub(#[from] hf_hub::api::tokio::ApiError), - - #[error("failed to commit files: {0}")] - Commit(#[from] hf_hub::api::tokio::CommitError), - - // GCS-specific errors #[error("GCS authentication failed: {0}")] - GcsAuth(#[from] google_cloud_storage::client::google_cloud_auth::error::Error), + GcsAuth(String), - #[error("GCS operation failed: {0}")] - GcsStorage(#[from] google_cloud_storage::http::Error), - - // Common errors #[error("IO error: {0}")] Io(#[from] std::io::Error), + #[error("GCS error: {0}")] + Gcs(String), + + #[error("HuggingFace Hub API error: {0}")] + HubApi(#[from] hf_hub::api::tokio::ApiError), + + #[error("HuggingFace Hub commit error: {0}")] + HubCommit(#[from] CommitError), + #[error("JSON error: {0}")] Json(#[from] serde_json::Error), + + #[error("run-down service error: {0}")] + RunDown(String), } #[derive(Error, Debug)] pub enum DownloadError { - #[error("failed to connect to HF hub: {0}")] - HfHub(#[from] hf_hub::api::tokio::ApiError), - #[error("GCS authentication failed: {0}")] - GcsAuth(#[from] google_cloud_storage::client::google_cloud_auth::error::Error), - - #[error("GCS operation failed: {0}")] - GcsStorage(#[from] google_cloud_storage::http::Error), + GcsAuth(String), #[error("IO error: {0}")] Io(#[from] std::io::Error), + #[error("GCS error: {0}")] + Gcs(String), + #[error("JSON error: {0}")] Json(#[from] serde_json::Error), + + #[error("run-down service error: {0}")] + RunDown(String), } diff --git a/shared/data-provider/src/gcs.rs b/shared/data-provider/src/gcs.rs index 71f29e414..29df528e2 100644 --- a/shared/data-provider/src/gcs.rs +++ b/shared/data-provider/src/gcs.rs @@ -1,19 +1,11 @@ use crate::errors::{DownloadError, UploadError}; use chrono::{DateTime, Utc}; -use google_cloud_storage::client::{Client, ClientConfig}; -use google_cloud_storage::http::objects::upload::Media; -use google_cloud_storage::http::objects::upload::UploadObjectRequest; -use google_cloud_storage::http::objects::upload::UploadType; -use google_cloud_storage::http::objects::{ - download::Range, get::GetObjectRequest, list::ListObjectsRequest, -}; -use psyche_coordinator::model::{self, GcsRepo}; -use psyche_core::FixedString; +use google_cloud_gax::paginator::ItemPaginator; +use google_cloud_storage::client::{Storage, StorageControl}; use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; use tokio::runtime::Runtime; -use tokio::sync::mpsc; -use tracing::info; +use tracing::{debug, info}; /// Checkpoint manifest.json uploaded to GCS alongside safetensors files. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -51,9 +43,9 @@ pub struct GcsManifestMetadata { pub run_id: String, } -const MODEL_EXTENSIONS: [&str; 3] = [".safetensors", ".json", ".py"]; +pub(crate) const MODEL_EXTENSIONS: [&str; 3] = [".safetensors", ".json", ".py"]; -fn get_cache_base(bucket: &str) -> PathBuf { +pub(crate) fn get_cache_base(bucket: &str) -> PathBuf { // Use HF_HOME if set, otherwise fall back to ~/.cache std::env::var("HF_HOME") .map(PathBuf::from) @@ -67,7 +59,7 @@ fn get_cache_base(bucket: &str) -> PathBuf { .join(bucket) } -fn get_cache_dir( +pub(crate) fn get_cache_dir( bucket: &str, prefix: Option<&str>, step: u32, @@ -82,7 +74,7 @@ fn get_cache_dir( } } -fn get_cache_dir_no_manifest(bucket: &str, prefix: Option<&str>) -> PathBuf { +pub(crate) fn get_cache_dir_no_manifest(bucket: &str, prefix: Option<&str>) -> PathBuf { let base = get_cache_base(bucket); match prefix { @@ -91,7 +83,7 @@ fn get_cache_dir_no_manifest(bucket: &str, prefix: Option<&str>) -> PathBuf { } } -fn collect_cached_files( +pub(crate) fn collect_cached_files( cache_dir: &Path, manifest: &GcsCheckpointManifest, ) -> Option> { @@ -110,48 +102,43 @@ pub async fn download_model_from_gcs_async( bucket: &str, prefix: Option<&str>, ) -> Result, DownloadError> { - // Use authenticated client if GOOGLE_APPLICATION_CREDENTIALS is set, otherwise anonymous - let config = if std::env::var("GOOGLE_APPLICATION_CREDENTIALS").is_ok() { - info!("Using authenticated GCS client"); - ClientConfig::default().with_auth().await? - } else { - info!("Using anonymous GCS client"); - ClientConfig::default().anonymous() - }; - let client = Client::new(config); + // Automatically handles authentication via GOOGLE_APPLICATION_CREDENTIALS + let storage = Storage::builder() + .build() + .await + .map_err(|e| DownloadError::Gcs(e.to_string()))?; + + let storage_control = StorageControl::builder() + .build() + .await + .map_err(|e| DownloadError::Gcs(e.to_string()))?; let manifest_object_path = match prefix { Some(p) => format!("{}/manifest.json", p), None => "manifest.json".to_string(), }; - // Get manifest metadata to obtain generation number - let manifest_metadata = client - .get_object(&GetObjectRequest { - bucket: bucket.to_owned(), - object: manifest_object_path.clone(), - ..Default::default() - }) + // Try to get manifest - first check if it exists + let bucket_resource_name = format!("projects/_/buckets/{}", bucket); + let manifest_result = storage + .read_object(&bucket_resource_name, &manifest_object_path) + .send() .await; - match manifest_metadata { - Ok(object_meta) => { - let manifest_generation = object_meta.generation; - - // Download manifest content - let manifest_data = client - .download_object( - &GetObjectRequest { - bucket: bucket.to_owned(), - object: manifest_object_path, - ..Default::default() - }, - &Range::default(), - ) - .await?; + match manifest_result { + Ok(mut read_response) => { + // Read manifest content + let mut manifest_data = Vec::new(); + while let Some(chunk_result) = read_response.next().await { + let chunk = chunk_result.map_err(|e| DownloadError::Gcs(e.to_string()))?; + manifest_data.extend_from_slice(&chunk); + } let manifest: GcsCheckpointManifest = serde_json::from_slice(&manifest_data)?; + // Use step as generation proxy (1.5.x doesn't expose generation in same way) + let manifest_generation = manifest.metadata.step as i64; + info!( "Found manifest: step {}, epoch {}, generation {}", manifest.metadata.step, manifest.metadata.epoch, manifest_generation @@ -171,12 +158,19 @@ pub async fn download_model_from_gcs_async( cache_dir ); std::fs::create_dir_all(&cache_dir)?; - download_files_from_manifest(&client, bucket, prefix, &cache_dir, &manifest).await? + download_files_from_manifest(&storage, bucket, prefix, &cache_dir, &manifest) + .await? }; // Download config files (json, py) - skips if already cached - let config_files = - download_files_no_manifest(&client, bucket, prefix, &cache_dir, &[".json", ".py"]) - .await?; + let config_files = download_files_no_manifest( + &storage_control, + &storage, + bucket, + prefix, + &cache_dir, + &[".json", ".py"], + ) + .await?; files.extend(config_files); Ok(files) } @@ -185,19 +179,28 @@ pub async fn download_model_from_gcs_async( info!("No manifest found, downloading model without manifest"); let cache_dir = get_cache_dir_no_manifest(bucket, prefix); std::fs::create_dir_all(&cache_dir)?; - download_files_no_manifest(&client, bucket, prefix, &cache_dir, &MODEL_EXTENSIONS).await + download_files_no_manifest( + &storage_control, + &storage, + bucket, + prefix, + &cache_dir, + &MODEL_EXTENSIONS, + ) + .await } } } async fn download_files_from_manifest( - client: &Client, + storage: &Storage, bucket: &str, prefix: Option<&str>, cache_dir: &Path, manifest: &GcsCheckpointManifest, ) -> Result, DownloadError> { let mut downloaded_files = Vec::new(); + let bucket_resource_name = format!("projects/_/buckets/{}", bucket); for file_entry in &manifest.files { let object_name = match prefix { @@ -217,17 +220,17 @@ async fn download_files_from_manifest( bucket, object_name, file_entry.generation ); - let data = client - .download_object( - &GetObjectRequest { - bucket: bucket.to_owned(), - object: object_name, - generation: Some(file_entry.generation), - ..Default::default() - }, - &Range::default(), - ) - .await?; + let mut read_response = storage + .read_object(&bucket_resource_name, &object_name) + .send() + .await + .map_err(|e| DownloadError::Gcs(e.to_string()))?; + + let mut data = Vec::new(); + while let Some(chunk_result) = read_response.next().await { + let chunk = chunk_result.map_err(|e| DownloadError::Gcs(e.to_string()))?; + data.extend_from_slice(&chunk); + } std::fs::write(&local_path, &data)?; info!("Downloaded: {} ({} bytes)", file_entry.filename, data.len()); @@ -240,34 +243,34 @@ async fn download_files_from_manifest( /// Download model files by listing the bucket. Skips files that already exist in cache. /// Used for initial model download (no manifest) and to fetch config files (json, py) after manifest download. async fn download_files_no_manifest( - client: &Client, + storage_control: &StorageControl, + storage: &Storage, bucket: &str, prefix: Option<&str>, cache_dir: &Path, extensions: &[&str], ) -> Result, DownloadError> { let mut all_objects = vec![]; - let mut page_token: Option = None; - - loop { - let results = client - .list_objects(&ListObjectsRequest { - bucket: bucket.to_owned(), - prefix: prefix.map(|s| s.to_owned()), - page_token: page_token.clone(), - ..Default::default() - }) - .await?; - for obj in results.items.iter().flatten() { - if extensions.iter().any(|ext| obj.name.ends_with(ext)) { - all_objects.push(obj.name.clone()); - } - } + let parent_name = format!("projects/_/buckets/{}", bucket); + debug!( + "Listing objects in GCS bucket: {}, parent: {}", + bucket, parent_name + ); + let mut list_request = storage_control.list_objects().set_parent(parent_name); + if let Some(p) = prefix { + list_request = list_request.set_prefix(p.to_string()); + } - match results.next_page_token { - Some(token) => page_token = Some(token), - None => break, + let mut stream = list_request.by_item(); + while let Some(obj) = stream + .next() + .await + .transpose() + .map_err(|e| DownloadError::Gcs(e.to_string()))? + { + if extensions.iter().any(|ext| obj.name.ends_with(ext)) { + all_objects.push(obj.name); } } @@ -293,22 +296,21 @@ async fn download_files_no_manifest( info!("Downloading: gs://{}/{}", bucket, object_name); - let data = client - .download_object( - &GetObjectRequest { - bucket: bucket.to_owned(), - object: object_name.clone(), - ..Default::default() - }, - &Range::default(), - ) - .await?; + let bucket_resource_name = format!("projects/_/buckets/{}", bucket); + let mut read_response = storage + .read_object(&bucket_resource_name, &object_name) + .send() + .await + .map_err(|e| DownloadError::Gcs(e.to_string()))?; + + let mut data = Vec::new(); + while let Some(chunk_result) = read_response.next().await { + let chunk = chunk_result.map_err(|e| DownloadError::Gcs(e.to_string()))?; + data.extend_from_slice(&chunk); + } - // Write to cache std::fs::write(&local_path, &data)?; - info!("Downloaded: {} ({} bytes)", filename, data.len()); - downloaded_files.push(local_path); } @@ -328,120 +330,102 @@ pub async fn upload_to_gcs( manifest_metadata: GcsManifestMetadata, local: Vec, step: u64, - tx_checkpoint: mpsc::UnboundedSender, + cancellation_token: tokio_util::sync::CancellationToken, ) -> Result<(), UploadError> { - let GcsUploadInfo { - gcs_bucket, - gcs_prefix, - } = gcs_info; - - let GcsManifestMetadata { epoch, run_id } = manifest_metadata; - - info!(bucket = gcs_bucket, "Uploading checkpoint to GCS"); - - let config = if std::env::var("GOOGLE_APPLICATION_CREDENTIALS").is_ok() { - info!("Using authenticated GCS client"); - ClientConfig::default().with_auth().await? - } else { - info!("Using anonymous GCS client"); - ClientConfig::default().anonymous() - }; - let client = Client::new(config); + let storage = Storage::builder() + .build() + .await + .map_err(|e| UploadError::Gcs(e.to_string()))?; let mut manifest = GcsCheckpointManifest { metadata: ManifestMetadata { timestamp: Utc::now(), - epoch, + epoch: manifest_metadata.epoch, step: step as u32, - run_id, + run_id: manifest_metadata.run_id, }, files: Vec::new(), }; - for path in local { + for path in local + .iter() + .filter(|p| p.extension() == Some("safetensors".as_ref())) + { + if cancellation_token.is_cancelled() { + info!("Upload cancelled before uploading {}", path.display()); + return Ok(()); + } + let file_name = path .file_name() - .ok_or_else(|| UploadError::NotAFile(path.clone()))? - .to_str() + .and_then(|n| n.to_str()) .ok_or_else(|| UploadError::InvalidFilename(path.clone()))?; - - // Only upload safetensors files - if !file_name.ends_with(".safetensors") { - continue; - } - - let object_name = match &gcs_prefix { - Some(p) => format!("{}/{}", p, file_name), + let object_name = match &gcs_info.gcs_prefix { + Some(p) => format!("{}/{}", p.trim_end_matches('/'), file_name), None => file_name.to_string(), }; + let bucket_resource_name = format!("projects/_/buckets/{}", gcs_info.gcs_bucket); - let size = std::fs::metadata(&path)?.len(); - let data = tokio::fs::read(&path).await?; - - let upload_type = UploadType::Simple(Media::new(object_name.clone())); - let uploaded = client - .upload_object( - &UploadObjectRequest { - bucket: gcs_bucket.clone(), - ..Default::default() - }, - data, - &upload_type, - ) - .await?; + let data_vec = tokio::fs::read(&path).await?; + let size = data_vec.len() as u64; + let data = bytes::Bytes::from(data_vec); + + let upload_future = storage + .write_object(&bucket_resource_name, &object_name, data) + .send_unbuffered(); + + let uploaded_file = tokio::select! { + biased; + + _ = cancellation_token.cancelled() => { + info!("Upload cancelled during upload of {}", path.display()); + return Ok(()); + } + result = upload_future => { + result.map_err(|e| UploadError::Gcs(e.to_string()))? + } + }; info!( - bucket = gcs_bucket, + bucket = gcs_info.gcs_bucket, object = object_name, - size = uploaded.size, - generation = uploaded.generation, - "Uploaded file to GCS" + size = uploaded_file.size, + "Successfully uploaded file to GCS" ); manifest.files.push(ManifestFileEntry { filename: file_name.to_string(), - generation: uploaded.generation, + generation: uploaded_file.generation, size_bytes: size, }); } // Upload the manifest file - let manifest_path = match &gcs_prefix { + let manifest_path = match &gcs_info.gcs_prefix { Some(p) => format!("{}/manifest.json", p), None => "manifest.json".to_string(), }; let manifest_json = serde_json::to_string_pretty(&manifest)?; + let manifest_bytes = bytes::Bytes::from(manifest_json.into_bytes()); - let upload_type = UploadType::Simple(Media::new(manifest_path.clone())); - client - .upload_object( - &UploadObjectRequest { - bucket: gcs_bucket.clone(), - ..Default::default() - }, - manifest_json.into_bytes(), - &upload_type, - ) - .await?; + let bucket_resource_name = format!("projects/_/buckets/{}", gcs_info.gcs_bucket); + storage + .write_object(&bucket_resource_name, &manifest_path, manifest_bytes) + .send_unbuffered() + .await + .map_err(|e| UploadError::Gcs(e.to_string()))?; info!( - bucket = gcs_bucket, + bucket = gcs_info.gcs_bucket, object = manifest_path, "Uploaded manifest to GCS" ); info!( "Upload to GCS complete at gs://{}/{}", - gcs_bucket, - gcs_prefix.as_deref().unwrap_or("") + gcs_info.gcs_bucket, + gcs_info.gcs_prefix.as_deref().unwrap_or("") ); - tx_checkpoint - .send(model::Checkpoint::Gcs(GcsRepo { - bucket: FixedString::from_str_truncated(&gcs_bucket), - prefix: gcs_prefix.map(|p| FixedString::from_str_truncated(&p)), - })) - .map_err(|_| UploadError::SendCheckpoint)?; - Ok(()) } diff --git a/shared/data-provider/src/gcs_signed.rs b/shared/data-provider/src/gcs_signed.rs new file mode 100644 index 000000000..829dbbc8a --- /dev/null +++ b/shared/data-provider/src/gcs_signed.rs @@ -0,0 +1,428 @@ +use crate::errors::{DownloadError, UploadError}; +use crate::gcs::{ + GcsCheckpointManifest, GcsManifestMetadata, MODEL_EXTENSIONS, ManifestFileEntry, + ManifestMetadata, collect_cached_files, get_cache_dir, get_cache_dir_no_manifest, +}; +use crate::run_down::{DownloadUrlEntry, RunDownClient}; +use chrono::Utc; +use futures::TryStreamExt; +use std::path::{Path, PathBuf}; +use tokio::io::AsyncWriteExt; +use tracing::{info, warn}; + +pub async fn upload_to_gcs_signed( + run_down: &RunDownClient, + manifest_metadata: GcsManifestMetadata, + local: Vec, + step: u64, + cancellation_token: tokio_util::sync::CancellationToken, +) -> Result<(), UploadError> { + let http = reqwest::Client::new(); + + let mut manifest = GcsCheckpointManifest { + metadata: ManifestMetadata { + timestamp: Utc::now(), + epoch: manifest_metadata.epoch, + step: step as u32, + run_id: manifest_metadata.run_id, + }, + files: Vec::new(), + }; + + for path in local + .iter() + .filter(|p| p.extension() == Some("safetensors".as_ref())) + { + if cancellation_token.is_cancelled() { + info!("Upload cancelled before uploading {}", path.display()); + return Ok(()); + } + + let file_name = path + .file_name() + .and_then(|n| n.to_str()) + .ok_or_else(|| UploadError::InvalidFilename(path.clone()))?; + + let file = tokio::fs::File::open(&path).await?; + let size = file.metadata().await?.len(); + + let upload_url = run_down + .get_upload_url(file_name) + .await + .map_err(|e| UploadError::RunDown(e.to_string()))?; + + info!(file = file_name, size, "Uploading file via signed URL"); + + let upload_future = http + .put(&upload_url.url) + .header("Content-Type", "application/octet-stream") + .header("Content-Length", size) + .body(reqwest::Body::from(file)) + .send(); + + let response = tokio::select! { + biased; + + _ = cancellation_token.cancelled() => { + info!("Upload cancelled during upload of {}", path.display()); + return Ok(()); + } + result = upload_future => { + result.map_err(|e| UploadError::RunDown(e.to_string()))? + } + }; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response.text().await.unwrap_or_default(); + return Err(UploadError::RunDown(format!( + "Signed URL upload failed for {}: {} {}", + file_name, status, error_text + ))); + } + + let generation = match response + .headers() + .get("x-goog-generation") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + { + Some(g) => g, + None => { + warn!( + file = file_name, + "x-goog-generation header missing or invalid in upload response, using 0" + ); + 0 + } + }; + + info!( + file = file_name, + size, generation, "Successfully uploaded file via signed URL" + ); + + manifest.files.push(ManifestFileEntry { + filename: file_name.to_string(), + generation, + size_bytes: size, + }); + } + + let manifest_json = serde_json::to_string_pretty(&manifest)?; + let manifest_bytes = manifest_json.into_bytes(); + + let manifest_upload_url = run_down + .get_upload_url("manifest.json") + .await + .map_err(|e| UploadError::RunDown(e.to_string()))?; + + let response = http + .put(&manifest_upload_url.url) + .header("Content-Type", "application/json") + .body(manifest_bytes) + .send() + .await + .map_err(|e| UploadError::RunDown(e.to_string()))?; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response.text().await.unwrap_or_default(); + return Err(UploadError::RunDown(format!( + "Signed URL upload failed for manifest.json: {} {}", + status, error_text + ))); + } + + info!( + run_id = run_down.run_id(), + "Upload via signed URLs complete" + ); + + Ok(()) +} + +pub async fn download_model_from_gcs_signed_async( + run_down: &RunDownClient, +) -> Result, DownloadError> { + let http = reqwest::Client::new(); + let run_id = run_down.run_id(); + + let download_response = run_down + .get_download_urls() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))?; + + info!( + "Got {} download URLs from run-down for run {}", + download_response.urls.len(), + run_id + ); + + let manifest_entry = download_response + .urls + .iter() + .find(|e| e.path.ends_with("manifest.json")); + + let cache_key = run_id; + + match manifest_entry { + Some(manifest_entry) => { + let response = http + .get(&manifest_entry.url) + .send() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))?; + + if !response.status().is_success() { + return Err(DownloadError::RunDown(format!( + "Failed to download manifest.json: {}", + response.status() + ))); + } + + // Get GCS generation number from manifest response + let manifest_generation = response + .headers() + .get("x-goog-generation") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()) + .unwrap_or(0); + + let manifest_data = response + .bytes() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))?; + + let manifest: GcsCheckpointManifest = serde_json::from_slice(&manifest_data)?; + + info!( + "Found manifest: step {}, epoch {}, generation {}", + manifest.metadata.step, manifest.metadata.epoch, manifest_generation + ); + + let cache_dir = + get_cache_dir(cache_key, None, manifest.metadata.step, manifest_generation); + + let mut files = if let Some(cached) = collect_cached_files(&cache_dir, &manifest) { + info!("Using cached checkpoint at {:?}", cache_dir); + cached + } else { + info!("Downloading checkpoint via signed URLs to {:?}", cache_dir); + std::fs::create_dir_all(&cache_dir)?; + download_files_from_signed_urls( + &http, + &download_response.urls, + &cache_dir, + &manifest, + ) + .await? + }; + + let config_files = download_non_manifest_files_from_signed_urls( + &http, + &download_response.urls, + &cache_dir, + &[".json", ".py"], + &manifest, + ) + .await?; + files.extend(config_files); + + Ok(files) + } + None => { + info!("No manifest found in signed URLs, downloading all model files"); + let cache_dir = get_cache_dir_no_manifest(cache_key, None); + std::fs::create_dir_all(&cache_dir)?; + download_all_model_files_from_signed_urls( + &http, + &download_response.urls, + &cache_dir, + &MODEL_EXTENSIONS, + ) + .await + } + } +} + +async fn download_files_from_signed_urls( + http: &reqwest::Client, + urls: &[DownloadUrlEntry], + cache_dir: &Path, + manifest: &GcsCheckpointManifest, +) -> Result, DownloadError> { + let mut downloaded_files = Vec::new(); + + for file_entry in &manifest.files { + let local_path = cache_dir.join(&file_entry.filename); + + if local_path.exists() { + info!("Using cached: {}", file_entry.filename); + downloaded_files.push(local_path); + continue; + } + + let url_entry = urls + .iter() + .find(|e| e.path.ends_with(&file_entry.filename)) + .ok_or_else(|| { + DownloadError::RunDown(format!( + "No signed URL found for file: {}", + file_entry.filename + )) + })?; + + info!("Downloading via signed URL: {}", file_entry.filename); + + let response = http + .get(&url_entry.url) + .send() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))?; + + if !response.status().is_success() { + return Err(DownloadError::RunDown(format!( + "Failed to download {}: {}", + file_entry.filename, + response.status() + ))); + } + + let mut stream = response.bytes_stream(); + let mut file = tokio::fs::File::create(&local_path) + .await + .map_err(DownloadError::Io)?; + while let Some(chunk) = stream + .try_next() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))? + { + file.write_all(&chunk).await.map_err(DownloadError::Io)?; + } + info!("Downloaded: {}", file_entry.filename); + downloaded_files.push(local_path); + } + + Ok(downloaded_files) +} + +async fn download_non_manifest_files_from_signed_urls( + http: &reqwest::Client, + urls: &[DownloadUrlEntry], + cache_dir: &Path, + extensions: &[&str], + manifest: &GcsCheckpointManifest, +) -> Result, DownloadError> { + let manifest_filenames: std::collections::HashSet<&str> = + manifest.files.iter().map(|f| f.filename.as_str()).collect(); + + let mut downloaded_files = Vec::new(); + + for url_entry in urls { + let filename = url_entry.path.rsplit('/').next().unwrap_or(&url_entry.path); + + if manifest_filenames.contains(filename) { + continue; + } + + if !extensions.iter().any(|ext| filename.ends_with(ext)) { + continue; + } + + let local_path = cache_dir.join(filename); + if local_path.exists() { + info!("Using cached: {}", filename); + downloaded_files.push(local_path); + continue; + } + + info!("Downloading config via signed URL: {}", filename); + + let response = http + .get(&url_entry.url) + .send() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))?; + + if !response.status().is_success() { + return Err(DownloadError::RunDown(format!( + "Failed to download {}: {}", + filename, + response.status() + ))); + } + + let mut stream = response.bytes_stream(); + let mut file = tokio::fs::File::create(&local_path) + .await + .map_err(DownloadError::Io)?; + while let Some(chunk) = stream + .try_next() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))? + { + file.write_all(&chunk).await.map_err(DownloadError::Io)?; + } + info!("Downloaded: {}", filename); + downloaded_files.push(local_path); + } + + Ok(downloaded_files) +} + +async fn download_all_model_files_from_signed_urls( + http: &reqwest::Client, + urls: &[DownloadUrlEntry], + cache_dir: &Path, + extensions: &[&str], +) -> Result, DownloadError> { + let mut downloaded_files = Vec::new(); + + for url_entry in urls { + let filename = url_entry.path.rsplit('/').next().unwrap_or(&url_entry.path); + + if !extensions.iter().any(|ext| filename.ends_with(ext)) { + continue; + } + + let local_path = cache_dir.join(filename); + if local_path.exists() { + info!("Using cached: {}", filename); + downloaded_files.push(local_path); + continue; + } + + info!("Downloading via signed URL: {}", filename); + + let response = http + .get(&url_entry.url) + .send() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))?; + + if !response.status().is_success() { + return Err(DownloadError::RunDown(format!( + "Failed to download {}: {}", + filename, + response.status() + ))); + } + + let mut stream = response.bytes_stream(); + let mut file = tokio::fs::File::create(&local_path) + .await + .map_err(DownloadError::Io)?; + while let Some(chunk) = stream + .try_next() + .await + .map_err(|e| DownloadError::RunDown(e.to_string()))? + { + file.write_all(&chunk).await.map_err(DownloadError::Io)?; + } + info!("Downloaded: {}", filename); + downloaded_files.push(local_path); + } + + Ok(downloaded_files) +} diff --git a/shared/data-provider/src/http.rs b/shared/data-provider/src/http.rs index 5417f8601..860c26ca6 100644 --- a/shared/data-provider/src/http.rs +++ b/shared/data-provider/src/http.rs @@ -2,15 +2,15 @@ use std::{str::FromStr, time::Duration}; use anyhow::{Context, Result, anyhow, bail}; use futures::future::join_all; -use google_cloud_storage::http::objects::list::ListObjectsRequest; use psyche_coordinator::model::HttpTrainingDataLocation; use psyche_core::{BatchId, Shuffle, TokenSize}; use rand::seq::SliceRandom; use rand_chacha::ChaCha8Rng; use rand_chacha::rand_core::SeedableRng; use reqwest::IntoUrl; +use serde::Deserialize; use tokio::task::JoinHandle; -use tracing::{info, trace}; +use tracing::{debug, info, trace}; use crate::{ TokenizedData, @@ -20,6 +20,22 @@ use crate::{ const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_millis(5000); +/// Response from GCS JSON API list objects endpoint +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct GcsListResponse { + items: Option>, + next_page_token: Option, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct GcsObject { + name: String, + size: Option, + media_link: Option, +} + #[derive(Clone, Copy, Debug)] struct SequencePointer { file_index: usize, @@ -86,49 +102,81 @@ impl FileURLs { } pub async fn from_gcp_bucket(bucket_name: &str, directory: Option) -> Result { - let config = google_cloud_storage::client::ClientConfig::default().anonymous(); - let client = google_cloud_storage::client::Client::new(config); - let mut data_files_matching_directory = { - let mut all_results = vec![]; - // the outer option is if we should continue looping - // the inner option is if we have a "next page token" - let mut next_page_token: Option> = Some(None); - - while let Some(maybe_next_page_token) = next_page_token { - let this_results = client - .list_objects(&ListObjectsRequest { - bucket: bucket_name.to_owned(), - prefix: directory.clone(), - page_token: maybe_next_page_token, - ..Default::default() - }) - .await?; - all_results.extend(this_results.items.iter().flatten().filter_map(|obj| { - let file_ext = obj.name.split('.').next_back()?; + debug!( + "http: from_gcp_bucket: bucket_name={}, directory={:?}", + bucket_name, directory + ); + + // Use the public GCS JSON API directly - no credentials required for public buckets + let client = reqwest::Client::new(); + let mut all_results: Vec<(reqwest::Url, u64)> = vec![]; + let mut page_token: Option = None; + + loop { + let mut url = format!( + "https://storage.googleapis.com/storage/v1/b/{}/o", + urlencoding::encode(bucket_name) + ); + + let mut params = vec![]; + if let Some(ref prefix) = directory { + params.push(format!("prefix={}", urlencoding::encode(prefix))); + } + if let Some(ref token) = page_token { + params.push(format!("pageToken={}", urlencoding::encode(token))); + } + if !params.is_empty() { + url = format!("{}?{}", url, params.join("&")); + } + + let response = client + .get(&url) + .send() + .await + .with_context(|| format!("Failed to list GCS bucket {}", bucket_name))?; + + if !response.status().is_success() { + bail!( + "Failed to list GCS bucket {}: HTTP {}", + bucket_name, + response.status() + ); + } + + let list_response: GcsListResponse = response + .json() + .await + .with_context(|| "Failed to parse GCS list response")?; + + if let Some(items) = list_response.items { + for obj in items { + let file_ext = obj.name.split('.').next_back().unwrap_or(""); if !DATA_FILE_EXTENSIONS.contains(&file_ext) { - return None; + continue; + } + + if let Some(ref media_link) = obj.media_link { + if let Ok(full_url) = media_link.parse::() { + let size = obj + .size + .as_ref() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0); + all_results.push((full_url, size)); + } } + } + } - Some( - obj.media_link - .parse::() - .map(|full_url| (full_url, obj.size as u64)) - .map_err(anyhow::Error::from), - ) - })); - - // if we have a token, Some(Some(String)), - // if not, None - next_page_token = this_results.next_page_token.map(Some) + match list_response.next_page_token { + Some(token) => page_token = Some(token), + None => break, } - all_results } - .into_iter() - .collect::>>()?; - data_files_matching_directory.sort_by(|a, b| a.0.cmp(&b.0)); + all_results.sort_by(|a, b| a.0.cmp(&b.0)); - Ok(Self(data_files_matching_directory)) + Ok(Self(all_results)) } pub async fn from_location(location: &HttpTrainingDataLocation) -> Result { diff --git a/shared/data-provider/src/hub.rs b/shared/data-provider/src/hub.rs index 13a575b84..1cc15f2ab 100644 --- a/shared/data-provider/src/hub.rs +++ b/shared/data-provider/src/hub.rs @@ -1,16 +1,10 @@ use crate::errors::UploadError; -use crate::hub::model::HubRepo; +use futures::future::try_join_all; use hf_hub::{ Cache, Repo, RepoType, - api::{ - Siblings, - tokio::{ApiError, UploadSource}, - }, + api::{Siblings, tokio::ApiError}, }; -use psyche_coordinator::model; -use psyche_core::FixedString; use std::{path::PathBuf, time::Instant}; -use tokio::sync::mpsc; use tracing::{error, info}; const MODEL_EXTENSIONS: [&str; 3] = [".safetensors", ".json", ".py"]; @@ -58,26 +52,22 @@ async fn download_repo_async( .collect::>(); let mut ret: Vec = Vec::new(); for chunk in siblings.chunks(max_concurrent_downloads.unwrap_or(siblings.len())) { - let futures = chunk - .iter() - .map(|x| async { - let start_time = Instant::now(); - tracing::debug!(filename = x.rfilename, "Starting file download from hub"); - let res = api.get(&x.rfilename).await; - if res.is_ok() { - let duration_secs = (Instant::now() - start_time).as_secs_f32(); - tracing::info!( - filename = x.rfilename, - duration_secs = duration_secs, - "Finished downloading file from hub" - ); - } - res - }) - .collect::>(); - for future in futures { - ret.push(future.await?); - } + let futures = chunk.iter().map(|x| async { + let start_time = Instant::now(); + tracing::debug!(filename = x.rfilename, "Starting file download from hub"); + let res = api.get(&x.rfilename).await; + if res.is_ok() { + let duration_secs = (Instant::now() - start_time).as_secs_f32(); + tracing::info!( + filename = x.rfilename, + duration_secs = duration_secs, + "Finished downloading file from hub" + ); + } + res + }); + let chunk_results = try_join_all(futures).await?; + ret.extend(chunk_results); } Ok(ret) } @@ -203,63 +193,72 @@ pub async fn upload_to_hub( hub_info: HubUploadInfo, local: Vec, step: u64, - tx_checkpoint: mpsc::UnboundedSender, + cancellation_token: tokio_util::sync::CancellationToken, ) -> Result<(), UploadError> { let HubUploadInfo { hub_repo, hub_token, } = hub_info; - info!(repo = hub_repo, "Uploading checkpoint to HuggingFace"); + if cancellation_token.is_cancelled() { + return Ok(()); + } + + // Collect all safetensors files to upload in a single commit + let files_to_upload: Vec<_> = local + .iter() + .filter(|p| p.extension() == Some("safetensors".as_ref())) + .map(|path| -> Result<_, UploadError> { + let file_name = path + .file_name() + .ok_or_else(|| UploadError::NotAFile(path.clone()))? + .to_str() + .ok_or_else(|| UploadError::InvalidFilename(path.clone()))? + .to_string(); + Ok((path.clone().into(), file_name)) + }) + .collect::, _>>()?; + + if files_to_upload.is_empty() { + info!(repo = hub_repo, "No safetensors files to upload"); + return Ok(()); + } + + let file_names: Vec<_> = files_to_upload + .iter() + .map(|(_, name)| name.clone()) + .collect(); + info!( + repo = hub_repo, + file_count = files_to_upload.len(), + "Uploading checkpoint to HuggingFace" + ); let api = hf_hub::api::tokio::ApiBuilder::new() - .with_token(Some(hub_token.clone())) + .with_token(Some(hub_token)) .build()?; let repo = Repo::model(hub_repo.clone()); let api_repo = api.repo(repo); - let files: Result, _> = local - .into_iter() - .map(|path| { - path.file_name() - .ok_or(UploadError::NotAFile(path.clone())) - .and_then(|name| { - name.to_str() - .ok_or(UploadError::InvalidFilename(path.clone())) - .map(|s| s.to_string()) - }) - .map(|name| (path.into(), name)) - }) - .collect(); - - let files = files?; - - let commit_info = api_repo - .upload_files(files, Some(format!("step {step}")), None, false) - .await - .map_err(|e| { - error!( - repo = hub_repo, - error = ?e, - "Failed to upload files to HuggingFace" - ); - e - })?; + let upload_future = + api_repo.upload_files(files_to_upload, Some(format!("step {step}")), None, false); - let revision = commit_info.oid; + tokio::select! { + biased; - info!( - repo = hub_repo, - revision = revision, - "Upload to HuggingFace complete" - ); + _ = cancellation_token.cancelled() => { + info!(repo = hub_repo, "Upload to HuggingFace cancelled"); + return Ok(()); + } + result = upload_future => { + result.map_err(|e| { + error!(repo = hub_repo, error = ?e, "Failed to upload files"); + e + })?; + } + } - tx_checkpoint - .send(model::Checkpoint::Hub(HubRepo { - repo_id: FixedString::from_str_truncated(&hub_repo), - revision: Some(FixedString::from_str_truncated(&revision)), - })) - .map_err(|_| UploadError::SendCheckpoint)?; + info!(repo = hub_repo, files = ?file_names, "Upload to HuggingFace complete"); Ok(()) } diff --git a/shared/data-provider/src/lib.rs b/shared/data-provider/src/lib.rs index 0044d77d2..d6d46850f 100644 --- a/shared/data-provider/src/lib.rs +++ b/shared/data-provider/src/lib.rs @@ -4,11 +4,13 @@ mod dummy; mod errors; mod file_extensions; mod gcs; +mod gcs_signed; pub mod http; mod hub; mod local; mod preprocessed; mod remote; +pub mod run_down; mod traits; mod weighted; @@ -21,6 +23,7 @@ pub use gcs::{ GcsCheckpointManifest, GcsManifestMetadata, GcsUploadInfo, ManifestFileEntry, ManifestMetadata, download_model_from_gcs_async, download_model_from_gcs_sync, upload_to_gcs, }; +pub use gcs_signed::{download_model_from_gcs_signed_async, upload_to_gcs_signed}; pub use hub::{ HubUploadInfo, download_dataset_repo_async, download_dataset_repo_sync, download_model_repo_async, download_model_repo_sync, upload_to_hub, @@ -29,5 +32,6 @@ pub use local::LocalDataProvider; pub use parquet::record::{ListAccessor, MapAccessor, RowAccessor}; pub use preprocessed::PreprocessedDataProvider; pub use remote::{DataProviderTcpClient, DataProviderTcpServer, DataServerTui}; +pub use run_down::RunDownClient; pub use traits::{LengthKnownDataProvider, TokenizedData, TokenizedDataProvider}; pub use weighted::{WeightedDataProvider, http::WeightedHttpProvidersConfig}; diff --git a/shared/data-provider/src/run_down.rs b/shared/data-provider/src/run_down.rs new file mode 100644 index 000000000..adc5c3018 --- /dev/null +++ b/shared/data-provider/src/run_down.rs @@ -0,0 +1,177 @@ +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; +use tracing::info; + +const DEFAULT_RUN_DOWN_BASE_URL: &str = "https://run-down.nousresearch.com/v1"; +const SIGNED_URL_EXPIRY_SECONDS: u64 = 3600; + +fn base_url() -> String { + std::env::var("RUN_DOWN_URL").unwrap_or_else(|_| DEFAULT_RUN_DOWN_BASE_URL.to_string()) +} + +/// Client for the Nous run-down service that provides signed URLs for GCS checkpoint +/// upload/download. Uses a generic signing function to decouple from specific wallet +/// implementations. +type SignFn = dyn Fn(&[u8]) -> Vec + Send + Sync; + +pub struct RunDownClient { + http: reqwest::Client, + run_id: String, + wallet_address: String, + sign_fn: Arc, +} + +impl std::fmt::Debug for RunDownClient { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RunDownClient") + .field("run_id", &self.run_id) + .finish() + } +} + +impl RunDownClient { + pub fn new( + run_id: String, + wallet_address: String, + sign_fn: impl Fn(&[u8]) -> Vec + Send + Sync + 'static, + ) -> Self { + Self { + http: reqwest::Client::new(), + run_id, + wallet_address, + sign_fn: Arc::new(sign_fn), + } + } + + pub fn run_id(&self) -> &str { + &self.run_id + } + + fn generate_signature(&self, expires_in_seconds: u64, nonce: u64) -> String { + let message = format!( + "nous-run-down-service:{}:{}:{}:{}", + self.run_id, self.wallet_address, expires_in_seconds, nonce + ); + let signature_bytes = (self.sign_fn)(message.as_bytes()); + bs58::encode(&signature_bytes).into_string() + } + + fn nonce() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64 + } + + /// Get a signed upload URL for a single file. + pub async fn get_upload_url(&self, filename: &str) -> Result { + let expires_in_seconds = SIGNED_URL_EXPIRY_SECONDS; + let nonce = Self::nonce(); + let signature = self.generate_signature(expires_in_seconds, nonce); + + let url = format!("{}/upload/{}", base_url(), self.run_id); + let body = serde_json::json!({ + "walletAddress": self.wallet_address, + "filename": filename, + "expiresInSeconds": expires_in_seconds, + "nonce": nonce.to_string(), + }); + + info!(filename, url, "Requesting signed upload URL from run-down"); + + let response = self + .http + .post(&url) + .header("X-Solana-Signature", &signature) + .json(&body) + .send() + .await + .map_err(|e| RunDownError::Request(e.to_string()))?; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response.text().await.unwrap_or_default(); + return Err(RunDownError::Api(format!( + "Upload URL request failed with status {}: {}", + status, error_text + ))); + } + + response + .json() + .await + .map_err(|e| RunDownError::Parse(e.to_string())) + } + + /// Get signed download URLs for all files in the run. + pub async fn get_download_urls(&self) -> Result { + let expires_in_seconds = SIGNED_URL_EXPIRY_SECONDS; + let nonce = Self::nonce(); + let signature = self.generate_signature(expires_in_seconds, nonce); + + let url = format!("{}/download/{}", base_url(), self.run_id); + let body = serde_json::json!({ + "walletAddress": self.wallet_address, + "expiresInSeconds": expires_in_seconds, + "nonce": nonce.to_string(), + }); + + info!(url, "Requesting signed download URLs from run-down"); + + let response = self + .http + .post(&url) + .header("X-Solana-Signature", &signature) + .json(&body) + .send() + .await + .map_err(|e| RunDownError::Request(e.to_string()))?; + + if !response.status().is_success() { + let status = response.status(); + let error_text = response.text().await.unwrap_or_default(); + return Err(RunDownError::Api(format!( + "Download URLs request failed with status {}: {}", + status, error_text + ))); + } + + response + .json() + .await + .map_err(|e| RunDownError::Parse(e.to_string())) + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct UploadUrlResponse { + pub url: String, + pub expires_at: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DownloadUrlEntry { + pub path: String, + pub url: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DownloadUrlsResponse { + pub urls: Vec, + pub expires_at: String, +} + +#[derive(Debug, thiserror::Error)] +pub enum RunDownError { + #[error("run-down request failed: {0}")] + Request(String), + + #[error("run-down API error: {0}")] + Api(String), + + #[error("failed to parse run-down response: {0}")] + Parse(String), +} diff --git a/shared/watcher/src/traits.rs b/shared/watcher/src/traits.rs index 1e13e0ef8..96dacdbed 100644 --- a/shared/watcher/src/traits.rs +++ b/shared/watcher/src/traits.rs @@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize}; pub enum OpportunisticData { WitnessStep(Witness, WitnessMetadata), WarmupStep(Witness), + CooldownStep(Witness), } impl OpportunisticData { @@ -14,6 +15,7 @@ impl OpportunisticData { match self { OpportunisticData::WitnessStep(..) => "witness", OpportunisticData::WarmupStep(..) => "warmup", + OpportunisticData::CooldownStep(..) => "cooldown", } } } diff --git a/tools/rust-tools/run-manager/src/commands/run/update_config.rs b/tools/rust-tools/run-manager/src/commands/run/update_config.rs index 641577307..4a22ef02f 100644 --- a/tools/rust-tools/run-manager/src/commands/run/update_config.rs +++ b/tools/rust-tools/run-manager/src/commands/run/update_config.rs @@ -92,9 +92,9 @@ impl Command for CommandUpdateConfig { let Model::LLM(mut llm) = model.unwrap_or(coordinator_account_state.state.coordinator.model); match llm.checkpoint { - Checkpoint::P2P(hub_repo) | Checkpoint::Dummy(hub_repo) => { - llm.checkpoint = Checkpoint::Hub(hub_repo) - } + Checkpoint::P2P(hub_repo) => llm.checkpoint = Checkpoint::Hub(hub_repo), + Checkpoint::P2PGcs(gcs_repo) => llm.checkpoint = Checkpoint::Gcs(gcs_repo), + Checkpoint::Dummy(hub_repo) => llm.checkpoint = Checkpoint::Hub(hub_repo), _ => {} } Some(Model::LLM(llm)) diff --git a/tools/rust-tools/run-manager/src/docker/coordinator_client.rs b/tools/rust-tools/run-manager/src/docker/coordinator_client.rs index a3d2cacef..3c794f2e0 100644 --- a/tools/rust-tools/run-manager/src/docker/coordinator_client.rs +++ b/tools/rust-tools/run-manager/src/docker/coordinator_client.rs @@ -4,6 +4,7 @@ use anchor_client::solana_sdk::{ use anchor_lang::AccountDeserialize; use anyhow::{Context, Result}; use psyche_coordinator::RunState; +use psyche_coordinator::model::{Checkpoint, LLM, Model}; use psyche_solana_authorizer::state::Authorization; use psyche_solana_coordinator::{ CoordinatorInstance, coordinator_account_from_bytes, find_coordinator_instance, @@ -213,6 +214,22 @@ impl CoordinatorClient { Ok(image_name) } + /// Get the checkpoint type configured for this run + pub fn get_checkpoint_type(&self, run_id: &str) -> Result { + let instance = self.fetch_coordinator_data(run_id)?; + + let coordinator_account_data = self + .rpc_client + .get_account(&instance.coordinator_account) + .context("RPC error: failed to get coordinator account")?; + + let coordinator_account = coordinator_account_from_bytes(&coordinator_account_data.data) + .context("Failed to deserialize CoordinatorAccount")?; + + let Model::LLM(LLM { checkpoint, .. }) = coordinator_account.state.coordinator.model; + Ok(checkpoint) + } + pub fn get_all_runs(&self) -> Result> { // Fetch all CoordinatorInstance accounts that are owned by the program let accounts = self diff --git a/tools/rust-tools/run-manager/src/docker/manager.rs b/tools/rust-tools/run-manager/src/docker/manager.rs index fe09a07e2..57e7cd50d 100644 --- a/tools/rust-tools/run-manager/src/docker/manager.rs +++ b/tools/rust-tools/run-manager/src/docker/manager.rs @@ -2,6 +2,7 @@ use anchor_client::solana_sdk::bs58; use anchor_client::solana_sdk::pubkey::Pubkey; use anchor_client::solana_sdk::signature::{EncodableKey, Keypair, Signer}; use anyhow::{Context, Result, anyhow, bail}; +use psyche_coordinator::model::Checkpoint; use std::io::{BufRead, BufReader, Cursor}; use std::path::PathBuf; use std::process::{Command, Stdio}; @@ -60,7 +61,7 @@ impl RunManager { info!("Using coordinator program ID: {}", coordinator_program_id); let rpc = get_env_var("RPC")?; - let scratch_dir = std::env::var("SCRATCH_DIR").ok(); + let scratch_dir = get_env_var("SCRATCH_DIR").ok(); let coordinator_client = CoordinatorClient::new(rpc, coordinator_program_id); @@ -314,7 +315,36 @@ impl RunManager { Ok(()) } + /// Validate that required credentials are available based on checkpoint type + fn validate_credentials(&self) -> Result<()> { + let checkpoint = self.coordinator_client.get_checkpoint_type(&self.run_id)?; + + match checkpoint { + Checkpoint::Gcs(_) | Checkpoint::P2PGcs(_) => { + info!("GCS checkpointing uses run-down signed URLs"); + } + Checkpoint::Hub(_) | Checkpoint::P2P(_) => { + // HF_TOKEN should be in the env file + if std::env::var("HF_TOKEN").is_err() { + bail!( + "This run uses HuggingFace checkpointing but HF_TOKEN is not set. \ + Please set HF_TOKEN in your environment or env file." + ); + } + info!("HuggingFace token validated for checkpoint upload"); + } + Checkpoint::Ephemeral | Checkpoint::Dummy(_) => { + // No credentials needed for ephemeral or dummy checkpoints + info!("No checkpoint credentials required for this run"); + } + } + + Ok(()) + } + pub async fn run(&self, entrypoint: Option) -> Result<()> { + self.validate_credentials()?; + loop { let docker_tag = self.prepare_image().await?; info!("Starting container..."); diff --git a/website/backend/src/coordinatorChainLoop.ts b/website/backend/src/coordinatorChainLoop.ts index 21aa04dc2..687e407da 100644 --- a/website/backend/src/coordinatorChainLoop.ts +++ b/website/backend/src/coordinatorChainLoop.ts @@ -342,6 +342,17 @@ export async function startWatchCoordinatorChainLoop( }) break } + case 'cooldown_witness': { + const runPdaAddr = i.accounts[1].toString() + const coordinatorAddr = i.accounts[2].toString() + runUpdates.getAndTouchCurrentRun({ + runPdaAddr, + coordinatorAddr, + decoded, + tx, + }) + break + } case 'update_client_version': { const runPdaAddr = i.accounts[1].toString() const coordinatorAddr = i.accounts[2].toString()