xmanager-slurm 0.4.0__tar.gz → 0.4.1__tar.gz

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: xmanager-slurm
-Version: 0.4.0
+Version: 0.4.1
 Summary: Slurm backend for XManager.
 Project-URL: GitHub, https://github.com/jessefarebro/xm-slurm
 Author-email: Jesse Farebrother <jfarebro@cs.mcgill.ca>

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/examples/job-array-sweep/launch.py

@@ -46,7 +46,7 @@ async def main(_):
                 )
             )
 
-        args = [{"args": {"scale": scale}} for scale in range(3)]
+        args = [xm_slurm.JobArgs(args={"scale": scale}) for scale in range(3)]
         wus = await experiment.add(make_job, args)
 
         for wu, status in zip(wus, await asyncio.gather(*[wu.get_status() for wu in wus])):

xmanager_slurm-0.4.1/examples/job-dependencies/eval.py

@@ -0,0 +1,13 @@
+import numpy as np
+from absl import app, flags
+
+INPUT_FILE = flags.DEFINE_string("input_file", "result.npy", "Input file path")
+
+
+def main(_):
+    result = np.load(INPUT_FILE.value)
+    print(f"Received result: {result}")
+
+
+if __name__ == "__main__":
+    app.run(main)

xmanager_slurm-0.4.1/examples/job-dependencies/launch.py

@@ -0,0 +1,104 @@
+import asyncio
+import datetime as dt
+import pathlib
+
+from absl import app
+from xmanager import xm
+
+import xm_slurm
+import xm_slurm.contrib.clusters
+
+
+@xm.run_in_asyncio_loop
+async def main(_):
+    async with xm_slurm.create_experiment("My Experiment") as experiment:
+        # Step 1: Specify the executor specification
+        executor_spec = xm_slurm.Slurm.Spec(tag="ghcr.io/jessefarebro/xm-slurm/test:latest")
+
+        # Step 2: Specify the executable and package it
+        [train_executable, eval_executable] = experiment.package(
+            [
+                xm_slurm.uv_container(
+                    executor_spec=executor_spec,
+                    entrypoint=xm.CommandList(["train.py"]),
+                ),
+                xm_slurm.uv_container(
+                    executor_spec=executor_spec,
+                    entrypoint=xm.CommandList(["eval.py"]),
+                ),
+            ],
+        )
+
+        workdir = pathlib.Path(f"/scratch/xm-slurm-examples/{experiment.experiment_id}")
+
+        # Step 4: Schedule train job
+        train_executor = xm_slurm.Slurm(
+            requirements=xm_slurm.JobRequirements(
+                CPU=1,
+                RAM=1.0 * xm.GiB,
+                GPU=1,
+                replicas=1,
+                cluster=xm_slurm.contrib.clusters.mila(),
+            ),
+            time=dt.timedelta(hours=1),
+        )
+
+        async def make_train_job(wu: xm.WorkUnit, args):
+            await wu.add(
+                xm.Job(
+                    executable=train_executable,
+                    executor=train_executor,
+                    args=xm.merge_args(
+                        [
+                            "--output_file",
+                            (workdir / f"{wu.work_unit_id}" / "result.npy").as_posix(),
+                        ],
+                        args,
+                    ),
+                ),
+            )
+
+        train_wus = await experiment.add(
+            make_train_job,
+            args=[xm_slurm.JobArgs(args=["--seed", seed]) for seed in range(5)],
+        )
+
+        # Step 5: Schedule eval job
+        eval_executor = xm_slurm.Slurm(
+            requirements=xm_slurm.JobRequirements(
+                CPU=1,
+                RAM=1.0 * xm.GiB,
+                GPU=1,
+                replicas=1,
+                cluster=xm_slurm.contrib.clusters.mila(),
+            ),
+            time=dt.timedelta(hours=1),
+        )
+
+        eval_wus = await experiment.add(
+            xm.Job(
+                executable=eval_executable,
+                executor=eval_executor,
+            ),
+            args=[
+                xm_slurm.JobArgs(
+                    args=[
+                        "--input_file",
+                        (workdir / f"{wu.work_unit_id}" / "result.npy").as_posix(),
+                    ]
+                )
+                for wu in train_wus
+            ],
+            dependency=[train_wu.after_completed() for train_wu in train_wus],
+        )
+
+        for wu in asyncio.as_completed([
+            *[train_wu.wait_until_complete() for train_wu in train_wus],
+            *[eval_wu.wait_until_complete() for eval_wu in eval_wus],
+        ]):
+            wu = await wu
+            print(f"Work Unit {wu!r} finished executing with status {await wu.get_status()}")
+
+
+if __name__ == "__main__":
+    app.run(main)

xmanager_slurm-0.4.1/examples/job-dependencies/pyproject.toml

@@ -0,0 +1,6 @@
+[project]
+name = "xm-slurm-example"
+description = "XManager Slurm test project"
+version = "0.0.1"
+requires-python = ">=3.10"
+dependencies = ["numpy", "absl-py"]

xmanager_slurm-0.4.1/examples/job-dependencies/train.py

@@ -0,0 +1,19 @@
+import pathlib
+
+import numpy as np
+from absl import app, flags
+
+OUTPUT_FILE = flags.DEFINE_string("output_file", "result.npy", "Output file path")
+SEED = flags.DEFINE_integer("seed", 0, "Random seed")
+
+
+def main(_):
+    np.random.seed(SEED.value)
+
+    pathlib.Path(OUTPUT_FILE.value).parent.mkdir(parents=True, exist_ok=True)
+    result = np.random.random((32,))
+    np.save(OUTPUT_FILE.value, result)
+
+
+if __name__ == "__main__":
+    app.run(main)

xmanager_slurm-0.4.1/examples/job-dependencies/uv.lock

@@ -0,0 +1,86 @@
+version = 1
+requires-python = ">=3.10"
+
+[[package]]
+name = "absl-py"
+version = "2.1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/7a/8f/fc001b92ecc467cc32ab38398bd0bfb45df46e7523bf33c2ad22a505f06e/absl-py-2.1.0.tar.gz", hash = "sha256:7820790efbb316739cde8b4e19357243fc3608a152024288513dd968d7d959ff", size = 118055 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a2/ad/e0d3c824784ff121c03cc031f944bc7e139a8f1870ffd2845cc2dd76f6c4/absl_py-2.1.0-py3-none-any.whl", hash = "sha256:526a04eadab8b4ee719ce68f204172ead1027549089702d99b9059f129ff1308", size = 133706 },
+]
+
+[[package]]
+name = "numpy"
+version = "2.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/4b/d1/8a730ea07f4a37d94f9172f4ce1d81064b7a64766b460378be278952de75/numpy-2.1.2.tar.gz", hash = "sha256:13532a088217fa624c99b843eeb54640de23b3414b14aa66d023805eb731066c", size = 18878063 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/1c/a2/40a76d357f168e9f9f06d6cc2c8e22dd5fb2bfbe63fe2c433057258c145a/numpy-2.1.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:30d53720b726ec36a7f88dc873f0eec8447fbc93d93a8f079dfac2629598d6ee", size = 21150947 },
+    { url = "https://files.pythonhosted.org/packages/b5/d0/ba271ea9108d7278d3889a7eb38d77370a88713fb94339964e71ac184d4a/numpy-2.1.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:e8d3ca0a72dd8846eb6f7dfe8f19088060fcb76931ed592d29128e0219652884", size = 13758184 },
+    { url = "https://files.pythonhosted.org/packages/7c/b9/5c6507439cd756201010f7937bf90712c2469052ae094584af14557dd64f/numpy-2.1.2-cp310-cp310-macosx_14_0_arm64.whl", hash = "sha256:fc44e3c68ff00fd991b59092a54350e6e4911152682b4782f68070985aa9e648", size = 5354091 },
+    { url = "https://files.pythonhosted.org/packages/60/21/7938cf724d9e84e45fb886f3fc794ab431d71facfebc261e3e9f19f3233a/numpy-2.1.2-cp310-cp310-macosx_14_0_x86_64.whl", hash = "sha256:7c1c60328bd964b53f8b835df69ae8198659e2b9302ff9ebb7de4e5a5994db3d", size = 6887169 },
+    { url = "https://files.pythonhosted.org/packages/09/8d/42a124657f5d31902fca73921b25a0d022cead2b32ce7e6975762cd2995a/numpy-2.1.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6cdb606a7478f9ad91c6283e238544451e3a95f30fb5467fbf715964341a8a86", size = 13888165 },
+    { url = "https://files.pythonhosted.org/packages/fb/25/ba023652a39a2c127200e85aed975fc6119b421e2c348e5d0171e2046edb/numpy-2.1.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d666cb72687559689e9906197e3bec7b736764df6a2e58ee265e360663e9baf7", size = 16326954 },
+    { url = "https://files.pythonhosted.org/packages/34/58/23e6b07fad492b7c47cf09cd8bad6983658f0f925b6c535fd008e3e86274/numpy-2.1.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:c6eef7a2dbd0abfb0d9eaf78b73017dbfd0b54051102ff4e6a7b2980d5ac1a03", size = 16702916 },
+    { url = "https://files.pythonhosted.org/packages/91/24/37b5cf2dc7d385ac97f7b7fe50cba312abb70a2a5eac74c23af028811f73/numpy-2.1.2-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:12edb90831ff481f7ef5f6bc6431a9d74dc0e5ff401559a71e5e4611d4f2d466", size = 14384372 },
+    { url = "https://files.pythonhosted.org/packages/ea/ec/0f6d471058a01d1a05a50d2793898de1549280fa715a8537987ee866b5d9/numpy-2.1.2-cp310-cp310-win32.whl", hash = "sha256:a65acfdb9c6ebb8368490dbafe83c03c7e277b37e6857f0caeadbbc56e12f4fb", size = 6535361 },
+    { url = "https://files.pythonhosted.org/packages/c2/3d/293cc5927f916a7bc6bf74da8f6defab63d1b13f0959d7e21878ad8a20d8/numpy-2.1.2-cp310-cp310-win_amd64.whl", hash = "sha256:860ec6e63e2c5c2ee5e9121808145c7bf86c96cca9ad396c0bd3e0f2798ccbe2", size = 12865501 },
+    { url = "https://files.pythonhosted.org/packages/aa/9c/9a6ec3ae89cd0648d419781284308f2956d2a61d932b5ac9682c956a171b/numpy-2.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:b42a1a511c81cc78cbc4539675713bbcf9d9c3913386243ceff0e9429ca892fe", size = 21154845 },
+    { url = "https://files.pythonhosted.org/packages/02/69/9f05c4ecc75fabf297b17743996371b4c3dfc4d92e15c5c38d8bb3db8d74/numpy-2.1.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:faa88bc527d0f097abdc2c663cddf37c05a1c2f113716601555249805cf573f1", size = 13789409 },
+    { url = "https://files.pythonhosted.org/packages/34/4e/f95c99217bf77bbfaaf660d693c10bd0dc03b6032d19316d316088c9e479/numpy-2.1.2-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:c82af4b2ddd2ee72d1fc0c6695048d457e00b3582ccde72d8a1c991b808bb20f", size = 5352097 },
+    { url = "https://files.pythonhosted.org/packages/06/13/f5d87a497c16658e9af8920449b0b5692b469586b8231340c672962071c5/numpy-2.1.2-cp311-cp311-macosx_14_0_x86_64.whl", hash = "sha256:13602b3174432a35b16c4cfb5de9a12d229727c3dd47a6ce35111f2ebdf66ff4", size = 6891195 },
+    { url = "https://files.pythonhosted.org/packages/6c/89/691ac07429ac061b344d5e37fa8e94be51a6017734aea15f2d9d7c6d119a/numpy-2.1.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1ebec5fd716c5a5b3d8dfcc439be82a8407b7b24b230d0ad28a81b61c2f4659a", size = 13895153 },
+    { url = "https://files.pythonhosted.org/packages/23/69/538317f0d925095537745f12aced33be1570bbdc4acde49b33748669af96/numpy-2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e2b49c3c0804e8ecb05d59af8386ec2f74877f7ca8fd9c1e00be2672e4d399b1", size = 16338306 },
+    { url = "https://files.pythonhosted.org/packages/af/03/863fe7062c2106d3c151f7df9353f2ae2237c1dd6900f127a3eb1f24cb1b/numpy-2.1.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:2cbba4b30bf31ddbe97f1c7205ef976909a93a66bb1583e983adbd155ba72ac2", size = 16710893 },
+    { url = "https://files.pythonhosted.org/packages/70/77/0ad9efe25482009873f9660d29a40a8c41a6f0e8b541195e3c95c70684c5/numpy-2.1.2-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8e00ea6fc82e8a804433d3e9cedaa1051a1422cb6e443011590c14d2dea59146", size = 14398048 },
+    { url = "https://files.pythonhosted.org/packages/3e/0f/e785fe75544db9f2b0bb1c181e13ceff349ce49753d807fd9672916aa06d/numpy-2.1.2-cp311-cp311-win32.whl", hash = "sha256:5006b13a06e0b38d561fab5ccc37581f23c9511879be7693bd33c7cd15ca227c", size = 6533458 },
+    { url = "https://files.pythonhosted.org/packages/d4/96/450054662295125af861d48d2c4bc081dadcf1974a879b2104613157aa62/numpy-2.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:f1eb068ead09f4994dec71c24b2844f1e4e4e013b9629f812f292f04bd1510d9", size = 12870896 },
+    { url = "https://files.pythonhosted.org/packages/a0/7d/554a6838f37f3ada5a55f25173c619d556ae98092a6e01afb6e710501d70/numpy-2.1.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d7bf0a4f9f15b32b5ba53147369e94296f5fffb783db5aacc1be15b4bf72f43b", size = 20848077 },
+    { url = "https://files.pythonhosted.org/packages/b0/29/cb48a402ea879e645b16218718f3f7d9588a77d674a9dcf22e4c43487636/numpy-2.1.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b1d0fcae4f0949f215d4632be684a539859b295e2d0cb14f78ec231915d644db", size = 13493242 },
+    { url = "https://files.pythonhosted.org/packages/56/44/f899b0581766c230da42f751b7b8896d096640b19b312164c267e48d36cb/numpy-2.1.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:f751ed0a2f250541e19dfca9f1eafa31a392c71c832b6bb9e113b10d050cb0f1", size = 5089219 },
+    { url = "https://files.pythonhosted.org/packages/79/8f/b987070d45161a7a4504afc67ed38544ed2c0ed5576263599a0402204a9c/numpy-2.1.2-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:bd33f82e95ba7ad632bc57837ee99dba3d7e006536200c4e9124089e1bf42426", size = 6620167 },
+    { url = "https://files.pythonhosted.org/packages/c4/a7/af3329fda3c3ec31d9b650e42bbcd3422fc62a765cbb1405fde4177a0996/numpy-2.1.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1b8cde4f11f0a975d1fd59373b32e2f5a562ade7cde4f85b7137f3de8fbb29a0", size = 13604905 },
+    { url = "https://files.pythonhosted.org/packages/9b/b4/e3c7e6fab0f77fff6194afa173d1f2342073d91b1d3b4b30b17c3fb4407a/numpy-2.1.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6d95f286b8244b3649b477ac066c6906fbb2905f8ac19b170e2175d3d799f4df", size = 16041825 },
+    { url = "https://files.pythonhosted.org/packages/e9/50/6828e66a78aa03147c111f84d55f33ce2dde547cb578d6744a3b06a0124b/numpy-2.1.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:ab4754d432e3ac42d33a269c8567413bdb541689b02d93788af4131018cbf366", size = 16409541 },
+    { url = "https://files.pythonhosted.org/packages/bf/72/66af7916d9c3c6dbfbc8acdd4930c65461e1953374a2bc43d00f948f004a/numpy-2.1.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e585c8ae871fd38ac50598f4763d73ec5497b0de9a0ab4ef5b69f01c6a046142", size = 14081134 },
+    { url = "https://files.pythonhosted.org/packages/dc/5a/59a67d84f33fe00ae74f0b5b69dd4f93a586a4aba7f7e19b54b2133db038/numpy-2.1.2-cp312-cp312-win32.whl", hash = "sha256:9c6c754df29ce6a89ed23afb25550d1c2d5fdb9901d9c67a16e0b16eaf7e2550", size = 6237784 },
+    { url = "https://files.pythonhosted.org/packages/4c/79/73735a6a5dad6059c085f240a4e74c9270feccd2bc66e4d31b5ca01d329c/numpy-2.1.2-cp312-cp312-win_amd64.whl", hash = "sha256:456e3b11cb79ac9946c822a56346ec80275eaf2950314b249b512896c0d2505e", size = 12568254 },
+    { url = "https://files.pythonhosted.org/packages/16/72/716fa1dbe92395a9a623d5049203ff8ddb0cfce65b9df9117c3696ccc011/numpy-2.1.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:a84498e0d0a1174f2b3ed769b67b656aa5460c92c9554039e11f20a05650f00d", size = 20834690 },
+    { url = "https://files.pythonhosted.org/packages/1e/fb/3e85a39511586053b5c6a59a643879e376fae22230ebfef9cfabb0e032e2/numpy-2.1.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4d6ec0d4222e8ffdab1744da2560f07856421b367928026fb540e1945f2eeeaf", size = 13507474 },
+    { url = "https://files.pythonhosted.org/packages/35/eb/5677556d9ba13436dab51e129f98d4829d95cd1b6bd0e199c14485a4bdb9/numpy-2.1.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:259ec80d54999cc34cd1eb8ded513cb053c3bf4829152a2e00de2371bd406f5e", size = 5074742 },
+    { url = "https://files.pythonhosted.org/packages/3e/c5/6c5ef5ba41b65a7e51bed50dbf3e1483eb578055633dd013e811a28e96a1/numpy-2.1.2-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:675c741d4739af2dc20cd6c6a5c4b7355c728167845e3c6b0e824e4e5d36a6c3", size = 6606787 },
+    { url = "https://files.pythonhosted.org/packages/08/ac/f2f29dd4fd325b379c7dc932a0ebab22f0e031dbe80b2f6019b291a3a544/numpy-2.1.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:05b2d4e667895cc55e3ff2b56077e4c8a5604361fc21a042845ea3ad67465aa8", size = 13601333 },
+    { url = "https://files.pythonhosted.org/packages/44/26/63f5f4e5089654dfb858f4892215ed968cd1a68e6f4a83f9961f84f855cb/numpy-2.1.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:43cca367bf94a14aca50b89e9bc2061683116cfe864e56740e083392f533ce7a", size = 16038090 },
+    { url = "https://files.pythonhosted.org/packages/1d/21/015e0594de9c3a8d5edd24943d2bd23f102ec71aec026083f822f86497e2/numpy-2.1.2-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:76322dcdb16fccf2ac56f99048af32259dcc488d9b7e25b51e5eca5147a3fb98", size = 16410865 },
+    { url = "https://files.pythonhosted.org/packages/df/01/c1bcf9e6025d79077fbf3f3ee503b50aa7bfabfcd8f4b54f5829f4c00f3f/numpy-2.1.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:32e16a03138cabe0cb28e1007ee82264296ac0983714094380b408097a418cfe", size = 14078077 },
+    { url = "https://files.pythonhosted.org/packages/ba/06/db9d127d63bd11591770ba9f3d960f8041e0f895184b9351d4b1b5b56983/numpy-2.1.2-cp313-cp313-win32.whl", hash = "sha256:242b39d00e4944431a3cd2db2f5377e15b5785920421993770cddb89992c3f3a", size = 6234904 },
+    { url = "https://files.pythonhosted.org/packages/a9/96/9f61f8f95b6e0ea0aa08633b704c75d1882bdcb331bdf8bfd63263b25b00/numpy-2.1.2-cp313-cp313-win_amd64.whl", hash = "sha256:f2ded8d9b6f68cc26f8425eda5d3877b47343e68ca23d0d0846f4d312ecaa445", size = 12561910 },
+    { url = "https://files.pythonhosted.org/packages/36/b8/033f627821784a48e8f75c218033471eebbaacdd933f8979c79637a1b44b/numpy-2.1.2-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:2ffef621c14ebb0188a8633348504a35c13680d6da93ab5cb86f4e54b7e922b5", size = 20857719 },
+    { url = "https://files.pythonhosted.org/packages/96/46/af5726fde5b74ed83f2f17a73386d399319b7ed4d51279fb23b721d0816d/numpy-2.1.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ad369ed238b1959dfbade9018a740fb9392c5ac4f9b5173f420bd4f37ba1f7a0", size = 13518826 },
+    { url = "https://files.pythonhosted.org/packages/db/6e/8ce677edf36da1c4dae80afe5529f47690697eb55b4864673af260ccea7b/numpy-2.1.2-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:d82075752f40c0ddf57e6e02673a17f6cb0f8eb3f587f63ca1eaab5594da5b17", size = 5115036 },
+    { url = "https://files.pythonhosted.org/packages/6a/ba/3cce44fb1b8438042c11847048812a776f75ee0e7070179c22e4cfbf420c/numpy-2.1.2-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:1600068c262af1ca9580a527d43dc9d959b0b1d8e56f8a05d830eea39b7c8af6", size = 6628641 },
+    { url = "https://files.pythonhosted.org/packages/59/c8/e722998720ccbd35ffbcf1d1b8ed0aa2304af88d3f1c38e06ebf983599b3/numpy-2.1.2-cp313-cp313t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a26ae94658d3ba3781d5e103ac07a876b3e9b29db53f68ed7df432fd033358a8", size = 13574803 },
+    { url = "https://files.pythonhosted.org/packages/7c/8e/fc1fdd83a55476765329ac2913321c4aed5b082a7915095628c4ca30ea72/numpy-2.1.2-cp313-cp313t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:13311c2db4c5f7609b462bc0f43d3c465424d25c626d95040f073e30f7570e35", size = 16021174 },
+    { url = "https://files.pythonhosted.org/packages/2a/b6/a790742aa88067adb4bd6c89a946778c1417d4deaeafce3ca928f26d4c52/numpy-2.1.2-cp313-cp313t-musllinux_1_1_x86_64.whl", hash = "sha256:2abbf905a0b568706391ec6fa15161fad0fb5d8b68d73c461b3c1bab6064dd62", size = 16400117 },
+    { url = "https://files.pythonhosted.org/packages/48/6f/129e3c17e3befe7fefdeaa6890f4c4df3f3cf0831aa053802c3862da67aa/numpy-2.1.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:ef444c57d664d35cac4e18c298c47d7b504c66b17c2ea91312e979fcfbdfb08a", size = 14066202 },
+    { url = "https://files.pythonhosted.org/packages/73/c9/3e1d6bbe6d3d2e2c5a9483b24b2f29a229b323f62054278a3bba7fee11e5/numpy-2.1.2-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:bdd407c40483463898b84490770199d5714dcc9dd9b792f6c6caccc523c00952", size = 20981945 },
+    { url = "https://files.pythonhosted.org/packages/6e/62/989c4988bde1a8e08117fccc3bab73d2886421fb98cde597168714f3c54e/numpy-2.1.2-pp310-pypy310_pp73-macosx_14_0_x86_64.whl", hash = "sha256:da65fb46d4cbb75cb417cddf6ba5e7582eb7bb0b47db4b99c9fe5787ce5d91f5", size = 6750558 },
+    { url = "https://files.pythonhosted.org/packages/53/b1/00ef9f30975f1312a53257f68e57b4513d14d537e03d507e2773a684b1e8/numpy-2.1.2-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1c193d0b0238638e6fc5f10f1b074a6993cb13b0b431f64079a509d63d3aa8b7", size = 16141552 },
+    { url = "https://files.pythonhosted.org/packages/c0/ec/0c04903b48dfea6be1d7b47ba70f98709fb7198fd970784a1400c391d522/numpy-2.1.2-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:a7d80b2e904faa63068ead63107189164ca443b42dd1930299e0d1cb041cec2e", size = 12789924 },
+]
+
+[[package]]
+name = "xm-slurm-example"
+version = "0.0.1"
+source = { virtual = "." }
+dependencies = [
+    { name = "absl-py" },
+    { name = "numpy" },
+]
+
+[package.metadata]
+requires-dist = [
+    { name = "absl-py" },
+    { name = "numpy" },
+]

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "xmanager-slurm"
 description = "Slurm backend for XManager."
-version = "0.4.0"
+version = "0.4.1"
 # readme = "README.md"
 requires-python = ">=3.10"
 license = { text = "MIT" }
@@ -36,6 +36,9 @@ dependencies = [
 [project.urls]
 "GitHub" = "https://github.com/jessefarebro/xm-slurm"
 
+[project.scripts]
+xm = "xm_slurm.scripts.cli:main"
+
 [tool.uv]
 dev-dependencies = [
     "mypy>=1.11.2",

xmanager_slurm-0.4.1/tests/test_dependencies.py

@@ -0,0 +1,149 @@
+import datetime as dt
+from typing import Callable
+
+import pytest
+
+from xm_slurm import config, dependencies, execution
+
+SlurmHandleGenerator = Callable[[str], execution.SlurmHandle]
+
+
+@pytest.fixture
+def slurm_handle() -> SlurmHandleGenerator:
+    dummy_ssh_config = config.SlurmSSHConfig(host="localhost")
+
+    def _slurm_handle(job_id: str):
+        return execution.SlurmHandle(
+            experiment_id=0, ssh=dummy_ssh_config, slurm_job=job_id, job_name="job"
+        )
+
+    return _slurm_handle
+
+
+def test_slurm_job_dependency_and(slurm_handle: SlurmHandleGenerator):
+    dep1 = dependencies.SlurmJobDependencyAfter([slurm_handle("123")])
+    dep2 = dependencies.SlurmJobDependencyAfter([slurm_handle("456")])
+    combined_dep = dep1 & dep2
+    assert isinstance(combined_dep, dependencies.SlurmJobDependencyAND)
+    assert combined_dep.to_dependency_str() == "after:123,after:456"
+
+
+def test_slurm_job_dependency_or(slurm_handle: SlurmHandleGenerator):
+    dep1 = dependencies.SlurmJobDependencyAfter([slurm_handle("123")])
+    dep2 = dependencies.SlurmJobDependencyAfter([slurm_handle("456")])
+    combined_dep = dep1 | dep2
+    assert isinstance(combined_dep, dependencies.SlurmJobDependencyOR)
+    assert combined_dep.to_dependency_str() == "after:123?after:456"
+
+
+def test_slurm_job_dependency_mixing_logical_operations(slurm_handle: SlurmHandleGenerator):
+    dep1 = dependencies.SlurmJobDependencyAfter([slurm_handle("123")])
+    dep2 = dependencies.SlurmJobDependencyAfter([slurm_handle("456")])
+    dep3 = dependencies.SlurmJobDependencyAfter([slurm_handle("789")])
+    with pytest.raises(
+        dependencies.SlurmDependencyException,
+        match="Slurm only supports chaining dependencies with the same logical operator. ",
+    ):
+        dep1 & dep2 | dep3  # type: ignore
+
+    with pytest.raises(
+        dependencies.SlurmDependencyException,
+        match="Slurm only supports chaining dependencies with the same logical operator. ",
+    ):
+        dep1 | dep2 & dep3  # type: ignore
+
+
+def test_slurm_job_dependency_chaining_and(slurm_handle: SlurmHandleGenerator):
+    dep1 = dependencies.SlurmJobDependencyAfter([slurm_handle("1")])
+    dep2 = dependencies.SlurmJobDependencyAfter([slurm_handle("2")])
+    dep3 = dependencies.SlurmJobDependencyAfter([slurm_handle("3")])
+    dep4 = dependencies.SlurmJobDependencyAfter([slurm_handle("4")])
+    combined_dep = dep1 & dep2 & dep3 & dep4
+    assert isinstance(combined_dep, dependencies.SlurmJobDependencyAND)
+    assert combined_dep.to_dependency_str() == "after:1,after:2,after:3,after:4"
+
+
+def test_slurm_job_dependency_chaining_or(slurm_handle: SlurmHandleGenerator):
+    dep1 = dependencies.SlurmJobDependencyAfter([slurm_handle("1")])
+    dep2 = dependencies.SlurmJobDependencyAfter([slurm_handle("2")])
+    dep3 = dependencies.SlurmJobDependencyAfter([slurm_handle("3")])
+    dep4 = dependencies.SlurmJobDependencyAfter([slurm_handle("4")])
+    combined_dep = dep1 | dep2 | dep3 | dep4
+    assert isinstance(combined_dep, dependencies.SlurmJobDependencyOR)
+    assert combined_dep.to_dependency_str() == "after:1?after:2?after:3?after:4"
+
+
+def test_slurm_job_dependency_after(slurm_handle: SlurmHandleGenerator):
+    dep = dependencies.SlurmJobDependencyAfter([slurm_handle("123")])
+    assert dep.to_dependency_str() == "after:123"
+
+
+def test_slurm_job_dependency_after_with_time(slurm_handle: SlurmHandleGenerator):
+    dep = dependencies.SlurmJobDependencyAfter([slurm_handle("123")], time=dt.timedelta(minutes=10))
+    assert dep.to_dependency_str() == "after:123+10"
+
+
+def test_slurm_job_dependency_after_with_invalid_time(slurm_handle: SlurmHandleGenerator):
+    with pytest.raises(
+        dependencies.SlurmDependencyException, match="Time must be specified in exact minutes"
+    ):
+        dependencies.SlurmJobDependencyAfter([slurm_handle("123")], time=dt.timedelta(seconds=30))
+
+
+@pytest.mark.parametrize(
+    "dependency_cls,dependency_type",
+    [
+        (dependencies.SlurmJobDependencyAfter, "after"),
+        (dependencies.SlurmJobDependencyAfterAny, "afterany"),
+        (dependencies.SlurmJobDependencyAfterNotOK, "afternotok"),
+        (dependencies.SlurmJobDependencyAfterOK, "afterok"),
+    ],
+)
+def test_slurm_job_dependency_after_not_ok(
+    slurm_handle: SlurmHandleGenerator,
+    dependency_cls: type,
+    dependency_type: str,
+):
+    dep = dependency_cls([slurm_handle("123"), slurm_handle("456")])
+    assert dep.to_dependency_str() == f"{dependency_type}:123:456"
+
+
+@pytest.mark.parametrize(
+    "dependency_cls",
+    [
+        dependencies.SlurmJobDependencyAfter,
+        dependencies.SlurmJobDependencyAfterAny,
+        dependencies.SlurmJobDependencyAfterNotOK,
+        dependencies.SlurmJobDependencyAfterOK,
+    ],
+)
+def test_slurm_job_dependency_after_no_handles(dependency_cls: type):
+    with pytest.raises(
+        dependencies.SlurmDependencyException, match="Dependency doesn't have any handles."
+    ):
+        dependency_cls([])
+
+
+def test_dependency_flatten(slurm_handle: SlurmHandleGenerator):
+    dep1 = dependencies.SlurmJobDependencyAfter([slurm_handle("1")])
+    dep2 = dependencies.SlurmJobDependencyAfter([slurm_handle("2")])
+    dep3 = dependencies.SlurmJobDependencyAfter([slurm_handle("3")])
+    combined_dep = dep1 & dep2 & dep3
+    assert combined_dep.flatten() == (dep1, dep2, dep3)
+
+
+def test_dependency_traverse(slurm_handle: SlurmHandleGenerator):
+    dep1 = dependencies.SlurmJobDependencyAfterOK([slurm_handle("1")])
+    dep2 = dependencies.SlurmJobDependencyAfterOK([slurm_handle("2")])
+    dep3 = dependencies.SlurmJobDependencyAfterOK([slurm_handle("3")])
+    combined_dep = dep1 & dep2 & dep3
+
+    def traverse_fn(dep: dependencies.SlurmJobDependency):
+        if isinstance(dep, dependencies.SlurmJobDependencyAfterOK):
+            return dependencies.SlurmJobDependencyAfterNotOK(dep.handles)
+        return dep
+
+    transformed_combined_dep = combined_dep.traverse(traverse_fn)
+    for dep in transformed_combined_dep.flatten():
+        assert isinstance(dep, dependencies.SlurmJobDependencyAfterNotOK)
+    assert transformed_combined_dep.to_dependency_str() == "afternotok:1,afternotok:2,afternotok:3"

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/uv.lock

@@ -2196,7 +2196,7 @@ wheels = [
 
 [[package]]
 name = "xmanager-slurm"
-version = "0.3.2"
+version = "0.4.0"
 source = { editable = "." }
 dependencies = [
     { name = "asyncssh" },

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/xm_slurm/__init__.py

@@ -10,6 +10,7 @@ from xm_slurm.experiment import (
     get_current_work_unit,
     get_experiment,
 )
+from xm_slurm.job_blocks import JobArgs
 from xm_slurm.packageables import (
     conda_container,
     docker_container,
@@ -34,13 +35,14 @@ __all__ = [
     "get_current_experiment",
     "get_current_work_unit",
     "get_experiment",
+    "JobArgs",
     "JobRequirements",
     "mamba_container",
-    "uv_container",
     "python_container",
     "ResourceQuantity",
     "ResourceType",
     "Slurm",
-    "SlurmSpec",
     "SlurmExperiment",
+    "SlurmSpec",
+    "uv_container",
 ]

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/xm_slurm/api.py

@@ -322,7 +322,7 @@ class XManagerSqliteAPI(XManagerAPI):
             db_path = Path(os.environ["XM_SLURM_STATE_DIR"]) / "db.sqlite3"
         else:
             db_path = Path.home() / ".local" / "state" / "xm-slurm" / "db.sqlite3"
-        logging.debug("Looking for db at: ", db_path)
+        logger.debug("Looking for db at: ", db_path)
         db_path.parent.mkdir(parents=True, exist_ok=True)
         engine = create_engine(f"sqlite:///{db_path}")
         Base.metadata.create_all(engine)

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/xm_slurm/config.py

@@ -5,9 +5,10 @@ import getpass
 import json
 import os
 import pathlib
-from typing import Literal, Mapping, NamedTuple
+from typing import Callable, Literal, Mapping, NamedTuple
 
 import asyncssh
+from xmanager import xm
 
 from xm_slurm import constants
 
@@ -137,7 +138,7 @@ class SlurmSSHConfig:
         )
 
     def __hash__(self):
-        return hash((self.host, self.host_public_key, self.user, self.port))
+        return hash((type(self), self.host, self.host_public_key, self.user, self.port))
 
 
 @dataclasses.dataclass(frozen=True, kw_only=True)
@@ -176,6 +177,9 @@ class SlurmClusterConfig:
 
     features: Mapping["xm_slurm.FeatureType", str] = dataclasses.field(default_factory=dict)  # type: ignore # noqa: F821
 
+    # Function to validate the Slurm executor config
+    validate: Callable[[xm.Job], None] | None = None
+
     def __post_init__(self) -> None:
         for src, dst in self.mounts.items():
             if not isinstance(src, (str, os.PathLike)):
@@ -194,6 +198,7 @@ class SlurmClusterConfig:
 
     def __hash__(self):
         return hash((
+            type(self),
             self.ssh,
             self.cwd,
             self.prolog,

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/xm_slurm/constants.py

@@ -1,5 +1,9 @@
 import re
 
+SLURM_JOB_ID_REGEX = re.compile(
+    r"^(?P<jobid>\d+)(?:(?:\+(?P<componentid>\d+))|(?:_(?P<arraytaskid>\d+)))?$"
+)
+
 IMAGE_URI_REGEX = re.compile(
     r"^(?P<scheme>(?:[^:]+://)?)?(?P<domain>[^/]+)(?P<path>/[^:]*)?(?::(?P<tag>[^@]+))?@?(?P<digest>.+)?$"
 )

{xmanager_slurm-0.4.0 → xmanager_slurm-0.4.1}/xm_slurm/contrib/clusters/__init__.py

@@ -1,13 +1,20 @@
+import datetime as dt
+import logging
 import os
 
+from xmanager import xm
+
 from xm_slurm import config, resources
 from xm_slurm.contrib.clusters import drac
+from xm_slurm.executors import Slurm
 
 # ComputeCanada alias
 cc = drac
 
 __all__ = ["drac", "mila", "cc"]
 
+logger = logging.getLogger(__name__)
+
 
 def mila(
     *,
@@ -24,6 +31,23 @@ def mila(
             "/home/mila/${USER:0:1}/$USER/.ssh": "/home/mila/${USER:0:1}/$USER/.ssh",
         }
 
+    def validate(job: xm.Job) -> None:
+        assert isinstance(job.executor, Slurm)
+
+        wants_requeue_with_grace_period = (
+            job.executor.requeue and job.executor.timeout_signal_grace_period > dt.timedelta(0)
+        )
+        partition = job.executor.partition or "main"
+
+        if wants_requeue_with_grace_period and (
+            partition is None or not partition.endswith("-grace")
+        ):
+            logger.warning(
+                f"Job {job.name} wants requeue with grace period, but partition `{partition}` does not end with '-grace'. "
+                "Mila Cluster requires you specify a grace partition. "
+                "This may result in the job not being requeued properly."
+            )
+
     return config.SlurmClusterConfig(
         name="mila",
         ssh=config.SlurmSSHConfig(
@@ -58,4 +82,5 @@ def mila(
             resources.FeatureType.NVIDIA_MIG: "mig",
             resources.FeatureType.NVIDIA_NVLINK: "nvlink",
         },
+        validate=validate,
     )