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

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:

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

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

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)}")

xm_slurm/templates/slurm/fragments/monitor.bash.j2

@@ -1,4 +1,4 @@
-{% macro monitor(requeue_max_attempts, requeue_exit_code) -%}
+{% 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}"
@@ -7,30 +7,69 @@ __xm_slurm_wait_for_children() {
   fi
 
   # If there are no child jobs we should error out
-  if [ -z "$(jobs -p)" ]; then
+  children=( $(jobs -p) )
+{% raw %}
+  if [ ${#children[@]} -eq 0 ]; then
+{% endraw %}
     echo "ERROR: no child jobs exist..." >&2
-    exit -1
+    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 requeue_on_timeout %}
+  # Start a watchdog process to signal timeout.
+  sleep {{ requeue_timeout }} &
+  timeout_pid=$!
+{% endif %}
 
-    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
+{% 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 }}
-    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
+{% 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

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

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.dist-info → xmanager_slurm-0.4.9.dist-info}/METADATA

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

@@ -1,12 +1,12 @@
 xm_slurm/__init__.py,sha256=WgRn9HDYa5H3sfIH-HZu33liBOh98jM4GqcR349RaSY,1086
 xm_slurm/batching.py,sha256=GbKBsNz9w8gIc2fHLZpslC0e4K9YUfLXFHmjduRRCfQ,4385
-xm_slurm/config.py,sha256=NxlObwlhpjCqafs2CZuw1P5WFwhqayfblGHo6NGFi3Y,7072
+xm_slurm/config.py,sha256=PvdLG6fSNfbABbtwELbnytx85vWLPtyVfEJtoQsLf94,7084
 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=AicnjpY9ncLBlJrLzpngjmS23zfFZftYrflrWpgcyzE,30888
-xm_slurm/executors.py,sha256=fMtxGUCi4vEKmb_p4JEpqPUTh7L_f1LcR_TamMLAWNg,4667
+xm_slurm/execution.py,sha256=c0aV1h2tKQFyAGM6JLd16MWFgpRLKAbcutZz17xPUSw,31400
+xm_slurm/executors.py,sha256=bUgKcgtvf-nPGjcuHRzUAqD1r3_vwea_h-Y9MAB-Kqo,4887
 xm_slurm/experiment.py,sha256=94r0mhtUPUzw4eaUEz0kpsufC25wEGqlDhV4Fcr1ukY,39883
 xm_slurm/filesystem.py,sha256=4rKtq3t-KDgxJbSGt6JVyRJT_3lCN_vIKTcwKHpTo3I,4389
 xm_slurm/job_blocks.py,sha256=_F8CKCs5BQFj40a2-mjG71HfacvWoBXBDPDKEaKTbXc,616
@@ -15,12 +15,12 @@ xm_slurm/packageables.py,sha256=fPUvqF2IvJ2Hn6hodDdQwtx1Ze3sJ8U-BUbxDHauW-g,1239
 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=ftV7kyyr4EUcIjon0godTDhhSzdVjyDkF2sVHA8zyls,5024
+xm_slurm/utils.py,sha256=xtFvktaxr0z65sTdu6HhOVfyo0OAB9t-EYXWcYrQQEU,5958
 xm_slurm/api/__init__.py,sha256=cyao3LZ3uLftu1wIv1aN7Qvsl6gYzYpkxeehTHZ0fA8,1089
 xm_slurm/api/abc.py,sha256=-lS2OndnOuEiwNdr8ccQKkwMd1iDmKMmkBOSTvo5H5w,1816
 xm_slurm/api/models.py,sha256=_INVh0j-4-rRs0WASyg4fNB6NF1L1nUeGgQ6-XnbwsM,1610
 xm_slurm/api/sqlite/client.py,sha256=WykSIO7b14rRLy9qebbkiLKXy7EHU61jtoebLX17HMM,14124
-xm_slurm/api/web/client.py,sha256=GkQIaOAJc1MIz0mYVWUA97ZvhAgz_o4xXHvVdfs2ErA,6265
+xm_slurm/api/web/client.py,sha256=uO67Y7fnQ-w__Vm_A5BEuy7Qi8wQcWk3vIsBGEBkyfk,6261
 xm_slurm/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 xm_slurm/contrib/clusters/__init__.py,sha256=XFCVnkThiU3_8uA_tUgDByOBanXNHrxDvfmuptmQ2KE,2214
 xm_slurm/contrib/clusters/drac.py,sha256=ViLYerYBMSuZXnWVbz9RDIPPV7JA8BgBpgTfj1wPP28,5881
@@ -39,13 +39,13 @@ xm_slurm/templates/docker/uv.Dockerfile,sha256=L2UJMX2c8waMdrRhiqPytQe3pTBu6u5Pp
 xm_slurm/templates/slurm/entrypoint.bash.j2,sha256=MRdSVwgGrgQdpEhqfkP35IidgsblrtVXB1YWzvE9hkk,666
 xm_slurm/templates/slurm/job-array.bash.j2,sha256=smxmSSzBEUHm6MJF-nYPVVjK6CLKrb1fRxF_tfrzAX8,552
 xm_slurm/templates/slurm/job-group.bash.j2,sha256=Cp8YhNOxYqaOkl4MFjQlcaLMGZwdDh97m8OGT5RWbAo,1101
-xm_slurm/templates/slurm/job.bash.j2,sha256=d35VYHdAKkgVK8s4XnUDJwQR0gLnDWRJu-Ldz-qALmQ,1914
-xm_slurm/templates/slurm/fragments/monitor.bash.j2,sha256=HYqYhXsTv8TCed5UaGCZVGIYsqxSKHcnPyNNTHWNvxc,1279
+xm_slurm/templates/slurm/job.bash.j2,sha256=pNKir1tkmRTGDiGxlQ3DkUaW9Zos_gdkkXJC_xX5Cxo,1985
+xm_slurm/templates/slurm/fragments/monitor.bash.j2,sha256=BJ1brSjhESOe9VX_OYaPyy9-qE3uiFlzxp8ZkFcTw8Y,2504
 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=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,,
+xmanager_slurm-0.4.9.dist-info/METADATA,sha256=WWPRzVrTsK5t8kD732EIejSlNgQ8KP01Ln7eM1Mj4e4,1042
+xmanager_slurm-0.4.9.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+xmanager_slurm-0.4.9.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.9.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.9.dist-info/RECORD,,