xmanager-slurm 0.4.5__tar.gz → 0.4.7__tar.gz

xmanager_slurm-0.4.7/.pdm-build/.gitignore

@@ -0,0 +1 @@
+*

xmanager_slurm-0.4.7/.pdm-build/xm_slurm-0.1+editable.dist-info/METADATA

@@ -0,0 +1,125 @@
+Metadata-Version: 2.1
+Name: xm-slurm
+Version: 0.1+editable
+Author-Email: Jesse Farebrother <jfarebro@cs.mcgill.ca>
+License: MIT
+Requires-Python: >=3.10
+Requires-Dist: xmanager>=0.4.0
+Requires-Dist: asyncssh>=2.13.2
+Requires-Dist: humanize>=4.8.0
+Requires-Dist: jinja2>=3.1.2
+Requires-Dist: toml>=0.10.2
+Requires-Dist: rich>=13.5.2
+Requires-Dist: immutabledict>=3.0.0
+Requires-Dist: backoff>=2.2.1
+Requires-Dist: sqlalchemy>=2.0
+Requires-Dist: alembic>=1.13.1
+Requires-Dist: aiosqlite>=0.20.0
+Requires-Dist: pathspec>=0.11.2; extra == "gcp"
+Requires-Dist: google-cloud-storage>=2.11.0; extra == "gcp"
+Requires-Dist: google-cloud-build>=3.20.0; extra == "gcp"
+Requires-Dist: google-cloud-logging>=3.8.0; extra == "gcp"
+Requires-Dist: google-cloud-iam>=2.12.2; extra == "gcp"
+Requires-Dist: google-cloud-kms>=2.19.2; extra == "gcp"
+Requires-Dist: google-crc32c>=1.5.0; extra == "gcp"
+Requires-Dist: pytest>=7.4.3; extra == "test"
+Provides-Extra: gcp
+Provides-Extra: test
+Description-Content-Type: text/markdown
+
+# Slurm XManager
+
+This project adds support for Slurm to XManager. This is done through the use of Docker on the client and Singularity/Apptainer containers on the cluster.
+This provides the following benefits:
+
+1. All development can be done locally and launched on any Slurm cluster without any configuration.
+2. Reproducible experiments (e.g., containerized runtime, code checkpointing, etc.)
+3. Easy to configure distributed experiments, all configuration is in Python.
+4. Launch experiments on any XManager launcher, e.g., GCP or Kubernetes
+
+## Minimal Example
+
+Currently the only out of the box container type is a [PDM](https://pdm.fming.dev/latest/) container.
+To use the pdm container you can start a new project with `pdm init`. From there you'll need to implement a launch script. A launch script will be decomposed into three parts:
+
+1. Specifying an executor specification, this will be where the exectuable will be stored.
+2. Specifying an executable and packaging the executables.
+3. Specifying a job and job requirements.
+
+```py
+import datetime
+from pathlib import Path
+
+from xmanager import xm
+
+import xm_slurm
+from xm_slurm.apptainer import packageables
+
+
+@xm.run_in_asyncio_loop
+async def main():
+    title = FLAGS.title
+
+    async with xm_slurm.create_experiment(title) as experiment:
+        # Step 1: Specify the executor specification
+        executor_spec = xm_slurm.SlurmSpec(tag="ghcr.io/YOUR_GITHUB_USERNAME/YOUR_GITHUB_REPOSITORY/launch:latest")
+
+        # Step 2: Specify the executable and package it
+        [executable] = experiment.package(
+            [
+                packageables.pdm_container(
+                    executor_spec=executor_spec,
+                    entrypoint=xm.ModuleName("train"),
+                    annotations={
+                        "org.opencontainers.image.source": "https://github.com/YOUR_GITHUB_USERNAME/YOUR_GITHUB_REPOSITORY"
+                    },
+                    args={}, # SPECIFY COMMAND CLI ARGS FOR THE EXECUTABLE
+                    env_vars={}, # SPECIFY COMMON ENV VARS FOR THE EXECUTABLE
+                ),
+            ]
+        )
+
+        # Step 3: Specify the executor and add the job
+        executor = xm_slurm.Slurm(
+            requirements=xm.JobRequirements(
+                resources={xm.ResourceType.A100: 1},
+                RAM=8 * xm.GiB,
+                CPU=4,
+            ),
+            time=dt.timedelta(hours=24),
+            account="", # SLURM ACCOUNT
+            # SPECIFY OTHER SLURM ARGUMENTS
+        )
+
+        await experiment.add(
+            job=xm.Job(
+                executor=executor,
+                executable=executable,
+                args={}, # SPECIFY ARGS FOR THIS JOB
+                env_vars={}, # SPECIFY ENV VARS FOR THIS JOB
+            ),
+        )
+
+if __name__ == "__main__":
+    main()
+```
+
+### Specifying sweeps
+
+When calling `experiment.add` you can specify an additional keyword argument `args`. This argument can be a sequence of arguments that'll specify a sweep. For example, if you want to run a sweep over the learning rate you can specify the following:
+
+```py
+await experiment.add(
+    job=xm.Job(
+        executor=executor,
+        executable=executable,
+    ),
+    args=[
+        {"learning_rate": 0.1},
+        {"learning_rate": 0.01},
+        {"learning_rate": 0.001},
+    ],
+)
+```
+
+This will launch a job array with 3 jobs, each with a different learning rate.

xmanager_slurm-0.4.7/.pdm-build/xm_slurm-0.1+editable.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: pdm-backend (2.1.8)
+Root-Is-Purelib: true
+Tag: py3-none-any

xmanager_slurm-0.4.7/.pdm-build/xm_slurm.pth

@@ -0,0 +1 @@
+/workspaces/xm-slurm

xmanager_slurm-0.4.7/.python-version

@@ -0,0 +1 @@
+3.11

{xmanager_slurm-0.4.5 → xmanager_slurm-0.4.7}/PKG-INFO

@@ -1,10 +1,11 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.5
+Version: 0.4.7
 Summary: Slurm backend for XManager.
 Project-URL: GitHub, https://github.com/jessefarebro/xm-slurm
 Author-email: Jesse Farebrother <jfarebro@cs.mcgill.ca>
 License: MIT
+License-File: LICENSE.md
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
@@ -13,13 +14,16 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Python: >=3.10
-Requires-Dist: asyncssh>=2.13.2
+Requires-Dist: aiofile>=3.9.0
+Requires-Dist: asyncssh>=2.19.0
 Requires-Dist: backoff>=2.2.1
 Requires-Dist: cloudpickle>=3.0.0
+Requires-Dist: httpx>=0.28.1
 Requires-Dist: humanize>=4.8.0
 Requires-Dist: immutabledict>=3.0.0
 Requires-Dist: jinja2>=3.1.2
 Requires-Dist: more-itertools>=10.2.0
 Requires-Dist: rich>=13.5.2
 Requires-Dist: toml>=0.10.2
+Requires-Dist: wrapt>=1.16.0
 Requires-Dist: xmanager>=0.5.0

{xmanager_slurm-0.4.5 → xmanager_slurm-0.4.7}/examples/metadata/launch.py

@@ -25,9 +25,18 @@ def f(x):
 ```
         """
         experiment.context.annotations.title = "Cool Experiment #2"
-        experiment.context.artifacts.add(
-            xm_slurm.Artifact(name="wandb", uri="https://wandb.ai/...")
-        )
+        experiment.context.artifacts["wandb"] = "https://wandb.ai/..."
+        experiment.context.python_config = """
+seed = 0
+"""
+        experiment.context.graphviz_config = """
+digraph G {
+    A -> B;
+    A -> C;
+    B -> D;
+    C -> D;
+}
+"""
 
         # Step 1: Specify the executor specification
         executor_spec = xm_slurm.Slurm.Spec(tag="ghcr.io/jessefarebro/xm-slurm/test:latest")

{xmanager_slurm-0.4.5 → xmanager_slurm-0.4.7}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "xmanager-slurm"
 description = "Slurm backend for XManager."
-version = "0.4.5"
+version = "0.4.7"
 # readme = "README.md"
 requires-python = ">=3.10"
 license = { text = "MIT" }
@@ -22,7 +22,7 @@ classifiers = [
 
 dependencies = [
     "xmanager>=0.5.0",
-    "asyncssh>=2.13.2",
+    "asyncssh>=2.19.0",
     "humanize>=4.8.0",
     "jinja2>=3.1.2",
     "toml>=0.10.2",
@@ -31,6 +31,9 @@ dependencies = [
     "backoff>=2.2.1",
     "more-itertools>=10.2.0",
     "cloudpickle>=3.0.0",
+    "aiofile>=3.9.0",
+    "wrapt>=1.16.0",
+    "httpx>=0.28.1",
 ]
 
 [project.urls]
@@ -57,6 +60,8 @@ dev-dependencies = [
 [tool.ruff]
 indent-width = 4
 line-length = 100
+
+[tool.ruff.lint]
 # Don't fix unused imports
 unfixable = ["F401"]
 # Ignore unused variables

{xmanager_slurm-0.4.5 → xmanager_slurm-0.4.7}/uv.lock

@@ -15,6 +15,18 @@ 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 = "aiofile"
+version = "3.9.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "caio" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/67/e2/d7cb819de8df6b5c1968a2756c3cb4122d4fa2b8fc768b53b7c9e5edb646/aiofile-3.9.0.tar.gz", hash = "sha256:e5ad718bb148b265b6df1b3752c4d1d83024b93da9bd599df74b9d9ffcf7919b", size = 17943 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/50/25/da1f0b4dd970e52bf5a36c204c107e11a0c6d3ed195eba0bfbc664c312b2/aiofile-3.9.0-py3-none-any.whl", hash = "sha256:ce2f6c1571538cbdfa0143b04e16b208ecb0e9cb4148e528af8a640ed51cc8aa", size = 19539 },
+]
+
 [[package]]
 name = "aiofiles"
 version = "24.1.0"
@@ -157,17 +169,16 @@ wheels = [
 
 [[package]]
 name = "anyio"
-version = "4.6.0"
+version = "3.7.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "exceptiongroup", marker = "python_full_version < '3.11'" },
     { name = "idna" },
     { name = "sniffio" },
-    { name = "typing-extensions", marker = "python_full_version < '3.11'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/78/49/f3f17ec11c4a91fe79275c426658e509b07547f874b14c1a526d86a83fc8/anyio-4.6.0.tar.gz", hash = "sha256:137b4559cbb034c477165047febb6ff83f390fc3b20bf181c1fc0a728cb8beeb", size = 170983 }
+sdist = { url = "https://files.pythonhosted.org/packages/28/99/2dfd53fd55ce9838e6ff2d4dac20ce58263798bd1a0dbe18b3a9af3fcfce/anyio-3.7.1.tar.gz", hash = "sha256:44a3c9aba0f5defa43261a8b3efb97891f2bd7d804e0e1f56419befa1adfc780", size = 142927 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/9e/ef/7a4f225581a0d7886ea28359179cb861d7fbcdefad29663fc1167b86f69f/anyio-4.6.0-py3-none-any.whl", hash = "sha256:c7d2e9d63e31599eeb636c8c5c03a7e108d73b345f064f1c19fdc87b79036a9a", size = 89631 },
+    { url = "https://files.pythonhosted.org/packages/19/24/44299477fe7dcc9cb58d0a57d5a7588d6af2ff403fdd2d47a246c91a3246/anyio-3.7.1-py3-none-any.whl", hash = "sha256:91dee416e570e92c64041bd18b900d1d6fa78dff7048769ce5ac5ddad004fbb5", size = 80896 },
 ]
 
 [[package]]
@@ -190,15 +201,15 @@ wheels = [
 
 [[package]]
 name = "asyncssh"
-version = "2.17.0"
+version = "2.19.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "cryptography" },
     { name = "typing-extensions" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/32/34/007ba4c65dc15ad394ff50d6ce272bcc028a8824ab1658c5fe5c480be515/asyncssh-2.17.0.tar.gz", hash = "sha256:3b159c105aa388c1e2245c4faf483f540ada8cad99402281119100166e5edb3c", size = 515658 }
+sdist = { url = "https://files.pythonhosted.org/packages/75/85/c723aa7dd69570a31b6638c1405e659712f18f569f280da2da27989445d3/asyncssh-2.19.0.tar.gz", hash = "sha256:723dead4d068b558708dc66a4ca7e7a93a813aa9416036eccb9af4c03ae2cf30", size = 533702 }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/58/ed/c4b731b4b5c76c2d21e7411866e8e10653bc9c1fcc7a9a5181852b6e6de6/asyncssh-2.17.0-py3-none-any.whl", hash = "sha256:70bfe4081793255b0115f7a72efd92d5bad8562639302531223bb00188db9747", size = 363736 },
+    { url = "https://files.pythonhosted.org/packages/95/8f/f0749af566fa39204f6380473e6f9632b804daeb0ecb24cc7de1fc9f2717/asyncssh-2.19.0-py3-none-any.whl", hash = "sha256:bb82ac30ff0cb4393fbaf1114e606ad7a4f13d6c4bdaed423c033ee26b455228", size = 372704 },
 ]
 
 [[package]]
@@ -249,6 +260,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a4/07/14f8ad37f2d12a5ce41206c21820d8cb6561b728e51fad4530dff0552a67/cachetools-5.5.0-py3-none-any.whl", hash = "sha256:02134e8439cdc2ffb62023ce1debca2944c3f289d66bb17ead3ab3dede74b292", size = 9524 },
 ]
 
+[[package]]
+name = "caio"
+version = "0.9.17"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/88/cf/59b868909a85ad9eca985ad5bbeb3c0a8cd435e50ae12def770d16911753/caio-0.9.17.tar.gz", hash = "sha256:8f30511526814d961aeef389ea6885273abe6c655f1e08abbadb95d12fdd9b4f", size = 25001 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/db/c3/17bc41b7c795d91d58ee7a70ad98e23f1ba0d50bdeadd82173bb02cddc8e/caio-0.9.17-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:3f69395fdd45c115b2ef59732e3c8664722a2b51de2d6eedb3d354b2f5f3be3c", size = 37849 },
+    { url = "https://files.pythonhosted.org/packages/bf/3f/0ae9f69deb3dc96b20bf084cc262b438a154d5de08a064628272713ff239/caio-0.9.17-cp310-cp310-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3028b746e9ec7f6d6ebb386a7fd8caf0eebed5d6e6b4f18c8ef25861934b1673", size = 79417 },
+    { url = "https://files.pythonhosted.org/packages/52/44/a79c7004a9562a176d78437816c736c34ab9fb6233a4b8164eb25628a09c/caio-0.9.17-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:079730a353bbde03796fab681e969472eace09ffbe5000e584868a7fe389ba6f", size = 37842 },
+    { url = "https://files.pythonhosted.org/packages/57/e2/1d04e506a5fd735856f0bb95b4d03b800947bd43c98193ee57d37070a51e/caio-0.9.17-cp311-cp311-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:549caa51b475877fe32856a26fe937366ae7a1c23a9727005b441db9abb12bcc", size = 79985 },
+    { url = "https://files.pythonhosted.org/packages/c3/b8/37dcee4bc4fae1701a86373a297bbca797f6b7bfe5f85993a11049649c63/caio-0.9.17-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:0ddb253b145a53ecca76381677ce465bc5efeaecb6aaf493fac43ae79659f0fb", size = 37909 },
+    { url = "https://files.pythonhosted.org/packages/80/f5/5e993120daeb4ec084f5f84c118bbd48b65379f63ed56919bf224e0eab42/caio-0.9.17-cp312-cp312-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b3e320b0ea371c810359934f8e8fe81777c493cc5fb4d41de44277cbe7336e74", size = 81709 },
+    { url = "https://files.pythonhosted.org/packages/58/72/3f4895adb1d23b0ac1d8afc748405a2ad3c77d8d0f23b05a64ff583c11e5/caio-0.9.17-py3-none-any.whl", hash = "sha256:c55d4dc6b3a36f93237ecd6360e1c131c3808bc47d4191a130148a99b80bb311", size = 19062 },
+]
+
 [[package]]
 name = "certifi"
 version = "2024.8.30"
@@ -374,7 +400,7 @@ name = "click"
 version = "8.1.7"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "colorama", marker = "platform_system == 'Windows'" },
+    { name = "colorama", marker = "sys_platform == 'win32'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 }
 wheels = [
@@ -926,6 +952,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/95/04/ff642e65ad6b90db43e668d70ffb6736436c7ce41fcc549f4e9472234127/h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761", size = 58259 },
 ]
 
+[[package]]
+name = "httpcore"
+version = "1.0.7"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "certifi" },
+    { name = "h11" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/6a/41/d7d0a89eb493922c37d343b607bc1b5da7f5be7e383740b4753ad8943e90/httpcore-1.0.7.tar.gz", hash = "sha256:8551cb62a169ec7162ac7be8d4817d561f60e08eaa485234898414bb5a8a0b4c", size = 85196 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/87/f5/72347bc88306acb359581ac4d52f23c0ef445b57157adedb9aee0cd689d2/httpcore-1.0.7-py3-none-any.whl", hash = "sha256:a3fff8f43dc260d5bd363d9f9cf1830fa3a458b332856f34282de498ed420edd", size = 78551 },
+]
+
 [[package]]
 name = "httplib2"
 version = "0.22.0"
@@ -938,6 +977,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a8/6c/d2fbdaaa5959339d53ba38e94c123e4e84b8fbc4b84beb0e70d7c1608486/httplib2-0.22.0-py3-none-any.whl", hash = "sha256:14ae0a53c1ba8f3d37e9e27cf37eabb0fb9980f435ba405d546948b009dd64dc", size = 96854 },
 ]
 
+[[package]]
+name = "httpx"
+version = "0.28.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "anyio" },
+    { name = "certifi" },
+    { name = "httpcore" },
+    { name = "idna" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b1/df/48c586a5fe32a0f01324ee087459e112ebb7224f646c0b5023f5e79e9956/httpx-0.28.1.tar.gz", hash = "sha256:75e98c5f16b0f35b567856f597f06ff2270a374470a5c2392242528e3e3e42fc", size = 141406 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl", hash = "sha256:d909fcccc110f8c7faf814ca82a9a4d816bc5a6dbfea25d6591d6985b8ba59ad", size = 73517 },
+]
+
 [[package]]
 name = "humanize"
 version = "4.10.0"
@@ -2196,18 +2250,21 @@ wheels = [
 
 [[package]]
 name = "xmanager-slurm"
-version = "0.4.5"
+version = "0.4.7"
 source = { editable = "." }
 dependencies = [
+    { name = "aiofile" },
     { name = "asyncssh" },
     { name = "backoff" },
     { name = "cloudpickle" },
+    { name = "httpx" },
     { name = "humanize" },
     { name = "immutabledict" },
     { name = "jinja2" },
     { name = "more-itertools" },
     { name = "rich" },
     { name = "toml" },
+    { name = "wrapt" },
     { name = "xmanager" },
 ]
 
@@ -2228,15 +2285,18 @@ dev = [
 
 [package.metadata]
 requires-dist = [
-    { name = "asyncssh", specifier = ">=2.13.2" },
+    { name = "aiofile", specifier = ">=3.9.0" },
+    { name = "asyncssh", specifier = ">=2.19.0" },
     { name = "backoff", specifier = ">=2.2.1" },
     { name = "cloudpickle", specifier = ">=3.0.0" },
+    { name = "httpx", specifier = ">=0.28.1" },
     { name = "humanize", specifier = ">=4.8.0" },
     { name = "immutabledict", specifier = ">=3.0.0" },
     { name = "jinja2", specifier = ">=3.1.2" },
     { name = "more-itertools", specifier = ">=10.2.0" },
     { name = "rich", specifier = ">=13.5.2" },
     { name = "toml", specifier = ">=0.10.2" },
+    { name = "wrapt", specifier = ">=1.16.0" },
     { name = "xmanager", specifier = ">=0.5.0" },
 ]
 

{xmanager_slurm-0.4.5 → xmanager_slurm-0.4.7}/xm_slurm/__init__.py

@@ -3,7 +3,6 @@ import logging
 from xm_slurm.executables import Dockerfile, DockerImage
 from xm_slurm.executors import Slurm, SlurmSpec
 from xm_slurm.experiment import (
-    Artifact,
     SlurmExperiment,
     create_experiment,
     get_current_experiment,
@@ -25,7 +24,6 @@ logging.getLogger("asyncssh").setLevel(logging.WARN)
 logging.getLogger("httpx").setLevel(logging.WARN)
 
 __all__ = [
-    "Artifact",
     "conda_container",
     "create_experiment",
     "docker_container",

xmanager_slurm-0.4.7/xm_slurm/api/__init__.py

@@ -0,0 +1,33 @@
+import functools
+import logging
+import os
+
+from xm_slurm.api import models
+from xm_slurm.api.abc import XManagerAPI
+from xm_slurm.api.sqlite import client as sqlite_client
+from xm_slurm.api.web import client as web_client
+
+logger = logging.getLogger(__name__)
+
+
+@functools.cache
+def client() -> XManagerAPI:
+    backend = os.environ.get("XM_SLURM_API_BACKEND", "sqlite").lower()
+    match backend:
+        case "rest":
+            if "XM_SLURM_REST_API_BASE_URL" not in os.environ:
+                raise ValueError("XM_SLURM_REST_API_BASE_URL is not set")
+            if "XM_SLURM_REST_API_TOKEN" not in os.environ:
+                raise ValueError("XM_SLURM_REST_API_TOKEN is not set")
+
+            return web_client.XManagerWebAPI(
+                base_url=os.environ["XM_SLURM_REST_API_BASE_URL"],
+                token=os.environ["XM_SLURM_REST_API_TOKEN"],
+            )
+        case "sqlite":
+            return sqlite_client.XManagerSqliteAPI()
+        case _:
+            raise ValueError(f"Invalid XM_SLURM_API_BACKEND: {backend}")
+
+
+__all__ = ["client", "XManagerAPI", "models"]

xmanager_slurm-0.4.7/xm_slurm/api/abc.py

@@ -0,0 +1,65 @@
+import abc
+
+from xm_slurm.api import models
+
+
+class XManagerAPI(abc.ABC):
+    @abc.abstractmethod
+    def get_experiment(self, xid: int) -> models.Experiment:
+        pass
+
+    @abc.abstractmethod
+    def delete_experiment(self, experiment_id: int) -> None:
+        pass
+
+    @abc.abstractmethod
+    def insert_experiment(self, experiment: models.ExperimentPatch) -> int:
+        pass
+
+    @abc.abstractmethod
+    def update_experiment(
+        self, experiment_id: int, experiment_patch: models.ExperimentPatch
+    ) -> None:
+        pass
+
+    @abc.abstractmethod
+    def insert_job(self, experiment_id: int, work_unit_id: int, job: models.SlurmJob) -> None:
+        pass
+
+    @abc.abstractmethod
+    def insert_work_unit(self, experiment_id: int, work_unit: models.WorkUnitPatch) -> None:
+        pass
+
+    @abc.abstractmethod
+    def update_work_unit(
+        self, experiment_id: int, work_unit_id: int, patch: models.ExperimentUnitPatch
+    ) -> None:
+        pass
+
+    @abc.abstractmethod
+    def delete_work_unit_artifact(self, experiment_id: int, work_unit_id: int, name: str) -> None:
+        pass
+
+    @abc.abstractmethod
+    def insert_work_unit_artifact(
+        self, experiment_id: int, work_unit_id: int, artifact: models.Artifact
+    ) -> None:
+        pass
+
+    @abc.abstractmethod
+    def delete_experiment_artifact(self, experiment_id: int, name: str) -> None:
+        pass
+
+    @abc.abstractmethod
+    def insert_experiment_artifact(self, experiment_id: int, artifact: models.Artifact) -> None:
+        pass
+
+    @abc.abstractmethod
+    def insert_experiment_config_artifact(
+        self, experiment_id: int, artifact: models.ConfigArtifact
+    ) -> None:
+        pass
+
+    @abc.abstractmethod
+    def delete_experiment_config_artifact(self, experiment_id: int, name: str) -> None:
+        pass

xmanager_slurm-0.4.7/xm_slurm/api/models.py

@@ -0,0 +1,70 @@
+import dataclasses
+import enum
+import typing as tp
+
+
+class ExperimentUnitRole(enum.Enum):
+    WORK_UNIT = enum.auto()
+    AUX_UNIT = enum.auto()
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class ExperimentPatch:
+    title: str | None = None
+    description: str | None = None
+    note: str | None = None
+    tags: list[str] | None = None
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class SlurmJob:
+    name: str
+    slurm_job_id: str
+    slurm_ssh_config: str
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class Artifact:
+    name: str
+    uri: str
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class ConfigArtifact:
+    name: tp.Literal["GRAPHVIZ", "PYTHON"]
+    uri: str
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class ExperimentUnit:
+    identity: str
+    args: str | None = None
+    jobs: list[SlurmJob] = dataclasses.field(default_factory=list)
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class ExperimentUnitPatch:
+    identity: str | None = None
+    args: str | None = None
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class WorkUnit(ExperimentUnit):
+    wid: int
+    artifacts: list[Artifact] = dataclasses.field(default_factory=list)
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class WorkUnitPatch(ExperimentUnitPatch):
+    wid: int
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True)
+class Experiment:
+    title: str
+    description: str | None
+    note: str | None
+    tags: list[str] | None
+
+    work_units: list[WorkUnit] = dataclasses.field(default_factory=list)
+    artifacts: list[Artifact] = dataclasses.field(default_factory=list)