xmanager-slurm 0.4.7__tar.gz → 0.4.9__tar.gz

{xmanager_slurm-0.4.7 → xmanager_slurm-0.4.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.7
+Version: 0.4.9
 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.9/examples/job-timeout/launch.py

@@ -0,0 +1,53 @@
+import datetime as dt
+
+from absl import app
+from xmanager import xm
+
+import xm_slurm
+import xm_slurm.contrib.clusters
+
+
+@xm.run_in_asyncio_loop
+async def main(_):
+    async with xm_slurm.create_experiment("My Experiment") as experiment:
+        # Step 1: Specify the executor specification
+        executor_spec = xm_slurm.Slurm.Spec(tag="ghcr.io/jessefarebro/xm-slurm/test:latest")
+
+        # Step 2: Specify the executable and package it
+        [executable] = experiment.package(
+            [
+                xm_slurm.uv_container(
+                    executor_spec=executor_spec,
+                    # Equivalent of `-m rich.status`
+                    entrypoint=xm.CommandList(["main.py", "900"]),
+                ),
+            ],
+        )
+
+        # Step 3: Construct requirements & executor
+        requirements = xm_slurm.JobRequirements(
+            CPU=1,
+            RAM=1.0 * xm.GiB,
+            GPU=1,
+            replicas=1,
+            cluster=xm_slurm.contrib.clusters.mila(),
+        )
+        executor = xm_slurm.Slurm(
+            requirements=requirements,
+            time=dt.timedelta(minutes=10),
+        )
+
+        # Step 4: Schedule job
+        wu = await experiment.add(
+            xm.Job(
+                executable=executable,
+                executor=executor,
+            )
+        )
+
+        await wu.wait_until_complete()
+        print(f"Job finished executing with status {await wu.get_status()}")
+
+
+if __name__ == "__main__":
+    app.run(main)

xmanager_slurm-0.4.9/examples/job-timeout/main.py

@@ -0,0 +1,11 @@
+import sys
+import time
+
+
+def main():
+    print("Starting job...")
+    time.sleep(int(sys.argv[1]))
+
+
+if __name__ == "__main__":
+    main()

xmanager_slurm-0.4.9/examples/uv/pyproject.toml

@@ -0,0 +1,6 @@
+[project]
+name = "xm-slurm-example"
+description = "XManager Slurm test project"
+version = "0.0.1"
+requires-python = ">=3.10"
+dependencies = ["rich"]

xmanager_slurm-0.4.9/examples/uv/uv.lock

@@ -0,0 +1,66 @@
+version = 1
+requires-python = ">=3.10"
+
+[[package]]
+name = "markdown-it-py"
+version = "3.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mdurl" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/38/71/3b932df36c1a044d397a1f92d1cf91ee0a503d91e470cbd670aa66b07ed0/markdown-it-py-3.0.0.tar.gz", hash = "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb", size = 74596 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/42/d7/1ec15b46af6af88f19b8e5ffea08fa375d433c998b8a7639e76935c14f1f/markdown_it_py-3.0.0-py3-none-any.whl", hash = "sha256:355216845c60bd96232cd8d8c40e8f9765cc86f46880e43a8fd22dc1a1a8cab1", size = 87528 },
+]
+
+[[package]]
+name = "mdurl"
+version = "0.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 },
+]
+
+[[package]]
+name = "pygments"
+version = "2.18.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/8e/62/8336eff65bcbc8e4cb5d05b55faf041285951b6e80f33e2bff2024788f31/pygments-2.18.0.tar.gz", hash = "sha256:786ff802f32e91311bff3889f6e9a86e81505fe99f2735bb6d60ae0c5004f199", size = 4891905 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f7/3f/01c8b82017c199075f8f788d0d906b9ffbbc5a47dc9918a945e13d5a2bda/pygments-2.18.0-py3-none-any.whl", hash = "sha256:b8e6aca0523f3ab76fee51799c488e38782ac06eafcf95e7ba832985c8e7b13a", size = 1205513 },
+]
+
+[[package]]
+name = "rich"
+version = "13.9.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "markdown-it-py" },
+    { name = "pygments" },
+    { name = "typing-extensions", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/aa/9e/1784d15b057b0075e5136445aaea92d23955aad2c93eaede673718a40d95/rich-13.9.2.tar.gz", hash = "sha256:51a2c62057461aaf7152b4d611168f93a9fc73068f8ded2790f29fe2b5366d0c", size = 222843 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/67/91/5474b84e505a6ccc295b2d322d90ff6aa0746745717839ee0c5fb4fdcceb/rich-13.9.2-py3-none-any.whl", hash = "sha256:8c82a3d3f8dcfe9e734771313e606b39d8247bb6b826e196f4914b333b743cf1", size = 242117 },
+]
+
+[[package]]
+name = "typing-extensions"
+version = "4.12.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/df/db/f35a00659bc03fec321ba8bce9420de607a1d37f8342eee1863174c69557/typing_extensions-4.12.2.tar.gz", hash = "sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8", size = 85321 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/26/9f/ad63fc0248c5379346306f8668cda6e2e2e9c95e01216d2b8ffd9ff037d0/typing_extensions-4.12.2-py3-none-any.whl", hash = "sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d", size = 37438 },
+]
+
+[[package]]
+name = "xm-slurm-example"
+version = "0.0.1"
+source = { virtual = "." }
+dependencies = [
+    { name = "rich" },
+]
+
+[package.metadata]
+requires-dist = [{ name = "rich" }]

{xmanager_slurm-0.4.7 → xmanager_slurm-0.4.9}/pyproject.toml

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

{xmanager_slurm-0.4.7 → xmanager_slurm-0.4.9}/uv.lock

@@ -1,4 +1,5 @@
 version = 1
+revision = 1
 requires-python = ">=3.10"
 resolution-markers = [
     "python_full_version < '3.11'",
@@ -2250,7 +2251,7 @@ wheels = [
 
 [[package]]
 name = "xmanager-slurm"
-version = "0.4.7"
+version = "0.4.9"
 source = { editable = "." }
 dependencies = [
     { name = "aiofile" },

{xmanager_slurm-0.4.7 → xmanager_slurm-0.4.9}/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.7 → xmanager_slurm-0.4.9}/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.7 → xmanager_slurm-0.4.9}/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]:
+        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,6 +466,7 @@ class SlurmExecutionClient:
         return False
 
     @functools.cache
+    @utils.reawaitable
     async def _state_dir(self, ssh_config: SlurmSSHConfig) -> pathlib.Path:
         state_dirs = [
             ("XM_SLURM_STATE_DIR", ""),

{xmanager_slurm-0.4.7 → xmanager_slurm-0.4.9}/xm_slurm/executors.py

@@ -57,8 +57,13 @@ class Slurm(xm.Executor):
 
     requeue: bool = True  # Is this job ellible for requeueing?
     requeue_on_exit_code: int = 42  # The exit code that triggers requeueing
+    requeue_on_timeout: bool = True  # Should the job requeue upon timeout minus the grace period
     requeue_max_attempts: int = 5  # How many times to attempt requeueing
 
+    @property
+    def requeue_timeout(self) -> dt.timedelta:
+        return self.time - self.timeout_signal_grace_period
+
     def __post_init__(self) -> None:
         if not isinstance(self.time, dt.timedelta):
             raise TypeError(f"time must be a `datetime.timedelta`, got {type(self.time)}")

xmanager_slurm-0.4.9/xm_slurm/templates/slurm/fragments/monitor.bash.j2

@@ -0,0 +1,76 @@
+{% macro monitor(requeue_max_attempts, requeue_exit_code, requeue_on_timeout, requeue_timeout) -%}
+__xm_slurm_wait_for_children() {
+  if [[ -n "${SLURM_ARRAY_JOB_ID:-}" ]]; then
+    local -r JOB_ID="${SLURM_ARRAY_JOB_ID}_${SLURM_ARRAY_TASK_ID}"
+  else
+    local -r JOB_ID="${SLURM_JOB_ID}"
+  fi
+
+  # If there are no child jobs we should error out
+  children=( $(jobs -p) )
+{% raw %}
+  if [ ${#children[@]} -eq 0 ]; then
+{% endraw %}
+    echo "ERROR: no child jobs exist..." >&2
+    exit 1
+  fi
+
+{% if requeue_on_timeout %}
+  # Start a watchdog process to signal timeout.
+  sleep {{ requeue_timeout }} &
+  timeout_pid=$!
+{% endif %}
+
+{% raw %}
+  while [ ${#children[@]} -gt 0 ]; do
+{% endraw %}
+    echo "INFO: Waiting for child processes to finish..."
+{% if requeue_on_timeout %}
+    # Wait on either one of the child processes or the timeout process.
+    wait -n -p child_pid "${children[@]}" "${timeout_pid}"
+{% else %}
+    wait -n -p child_pid "${children[@]}"
+{% endif %}
+    local child_exit_code=$?
+
+{% if requeue_on_timeout %}
+    # If the finished process is the watchdog, trigger the timeout handling.
+    if [ "${child_pid}" = "${timeout_pid}" ]; then
+      echo "INFO: Timeout of {{ requeue_timeout }} seconds reached. Killing remaining processes: ${children[*]}" >&2
+      kill "${children[@]}" 2>/dev/null || true
+      scontrol requeue "${JOB_ID}"
+      exit {{ requeue_exit_code }}
+    fi
+{% endif %}
+
+    echo "INFO: Process ${child_pid} finished with exit code ${child_exit_code}."
+
+    # Handle the exit code of the finished process.
+    if [ "${child_exit_code}" -eq "{{ requeue_exit_code }}" ] && [ "${SLURM_RESTART_COUNT:-0}" -le "{{ requeue_max_attempts }}" ]; then
+      echo "INFO: Received requeue exit code {{ requeue_exit_code }} from process ${child_pid}. Requeuing Slurm job ${JOB_ID} after ${SLURM_RESTART_COUNT-0} restarts." >&2
+      scontrol requeue "${JOB_ID}"
+      exit {{ requeue_exit_code }}
+    elif [ "${child_exit_code}" -ne 0 ]; then
+      echo "ERROR: Process ${child_pid} exited with code ${child_exit_code}." >&2
+      exit "${child_exit_code}"
+    fi
+
+    # Remove the finished PID from the array in a concise way.
+    for i in "${!children[@]}"; do
+      if [ "${children[i]}" = "$child_pid" ]; then
+        unset 'children[i]'
+        break
+      fi
+    done
+
+    # Reindex the array.
+    children=( "${children[@]}" )
+  done
+
+{% if requeue_on_timeout %}
+  kill "$timeout_pid" 2>/dev/null || true
+{% endif %}
+}
+
+__xm_slurm_wait_for_children
+{%- endmacro %}

{xmanager_slurm-0.4.7 → xmanager_slurm-0.4.9}/xm_slurm/templates/slurm/job.bash.j2

@@ -73,7 +73,7 @@ echo "[INFO] Start timestamp: $(date)"
 
 {% block monitor -%}
 {% from 'fragments/monitor.bash.j2' import monitor %}
-{{ monitor(job.executor.requeue_max_attempts, job.executor.requeue_on_exit_code) }}
+{{ monitor(job.executor.requeue_max_attempts, job.executor.requeue_on_exit_code, job.executor.requeue_on_timeout, job.executor.requeue_timeout.seconds) }}
 {%- endblock monitor %}
 
 

{xmanager_slurm-0.4.7 → xmanager_slurm-0.4.9}/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

xmanager_slurm-0.4.7/xm_slurm/templates/slurm/fragments/monitor.bash.j2

@@ -1,37 +0,0 @@
-{% macro monitor(requeue_max_attempts, requeue_exit_code) -%}
-__xm_slurm_wait_for_children() {
-  if [[ -n "${SLURM_ARRAY_JOB_ID:-}" ]]; then
-    local -r JOB_ID="${SLURM_ARRAY_JOB_ID}_${SLURM_ARRAY_TASK_ID}"
-  else
-    local -r JOB_ID="${SLURM_JOB_ID}"
-  fi
-
-  # If there are no child jobs we should error out
-  if [ -z "$(jobs -p)" ]; then
-    echo "ERROR: no child jobs exist..." >&2
-    exit -1
-  fi
-
-  # Loop through all job IDs in the background job list and wait for them to finish
-  for job in "$(jobs -p)"; do
-    echo "INFO: Waiting for job ${job} to finish..."
-    set +e
-    wait "${job}"
-    local -r JOB_EXIT_CODE="${?}"
-    set -e
-
-    if [ "${JOB_EXIT_CODE}" -eq "{{ requeue_exit_code }}" ] && [ "${SLURM_RESTART_COUNT-0}" -le "{{ requeue_max_attempts }}" ]; then
-      echo "INFO: Received requeue exit code {{ requeue_exit_code }} from job ${job}. Requeing Slurm job ${JOB_ID} after ${SLURM_RESTART_COUNT-0} restarts." >&2
-      scontrol requeue "${JOB_ID}"
-      exit {{ requeue_exit_code }}
-    elif [ "${JOB_EXIT_CODE}" -ne 0 ]; then
-      echo "ERROR: Job ${job} exited with code ${JOB_EXIT_CODE}." >&2
-      exit "${JOB_EXIT_CODE}"
-    else
-      echo "INFO: Job ${job} exited successfully." >&2
-    fi
-  done
-}
-
-__xm_slurm_wait_for_children
-{%- endmacro %}