xmanager-slurm 0.4.1__py3-none-any.whl → 0.4.2__py3-none-any.whl

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,
     )

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 | {

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__":

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 %}

{xmanager_slurm-0.4.1.dist-info → xmanager_slurm-0.4.2.dist-info}/METADATA

@@ -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.dist-info → xmanager_slurm-0.4.2.dist-info}/RECORD

@@ -11,12 +11,12 @@ xm_slurm/executors.py,sha256=fMtxGUCi4vEKmb_p4JEpqPUTh7L_f1LcR_TamMLAWNg,4667
 xm_slurm/experiment.py,sha256=trHapcYxPNKofzSqu7KZawML59tZ8FVjoEZYe2Wal7w,44521
 xm_slurm/job_blocks.py,sha256=_F8CKCs5BQFj40a2-mjG71HfacvWoBXBDPDKEaKTbXc,616
 xm_slurm/packageables.py,sha256=YZFTL6UWx9A_zyztTy1THUlj3pW1rA0cBPHJxD1LOJk,12884
-xm_slurm/resources.py,sha256=EaYDATVudrEDPKKdSZoWgfqPiidc6DMjIctmzLQmiH0,5683
+xm_slurm/resources.py,sha256=tET3TPOQ8nXYE_SxAs2fiHt9UKJsCLW1vFktJTH0xG4,5722
 xm_slurm/status.py,sha256=WTWiDHi-ZHtwHRnDP0cGa-27zTSm6LkA-GCKsN-zBgg,6916
 xm_slurm/types.py,sha256=TsVykDm-LazVkrjeJrTwCMs4Q8APKhy7BTk0yKIhFNg,805
 xm_slurm/utils.py,sha256=ESjOkGT7bRSzIeZrUtZplSHP4oaH6VZ92y2woYdcyKM,2239
 xm_slurm/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-xm_slurm/contrib/clusters/__init__.py,sha256=vugR50D9fPJQN5bTd7cSArDGrA6pC-YJHMXrEyvr_Uw,2980
+xm_slurm/contrib/clusters/__init__.py,sha256=JI_zTMfozanpfbBcNjPpgGi5Ppc6jyjM05X3rvFODLs,2321
 xm_slurm/contrib/clusters/drac.py,sha256=tJeQFWFIpeZ1gD3j6AAJssNoLSiDkB-3lz1_ObnkRhc,5905
 xm_slurm/experimental/parameter_controller.py,sha256=b5LfglHV307F6QcPrHeZX5GJBtyOK9aQydke_SZ3Wto,8457
 xm_slurm/packaging/__init__.py,sha256=dh307yLpUT9KN7rJ1e9fYC6hegGKfZcGboUq9nGpDVQ,233
@@ -25,7 +25,7 @@ xm_slurm/packaging/registry.py,sha256=GrdmQg9MgSo38OiqOzMKWSkQyBuyryOfc3zcdgZ4CU
 xm_slurm/packaging/router.py,sha256=yPbdA9clrhly97cLgDsSRZG2LZRKE-oz8Hhdb7WtYqk,2070
 xm_slurm/packaging/utils.py,sha256=KI5s32rNTCfgwzY_7Ghck27jHKvKg5sl5_NEEqJbJqI,3999
 xm_slurm/scripts/_cloudpickle.py,sha256=dlJYf2SceOuUn8wi-ozuoYAQg71wqD2MUVOUCyOwWIY,647
-xm_slurm/scripts/cli.py,sha256=ZXqYOs8X23TYDdKxvV-wIa-0mTfpxSl4_Pli6TiKI7s,1435
+xm_slurm/scripts/cli.py,sha256=xA4SqcMtX_NXXdUDgJ47qNHw2uGvmn_JA3XiDXk-jFA,2152
 xm_slurm/templates/docker/docker-bake.hcl.j2,sha256=ClsFpj91Mr1VfA8L6eqBG3HQz0Z8VenF6mEfmAhQgUo,1498
 xm_slurm/templates/docker/mamba.Dockerfile,sha256=Sgxr5IA5T-pT1Shumb5k3JngoG4pgCdBXjzqslFJdZI,753
 xm_slurm/templates/docker/python.Dockerfile,sha256=U4b4QVkopckQ0o9jJIE7d_M6TvExEYlYDirNwCoZ7W4,865
@@ -35,10 +35,10 @@ xm_slurm/templates/slurm/job-group.bash.j2,sha256=UkjfBE7jg9mepcUWaHZEAjkiXsIM1j
 xm_slurm/templates/slurm/job.bash.j2,sha256=v0xGYzagDdWW6Tg44qobGJLNSUP1Cf4CcekrPibYdrE,1864
 xm_slurm/templates/slurm/fragments/monitor.bash.j2,sha256=HYqYhXsTv8TCed5UaGCZVGIYsqxSKHcnPyNNTHWNvxc,1279
 xm_slurm/templates/slurm/fragments/proxy.bash.j2,sha256=VJLglZo-Nvx9R-qe3rHTxr07CylTQ6Z9NwBzvIpAZrA,814
-xm_slurm/templates/slurm/runtimes/apptainer.bash.j2,sha256=dMntzelhs8DqKyIpO9S6wzMfH2PDevmgvyjCW8Xc2dY,3222
+xm_slurm/templates/slurm/runtimes/apptainer.bash.j2,sha256=ggSsAxv-2_Ct3hSxFJgwgwa3Wu8xH3JqLxWtJOYYrsA,3253
 xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=xKXYFvQvazMx0PgvmlRXR6eecoiBUl8y52dIzQtWkBE,1469
-xmanager_slurm-0.4.1.dist-info/METADATA,sha256=3mT4XIm8evv-5qw7oney4nYn3IasIA_l1rWz86XNOY8,954
-xmanager_slurm-0.4.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
-xmanager_slurm-0.4.1.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
-xmanager_slurm-0.4.1.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
-xmanager_slurm-0.4.1.dist-info/RECORD,,
+xmanager_slurm-0.4.2.dist-info/METADATA,sha256=FqBIDTKXc5mrv_R5WDxlOvs7l87Cxrv2bAmnQuJlq2k,954
+xmanager_slurm-0.4.2.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+xmanager_slurm-0.4.2.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.2.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.2.dist-info/RECORD,,