xmanager-slurm 0.4.11__tar.gz → 0.4.12__tar.gz

{xmanager_slurm-0.4.11 → xmanager_slurm-0.4.12}/.vscode/settings.json

@@ -2,4 +2,5 @@
   "search.exclude": {
     "uv.lock": true,
   },
+  "python.analysis.typeCheckingMode": "basic",
 }

{xmanager_slurm-0.4.11 → xmanager_slurm-0.4.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.11
+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.11 → xmanager_slurm-0.4.12}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "xmanager-slurm"
 description = "Slurm backend for XManager."
-version = "0.4.11"
+version = "0.4.12"
 # readme = "README.md"
 requires-python = ">=3.10"
 license = { text = "MIT" }
@@ -27,7 +27,6 @@ dependencies = [
     "jinja2>=3.1.2",
     "toml>=0.10.2",
     "rich>=13.5.2",
-    "immutabledict>=3.0.0",
     "backoff>=2.2.1",
     "more-itertools>=10.2.0",
     "cloudpickle>=3.0.0",

{xmanager_slurm-0.4.11 → xmanager_slurm-0.4.12}/uv.lock

@@ -2251,7 +2251,7 @@ wheels = [
 
 [[package]]
 name = "xmanager-slurm"
-version = "0.4.11"
+version = "0.4.12"
 source = { editable = "." }
 dependencies = [
     { name = "aiofile" },
@@ -2260,7 +2260,6 @@ dependencies = [
     { name = "cloudpickle" },
     { name = "httpx" },
     { name = "humanize" },
-    { name = "immutabledict" },
     { name = "jinja2" },
     { name = "more-itertools" },
     { name = "rich" },
@@ -2292,7 +2291,6 @@ requires-dist = [
     { 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" },

{xmanager_slurm-0.4.11 → xmanager_slurm-0.4.12}/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,

{xmanager_slurm-0.4.11 → xmanager_slurm-0.4.12}/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.
 

{xmanager_slurm-0.4.11 → xmanager_slurm-0.4.12}/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:

{xmanager_slurm-0.4.11 → xmanager_slurm-0.4.12}/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: