xmanager-slurm 0.4.6__tar.gz → 0.4.8__tar.gz

{xmanager_slurm-0.4.6 → xmanager_slurm-0.4.8}/PKG-INFO

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

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

{xmanager_slurm-0.4.6 → xmanager_slurm-0.4.8}/uv.lock

@@ -2250,7 +2250,7 @@ wheels = [
 
 [[package]]
 name = "xmanager-slurm"
-version = "0.4.6"
+version = "0.4.8"
 source = { editable = "." }
 dependencies = [
     { name = "aiofile" },

{xmanager_slurm-0.4.6 → xmanager_slurm-0.4.8}/xm_slurm/api/web/client.py

@@ -33,7 +33,7 @@ class XManagerWebAPI(XManagerAPI):
         self.client = httpx.Client(headers={"Authorization": f"Bearer {token}"}, verify=False)
 
     def _make_url(self, path: str) -> str:
-        return f"{self.base_url}/api{path}"
+        return f"{self.base_url}{path}"
 
     @with_backoff
     def get_experiment(self, xid: int) -> models.Experiment:

{xmanager_slurm-0.4.6 → xmanager_slurm-0.4.8}/xm_slurm/config.py

@@ -119,8 +119,9 @@ class SlurmSSHConfig:
             config=None,
             kbdint_auth=False,
             disable_trivial_auth=True,
+            known_hosts=self.known_hosts,
         )
-        options.prepare(last_config=self.config, known_hosts=self.known_hosts)
+        options.prepare(last_config=self.config)
         return options
 
     def serialize(self):

{xmanager_slurm-0.4.6 → xmanager_slurm-0.4.8}/xm_slurm/execution.py

@@ -235,11 +235,10 @@ class SlurmHandle(_BatchedSlurmHandle, tp.Generic[SlurmJobT]):
         return await self._batched_get_state(self.ssh, self.slurm_job)
 
     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"
-
+        self, *, num_lines: int, block_size: int, wait: bool, follow: bool, raw: bool = False
+    ) -> tp.AsyncGenerator[tp.Union[str, ConsoleRenderable], None]:
+        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)
 
         if wait:
@@ -249,7 +248,7 @@ class SlurmHandle(_BatchedSlurmHandle, tp.Generic[SlurmJobT]):
         file_size = await fs.size(file)
         assert file_size is not None
 
-        async with await fs.open(file, "rb") as remote_file:  # type: ignore
+        async with await fs.open(file, "rb") as remote_file:
             data = b""
             lines = []
             position = file_size
@@ -257,27 +256,40 @@ class SlurmHandle(_BatchedSlurmHandle, tp.Generic[SlurmJobT]):
             while len(lines) <= num_lines and position > 0:
                 read_size = min(block_size, position)
                 position -= read_size
-                await remote_file.seek(position)  # type: ignore
+                await remote_file.seek(position)
                 chunk = await remote_file.read(read_size)
                 data = chunk + data
                 lines = data.splitlines()
 
             if position <= 0:
-                yield Rule("[bold red]BEGINNING OF FILE[/bold red]")
+                if raw:
+                    yield "\033[31mBEGINNING OF FILE\033[0m\n"
+                else:
+                    yield Rule("[bold red]BEGINNING OF FILE[/bold red]")
             for line in lines[-num_lines:]:
-                yield Text.from_ansi(line.decode("utf-8", errors="replace"))
+                if raw:
+                    yield line.decode("utf-8", errors="replace") + "\n"
+                else:
+                    yield Text.from_ansi(line.decode("utf-8", errors="replace"))
 
             if (await self.get_state()) not in status.SlurmActiveJobStates:
-                yield Rule("[bold red]END OF FILE[/bold red]")
-                return
+                if raw:
+                    yield "\033[31mEND OF FILE\033[0m\n"
+                    return
+                else:
+                    yield Rule("[bold red]END OF FILE[/bold red]")
+                    return
 
             if not follow:
                 return
 
-            await remote_file.seek(file_size)  # type: ignore
+            await remote_file.seek(file_size)
             while True:
                 if new_data := (await remote_file.read(block_size)):
-                    yield Text.from_ansi(new_data.decode("utf-8", errors="replace"))
+                    if raw:
+                        yield new_data.decode("utf-8", errors="replace")
+                    else:
+                        yield Text.from_ansi(new_data.decode("utf-8", errors="replace"))
                 else:
                     await asyncio.sleep(0.25)
 
@@ -454,10 +466,28 @@ class SlurmExecutionClient:
         return False
 
     @functools.cache
+    @utils.reawaitable
     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 +676,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())

{xmanager_slurm-0.4.6 → xmanager_slurm-0.4.8}/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 → xmanager_slurm-0.4.8}/xm_slurm/utils.py

@@ -1,3 +1,4 @@
+import asyncio
 import functools
 import logging
 import os
@@ -13,10 +14,39 @@ import typing as tp
 from xmanager import xm
 
 T = tp.TypeVar("T")
+P = tp.ParamSpec("P")
 
 logger = logging.getLogger(__name__)
 
 
+class CachedAwaitable(tp.Awaitable[T]):
+    def __init__(self, awaitable: tp.Awaitable[T]):
+        self.awaitable = awaitable
+        self.result: asyncio.Future[T] | None = None
+
+    def __await__(self):
+        if not self.result:
+            future = asyncio.get_event_loop().create_future()
+            self.result = future
+            try:
+                result = yield from self.awaitable.__await__()
+                future.set_result(result)
+            except Exception as e:
+                future.set_exception(e)
+
+        if not self.result.done():
+            yield from self.result
+        return self.result.result()
+
+
+def reawaitable(f: tp.Callable[P, tp.Awaitable[T]]) -> tp.Callable[P, CachedAwaitable[T]]:
+    @functools.wraps(f)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> CachedAwaitable[T]:
+        return CachedAwaitable(f(*args, **kwargs))
+
+    return wrapper
+
+
 @functools.cache
 def find_project_root() -> pathlib.Path:
     launch_script_path: pathlib.Path | None = None