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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: xmanager-slurm
-Version: 0.4.1
+Version: 0.4.2
 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.1 → xmanager_slurm-0.4.2}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 [project]
 name = "xmanager-slurm"
 description = "Slurm backend for XManager."
-version = "0.4.1"
+version = "0.4.2"
 # readme = "README.md"
 requires-python = ">=3.10"
 license = { text = "MIT" }

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

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

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

@@ -31,23 +31,6 @@ 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(
@@ -77,10 +60,11 @@ def mila(
             resources.ResourceType.A100: "a100",
             resources.ResourceType.A100_80GIB: "a100l",
             resources.ResourceType.A6000: "a6000",
+            resources.ResourceType.L40S: "l40s",
+            resources.ResourceType.H100: "h100",
         },
         features={
             resources.FeatureType.NVIDIA_MIG: "mig",
             resources.FeatureType.NVIDIA_NVLINK: "nvlink",
         },
-        validate=validate,
     )

{xmanager_slurm-0.4.1 → xmanager_slurm-0.4.2}/xm_slurm/resources.py

@@ -33,6 +33,7 @@ class ResourceType(enum.IntEnum):
     A6000 = 1033
 
     H100 = 1040
+    L40S = 1041
 
 
 AcceleratorType = set([
@@ -47,6 +48,7 @@ AcceleratorType = set([
     ResourceType.A5000,
     ResourceType.A6000,
     ResourceType.H100,
+    ResourceType.L40S,
 ])
 
 assert AcceleratorType | {

{xmanager_slurm-0.4.1 → xmanager_slurm-0.4.2}/xm_slurm/scripts/cli.py

@@ -1,4 +1,5 @@
 import argparse
+import sys
 
 from xmanager import xm
 
@@ -8,13 +9,26 @@ from xm_slurm.console import console
 
 async def logs(
     experiment_id: int,
-    wid: int,
     *,
+    wid: int | None,
+    identity: str | None,
     follow: bool = True,
     num_lines: int = 10,
     block_size: int = 1024,
 ):
-    wu = xm_slurm.get_experiment(experiment_id).work_units()[wid]
+    xp = xm_slurm.get_experiment(experiment_id)
+
+    if wid is not None:
+        wu = xp.work_units()[wid]
+    elif identity is not None:
+        wu = xp._get_work_unit_by_identity(identity)
+        if wu is None:
+            console.print(f"[red]Work Unit with identity {identity} not found.[/red]")
+            sys.exit(1)
+    else:
+        raise ValueError("Must specify either wid or identity.")
+    assert wu is not None
+
     async for log in wu.logs(num_lines=num_lines, block_size=block_size, wait=True, follow=follow):
         console.print(log, end="\n")
 
@@ -26,7 +40,12 @@ async def main():
 
     logs_parser = subparsers.add_parser("logs", help="Display logs for a specific experiment.")
     logs_parser.add_argument("xid", type=int, help="Experiment ID.")
-    logs_parser.add_argument("wid", type=int, help="Work Unit ID.")
+
+    # Create a mutually exclusive group for wid and identity
+    group = logs_parser.add_mutually_exclusive_group()
+    group.add_argument("--wid", type=int, help="Work Unit ID.")
+    group.add_argument("--identity", type=str, help="Work Unit identity.")
+
     logs_parser.add_argument(
         "-n",
         "--n-lines",
@@ -45,7 +64,13 @@ async def main():
     args = parser.parse_args()
     match args.subcommand:
         case "logs":
-            await logs(args.xid, args.wid, follow=args.follow, num_lines=args.n_lines)
+            await logs(
+                args.xid,
+                wid=args.wid,
+                identity=args.identity,
+                follow=args.follow,
+                num_lines=args.n_lines,
+            )
 
 
 if __name__ == "__main__":

{xmanager_slurm-0.4.1 → xmanager_slurm-0.4.2}/xm_slurm/templates/slurm/runtimes/apptainer.bash.j2

@@ -6,6 +6,7 @@ export {{ key }}="{{ value }}"
 {% else %}
 export APPTAINERENV_{{ key }}="{{ value }}"
 export SINGULARITYENV_{{ key }}="{{ value }}"
+export {{ key }}="{{ value }}"
 {% endif %}
 {% endfor %}
 {%- endmacro %}