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

xm_slurm/utils.py

@@ -1,51 +1,20 @@
 import functools
+import logging
 import os
 import pathlib
+import pty
+import re
+import select
+import shutil
+import subprocess
 import sys
-from typing import Callable, Hashable, Iterable, MutableSet, TypeVar
+import typing as tp
 
-T = TypeVar("T")
+from xmanager import xm
 
+T = tp.TypeVar("T")
 
-class UserSet(Hashable, MutableSet[T]):
-    __hash__ = MutableSet._hash
-
-    def __init__(
-        self,
-        iterable: Iterable[T] = (),
-        /,
-        on_add: Callable[[T], None] = lambda x: None,
-        on_remove: Callable[[T], None] = lambda x: None,
-        on_discard: Callable[[T], None] = lambda x: None,
-    ):
-        self.data = set(iterable)
-        self._on_add = on_add
-        self._on_remove = on_remove
-        self._on_discard = on_discard
-
-    def __contains__(self, value):
-        return value in self.data
-
-    def __iter__(self):
-        return iter(self.data)
-
-    def __len__(self):
-        return len(self.data)
-
-    def __repr__(self):
-        return repr(self.data)
-
-    def add(self, value: T):
-        self.data.add(value)
-        self._on_add(value)
-
-    def remove(self, value: T):
-        self.data.remove(value)
-        self._on_remove(value)
-
-    def discard(self, value: T):
-        self.data.discard(value)
-        self._on_discard(value)
+logger = logging.getLogger(__name__)
 
 
 @functools.cache
@@ -75,3 +44,115 @@ def find_project_root() -> pathlib.Path:
         pdir = pdir.parent
 
     raise RuntimeError(f"Could not find project root from {sys.argv[0]}. Please specify `context`.")
+
+
+# Cursor commands to filter out from the command data stream
+_CURSOR_ESCAPE_SEQUENCES_REGEX = re.compile(
+    rb"\x1b\[\?25[hl]"  # Matches cursor show/hide commands (CSI ?25h and CSI ?25l)
+    rb"|\x1b\[[0-9;]*[Hf]"  # Matches cursor position commands (CSI n;mH and CSI n;mf)
+    rb"|\x1b\[s"  # Matches cursor save position (CSI s)
+    rb"|\x1b\[u"  # Matches cursor restore position (CSI u)
+    rb"|\x1b\[2J"  # Matches clear screen (CSI 2J)
+    rb"|\x1b\[K"  # Matches clear line (CSI K)
+)
+
+
+def run_command(
+    args: tp.Sequence[str] | xm.SequentialArgs,
+    env: dict[str, str] | None = None,
+    tty: bool = False,
+    cwd: str | os.PathLike[str] | None = None,
+    stdin: tp.IO[tp.AnyStr] | str | None = None,
+    check: bool = False,
+    return_stdout: bool = False,
+    return_stderr: bool = False,
+) -> subprocess.CompletedProcess[str]:
+    if isinstance(args, xm.SequentialArgs):
+        args = args.to_list()
+    args = list(args)
+
+    executable = shutil.which(args[0])
+    if not executable:
+        raise RuntimeError(f"Couldn't find executable {args[0]}")
+    executable = pathlib.Path(executable)
+
+    subprocess_env = os.environ.copy() | (env if env else {})
+    if executable.name == "docker" and args[1] == "buildx":
+        subprocess_env |= {"DOCKER_CLI_EXPERIMENTAL": "enabled"}
+
+    logger.debug(f"command: {' '.join(args)}")
+
+    stdout_master, stdout_slave = pty.openpty()
+    stderr_master, stderr_slave = pty.openpty()
+
+    stdout_data, stderr_data = b"", b""
+    with subprocess.Popen(
+        executable=executable,
+        args=args,
+        shell=False,
+        text=True,
+        bufsize=0,
+        stdin=subprocess.PIPE if stdin else None,
+        stdout=stdout_slave,
+        stderr=stderr_slave,
+        start_new_session=True,
+        close_fds=True,
+        cwd=cwd,
+        env=subprocess_env,
+    ) as process:
+        os.close(stdout_slave)
+        os.close(stderr_slave)
+
+        if stdin and process.stdin:
+            process.stdin.write(stdin if isinstance(stdin, str) else tp.cast(str, stdin.read()))
+            process.stdin.close()
+
+        fds = [stdout_master, stderr_master]
+        while fds:
+            rlist, _, _ = select.select(fds, [], [])
+            for fd in rlist:
+                try:
+                    data = os.read(fd, 1024)
+                except OSError:
+                    data = None
+
+                if not data:
+                    os.close(fd)
+                    fds.remove(fd)
+                    continue
+
+                data = _CURSOR_ESCAPE_SEQUENCES_REGEX.sub(b"", data)
+
+                if fd == stdout_master:
+                    if return_stdout:
+                        stdout_data += data
+                    if tty:
+                        os.write(pty.STDOUT_FILENO, data)
+                elif fd == stderr_master:
+                    if return_stderr:
+                        stderr_data += data
+                    if tty:
+                        os.write(pty.STDERR_FILENO, data)
+                else:
+                    raise RuntimeError("Unexpected file descriptor")
+
+    stdout = stdout_data.decode(errors="replace") if stdout_data else ""
+    stderr = stderr_data.decode(errors="replace") if stderr_data else ""
+
+    logger.debug(f"return code: {process.returncode}")
+    if stdout:
+        logger.debug(f"stdout: {stdout}")
+    if stderr:
+        logger.debug(f"stderr: {stderr}")
+
+    retcode = process.poll()
+    assert retcode is not None
+
+    if check and retcode:
+        raise subprocess.CalledProcessError(retcode, process.args)
+    return subprocess.CompletedProcess(
+        process.args,
+        retcode,
+        stdout=stdout,
+        stderr=stderr,
+    )

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

@@ -1,10 +1,11 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: xmanager-slurm
-Version: 0.4.5
+Version: 0.4.6
 Summary: Slurm backend for XManager.
 Project-URL: GitHub, https://github.com/jessefarebro/xm-slurm
 Author-email: Jesse Farebrother <jfarebro@cs.mcgill.ca>
 License: MIT
+License-File: LICENSE.md
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
@@ -13,13 +14,16 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Requires-Python: >=3.10
-Requires-Dist: asyncssh>=2.13.2
+Requires-Dist: aiofile>=3.9.0
+Requires-Dist: asyncssh>=2.19.0
 Requires-Dist: backoff>=2.2.1
 Requires-Dist: cloudpickle>=3.0.0
+Requires-Dist: httpx>=0.28.1
 Requires-Dist: humanize>=4.8.0
 Requires-Dist: immutabledict>=3.0.0
 Requires-Dist: jinja2>=3.1.2
 Requires-Dist: more-itertools>=10.2.0
 Requires-Dist: rich>=13.5.2
 Requires-Dist: toml>=0.10.2
+Requires-Dist: wrapt>=1.16.0
 Requires-Dist: xmanager>=0.5.0

xmanager_slurm-0.4.6.dist-info/RECORD

@@ -0,0 +1,51 @@
+xm_slurm/__init__.py,sha256=WgRn9HDYa5H3sfIH-HZu33liBOh98jM4GqcR349RaSY,1086
+xm_slurm/batching.py,sha256=GbKBsNz9w8gIc2fHLZpslC0e4K9YUfLXFHmjduRRCfQ,4385
+xm_slurm/config.py,sha256=NxlObwlhpjCqafs2CZuw1P5WFwhqayfblGHo6NGFi3Y,7072
+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/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
+xm_slurm/job_blocks.py,sha256=_F8CKCs5BQFj40a2-mjG71HfacvWoBXBDPDKEaKTbXc,616
+xm_slurm/metadata_context.py,sha256=mksVRbVUuistL1uE7TC-fkW-Y69On52jN_svP1e1kiQ,7841
+xm_slurm/packageables.py,sha256=fPUvqF2IvJ2Hn6hodDdQwtx1Ze3sJ8U-BUbxDHauW-g,12398
+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/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/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
+xm_slurm/experimental/parameter_controller.py,sha256=b5LfglHV307F6QcPrHeZX5GJBtyOK9aQydke_SZ3Wto,8457
+xm_slurm/packaging/__init__.py,sha256=dh307yLpUT9KN7rJ1e9fYC6hegGKfZcGboUq9nGpDVQ,233
+xm_slurm/packaging/docker.py,sha256=ZNlTWyBQVJQ1kzOhWNqva5g8F31gAbmWL2-DFY9FJ-w,13677
+xm_slurm/packaging/registry.py,sha256=GrdmQg9MgSo38OiqOzMKWSkQyBuyryOfc3zcdgZ4CUE,1148
+xm_slurm/packaging/router.py,sha256=yPbdA9clrhly97cLgDsSRZG2LZRKE-oz8Hhdb7WtYqk,2070
+xm_slurm/packaging/utils.py,sha256=6EAb17zKQQeuyNY2EV9AoW1RvnDGrQwmIT9wtQEsC4c,632
+xm_slurm/scripts/_cloudpickle.py,sha256=dlJYf2SceOuUn8wi-ozuoYAQg71wqD2MUVOUCyOwWIY,647
+xm_slurm/scripts/cli.py,sha256=nnfNF2FAVLrm-4KuUooFfSWxX7hOjmXpI7k5f3Sn5us,2349
+xm_slurm/templates/docker/docker-bake.hcl.j2,sha256=7qSJl2VN5poz-Hh8Gjo7--qR-k3lmfGtBu2mNbfG2uA,1499
+xm_slurm/templates/docker/mamba.Dockerfile,sha256=Sgxr5IA5T-pT1Shumb5k3JngoG4pgCdBXjzqslFJdZI,753
+xm_slurm/templates/docker/python.Dockerfile,sha256=U4b4QVkopckQ0o9jJIE7d_M6TvExEYlYDirNwCoZ7W4,865
+xm_slurm/templates/docker/uv.Dockerfile,sha256=L2UJMX2c8waMdrRhiqPytQe3pTBu6u5PpMhJYsKkbEg,1040
+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/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,,

{xmanager_slurm-0.4.5.dist-info → xmanager_slurm-0.4.6.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.26.3
+Generator: hatchling 1.27.0
 Root-Is-Purelib: true
 Tag: py3-none-any