xmanager-slurm 0.4.6__py3-none-any.whl → 0.4.7__py3-none-any.whl

xm_slurm/execution.py

@@ -237,8 +237,8 @@ class SlurmHandle(_BatchedSlurmHandle, tp.Generic[SlurmJobT]):
     async def logs(
         self, *, num_lines: int, block_size: int, wait: bool, follow: bool
     ) -> tp.AsyncGenerator[ConsoleRenderable, None]:
-        statedir = await get_client()._state_dir(self.ssh)
-        file = statedir / f"{self.experiment_id}/slurm-{self.slurm_job.job_id}.out"
+        experiment_dir = await get_client().experiment_dir(self.ssh, self.experiment_id)
+        file = experiment_dir / f"slurm-{self.slurm_job.job_id}.out"
 
         fs = await get_client().fs(self.ssh)
 
@@ -455,9 +455,26 @@ class SlurmExecutionClient:
 
     @functools.cache
     async def _state_dir(self, ssh_config: SlurmSSHConfig) -> pathlib.Path:
-        cmd = await self.run(ssh_config, "printenv HOME", check=True)
-        assert isinstance(cmd.stdout, str)
-        return pathlib.Path(cmd.stdout.strip()) / ".local" / "state" / "xm-slurm"
+        state_dirs = [
+            ("XM_SLURM_STATE_DIR", ""),
+            ("XDG_STATE_HOME", "xm-slurm"),
+            ("HOME", ".local/state/xm-slurm"),
+        ]
+
+        for env_var, subpath in state_dirs:
+            cmd = await self.run(ssh_config, f"printenv {env_var}", check=False)
+            assert isinstance(cmd.stdout, str)
+            if cmd.returncode == 0:
+                return pathlib.Path(cmd.stdout.strip()) / subpath
+
+        raise SlurmExecutionError(
+            "Failed to find a valid state directory for XManager. "
+            "We weren't able to resolve any of the following paths: "
+            f"{', '.join(env_var + ('/' + subpath if subpath else '') for env_var, subpath in state_dirs)}."
+        )
+
+    async def experiment_dir(self, ssh_config: SlurmSSHConfig, experiment_id: int) -> pathlib.Path:
+        return (await self._state_dir(ssh_config)) / f"{experiment_id:08d}"
 
     async def run(
         self,
@@ -646,7 +663,7 @@ class SlurmExecutionClient:
 
         fs = await self.fs(cluster.ssh)
 
-        template_dir = (await self._state_dir(cluster.ssh)) / f"{experiment_id}"
+        template_dir = await self.experiment_dir(cluster.ssh, experiment_id)
 
         await fs.makedirs(template_dir, exist_ok=True)
         await fs.write(template_dir / submission_script_path, submission_script.encode())

xm_slurm/templates/slurm/runtimes/podman.bash.j2

@@ -3,7 +3,6 @@ time podman pull \
 {% if job.executable.credentials %}
   --creds {{ job.executable.credentials.username }}:{{ job.executable.credentials.password }} \
 {% endif %}
-  --retry 3 \
   {{ job.executable.image }}
 
 cat << 'ENTRYPOINT_EOF' > "$SLURM_TMPDIR"/xm-slurm-entrypoint.sh

{xmanager_slurm-0.4.6.dist-info → xmanager_slurm-0.4.7.dist-info}/METADATA

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

@@ -5,7 +5,7 @@ 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
 xm_slurm/executables.py,sha256=fGmrFBl-258bMn6ip5adYeM7xxUHAeIbDN9zD2FDGtY,6373
-xm_slurm/execution.py,sha256=MRv7MjFVDhGbUrq28kpW82kOYvKdb0tS7tV5jh8g_I0,30163
+xm_slurm/execution.py,sha256=AicnjpY9ncLBlJrLzpngjmS23zfFZftYrflrWpgcyzE,30888
 xm_slurm/executors.py,sha256=fMtxGUCi4vEKmb_p4JEpqPUTh7L_f1LcR_TamMLAWNg,4667
 xm_slurm/experiment.py,sha256=94r0mhtUPUzw4eaUEz0kpsufC25wEGqlDhV4Fcr1ukY,39883
 xm_slurm/filesystem.py,sha256=4rKtq3t-KDgxJbSGt6JVyRJT_3lCN_vIKTcwKHpTo3I,4389
@@ -43,9 +43,9 @@ xm_slurm/templates/slurm/job.bash.j2,sha256=d35VYHdAKkgVK8s4XnUDJwQR0gLnDWRJu-Ld
 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=lE2EWVCK2O-n08RL4_MJYIikVTvODjcYKuv7Eh73Q2w,1932
-xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=yUOUTzGNyVy1IQSgvY9wn9pE-reB1SYNJmuV2ji8l6w,1162
-xmanager_slurm-0.4.6.dist-info/METADATA,sha256=Zo382SdJ2QWrzEVwNYQX7DXRhhcv9HZA8SPRuWtT0W4,1042
-xmanager_slurm-0.4.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-xmanager_slurm-0.4.6.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
-xmanager_slurm-0.4.6.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
-xmanager_slurm-0.4.6.dist-info/RECORD,,
+xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=3j7K5eyXt_WhXK0EoMlxnhlmFVJ2JyxRKbsMRaDqzSs,1148
+xmanager_slurm-0.4.7.dist-info/METADATA,sha256=mRwFTuy53rP0lDYJ9xMXx6E8Ed7sZAzjSEcQpjvV3GQ,1042
+xmanager_slurm-0.4.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xmanager_slurm-0.4.7.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.7.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.7.dist-info/RECORD,,