xmanager-slurm 0.4.11__py3-none-any.whl → 0.4.13__py3-none-any.whl

xm_slurm/api/sqlite/client.py

@@ -153,7 +153,7 @@ class XManagerSqliteAPI(XManagerAPI):
             artifacts = [
                 models.Artifact(name=artifact.name, uri=artifact.uri)
                 for artifact in experiment.artifacts
-            ] + [models.Artifact(name=config.name, uri=config.uri) for config in experiment.configs]
+            ]
 
             return models.Experiment(
                 title=experiment.title,

xm_slurm/contrib/clusters/drac.py

@@ -4,7 +4,7 @@ from typing import Literal
 from xm_slurm import config
 from xm_slurm.resources import FeatureType, ResourceType
 
-__all__ = ["narval", "beluga", "cedar", "graham"]
+__all__ = ["beluga", "cedar", "fir", "graham", "narval"]
 
 
 def _drac_cluster(
@@ -121,6 +121,36 @@ def beluga(
     )
 
 
+def rorqual(
+    *,
+    user: str | None = None,
+    account: str | None = None,
+    proxy: Literal["submission-host"] | str | None = None,
+    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """DRAC Beluga Cluster (https://docs.alliancecan.ca/wiki/Rorqual/en)."""
+    modules = []
+    if proxy != "submission-host":
+        modules.append("httpproxy")
+
+    return _drac_cluster(
+        name="rorqual",
+        host="robot.rorqual.alliancecan.ca",
+        host_public_key=config.PublicKey(
+            "ssh-ed25519", "AAAAC3NzaC1lZDI1NTE5AAAAINME5e9bifKZbuKKOQSpe3xrvC4g1b0QLMYj+AXBQGJe"
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        proxy=proxy,
+        modules=modules,
+        resources={ResourceType.H100: "h100"},
+        features={
+            FeatureType.NVIDIA_NVLINK: "nvlink",
+        },
+    )
+
+
 def cedar(
     *,
     user: str | None = None,
@@ -146,6 +176,27 @@ def cedar(
     )
 
 
+def fir(
+    *,
+    user: str | None = None,
+    account: str | None = None,
+    mounts: dict[os.PathLike[str] | str, os.PathLike[str] | str] | None = None,
+) -> config.SlurmClusterConfig:
+    """DRAC Fir Cluster (https://docs.alliancecan.ca/wiki/Fir/en)."""
+    return _drac_cluster(
+        name="fir",
+        host="robot.fir.alliancecan.ca",
+        host_public_key=config.PublicKey(
+            "ssh-ed25519",
+            "AAAAC3NzaC1lZDI1NTE5AAAAIJtenyJz+inwobvlJntWYFNu+ANcVWNcOHRKcEN6zmDo",
+        ),
+        user=user,
+        account=account,
+        mounts=mounts,
+        resources={ResourceType.H100: "h100"},
+    )
+
+
 def graham(
     *,
     user: str | None = None,

xm_slurm/packageables.py

@@ -3,7 +3,6 @@ import pathlib
 import sys
 from typing import Literal, Mapping, Sequence
 
-import immutabledict
 from xmanager import xm
 
 from xm_slurm import job_blocks, utils
@@ -15,7 +14,7 @@ def docker_image(
     *,
     image: str,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a packageable for a pre-built Docker image.
 
@@ -30,7 +29,7 @@ def docker_image(
         executor_spec=SlurmSpec(),
         executable_spec=DockerImage(image=image),
         args=xm.SequentialArgs.from_collection(args),
-        env_vars=dict(env_vars),
+        env_vars=env_vars or {},
     )
 
 
@@ -41,11 +40,11 @@ def docker_container(
     context: pathlib.Path | None = None,
     target: str | None = None,
     ssh: Sequence[str] | Literal[True] | None = None,
-    build_args: Mapping[str, str] = immutabledict.immutabledict(),
+    build_args: Mapping[str, str] | None = None,
     cache_from: str | Sequence[str] | None = None,
-    labels: Mapping[str, str] = immutabledict.immutabledict(),
+    labels: Mapping[str, str] | None = None,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a Docker container packageable from a dockerfile.
 
@@ -89,12 +88,12 @@ def docker_container(
             context=context,
             target=target,
             ssh=ssh,
-            build_args=build_args,
+            build_args=build_args or {},
             cache_from=cache_from,
-            labels=labels,
+            labels=labels or {},
         ),
         args=xm.SequentialArgs.from_collection(args),
-        env_vars=dict(env_vars),
+        env_vars=env_vars or {},
     )
 
 
@@ -108,10 +107,10 @@ def python_container(
     extra_system_packages: Sequence[str] = (),
     extra_python_packages: Sequence[str] = (),
     cache_from: str | Sequence[str] | None = None,
-    labels: Mapping[str, str] = immutabledict.immutabledict(),
+    labels: Mapping[str, str] | None = None,
     ssh: Sequence[str] | Literal[True] | None = None,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a Python container from a base image using pip from a `requirements.txt` file.
 
@@ -183,10 +182,10 @@ def mamba_container(
     environment: pathlib.Path | None = None,
     base_image: str = "gcr.io/distroless/base-debian10",
     cache_from: str | Sequence[str] | None = None,
-    labels: Mapping[str, str] = immutabledict.immutabledict(),
+    labels: Mapping[str, str] | None = None,
     ssh: Sequence[str] | Literal[True] | None = None,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a Conda container from a base image using mamba from a `environment.yml` file.
 
@@ -253,10 +252,10 @@ def uv_container(
     extra_system_packages: Sequence[str] = (),
     extra_python_packages: Sequence[str] = (),
     cache_from: str | Sequence[str] | None = None,
-    labels: Mapping[str, str] = immutabledict.immutabledict(),
+    labels: Mapping[str, str] | None = None,
     ssh: Sequence[str] | Literal[True] | None = None,
     args: xm.UserArgs | None = None,
-    env_vars: Mapping[str, str] = immutabledict.immutabledict(),
+    env_vars: Mapping[str, str] | None = None,
 ) -> xm.Packageable:
     """Creates a Python container from a base image using uv from a `uv.lock` file.
 

xm_slurm/resources.py

@@ -3,8 +3,6 @@ import itertools
 import math
 from typing import Mapping
 
-import immutabledict
-
 from xm_slurm import config
 
 
@@ -76,7 +74,7 @@ class JobRequirements:
     def __init__(
         self,
         *,
-        resources: Mapping[ResourceType | str, ResourceQuantity] = immutabledict.immutabledict(),
+        resources: Mapping[ResourceType | str, ResourceQuantity] | None = None,
         replicas: int = 1,
         location: str | None = None,
         cluster: config.SlurmClusterConfig | None = None,
@@ -87,6 +85,9 @@ class JobRequirements:
         self.accelerator = None
         self.cluster = cluster
 
+        if resources is None:
+            resources = {}
+
         self.task_requirements: dict[ResourceType | str, ResourceQuantity] = {}
         for resource_name, value in itertools.chain(resources.items(), kw_resources.items()):
             match resource_name:

xm_slurm/scripts/cli.py

@@ -19,7 +19,13 @@ async def logs(
     xp = xm_slurm.get_experiment(experiment_id)
 
     if wid is not None:
-        wu = xp.work_units()[wid]
+        wus = xp.work_units()
+        if wid not in wus:
+            console.print(
+                f"[red]Work Unit ID {wid} not found for experiment {experiment_id} with {len(wus)} work units.[/red]"
+            )
+            sys.exit(1)
+        wu = wus[wid]
     elif identity is not None:
         wu = xp._get_work_unit_by_identity(identity)
         if wu is None:

{xmanager_slurm-0.4.11.dist-info → xmanager_slurm-0.4.13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.11
+Version: 0.4.13
 Summary: Slurm backend for XManager.
 Project-URL: GitHub, https://github.com/jessefarebro/xm-slurm
 Author-email: Jesse Farebrother <jfarebro@cs.mcgill.ca>
@@ -20,7 +20,6 @@ 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

{xmanager_slurm-0.4.11.dist-info → xmanager_slurm-0.4.13.dist-info}/RECORD

@@ -11,19 +11,19 @@ xm_slurm/experiment.py,sha256=94r0mhtUPUzw4eaUEz0kpsufC25wEGqlDhV4Fcr1ukY,39883
 xm_slurm/filesystem.py,sha256=4rKtq3t-KDgxJbSGt6JVyRJT_3lCN_vIKTcwKHpTo3I,4389
 xm_slurm/job_blocks.py,sha256=_F8CKCs5BQFj40a2-mjG71HfacvWoBXBDPDKEaKTbXc,616
 xm_slurm/metadata_context.py,sha256=mksVRbVUuistL1uE7TC-fkW-Y69On52jN_svP1e1kiQ,7841
-xm_slurm/packageables.py,sha256=fPUvqF2IvJ2Hn6hodDdQwtx1Ze3sJ8U-BUbxDHauW-g,12398
-xm_slurm/resources.py,sha256=tET3TPOQ8nXYE_SxAs2fiHt9UKJsCLW1vFktJTH0xG4,5722
+xm_slurm/packageables.py,sha256=K6vNhLvASdnqsc8vXlT3h9cObJpC9Rbw93pUBvBwapQ,12209
+xm_slurm/resources.py,sha256=T7uje3E6oWbZSrsxykgW-40DE-Bvw_NWDM2qXbw2rgI,5740
 xm_slurm/status.py,sha256=WTWiDHi-ZHtwHRnDP0cGa-27zTSm6LkA-GCKsN-zBgg,6916
 xm_slurm/types.py,sha256=TsVykDm-LazVkrjeJrTwCMs4Q8APKhy7BTk0yKIhFNg,805
 xm_slurm/utils.py,sha256=xtFvktaxr0z65sTdu6HhOVfyo0OAB9t-EYXWcYrQQEU,5958
 xm_slurm/api/__init__.py,sha256=cyao3LZ3uLftu1wIv1aN7Qvsl6gYzYpkxeehTHZ0fA8,1089
 xm_slurm/api/abc.py,sha256=-lS2OndnOuEiwNdr8ccQKkwMd1iDmKMmkBOSTvo5H5w,1816
 xm_slurm/api/models.py,sha256=_INVh0j-4-rRs0WASyg4fNB6NF1L1nUeGgQ6-XnbwsM,1610
-xm_slurm/api/sqlite/client.py,sha256=WykSIO7b14rRLy9qebbkiLKXy7EHU61jtoebLX17HMM,14124
+xm_slurm/api/sqlite/client.py,sha256=jAesCKDuYwnNcAxwJk_1b1TB8cT_QGbSjo1UE3mZjEQ,14037
 xm_slurm/api/web/client.py,sha256=uO67Y7fnQ-w__Vm_A5BEuy7Qi8wQcWk3vIsBGEBkyfk,6261
 xm_slurm/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xm_slurm/contrib/clusters/__init__.py,sha256=XFCVnkThiU3_8uA_tUgDByOBanXNHrxDvfmuptmQ2KE,2214
-xm_slurm/contrib/clusters/drac.py,sha256=ViLYerYBMSuZXnWVbz9RDIPPV7JA8BgBpgTfj1wPP28,5881
+xm_slurm/contrib/clusters/drac.py,sha256=vY3dxrNUk12H9Gq-tuCcqo2YcdTGq-4LJnQF6DzD4_k,7431
 xm_slurm/experimental/parameter_controller.py,sha256=b5LfglHV307F6QcPrHeZX5GJBtyOK9aQydke_SZ3Wto,8457
 xm_slurm/packaging/__init__.py,sha256=dh307yLpUT9KN7rJ1e9fYC6hegGKfZcGboUq9nGpDVQ,233
 xm_slurm/packaging/docker.py,sha256=-DWcB9qqbeHmIEqyfF0-v6xOT25ae90u2x-QZ7kluOw,13579
@@ -31,7 +31,7 @@ xm_slurm/packaging/registry.py,sha256=GrdmQg9MgSo38OiqOzMKWSkQyBuyryOfc3zcdgZ4CU
 xm_slurm/packaging/router.py,sha256=yPbdA9clrhly97cLgDsSRZG2LZRKE-oz8Hhdb7WtYqk,2070
 xm_slurm/packaging/utils.py,sha256=6EAb17zKQQeuyNY2EV9AoW1RvnDGrQwmIT9wtQEsC4c,632
 xm_slurm/scripts/_cloudpickle.py,sha256=dlJYf2SceOuUn8wi-ozuoYAQg71wqD2MUVOUCyOwWIY,647
-xm_slurm/scripts/cli.py,sha256=nnfNF2FAVLrm-4KuUooFfSWxX7hOjmXpI7k5f3Sn5us,2349
+xm_slurm/scripts/cli.py,sha256=zzsQpvkx9VThAeQPM34iDK9wAWfCVCIIvLNI12UaMhw,2577
 xm_slurm/templates/docker/docker-bake.hcl.j2,sha256=7qSJl2VN5poz-Hh8Gjo7--qR-k3lmfGtBu2mNbfG2uA,1499
 xm_slurm/templates/docker/mamba.Dockerfile,sha256=Sgxr5IA5T-pT1Shumb5k3JngoG4pgCdBXjzqslFJdZI,753
 xm_slurm/templates/docker/python.Dockerfile,sha256=U4b4QVkopckQ0o9jJIE7d_M6TvExEYlYDirNwCoZ7W4,865
@@ -44,8 +44,8 @@ xm_slurm/templates/slurm/fragments/monitor.bash.j2,sha256=BJ1brSjhESOe9VX_OYaPyy
 xm_slurm/templates/slurm/fragments/proxy.bash.j2,sha256=VJLglZo-Nvx9R-qe3rHTxr07CylTQ6Z9NwBzvIpAZrA,814
 xm_slurm/templates/slurm/runtimes/apptainer.bash.j2,sha256=lE2EWVCK2O-n08RL4_MJYIikVTvODjcYKuv7Eh73Q2w,1932
 xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=3j7K5eyXt_WhXK0EoMlxnhlmFVJ2JyxRKbsMRaDqzSs,1148
-xmanager_slurm-0.4.11.dist-info/METADATA,sha256=f7oAbg0CGJLT3iQUERwjnB4IdhrDzgDoYEL2izYTA0o,1043
-xmanager_slurm-0.4.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xmanager_slurm-0.4.11.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
-xmanager_slurm-0.4.11.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
-xmanager_slurm-0.4.11.dist-info/RECORD,,
+xmanager_slurm-0.4.13.dist-info/METADATA,sha256=E8QV408HDZOHu4OhsV8Xq4pon_1A5Qi3RLABBAS9iDM,1007
+xmanager_slurm-0.4.13.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xmanager_slurm-0.4.13.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.13.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.13.dist-info/RECORD,,