xmanager-slurm 0.4.10__py3-none-any.whl → 0.4.12__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/config.py

@@ -62,9 +62,14 @@ class SlurmSSHConfig:
         if self.host_public_key is None:
             return None
 
-        return asyncssh.import_known_hosts(
-            f"[{self.host}]:{self.port} {self.host_public_key.algorithm} {self.host_public_key.key}"
+        hostname = f"{self.host}"
+        if self.port is not None and self.port != asyncssh.DEFAULT_PORT:
+            hostname = f"[{hostname}]:{self.port}"
+
+        known_hosts = asyncssh.import_known_hosts(
+            f"{hostname} {self.host_public_key.algorithm} {self.host_public_key.key}"
         )
+        return known_hosts
 
     @functools.cached_property
     def config(self) -> asyncssh.config.SSHConfig:
@@ -117,9 +122,9 @@ class SlurmSSHConfig:
     def connection_options(self) -> asyncssh.SSHClientConnectionOptions:
         options = asyncssh.SSHClientConnectionOptions(
             config=None,
-            kbdint_auth=False,
             disable_trivial_auth=True,
             known_hosts=self.known_hosts,
+            server_host_key_algs=self.host_public_key.algorithm if self.host_public_key else None,
         )
         options.prepare(last_config=self.config)
         return options

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.10.dist-info → xmanager_slurm-0.4.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.10
+Version: 0.4.12
 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.10.dist-info → xmanager_slurm-0.4.12.dist-info}/RECORD

@@ -1,6 +1,6 @@
 xm_slurm/__init__.py,sha256=WgRn9HDYa5H3sfIH-HZu33liBOh98jM4GqcR349RaSY,1086
 xm_slurm/batching.py,sha256=GbKBsNz9w8gIc2fHLZpslC0e4K9YUfLXFHmjduRRCfQ,4385
-xm_slurm/config.py,sha256=PvdLG6fSNfbABbtwELbnytx85vWLPtyVfEJtoQsLf94,7084
+xm_slurm/config.py,sha256=i9WuxjfGBgVoHHDNk3BvO7LCwpBTJeRnOiFvTU-FHrk,7330
 xm_slurm/console.py,sha256=UpMqeJ0C8i0pkue1AHnnyyX0bFJ9zZeJ7HBR6yhuA8A,54
 xm_slurm/constants.py,sha256=zefVtlFdflgSolie5g_rVxWV-Zpydxapchm3y0a2FDc,999
 xm_slurm/dependencies.py,sha256=-5gN_tpfs3dOA7H5_MIHO2ratb7F5Pm_yjkR5rZcgI8,6421
@@ -11,15 +11,15 @@ 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
@@ -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.10.dist-info/METADATA,sha256=jsnzz5ge4NpBIC_V3R3KYVm9AyG3cDH96YDUZ-2ZeDc,1043
-xmanager_slurm-0.4.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xmanager_slurm-0.4.10.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
-xmanager_slurm-0.4.10.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
-xmanager_slurm-0.4.10.dist-info/RECORD,,
+xmanager_slurm-0.4.12.dist-info/METADATA,sha256=ttPSMz6bQs8BEQVCurQ9lFWV7MuzLOHWU7lISOOoufA,1007
+xmanager_slurm-0.4.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xmanager_slurm-0.4.12.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.12.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.12.dist-info/RECORD,,