diff --git a/CHANGELOG.md b/CHANGELOG.md index f6779face..b25257e69 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,19 @@ # Dev +## Breaking change + +### Memory resource in megabytes + +* Automatically detected resource "mem" that is the size of RAM of a worker is now using megabytes as a unit. + i.e. `--resource mem=100` asks now for 100 MiB (previously 100 bytes). + ## New features +### Non-integer resource requests +* You may now ask of non-integer amount of a resource. e.g. for 0.5 of GPU. + This enables resource sharing on the logical level of HyperQueue scheduler and allows to utilize remaining part the resource + by another tasks. + ### Job submission * You can now specify `cleanup modes` when passing `stdout`/`stderr` paths to tasks. Cleanup mode decides what should happen with the file once the task has finished executing. Currently, a single cleanup mode is implemented, which removes @@ -10,6 +22,7 @@ the file if the task has finished successfully: $ hq submit --stdout=out.txt:rm-if-finished /my-program ``` + # v0.16.0 ## New features diff --git a/Cargo.lock b/Cargo.lock index 1b9f01ca8..4773a7e38 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4,9 +4,9 @@ version = 3 [[package]] name = "addr2line" -version = "0.20.0" +version = "0.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4fa78e18c64fce05e902adecd7a5eed15a5e0a3439f7b0e169f0252214865e3" +checksum = "8a30b2e23b9e17a9f90641c7ab1549cd9b44f296d3ccbf309d2863cfe398a0cb" dependencies = [ "gimli", ] @@ -30,9 +30,9 @@ dependencies = [ [[package]] name = "aho-corasick" -version = "1.0.2" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43f6cb1bf222025340178f382c426f13757b2960e89779dfcb319c32542a5a41" +checksum = "ea5d730647d4fadd988536d06fecce94b7b4f2a7efdae548f1cf4b63205518ab" dependencies = [ "memchr", ] @@ -54,24 +54,23 @@ dependencies = [ [[package]] name = "anstream" -version = "0.3.2" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ca84f3628370c59db74ee214b3263d58f9aadd9b4fe7e711fd87dc452b7f163" +checksum = "b1f58811cfac344940f1a400b6e6231ce35171f614f26439e80f8c1465c5cc0c" dependencies = [ "anstyle", "anstyle-parse", "anstyle-query", "anstyle-wincon", "colorchoice", - "is-terminal", "utf8parse", ] [[package]] name = "anstyle" -version = "1.0.1" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3a30da5c5f2d5e72842e00bcb57657162cdabef0931f40e2deb9b4140440cecd" +checksum = "b84bf0a05bbb2a83e5eb6fa36bb6e87baa08193c35ff52bbf6b38d8af2890e46" [[package]] name = "anstyle-parse" @@ -93,9 +92,9 @@ dependencies = [ [[package]] name = "anstyle-wincon" -version = "1.0.1" +version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "180abfa45703aebe0093f79badacc01b8fd4ea2e35118747e5811127f926e188" +checksum = "58f54d10c6dfa51283a066ceab3ec1ab78d13fae00aa49243a45e4571fb79dfd" dependencies = [ "anstyle", "windows-sys 0.48.0", @@ -103,9 +102,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.72" +version = "1.0.75" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b13c32d80ecc7ab747b80c3784bce54ee8a7a0cc4fbda9bf4cda2cf6fe90854" +checksum = "a4668cab20f66d8d020e1fbc0ebe47217433c1b6c8f2040faf858554e394ace6" [[package]] name = "arrayvec" @@ -115,9 +114,9 @@ checksum = "96d30a06541fbafbc7f82ed10c06164cfbd2c401138f6addd8404629c4b16711" [[package]] name = "async-compression" -version = "0.4.1" +version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62b74f44609f0f91493e3082d3734d98497e094777144380ea4db9f9905dd5b6" +checksum = "bb42b2197bf15ccb092b62c74515dbd8b86d0effd934795f6687c93b6e679a2c" dependencies = [ "flate2", "futures-core", @@ -145,9 +144,9 @@ checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" [[package]] name = "backtrace" -version = "0.3.68" +version = "0.3.69" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4319208da049c43661739c5fade2ba182f09d1dc2299b32298d3a31692b17e12" +checksum = "2089b7e3f35b9dd2d0ed921ead4f6d318c27680d4a5bd167b3ee120edb105837" dependencies = [ "addr2line", "cc", @@ -175,9 +174,9 @@ checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" [[package]] name = "bitflags" -version = "2.3.3" +version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "630be753d4e58660abd17930c71b647fe46c27ea6b63cc59e1e3851406972e42" +checksum = "b4682ae6287fcf752ecaabbfcc7b6f9b72aa33933dc23a554d853aea8eea8635" [[package]] name = "brownstone" @@ -190,9 +189,9 @@ dependencies = [ [[package]] name = "bstr" -version = "1.6.0" +version = "1.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6798148dccfbff0fae41c7574d2fa8f1ef3492fba0face179de5d8d447d67b05" +checksum = "4c2f7349907b712260e64b0afe2f84692af14a454be26187d9df565c7f69266a" dependencies = [ "memchr", "regex-automata", @@ -201,9 +200,9 @@ dependencies = [ [[package]] name = "bumpalo" -version = "3.13.0" +version = "3.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a3e2c3daef883ecc1b5d58c15adae93470a91d425f3532ba1695849656af3fc1" +checksum = "7f30e7476521f6f8af1a1c4c0b8cc94f0bee37d91763d0ca2665f299b6cd8aec" [[package]] name = "byteorder" @@ -213,9 +212,9 @@ checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" [[package]] name = "bytes" -version = "1.4.0" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "89b2fd2a0dcf38d7971e2194b6b6eebab45ae01067456a7fd93d5547a61b70be" +checksum = "a2bd12c1caf447e69cd4528f47f94d203fd2582878ecb9e9465484c4148a8223" [[package]] name = "cassowary" @@ -231,9 +230,12 @@ checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" [[package]] name = "cc" -version = "1.0.79" +version = "1.0.83" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50d30906286121d95be3d479533b458f87493b30a4b5f79a607db8f5d11aa91f" +checksum = "f1174fb0b6ec23863f8b971027804a42614e347eafb0a95bf0b12cdae21fc4d0" +dependencies = [ + "libc", +] [[package]] name = "cfg-if" @@ -243,18 +245,17 @@ checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" [[package]] name = "chrono" -version = "0.4.26" +version = "0.4.31" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec837a71355b28f6556dbd569b37b3f363091c0bd4b2e735674521b4c5fd9bc5" +checksum = "7f2c685bad3eb3d45a01354cedb7d5faa66194d1d58ba6e267a8de788f79db38" dependencies = [ "android-tzdata", "iana-time-zone", "js-sys", "num-traits", "serde", - "time", "wasm-bindgen", - "winapi", + "windows-targets 0.48.5", ] [[package]] @@ -280,20 +281,19 @@ dependencies = [ [[package]] name = "clap" -version = "4.3.19" +version = "4.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fd304a20bff958a57f04c4e96a2e7594cc4490a0e809cbd48bb6437edaa452d" +checksum = "824956d0dca8334758a5b7f7e50518d66ea319330cbceedcf76905c2f6ab30e3" dependencies = [ "clap_builder", "clap_derive", - "once_cell", ] [[package]] name = "clap_builder" -version = "4.3.19" +version = "4.4.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "01c6a3f08f1fe5662a35cfe393aec09c4df95f60ee93b7556505260f75eee9e1" +checksum = "122ec64120a49b4563ccaedcbea7818d069ed8e9aa6d829b82d8a4128936b2ab" dependencies = [ "anstream", "anstyle", @@ -303,30 +303,30 @@ dependencies = [ [[package]] name = "clap_complete" -version = "4.3.2" +version = "4.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fc443334c81a804575546c5a8a79b4913b50e28d69232903604cada1de817ce" +checksum = "8baeccdb91cd69189985f87f3c7e453a3a451ab5746cf3be6acc92120bd16d24" dependencies = [ - "clap 4.3.19", + "clap 4.4.5", ] [[package]] name = "clap_derive" -version = "4.3.12" +version = "4.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "54a9bb5758fc5dfe728d1019941681eccaf0cf8a4189b692a0ee2f2ecf90a050" +checksum = "0862016ff20d69b84ef8247369fabf5c008a7417002411897d40ee1f4532b873" dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", ] [[package]] name = "clap_lex" -version = "0.5.0" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2da6da31387c7e4ef160ffab6d5e7f00c42626fe39aea70a7b0f1773f7dd6c1b" +checksum = "cd7cc57abe963c6d3b9d8be5b06ba7c8957a930305ca90304f24ef040aa6f961" [[package]] name = "cli-table" @@ -387,6 +387,12 @@ dependencies = [ "unicode-xid", ] +[[package]] +name = "convert_case" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6245d59a3e82a7fc217c5828a6692dbc6dfb63a0c8c90495621f7b9d79704a0e" + [[package]] name = "core-foundation-sys" version = "0.8.4" @@ -395,9 +401,9 @@ checksum = "e496a50fda8aacccc86d7529e2c1e0892dbd0f898a6b5645b5561b89c3210efa" [[package]] name = "core_affinity" -version = "0.8.0" +version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4436406e93f52cce33bfba4be067a9f7229da44a634c385e4b22cdfaca5f84cc" +checksum = "622892f5635ce1fc38c8f16dfc938553ed64af482edb5e150bf4caedbfcb2304" dependencies = [ "libc", "num_cpus", @@ -449,16 +455,6 @@ dependencies = [ "itertools", ] -[[package]] -name = "crossbeam-channel" -version = "0.5.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a33c2bf77f2df06183c3aa30d1e96c0695a313d4f9c453cc3762a6db39f99200" -dependencies = [ - "cfg-if", - "crossbeam-utils", -] - [[package]] name = "crossbeam-deque" version = "0.8.3" @@ -621,8 +617,10 @@ version = "0.99.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4fb810d30a7c1953f91334de7244731fc3f3c10d7fe163338a35b9f640960321" dependencies = [ + "convert_case", "proc-macro2", "quote", + "rustc_version", "syn 1.0.109", ] @@ -691,9 +689,9 @@ checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" [[package]] name = "errno" -version = "0.3.1" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4bcfec3a70f97c962c307b2d2c56e358cf1d00b558d74262b5f929ee8cc7e73a" +checksum = "136526188508e25c6fef639d7927dfb3e0e3084488bf202267829cf7fc23dbdd" dependencies = [ "errno-dragonfly", "libc", @@ -712,15 +710,15 @@ dependencies = [ [[package]] name = "fiat-crypto" -version = "0.1.20" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e825f6987101665dea6ec934c09ec6d721de7bc1bf92248e1d5810c8cd636b77" +checksum = "d0870c84016d4b481be5c9f323c24f65e31e901ae618f0e80f4308fb00de1d2d" [[package]] name = "flate2" -version = "1.0.26" +version = "1.0.27" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3b9429470923de8e8cbd4d2dc513535400b4b3fef0319fb5c4e1f520a7bef743" +checksum = "c6c98ee8095e9d1dcbf2fcc6d95acccb90d1c81db1e44725c6a984b1dbdfb010" dependencies = [ "crc32fast", "miniz_oxide", @@ -794,7 +792,7 @@ checksum = "89ca545a94061b6365f2c7355b4b32bd20df3ff95f02da9329b34ccc3bd6ee72" dependencies = [ "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", ] [[package]] @@ -843,7 +841,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0176e0459c2e4a1fe232f984bca6890e681076abb9934f6cea7c326f3fc47818" dependencies = [ "libc", - "windows-targets 0.48.1", + "windows-targets 0.48.5", ] [[package]] @@ -854,7 +852,7 @@ checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427" dependencies = [ "cfg-if", "libc", - "wasi 0.11.0+wasi-snapshot-preview1", + "wasi", ] [[package]] @@ -865,14 +863,14 @@ checksum = "ba330b70a5341d3bc730b8e205aaee97ddab5d9c448c4f51a7c2d924266fa8f9" dependencies = [ "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", ] [[package]] name = "gimli" -version = "0.27.3" +version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6c80984affa11d98d1b88b66ac8853f143217b399d3c74116778ff8fdb4ed2e" +checksum = "6fb8d784f27acf97159b40fc4db5ecd8aa23b9ad5ef69cdd136d3bc80665f0c0" [[package]] name = "glob" @@ -921,9 +919,9 @@ dependencies = [ [[package]] name = "hermit-abi" -version = "0.3.2" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "443144c8cdadd93ebf52ddb4056d257f5b52c04d3c804e657d19eb73fc33668b" +checksum = "d77f7ec81a6d05a3abb01ab6eb7590f6083d08449fe5a1c8b1e620283546ccb7" [[package]] name = "hex" @@ -950,7 +948,7 @@ dependencies = [ "bytes", "chrono", "chumsky", - "clap 4.3.19", + "clap 4.4.5", "clap_complete", "cli-table", "colored", @@ -968,7 +966,7 @@ dependencies = [ "insta", "jemallocator", "log", - "nix 0.26.2", + "nix 0.26.4", "nom", "nom-supreme", "num_cpus", @@ -1075,9 +1073,9 @@ dependencies = [ [[package]] name = "insta" -version = "1.31.0" +version = "1.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0770b0a3d4c70567f0d58331f3088b0e4c4f56c9b8d764efe654b4a5d46de3a" +checksum = "a3e02c584f4595792d09509a94cdb92a3cef7592b1eb2d9877ee6f527062d0ea" dependencies = [ "console", "lazy_static", @@ -1123,7 +1121,7 @@ version = "0.4.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cb0889898416213fab133e1d33a0e5858a48177452750691bde3666d0fdbaf8b" dependencies = [ - "hermit-abi 0.3.2", + "hermit-abi 0.3.3", "rustix", "windows-sys 0.48.0", ] @@ -1186,9 +1184,9 @@ checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" [[package]] name = "libc" -version = "0.2.147" +version = "0.2.148" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4668fb0ea861c1df094127ac5f1da3409a82116a4ba74fca2e58ef927159bb3" +checksum = "9cdc71e17332e86d2e1d38c1f99edcb6288ee11b815fb1a4b049eaa2114d369b" [[package]] name = "linked-hash-map" @@ -1198,9 +1196,9 @@ checksum = "0717cef1bc8b636c6e1c1bbdefc09e6322da8a9321966e8928ef80d20f7f770f" [[package]] name = "linux-raw-sys" -version = "0.4.3" +version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09fc20d2ca12cb9f044c93e3bd6d32d523e6e2ec3db4f7b2939cd99026ecd3f0" +checksum = "1a9bad9f94746442c783ca431b22403b519cd7fbeed0533fdd6328b2f2212128" [[package]] name = "lock_api" @@ -1214,9 +1212,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.19" +version = "0.4.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b06a4cde4c0f271a446782e3eff8de789548ce57dbc8eca9292c27f4a42004b4" +checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" [[package]] name = "mach" @@ -1229,9 +1227,9 @@ dependencies = [ [[package]] name = "memchr" -version = "2.5.0" +version = "2.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2dffe52ecf27772e601905b7522cb4ef790d2cc203488bbd0e2fe85fcb74566d" +checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c" [[package]] name = "memoffset" @@ -1282,7 +1280,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "927a765cd3fc26206e66b296465fa9d3e5ab003e651c1b3c060e7956d96b19d2" dependencies = [ "libc", - "wasi 0.11.0+wasi-snapshot-preview1", + "wasi", "windows-sys 0.48.0", ] @@ -1301,16 +1299,15 @@ dependencies = [ [[package]] name = "nix" -version = "0.26.2" +version = "0.26.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bfdda3d196821d6af13126e40375cdf7da646a96114af134d5f417a9a1dc8e1a" +checksum = "598beaf3cc6fdd9a5dfb1630c2800c7acd31df7aaf0f565796fba2b53ca1af1b" dependencies = [ "bitflags 1.3.2", "cfg-if", "libc", "memoffset 0.7.1", "pin-utils", - "static_assertions", ] [[package]] @@ -1361,7 +1358,7 @@ version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4161fcb6d602d4d2081af7c3a45852d875a03dd337a6bfdd6e06407b61342a43" dependencies = [ - "hermit-abi 0.3.2", + "hermit-abi 0.3.3", "libc", ] @@ -1373,9 +1370,9 @@ checksum = "b8f8bdf33df195859076e54ab11ee78a1b208382d3a26ec40d142ffc1ecc49ef" [[package]] name = "object" -version = "0.31.1" +version = "0.32.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bda667d9f2b5051b8833f59f3bf748b28ef54f850f4fcb389a252aa383866d1" +checksum = "9cf5f9dd3933bd50a9e1f149ec995f39ae2c496d31fd772c1fd45ebc27e902b0" dependencies = [ "memchr", ] @@ -1400,9 +1397,9 @@ checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" [[package]] name = "orion" -version = "0.17.5" +version = "0.17.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b11468cc6afd61a126fe3f91cc4cc8a0dbe7917d0a4b5e8357ba91cc47444462" +checksum = "7abdb10181903c8c4b016ba45d6d6d5af1a1e2a461aa4763a83b87f5df4695e5" dependencies = [ "ct-codecs", "fiat-crypto", @@ -1462,7 +1459,7 @@ dependencies = [ "libc", "redox_syscall 0.3.5", "smallvec", - "windows-targets 0.48.1", + "windows-targets 0.48.5", ] [[package]] @@ -1492,9 +1489,9 @@ dependencies = [ [[package]] name = "pin-project-lite" -version = "0.2.10" +version = "0.2.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c40d25201921e5ff0c862a505c6557ea88568a4e3ace775ab55e93f2f4f9d57" +checksum = "8afb450f006bf6385ca15ef45d71d2288452bc3683ce2e2cacc0d18e4be60b58" [[package]] name = "pin-utils" @@ -1560,9 +1557,9 @@ checksum = "dc375e1527247fe1a97d8b7156678dfe7c1af2fc075c9a4db3690ecd2a148068" [[package]] name = "proc-macro2" -version = "1.0.66" +version = "1.0.67" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "18fb31db3f9bddb2ea821cde30a9f70117e3f119938b5ee630b7403aa6e2ead9" +checksum = "3d433d9f1a3e8c1263d9456598b16fec66f4acc9a74dacffd35c7bb09b3a1328" dependencies = [ "unicode-ident", ] @@ -1701,9 +1698,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.32" +version = "1.0.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50f3b39ccfb720540debaa0164757101c08ecb8d326b15358ce76a62c7e85965" +checksum = "5267fca4496028628a95160fc423a33e8b2e6af8a5302579e322e4b520293cae" dependencies = [ "proc-macro2", ] @@ -1768,9 +1765,9 @@ dependencies = [ [[package]] name = "rayon" -version = "1.7.0" +version = "1.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d2df5196e37bcc87abebc0053e20787d73847bb33134a69841207dd0a47f03b" +checksum = "9c27db03db7734835b3f53954b534c91069375ce6ccaa2e065441e07d9b6cdb1" dependencies = [ "either", "rayon-core", @@ -1778,14 +1775,12 @@ dependencies = [ [[package]] name = "rayon-core" -version = "1.11.0" +version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4b8f95bd6966f5c87776639160a66bd8ab9895d9d4ab01ddba9fc60661aebe8d" +checksum = "5ce3fb6ad83f861aac485e76e1985cd109d9a3713802152be56c3b1f0e0658ed" dependencies = [ - "crossbeam-channel", "crossbeam-deque", "crossbeam-utils", - "num_cpus", ] [[package]] @@ -1837,9 +1832,9 @@ dependencies = [ [[package]] name = "regex" -version = "1.9.1" +version = "1.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2eae68fc220f7cf2532e4494aded17545fce192d59cd996e0fe7887f4ceb575" +checksum = "697061221ea1b4a94a624f67d0ae2bfe4e22b8a17b6a192afb11046542cc8c47" dependencies = [ "aho-corasick", "memchr", @@ -1849,9 +1844,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.3.4" +version = "0.3.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7b6d6190b7594385f61bd3911cd1be99dfddcfc365a4160cc2ab5bff4aed294" +checksum = "c2f401f4955220693b56f8ec66ee9c78abffd8d1c4f23dc41a23839eb88f0795" dependencies = [ "aho-corasick", "memchr", @@ -1860,9 +1855,9 @@ dependencies = [ [[package]] name = "regex-syntax" -version = "0.7.4" +version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e5ea92a5b6195c6ef2a0295ea818b312502c6fc94dde986c5553242e18fd4ce2" +checksum = "dbb5fb1acd8a1a18b3dd5be62d25485eb770e05afb408a9627d14d451bae12da" [[package]] name = "remove_dir_all" @@ -1913,13 +1908,22 @@ version = "0.1.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" +[[package]] +name = "rustc_version" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa0f585226d2e68097d4f95d113b15b83a82e819ab25717ec0590d9584ef366" +dependencies = [ + "semver", +] + [[package]] name = "rustix" -version = "0.38.4" +version = "0.38.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0a962918ea88d644592894bc6dc55acc6c0956488adcebbfb6e273506b7fd6e5" +checksum = "747c788e9ce8e92b12cd485c49ddf90723550b654b32508f979b71a7b1ecda4f" dependencies = [ - "bitflags 2.3.3", + "bitflags 2.4.0", "errno", "libc", "linux-raw-sys", @@ -1947,11 +1951,17 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "semver" +version = "1.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad977052201c6de01a8ef2aa3378c4bd23217a056337d1d6da40468d267a4fb0" + [[package]] name = "serde" -version = "1.0.178" +version = "1.0.188" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "60363bdd39a7be0266a520dab25fdc9241d2f987b08a01e01f0ec6d06a981348" +checksum = "cf9e0fcba69a370eed61bcf2b728575f726b50b55cba78064753d708ddc7549e" dependencies = [ "serde_derive", ] @@ -1986,20 +1996,20 @@ dependencies = [ [[package]] name = "serde_derive" -version = "1.0.178" +version = "1.0.188" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f28482318d6641454cb273da158647922d1be6b5a2fcc6165cd89ebdd7ed576b" +checksum = "4eca7ac642d82aa35b60049a6eccb4be6be75e599bd2e9adb5f875a737654af2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", ] [[package]] name = "serde_json" -version = "1.0.104" +version = "1.0.107" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "076066c5f1078eac5b722a31827a8832fe108bed65dfa75e233c89f8206e976c" +checksum = "6b420ce6e3d8bd882e9b243c6eed35dbc9a6110c9769e74b584e0d68d1f20c65" dependencies = [ "itoa", "ryu", @@ -2051,36 +2061,36 @@ checksum = "420acb44afdae038210c99e69aae24109f32f15500aa708e81d46c9f29d55fcf" [[package]] name = "slab" -version = "0.4.8" +version = "0.4.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6528351c9bc8ab22353f9d776db39a20288e8d6c37ef8cfe3317cf875eecfc2d" +checksum = "8f92a496fb766b417c996b9c5e57daf2f7ad3b0bebe1ccfca4856390e3d3bb67" dependencies = [ "autocfg", ] [[package]] name = "smallvec" -version = "1.11.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62bb4feee49fdd9f707ef802e22365a35de4b7b299de4763d44bfea899442ff9" +checksum = "942b4a808e05215192e39f4ab80813e599068285906cc91aa64f923db842bd5a" dependencies = [ "serde", ] [[package]] name = "smawk" -version = "0.3.1" +version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f67ad224767faa3c7d8b6d91985b78e70a1324408abcb1cfcc2be4c06bc06043" +checksum = "b7c388c1b5e93756d0c740965c41e8822f866621d41acbdf6336a6a168f8840c" [[package]] name = "socket2" -version = "0.4.9" +version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "64a4a911eed85daf18834cfaa86a79b7d266ff93ff5ba14005426219480ed662" +checksum = "4031e820eb552adee9295814c0ced9e5cf38ddf1e8b7d566d6de8e2538ea989e" dependencies = [ "libc", - "winapi", + "windows-sys 0.48.0", ] [[package]] @@ -2096,12 +2106,6 @@ dependencies = [ "winapi", ] -[[package]] -name = "static_assertions" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" - [[package]] name = "strsim" version = "0.10.0" @@ -2127,9 +2131,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.27" +version = "2.0.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b60f673f44a8255b9c8c657daf66a596d435f2da81a555b06dc644d080ba45e0" +checksum = "7303ef2c05cd654186cb250d29049a24840ca25d2747c25c0381c8d9e2f582e8" dependencies = [ "proc-macro2", "quote", @@ -2142,12 +2146,13 @@ version = "0.1.0" dependencies = [ "anyhow", "bincode", - "bitflags 2.3.3", + "bitflags 2.4.0", "bstr", "byteorder", "bytes", "criterion", "derive_builder", + "derive_more", "env_logger", "futures", "fxhash", @@ -2155,7 +2160,7 @@ dependencies = [ "hashbrown 0.14.0", "hex", "log", - "nix 0.26.2", + "nix 0.26.4", "orion", "priority-queue", "psutil", @@ -2186,9 +2191,9 @@ dependencies = [ [[package]] name = "termcolor" -version = "1.2.0" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be55cf8942feac5c765c2c993422806843c9a9a45d4d5c407ad6dd2ea95eb9b6" +checksum = "6093bad37da69aab9d123a8091e4be0aa4a03e4d601ec641c327398315f62b64" dependencies = [ "winapi-util", ] @@ -2233,22 +2238,22 @@ checksum = "aac81b6fd6beb5884b0cf3321b8117e6e5d47ecb6fc89f414cfdcca8b2fe2dd8" [[package]] name = "thiserror" -version = "1.0.44" +version = "1.0.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "611040a08a0439f8248d1990b111c95baa9c704c805fa1f62104b39655fd7f90" +checksum = "9d6d7a740b8a666a7e828dd00da9c0dc290dff53154ea77ac109281de90589b7" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.44" +version = "1.0.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "090198534930841fab3a5d1bb637cde49e339654e606195f8d9c76eeb081dc96" +checksum = "49922ecae66cc8a249b77e68d1d0623c1b2c514f0060c27cdc68bd62a1219d35" dependencies = [ "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", ] [[package]] @@ -2261,17 +2266,6 @@ dependencies = [ "once_cell", ] -[[package]] -name = "time" -version = "0.1.45" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b797afad3f312d1c66a56d11d0316f916356d11bd158fbc6ca6389ff6bf805a" -dependencies = [ - "libc", - "wasi 0.10.0+wasi-snapshot-preview1", - "winapi", -] - [[package]] name = "tinytemplate" version = "1.2.1" @@ -2284,11 +2278,10 @@ dependencies = [ [[package]] name = "tokio" -version = "1.29.1" +version = "1.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "532826ff75199d5833b9d2c5fe410f29235e25704ee5f0ef599fb51c21f4a4da" +checksum = "17ed6077ed6cd6c74735e21f37eb16dc3935f96878b1fe961074089cc80893f9" dependencies = [ - "autocfg", "backtrace", "bytes", "libc", @@ -2310,14 +2303,14 @@ checksum = "630bdcf245f78637c13ec01ffae6187cca34625e8c63150d424b59e55af2675e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", ] [[package]] name = "tokio-util" -version = "0.7.8" +version = "0.7.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "806fe8c2c87eccc8b3267cbae29ed3ab2d0bd37fca70ab622e46aaa9375ddb7d" +checksum = "1d68074620f57a0b21594d9735eb2e98ab38b17f80d3fcb189fca266771ca60d" dependencies = [ "bytes", "futures-core", @@ -2329,9 +2322,9 @@ dependencies = [ [[package]] name = "toml" -version = "0.7.6" +version = "0.7.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c17e963a819c331dcacd7ab957d80bc2b9a9c1e71c804826d2f283dd65306542" +checksum = "dd79e69d3b627db300ff956027cc6c3798cef26d22526befdfcd12feeb6d2257" dependencies = [ "serde", "serde_spanned", @@ -2350,9 +2343,9 @@ dependencies = [ [[package]] name = "toml_edit" -version = "0.19.14" +version = "0.19.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8123f27e969974a3dfba720fdb560be359f57b44302d280ba72e76a74480e8a" +checksum = "1b5bb770da30e5cbfde35a2d7b9b8a2c4b8ef89548a7a6aeab5c9a576e3e7421" dependencies = [ "indexmap 2.0.0", "serde", @@ -2381,7 +2374,7 @@ checksum = "5f4f31f56159e98206da9efd823404b79b6ef3143b4a7ab76e67b1751b25a4ab" dependencies = [ "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", ] [[package]] @@ -2453,9 +2446,9 @@ checksum = "ccb97dac3243214f8d8507998906ca3e2e0b900bf9bf4870477f125b82e68f6e" [[package]] name = "unicode-ident" -version = "1.0.11" +version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "301abaae475aa91687eb82514b328ab47a211a533026cb25fc3e519b86adfc3c" +checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b" [[package]] name = "unicode-linebreak" @@ -2471,9 +2464,9 @@ checksum = "1dd624098567895118886609431a7c3b8f516e41d30e0643f03d94592a147e36" [[package]] name = "unicode-width" -version = "0.1.10" +version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c0edd1e5b14653f783770bce4a4dabb4a5108a5370a5f5d8cfe8710c361f6c8b" +checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" [[package]] name = "unicode-xid" @@ -2507,20 +2500,14 @@ checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" [[package]] name = "walkdir" -version = "2.3.3" +version = "2.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "36df944cda56c7d8d8b7496af378e6b16de9284591917d307c9b4d313c44e698" +checksum = "d71d857dc86794ca4c280d616f7da00d2dbfd8cd788846559a6813e6aa4b54ee" dependencies = [ "same-file", "winapi-util", ] -[[package]] -name = "wasi" -version = "0.10.0+wasi-snapshot-preview1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a143597ca7c7793eff794def352d41792a93c481eb1042423ff7ff72ba2c31f" - [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" @@ -2548,7 +2535,7 @@ dependencies = [ "once_cell", "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", "wasm-bindgen-shared", ] @@ -2570,7 +2557,7 @@ checksum = "54681b18a46765f095758388f2d0cf16eb8d4169b639ab575a8f5693af210c7b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.27", + "syn 2.0.37", "wasm-bindgen-backend", "wasm-bindgen-shared", ] @@ -2609,9 +2596,9 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +checksum = "f29e6f9198ba0d26b4c9f07dbe6f9ed633e1f3d5b8b414090084349e46a52596" dependencies = [ "winapi", ] @@ -2628,7 +2615,7 @@ version = "0.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e686886bc078bc1b0b600cac0147aadb815089b6e4da64016cbd754b6342700f" dependencies = [ - "windows-targets 0.48.1", + "windows-targets 0.48.5", ] [[package]] @@ -2646,7 +2633,7 @@ version = "0.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "677d2418bec65e3338edb076e806bc1ec15693c5d0104683f2efe857f61056a9" dependencies = [ - "windows-targets 0.48.1", + "windows-targets 0.48.5", ] [[package]] @@ -2666,17 +2653,17 @@ dependencies = [ [[package]] name = "windows-targets" -version = "0.48.1" +version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05d4b17490f70499f20b9e791dcf6a299785ce8af4d709018206dc5b4953e95f" +checksum = "9a2fa6e2155d7247be68c096456083145c183cbbbc2764150dda45a87197940c" dependencies = [ - "windows_aarch64_gnullvm 0.48.0", - "windows_aarch64_msvc 0.48.0", - "windows_i686_gnu 0.48.0", - "windows_i686_msvc 0.48.0", - "windows_x86_64_gnu 0.48.0", - "windows_x86_64_gnullvm 0.48.0", - "windows_x86_64_msvc 0.48.0", + "windows_aarch64_gnullvm 0.48.5", + "windows_aarch64_msvc 0.48.5", + "windows_i686_gnu 0.48.5", + "windows_i686_msvc 0.48.5", + "windows_x86_64_gnu 0.48.5", + "windows_x86_64_gnullvm 0.48.5", + "windows_x86_64_msvc 0.48.5", ] [[package]] @@ -2687,9 +2674,9 @@ checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" [[package]] name = "windows_aarch64_gnullvm" -version = "0.48.0" +version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "91ae572e1b79dba883e0d315474df7305d12f569b400fcf90581b06062f7e1bc" +checksum = "2b38e32f0abccf9987a4e3079dfb67dcd799fb61361e53e2882c3cbaf0d905d8" [[package]] name = "windows_aarch64_msvc" @@ -2699,9 +2686,9 @@ checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" [[package]] name = "windows_aarch64_msvc" -version = "0.48.0" +version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2ef27e0d7bdfcfc7b868b317c1d32c641a6fe4629c171b8928c7b08d98d7cf3" +checksum = "dc35310971f3b2dbbf3f0690a219f40e2d9afcf64f9ab7cc1be722937c26b4bc" [[package]] name = "windows_i686_gnu" @@ -2711,9 +2698,9 @@ checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" [[package]] name = "windows_i686_gnu" -version = "0.48.0" +version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "622a1962a7db830d6fd0a69683c80a18fda201879f0f447f065a3b7467daa241" +checksum = "a75915e7def60c94dcef72200b9a8e58e5091744960da64ec734a6c6e9b3743e" [[package]] name = "windows_i686_msvc" @@ -2723,9 +2710,9 @@ checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" [[package]] name = "windows_i686_msvc" -version = "0.48.0" +version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4542c6e364ce21bf45d69fdd2a8e455fa38d316158cfd43b3ac1c5b1b19f8e00" +checksum = "8f55c233f70c4b27f66c523580f78f1004e8b5a8b659e05a4eb49d4166cca406" [[package]] name = "windows_x86_64_gnu" @@ -2735,9 +2722,9 @@ checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" [[package]] name = "windows_x86_64_gnu" -version = "0.48.0" +version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca2b8a661f7628cbd23440e50b05d705db3686f894fc9580820623656af974b1" +checksum = "53d40abd2583d23e4718fddf1ebec84dbff8381c07cae67ff7768bbf19c6718e" [[package]] name = "windows_x86_64_gnullvm" @@ -2747,9 +2734,9 @@ checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" [[package]] name = "windows_x86_64_gnullvm" -version = "0.48.0" +version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7896dbc1f41e08872e9d5e8f8baa8fdd2677f29468c4e156210174edc7f7b953" +checksum = "0b7b52767868a23d5bab768e390dc5f5c55825b6d30b86c844ff2dc7414044cc" [[package]] name = "windows_x86_64_msvc" @@ -2759,15 +2746,15 @@ checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" [[package]] name = "windows_x86_64_msvc" -version = "0.48.0" +version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a515f5799fe4961cb532f983ce2b23082366b898e52ffbce459c86f67c8378a" +checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" [[package]] name = "winnow" -version = "0.5.1" +version = "0.5.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25b5872fa2e10bd067ae946f927e726d7d603eaeb6e02fa6a350e0722d2b8c11" +checksum = "7c2e3184b9c4e92ad5167ca73039d0c42476302ab603e2fec4487511f38ccefc" dependencies = [ "memchr", ] diff --git a/crates/hyperqueue/src/client/commands/submit/command.rs b/crates/hyperqueue/src/client/commands/submit/command.rs index 5058bc8c0..e61cf2f2b 100644 --- a/crates/hyperqueue/src/client/commands/submit/command.rs +++ b/crates/hyperqueue/src/client/commands/submit/command.rs @@ -15,7 +15,7 @@ use tako::gateway::{ ResourceRequest, ResourceRequestEntries, ResourceRequestEntry, ResourceRequestVariants, }; use tako::program::{FileOnCloseBehavior, ProgramDefinition, StdioDef}; -use tako::resources::{AllocationRequest, NumOfNodes, CPU_RESOURCE_NAME}; +use tako::resources::{AllocationRequest, NumOfNodes, ResourceAmount, CPU_RESOURCE_NAME}; use super::directives::parse_hq_directives; use crate::client::commands::submit::directives::parse_hq_directives_from_file; @@ -463,7 +463,7 @@ impl JobSubmitOpts { 0, ResourceRequestEntry { resource: CPU_RESOURCE_NAME.to_string(), - policy: AllocationRequest::Compact(1), + policy: AllocationRequest::Compact(ResourceAmount::new_units(1)), }, ) } diff --git a/crates/hyperqueue/src/client/commands/submit/defs.rs b/crates/hyperqueue/src/client/commands/submit/defs.rs index ffd3968d3..7f3b44cf2 100644 --- a/crates/hyperqueue/src/client/commands/submit/defs.rs +++ b/crates/hyperqueue/src/client/commands/submit/defs.rs @@ -12,13 +12,14 @@ use std::path::PathBuf; use std::time::Duration; use tako::gateway::{ResourceRequest, ResourceRequestEntries, ResourceRequestEntry}; use tako::program::FileOnCloseBehavior; -use tako::resources::{AllocationRequest, NumOfNodes}; +use tako::resources::{AllocationRequest, NumOfNodes, ResourceAmount}; use tako::{Map, Priority}; #[derive(Deserialize)] #[serde(untagged)] -pub enum IntOrString { - Int(u64), +enum PrimitiveType { + Int(u32), + Float(f32), String(String), } @@ -55,13 +56,14 @@ fn deserialize_resource_entries<'de, D>(deserializer: D) -> Result, { - let tmp = Map::::deserialize(deserializer)?; + let tmp = Map::::deserialize(deserializer)?; let mut result = ResourceRequestEntries::new(); for (k, v) in tmp { let policy = match v { - IntOrString::Int(n) => AllocationRequest::Compact(n), - IntOrString::String(s) => { + PrimitiveType::Int(n) => AllocationRequest::Compact(ResourceAmount::new_units(n)), + PrimitiveType::Float(f) => AllocationRequest::Compact(ResourceAmount::from_float(f)), + PrimitiveType::String(s) => { parse_allocation_request(&s).map_err(serde::de::Error::custom)? } }; diff --git a/crates/hyperqueue/src/client/output/cli.rs b/crates/hyperqueue/src/client/output/cli.rs index e4de5e977..eec314032 100644 --- a/crates/hyperqueue/src/client/output/cli.rs +++ b/crates/hyperqueue/src/client/output/cli.rs @@ -8,7 +8,7 @@ use crate::client::job::WorkerMap; use crate::client::output::outputs::{Output, OutputStream, MAX_DISPLAYED_WORKERS}; use crate::client::status::{get_task_status, job_status, Status}; use crate::common::env::is_hq_env; -use crate::common::format::{human_duration, human_size}; +use crate::common::format::{human_duration, human_mem_amount, human_size}; use crate::common::manager::info::GetManagerInfo; use crate::server::autoalloc::{Allocation, AllocationState}; use crate::server::job::{JobTaskCounters, JobTaskInfo, JobTaskState, StartedTaskData}; @@ -1397,7 +1397,7 @@ fn resources_summary(resources: &ResourceDescriptor, multiline: bool) -> String let special_format = |descriptor: &ResourceDescriptorItem| -> Option { if descriptor.name == tako::resources::MEM_RESOURCE_NAME { if let ResourceDescriptorKind::Sum { size } = descriptor.kind { - return Some(human_size(size)); + return Some(human_mem_amount(size)); } } None @@ -1515,7 +1515,7 @@ mod tests { fn test_resources_summary_mem() { let d = ResourceDescriptor::new(vec![ResourceDescriptorItem { name: MEM_RESOURCE_NAME.into(), - kind: res_kind_sum(4 * 1024 * 1024 * 1024 + 123 * 1024 * 1024), + kind: res_kind_sum(4 * 1024 + 123), }]); assert_eq!(resources_summary(&d, false), "mem 4.12 GiB"); } diff --git a/crates/hyperqueue/src/client/resources.rs b/crates/hyperqueue/src/client/resources.rs index 0e6480fb0..986895886 100644 --- a/crates/hyperqueue/src/client/resources.rs +++ b/crates/hyperqueue/src/client/resources.rs @@ -1,75 +1,59 @@ -use nom::branch::alt; -use nom::bytes::complete::take_while; -use nom::character::complete::{anychar, char, multispace0, multispace1}; -use nom::combinator::{map, map_res, opt, recognize, verify}; -use nom::sequence::{pair, preceded, separated_pair, tuple}; -use nom_supreme::tag::complete::tag; -use nom_supreme::ParserExt; +use chumsky::primitive::{choice, filter, just}; +use chumsky::text::TextParser; +use chumsky::Parser; -use tako::resources::{AllocationRequest, ResourceAmount}; +use tako::resources::AllocationRequest; -use crate::common::parser::{consume_all, p_u64, NomResult}; +use crate::common::parser2::{ + all_consuming, parse_exact_string, parse_resource_amount, CharParser, ParseError, +}; use crate::worker::parser::{is_valid_resource_char, is_valid_starting_resource_char}; -fn p_allocation_request(input: &str) -> NomResult { - alt(( - map(tag("all"), |_| AllocationRequest::All), - map_res( - tuple(( - p_u64, - opt(preceded( - multispace1, - alt((tag("compact!"), tag("compact"), tag("scatter"))), - )), - )), - |(count, policy)| { - let count = count as ResourceAmount; - if count == 0 { - return Err(anyhow::anyhow!("Requesting zero resources is not allowed")); - } - Ok(match policy { - None | Some("compact") => AllocationRequest::Compact(count), - Some("compact!") => AllocationRequest::ForceCompact(count), - Some("scatter") => AllocationRequest::Scatter(count), +fn p_allocation_request() -> impl CharParser { + parse_exact_string("all") + .map(|_| AllocationRequest::All) + .or(parse_resource_amount() + .padded() + .then(choice((just("compact!"), just("compact"), just("scatter"))).or_not()) + .try_map(|(amount, policy), span| { + let alloc = match policy { + None | Some("compact") => AllocationRequest::Compact(amount), + Some("compact!") => AllocationRequest::ForceCompact(amount), + Some("scatter") => AllocationRequest::Scatter(amount), _ => unreachable!(), - }) - }, - ), - ))(input) + }; + alloc + .validate() + .map_err(|e| ParseError::custom(span, e.to_string()))?; + Ok(alloc) + })) } -/// Parses a resource identifier. -/// It has to be an alphanumeric string beginning with a letter. It can also a slash. -fn p_resource_identifier(input: &str) -> NomResult<&str> { - recognize(pair( - verify(anychar, |&c| is_valid_starting_resource_char(c)).context("Letter"), - take_while(is_valid_resource_char).context("Letter, digit or slash"), - ))(input) +fn p_resource_identifier() -> impl CharParser { + filter(|&c| is_valid_starting_resource_char(c)) + .chain(filter(|&c| is_valid_resource_char(c)).repeated()) + .map(|s| s.iter().collect::()) } -fn p_resource_request(input: &str) -> NomResult<(String, AllocationRequest)> { - map( - separated_pair( - p_resource_identifier.context("Resource identifier"), - tuple((multispace0, char('='), multispace0)), - p_allocation_request.context("Resource amount"), - ), - |(name, value)| (name.to_string(), value), - )(input) +fn p_resource_request() -> impl CharParser<(String, AllocationRequest)> { + p_resource_identifier() + .then_ignore(just("=").padded()) + .then(p_allocation_request()) } pub fn parse_resource_request(input: &str) -> anyhow::Result<(String, AllocationRequest)> { - consume_all(p_resource_request, input) + all_consuming(p_resource_request()).parse_text(input) } pub fn parse_allocation_request(input: &str) -> anyhow::Result { - consume_all(p_allocation_request, input) + all_consuming(p_allocation_request()).parse_text(input) } #[cfg(test)] mod test { use super::*; - use crate::tests::utils::check_parse_error; + use crate::tests::utils::expect_parser_error; + use tako::resources::ResourceAmount; #[test] fn test_parse_resource_request() { @@ -79,31 +63,39 @@ mod test { ); assert_eq!( parse_resource_request("ab1c=10").unwrap(), - ("ab1c".to_string(), AllocationRequest::Compact(10)) + ( + "ab1c".to_string(), + AllocationRequest::Compact(ResourceAmount::new_units(10)) + ) ); assert_eq!( parse_resource_request("a=5_000 compact").unwrap(), - ("a".to_string(), AllocationRequest::Compact(5000)) + ( + "a".to_string(), + AllocationRequest::Compact(ResourceAmount::new_units(5000)) + ) ); assert_eq!( parse_resource_request("cpus=351 scatter").unwrap(), - ("cpus".to_string(), AllocationRequest::Scatter(351)) + ( + "cpus".to_string(), + AllocationRequest::Scatter(ResourceAmount::new_units(351)) + ) ); assert_eq!( parse_resource_request("cpus=11 compact!").unwrap(), - ("cpus".to_string(), AllocationRequest::ForceCompact(11)) + ( + "cpus".to_string(), + AllocationRequest::ForceCompact(ResourceAmount::new_units(11)) + ) ); } #[test] fn test_parse_no_name() { - check_parse_error( - p_resource_request, - "=1", - r#"Parse error -expected Resource identifier at character 0: "=1" -expected Letter at character 0: "=1" - expected: Verify at character 0: "=1""#, + assert_eq!( + expect_parser_error(all_consuming(p_resource_request()), "=1"), + "Unexpected token found, expected something else:\n =1\n |\n --- Unexpected token `=`\n", ); } @@ -111,53 +103,42 @@ expected Letter at character 0: "=1" fn test_parse_identifier_slash() { assert_eq!( parse_resource_request("a/b=1").unwrap(), - ("a/b".to_string(), AllocationRequest::Compact(1)) + ( + "a/b".to_string(), + AllocationRequest::Compact(ResourceAmount::new_units(1)) + ) ); } #[test] fn test_parse_identifier_start_with_digit() { - check_parse_error( - p_resource_request, - "1a=1", - r#"Parse error -expected Resource identifier at character 0: "1a=1" -expected Letter at character 0: "1a=1" - expected: Verify at character 0: "1a=1""#, + assert_eq!( + expect_parser_error(all_consuming(p_resource_request()), "1a=1"), + "Unexpected token found, expected something else:\n 1a=1\n |\n --- Unexpected token `1`\n", ); } #[test] fn test_parse_zero_resources() { - check_parse_error( - p_resource_request, - "aa=0", - r#"Parse error -expected Resource amount at character 3: "0" - "Requesting zero resources is not allowed" at character 3: "0""#, + assert_eq!( + expect_parser_error(all_consuming(p_resource_request()), "aa=0"), + "Unexpected end of input found, expected something else:\n aa=0\n |\n --- Error: Zero resources cannot be requested\n", ); } #[test] fn test_parse_resource_request_error() { - check_parse_error( - p_resource_request, - "", - r#"Parse error -expected Resource identifier at the end of input -expected Letter at the end of input - expected something at the end of input"#, + assert_eq!( + expect_parser_error(all_consuming(p_resource_request()), ""), + "Unexpected end of input found, expected something else:\n(the input was empty)", ); - check_parse_error( - p_resource_request, - "a", - r#"Parse error -expected "=" at the end of input"#, + assert_eq!( + expect_parser_error(all_consuming(p_resource_request()), "a"), + "Unexpected end of input found, expected =:\n a\n |\n --- Unexpected end of input\n", ); - check_parse_error( - p_resource_request, - "a=x", - "Parse error\nexpected Resource amount at character 2: \"x\"\n expected one of the following 2 variants:\n expected \"all\" at character 2: \"x\"\n or\n expected integer at character 2: \"x\"", + assert_eq!( + expect_parser_error(all_consuming(p_resource_request()), "a=x"), + "Unexpected token found, expected all:\n a=x\n |\n --- Unexpected token `x`\n", ); } } diff --git a/crates/hyperqueue/src/common/format.rs b/crates/hyperqueue/src/common/format.rs index 37748b17c..410e4f131 100644 --- a/crates/hyperqueue/src/common/format.rs +++ b/crates/hyperqueue/src/common/format.rs @@ -1,3 +1,4 @@ +use tako::resources::ResourceAmount; use tako::worker::ServerLostPolicy; pub fn human_duration(duration: chrono::Duration) -> String { @@ -26,6 +27,15 @@ pub fn human_size(size: u64) -> String { } } +pub fn human_mem_amount(amount: ResourceAmount) -> String { + let f = amount.as_f32(); + if f < 512f32 { + format!("{:.2} MiB", f) + } else { + format!("{:.2} GiB", f / 1024.0) + } +} + pub fn server_lost_policy_to_str(policy: &ServerLostPolicy) -> &str { match policy { ServerLostPolicy::Stop => "stop", diff --git a/crates/hyperqueue/src/common/parser2.rs b/crates/hyperqueue/src/common/parser2.rs index 33bf6a1ad..caee92e6e 100644 --- a/crates/hyperqueue/src/common/parser2.rs +++ b/crates/hyperqueue/src/common/parser2.rs @@ -1,10 +1,11 @@ use anyhow::anyhow; use chumsky::error::{Simple, SimpleReason}; -use chumsky::primitive::end; +use chumsky::primitive::{end, just}; use chumsky::text::ident; use chumsky::{Error, Parser, Span}; use colored::Color; use std::ops::Range; +use tako::resources::{ResourceAmount, FRACTIONS_MAX_DIGITS}; #[derive(Clone, Debug, PartialEq)] pub struct ParseError { @@ -230,12 +231,25 @@ pub fn parse_u32() -> impl CharParser { }) } -/// Parse 8-byte integer. -pub fn parse_u64() -> impl CharParser { - parse_integer_string().try_map(|p, span| { - p.parse::() - .map_err(|_| ParseError::custom(span, "Cannot parse as 8-byte unsigned integer")) - }) +/// Parse int32 or float with FRACTIONS_MAX_DIGITS precision (e.g. "2", "3.3", "0.0001") +pub fn parse_resource_amount() -> impl CharParser { + parse_u32() + .then(just('.').ignore_then(parse_integer_string()).or_not()) + .try_map(|(units, frac), span| { + let f = if let Some(s) = frac { + if s.len() > FRACTIONS_MAX_DIGITS { + return Err(ParseError::custom(span, "Resource precision exceeded")); + } + let mut n = s.parse::().unwrap(); + for _ in s.len()..FRACTIONS_MAX_DIGITS { + n *= 10; + } + n + } else { + 0 + }; + Ok(ResourceAmount::new(units, f)) + }) } /// Parses an exact string. @@ -283,6 +297,45 @@ mod tests { assert_eq!(parse_u32().parse_text("1019").unwrap(), 1019); } + #[test] + fn test_parse_resource_amount() { + assert_eq!( + parse_resource_amount().parse_text("0").unwrap(), + ResourceAmount::ZERO + ); + assert_eq!( + parse_resource_amount().parse_text("123").unwrap(), + ResourceAmount::new_units(123) + ); + assert_eq!( + parse_resource_amount().parse_text("0.01").unwrap(), + ResourceAmount::new(0, 100) + ); + assert_eq!( + parse_resource_amount().parse_text("3.").unwrap(), + ResourceAmount::new(3, 0) + ); + assert_eq!( + parse_resource_amount().parse_text("100.0001").unwrap(), + ResourceAmount::new(100, 1) + ); + assert_eq!( + parse_resource_amount().parse_text("100.9").unwrap(), + ResourceAmount::new(100, 9000) + ); + assert_eq!( + parse_resource_amount().parse_text("0.346").unwrap(), + ResourceAmount::new(0, 3460) + ); + let r = parse_resource_amount().parse_text("0.12345"); + insta::assert_snapshot!(r.unwrap_err().to_string().as_str(), @r###" + Unexpected end of input found, expected something else: + 0.12345 + | + --- Resource precision exceeded + "###); + } + #[test] fn test_parse_u32_empty() { insta::assert_snapshot!(expect_parser_error(parse_u32(), ""), @r###" diff --git a/crates/hyperqueue/src/server/autoalloc/process.rs b/crates/hyperqueue/src/server/autoalloc/process.rs index 7a142fb1d..a288eda01 100644 --- a/crates/hyperqueue/src/server/autoalloc/process.rs +++ b/crates/hyperqueue/src/server/autoalloc/process.rs @@ -899,6 +899,7 @@ mod tests { use crate::server::state::StateRef; use crate::tests::utils::create_hq_state; use crate::transfer::messages::{JobDescription, PinMode, TaskDescription}; + use tako::resources::ResourceAmount; #[tokio::test] async fn fill_backlog() { @@ -1743,7 +1744,7 @@ mod tests { min_time, resources: smallvec![ResourceRequestEntry { resource: CPU_RESOURCE_NAME.to_string(), - policy: AllocationRequest::Compact(1), + policy: AllocationRequest::Compact(ResourceAmount::new_units(1)), }], }); diff --git a/crates/hyperqueue/src/server/client/submit.rs b/crates/hyperqueue/src/server/client/submit.rs index e8f4b4f3c..18c0466d7 100644 --- a/crates/hyperqueue/src/server/client/submit.rs +++ b/crates/hyperqueue/src/server/client/submit.rs @@ -533,7 +533,7 @@ mod tests { min_time: Duration::from_secs(2), resources: smallvec![ResourceRequestEntry { resource: CPU_RESOURCE_NAME.to_string(), - policy: AllocationRequest::Compact(cpu_count as ResourceAmount), + policy: AllocationRequest::Compact(ResourceAmount::new_units(cpu_count)), }], }), pin_mode: PinMode::None, diff --git a/crates/hyperqueue/src/worker/hwdetect.rs b/crates/hyperqueue/src/worker/hwdetect.rs index b9271d587..699e111df 100644 --- a/crates/hyperqueue/src/worker/hwdetect.rs +++ b/crates/hyperqueue/src/worker/hwdetect.rs @@ -9,8 +9,9 @@ use nom_supreme::tag::complete::tag; use tako::hwstats::GpuFamily; use tako::internal::has_unique_elements; use tako::resources::{ - ResourceDescriptorItem, ResourceDescriptorKind, ResourceIndex, ResourceLabel, - AMD_GPU_RESOURCE_NAME, MEM_RESOURCE_NAME, NVIDIA_GPU_RESOURCE_NAME, + ResourceAmount, ResourceDescriptorItem, ResourceDescriptorKind, ResourceFractions, + ResourceIndex, ResourceLabel, AMD_GPU_RESOURCE_NAME, FRACTIONS_PER_UNIT, MEM_RESOURCE_NAME, + NVIDIA_GPU_RESOURCE_NAME, }; use tako::{format_comma_delimited, Set}; @@ -128,9 +129,16 @@ pub fn detect_additional_resources( if !has_resource(items, MEM_RESOURCE_NAME) { if let Ok(mem) = read_linux_memory() { log::info!("Detected {mem}B of memory ({})", human_size(mem)); + let units = mem / (1024 * 1024); + let fractions = ((mem % (1024 * 1024)) * FRACTIONS_PER_UNIT as u64) / (1024 * 1024); items.push(ResourceDescriptorItem { name: MEM_RESOURCE_NAME.to_string(), - kind: ResourceDescriptorKind::Sum { size: mem }, + kind: ResourceDescriptorKind::Sum { + size: ResourceAmount::new( + units.try_into().unwrap(), + fractions as ResourceFractions, + ), + }, }); } } diff --git a/crates/hyperqueue/src/worker/parser.rs b/crates/hyperqueue/src/worker/parser.rs index f97270caa..f50b0e7b2 100644 --- a/crates/hyperqueue/src/worker/parser.rs +++ b/crates/hyperqueue/src/worker/parser.rs @@ -3,11 +3,11 @@ use chumsky::text::TextParser; use chumsky::{Error, Parser}; use tako::resources::{ - DescriptorError, ResourceAmount, ResourceDescriptorItem, ResourceDescriptorKind, + DescriptorError, ResourceDescriptorItem, ResourceDescriptorKind, ResourceUnits, }; use crate::common::parser2::{ - all_consuming, parse_exact_string, parse_u32, parse_u64, CharParser, ParseError, + all_consuming, parse_exact_string, parse_resource_amount, parse_u32, CharParser, ParseError, }; pub fn parse_cpu_definition(input: &str) -> anyhow::Result { @@ -53,8 +53,8 @@ fn parse_resource_group_x_notation() -> impl CharParser .then(just('x').ignore_then(parse_u32())) .map(|(groups, group_size)| { ResourceDescriptorKind::regular_sockets( - groups as ResourceAmount, - group_size as ResourceAmount, + groups as ResourceUnits, + group_size as ResourceUnits, ) }) } @@ -148,7 +148,7 @@ fn parse_resource_sum() -> impl CharParser { let start = parse_exact_string("sum").then(just('(').padded()); let end = just(')').padded(); - let value = parse_u64() + let value = parse_resource_amount() .labelled("sum") .map(|size| ResourceDescriptorKind::Sum { size }); @@ -212,6 +212,7 @@ mod test { use tako::internal::tests::utils::shared::{ res_kind_groups, res_kind_list, res_kind_range, res_kind_sum, }; + use tako::resources::ResourceAmount; use super::*; @@ -488,18 +489,29 @@ mod test { #[test] fn test_parse_resource_def_sum() { check_item( - parse_resource_definition("mem=sum(1000_3000_2000)"), + parse_resource_definition("mem=sum(1_3000_2000)"), "mem", - res_kind_sum(1000_3000_2000), + res_kind_sum(1_3000_2000), + ); + } + + #[test] + fn test_parse_resource_def_sum_frac() { + check_item( + parse_resource_definition("mem=sum(1.5)"), + "mem", + ResourceDescriptorKind::Sum { + size: ResourceAmount::new(1, 5000), + }, ); } #[test] fn test_parse_resource_def_sum_whitespace() { check_item( - parse_resource_definition(" mem = sum ( 1000_3000_2000 ) "), + parse_resource_definition(" mem = sum ( 1_3000_2000 ) "), "mem", - res_kind_sum(1000_3000_2000), + res_kind_sum(1_3000_2000), ); } diff --git a/crates/hyperqueue/src/worker/start.rs b/crates/hyperqueue/src/worker/start.rs index 09984cbba..723bb53cb 100644 --- a/crates/hyperqueue/src/worker/start.rs +++ b/crates/hyperqueue/src/worker/start.rs @@ -235,16 +235,17 @@ fn insert_resources_into_env(ctx: &LaunchContext, program: &mut ProgramDefinitio } for alloc in &ctx.allocation().resources { - let resource_name = resource_map.get_name(alloc.resource).unwrap(); + let resource_name = resource_map.get_name(alloc.resource_id).unwrap(); if let Some(labels) = allocation_to_labels(alloc, ctx) { if resource_name == CPU_RESOURCE_NAME { /* Extra variables for CPUS */ program.env.insert(HQ_CPUS.into(), labels.clone().into()); - if !program.env.contains_key(b"OMP_NUM_THREADS".as_bstr()) { - program.env.insert( - "OMP_NUM_THREADS".into(), - alloc.value.amount().to_string().into(), - ); + if !program.env.contains_key(b"OMP_NUM_THREADS".as_bstr()) + && alloc.amount.fractions() == 0 + { + program + .env + .insert("OMP_NUM_THREADS".into(), alloc.amount.to_string().into()); } } if resource_name == NVIDIA_GPU_RESOURCE_NAME { @@ -286,13 +287,12 @@ fn resource_env_var_name(prefix: &str, resource_name: &str) -> BString { fn allocation_to_labels(allocation: &ResourceAllocation, ctx: &LaunchContext) -> Option { let label_map = ctx.get_resource_label_map(); - allocation.value.indices().map(|indices| { - format_comma_delimited( - indices - .iter() - .map(|index| label_map.get_label(allocation.resource, *index)), - ) - }) + if allocation.indices.is_empty() { + return None; + } + Some(format_comma_delimited(allocation.indices.iter().map( + |idx| label_map.get_label(allocation.resource_id, idx.index), + ))) } fn pin_program( @@ -305,7 +305,7 @@ fn pin_program( allocation .resources .iter() - .find(|r| r.resource == CPU_RESOURCE_ID) + .find(|r| r.resource_id == CPU_RESOURCE_ID) .and_then(|r| allocation_to_labels(r, ctx)) }; match pin_mode { diff --git a/crates/pyhq/python/hyperqueue/ffi/protocol.py b/crates/pyhq/python/hyperqueue/ffi/protocol.py index eabf9ef6c..3a4709f66 100644 --- a/crates/pyhq/python/hyperqueue/ffi/protocol.py +++ b/crates/pyhq/python/hyperqueue/ffi/protocol.py @@ -6,14 +6,14 @@ class ResourceRequest: n_nodes: int = 0 - resources: Dict[str, Union[int, str]] = dataclasses.field(default_factory=dict) + resources: Dict[str, Union[int, float, str]] = dataclasses.field(default_factory=dict) def __init__( self, *, n_nodes=0, - cpus: Union[int, str] = 1, - resources: Optional[Dict[str, Union[int, str]]] = None, + cpus: Union[int, float, str] = 1, + resources: Optional[Dict[str, Union[int, float, str]]] = None, ): self.n_nodes = n_nodes if resources is None: diff --git a/crates/pyhq/src/client/job.rs b/crates/pyhq/src/client/job.rs index 2fe287c54..b7c50e804 100644 --- a/crates/pyhq/src/client/job.rs +++ b/crates/pyhq/src/client/job.rs @@ -29,7 +29,8 @@ use crate::{borrow_mut, run_future, ClientContextPtr, FromPyObject, PyJobId, PyT #[derive(Debug, FromPyObject)] enum AllocationValue { - Int(u64), + Int(u32), + Float(f32), String(String), } @@ -166,8 +167,13 @@ fn build_task_desc(desc: TaskDescription, submit_dir: &Path) -> anyhow::Result { - AllocationRequest::Compact(value as ResourceAmount) + AllocationValue::Int(value) => AllocationRequest::Compact( + ResourceAmount::new_units(value), + ), + AllocationValue::Float(value) => { + AllocationRequest::Compact(ResourceAmount::from_float( + value, + )) } AllocationValue::String(str) => { parse_allocation_request(&str)? diff --git a/crates/tako/Cargo.toml b/crates/tako/Cargo.toml index df00287b1..5bd20981f 100644 --- a/crates/tako/Cargo.toml +++ b/crates/tako/Cargo.toml @@ -36,6 +36,7 @@ bitflags = "2" psutil = "3.2.1" fxhash = "0.2.1" thin-vec = "0.2.8" +derive_more = "0.99" [dev-dependencies] anyhow = { workspace = true } diff --git a/crates/tako/benches/benchmarks/worker.rs b/crates/tako/benches/benchmarks/worker.rs index 2d18120de..9b89b4f11 100644 --- a/crates/tako/benches/benchmarks/worker.rs +++ b/crates/tako/benches/benchmarks/worker.rs @@ -20,6 +20,7 @@ use tokio::sync::Notify; use tako::internal::worker::state::{TaskMap, WorkerStateRef}; use tako::internal::worker::task::Task; +use tako::resources::ResourceAmount; use tako::TaskId; use crate::create_worker; @@ -199,11 +200,11 @@ fn bench_resource_queue_release_allocation(c: &mut BenchmarkGroup) { smallvec![ ResourceRequestEntry { resource_id: 0.into(), - request: AllocationRequest::Compact(64), + request: AllocationRequest::Compact(ResourceAmount::new_units(64)), }, ResourceRequestEntry { resource_id: 1.into(), - request: AllocationRequest::Compact(2), + request: AllocationRequest::Compact(ResourceAmount::new_units(2)), }, ], )]); @@ -241,11 +242,15 @@ fn bench_resource_queue_start_tasks(c: &mut BenchmarkGroup) { smallvec![ ResourceRequestEntry { resource_id: 0.into(), - request: AllocationRequest::Compact(64), + request: AllocationRequest::Compact( + ResourceAmount::new_units(64) + ), }, ResourceRequestEntry { resource_id: 1.into(), - request: AllocationRequest::Compact(2), + request: AllocationRequest::Compact( + ResourceAmount::new_units(2) + ), }, ], )]); diff --git a/crates/tako/src/gateway.rs b/crates/tako/src/gateway.rs index 8b40335e3..0f6e261fc 100644 --- a/crates/tako/src/gateway.rs +++ b/crates/tako/src/gateway.rs @@ -3,7 +3,9 @@ use serde::{Deserialize, Serialize, Serializer}; use crate::internal::messages::common::TaskFailInfo; use crate::internal::messages::worker::WorkerOverview; use crate::internal::worker::configuration::WorkerConfiguration; -use crate::resources::{AllocationRequest, NumOfNodes, ResourceDescriptor, CPU_RESOURCE_NAME}; +use crate::resources::{ + AllocationRequest, NumOfNodes, ResourceAmount, ResourceDescriptor, CPU_RESOURCE_NAME, +}; use crate::task::SerializedTaskContext; use crate::{Priority, TaskId, WorkerId}; use smallvec::{smallvec, SmallVec}; @@ -35,7 +37,7 @@ impl Default for ResourceRequest { n_nodes: 0, resources: smallvec![ResourceRequestEntry { resource: CPU_RESOURCE_NAME.to_string(), - policy: AllocationRequest::Compact(1), + policy: AllocationRequest::Compact(ResourceAmount::new_units(1)), }], min_time: Default::default(), } diff --git a/crates/tako/src/internal/common/resources/allocation.rs b/crates/tako/src/internal/common/resources/allocation.rs index 91c4a42d0..d14cf423c 100644 --- a/crates/tako/src/internal/common/resources/allocation.rs +++ b/crates/tako/src/internal/common/resources/allocation.rs @@ -1,121 +1,96 @@ -use crate::format_comma_delimited; use crate::internal::common::resources::{ResourceAmount, ResourceId, ResourceIndex}; -use crate::internal::worker::resources::counts::ResourceCountVec; +use crate::resources::ResourceFractions; use smallvec::SmallVec; #[derive(Debug)] -pub enum AllocationValue { - Indices(SmallVec<[ResourceIndex; 2]>), - Sum(ResourceAmount), -} - -impl AllocationValue { - pub fn new_indices(mut indices: SmallVec<[ResourceIndex; 2]>) -> Self { - indices.sort_unstable(); - AllocationValue::Indices(indices) - } - - pub fn new_sum(size: ResourceAmount) -> Self { - AllocationValue::Sum(size) - } - - pub fn indices(&self) -> Option<&[ResourceIndex]> { - match self { - AllocationValue::Indices(indices) => Some(indices), - AllocationValue::Sum(_) => None, - } - } - - pub fn to_comma_delimited_list(&self) -> Option { - self.indices().map(format_comma_delimited) - } - - pub fn amount(&self) -> ResourceAmount { - match self { - AllocationValue::Indices(indices) => indices.len() as ResourceAmount, - AllocationValue::Sum(amount) => *amount, - } - } +pub struct AllocationIndex { + pub index: ResourceIndex, + pub group_idx: u32, + pub fractions: ResourceFractions, } #[derive(Debug)] pub struct ResourceAllocation { - pub resource: ResourceId, - pub value: AllocationValue, + pub resource_id: ResourceId, + pub amount: ResourceAmount, + // INVARIANT: indices are sorted by .fractions, i.e. non-whole allocations are at the end + pub indices: SmallVec<[AllocationIndex; 1]>, } -pub type ResourceAllocations = SmallVec<[ResourceAllocation; 2]>; +impl ResourceAllocation { + pub fn resource_indices(&self) -> impl Iterator + '_ { + self.indices.iter().map(|x| x.index) + } +} + +pub type ResourceAllocations = Vec; #[derive(Debug)] pub struct Allocation { pub nodes: Vec, pub resources: ResourceAllocations, - pub counts: ResourceCountVec, // Store counts profile - it is used when Allocation is freed, to find the right profile +} + +impl Default for Allocation { + fn default() -> Self { + Self::new() + } } impl Allocation { - #[inline] - pub fn new( - nodes: Vec, - resources: ResourceAllocations, - counts: ResourceCountVec, - ) -> Self { + pub fn new() -> Self { Allocation { - nodes, - resources, - counts, + nodes: Vec::new(), + resources: Vec::new(), } } + pub fn add_resource_allocation(&mut self, ra: ResourceAllocation) { + self.resources.push(ra); + } + pub fn resource_allocation(&self, id: ResourceId) -> Option<&ResourceAllocation> { - self.resources.iter().find(|r| r.resource == id) + self.resources.iter().find(|r| r.resource_id == id) } } #[cfg(test)] mod tests { - use crate::internal::common::resources::AllocationValue; - use crate::resources::{Allocation, ResourceAmount, ResourceIndex}; - use crate::Set; + use crate::internal::common::resources::allocation::AllocationIndex; + use crate::internal::common::resources::ResourceId; + use crate::resources::{ + Allocation, ResourceAllocation, ResourceAmount, ResourceIndex, ResourceUnits, + }; impl Allocation { - pub fn get_indices(&self, idx: u32) -> Vec { - let a = self - .resources - .iter() - .find(|r| r.resource == idx.into()) - .unwrap(); - match &a.value { - AllocationValue::Indices(x) => x.to_vec(), - AllocationValue::Sum(_) => panic!("Sum not indices"), + pub fn new_simple(counts: &[ResourceUnits]) -> Self { + Allocation { + nodes: Vec::new(), + resources: counts + .iter() + .enumerate() + .filter(|(_id, c)| **c > 0) + .map(|(id, c)| ResourceAllocation { + resource_id: ResourceId::new(id as u32), + amount: ResourceAmount::new_units(*c), + indices: (0..*c) + .map(|x| AllocationIndex { + index: ResourceIndex::new(x), + group_idx: 0, + fractions: 0, + }) + .collect(), + }) + .collect(), } } - - pub fn get_sum(&self, idx: u32) -> ResourceAmount { + pub fn get_indices(&self, idx: u32) -> Vec { let a = self .resources .iter() - .find(|r| r.resource == idx.into()) + .find(|r| r.resource_id == idx.into()) .unwrap(); - match &a.value { - AllocationValue::Indices(_) => panic!("Indices not sum"), - AllocationValue::Sum(s) => *s, - } - } - } - - impl AllocationValue { - pub fn get_checked_indices(&self) -> Vec { - match self { - AllocationValue::Indices(indices) => { - let v: Vec = indices.iter().map(|x| x.as_num()).collect(); - assert_eq!(v.iter().collect::>().len(), v.len()); - v - } - AllocationValue::Sum(_) => { - panic!("Not indices") - } - } + a.indices.iter().map(|a| a.index).collect() } } } diff --git a/crates/tako/src/internal/common/resources/amount.rs b/crates/tako/src/internal/common/resources/amount.rs new file mode 100644 index 000000000..ffd03f15d --- /dev/null +++ b/crates/tako/src/internal/common/resources/amount.rs @@ -0,0 +1,135 @@ +use derive_more::{Add, AddAssign, Sub, SubAssign, Sum}; +use serde::{Deserialize, Serialize}; + +pub type ResourceUnits = u32; +pub type ResourceFractions = u32; + +pub const FRACTIONS_PER_UNIT: ResourceFractions = 10_000; +pub const FRACTIONS_MAX_DIGITS: usize = FRACTIONS_PER_UNIT.ilog10() as usize; + +#[derive( + Debug, + Serialize, + Clone, + Copy, + Hash, + Eq, + Deserialize, + PartialEq, + Ord, + PartialOrd, + AddAssign, + SubAssign, + Sub, + Add, + Sum, +)] +pub struct ResourceAmount(u64); + +impl ResourceAmount { + pub const ZERO: ResourceAmount = ResourceAmount(0); + + pub fn new(units: ResourceUnits, fractions: ResourceFractions) -> Self { + assert!(fractions < FRACTIONS_PER_UNIT); + ResourceAmount(units as u64 * FRACTIONS_PER_UNIT as u64 + fractions as u64) + } + + pub fn new_units(units: ResourceUnits) -> Self { + ResourceAmount(units as u64 * FRACTIONS_PER_UNIT as u64) + } + + pub fn from_float(value: f32) -> Self { + ResourceAmount((value * FRACTIONS_PER_UNIT as f32).ceil() as u64) + } + + pub fn new_fractions(fractions: ResourceFractions) -> Self { + assert!(fractions < FRACTIONS_PER_UNIT); + ResourceAmount(fractions as u64) + } + + pub fn is_zero(&self) -> bool { + self.0 == 0 + } + + pub fn units(&self) -> ResourceUnits { + (self.0 / (FRACTIONS_PER_UNIT as u64)) as ResourceUnits + } + + pub fn fractions(&self) -> ResourceFractions { + (self.0 % (FRACTIONS_PER_UNIT as u64)) as ResourceFractions + } + + pub fn split(&self) -> (ResourceUnits, ResourceFractions) { + (self.units(), self.fractions()) + } + + pub fn total_fractions(&self) -> u64 { + self.0 + } + + pub fn as_f32(&self) -> f32 { + self.0 as f32 / FRACTIONS_PER_UNIT as f32 + } +} + +impl std::fmt::Display for ResourceAmount { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + let fractions = self.fractions(); + write!(f, "{}", self.units())?; + if fractions != 0 { + let num = format!("{:01$}", fractions, 4); + write!(f, ".{}", num.trim_end_matches('0'))?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + impl ResourceAmount { + pub fn assert_eq_units(&self, units: ResourceUnits) { + assert_eq!(self.units(), units); + assert_eq!(self.fractions(), 0); + } + } + + #[test] + pub fn test_amount_add() { + let r1 = ResourceAmount::new(10, 1234); + let r2 = ResourceAmount::new(2, 4321); + let r3 = ResourceAmount::new(7, 9999); + assert_eq!(r1 + r2, ResourceAmount::new(12, 5555)); + assert_eq!(r1 + r3, ResourceAmount::new(18, 1233)); + assert_eq!(r1 + ResourceAmount::ZERO, r1); + } + + #[test] + pub fn test_amount_from_float() { + assert_eq!( + ResourceAmount::from_float(2.0f32), + ResourceAmount::new_units(2) + ); + assert_eq!( + ResourceAmount::from_float(1.1f32), + ResourceAmount::new(1, 1000) + ); + assert_eq!( + ResourceAmount::from_float(20.0001f32), + ResourceAmount::new(20, 1) + ); + } + + #[test] + pub fn test_amount_display() { + assert_eq!(ResourceAmount::new(0, 0).to_string(), "0"); + assert_eq!(ResourceAmount::new(0, 1).to_string(), "0.0001"); + assert_eq!(ResourceAmount::new(500, 0).to_string(), "500"); + assert_eq!(ResourceAmount::new(500, 123).to_string(), "500.0123"); + assert_eq!(ResourceAmount::new(500, 9999).to_string(), "500.9999"); + assert_eq!(ResourceAmount::new(1, 1000).to_string(), "1.1"); + assert_eq!(ResourceAmount::new(1, 2200).to_string(), "1.22"); + assert_eq!(ResourceAmount::new(1, 3410).to_string(), "1.341"); + } +} diff --git a/crates/tako/src/internal/common/resources/descriptor.rs b/crates/tako/src/internal/common/resources/descriptor.rs index a86f01839..678e7e144 100644 --- a/crates/tako/src/internal/common/resources/descriptor.rs +++ b/crates/tako/src/internal/common/resources/descriptor.rs @@ -1,4 +1,6 @@ -use crate::internal::common::resources::{ResourceAmount, ResourceIndex, ResourceLabel}; +use crate::internal::common::resources::{ + ResourceAmount, ResourceIndex, ResourceLabel, ResourceUnits, +}; use crate::internal::common::utils::{format_comma_delimited, has_unique_elements}; use crate::internal::common::Set; use serde::{Deserialize, Serialize}; @@ -33,11 +35,11 @@ pub enum ResourceDescriptorKind { } impl ResourceDescriptorKind { - pub fn regular_sockets(n_sockets: ResourceAmount, socket_size: ResourceAmount) -> Self { + pub fn regular_sockets(n_sockets: ResourceUnits, socket_size: ResourceUnits) -> Self { assert!(n_sockets > 0); assert!(socket_size > 0); if n_sockets == 1 { - Self::simple_indices(socket_size as u32) + Self::simple_indices(socket_size) } else { let mut sockets = Vec::with_capacity(n_sockets as usize); let mut i = 0; @@ -106,7 +108,7 @@ impl ResourceDescriptorKind { }) } - pub fn simple_indices(size: u32) -> Self { + pub fn simple_indices(size: ResourceUnits) -> Self { assert!(size > 0); ResourceDescriptorKind::Range { start: ResourceIndex::from(0), @@ -125,14 +127,16 @@ impl ResourceDescriptorKind { pub fn size(&self) -> ResourceAmount { match self { - ResourceDescriptorKind::List { values } => values.len() as ResourceAmount, + ResourceDescriptorKind::List { values } => { + ResourceAmount::new_units(values.len() as ResourceUnits) + } ResourceDescriptorKind::Range { start, end } if end >= start => { - (end.as_num() + 1 - start.as_num()) as ResourceAmount + ResourceAmount::new_units(end.as_num() + 1 - start.as_num()) } - ResourceDescriptorKind::Range { .. } => 0, + ResourceDescriptorKind::Range { .. } => ResourceAmount::ZERO, ResourceDescriptorKind::Sum { size } => *size, ResourceDescriptorKind::Groups { groups } => { - groups.iter().map(|x| x.len() as ResourceAmount).sum() + ResourceAmount::new_units(groups.iter().map(|x| x.len() as ResourceUnits).sum()) } } } @@ -218,7 +222,7 @@ impl ResourceDescriptor { .validate() .map_err(|e| format!("Invalid resource definition for {}: {:?}", item.name, e))?; - if item.kind.size() == 0 { + if item.kind.size().is_zero() { return Err(format!("Resource {} is empty", item.name).into()); } if item.name == "cpus" { @@ -248,20 +252,22 @@ mod tests { } } - pub fn sum(name: &str, size: u64) -> Self { + pub fn sum(name: &str, size: u32) -> Self { ResourceDescriptorItem { name: name.to_string(), - kind: ResourceDescriptorKind::Sum { size: size.into() }, + kind: ResourceDescriptorKind::Sum { + size: ResourceAmount::new_units(size), + }, } } } impl ResourceDescriptor { - pub fn simple(n_cpus: ResourceAmount) -> Self { + pub fn simple(n_cpus: ResourceUnits) -> Self { Self::sockets(1, n_cpus) } - pub fn sockets(n_sockets: ResourceAmount, n_cpus_per_socket: ResourceAmount) -> Self { + pub fn sockets(n_sockets: ResourceUnits, n_cpus_per_socket: ResourceUnits) -> Self { ResourceDescriptor::new(vec![ResourceDescriptorItem { name: CPU_RESOURCE_NAME.to_string(), kind: ResourceDescriptorKind::regular_sockets(n_sockets, n_cpus_per_socket), diff --git a/crates/tako/src/internal/common/resources/mod.rs b/crates/tako/src/internal/common/resources/mod.rs index acca78c9e..79cb0c715 100644 --- a/crates/tako/src/internal/common/resources/mod.rs +++ b/crates/tako/src/internal/common/resources/mod.rs @@ -1,11 +1,12 @@ pub mod allocation; +pub mod amount; pub mod descriptor; pub mod map; pub mod request; use crate::define_id_type; use crate::internal::common::index::IndexVec; -pub use allocation::{Allocation, AllocationValue, ResourceAllocation, ResourceAllocations}; +pub use allocation::{Allocation, ResourceAllocation, ResourceAllocations}; pub use descriptor::{ DescriptorError, ResourceDescriptor, ResourceDescriptorItem, ResourceDescriptorKind, }; @@ -18,14 +19,13 @@ pub use request::{ ResourceRequestVariants, TimeRequest, }; +pub use amount::{ResourceAmount, ResourceFractions, ResourceUnits}; + pub type NumOfNodes = u32; // Identifies a globally unique Resource request stored in Core. define_id_type!(ResourceId, u32); -/// Represents some amount within a single generic resource (e.g. 100 MiB of memory). -pub type ResourceAmount = u64; - // Represents an index within a single generic resource (e.g. GPU with ID 1). define_id_type!(ResourceIndex, u32); diff --git a/crates/tako/src/internal/common/resources/request.rs b/crates/tako/src/internal/common/resources/request.rs index a6843589c..faf22395b 100644 --- a/crates/tako/src/internal/common/resources/request.rs +++ b/crates/tako/src/internal/common/resources/request.rs @@ -19,17 +19,27 @@ pub enum AllocationRequest { impl AllocationRequest { pub fn validate(&self) -> crate::Result<()> { + let check_nonzero = |amount: &ResourceAmount| { + if amount.is_zero() { + Err(DsError::GenericError( + "Zero resources cannot be requested".to_string(), + )) + } else { + Ok(()) + } + }; match &self { - AllocationRequest::Scatter(n_cpus) - | AllocationRequest::ForceCompact(n_cpus) - | AllocationRequest::Compact(n_cpus) => { - if *n_cpus == 0 { + AllocationRequest::ForceCompact(amount) => check_nonzero(amount).and_then(|_| { + if amount.fractions() > 0 { Err(DsError::GenericError( - "Zero resources cannot be requested".to_string(), + "ForceCompact have to use whole resource counts".to_string(), )) } else { Ok(()) } + }), + AllocationRequest::Scatter(amount) | AllocationRequest::Compact(amount) => { + check_nonzero(amount) } AllocationRequest::All => Ok(()), } @@ -40,7 +50,7 @@ impl AllocationRequest { AllocationRequest::Compact(amount) | AllocationRequest::ForceCompact(amount) | AllocationRequest::Scatter(amount) => *amount, - AllocationRequest::All => 1, + AllocationRequest::All => ResourceAmount::new_units(1), } } diff --git a/crates/tako/src/internal/messages/worker.rs b/crates/tako/src/internal/messages/worker.rs index 16727675c..fc4285d7a 100644 --- a/crates/tako/src/internal/messages/worker.rs +++ b/crates/tako/src/internal/messages/worker.rs @@ -4,6 +4,7 @@ use std::time::Duration; use crate::hwstats::WorkerHwStateMessage; use crate::internal::common::resources::{ResourceAmount, ResourceIndex}; use crate::internal::messages::common::TaskFailInfo; +use crate::resources::ResourceFractions; use crate::task::SerializedTaskContext; use crate::{InstanceId, Priority}; use crate::{TaskId, WorkerId}; @@ -106,18 +107,12 @@ pub struct StealResponseMsg { pub responses: Vec<(TaskId, StealResponse)>, } -#[derive(Serialize, Deserialize, Debug, Clone)] -#[cfg_attr(test, derive(Eq, PartialEq))] -pub enum TaskResourceAllocationValue { - Indices(Vec), - Sum(ResourceAmount), -} - #[derive(Serialize, Deserialize, Debug, Clone)] #[cfg_attr(test, derive(Eq, PartialEq))] pub struct ResourceAllocation { pub resource: String, - pub value: TaskResourceAllocationValue, + pub indices: Vec<(ResourceIndex, ResourceFractions)>, + pub amount: ResourceAmount, } #[derive(Serialize, Deserialize, Debug, Clone)] diff --git a/crates/tako/src/internal/server/core.rs b/crates/tako/src/internal/server/core.rs index b5cb1efb0..007b76ae2 100644 --- a/crates/tako/src/internal/server/core.rs +++ b/crates/tako/src/internal/server/core.rs @@ -517,8 +517,8 @@ mod tests { op: F, ) { for task_id in task_ids { - let task_id: TaskId = task_id.clone().into(); - if !op(&self.get_task(task_id)) { + let task_id: TaskId = (*task_id).into(); + if !op(self.get_task(task_id)) { panic!("Task {} does not satisfy the condition", task_id); } } @@ -530,8 +530,8 @@ mod tests { op: F, ) { for worker_id in worker_ids { - let worker_id: WorkerId = worker_id.clone().into(); - if !op(&self.get_worker_by_id_or_panic(worker_id)) { + let worker_id: WorkerId = (*worker_id).into(); + if !op(self.get_worker_by_id_or_panic(worker_id)) { panic!("Worker {} does not satisfy the condition", worker_id); } } diff --git a/crates/tako/src/internal/server/reactor.rs b/crates/tako/src/internal/server/reactor.rs index d35bf9d48..9bfc0ee65 100644 --- a/crates/tako/src/internal/server/reactor.rs +++ b/crates/tako/src/internal/server/reactor.rs @@ -273,8 +273,6 @@ pub(crate) fn on_task_finished( if let Some(task) = tasks.find_task_mut(msg.id) { log::debug!("Task id={} finished on worker={}", task.id, worker_id); - if task.configuration.resources.is_multi_node() {} - assert!(task.is_assigned_or_stolen_from(worker_id)); match &task.state { diff --git a/crates/tako/src/internal/server/workerload.rs b/crates/tako/src/internal/server/workerload.rs index 8f0f73b41..2b332a771 100644 --- a/crates/tako/src/internal/server/workerload.rs +++ b/crates/tako/src/internal/server/workerload.rs @@ -25,7 +25,10 @@ impl WorkerResources { } pub(crate) fn get(&self, resource_id: ResourceId) -> ResourceAmount { - self.n_resources.get(resource_id).copied().unwrap_or(0) + self.n_resources + .get(resource_id) + .copied() + .unwrap_or(ResourceAmount::ZERO) } pub(crate) fn from_description( @@ -42,7 +45,8 @@ impl WorkerResources { .max() .unwrap_or(0); - let mut n_resources: ResourceVec = IndexVec::filled(0, resource_count); + let mut n_resources: ResourceVec = + IndexVec::filled(ResourceAmount::ZERO, resource_count); for descriptor in &resource_desc.resources { let position = resource_map.get_index(&descriptor.name).unwrap(); @@ -61,13 +65,9 @@ impl WorkerResources { } pub(crate) fn is_capable_to_run(&self, rqv: &ResourceRequestVariants) -> bool { - rqv.requests().iter().any(|rq| { - rq.entries().iter().all(|r| { - let ask = r.request.min_amount(); - let has = self.get(r.resource_id); - ask <= has - }) - }) + rqv.requests() + .iter() + .any(|rq| self.is_capable_to_run_request(rq)) } pub(crate) fn to_transport(&self) -> WorkerResourceCounts { @@ -85,23 +85,23 @@ impl WorkerResources { } } - pub fn difficulty_score(&self, request: &ResourceRequest) -> u32 { + pub fn difficulty_score(&self, request: &ResourceRequest) -> u64 { let mut result = 0; for entry in request.entries() { let count = self .n_resources .get(entry.resource_id) .copied() - .unwrap_or(0); - if count == 0 { + .unwrap_or(ResourceAmount::ZERO); + if count.is_zero() { return 0; } - result += ((entry.request.amount(count) * 512) / (count * 512)) as u32; + result += entry.request.amount(count).total_fractions() / count.total_fractions(); } result } - pub fn difficulty_score_of_rqv(&self, rqv: &ResourceRequestVariants) -> u32 { + pub fn difficulty_score_of_rqv(&self, rqv: &ResourceRequestVariants) -> u64 { rqv.requests() .iter() .map(|r| self.difficulty_score(r)) @@ -128,7 +128,7 @@ pub struct WorkerLoad { impl WorkerLoad { pub(crate) fn new(worker_resources: &WorkerResources) -> WorkerLoad { WorkerLoad { - n_resources: IndexVec::filled(0, worker_resources.n_resources.len()), + n_resources: IndexVec::filled(ResourceAmount::ZERO, worker_resources.n_resources.len()), non_first_rq: Default::default(), round_robin_counter: 0, } @@ -199,7 +199,10 @@ impl WorkerLoad { } pub(crate) fn get(&self, resource_id: ResourceId) -> ResourceAmount { - self.n_resources.get(resource_id).copied().unwrap_or(0) + self.n_resources + .get(resource_id) + .copied() + .unwrap_or(ResourceAmount::ZERO) } pub(crate) fn have_immediate_resources_for_rq( @@ -246,16 +249,20 @@ impl WorkerLoad { pub(crate) fn load_wrt_request(&self, wr: &WorkerResources, request: &ResourceRequest) -> u32 { let mut result = 0; for entry in request.entries() { - let count = wr.n_resources.get(entry.resource_id).copied().unwrap_or(0); - if count == 0 { + let count = wr + .n_resources + .get(entry.resource_id) + .copied() + .unwrap_or(ResourceAmount::ZERO); + if count.is_zero() { return 0; } let load = self .n_resources .get(entry.resource_id) .copied() - .unwrap_or(0); - result += ((load * 512) / (count * 512)) as u32; + .unwrap_or(ResourceAmount::ZERO); + result += (load.total_fractions() * 512 / count.total_fractions()) as u32; } result } @@ -294,18 +301,23 @@ mod tests { use crate::internal::server::workerload::{ ResourceRequestLowerBound, WorkerLoad, WorkerResources, }; - use crate::internal::tests::utils::resources::{cpus_compact, ResBuilder}; + use crate::internal::tests::utils::resources::{cpus_compact, ra_builder, ResBuilder}; + use crate::resources::{ResourceAmount, ResourceUnits}; use crate::TaskId; use smallvec::smallvec; + pub fn wr_builder(units: &[ResourceUnits]) -> WorkerResources { + WorkerResources { + n_resources: ra_builder(units), + } + } + #[test] fn worker_load_check_lb() { - let wr = WorkerResources { - n_resources: vec![2, 10, 100, 5].into(), - }; + let wr = wr_builder(&[2, 10, 100, 5]); let load = WorkerLoad::new(&wr); let load2 = WorkerLoad { - n_resources: vec![0, 9, 0, 0, 0, 0].into(), + n_resources: ra_builder(&[0, 9, 0, 0, 0, 0]), non_first_rq: Default::default(), round_robin_counter: 0, }; @@ -347,49 +359,58 @@ mod tests { #[test] fn worker_load_variants() { - let wr = WorkerResources { - n_resources: vec![13, 4, 5].into(), - }; + let wr = wr_builder(&[13, 4, 5]); let mut load = WorkerLoad::new(&wr); let rq1 = ResBuilder::default().add(0, 2).add(1, 2).finish(); let rq2 = ResBuilder::default().add(0, 4).finish(); let rqv = ResourceRequestVariants::new(smallvec![rq1, rq2]); load.add_request(TaskId::new(1), &rqv, &wr); - assert_eq!(load.n_resources, vec![2, 2, 0].into()); + assert_eq!(load.n_resources, ra_builder(&[2, 2, 0])); assert!(load.non_first_rq.is_empty()); load.add_request(TaskId::new(2), &rqv, &wr); - assert_eq!(load.n_resources, vec![4, 4, 0].into()); + assert_eq!(load.n_resources, ra_builder(&[4, 4, 0])); assert!(load.non_first_rq.is_empty()); load.add_request(TaskId::new(3), &rqv, &wr); - assert_eq!(load.n_resources, vec![8, 4, 0].into()); + assert_eq!(load.n_resources, ra_builder(&[8, 4, 0])); assert_eq!(load.non_first_rq.len(), 1); assert_eq!(load.non_first_rq.get(&TaskId::new(3)), Some(&1)); load.add_request(TaskId::new(4), &rqv, &wr); - assert_eq!(load.n_resources, vec![12, 4, 0].into()); + assert_eq!(load.n_resources, ra_builder(&[12, 4, 0])); assert_eq!(load.non_first_rq.len(), 2); assert_eq!(load.non_first_rq.get(&TaskId::new(4)), Some(&1)); load.add_request(TaskId::new(5), &rqv, &wr); assert!( - load.n_resources == vec![16, 4, 0].into() || load.n_resources == vec![14, 6, 0].into() + load.n_resources == ra_builder(&[16, 4, 0]) + || load.n_resources == ra_builder(&[14, 6, 0]) ); let resources = load.n_resources.clone(); load.remove_request(TaskId::new(3), &rqv, &wr); assert_eq!( load.n_resources, - vec![resources[0.into()] - 4, resources[1.into()], 0].into() + vec![ + resources[0.into()] - ResourceAmount::new_units(4), + resources[1.into()], + ResourceAmount::ZERO, + ] + .into() ); assert!(load.non_first_rq.get(&TaskId::new(3)).is_none()); load.remove_request(TaskId::new(1), &rqv, &wr); assert_eq!( load.n_resources, - vec![resources[0.into()] - 6, resources[1.into()] - 2, 0].into() + vec![ + resources[0.into()] - ResourceAmount::new_units(6), + resources[1.into()] - ResourceAmount::new_units(2), + ResourceAmount::ZERO, + ] + .into() ); assert!(load.non_first_rq.get(&TaskId::new(1)).is_none()); } diff --git a/crates/tako/src/internal/tests/integration/test_resources.rs b/crates/tako/src/internal/tests/integration/test_resources.rs index 1301f8837..07f22c001 100644 --- a/crates/tako/src/internal/tests/integration/test_resources.rs +++ b/crates/tako/src/internal/tests/integration/test_resources.rs @@ -90,10 +90,7 @@ async fn test_submit_2_sleeps_on_separated_2() { ) .await; - let workers = handler - .start_workers(|| Default::default(), 3) - .await - .unwrap(); + let workers = handler.start_workers(Default::default, 3).await.unwrap(); let worker_ids: Vec = workers.iter().map(|x| x.id).collect(); wait_for_task_start(&mut handler, 1).await; diff --git a/crates/tako/src/internal/tests/integration/test_worker.rs b/crates/tako/src/internal/tests/integration/test_worker.rs index 4ca7ae03d..1775a75a4 100644 --- a/crates/tako/src/internal/tests/integration/test_worker.rs +++ b/crates/tako/src/internal/tests/integration/test_worker.rs @@ -144,10 +144,7 @@ async fn test_panic_on_worker_lost() { #[tokio::test] async fn test_lost_worker_with_tasks_continue() { run_test(Default::default(), |mut handler| async move { - let _workers = handler - .start_workers(|| Default::default(), 2) - .await - .unwrap(); + let _workers = handler.start_workers(Default::default, 2).await.unwrap(); let task_ids = handler .submit(GraphBuilder::singleton(simple_task(&["sleep", "1"], 1))) diff --git a/crates/tako/src/internal/tests/integration/utils/task.rs b/crates/tako/src/internal/tests/integration/utils/task.rs index 7fe97dedb..eb655b247 100644 --- a/crates/tako/src/internal/tests/integration/utils/task.rs +++ b/crates/tako/src/internal/tests/integration/utils/task.rs @@ -168,7 +168,7 @@ pub struct ResourceRequestConfig { } impl ResourceRequestConfigBuilder { - pub fn cpus(self, n_cpus: ResourceAmount) -> Self { + pub fn cpus>(self, n_cpus: A) -> Self { self.add_compact("cpus", n_cpus) } @@ -183,13 +183,13 @@ impl ResourceRequestConfigBuilder { }) } - pub fn add_compact(mut self, name: &str, amount: ResourceAmount) -> Self { - self._add(name, AllocationRequest::Compact(amount)); + pub fn add_compact>(mut self, name: &str, amount: A) -> Self { + self._add(name, AllocationRequest::Compact(amount.into())); self } - pub fn add_force_compact(mut self, name: &str, amount: ResourceAmount) -> Self { - self._add(name, AllocationRequest::ForceCompact(amount)); + pub fn add_force_compact>(mut self, name: &str, amount: A) -> Self { + self._add(name, AllocationRequest::ForceCompact(amount.into())); self } } diff --git a/crates/tako/src/internal/tests/integration/utils/worker.rs b/crates/tako/src/internal/tests/integration/utils/worker.rs index 7135d8db9..fbdc3ff19 100644 --- a/crates/tako/src/internal/tests/integration/utils/worker.rs +++ b/crates/tako/src/internal/tests/integration/utils/worker.rs @@ -273,7 +273,7 @@ impl TaskLauncher for TestTaskLauncher { ctx.allocation(), ctx.body().len(), ); - rmp_serde::from_slice(&ctx.body())? + rmp_serde::from_slice(ctx.body())? }; Ok(TaskLaunchData::from_future(Box::pin(async move { diff --git a/crates/tako/src/internal/tests/test_query.rs b/crates/tako/src/internal/tests/test_query.rs index e1e6af8a5..3241ee9d7 100644 --- a/crates/tako/src/internal/tests/test_query.rs +++ b/crates/tako/src/internal/tests/test_query.rs @@ -95,7 +95,7 @@ fn test_query_enough_workers2() { rt.schedule(); let r = compute_new_worker_query( - &rt.core(), + rt.core(), &[ WorkerTypeQuery { descriptor: ResourceDescriptor::simple(2), @@ -125,7 +125,7 @@ fn test_query_not_enough_workers3() { rt.schedule(); let r = compute_new_worker_query( - &rt.core(), + rt.core(), &[ WorkerTypeQuery { descriptor: ResourceDescriptor::simple(2), @@ -155,7 +155,7 @@ fn test_query_many_workers_needed() { rt.schedule(); let r = compute_new_worker_query( - &rt.core(), + rt.core(), &[ WorkerTypeQuery { descriptor: ResourceDescriptor::simple(2), @@ -201,7 +201,7 @@ fn test_query_multi_node_tasks() { rt.schedule(); let r = compute_new_worker_query( - &rt.core(), + rt.core(), &[ WorkerTypeQuery { descriptor: ResourceDescriptor::simple(1), diff --git a/crates/tako/src/internal/tests/test_reactor.rs b/crates/tako/src/internal/tests/test_reactor.rs index 22af47de8..60046a105 100644 --- a/crates/tako/src/internal/tests/test_reactor.rs +++ b/crates/tako/src/internal/tests/test_reactor.rs @@ -30,7 +30,7 @@ use crate::internal::tests::utils::task::{task, task_running_msg, task_with_deps use crate::internal::tests::utils::workflows::{submit_example_1, submit_example_3}; use crate::internal::tests::utils::{env, schedule}; use crate::internal::worker::configuration::OverviewConfiguration; -use crate::resources::{ResourceDescriptorItem, ResourceMap}; +use crate::resources::{ResourceAmount, ResourceDescriptorItem, ResourceMap}; use crate::worker::{ServerLostPolicy, WorkerConfiguration}; use crate::{TaskId, WorkerId}; @@ -71,12 +71,15 @@ fn test_worker_add() { assert_eq!(new_w.len(), 1); assert_eq!(new_w[0].0.as_num(), 402); - assert_eq!(new_w[0].1.resources.resources[0].kind.size(), 4); + assert_eq!( + new_w[0].1.resources.resources[0].kind.size(), + ResourceAmount::new_units(4) + ); assert!( matches!(comm.take_broadcasts(1)[0], ToWorkerMessage::NewWorker(NewWorkerMsg { worker_id: WorkerId(402), address: ref _a, resources: ref r, - }) if r.n_resources == vec![4]) + }) if r.n_resources == vec![ResourceAmount::new_units(4)]) ); comm.check_need_scheduling(); @@ -129,7 +132,7 @@ fn test_worker_add() { assert!( matches!(comm.take_broadcasts(1)[0], ToWorkerMessage::NewWorker(NewWorkerMsg { worker_id: WorkerId(502), address: ref a, resources: ref r, - }) if a == "test2:123" && r.n_resources == vec![5, 0, 100_000_000]) + }) if a == "test2:123" && r.n_resources == vec![ResourceAmount::new_units(5), ResourceAmount::new_units(0), ResourceAmount::new_units(100_000_000)]) ); comm.check_need_scheduling(); comm.emptiness_check(); @@ -154,12 +157,12 @@ fn test_submit_jobs() { assert_eq!(t1.get_unfinished_deps(), 0); assert_eq!(t2.get_unfinished_deps(), 1); - check_task_consumers_exact(&t1, &[t2]); + check_task_consumers_exact(t1, &[t2]); let t3 = task(604); - let t4 = task_with_deps(602, &[&t1, &t3], 1); + let t4 = task_with_deps(602, &[t1, &t3], 1); let t5 = task_with_deps(603, &[&t3], 1); - let t6 = task_with_deps(601, &[&t3, &t4, &t5, &t2], 1); + let t6 = task_with_deps(601, &[&t3, &t4, &t5, t2], 1); on_new_tasks(&mut core, &mut comm, vec![t6, t3, t4, t5]); comm.check_need_scheduling(); @@ -170,10 +173,10 @@ fn test_submit_jobs() { let t4 = core.get_task(602.into()); let t6 = core.get_task(601.into()); - check_task_consumers_exact(&t1, &[t2, t4]); + check_task_consumers_exact(t1, &[t2, t4]); assert_eq!(t1.get_unfinished_deps(), 0); - check_task_consumers_exact(&t2, &[t6]); + check_task_consumers_exact(t2, &[t6]); assert_eq!(t1.get_unfinished_deps(), 0); assert_eq!(t2.get_unfinished_deps(), 1); diff --git a/crates/tako/src/internal/tests/test_scheduler_mn.rs b/crates/tako/src/internal/tests/test_scheduler_mn.rs index 218a30363..f89442f48 100644 --- a/crates/tako/src/internal/tests/test_scheduler_mn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_mn.rs @@ -100,9 +100,9 @@ fn test_schedule_mn_simple() { }; let task3 = core.get_task(3.into()); - let ws3 = test_mn_task(&task3, &mut comm, true); + let ws3 = test_mn_task(task3, &mut comm, true); let task4 = core.get_task(4.into()); - let ws4 = test_mn_task(&task4, &mut comm, true); + let ws4 = test_mn_task(task4, &mut comm, true); for w in &ws4 { assert!(!ws3.contains(w)); } @@ -122,7 +122,7 @@ fn test_schedule_mn_simple() { core.sanity_check(); let task2 = core.get_task(2.into()); - let ws2 = test_mn_task(&task2, &mut comm, false); + let ws2 = test_mn_task(task2, &mut comm, false); comm.emptiness_check(); finish_on_worker(&mut core, 3, ws2[0], 0); diff --git a/crates/tako/src/internal/tests/test_scheduler_sn.rs b/crates/tako/src/internal/tests/test_scheduler_sn.rs index 8437f4957..0538ead2b 100644 --- a/crates/tako/src/internal/tests/test_scheduler_sn.rs +++ b/crates/tako/src/internal/tests/test_scheduler_sn.rs @@ -14,7 +14,7 @@ use crate::internal::tests::utils::schedule::{ use crate::internal::tests::utils::task::task; use crate::internal::tests::utils::task::TaskBuilder; use crate::internal::tests::utils::workflows::submit_example_1; -use crate::resources::{ResourceAmount, ResourceDescriptorItem}; +use crate::resources::{ResourceAmount, ResourceDescriptorItem, ResourceUnits}; use crate::{TaskId, WorkerId}; use std::time::Duration; @@ -23,7 +23,7 @@ fn test_no_deps_scattering_1() { let mut core = Core::default(); create_test_workers(&mut core, &[5, 5, 5]); - let tasks: Vec = (1..=4).map(|i| task(i)).collect(); + let tasks: Vec = (1..=4).map(task).collect(); submit_test_tasks(&mut core, tasks); let mut scheduler = create_test_scheduler(); @@ -38,9 +38,9 @@ fn test_no_deps_scattering_1() { assert_eq!(m1.len() + m2.len() + m3.len(), 4); assert!( - (m1.len() == 4 && m2.len() == 0 && m3.len() == 0) - || (m1.len() == 0 && m2.len() == 4 && m3.len() == 0) - || (m1.len() == 0 && m2.len() == 0 && m3.len() == 4) + (m1.len() == 4 && m2.is_empty() && m3.is_empty()) + || (m1.is_empty() && m2.len() == 4 && m3.is_empty()) + || (m1.is_empty() && m2.is_empty() && m3.len() == 4) ); } @@ -85,7 +85,7 @@ fn test_no_deps_distribute_without_balance() { let mut core = Core::default(); create_test_workers(&mut core, &[10, 10, 10]); - let tasks: Vec = (1..=150).map(|i| task(i)).collect(); + let tasks: Vec = (1..=150).map(task).collect(); submit_test_tasks(&mut core, tasks); let mut scheduler = create_test_scheduler(); @@ -115,8 +115,8 @@ fn test_no_deps_distribute_with_balance() { assert!(w.is_underloaded()); } - let mut active_ids: Set = (1..301).into_iter().map(|id| id.into()).collect(); - let tasks: Vec = (1..301).map(|i| task(i)).collect(); + let mut active_ids: Set = (1..301).map(|id| id.into()).collect(); + let tasks: Vec = (1..301).map(task).collect(); submit_test_tasks(&mut core, tasks); let mut scheduler = create_test_scheduler(); @@ -323,6 +323,7 @@ fn test_resource_balancing1() { #[test] fn test_resource_balancing2() { + let u = ResourceAmount::new_units; let mut rt = TestEnv::new(); rt.new_workers(&[12, 12, 12]); @@ -330,13 +331,14 @@ fn test_resource_balancing2() { 4, 4, 4, /* 12 */ 4, 4, 4, /* 12 */ 4, 4, 4, /* 12 */ ]]); rt.balance(); - assert_eq!(rt.worker_load(100).get(0.into()), 12); - assert_eq!(rt.worker_load(101).get(0.into()), 12); - assert_eq!(rt.worker_load(102).get(0.into()), 12); + assert_eq!(rt.worker_load(100).get(0.into()), u(12)); + assert_eq!(rt.worker_load(101).get(0.into()), u(12)); + assert_eq!(rt.worker_load(102).get(0.into()), u(12)); } #[test] fn test_resource_balancing3() { + let u = ResourceAmount::new_units; let mut rt = TestEnv::new(); rt.new_workers(&[12, 12, 12]); @@ -348,13 +350,14 @@ fn test_resource_balancing3() { &[6], // 6 ]); rt.balance(); - assert!(rt.worker_load(100).get(0.into()) >= 12); - assert!(rt.worker_load(101).get(0.into()) >= 12); - assert!(rt.worker_load(102).get(0.into()) >= 12); + assert!(rt.worker_load(100).get(0.into()) >= u(12)); + assert!(rt.worker_load(101).get(0.into()) >= u(12)); + assert!(rt.worker_load(102).get(0.into()) >= u(12)); } #[test] fn test_resource_balancing4() { + let u = ResourceAmount::new_units; let mut rt = TestEnv::new(); rt.new_workers(&[12, 12, 12]); @@ -362,13 +365,14 @@ fn test_resource_balancing4() { 2, 4, 2, 4, /* 12 */ 4, 4, 4, /* 12 */ 2, 2, 2, 2, 4, /* 12 */ ]]); rt.balance(); - assert_eq!(rt.worker_load(100).get(0.into()), 12); - assert_eq!(rt.worker_load(101).get(0.into()), 12); - assert_eq!(rt.worker_load(102).get(0.into()), 12); + assert_eq!(rt.worker_load(100).get(0.into()), u(12)); + assert_eq!(rt.worker_load(101).get(0.into()), u(12)); + assert_eq!(rt.worker_load(102).get(0.into()), u(12)); } #[test] fn test_resource_balancing5() { + let u = ResourceAmount::new_units; let mut rt = TestEnv::new(); rt.new_workers(&[12, 12, 12]); @@ -376,11 +380,12 @@ fn test_resource_balancing5() { 2, 4, 2, 4, 2, /* 14 */ 4, 4, 4, /* 12 */ 4, 4, 4, /* 12 */ ]]); rt.balance(); - rt.check_worker_load_lower_bounds(&[10, 12, 12]); + rt.check_worker_load_lower_bounds(&[u(10), u(12), u(12)]); } #[test] fn test_resource_balancing6() { + let u = ResourceAmount::new_units; let mut rt = TestEnv::new(); rt.new_workers(&[12, 12, 12, 12, 12]); rt.new_assigned_tasks_cpus(&[ @@ -391,11 +396,12 @@ fn test_resource_balancing6() { &[12], ]); rt.balance(); - rt.check_worker_load_lower_bounds(&[12, 12, 12]); + rt.check_worker_load_lower_bounds(&[u(12), u(12), u(12)]); } #[test] fn test_resource_balancing7() { + let u = ResourceAmount::new_units; let mut rt = TestEnv::new(); rt.new_workers(&[12, 12, 12]); @@ -414,19 +420,20 @@ fn test_resource_balancing7() { rt.new_task_running(TaskBuilder::new(33).cpus_compact(4), 102); rt.schedule(); - rt.check_worker_load_lower_bounds(&[12, 12, 12]); + rt.check_worker_load_lower_bounds(&[u(12), u(12), u(12)]); } #[test] fn test_resources_blocked_workers() { + let u = ResourceAmount::new_units; let mut rt = TestEnv::new(); rt.new_workers(&[4, 8, 2]); rt.new_assigned_tasks_cpus(&[&[4, 4, 4, 4, 4]]); rt.balance(); - assert!(rt.worker_load(100).get(0.into()) >= 4); - assert!(rt.worker_load(101).get(0.into()) >= 8); - assert_eq!(rt.worker_load(102).get(0.into()), 0); + assert!(rt.worker_load(100).get(0.into()) >= u(4)); + assert!(rt.worker_load(101).get(0.into()) >= u(8)); + assert_eq!(rt.worker_load(102).get(0.into()), u(0)); assert!(!rt.worker(100).is_parked()); assert!(!rt.worker(101).is_parked()); @@ -456,14 +463,15 @@ fn test_resources_blocked_workers() { #[test] fn test_resources_no_workers1() { + let u = ResourceAmount::new_units; let mut rt = TestEnv::new(); rt.new_workers(&[4, 8, 2]); rt.new_ready_tasks_cpus(&[8, 8, 16, 24]); rt.schedule(); - assert_eq!(rt.worker_load(100).get(0.into()), 0); - assert_eq!(rt.worker_load(101).get(0.into()), 16); - assert_eq!(rt.worker_load(102).get(0.into()), 0); + assert_eq!(rt.worker_load(100).get(0.into()), u(0)); + assert_eq!(rt.worker_load(101).get(0.into()), u(16)); + assert_eq!(rt.worker_load(102).get(0.into()), u(0)); let (sn, mn) = rt.core().take_sleeping_tasks(); assert_eq!(sn.len(), 2); @@ -472,7 +480,7 @@ fn test_resources_no_workers1() { #[test] fn test_resources_no_workers2() { - fn check(task_cpu_counts: &[ResourceAmount]) { + fn check(task_cpu_counts: &[ResourceUnits]) { println!("Checking order {:?}", task_cpu_counts); let mut rt = TestEnv::new(); @@ -487,15 +495,15 @@ fn test_resources_no_workers2() { rt.new_ready_tasks_cpus(task_cpu_counts); rt.schedule(); - assert_eq!(rt.worker_load(100).get(0.into()), 0); - assert_eq!(rt.worker_load(101).get(0.into()), 0); - assert_eq!(rt.worker_load(102).get(0.into()), 0); + assert!(rt.worker_load(100).get(0.into()).is_zero()); + assert!(rt.worker_load(101).get(0.into()).is_zero()); + assert!(rt.worker_load(102).get(0.into()).is_zero()); rt.new_workers(&[9, 10]); rt.schedule(); - assert_eq!(rt.worker_load(100).get(0.into()), 0); - assert_eq!(rt.worker_load(101).get(0.into()), 0); - assert_eq!(rt.worker_load(102).get(0.into()), 0); + assert!(rt.worker_load(100).get(0.into()).is_zero()); + assert!(rt.worker_load(101).get(0.into()).is_zero()); + assert!(rt.worker_load(102).get(0.into()).is_zero()); assert_eq!(rt.worker(103).sn_tasks().len(), 1); assert_eq!(rt.worker(104).sn_tasks().len(), 1); @@ -570,7 +578,7 @@ fn test_generic_resource_assign2() { None, vec![ ResourceDescriptorItem::range("Res0", 1, 10), - ResourceDescriptorItem::sum("Res1", 1000_000), + ResourceDescriptorItem::sum("Res1", 1_000_000), ], ), ]); @@ -646,7 +654,7 @@ fn test_generic_resource_balance1() { None, vec![ ResourceDescriptorItem::range("Res0", 1, 10), - ResourceDescriptorItem::sum("Res1", 1000_000), + ResourceDescriptorItem::sum("Res1", 1_000_000), ], ), ]); @@ -695,7 +703,7 @@ fn test_generic_resource_balance2() { None, vec![ ResourceDescriptorItem::range("Res0", 1, 10), - ResourceDescriptorItem::sum("Res1", 1000_000), + ResourceDescriptorItem::sum("Res1", 1_000_000), ], ), ]); diff --git a/crates/tako/src/internal/tests/test_worker.rs b/crates/tako/src/internal/tests/test_worker.rs index bb4706982..2fcf963af 100644 --- a/crates/tako/src/internal/tests/test_worker.rs +++ b/crates/tako/src/internal/tests/test_worker.rs @@ -3,7 +3,7 @@ use crate::internal::messages::worker::{ ComputeTaskMsg, NewWorkerMsg, ToWorkerMessage, WorkerResourceCounts, }; use crate::internal::server::workerload::WorkerResources; -use crate::internal::tests::utils::resources::ResourceRequestBuilder; +use crate::internal::tests::utils::resources::{ra_builder, ResourceRequestBuilder}; use crate::internal::worker::comm::WorkerComm; use crate::internal::worker::configuration::OverviewConfiguration; use crate::internal::worker::rpc::process_worker_message; @@ -13,6 +13,7 @@ use crate::resources::{ResourceDescriptor, ResourceMap}; use crate::worker::{ServerLostPolicy, WorkerConfiguration}; use crate::{Set, TaskId, WorkerId}; use smallvec::smallvec; +use std::ops::Deref; use std::time::Duration; use tokio::sync::oneshot::Receiver; @@ -65,7 +66,7 @@ fn create_test_worker_state(config: WorkerConfiguration) -> WorkerStateRef { config, None, resource_map, - Box::new(TestLauncher::default()), + Box::new(TestLauncher), "testuid".to_string(), ) } @@ -158,12 +159,12 @@ fn test_worker_other_workers() { assert!(state.ready_task_queue.worker_resources().is_empty()); let r1 = WorkerResourceCounts { - n_resources: vec![2, 0, 1], + n_resources: ra_builder(&[2, 0, 1]).deref().clone(), }; let wr1 = WorkerResources::from_transport(r1.clone()); let r2 = WorkerResourceCounts { - n_resources: vec![2, 1], + n_resources: ra_builder(&[2, 1]).deref().clone(), }; let wr2 = WorkerResources::from_transport(r2.clone()); diff --git a/crates/tako/src/internal/tests/utils/env.rs b/crates/tako/src/internal/tests/utils/env.rs index a8625103d..46a19d183 100644 --- a/crates/tako/src/internal/tests/utils/env.rs +++ b/crates/tako/src/internal/tests/utils/env.rs @@ -18,7 +18,9 @@ use crate::internal::tests::utils::schedule; use crate::internal::tests::utils::task::TaskBuilder; use crate::internal::transfer::auth::{deserialize, serialize}; use crate::internal::worker::configuration::OverviewConfiguration; -use crate::resources::{ResourceAmount, ResourceDescriptorItem, ResourceDescriptorKind}; +use crate::resources::{ + ResourceAmount, ResourceDescriptorItem, ResourceDescriptorKind, ResourceUnits, +}; use crate::task::SerializedTaskContext; use crate::worker::{ServerLostPolicy, WorkerConfiguration}; use crate::{TaskId, WorkerId}; @@ -111,7 +113,7 @@ impl TestEnv { gpu_families: Default::default(), }), idle_timeout: None, - time_limit: time_limit.clone(), + time_limit: *time_limit, on_server_lost: ServerLostPolicy::Stop, extra: Default::default(), }; @@ -126,7 +128,7 @@ impl TestEnv { self.new_workers_ext(&defs); } - pub fn new_ready_tasks_cpus(&mut self, tasks: &[ResourceAmount]) -> Vec { + pub fn new_ready_tasks_cpus(&mut self, tasks: &[ResourceUnits]) -> Vec { let tasks: Vec<_> = tasks .iter() .map(|n_cpus| { @@ -151,7 +153,7 @@ impl TestEnv { self._test_assign(task_id.into(), worker_id.into()); } - pub fn new_assigned_tasks_cpus(&mut self, tasks: &[&[ResourceAmount]]) { + pub fn new_assigned_tasks_cpus(&mut self, tasks: &[&[ResourceUnits]]) { for (i, tdefs) in tasks.iter().enumerate() { let w_id = WorkerId::new(100 + i as u32); let task_ids = self.new_ready_tasks_cpus(tdefs); diff --git a/crates/tako/src/internal/tests/utils/resources.rs b/crates/tako/src/internal/tests/utils/resources.rs index 4162ff812..829265d43 100644 --- a/crates/tako/src/internal/tests/utils/resources.rs +++ b/crates/tako/src/internal/tests/utils/resources.rs @@ -1,10 +1,16 @@ use std::time::Duration; use crate::internal::common::resources::request::{ResourceRequest, ResourceRequestEntry}; -use crate::internal::common::resources::{ResourceId, ResourceRequestVariants}; -use crate::resources::{AllocationRequest, NumOfNodes, ResourceAmount}; +use crate::internal::common::resources::{ResourceId, ResourceRequestVariants, ResourceVec}; +use crate::resources::{AllocationRequest, NumOfNodes, ResourceAmount, ResourceUnits}; pub use ResourceRequestBuilder as ResBuilder; +impl From for ResourceAmount { + fn from(val: u32) -> Self { + ResourceAmount::new_units(val) + } +} + #[derive(Default, Clone)] pub struct ResourceRequestBuilder { n_nodes: NumOfNodes, @@ -13,8 +19,8 @@ pub struct ResourceRequestBuilder { } impl ResourceRequestBuilder { - pub fn add>(self, id: Id, amount: ResourceAmount) -> Self { - self.add_compact(id, amount) + pub fn add, A: Into>(self, id: Id, units: A) -> Self { + self.add_compact(id, units) } pub fn n_nodes(mut self, n_nodes: NumOfNodes) -> Self { @@ -22,33 +28,41 @@ impl ResourceRequestBuilder { self } - pub fn cpus(self, count: ResourceAmount) -> Self { + pub fn cpus>(self, count: A) -> Self { self.add(0, count) } fn _add(&mut self, id: ResourceId, request: AllocationRequest) { self.resources.push(ResourceRequestEntry { - resource_id: id.into(), + resource_id: id, request, }); } - pub fn add_compact>(mut self, id: Id, amount: ResourceAmount) -> Self { - self._add(id.into(), AllocationRequest::Compact(amount)); + pub fn add_compact, A: Into>( + mut self, + id: Id, + amount: A, + ) -> Self { + self._add(id.into(), AllocationRequest::Compact(amount.into())); self } - pub fn add_force_compact>( + pub fn add_force_compact, A: Into>( mut self, id: Id, - amount: ResourceAmount, + amount: A, ) -> Self { - self._add(id.into(), AllocationRequest::ForceCompact(amount)); + self._add(id.into(), AllocationRequest::ForceCompact(amount.into())); self } - pub fn add_scatter>(mut self, id: Id, amount: ResourceAmount) -> Self { - self._add(id.into(), AllocationRequest::Scatter(amount)); + pub fn add_scatter, A: Into>( + mut self, + id: Id, + amount: A, + ) -> Self { + self._add(id.into(), AllocationRequest::Scatter(amount.into())); self } @@ -69,7 +83,7 @@ impl ResourceRequestBuilder { 0, ResourceRequestEntry { resource_id: 0.into(), - request: AllocationRequest::Compact(1), + request: AllocationRequest::Compact(ResourceAmount::new_units(1)), }, ) } @@ -81,17 +95,15 @@ impl ResourceRequestBuilder { } } -pub fn cpus_compact(count: ResourceAmount) -> ResBuilder { +pub fn cpus_compact(count: ResourceUnits) -> ResBuilder { ResBuilder::default().add(0, count) } -/* -pub fn cpus_force_compact(count: NumOfCpus) -> ResBuilder { - ResBuilder::default().cpus(CpuRequest::ForceCompact(count)) -} -pub fn cpus_scatter(count: NumOfCpus) -> ResBuilder { - ResBuilder::default().cpus(CpuRequest::Scatter(count)) -} -pub fn cpus_all() -> ResBuilder { - ResBuilder::default().cpus(CpuRequest::All) + +pub fn ra_builder(units: &[ResourceUnits]) -> ResourceVec { + let vec: Vec = units + .iter() + .copied() + .map(ResourceAmount::new_units) + .collect(); + vec.into() } -*/ diff --git a/crates/tako/src/internal/tests/utils/schedule.rs b/crates/tako/src/internal/tests/utils/schedule.rs index 1e58182bf..e1b595b4b 100644 --- a/crates/tako/src/internal/tests/utils/schedule.rs +++ b/crates/tako/src/internal/tests/utils/schedule.rs @@ -48,7 +48,7 @@ pub fn new_test_worker( on_new_worker(core, &mut TestComm::default(), worker); } -pub fn create_test_worker(core: &mut Core, worker_id: WorkerId, cpus: u64) { +pub fn create_test_worker(core: &mut Core, worker_id: WorkerId, cpus: u32) { let wcfg = create_test_worker_config(worker_id, ResourceDescriptor::simple(cpus)); new_test_worker( core, @@ -58,7 +58,7 @@ pub fn create_test_worker(core: &mut Core, worker_id: WorkerId, cpus: u64) { ); } -pub fn create_test_workers(core: &mut Core, cpus: &[u64]) { +pub fn create_test_workers(core: &mut Core, cpus: &[u32]) { for (i, c) in cpus.iter().enumerate() { let worker_id = WorkerId::new((100 + i) as u32); create_test_worker(core, worker_id, *c); @@ -88,8 +88,8 @@ pub(crate) fn force_assign_mn( ) { core.remove_from_ready_to_assign(task_id); let (task_map, worker_map) = core.split_tasks_workers_mut(); - let mut task = task_map.get_task_mut(task_id); - scheduler.assign_multinode(worker_map, &mut task, workers); + let task = task_map.get_task_mut(task_id); + scheduler.assign_multinode(worker_map, task, workers); } pub(crate) fn start_mn_task_on_worker(core: &mut Core, task_id: TaskId, worker_ids: Vec) { @@ -98,7 +98,7 @@ pub(crate) fn start_mn_task_on_worker(core: &mut Core, task_id: TaskId, worker_i force_assign_mn( core, &mut scheduler, - worker_ids.into_iter().map(|x| x.into()).collect(), + worker_ids.into_iter().collect(), task_id, ); scheduler.finish_scheduling(core, &mut comm); diff --git a/crates/tako/src/internal/tests/utils/shared.rs b/crates/tako/src/internal/tests/utils/shared.rs index 492a010e0..7d8bf7f2b 100644 --- a/crates/tako/src/internal/tests/utils/shared.rs +++ b/crates/tako/src/internal/tests/utils/shared.rs @@ -1,7 +1,7 @@ use crate::internal::worker::resources::allocator::ResourceAllocator; use crate::internal::worker::resources::map::ResourceLabelMap; use crate::resources::{ - ResourceDescriptor, ResourceDescriptorItem, ResourceDescriptorKind, ResourceMap, + ResourceAmount, ResourceDescriptor, ResourceDescriptorItem, ResourceDescriptorKind, ResourceMap, }; pub fn res_kind_range(start: u32, end: u32) -> ResourceDescriptorKind { @@ -26,8 +26,10 @@ pub fn res_kind_groups(groups: &[Vec<&str>]) -> ResourceDescriptorKind { } } -pub fn res_kind_sum(size: u64) -> ResourceDescriptorKind { - ResourceDescriptorKind::Sum { size } +pub fn res_kind_sum(size: u32) -> ResourceDescriptorKind { + ResourceDescriptorKind::Sum { + size: ResourceAmount::new_units(size), + } } pub fn res_item(name: &str, kind: ResourceDescriptorKind) -> ResourceDescriptorItem { diff --git a/crates/tako/src/internal/tests/utils/task.rs b/crates/tako/src/internal/tests/utils/task.rs index 443864936..d552013d7 100644 --- a/crates/tako/src/internal/tests/utils/task.rs +++ b/crates/tako/src/internal/tests/utils/task.rs @@ -72,7 +72,7 @@ impl TaskBuilder { self } - pub fn cpus_compact(mut self, count: ResourceAmount) -> TaskBuilder { + pub fn cpus_compact>(mut self, count: A) -> TaskBuilder { self.resources_builder = self.resources_builder.cpus(count); self } @@ -82,10 +82,10 @@ impl TaskBuilder { self } - pub fn add_resource>( + pub fn add_resource, A: Into>( mut self, id: Id, - amount: ResourceAmount, + amount: A, ) -> TaskBuilder { self.resources_builder = self.resources_builder.add(id, amount); self diff --git a/crates/tako/src/internal/worker/reactor.rs b/crates/tako/src/internal/worker/reactor.rs index d86bf227d..0227296c5 100644 --- a/crates/tako/src/internal/worker/reactor.rs +++ b/crates/tako/src/internal/worker/reactor.rs @@ -6,13 +6,14 @@ use crate::internal::worker::taskenv::TaskEnv; use crate::launcher::{LaunchContext, StopReason, TaskFuture, TaskLaunchData, TaskResult}; use crate::TaskId; use futures::future::Either; +use std::rc::Rc; use tokio::sync::oneshot; pub(crate) fn run_task( state: &mut WorkerState, state_ref: &WorkerStateRef, task_id: TaskId, - allocation: Allocation, + allocation: Rc, resource_index: usize, ) { log::debug!("Task={} assigned", task_id); diff --git a/crates/tako/src/internal/worker/resources/allocator.rs b/crates/tako/src/internal/worker/resources/allocator.rs index 6d55a906b..29cd339be 100644 --- a/crates/tako/src/internal/worker/resources/allocator.rs +++ b/crates/tako/src/internal/worker/resources/allocator.rs @@ -1,27 +1,25 @@ use crate::internal::common::resources::request::{ AllocationRequest, ResourceRequest, ResourceRequestEntry, ResourceRequestVariants, }; -use crate::internal::common::resources::{ - ResourceAllocation, ResourceAllocations, ResourceId, ResourceVec, -}; +use crate::internal::common::resources::{ResourceId, ResourceVec}; use crate::internal::server::workerload::WorkerResources; -use crate::internal::worker::resources::counts::{ - resource_count_add_at, ResourceCount, ResourceCountVec, -}; +use crate::internal::worker::resources::concise::{ConciseFreeResources, ConciseResourceState}; use crate::internal::worker::resources::map::ResourceLabelMap; use crate::internal::worker::resources::pool::ResourcePool; -use crate::resources::{Allocation, ResourceAmount, ResourceDescriptor, ResourceMap}; -use smallvec::smallvec; +use crate::resources::{ + Allocation, ResourceAmount, ResourceDescriptor, ResourceMap, ResourceUnits, +}; +use std::rc::Rc; use std::time::Duration; pub struct ResourceAllocator { pools: ResourceVec, - free_resources: ResourceCountVec, + free_resources: ConciseFreeResources, remaining_time: Option, blocked_requests: Vec, higher_priority_blocked_requests: usize, - running_tasks: Vec, // TODO: Rework on multiset? + running_tasks: Vec>, // TODO: Rework on multiset? own_resources: WorkerResources, } @@ -31,7 +29,7 @@ struct BlockedRequest { request: ResourceRequest, /// Reachable state of free resources AFTER some of currently running tasks is finished AND /// this resource request is enabled - witnesses: Vec, + witnesses: Vec, } impl ResourceAllocator { @@ -59,8 +57,13 @@ impl ResourceAllocator { pools[idx] = ResourcePool::new(&item.kind, idx, label_map); } - let free_resources = - ResourceCountVec::new(pools.iter().map(|p| p.count()).collect::>().into()); + let free_resources = ConciseFreeResources::new( + pools + .iter() + .map(|p| p.concise_state()) + .collect::>() + .into(), + ); ResourceAllocator { pools, @@ -83,34 +86,24 @@ impl ResourceAllocator { self.blocked_requests.clear(); } - fn claim_resources(&mut self, counts: ResourceCountVec) -> Allocation { - assert!(self.pools.len() >= counts.len()); - let mut allocations = ResourceAllocations::new(); - for (i, (pool, count)) in self.pools.iter_mut().zip(counts.all_counts()).enumerate() { - if count.iter().sum::() > 0 { - allocations.push(ResourceAllocation { - resource: ResourceId::new(i as u32), - value: pool.claim_resources(count), - }); - } + fn release_allocation_helper(&mut self, allocation: &Allocation) { + self.free_resources.add(allocation); + for al in &allocation.resources { + self.pools[al.resource_id].release_allocation(al); } - Allocation::new(Vec::new(), allocations, counts) } - pub fn release_allocation(&mut self, allocation: Allocation) { - for al in allocation.resources { - self.pools[al.resource].release_allocation(al.value); - } - self.free_resources.add(&allocation.counts); + pub fn release_allocation(&mut self, allocation: Rc) { + self.release_allocation_helper(&allocation); let position = self .running_tasks .iter() - .position(|c| c == &allocation.counts) + .position(|a| Rc::ptr_eq(a, &allocation)) .unwrap(); self.running_tasks.swap_remove(position); } - pub fn take_running_allocations(self) -> Vec { + pub fn take_running_allocations(self) -> Vec> { self.running_tasks } @@ -120,103 +113,47 @@ impl ResourceAllocator { fn has_resources_for_entry( pool: &ResourcePool, - free: &ResourceCount, + free: &ConciseResourceState, policy: &AllocationRequest, ) -> bool { - let sum = free.iter().sum::(); + let max_alloc = free.amount_max_alloc(); match policy { AllocationRequest::Compact(amount) | AllocationRequest::Scatter(amount) => { - *amount <= sum + *amount <= max_alloc } AllocationRequest::ForceCompact(amount) => { - if *amount > sum { + if amount.is_zero() { + return true; + } + if *amount > max_alloc { return false; } - if *amount == 0 { - return true; + let (units, fractions) = amount.split(); + if fractions > 0 { + // Fractions and ForceCompact is not supported in the current version + // It is checked in resource request validation, this code should be unreachable + unreachable!() } - let socket_size = (((amount - 1) / pool.min_group_size()) as usize) + 1; - if free.len() < socket_size { + let socket_count = (((units - 1) / pool.min_group_size()) as usize) + 1; + if free.n_groups() < socket_count { return false; } - let mut free = free.clone(); - free.sort_unstable(); - let sum = free.iter().rev().take(socket_size).sum(); - *amount <= sum - } - AllocationRequest::All => sum == pool.full_size(), - } - } + let mut groups: Vec = free.groups().to_vec(); + groups.sort_unstable(); + let sum: ResourceUnits = groups.iter().rev().take(socket_count).sum(); - fn try_allocate_compact(free: &mut ResourceCount, mut amount: ResourceAmount) -> ResourceCount { - let mut result = ResourceCount::new(); - loop { - if let Some((i, c)) = free - .iter_mut() - .enumerate() - .filter(|(_i, c)| **c >= amount) - .min_by_key(|(_i, c)| **c) - { - resource_count_add_at(&mut result, i, amount); - *c -= amount; - return result; - } else { - let (i, c) = free - .iter_mut() - .enumerate() - .max_by_key(|(_i, c)| **c) - .unwrap(); - resource_count_add_at(&mut result, i, *c); - amount -= *c; - *c = 0; - } - } - } - - fn try_allocate_scatter(free: &mut ResourceCount, mut amount: ResourceAmount) -> ResourceCount { - assert!(free.iter().sum::() >= amount); - let mut result = ResourceCount::new(); - - let mut idx = 0; - while amount > 0 { - if free[idx] > 0 { - free[idx] -= 1; - amount -= 1; - resource_count_add_at(&mut result, idx, 1); - } - idx = (idx + 1) % free.len(); - } - result - } - - fn allocate_entry(free: &mut ResourceCount, policy: &AllocationRequest) -> ResourceCount { - match policy { - AllocationRequest::Compact(amount) - | AllocationRequest::ForceCompact(amount) - | AllocationRequest::Scatter(amount) - if free.len() == 1 => - { - free[0] -= *amount; - smallvec![*amount] - } - AllocationRequest::Compact(amount) | AllocationRequest::ForceCompact(amount) => { - Self::try_allocate_compact(free, *amount) - } - AllocationRequest::Scatter(amount) => Self::try_allocate_scatter(free, *amount), - AllocationRequest::All => { - let result = free.clone(); - free.fill(0); - result + units <= sum } + AllocationRequest::All => max_alloc == pool.full_size(), } } fn compute_witness<'b>( pools: &[ResourcePool], - free: &ResourceCountVec, + free: &ConciseFreeResources, request: &ResourceRequest, - running: impl Iterator, - ) -> Option { + running: impl Iterator>, + ) -> Option { let mut free = free.clone(); if Self::has_resources_for_request(pools, &free, request) { return Some(free); @@ -232,7 +169,7 @@ impl ResourceAllocator { fn has_resources_for_request( pools: &[ResourcePool], - free: &ResourceCountVec, + free: &ConciseFreeResources, request: &ResourceRequest, ) -> bool { request.entries().iter().all(|entry| { @@ -247,38 +184,34 @@ impl ResourceAllocator { fn compute_witnesses( pools: &[ResourcePool], - free_resources: &ResourceCountVec, - running: &mut [ResourceCountVec], + free_resources: &ConciseFreeResources, + running: &mut [Rc], request: &ResourceRequest, - ) -> Vec { + ) -> Vec { let mut full = free_resources.clone(); for running in running.iter() { full.add(running); } - let mut witnesses = Vec::with_capacity(free_resources.len() * 2); + let mut witnesses: Vec = + Vec::with_capacity(2 * free_resources.n_resources()); - let compute_witnesses = |witnesses: &mut Vec, - running: &[ResourceCountVec]| { + let compute_witnesses = |witnesses: &mut Vec, + running: &[Rc]| { let w = Self::compute_witness(pools, free_resources, request, running.iter()); witnesses.push(w.unwrap()); let w = Self::compute_witness(pools, free_resources, request, running.iter().rev()); witnesses.push(w.unwrap()); }; - for i in 0..free_resources.len() { + for i in 0..free_resources.n_resources() { let idx = ResourceId::from(i as u32); + /*running + .sort_unstable_by_key(|x| (x.get_sum(idx), -(x.fraction(&full) * 10_000.0) as u32));*/ running.sort_unstable_by_key(|x| { - ( - x.get(idx).iter().sum::(), - (x.fraction(&full) * 10_000.0) as u32, - ) - }); - running.sort_unstable_by_key(|x| { - ( - x.get(idx).iter().sum::(), - -(x.fraction(&full) * 10_000.0) as u32, - ) + x.resource_allocation(idx) + .map(|a| a.amount) + .unwrap_or(ResourceAmount::ZERO) }); compute_witnesses(&mut witnesses, running); } @@ -292,7 +225,7 @@ impl ResourceAllocator { }); } - fn check_blocked_request(&mut self, allocation: &ResourceCountVec) -> bool { + fn check_blocked_request(&mut self, allocation: &Allocation) -> bool { for blocked in &mut self.blocked_requests[..self.higher_priority_blocked_requests] { if blocked.witnesses.is_empty() { blocked.witnesses = Self::compute_witnesses( @@ -304,7 +237,7 @@ impl ResourceAllocator { } for witness in &blocked.witnesses { let mut free = witness.clone(); - assert!(free.remove(allocation)); + free.remove(allocation); if !Self::has_resources_for_request(&self.pools, &free, &blocked.request) { return false; } @@ -316,14 +249,15 @@ impl ResourceAllocator { pub fn try_allocate( &mut self, request: &ResourceRequestVariants, - ) -> Option<(Allocation, usize)> { - request.requests().iter().enumerate().find_map(|(i, r)| { - self.try_allocate_counts(r) - .map(|c| (self.claim_resources(c), i)) - }) + ) -> Option<(Rc, usize)> { + request + .requests() + .iter() + .enumerate() + .find_map(|(i, r)| self.try_allocate_variant(r).map(|c| (c, i))) } - fn try_allocate_counts(&mut self, request: &ResourceRequest) -> Option { + fn try_allocate_variant(&mut self, request: &ResourceRequest) -> Option> { if let Some(remaining_time) = self.remaining_time { if remaining_time < request.min_time() { return None; @@ -339,24 +273,21 @@ impl ResourceAllocator { return None; } - let mut allocation = ResourceCountVec::default(); + let mut allocation = Allocation::new(); for entry in request.entries() { - allocation.set( - entry.resource_id, - Self::allocate_entry( - self.free_resources.get_mut(entry.resource_id), - &entry.request, - ), - ) + let pool = self.pools.get_mut(entry.resource_id.as_usize()).unwrap(); + allocation + .add_resource_allocation(pool.claim_resources(entry.resource_id, &entry.request)) } + self.free_resources.remove(&allocation); if !self.check_blocked_request(&allocation) { - self.free_resources.add(&allocation); + self.release_allocation_helper(&allocation); return None; } - - self.running_tasks.push(allocation.clone()); - Some(allocation) + let allocation_rc = Rc::new(allocation); + self.running_tasks.push(allocation_rc.clone()); + Some(allocation_rc) } pub fn difficulty_score(&self, entry: &ResourceRequestEntry) -> f32 { @@ -364,20 +295,22 @@ impl ResourceAllocator { .pools .get(entry.resource_id) .map(|x| x.full_size()) - .unwrap_or(0); - if size == 0 { + .unwrap_or(ResourceAmount::ZERO); + if size.is_zero() { 0.0f32 } else { match entry.request { AllocationRequest::Compact(amount) | AllocationRequest::Scatter(amount) => { - amount as f32 / size as f32 + amount.total_fractions() as f32 / size.total_fractions() as f32 } AllocationRequest::ForceCompact(amount) if self.pools[entry.resource_id].n_groups() == 1 => { - amount as f32 / size as f32 + amount.total_fractions() as f32 / size.total_fractions() as f32 + } + AllocationRequest::ForceCompact(amount) => { + (amount.total_fractions() * 2) as f32 / size.total_fractions() as f32 } - AllocationRequest::ForceCompact(amount) => (amount * 2) as f32 / size as f32, AllocationRequest::All => 2.0, } } @@ -385,29 +318,29 @@ impl ResourceAllocator { pub fn validate(&self) { #[cfg(debug_assertions)] - for (pool, count) in self.pools.iter().zip(self.free_resources.all_counts()) { + for (pool, state) in self.pools.iter().zip(self.free_resources.all_states()) { pool.validate(); - assert_eq!(&pool.count(), count); + assert_eq!(pool.concise_state().strip_zeros(), state.strip_zeros()); } } } #[cfg(test)] mod tests { + use crate::internal::common::resources::descriptor::{ ResourceDescriptor, ResourceDescriptorKind, }; - use crate::internal::common::resources::{ - Allocation, AllocationValue, ResourceId, ResourceRequestVariants, - }; + use crate::internal::common::resources::{Allocation, ResourceId, ResourceRequestVariants}; use crate::internal::tests::utils::resources::{cpus_compact, ResBuilder}; use crate::internal::tests::utils::shared::res_allocator_from_descriptor; use crate::internal::tests::utils::sorted_vec; use crate::internal::worker::resources::allocator::ResourceAllocator; - use crate::internal::worker::resources::counts::ResourceCountVec; + use crate::internal::worker::resources::concise::{ConciseFreeResources, ConciseResourceState}; use crate::internal::worker::resources::pool::ResourcePool; - use crate::resources::{ResourceAmount, ResourceDescriptorItem}; - use smallvec::smallvec; + use crate::resources::{ResourceAmount, ResourceDescriptorItem, ResourceIndex, ResourceUnits}; + + use std::rc::Rc; use std::time::Duration; impl ResourceAllocator { @@ -423,8 +356,8 @@ mod tests { } pub fn simple_descriptor( - n_sockets: ResourceAmount, - socket_size: ResourceAmount, + n_sockets: ResourceUnits, + socket_size: ResourceUnits, ) -> ResourceDescriptor { ResourceDescriptor::new(vec![ResourceDescriptorItem { name: "cpus".to_string(), @@ -437,8 +370,8 @@ mod tests { } fn simple_allocator( - free: &[ResourceAmount], - running: &[&[ResourceAmount]], + free: &[ResourceUnits], + running: &[&[ResourceUnits]], remaining_time: Option, ) -> ResourceAllocator { let mut names = vec!["cpus".to_string()]; @@ -455,7 +388,7 @@ mod tests { } ResourceDescriptorItem { name: names[i].clone(), - kind: ResourceDescriptorKind::simple_indices(total as u32), + kind: ResourceDescriptorKind::simple_indices(total), } }) .collect(); @@ -463,10 +396,12 @@ mod tests { let mut ac = res_allocator_from_descriptor(descriptor); for r in running { - let c = ResourceCountVec::new_simple(r); - assert!(ac.free_resources.remove(&c)); - ac.claim_resources(c.clone()); - ac.running_tasks.push(c); + assert_eq!(r.len(), 1); // As far I see, only 1 element array is now used in tests + let units = r[0]; + let rq = ResBuilder::default() + .add(0, ResourceAmount::new_units(units)) + .finish(); + ac.try_allocate_variant(&rq).unwrap(); } ac.init_allocator(remaining_time); ac @@ -474,20 +409,21 @@ mod tests { fn simple_alloc( allocator: &mut ResourceAllocator, - counts: &[ResourceAmount], + counts: &[ResourceUnits], expect_pass: bool, ) { let mut builder = ResBuilder::default(); for (i, count) in counts.iter().enumerate() { if *count > 0 { - builder = builder.add(ResourceId::from(i as u32), *count); + builder = builder.add( + ResourceId::from(i as u32), + ResourceAmount::new_units(*count), + ); } } - let al = allocator.try_allocate_counts(&builder.finish()); + let al = allocator.try_allocate_variant(&builder.finish()); if expect_pass { - let r = al.unwrap(); - r.assert_eq(counts); - allocator.claim_resources(r); + al.unwrap(); } else { assert!(al.is_none()); } @@ -507,7 +443,7 @@ mod tests { &Default::default(), ), ]; - let free = ResourceCountVec::new_simple(&[1, 2]); + let free = ConciseFreeResources::new_simple(&[1, 2]); let rq = ResBuilder::default().add(0, 3).add(1, 1).finish(); assert!(ResourceAllocator::compute_witness(&pools, &free, &rq, [].iter()).is_none()); @@ -515,7 +451,7 @@ mod tests { &pools, &mut free.clone(), &rq, - [ResourceCountVec::new_simple(&[1])].iter() + [Rc::new(Allocation::new_simple(&[1]))].iter() ) .is_none()); assert_eq!( @@ -524,16 +460,16 @@ mod tests { &free, &rq, [ - ResourceCountVec::new_simple(&[1]), - ResourceCountVec::new_simple(&[1]), - ResourceCountVec::new_simple(&[1]) + Rc::new(Allocation::new_simple(&[1])), + Rc::new(Allocation::new_simple(&[1])), + Rc::new(Allocation::new_simple(&[1])) ] .iter() ), - Some(ResourceCountVec::new_simple(&[3, 2])) + Some(ConciseFreeResources::new_simple(&[3, 2])) ); - let free = ResourceCountVec::new_simple(&[4, 2]); + let free = ConciseFreeResources::new_simple(&[4, 2]); assert_eq!( ResourceAllocator::compute_witness(&pools, &free, &rq, [].iter()), Some(free.clone()) @@ -546,15 +482,15 @@ mod tests { &free, &rq, [ - ResourceCountVec::new_simple(&[2]), - ResourceCountVec::new_simple(&[2]), - ResourceCountVec::new_simple(&[2]), - ResourceCountVec::new_simple(&[0, 3]), - ResourceCountVec::new_simple(&[0, 3]) + Rc::new(Allocation::new_simple(&[2])), + Rc::new(Allocation::new_simple(&[2])), + Rc::new(Allocation::new_simple(&[2])), + Rc::new(Allocation::new_simple(&[0, 3])), + Rc::new(Allocation::new_simple(&[0, 3])) ] .iter() ), - Some(ResourceCountVec::new_simple(&[10, 5])) + Some(ConciseFreeResources::new_simple(&[10, 5])) ); } @@ -578,9 +514,9 @@ mod tests { fn test_allocator_priority_check1() { // Running: [1,1,1], free: [1], try: [2p1] [1p0] let mut allocator = simple_allocator(&[1], &[&[1], &[1], &[1]], None); - simple_alloc(&mut allocator, &[2], false); allocator.free_resources.assert_eq(&[1]); + allocator.validate(); allocator.close_priority_level(); @@ -731,13 +667,9 @@ mod tests { let (al, idx) = allocator.try_allocate(&rq).unwrap(); assert_eq!(idx, 0); assert_eq!(al.resources.len(), 1); - assert_eq!(al.resources[0].resource, ResourceId::new(0)); - assert_eq!(al.resources[0].value.get_checked_indices().len(), 3); - assert!(al.resources[0] - .value - .get_checked_indices() - .iter() - .all(|x| *x < 4)); + assert_eq!(al.resources[0].resource_id, ResourceId::new(0)); + assert_eq!(al.resources[0].indices.len(), 3); + assert!(al.resources[0].resource_indices().all(|x| x.as_num() < 4)); let rq = cpus_compact(2).finish_v(); assert!(allocator.try_allocate(&rq).is_none()); @@ -751,15 +683,18 @@ mod tests { let rq = cpus_compact(4).finish_v(); let (al, _idx) = allocator.try_allocate(&rq).unwrap(); assert_eq!( - al.resources[0].value.get_checked_indices(), - vec![0, 1, 2, 3] + al.resources[0].resource_indices().collect::>(), + vec![3.into(), 2.into(), 1.into(), 0.into()] ); assert_eq!(allocator.running_tasks.len(), 1); allocator.release_allocation(al); assert_eq!(allocator.running_tasks.len(), 0); - assert_eq!(allocator.free_resources.get(0.into())[0], 4); - let v: Vec = vec![4]; - assert_eq!(allocator.pools[ResourceId::new(0)].count().to_vec(), v); + allocator + .free_resources + .get(0.into()) + .amount_sum() + .assert_eq_units(4); + assert_eq!(allocator.free_resources.get(0.into()).n_groups(), 1); allocator.init_allocator(None); @@ -783,12 +718,12 @@ mod tests { assert!(allocator.try_allocate(&rq2).is_none()); let mut v = Vec::new(); - v.append(&mut al1.resources[0].value.get_checked_indices()); + v.extend(al1.resources[0].resource_indices()); assert_eq!(v.len(), 1); - v.append(&mut al5.resources[0].value.get_checked_indices()); + v.extend(al5.resources[0].resource_indices()); assert_eq!(v.len(), 3); - v.append(&mut al3.resources[0].value.get_checked_indices()); - assert_eq!(sorted_vec(v), vec![0, 1, 2, 3]); + v.extend(al3.resources[0].resource_indices()); + assert_eq!(sorted_vec(v), vec![0.into(), 1.into(), 2.into(), 3.into()]); allocator.validate(); } @@ -857,13 +792,13 @@ mod tests { assert_eq!( al.resource_allocation(0.into()) .unwrap() - .value - .get_checked_indices(), - (0..24u32).collect::>() + .resource_indices() + .collect::>(), + (0..24u32).rev().map(ResourceIndex::new).collect::>() ); - assert_eq!(allocator.get_current_free(0), 0); + assert!(allocator.get_current_free(0).is_zero()); allocator.release_allocation(al); - assert_eq!(allocator.get_current_free(0), 24); + assert_eq!(allocator.get_current_free(0), ResourceAmount::new_units(24)); allocator.validate(); } @@ -877,13 +812,13 @@ mod tests { assert_eq!( al.resource_allocation(0.into()) .unwrap() - .value - .get_checked_indices(), - (0..24u32).collect::>() + .resource_indices() + .collect::>(), + (0..24u32).map(ResourceIndex::new).collect::>() ); - assert_eq!(allocator.get_current_free(0), 0); + assert!(allocator.get_current_free(0).is_zero()); allocator.release_allocation(al); - assert_eq!(allocator.get_current_free(0), 24); + assert_eq!(allocator.get_current_free(0), ResourceAmount::new_units(24)); allocator.init_allocator(None); @@ -1018,7 +953,9 @@ mod tests { }, ResourceDescriptorItem { name: "res1".to_string(), - kind: ResourceDescriptorKind::Sum { size: 100_000_000 }, + kind: ResourceDescriptorKind::Sum { + size: ResourceAmount::new_units(100_000_000), + }, }, ResourceDescriptorItem { name: "res2".to_string(), @@ -1033,13 +970,13 @@ mod tests { assert_eq!( allocator.free_resources, - ResourceCountVec::new( + ConciseFreeResources::new( vec![ - smallvec![4], - smallvec![96], - smallvec![100_000_000], - smallvec![2], - smallvec![2] + ConciseResourceState::new_simple(&[4]), + ConciseResourceState::new_simple(&[96]), + ConciseResourceState::new_simple(&[100_000_000]), + ConciseResourceState::new_simple(&[2]), + ConciseResourceState::new_simple(&[2]), ] .into() ) @@ -1048,45 +985,42 @@ mod tests { let rq: ResourceRequestVariants = cpus_compact(1) .add(4, 1) .add(1, 12) - .add(2, 1000_000) + .add(2, 1_000_000) .finish_v(); rq.validate().unwrap(); let (al, _) = allocator.try_allocate(&rq).unwrap(); assert_eq!(al.resources.len(), 4); - assert_eq!(al.resources[0].resource.as_num(), 0); - - assert_eq!(al.resources[1].resource.as_num(), 1); - assert!(matches!( - &al.resources[1].value, - AllocationValue::Indices(indices) if indices.len() == 12 - )); - - assert_eq!(al.resources[2].resource.as_num(), 2); - assert!(matches!( - &al.resources[2].value, - AllocationValue::Sum(1000_000) - )); - - assert_eq!(al.resources[3].resource.as_num(), 4); - assert!(matches!( - &al.resources[3].value, - AllocationValue::Indices(indices) if indices.len() == 1 - )); - - assert_eq!(allocator.get_current_free(1), 84); - assert_eq!(allocator.get_current_free(2), 99_000_000); - assert_eq!(allocator.get_current_free(3), 2); - assert_eq!(allocator.get_current_free(4), 1); + assert_eq!(al.resources[0].resource_id.as_num(), 0); + + assert_eq!(al.resources[1].resource_id.as_num(), 1); + assert_eq!(al.resources[1].indices.len(), 12); + + assert_eq!(al.resources[2].resource_id.as_num(), 2); + assert_eq!(al.resources[2].amount, ResourceAmount::new_units(1_000_000)); + + assert_eq!(al.resources[3].resource_id.as_num(), 4); + assert_eq!(al.resources[3].indices.len(), 1); + + assert_eq!(allocator.get_current_free(1), ResourceAmount::new_units(84)); + assert_eq!( + allocator.get_current_free(2), + ResourceAmount::new_units(99_000_000) + ); + assert_eq!(allocator.get_current_free(3), ResourceAmount::new_units(2)); + assert_eq!(allocator.get_current_free(4), ResourceAmount::new_units(1)); let rq = cpus_compact(1).add(4, 2).finish_v(); assert!(allocator.try_allocate(&rq).is_none()); allocator.release_allocation(al); - assert_eq!(allocator.get_current_free(1), 96); - assert_eq!(allocator.get_current_free(2), 100_000_000); - assert_eq!(allocator.get_current_free(3), 2); - assert_eq!(allocator.get_current_free(4), 2); + assert_eq!(allocator.get_current_free(1), ResourceAmount::new_units(96)); + assert_eq!( + allocator.get_current_free(2), + ResourceAmount::new_units(100_000_000) + ); + assert_eq!(allocator.get_current_free(3), ResourceAmount::new_units(2)); + assert_eq!(allocator.get_current_free(4), ResourceAmount::new_units(2)); allocator.init_allocator(None); assert!(allocator.try_allocate(&rq).is_some()); @@ -1116,4 +1050,315 @@ mod tests { allocator.validate(); } + + #[test] + fn test_allocator_sum_max_fractions() { + let descriptor = ResourceDescriptor::new(vec![ResourceDescriptorItem { + name: "cpus".to_string(), + kind: ResourceDescriptorKind::Sum { + size: ResourceAmount::new(0, 300), + }, + }]); + let mut allocator = test_allocator(&descriptor); + allocator.init_allocator(None); + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(1, 0)) + .finish_v(); + assert!(allocator.try_allocate(&rq).is_none()); + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(0, 301)) + .finish_v(); + assert!(allocator.try_allocate(&rq).is_none()); + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(0, 250)) + .finish_v(); + assert!(allocator.try_allocate(&rq).is_some()); + } + + #[test] + fn test_allocator_indices_and_fractions() { + let descriptor = simple_descriptor(1, 4); + let mut allocator = test_allocator(&descriptor); + allocator.init_allocator(None); + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(4, 1)) + .finish_v(); + assert!(allocator.try_allocate(&rq).is_none()); + + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(2, 1500)) + .finish_v(); + let al1 = allocator.try_allocate(&rq).unwrap().0; + assert_eq!(al1.resources[0].indices.len(), 3); + assert_eq!(al1.resources[0].indices[0].fractions, 0); + assert_eq!(al1.resources[0].indices[1].fractions, 0); + assert_eq!(al1.resources[0].indices[2].fractions, 1500); + assert_eq!(al1.resources[0].amount, ResourceAmount::new(2, 1500)); + + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(0, 5200)) + .finish_v(); + let al2 = allocator.try_allocate(&rq).unwrap().0; + assert_eq!(al2.resources[0].indices.len(), 1); + assert_eq!(al2.resources[0].indices[0].fractions, 5200); + assert_eq!( + al2.resources[0].indices[0].index, + al1.resources[0].indices[2].index + ); + assert_eq!(al2.resources[0].amount, ResourceAmount::new(0, 5200)); + + let al3 = allocator.try_allocate(&rq).unwrap().0; + assert_eq!(al3.resources[0].indices.len(), 1); + assert_eq!(al3.resources[0].indices[0].fractions, 5200); + assert_ne!( + al3.resources[0].indices[0].index, + al1.resources[0].indices[2].index + ); + assert_eq!(al3.resources[0].amount, ResourceAmount::new(0, 5200)); + + assert!(allocator.try_allocate(&rq).is_none()); + + allocator.release_allocation(al1); + + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new(2, 9600) + ); + + allocator.release_allocation(al3); + allocator.release_allocation(al2); + + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new(4, 0) + ); + } + + #[test] + fn test_allocator_fractions_compactness() { + // Two 0.75 does not gives 1.5 + let descriptor = simple_descriptor(1, 2); + let mut allocator = test_allocator(&descriptor); + allocator.init_allocator(None); + let rq1 = ResBuilder::default() + .add(0, ResourceAmount::new(0, 7500)) + .finish_v(); + let rq2 = ResBuilder::default() + .add(0, ResourceAmount::new(0, 2500)) + .finish_v(); + + let al1 = allocator.try_allocate(&rq1).unwrap().0; + let al2 = allocator.try_allocate(&rq1).unwrap().0; + let al3 = allocator.try_allocate(&rq2).unwrap().0; + let al4 = allocator.try_allocate(&rq2).unwrap().0; + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new(0, 0) + ); + allocator.release_allocation(al1); + allocator.release_allocation(al2); + + allocator.init_allocator(None); + + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new(1, 5000) + ); + + let rq3 = ResBuilder::default() + .add(0, ResourceAmount::new(1, 5000)) + .finish_v(); + assert!(allocator.try_allocate(&rq3).is_none()); + allocator.release_allocation(al4); + allocator.init_allocator(None); + let al5 = allocator.try_allocate(&rq3).unwrap().0; + allocator.release_allocation(al3); + allocator.release_allocation(al5); + + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new(2, 0) + ); + } + + #[test] + fn test_allocator_groups_and_fractions_scatter() { + let descriptor = simple_descriptor(3, 2); + let mut allocator = test_allocator(&descriptor); + allocator.init_allocator(None); + let rq1 = ResBuilder::default() + .add_scatter(0, ResourceAmount::new(6, 1)) + .finish_v(); + assert!(allocator.try_allocate(&rq1).is_none()); + + let rq1 = ResBuilder::default() + .add_scatter(0, ResourceAmount::new(2, 5000)) + .finish_v(); + let al1 = allocator.try_allocate(&rq1).unwrap().0; + let al2 = allocator.try_allocate(&rq1).unwrap().0; + allocator.validate(); + let r1 = &al1.resources[0].indices; + let r2 = &al2.resources[0].indices; + assert_eq!(r1[2].fractions, 5000); + assert_eq!(r2[2].fractions, 5000); + assert_eq!(r1[2].group_idx, r2[2].group_idx); + allocator.release_allocation(al1); + allocator.release_allocation(al2); + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new_units(6) + ); + } + + #[test] + fn test_allocator_groups_and_fractions() { + let descriptor = simple_descriptor(3, 2); + let mut allocator = test_allocator(&descriptor); + allocator.init_allocator(None); + let rq1 = ResBuilder::default() + .add(0, ResourceAmount::new(6, 1)) + .finish_v(); + assert!(allocator.try_allocate(&rq1).is_none()); + + let rq1 = ResBuilder::default() + .add_compact(0, ResourceAmount::new(3, 5000)) + .finish_v(); + let al1 = allocator.try_allocate(&rq1).unwrap().0; + // pools [[1 1] [1 0.5] [0 0] + + let r1 = &al1.resources[0].indices; + assert_eq!(r1.len(), 4); + assert_eq!(r1[0].group_idx, r1[1].group_idx); + assert_eq!(r1[2].group_idx, r1[3].group_idx); + assert_ne!(r1[0].group_idx, r1[2].group_idx); + + let rq2 = ResBuilder::default() + .add_compact(0, ResourceAmount::new(0, 4000)) + .finish_v(); + let al2 = allocator.try_allocate(&rq2).unwrap().0; + // pools [[1 1] [0.1 1] [0 0] + + let r2 = &al2.resources[0].indices; + assert_eq!(r2.len(), 1); + assert_eq!(r2[0].group_idx, r1[2].group_idx); + + allocator.release_allocation(al1); + // pools [[1 1] [0.6 1] [1 1] + + let rq3 = ResBuilder::default() + .add_compact(0, ResourceAmount::new(2, 8000)) + .finish_v(); + let al3 = allocator.try_allocate(&rq3).unwrap().0; + // pools [[1 1] [0.6 0.2] [0 0] + + let r3 = &al3.resources[0].indices; + assert_eq!(r3.len(), 3); + assert_eq!(r3[0].group_idx, r3[1].group_idx); + assert_ne!(r3[0].group_idx, r3[2].group_idx); + assert_ne!(r3[0].group_idx, r2[0].group_idx); + assert_ne!(r3[1].group_idx, r2[0].group_idx); + assert_eq!(r3[2].group_idx, r2[0].group_idx); + assert_ne!(r3[2].index, r2[0].index); + + let rq4 = ResBuilder::default() + .add_compact(0, ResourceAmount::new(0, 7000)) + .finish_v(); + let al4 = allocator.try_allocate(&rq4).unwrap().0; + // pools [[1 0.3] [0.6 0.2] [0 0] + allocator.validate(); + + allocator.release_allocation(al2); + // pools [[1 0.3] [1 0.2] [0 0] + + let rq6 = ResBuilder::default() + .add_compact(0, ResourceAmount::new(2, 3000)) + .finish_v(); + let al6 = allocator.try_allocate(&rq6).unwrap().0; + let r5 = &al6.resources[0].indices; + assert_eq!(r5.len(), 3); + assert_eq!(r5[0].fractions, 0); + assert_eq!(r5[1].fractions, 0); + assert_eq!(r5[2].fractions, 3000); + assert_eq!(r5[0].group_idx, r5[2].group_idx); + assert_ne!(r5[1].group_idx, r5[0].group_idx); + allocator.validate(); + + allocator.release_allocation(al3); + allocator.release_allocation(al4); + allocator.release_allocation(al6); + + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new_units(6) + ); + } + + #[test] + fn test_allocator_sum_fractions() { + let descriptor = ResourceDescriptor::new(vec![ResourceDescriptorItem { + name: "cpus".to_string(), + kind: ResourceDescriptorKind::Sum { + size: ResourceAmount::new_units(2), + }, + }]); + let mut allocator = test_allocator(&descriptor); + allocator.init_allocator(None); + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(2, 3000)) + .finish_v(); + assert!(allocator.try_allocate(&rq).is_none()); + + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(1, 3000)) + .finish_v(); + let al1 = allocator.try_allocate(&rq).unwrap().0; + assert_eq!(al1.resources[0].indices.len(), 0); + assert_eq!(al1.resources[0].amount, ResourceAmount::new(1, 3000)); + allocator.validate(); + + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(0, 7001)) + .finish_v(); + assert!(allocator.try_allocate(&rq).is_none()); + + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(0, 7000)) + .finish_v(); + let al2 = allocator.try_allocate(&rq).unwrap().0; + assert_eq!(al2.resources[0].indices.len(), 0); + assert_eq!(al2.resources[0].amount, ResourceAmount::new(0, 7000)); + + allocator.release_allocation(al1); + + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(2, 0)) + .finish_v(); + assert!(allocator.try_allocate(&rq).is_none()); + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(1, 3001)) + .finish_v(); + assert!(allocator.try_allocate(&rq).is_none()); + + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(1, 0)) + .finish_v(); + let al3 = allocator.try_allocate(&rq).unwrap().0; + + let rq = ResBuilder::default() + .add(0, ResourceAmount::new(0, 2000)) + .finish_v(); + let al4 = allocator.try_allocate(&rq).unwrap().0; + + allocator.release_allocation(al4); + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new(0, 3000) + ); + allocator.release_allocation(al2); + allocator.release_allocation(al3); + assert_eq!( + allocator.pools[0.into()].concise_state().amount_sum(), + ResourceAmount::new_units(2) + ) + } } diff --git a/crates/tako/src/internal/worker/resources/concise.rs b/crates/tako/src/internal/worker/resources/concise.rs new file mode 100644 index 000000000..516ec456b --- /dev/null +++ b/crates/tako/src/internal/worker/resources/concise.rs @@ -0,0 +1,216 @@ +use crate::internal::common::resources::amount::FRACTIONS_PER_UNIT; +use crate::internal::common::resources::{ResourceId, ResourceVec}; +use crate::resources::{ + Allocation, ResourceAllocation, ResourceAmount, ResourceFractions, ResourceIndex, ResourceUnits, +}; +use crate::Map; +use smallvec::SmallVec; + +#[derive(Debug, Clone, Eq, PartialEq)] +pub(crate) struct ConciseResourceState { + free_units: SmallVec<[ResourceUnits; 1]>, + fractions: Map, +} + +impl ConciseResourceState { + pub fn new( + free_units: SmallVec<[ResourceUnits; 1]>, + fractions: Map, + ) -> Self { + ConciseResourceState { + free_units, + fractions, + } + } + + fn remove_fractions( + &mut self, + group_id: usize, + resource_idx: ResourceIndex, + fractions: ResourceFractions, + ) { + let old_f = self.fractions.entry(resource_idx).or_insert(0); + if *old_f < fractions { + *old_f = FRACTIONS_PER_UNIT + *old_f - fractions; + assert!(self.free_units[group_id] > 0); + self.free_units[group_id] -= 1; + } else { + *old_f -= fractions; + } + } + + pub fn remove(&mut self, resource_allocation: &ResourceAllocation) { + if self.free_units.len() == 1 { + let (units, fractions) = resource_allocation.amount.split(); + assert!(self.free_units[0] >= units); + self.free_units[0] -= units; + if fractions > 0 { + if resource_allocation.indices.is_empty() { + self.remove_fractions(0, ResourceIndex::new(0), fractions); + } else { + for idx in resource_allocation.indices.iter().rev() { + if idx.fractions == 0 { + break; + } + self.remove_fractions(0, idx.index, idx.fractions); + } + } + } + } else { + for idx in &resource_allocation.indices { + if idx.fractions == 0 { + assert!(self.free_units[idx.group_idx as usize] > 0); + self.free_units[idx.group_idx as usize] -= 1 + } else { + self.remove_fractions(idx.group_idx as usize, idx.index, idx.fractions); + } + } + } + } + + fn add_fractions( + &mut self, + group_id: usize, + resource_idx: ResourceIndex, + fractions: ResourceFractions, + ) { + let old_f = self.fractions.entry(resource_idx).or_insert(0); + *old_f += fractions; + if *old_f >= FRACTIONS_PER_UNIT { + *old_f -= FRACTIONS_PER_UNIT; + assert!(*old_f < FRACTIONS_PER_UNIT); + self.free_units[group_id] += 1; + } + } + + pub fn add(&mut self, resource_allocation: &ResourceAllocation) { + if self.free_units.len() == 1 { + let (units, fractions) = resource_allocation.amount.split(); + self.free_units[0] += units; + if fractions > 0 { + if resource_allocation.indices.is_empty() { + self.add_fractions(0, ResourceIndex::new(0), fractions); + } else { + for idx in resource_allocation.indices.iter().rev() { + if idx.fractions == 0 { + break; + } + self.add_fractions(0, idx.index, idx.fractions); + } + } + } + } else { + for idx in &resource_allocation.indices { + if idx.fractions == 0 { + self.free_units[idx.group_idx as usize] += 1 + } else { + self.add_fractions(idx.group_idx as usize, idx.index, idx.fractions); + } + } + } + } + + pub fn n_groups(&self) -> usize { + self.free_units.len() + } + + pub fn groups(&self) -> &[ResourceUnits] { + &self.free_units + } + + pub fn amount_max_alloc(&self) -> ResourceAmount { + let units = self.free_units.iter().sum(); + let fractions = *self.fractions.values().max().unwrap_or(&0); + ResourceAmount::new(units, fractions) + } + + #[cfg(test)] + pub fn amount_sum(&self) -> ResourceAmount { + let units = ResourceAmount::new_units(self.free_units.iter().sum()); + let fractions = self + .fractions + .values() + .map(|f| ResourceAmount::new_fractions(*f)) + .sum(); + units + fractions + } + + pub(crate) fn strip_zeros(&self) -> ConciseResourceState { + ConciseResourceState::new( + self.free_units.clone(), + self.fractions + .iter() + .filter_map(|(k, v)| if *v > 0 { Some((*k, *v)) } else { None }) + .collect(), + ) + } +} + +#[derive(Debug, Clone)] +#[cfg_attr(test, derive(Eq, PartialEq))] +pub(crate) struct ConciseFreeResources { + resources: ResourceVec, +} + +impl ConciseFreeResources { + pub fn new(resources: ResourceVec) -> Self { + ConciseFreeResources { resources } + } + + pub fn add(&mut self, allocation: &Allocation) { + for ra in &allocation.resources { + self.resources[ra.resource_id].add(ra); + } + } + + pub fn remove(&mut self, allocation: &Allocation) { + for ra in &allocation.resources { + self.resources[ra.resource_id].remove(ra); + } + } + + pub fn n_resources(&self) -> usize { + self.resources.len() + } + + pub fn get(&self, resource_id: ResourceId) -> &ConciseResourceState { + &self.resources[resource_id] + } + + /*pub fn get_mut(&mut self, resource_id: ResourceId) -> &mut ConciseResourceState { + &mut self.resources[resource_id] + }*/ + + pub fn all_states(&self) -> &[ConciseResourceState] { + &self.resources + } +} + +#[cfg(test)] +mod tests { + use crate::internal::worker::resources::concise::{ConciseFreeResources, ConciseResourceState}; + use crate::resources::ResourceUnits; + use crate::Map; + + impl ConciseFreeResources { + pub fn new_simple(counts: &[ResourceUnits]) -> Self { + Self::new( + counts + .iter() + .map(|c| ConciseResourceState::new_simple(&[*c])) + .collect::>() + .into(), + ) + } + + pub fn assert_eq(&self, counts: &[ResourceUnits]) { + assert_eq!(self, &Self::new_simple(counts)); + } + } + + impl ConciseResourceState { + pub fn new_simple(free_units: &[ResourceUnits]) -> Self { + ConciseResourceState::new(free_units.into(), Map::new()) + } + } +} diff --git a/crates/tako/src/internal/worker/resources/counts.rs b/crates/tako/src/internal/worker/resources/counts.rs deleted file mode 100644 index 1fba72df0..000000000 --- a/crates/tako/src/internal/worker/resources/counts.rs +++ /dev/null @@ -1,166 +0,0 @@ -use crate::internal::common::resources::{ResourceId, ResourceVec}; -use crate::resources::ResourceAmount; -use smallvec::{smallvec, SmallVec}; - -/// Resource counts of a one resource_id -/// In most cases len() == 1 -/// len() > 1 only if resource is group resource -pub type ResourceCount = SmallVec<[ResourceAmount; 1]>; - -/// Resource counts per resource_id -#[derive(Default, PartialEq, Eq, Debug, Clone)] -pub struct ResourceCountVec { - counts: ResourceVec, -} - -pub fn resource_count_add_at(count: &mut ResourceCount, group_idx: usize, value: ResourceAmount) { - if group_idx >= count.len() { - count.resize(group_idx + 1, 0); - } - count[group_idx] += value; -} - -fn resource_count_add(first: &mut ResourceCount, other: &ResourceCount) { - if first.len() < other.len() { - first.resize(other.len(), 0); - } - for (t, v) in first.iter_mut().zip(other.iter()) { - *t += v; - } -} - -fn resource_count_remove(first: &mut ResourceCount, other: &ResourceCount) -> bool { - if first.len() < other.len() { - first.resize(other.len(), 0); - } - let mut valid = true; - for (t, v) in first.iter_mut().zip(other.iter()) { - if *t < *v { - *t = 0; - valid = false; - } - *t -= v; - } - valid -} - -impl ResourceCountVec { - pub fn new(counts: ResourceVec) -> Self { - ResourceCountVec { counts } - } - - pub fn set(&mut self, resource_id: ResourceId, count: ResourceCount) { - if resource_id.as_num() as usize >= self.counts.len() { - self.counts - .resize(resource_id.as_num() as usize + 1, smallvec![]) - } - self.counts[resource_id] = count; - } - - pub fn get_mut(&mut self, resource_id: ResourceId) -> &mut ResourceCount { - &mut self.counts[resource_id] - } - - pub fn get(&self, resource_id: ResourceId) -> &ResourceCount { - &self.counts[resource_id] - } - - pub fn len(&self) -> usize { - self.counts.len() - } - - fn match_sizes(&mut self, other: &ResourceCountVec) { - if self.counts.len() < other.counts.len() { - self.counts.resize(other.counts.len(), smallvec![]); - } - } - - pub fn add(&mut self, other: &ResourceCountVec) { - self.match_sizes(other); - for (t, v) in self.counts.iter_mut().zip(other.counts.iter()) { - resource_count_add(t, v); - } - } - - pub fn remove(&mut self, other: &ResourceCountVec) -> bool { - self.match_sizes(other); - let mut valid = true; - for (t, v) in self.counts.iter_mut().zip(other.counts.iter()) { - valid &= resource_count_remove(t, v); - } - valid - } - - pub fn fraction(&self, bounds: &ResourceCountVec) -> f32 { - self.counts - .iter() - .zip(bounds.counts.iter()) - .map(|(x, y)| { - let sy: ResourceAmount = y.iter().sum(); - if sy == 0 { - 0f32 - } else { - let sx: ResourceAmount = x.iter().sum(); - sx as f32 / sy as f32 - } - }) - .sum() - } - - pub fn all_counts(&self) -> &[ResourceCount] { - &self.counts - } - /*pub fn from_request(request: ResourceRequest) -> Self { - ResourceCounter { counts: request } - }*/ -} - -#[cfg(test)] -mod tests { - use crate::internal::worker::resources::counts::ResourceCountVec; - use crate::resources::ResourceAmount; - use smallvec::smallvec; - - impl ResourceCountVec { - pub fn new_simple(counts: &[ResourceAmount]) -> Self { - let counts: Vec<_> = counts.into_iter().map(|x| smallvec![*x]).collect(); - Self::new(counts.into()) - } - - pub fn assert_eq(&self, counts: &[ResourceAmount]) { - assert_eq!(self, &ResourceCountVec::new_simple(counts)); - } - } - - #[test] - fn test_counts_add() { - let counts1 = vec![ - smallvec![], - smallvec![0, 1], - smallvec![1, 0], - smallvec![1024], - ]; - let mut v1 = ResourceCountVec::new(counts1.into()); - - let counts2 = vec![ - smallvec![21], - smallvec![2, 1], - smallvec![0, 3], - smallvec![444], - smallvec![2], - ]; - let v2 = ResourceCountVec::new(counts2.into()); - - let counts3 = vec![ - smallvec![21], - smallvec![2, 2], - smallvec![1, 3], - smallvec![1468], - smallvec![2], - ]; - let v3 = ResourceCountVec::new(counts3.into()); - - v1.add(&v2); - assert_eq!(v1, v3); - } -} diff --git a/crates/tako/src/internal/worker/resources/mod.rs b/crates/tako/src/internal/worker/resources/mod.rs index c1c8338d0..20b8da915 100644 --- a/crates/tako/src/internal/worker/resources/mod.rs +++ b/crates/tako/src/internal/worker/resources/mod.rs @@ -1,4 +1,4 @@ pub mod allocator; -pub mod counts; +pub mod concise; pub mod map; pub mod pool; diff --git a/crates/tako/src/internal/worker/resources/pool.rs b/crates/tako/src/internal/worker/resources/pool.rs index 280bd6489..0411ec798 100644 --- a/crates/tako/src/internal/worker/resources/pool.rs +++ b/crates/tako/src/internal/worker/resources/pool.rs @@ -1,35 +1,39 @@ +use crate::internal::common::resources::allocation::AllocationIndex; +use crate::internal::common::resources::amount::FRACTIONS_PER_UNIT; use crate::internal::common::resources::descriptor::ResourceDescriptorKind; -use crate::internal::common::resources::{ - AllocationValue, ResourceAmount, ResourceId, ResourceIndex, -}; +use crate::internal::common::resources::{ResourceAmount, ResourceId, ResourceIndex}; use crate::internal::common::Map; -use crate::internal::worker::resources::counts::ResourceCount; +use crate::internal::worker::resources::concise::ConciseResourceState; use crate::internal::worker::resources::map::ResourceLabelMap; +use crate::resources::{AllocationRequest, ResourceAllocation, ResourceFractions, ResourceUnits}; use crate::Set; -use smallvec::smallvec; + +use smallvec::{smallvec, SmallVec}; #[derive(Debug)] -pub struct IndicesResourcePool { +pub(crate) struct IndicesResourcePool { full_size: ResourceAmount, indices: Vec, + fractions: Map, } #[derive(Debug)] -pub struct GroupsResourcePool { +pub(crate) struct GroupsResourcePool { full_size: ResourceAmount, indices: Vec>, + fractions: Vec>, min_group_size: ResourceAmount, - reverse_map: Map, + //reverse_map: Map, } #[derive(Debug)] -pub struct SumResourcePool { +pub(crate) struct SumResourcePool { full_size: ResourceAmount, free: ResourceAmount, } #[derive(Debug)] -pub enum ResourcePool { +pub(crate) enum ResourcePool { Empty, Indices(IndicesResourcePool), Groups(GroupsResourcePool), @@ -52,7 +56,8 @@ impl ResourcePool { .expect("Resource label not found") }) .collect(), - full_size: values.len() as ResourceAmount, + full_size: ResourceAmount::new_units(values.len() as ResourceUnits), + fractions: Map::new(), }), ResourceDescriptorKind::Groups { groups } => { let groups: Vec> = groups @@ -71,17 +76,17 @@ impl ResourcePool { ResourcePool::Groups(GroupsResourcePool { indices: groups.clone(), - full_size: groups.iter().map(|g| g.len() as ResourceAmount).sum(), - reverse_map: groups - .iter() - .enumerate() - .flat_map(|(i, group)| group.iter().map(move |idx| (*idx, i))) - .collect(), - min_group_size: groups - .iter() - .map(|g| g.len() as ResourceAmount) - .min() - .unwrap_or(1), + full_size: ResourceAmount::new_units( + groups.iter().map(|g| g.len() as ResourceUnits).sum(), + ), + min_group_size: ResourceAmount::new_units( + groups + .iter() + .map(|g| g.len() as ResourceUnits) + .min() + .unwrap_or(1), + ), + fractions: groups.iter().map(|_| Map::new()).collect(), }) } ResourceDescriptorKind::Range { start, end } => { @@ -89,8 +94,9 @@ impl ResourcePool { .map(|id| id.into()) .collect(); ResourcePool::Indices(IndicesResourcePool { - full_size: indices.len() as ResourceAmount, + full_size: ResourceAmount::new_units(indices.len() as ResourceUnits), indices, + fractions: Map::new(), }) } ResourceDescriptorKind::Sum { size } => ResourcePool::Sum(SumResourcePool { @@ -102,19 +108,19 @@ impl ResourcePool { pub fn full_size(&self) -> ResourceAmount { match self { - ResourcePool::Empty => 0, + ResourcePool::Empty => ResourceAmount::ZERO, ResourcePool::Indices(pool) => pool.full_size, ResourcePool::Groups(pool) => pool.full_size, ResourcePool::Sum(pool) => pool.full_size, } } - pub fn min_group_size(&self) -> ResourceAmount { + pub fn min_group_size(&self) -> ResourceUnits { match self { ResourcePool::Empty => 0, - ResourcePool::Indices(pool) => pool.full_size, - ResourcePool::Groups(pool) => pool.min_group_size, - ResourcePool::Sum(pool) => pool.full_size, + ResourcePool::Indices(pool) => pool.full_size.units(), + ResourcePool::Groups(pool) => pool.min_group_size.units(), + ResourcePool::Sum(pool) => pool.full_size.units(), } } @@ -125,65 +131,324 @@ impl ResourcePool { } } - pub fn count(&self) -> ResourceCount { - match self { - ResourcePool::Empty => smallvec![], - ResourcePool::Indices(pool) => smallvec![pool.indices.len() as ResourceAmount], - ResourcePool::Sum(pool) => smallvec![pool.free], - ResourcePool::Groups(pool) => pool - .indices - .iter() - .map(|g| g.len() as ResourceAmount) - .collect(), + pub(crate) fn concise_state(&self) -> ConciseResourceState { + let (units, fractions) = match self { + ResourcePool::Empty => (smallvec![], Map::new()), + ResourcePool::Indices(pool) => ( + smallvec![pool.indices.len() as ResourceUnits], + pool.fractions.clone(), + ), + ResourcePool::Sum(pool) => { + let mut fractions: Map = Map::new(); + let frac = pool.free.fractions(); + if frac > 0 { + fractions.insert(ResourceIndex::new(0), frac); + } + (smallvec![pool.free.units()], fractions) + } + ResourcePool::Groups(pool) => { + let mut fractions = Map::new(); + for f in &pool.fractions { + fractions.extend(f); + } + ( + pool.indices + .iter() + .map(|g| g.len() as ResourceUnits) + .collect(), + fractions, + ) + } + }; + ConciseResourceState::new(units, fractions) + } + + fn claim_all_from_groups(pool: &mut GroupsResourcePool) -> SmallVec<[AllocationIndex; 1]> { + pool.indices + .iter_mut() + .enumerate() + .flat_map(|(group_id, group)| { + std::mem::take(group) + .into_iter() + .map(move |index| AllocationIndex { + index, + group_idx: group_id as u32, + fractions: 0, + }) + }) + .collect() + } + + fn claim_scatter_from_groups( + amount: ResourceAmount, + pool: &mut GroupsResourcePool, + ) -> SmallVec<[AllocationIndex; 1]> { + let mut indices: SmallVec<[AllocationIndex; 1]> = Default::default(); + let (mut units, mut fractions) = amount.split(); + + let mut group_idx = 0; + while units > 0 || fractions > 0 { + if units > 0 { + if let Some(index) = pool.indices[group_idx].pop() { + units -= 1; + indices.push(AllocationIndex { + index, + group_idx: group_idx as u32, + fractions: 0, + }) + } + } else if let Some((index, f)) = + Self::best_fraction_match(&mut pool.fractions[group_idx], fractions) + { + *f -= fractions; + indices.push(AllocationIndex { + index: *index, + group_idx: group_idx as u32, + fractions, + }); + fractions = 0; + } else if let Some(index) = pool.indices[group_idx].pop() { + pool.fractions[group_idx].insert(index, FRACTIONS_PER_UNIT - fractions); + indices.push(AllocationIndex { + index, + group_idx: group_idx as u32, + fractions, + }); + fractions = 0; + } + group_idx = (group_idx + 1) % pool.indices.len(); } + indices } - pub fn claim_resources(&mut self, amount: &ResourceCount) -> AllocationValue { - match self { - ResourcePool::Empty => unreachable!(), - ResourcePool::Indices(ref mut pool) => { - assert_eq!(amount.len(), 1); - AllocationValue::new_indices( - (0..amount[0]) - .map(|_| pool.indices.pop().unwrap()) - .collect(), + fn claim_compact_from_groups( + amount: ResourceAmount, + pool: &mut GroupsResourcePool, + ) -> SmallVec<[AllocationIndex; 1]> { + let mut indices = Default::default(); + let mut remaining = amount; + let mut fraction_idx: Option = None; + + let mut amounts: Vec<_> = pool + .indices + .iter() + .zip(pool.fractions.iter()) + .map(|(i, f)| { + ResourceAmount::new( + i.len() as ResourceUnits, + f.values().max().copied().unwrap_or(0), ) + }) + .collect(); + + loop { + if let Some((group_idx, _a)) = amounts + .iter() + .enumerate() + .filter(|(_i, a)| **a >= remaining) + .min_by_key(|(_i, a)| **a) + { + let (units, fractions) = remaining.split(); + Self::take_indices( + &mut pool.indices[group_idx], + group_idx as u32, + units, + &mut indices, + ); + if fractions > 0 { + indices.push( + if let Some((index, f)) = + Self::best_fraction_match(&mut pool.fractions[group_idx], fractions) + { + *f -= fractions; + AllocationIndex { + index: *index, + group_idx: group_idx as u32, + fractions, + } + } else { + let index = pool.indices[group_idx].pop().unwrap(); + pool.fractions[group_idx].insert(index, FRACTIONS_PER_UNIT - fractions); + AllocationIndex { + index, + group_idx: group_idx as u32, + fractions, + } + }, + ) + } + break; + } else { + let (group_idx, amount) = amounts + .iter_mut() + .enumerate() + .max_by_key(|(_i, a)| **a) + .unwrap(); + *amount = ResourceAmount::ZERO; + let (mut units, mut fractions) = remaining.split(); + let size = pool.indices[group_idx].len() as ResourceUnits; + units -= size; + Self::take_indices( + &mut pool.indices[group_idx], + group_idx as u32, + size, + &mut indices, + ); + if fractions > 0 { + if let Some((index, f)) = + Self::best_fraction_match(&mut pool.fractions[group_idx], fractions) + { + *f -= fractions; + fraction_idx = Some(indices.len()); + indices.push(AllocationIndex { + index: *index, + group_idx: group_idx as u32, + fractions, + }); + fractions = 0; + } + } + remaining = ResourceAmount::new(units, fractions) } - ResourcePool::Sum(ref mut pool) => { - assert_eq!(amount.len(), 1); - assert!(pool.free >= amount[0]); - pool.free -= amount[0]; - AllocationValue::new_sum(amount[0]) + } + + if let Some(allocation_idx) = fraction_idx { + let last = indices.len() - 1; + indices.swap(allocation_idx, last); + } + + indices + } + + pub fn take_indices( + pool_indices: &mut Vec, + group_id: u32, + units: ResourceUnits, + out: &mut SmallVec<[AllocationIndex; 1]>, + ) { + (0..units).for_each(|_| { + out.push(AllocationIndex { + index: pool_indices.pop().unwrap(), + group_idx: group_id, + fractions: 0, + }) + }) + } + + fn best_fraction_match( + frac_map: &mut Map, + fractions: ResourceFractions, + ) -> Option<(&ResourceIndex, &mut ResourceFractions)> { + frac_map + .iter_mut() + .filter(|(_, f)| **f >= fractions) + .min_by_key(|(_, f)| **f) + } + + pub(crate) fn claim_resources( + &mut self, + resource_id: ResourceId, + policy: &AllocationRequest, + ) -> ResourceAllocation { + let (amount, indices) = match self { + ResourcePool::Empty => unreachable!(), + ResourcePool::Indices(pool) => { + let amount = policy.amount(pool.full_size); + let (units, fractions) = amount.split(); + let mut indices = Default::default(); + Self::take_indices(&mut pool.indices, 0, units, &mut indices); + if fractions > 0 { + if let Some((r_idx, f)) = + Self::best_fraction_match(&mut pool.fractions, fractions) + { + let index = *r_idx; + *f -= fractions; + indices.push(AllocationIndex { + index, + group_idx: 0, + fractions, + }); + } else { + let index = pool.indices.pop().unwrap(); + indices.push(AllocationIndex { + index, + group_idx: 0, + fractions, + }); + pool.fractions.insert(index, FRACTIONS_PER_UNIT - fractions); + } + } + (amount, indices) } ResourcePool::Groups(pool) => { - let mut result = smallvec![]; - assert!(pool.indices.len() >= amount.len()); - for (g, a) in pool.indices.iter_mut().zip(amount) { - for _ in 0..*a { - result.push(g.pop().unwrap()) + match policy { + AllocationRequest::Compact(amount) + | AllocationRequest::ForceCompact(amount) => { + // We do neet need to distinguish between force compact and compact + // because we already know that here that allocation is possible + (*amount, Self::claim_compact_from_groups(*amount, pool)) + } + AllocationRequest::Scatter(amount) => { + (*amount, Self::claim_scatter_from_groups(*amount, pool)) } + AllocationRequest::All => (pool.full_size, Self::claim_all_from_groups(pool)), } - AllocationValue::new_indices(result) } + ResourcePool::Sum(pool) => { + let amount = policy.amount(pool.full_size); + pool.free -= amount; + (amount, Default::default()) + } + }; + ResourceAllocation { + resource_id, + amount, + indices, } } - pub fn release_allocation(&mut self, allocation: AllocationValue) { + pub fn release_allocation(&mut self, allocation: &ResourceAllocation) { match self { ResourcePool::Empty => unreachable!(), - ResourcePool::Indices(pool) => pool - .indices - .extend_from_slice(allocation.indices().unwrap()), - ResourcePool::Sum(pool) => match allocation { - AllocationValue::Sum(amount) => { - pool.free += amount; - assert!(pool.free <= pool.full_size); + ResourcePool::Indices(pool) => { + for index in allocation.indices.iter().rev() { + // Iterating reversly as indices was originaly taken by pop() + // Just to add small determinism, we return then in same order + assert_eq!(index.group_idx, 0); + if index.fractions == 0 { + pool.indices.push(index.index); + } else { + let f = pool.fractions.get_mut(&index.index).unwrap(); + *f += index.fractions; + if *f == FRACTIONS_PER_UNIT { + pool.fractions.remove(&index.index); + pool.indices.push(index.index); + } + } } - _ => unreachable!(), - }, + //pool.indices.extend(&allocation.resource_indices()) + } + ResourcePool::Sum(pool) => { + pool.free += allocation.amount; + assert!(pool.free <= pool.full_size); + assert!(allocation.indices.is_empty()); + } ResourcePool::Groups(pool) => { - for idx in allocation.indices().unwrap().iter() { - pool.indices[pool.reverse_map[idx]].push(*idx); + for index in allocation.indices.iter().rev() { + // Iterating reversly as indices was originaly taken by pop() + // Just to add small determinism, we return then in same order + if index.fractions == 0 { + pool.indices[index.group_idx as usize].push(index.index); + } else { + let f = pool.fractions[index.group_idx as usize] + .get_mut(&index.index) + .unwrap(); + *f += index.fractions; + if *f == FRACTIONS_PER_UNIT { + pool.fractions[index.group_idx as usize].remove(&index.index); + pool.indices[index.group_idx as usize].push(index.index); + } + } } } }; @@ -201,12 +466,29 @@ impl ResourcePool { Set::from_iter(pool.indices.iter()).len(), pool.indices.len() ); - assert!(pool.indices.len() <= pool.full_size as usize) + assert!(pool.indices.len() <= pool.full_size.units() as usize); + for index in &pool.indices { + assert!(!pool.fractions.contains_key(index)); + } + for f in pool.fractions.values() { + //assert!(*f > 0); + assert!(*f < FRACTIONS_PER_UNIT); + } } ResourcePool::Groups(pool) => { let sum: usize = pool.indices.iter().map(|x| x.len()).sum(); assert_eq!(Set::from_iter(pool.indices.iter().flatten()).len(), sum); - assert!(sum <= pool.full_size as usize) + assert!(sum <= pool.full_size.units() as usize); + for index in pool.indices.iter().flatten() { + for f in &pool.fractions { + assert!(!f.contains_key(index)); + } + } + for f in &pool.fractions { + for ff in f.values() { + assert!(*ff < FRACTIONS_PER_UNIT); + } + } } ResourcePool::Sum(pool) => { assert!(pool.free <= pool.full_size) @@ -219,16 +501,18 @@ impl ResourcePool { mod tests { use crate::internal::common::resources::ResourceAllocation; use crate::internal::worker::resources::pool::ResourcePool; - use crate::resources::ResourceAmount; + use crate::resources::{ResourceAmount, ResourceUnits}; impl ResourcePool { pub fn current_free(&self) -> ResourceAmount { match self { - ResourcePool::Empty => 0, - ResourcePool::Indices(pool) => pool.indices.len() as ResourceAmount, - ResourcePool::Groups(pool) => { - pool.indices.iter().map(|g| g.len() as ResourceAmount).sum() + ResourcePool::Empty => ResourceAmount::ZERO, + ResourcePool::Indices(pool) => { + ResourceAmount::new_units(pool.indices.len() as ResourceUnits) } + ResourcePool::Groups(pool) => ResourceAmount::new_units( + pool.indices.iter().map(|g| g.len()).sum::() as ResourceUnits, + ), ResourcePool::Sum(pool) => pool.free, } } @@ -244,13 +528,11 @@ mod tests { ResourcePool::Indices(_) => { vec![0] } - ResourcePool::Groups(pool) => { + ResourcePool::Groups(_pool) => { let mut result: Vec<_> = allocation - .value - .indices() - .unwrap() + .indices .iter() - .map(|idx| *pool.reverse_map.get(idx).unwrap()) + .map(|idx| idx.group_idx as usize) .collect(); result.sort(); result.dedup(); diff --git a/crates/tako/src/internal/worker/rpc.rs b/crates/tako/src/internal/worker/rpc.rs index 5553c6bff..2b84f4b7d 100644 --- a/crates/tako/src/internal/worker/rpc.rs +++ b/crates/tako/src/internal/worker/rpc.rs @@ -15,11 +15,11 @@ use tokio::time::timeout; use crate::comm::{ConnectionRegistration, RegisterWorker}; use crate::hwstats::WorkerHwStateMessage; use crate::internal::common::resources::map::ResourceMap; -use crate::internal::common::resources::{Allocation, AllocationValue}; +use crate::internal::common::resources::Allocation; use crate::internal::common::WrappedRcRefCell; use crate::internal::messages::worker::{ - FromWorkerMessage, StealResponseMsg, TaskResourceAllocation, TaskResourceAllocationValue, - ToWorkerMessage, WorkerOverview, WorkerRegistrationResponse, WorkerStopReason, + FromWorkerMessage, StealResponseMsg, TaskResourceAllocation, ToWorkerMessage, WorkerOverview, + WorkerRegistrationResponse, WorkerStopReason, }; use crate::internal::server::rpc::ConnectionDescriptor; use crate::internal::transfer::auth::{ @@ -518,15 +518,15 @@ fn resource_allocation_to_msg( .map( |alloc| crate::internal::messages::worker::ResourceAllocation { resource: resource_map - .get_name(alloc.resource) + .get_name(alloc.resource_id) .unwrap_or("unknown") .to_string(), - value: match &alloc.value { - AllocationValue::Indices(indices) => { - TaskResourceAllocationValue::Indices(indices.iter().cloned().collect()) - } - AllocationValue::Sum(amount) => TaskResourceAllocationValue::Sum(*amount), - }, + amount: alloc.amount, + indices: alloc + .indices + .iter() + .map(|i| (i.index, i.fractions)) + .collect(), }, ) .collect(), diff --git a/crates/tako/src/internal/worker/rqueue.rs b/crates/tako/src/internal/worker/rqueue.rs index 94a522f29..5d9fceb31 100644 --- a/crates/tako/src/internal/worker/rqueue.rs +++ b/crates/tako/src/internal/worker/rqueue.rs @@ -6,6 +6,7 @@ use crate::internal::worker::state::TaskMap; use crate::internal::worker::task::Task; use crate::{Priority, PriorityTuple, Set, TaskId, WorkerId}; use priority_queue::PriorityQueue; +use std::rc::Rc; use std::time::Duration; type QueuePriorityTuple = (Priority, Priority, Priority); // user priority, resource priority, scheduler priority @@ -72,7 +73,7 @@ impl ResourceWaitQueue { p } - pub fn release_allocation(&mut self, allocation: Allocation) { + pub fn release_allocation(&mut self, allocation: Rc) { self.allocator.release_allocation(allocation); } @@ -137,7 +138,7 @@ impl ResourceWaitQueue { &mut self, task_map: &TaskMap, remaining_time: Option, - ) -> Vec<(TaskId, Allocation, usize)> { + ) -> Vec<(TaskId, Rc, usize)> { self.allocator.init_allocator(remaining_time); let mut out = Vec::new(); while !self.try_start_tasks_helper(task_map, &mut out) { @@ -149,7 +150,7 @@ impl ResourceWaitQueue { fn try_start_tasks_helper( &mut self, _task_map: &TaskMap, - out: &mut Vec<(TaskId, Allocation, usize)>, + out: &mut Vec<(TaskId, Rc, usize)>, ) -> bool { let current_priority: QueuePriorityTuple = if let Some(Some(priority)) = self.queues.values().map(|qfr| qfr.current_priority()).max() @@ -186,10 +187,11 @@ mod tests { use crate::internal::common::resources::{ ResourceDescriptor, ResourceRequest, ResourceRequestVariants, }; - use crate::internal::tests::utils::resources::ResBuilder; use crate::internal::tests::utils::resources::{cpus_compact, ResourceRequestBuilder}; + use crate::internal::tests::utils::resources::{ra_builder, ResBuilder}; use crate::internal::worker::rqueue::ResourceWaitQueue; use crate::internal::worker::test_util::{worker_task, WorkerTaskBuilder}; + use std::ops::Deref; use std::time::Duration; use crate::internal::messages::worker::WorkerResourceCounts; @@ -465,7 +467,7 @@ mod tests { rq.new_worker( 400.into(), WorkerResources::from_transport(WorkerResourceCounts { - n_resources: vec![2, 0], + n_resources: ra_builder(&[2, 0]).deref().clone(), }), ); @@ -476,7 +478,7 @@ mod tests { rq.new_worker( 401.into(), WorkerResources::from_transport(WorkerResourceCounts { - n_resources: vec![2, 2], + n_resources: ra_builder(&[2, 2]).deref().clone(), }), ); assert_eq!(rq.resource_priority(&rq1), 0); @@ -487,7 +489,7 @@ mod tests { rq.new_worker( WorkerId::new(i), WorkerResources::from_transport(WorkerResourceCounts { - n_resources: vec![3, 0], + n_resources: ra_builder(&[3, 0]).deref().clone(), }), ); } @@ -540,7 +542,7 @@ mod tests { rq.new_worker( 400.into(), WorkerResources::from_transport(WorkerResourceCounts { - n_resources: vec![16, 2, 0, 1], + n_resources: ra_builder(&[16, 2, 0, 1]).deref().clone(), }), ); @@ -573,7 +575,7 @@ mod tests { rq.new_worker( 400.into(), WorkerResources::from_transport(WorkerResourceCounts { - n_resources: vec![16, 2, 0, 1], + n_resources: ra_builder(&[16, 2, 0, 1]).deref().clone(), }), ); @@ -606,7 +608,7 @@ mod tests { rq.new_worker( 400.into(), WorkerResources::from_transport(WorkerResourceCounts { - n_resources: vec![16, 2, 0, 1], + n_resources: ra_builder(&[16, 2, 0, 1]).deref().clone(), }), ); diff --git a/crates/tako/src/internal/worker/state.rs b/crates/tako/src/internal/worker/state.rs index cdd543bbc..50d2be3d3 100644 --- a/crates/tako/src/internal/worker/state.rs +++ b/crates/tako/src/internal/worker/state.rs @@ -1,3 +1,4 @@ +use std::rc::Rc; use std::sync::Arc; use std::time::Instant; @@ -216,7 +217,7 @@ impl WorkerState { self.comm.notify_start_task(); } - pub fn start_task(&mut self, task_id: TaskId, task_env: TaskEnv, allocation: Allocation) { + pub fn start_task(&mut self, task_id: TaskId, task_env: TaskEnv, allocation: Rc) { let task = self.get_task_mut(task_id); task.state = TaskState::Running(task_env, allocation); self.running_tasks.insert(task_id); diff --git a/crates/tako/src/internal/worker/task.rs b/crates/tako/src/internal/worker/task.rs index 92bf281e2..7ff825d91 100644 --- a/crates/tako/src/internal/worker/task.rs +++ b/crates/tako/src/internal/worker/task.rs @@ -3,11 +3,12 @@ use crate::internal::common::stablemap::ExtractKey; use crate::internal::messages::worker::ComputeTaskMsg; use crate::internal::worker::taskenv::TaskEnv; use crate::{InstanceId, Priority, TaskId, WorkerId}; +use std::rc::Rc; use std::time::Duration; pub enum TaskState { Waiting(u32), - Running(TaskEnv, Allocation), + Running(TaskEnv, Rc), } pub struct Task { diff --git a/crates/tako/src/internal/worker/test_util.rs b/crates/tako/src/internal/worker/test_util.rs index 832f87418..d8670a0d2 100644 --- a/crates/tako/src/internal/worker/test_util.rs +++ b/crates/tako/src/internal/worker/test_util.rs @@ -8,6 +8,7 @@ use crate::internal::worker::state::TaskMap; use crate::internal::worker::task::Task; use crate::{InstanceId, Priority, TaskId, WorkerId}; use smallvec::smallvec; +use std::rc::Rc; use std::time::Duration; pub struct WorkerTaskBuilder { @@ -97,7 +98,7 @@ impl ResourceQueueBuilder { self.queue.new_worker(worker_id, wr); } - pub fn start_tasks(&mut self) -> Map { + pub fn start_tasks(&mut self) -> Map> { self.queue .try_start_tasks(&self.task_map, None) .into_iter() @@ -105,7 +106,7 @@ impl ResourceQueueBuilder { .collect() } - pub fn start_tasks_duration(&mut self, duration: Duration) -> Map { + pub fn start_tasks_duration(&mut self, duration: Duration) -> Map> { self.queue .try_start_tasks(&self.task_map, Some(duration)) .into_iter() diff --git a/crates/tako/src/lib.rs b/crates/tako/src/lib.rs index 78b70d607..26a755420 100644 --- a/crates/tako/src/lib.rs +++ b/crates/tako/src/lib.rs @@ -28,16 +28,21 @@ pub type Result = std::result::Result; pub mod resources { pub use crate::internal::common::resources::{ - Allocation, AllocationRequest, AllocationValue, NumOfNodes, ResourceAllocation, - ResourceAmount, ResourceDescriptor, ResourceDescriptorItem, ResourceDescriptorKind, + Allocation, AllocationRequest, NumOfNodes, ResourceAllocation, ResourceAmount, + ResourceDescriptor, ResourceDescriptorItem, ResourceDescriptorKind, ResourceFractions, ResourceIndex, ResourceLabel, ResourceRequest, ResourceRequestEntries, - ResourceRequestEntry, ResourceRequestVariants, TimeRequest, AMD_GPU_RESOURCE_NAME, - CPU_RESOURCE_ID, CPU_RESOURCE_NAME, MEM_RESOURCE_NAME, NVIDIA_GPU_RESOURCE_NAME, + ResourceRequestEntry, ResourceRequestVariants, ResourceUnits, TimeRequest, + AMD_GPU_RESOURCE_NAME, CPU_RESOURCE_ID, CPU_RESOURCE_NAME, MEM_RESOURCE_NAME, + NVIDIA_GPU_RESOURCE_NAME, }; pub use crate::internal::common::resources::map::ResourceMap; pub use crate::internal::common::resources::descriptor::DescriptorError; + + pub use crate::internal::common::resources::amount::{ + FRACTIONS_MAX_DIGITS, FRACTIONS_PER_UNIT, + }; } pub mod server { diff --git a/docs/jobs/cresources.md b/docs/jobs/cresources.md index 1f48cb3d4..7cb364671 100644 --- a/docs/jobs/cresources.md +++ b/docs/jobs/cresources.md @@ -53,6 +53,7 @@ The following variables are created when a task is executed: * ``HQ_PIN`` - Is set to `taskset` or `omp` (depending on the used pin mode) if the task was pinned by HyperQueue (see below). * ``NUM_OMP_THREADS`` -- Set to number of cores assigned for task. (For compatibility with OpenMP). + This option is not set when you ask for a non-integer number of CPUs. ## Pinning diff --git a/docs/jobs/jobfile.md b/docs/jobs/jobfile.md index a179ae192..f255101fe 100644 --- a/docs/jobs/jobfile.md +++ b/docs/jobs/jobfile.md @@ -172,4 +172,12 @@ In the case that many tasks with such a configuration are submitted to a worker and 3 tasks in the second one. For a task with resource variants, HyperQueue sets variable `HQ_RESOURCE_VARIANT` -to an index of chosen variant (counted from 0) when a task is started. \ No newline at end of file +to an index of chosen variant (counted from 0) when a task is started. + + +## Non-integer resource amounts + +You may specify a resource number as float, e.g. `resources = { "foo" = 1.5 }`. +It is valid but internally the type if converted to float, that may for some numbers lead to +a rounding up when number is converted to 4-digit precision of resource amounts. +If you want to avoid this, put the number into parentheses, e.g. `resources = { "foo" = "1.5" }`. \ No newline at end of file diff --git a/docs/jobs/resources.md b/docs/jobs/resources.md index 612b7ec5e..7c8e6d7d6 100644 --- a/docs/jobs/resources.md +++ b/docs/jobs/resources.md @@ -120,7 +120,7 @@ The following resources are detected automatically if a resource of a given name # The worker will have resource gpus/nvidia=[2,3] ``` -* RAM of the node is detected as resource "mem" in bytes. +* RAM of the node is detected as resource "mem" in megabytes; i.e. `--resource mem=100` asks for 100 MiBs of the memory. If you want to see how is your system seen by a worker without actually starting it, you can start: @@ -141,7 +141,7 @@ When you submit a job, you can define a **resource requests** with the `--resour $ hq submit --resource = --resource = ... ``` -Where `NAME` is a name of the requested resource and the `AMOUNT` is a positive integer defining the +Where `NAME` is a name of the requested resource and the `AMOUNT` is a positive number defining the size of the request. Tasks with such resource requests will only be executed on workers that fulfill all the specified @@ -211,6 +211,32 @@ When strategy is not defined then ``compact`` is used as default. ```console $ hq submit --resource cpus="8 scatter" ... ``` + +### Non-integer allocation of resources + +Amount of the resource may be a non-integer number. +E.g. you may ask for 0.5 of a resource. It tells the scheduler that you want to utilize only half of the resource +and if another process asks for at most 0.5 of the resource, it may get the same resource. +This resource sharing is done on logical of HyperQueue and actual resource sharing is up to tasks. + +The precision for defining amount is four decimal places. Therefore, the minimal resource amount that you +can ask for is `0.0001`. + +For sum resources, the amount is simply removed from the pool as in the case of integer resources. + +In the case of an indexed resource, the partial resource is always taken from a single index. +It means that if there is an indexed resource with two indices that are both utilized on 0.75, +then a task that ask for 0.5 of this resource will not be started, despite there is available 0.5 of the resource in total, +because there is no single index that is free at least on 0.5. + +If non-integer is bigger than 1, than integer part is always satisfied as whole indices and rest is a part of another index. +E.g. when you ask for 2.5 of an indexed resource, you will get 2 complete indices and one index allocated on 50%. + + +!!! note + + In the current version, policy "compact!" is not allowed with non-integer amounts. + ### Resource environment variables When a task that has resource requests is executed, the following variables are passed to it for @@ -218,7 +244,8 @@ each resource request named ``: * `HQ_RESOURCE_REQUEST_` contains the amount of requested resources. * `HQ_RESOURCE_VALUES_` contains the specific resource values allocated for the task as a -comma-separated list. This variable is only filled for indexed resource pool. +comma-separated list. This variable is only filled for an indexed resource pool. +In case of non-integer amount, the partially allocated index is always the last index. The slash symbol (`/`) in resource name is normalized to underscore (`_`) when being used in the environment variable name. diff --git a/tests/output/test_json.py b/tests/output/test_json.py index de4a3771f..a23ffb05a 100644 --- a/tests/output/test_json.py +++ b/tests/output/test_json.py @@ -195,11 +195,11 @@ def test_print_job_detail_resources(hq_env: HqEnv): "min_time": 0.0, "n_nodes": 0, "resources": [ - {"request": {"Compact": 2}, "resource": "res1"}, - {"request": {"ForceCompact": 8}, "resource": "res2"}, + {"request": {"Compact": 2 * 10_000}, "resource": "res1"}, + {"request": {"ForceCompact": 8 * 10_000}, "resource": "res2"}, {"request": "All", "resource": "res3"}, - {"request": {"Scatter": 4}, "resource": "res4"}, - {"request": {"Compact": 1}, "resource": "cpus"}, + {"request": {"Scatter": 4 * 10_000}, "resource": "res4"}, + {"request": {"Compact": 1 * 10_000}, "resource": "cpus"}, ], } ] diff --git a/tests/pyapi/test_job.py b/tests/pyapi/test_job.py index 409ad38b5..cb6b7ab8f 100644 --- a/tests/pyapi/test_job.py +++ b/tests/pyapi/test_job.py @@ -231,13 +231,13 @@ def test_resources_task(hq_env: HqEnv): job.program( args=bash("echo Hello"), - resources=ResourceRequest(cpus="2", resources={"res0": 1}), + resources=ResourceRequest(cpus="2", resources={"res0": 1, "res1": 105.1}), ) submitted_job = client.submit(job) wait_for_job_state(hq_env, 1, "WAITING") table = hq_env.command(["task", "info", str(submitted_job.id), "0"], as_table=True) - table.check_row_value("Resources", "cpus: 2 compact\nres0: 1 compact") + table.check_row_value("Resources", "cpus: 2 compact\nres0: 1 compact\nres1: 105.1 compact") def test_priority_task(hq_env: HqEnv): diff --git a/tests/test_cpus.py b/tests/test_cpus.py index 57791f9f1..f7af18679 100644 --- a/tests/test_cpus.py +++ b/tests/test_cpus.py @@ -69,7 +69,7 @@ def test_job_num_of_cpus(hq_env: HqEnv): table = hq_env.command(["job", "info", "6"], as_table=True) table.check_row_value("Resources", "cpus: all") lst = read_list(default_task_output(job_id=6)) - assert list(range(12)) == lst + assert list(range(12)) == sorted(lst) def test_set_omp_num_threads(hq_env: HqEnv): diff --git a/tests/test_jobfile.py b/tests/test_jobfile.py index 434936644..c7be6d4cd 100644 --- a/tests/test_jobfile.py +++ b/tests/test_jobfile.py @@ -53,6 +53,12 @@ def test_job_file_submit_maximal(hq_env: HqEnv, tmp_path): id = 13 pin = "omp" command = ["sleep", "0"] + +[[task]] +id = 14 +command = ["sleep", "0"] +[[task.request]] +resources = { "gpus" = 1.1 } """) hq_env.command(["job", "submit-file", "job.toml"]) wait_for_job_state(hq_env, 1, "FINISHED") @@ -90,6 +96,9 @@ def test_job_file_submit_maximal(hq_env: HqEnv, tmp_path): table = hq_env.command(["task", "info", "1", "13"], as_table=True) table.check_row_value("Pin", "omp") + table = hq_env.command(["task", "info", "1", "14"], as_table=True) + table.check_row_value("Resources", "gpus: 1.1 compact") + def test_job_file_resource_variants1(hq_env: HqEnv, tmp_path): hq_env.start_server() diff --git a/tests/test_resources.py b/tests/test_resources.py index 25566ed54..609c7e0b2 100644 --- a/tests/test_resources.py +++ b/tests/test_resources.py @@ -1,3 +1,4 @@ +import collections import multiprocessing import time @@ -21,14 +22,16 @@ def test_worker_resources_display(hq_env: HqEnv): "fairy=sum(1000_1000)", "--resource", "shark=[1,3,5,2]", + "--resource", + "small=sum(0.5)", ], ) table = hq_env.command(["worker", "list"], as_table=True) - assert table.get_column_value("Resources") == ["cpus 4x2; fairy 10001000; potato 12; shark 4"] + assert table.get_column_value("Resources") == ["cpus 4x2; fairy 10001000; potato 12; shark 4; small 0.5"] table = hq_env.command(["worker", "info", "1"], as_table=True) print(table.get_row_value("Resources")) - assert table.get_row_value("Resources") == "cpus: 4x2\nfairy: 10001000\npotato: 12\nshark: 4" + assert table.get_row_value("Resources") == "cpus: 4x2\nfairy: 10001000\npotato: 12\nshark: 4\nsmall: 0.5" def test_task_resources_ignore_worker_without_resource(hq_env: HqEnv): @@ -130,6 +133,141 @@ def test_task_resources_range_multiple_allocated_values(hq_env: HqEnv): assert len(set(all_values)) == 6 +def test_task_resource_fractions_sum(hq_env: HqEnv): + hq_env.start_server() + hq_env.start_worker(cpus=4, args=["--resource", "foo=sum(2)"]) + + hq_env.command( + [ + "submit", + "--array=1-4", + "--resource", + "foo=0.5", + "--", + "bash", + "-c", + "date +%s%N\nsleep 1", + ] + ) + wait_for_job_state(hq_env, 1, "FINISHED") + times = [] + for task_id in range(1, 5): + with open(default_task_output(task_id=task_id)) as f: + times.append(int(f.readline().rstrip())) + times.sort() + assert ((times[-1] - times[0]) / 1000_000_000) < 0.15 + + +def test_task_resource_fractions_scatter(hq_env: HqEnv): + hq_env.start_server() + hq_env.start_worker(cpus=4, args=["--resource", "foo=[[a,b,c],[i, j, k],[x, y],[v, w]]"]) + + hq_env.command( + [ + "submit", + "--resource", + "foo=3.5", + "--", + "bash", + "-c", + "echo $HQ_RESOURCE_VALUES_foo", + ] + ) + wait_for_job_state(hq_env, 1, "FINISHED") + with open(default_task_output()) as f: + r = f.readline().rstrip().split(",") + assert len(r) == 4 + assert len(set(r)) == 4 + + +def test_task_resource_fractions_sum2(hq_env: HqEnv): + hq_env.start_server() + hq_env.start_worker(cpus=4, args=["--resource", "foo=sum(2.2)"]) + + hq_env.command( + [ + "submit", + "--array=1-4", + "--resource", + "foo=0.6", + "--", + "bash", + "-c", + "date +%s%N\nsleep 1", + ] + ) + wait_for_job_state(hq_env, 1, "FINISHED") + times = [] + for task_id in range(1, 5): + with open(default_task_output(task_id=task_id)) as f: + times.append(int(f.readline().rstrip())) + times.sort() + assert 0.99 < ((times[-1] - times[0]) / 1000_000_000) < 1.25 + + +def test_task_resource_fractions_sharing_small(hq_env: HqEnv): + hq_env.start_server() + hq_env.start_worker(cpus=4, args=["--resource", "foo=[a,b]"]) + + hq_env.command( + [ + "submit", + "--array=1-4", + "--resource", + "foo=0.4", + "--", + "bash", + "-c", + "date +%s%N\necho $HQ_RESOURCE_VALUES_foo\nsleep 1", + ] + ) + wait_for_job_state(hq_env, 1, "FINISHED") + times = [] + counts = collections.Counter() + for task_id in range(1, 5): + with open(default_task_output(task_id=task_id)) as f: + times.append(int(f.readline().rstrip())) + name = f.readline().strip() + assert name in ["a", "b"] + counts[name] += 1 + times.sort() + assert ((times[-1] - times[0]) / 1000_000_000) < 0.15 + assert counts == collections.Counter({"a": 2, "b": 2}) + + +def test_task_resource_fractions_sharing_larger(hq_env: HqEnv): + hq_env.start_server() + hq_env.start_worker(cpus=4, args=["--resource", "foo=[a0, a1, a2, a3, a4, a5, a6, a7, a8, a9]"]) + + hq_env.command( + [ + "submit", + "--array=1-4", + "--resource", + "foo=2.5", + "--", + "bash", + "-c", + "date +%s%N\necho $HQ_RESOURCE_VALUES_foo\nsleep 1", + ] + ) + wait_for_job_state(hq_env, 1, "FINISHED") + times = [] + counts_frac = collections.Counter() + counts_whole = collections.Counter() + for task_id in range(1, 5): + with open(default_task_output(task_id=task_id)) as f: + times.append(int(f.readline().rstrip())) + names = f.readline().strip().split(",") + for name in names[:-1]: + counts_whole[name] += 1 + counts_frac[names[-1]] += 1 + times.sort() + assert ((times[-1] - times[0]) / 1000_000_000) < 0.15 + assert tuple(counts_frac.values()) == (2, 2) + assert tuple(counts_whole.values()) == (1,) * 8 + + def test_task_resources_allocate_string(hq_env: HqEnv): hq_env.start_server() hq_env.start_worker(args=["--resource", "foo=[a,b,c]"])