xmanager-slurm 0.3.2__py3-none-any.whl → 0.4.1__py3-none-any.whl

xm_slurm/templates/docker/uv.Dockerfile

@@ -0,0 +1,35 @@
+# syntax=docker/dockerfile:1.4
+ARG BASE_IMAGE=docker.io/python:3.10-slim-bookworm
+FROM $BASE_IMAGE
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
+
+ARG EXTRA_SYSTEM_PACKAGES=""
+ARG EXTRA_PYTHON_PACKAGES=""
+
+WORKDIR /workspace
+
+ENV UV_PYTHON_DOWNLOADS=0
+ENV UV_COMPILE_BYTECODE=1
+ENV UV_LINK_MODE=copy
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends \
+    git $EXTRA_SYSTEM_PACKAGES \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN uv pip install --system pysocks $EXTRA_PYTHON_PACKAGES
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=bind,source=uv.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    --mount=type=ssh \
+    uv sync --frozen --no-install-project --no-dev --no-editable
+
+COPY --link . /workspace
+RUN --mount=type=cache,target=/root/.cache/uv \
+    --mount=type=ssh \
+    uv sync --frozen --no-dev
+
+ENV PATH="/workspace/.venv/bin:$PATH"
+
+ENTRYPOINT [ "uv", "run", "python" ]

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

@@ -24,6 +24,11 @@ __xm_slurm_wait_for_children() {
       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/templates/slurm/job-array.bash.j2

@@ -2,8 +2,7 @@
 {% block directives %}
 {{ super() -}}
 #SBATCH --array=0-{{ args | length - 1 }}
-#SBATCH --output=xm-%j-%a.stdout
-#SBATCH --error=xm-%j-%a.stderr
+#SBATCH --output=slurm-%A_%a.out
 {% endblock directives %}
 
 {% block bootstrap %}

xm_slurm/templates/slurm/job.bash.j2

@@ -2,8 +2,7 @@
 {% block directives %}
 #SBATCH --open-mode=append
 #SBATCH --export=NONE
-#SBATCH --output=xm-%j.stdout
-#SBATCH --error=xm-%j.stderr
+#SBATCH --output=slurm-%j.out
 #SBATCH --comment="{'xid': {{ experiment_id }}}"
 {% if cluster.account and not job.executor.account %}
 #SBATCH --account={{ cluster.account }}
@@ -16,8 +15,10 @@
 {% endif %}
 {% if identity %}
 #SBATCH --job-name=xm[{{ experiment_id }}.{{ identity }}]
-#SBATCH --dependency=singleton
 {% else %}
+{% if dependency %}
+#SBATCH {{ dependency.to_directive() }}
+{% endif %}
 #SBATCH --job-name=xm[{{ experiment_id }}]
 {% endif %}
 {% for directive in job.executor.to_directives() %}

xm_slurm/types.py

@@ -0,0 +1,23 @@
+import typing as tp
+
+InstanceT_contra = tp.TypeVar("InstanceT_contra", contravariant=True)
+GetterT_co = tp.TypeVar("GetterT_co", covariant=True)
+SetterT_co = tp.TypeVar("SetterT_co", contravariant=True)
+
+
+class Descriptor(tp.Protocol[GetterT_co, SetterT_co]):
+    def __set_name__(self, owner: tp.Type[tp.Any], name: str) -> None: ...
+
+    @tp.overload
+    def __get__(
+        self, instance: InstanceT_contra, owner: tp.Type[InstanceT_contra] | None = None
+    ) -> GetterT_co: ...
+
+    @tp.overload
+    def __get__(self, instance: None, owner: tp.Type[InstanceT_contra]) -> GetterT_co: ...
+
+    def __get__(
+        self, instance: InstanceT_contra | None, owner: tp.Type[InstanceT_contra] | None = None
+    ) -> GetterT_co: ...
+
+    def __set__(self, instance: tp.Any, value: SetterT_co) -> None: ...

xm_slurm/utils.py

@@ -35,17 +35,17 @@ class UserSet(Hashable, MutableSet[T]):
     def __repr__(self):
         return repr(self.data)
 
-    def add(self, item: T):
-        self.data.add(item)
-        self._on_add(item)
+    def add(self, value: T):
+        self.data.add(value)
+        self._on_add(value)
 
-    def remove(self, item: T):
-        self.data.remove(item)
-        self._on_remove(item)
+    def remove(self, value: T):
+        self.data.remove(value)
+        self._on_remove(value)
 
-    def discard(self, item: T):
-        self.data.discard(item)
-        self._on_discard(item)
+    def discard(self, value: T):
+        self.data.discard(value)
+        self._on_discard(value)
 
 
 @functools.cache
@@ -62,7 +62,15 @@ def find_project_root() -> pathlib.Path:
 
     pdir = launch_script_path.parent if launch_script_path else pathlib.Path.cwd().resolve()
     while pdir != pdir.parent:
-        if (pdir / "pyproject.toml").exists():
+        if (
+            (pdir / "pyproject.toml").exists()
+            or (pdir / "setup.py").exists()
+            or (pdir / "setup.cfg").exists()
+            or (pdir / "requirements.txt").exists()
+            or (pdir / "requirements.in").exists()
+            or (pdir / "uv.lock").exists()
+            or (pdir / ".venv").exists()
+        ):
             return pdir
         pdir = pdir.parent
 

xmanager_slurm-0.4.1.dist-info/METADATA

@@ -0,0 +1,26 @@
+Metadata-Version: 2.3
+Name: xmanager-slurm
+Version: 0.4.1
+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
+Classifier: Programming Language :: Python :: 3
+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: backoff>=2.2.1
+Requires-Dist: cloudpickle>=3.0.0
+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: xmanager>=0.5.0

xmanager_slurm-0.4.1.dist-info/RECORD

@@ -0,0 +1,44 @@
+xm_slurm/__init__.py,sha256=Ld2w7ofLlTieWOHP_Jb3f48-qtVQBjFXynxUm9WF8mc,1116
+xm_slurm/api.py,sha256=LeGgHz82t8Oay0Z1Ourv9-r-DBur3lhCUTnmmGhGFY4,18502
+xm_slurm/batching.py,sha256=GbKBsNz9w8gIc2fHLZpslC0e4K9YUfLXFHmjduRRCfQ,4385
+xm_slurm/config.py,sha256=GLLEkRLJxQW0urmHCLmwq_4ECmimEBQFl8Nz62SIo78,6787
+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=S3z8jSDL6AdyGYpzy_cCs03Mj0vgA4ZTqIe8APYor3E,6469
+xm_slurm/execution.py,sha256=i2oYH5RS-mHsHPwFDFZvo5qCudbgqBML-Hzq6DPNItw,25721
+xm_slurm/executors.py,sha256=fMtxGUCi4vEKmb_p4JEpqPUTh7L_f1LcR_TamMLAWNg,4667
+xm_slurm/experiment.py,sha256=trHapcYxPNKofzSqu7KZawML59tZ8FVjoEZYe2Wal7w,44521
+xm_slurm/job_blocks.py,sha256=_F8CKCs5BQFj40a2-mjG71HfacvWoBXBDPDKEaKTbXc,616
+xm_slurm/packageables.py,sha256=YZFTL6UWx9A_zyztTy1THUlj3pW1rA0cBPHJxD1LOJk,12884
+xm_slurm/resources.py,sha256=EaYDATVudrEDPKKdSZoWgfqPiidc6DMjIctmzLQmiH0,5683
+xm_slurm/status.py,sha256=WTWiDHi-ZHtwHRnDP0cGa-27zTSm6LkA-GCKsN-zBgg,6916
+xm_slurm/types.py,sha256=TsVykDm-LazVkrjeJrTwCMs4Q8APKhy7BTk0yKIhFNg,805
+xm_slurm/utils.py,sha256=ESjOkGT7bRSzIeZrUtZplSHP4oaH6VZ92y2woYdcyKM,2239
+xm_slurm/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+xm_slurm/contrib/clusters/__init__.py,sha256=vugR50D9fPJQN5bTd7cSArDGrA6pC-YJHMXrEyvr_Uw,2980
+xm_slurm/contrib/clusters/drac.py,sha256=tJeQFWFIpeZ1gD3j6AAJssNoLSiDkB-3lz1_ObnkRhc,5905
+xm_slurm/experimental/parameter_controller.py,sha256=b5LfglHV307F6QcPrHeZX5GJBtyOK9aQydke_SZ3Wto,8457
+xm_slurm/packaging/__init__.py,sha256=dh307yLpUT9KN7rJ1e9fYC6hegGKfZcGboUq9nGpDVQ,233
+xm_slurm/packaging/docker.py,sha256=TA8-QG09EdhF4K1ixrEboVFarF9LcURNHhzUXL-7Iqg,11518
+xm_slurm/packaging/registry.py,sha256=GrdmQg9MgSo38OiqOzMKWSkQyBuyryOfc3zcdgZ4CUE,1148
+xm_slurm/packaging/router.py,sha256=yPbdA9clrhly97cLgDsSRZG2LZRKE-oz8Hhdb7WtYqk,2070
+xm_slurm/packaging/utils.py,sha256=KI5s32rNTCfgwzY_7Ghck27jHKvKg5sl5_NEEqJbJqI,3999
+xm_slurm/scripts/_cloudpickle.py,sha256=dlJYf2SceOuUn8wi-ozuoYAQg71wqD2MUVOUCyOwWIY,647
+xm_slurm/scripts/cli.py,sha256=ZXqYOs8X23TYDdKxvV-wIa-0mTfpxSl4_Pli6TiKI7s,1435
+xm_slurm/templates/docker/docker-bake.hcl.j2,sha256=ClsFpj91Mr1VfA8L6eqBG3HQz0Z8VenF6mEfmAhQgUo,1498
+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=kYD32oUS1jUaARsNV1o6EFnIfLCNh5GMmck27b-5NRU,969
+xm_slurm/templates/slurm/job-array.bash.j2,sha256=iYtGMRDXgwwc2_8E3v4a30f3fKuq4zWgZHkxCXJ9iXc,567
+xm_slurm/templates/slurm/job-group.bash.j2,sha256=UkjfBE7jg9mepcUWaHZEAjkiXsIM1j_sLxLzxkteD-Y,1120
+xm_slurm/templates/slurm/job.bash.j2,sha256=v0xGYzagDdWW6Tg44qobGJLNSUP1Cf4CcekrPibYdrE,1864
+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=dMntzelhs8DqKyIpO9S6wzMfH2PDevmgvyjCW8Xc2dY,3222
+xm_slurm/templates/slurm/runtimes/podman.bash.j2,sha256=xKXYFvQvazMx0PgvmlRXR6eecoiBUl8y52dIzQtWkBE,1469
+xmanager_slurm-0.4.1.dist-info/METADATA,sha256=3mT4XIm8evv-5qw7oney4nYn3IasIA_l1rWz86XNOY8,954
+xmanager_slurm-0.4.1.dist-info/WHEEL,sha256=1yFddiXMmvYK7QYTqtRNtX66WJ0Mz8PYEiEUoOUUxRY,87
+xmanager_slurm-0.4.1.dist-info/entry_points.txt,sha256=_HLGmLgxuQLOPmF2gOFYDVq2HqtMVD_SzigHvUh8TCY,49
+xmanager_slurm-0.4.1.dist-info/licenses/LICENSE.md,sha256=IxstXr3MPHwTJ5jMrByHrQsR1ZAGQ2U_uz_4qzI_15Y,11756
+xmanager_slurm-0.4.1.dist-info/RECORD,,

{xmanager_slurm-0.3.2.dist-info → xmanager_slurm-0.4.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: pdm-backend (2.3.0)
+Generator: hatchling 1.25.0
 Root-Is-Purelib: true
 Tag: py3-none-any

xmanager_slurm-0.4.1.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+xm = xm_slurm.scripts.cli:main

xmanager_slurm-0.4.1.dist-info/licenses/LICENSE.md

@@ -0,0 +1,227 @@
+This work is dual-licensed under Apache 2.0 and MIT License attached below.
+You can choose between one of them if you use this work.
+
+---
+
+The MIT License (MIT)
+=====================
+
+Copyright © `2024` `Jesse Farebrother`
+
+Permission is hereby granted, free of charge, to any person
+obtaining a copy of this software and associated documentation
+files (the “Software”), to deal in the Software without
+restriction, including without limitation the rights to use,
+copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+OTHER DEALINGS IN THE SOFTWARE.
+
+---
+
+Apache License
+==============
+
+_Version 2.0, January 2004_  
+_&lt;<http://www.apache.org/licenses/>&gt;_
+
+### Terms and Conditions for use, reproduction, and distribution
+
+#### 1. Definitions
+
+“License” shall mean the terms and conditions for use, reproduction, and
+distribution as defined by Sections 1 through 9 of this document.
+
+“Licensor” shall mean the copyright owner or entity authorized by the copyright
+owner that is granting the License.
+
+“Legal Entity” shall mean the union of the acting entity and all other entities
+that control, are controlled by, or are under common control with that entity.
+For the purposes of this definition, “control” means **(i)** the power, direct or
+indirect, to cause the direction or management of such entity, whether by
+contract or otherwise, or **(ii)** ownership of fifty percent (50%) or more of the
+outstanding shares, or **(iii)** beneficial ownership of such entity.
+
+“You” (or “Your”) shall mean an individual or Legal Entity exercising
+permissions granted by this License.
+
+“Source” form shall mean the preferred form for making modifications, including
+but not limited to software source code, documentation source, and configuration
+files.
+
+“Object” form shall mean any form resulting from mechanical transformation or
+translation of a Source form, including but not limited to compiled object code,
+generated documentation, and conversions to other media types.
+
+“Work” shall mean the work of authorship, whether in Source or Object form, made
+available under the License, as indicated by a copyright notice that is included
+in or attached to the work (an example is provided in the Appendix below).
+
+“Derivative Works” shall mean any work, whether in Source or Object form, that
+is based on (or derived from) the Work and for which the editorial revisions,
+annotations, elaborations, or other modifications represent, as a whole, an
+original work of authorship. For the purposes of this License, Derivative Works
+shall not include works that remain separable from, or merely link (or bind by
+name) to the interfaces of, the Work and Derivative Works thereof.
+
+“Contribution” shall mean any work of authorship, including the original version
+of the Work and any modifications or additions to that Work or Derivative Works
+thereof, that is intentionally submitted to Licensor for inclusion in the Work
+by the copyright owner or by an individual or Legal Entity authorized to submit
+on behalf of the copyright owner. For the purposes of this definition,
+“submitted” means any form of electronic, verbal, or written communication sent
+to the Licensor or its representatives, including but not limited to
+communication on electronic mailing lists, source code control systems, and
+issue tracking systems that are managed by, or on behalf of, the Licensor for
+the purpose of discussing and improving the Work, but excluding communication
+that is conspicuously marked or otherwise designated in writing by the copyright
+owner as “Not a Contribution.”
+
+“Contributor” shall mean Licensor and any individual or Legal Entity on behalf
+of whom a Contribution has been received by Licensor and subsequently
+incorporated within the Work.
+
+#### 2. Grant of Copyright License
+
+Subject to the terms and conditions of this License, each Contributor hereby
+grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free,
+irrevocable copyright license to reproduce, prepare Derivative Works of,
+publicly display, publicly perform, sublicense, and distribute the Work and such
+Derivative Works in Source or Object form.
+
+#### 3. Grant of Patent License
+
+Subject to the terms and conditions of this License, each Contributor hereby
+grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free,
+irrevocable (except as stated in this section) patent license to make, have
+made, use, offer to sell, sell, import, and otherwise transfer the Work, where
+such license applies only to those patent claims licensable by such Contributor
+that are necessarily infringed by their Contribution(s) alone or by combination
+of their Contribution(s) with the Work to which such Contribution(s) was
+submitted. If You institute patent litigation against any entity (including a
+cross-claim or counterclaim in a lawsuit) alleging that the Work or a
+Contribution incorporated within the Work constitutes direct or contributory
+patent infringement, then any patent licenses granted to You under this License
+for that Work shall terminate as of the date such litigation is filed.
+
+#### 4. Redistribution
+
+You may reproduce and distribute copies of the Work or Derivative Works thereof
+in any medium, with or without modifications, and in Source or Object form,
+provided that You meet the following conditions:
+
+* **(a)** You must give any other recipients of the Work or Derivative Works a copy of
+this License; and
+* **(b)** You must cause any modified files to carry prominent notices stating that You
+changed the files; and
+* **(c)** You must retain, in the Source form of any Derivative Works that You distribute,
+all copyright, patent, trademark, and attribution notices from the Source form
+of the Work, excluding those notices that do not pertain to any part of the
+Derivative Works; and
+* **(d)** If the Work includes a “NOTICE” text file as part of its distribution, then any
+Derivative Works that You distribute must include a readable copy of the
+attribution notices contained within such NOTICE file, excluding those notices
+that do not pertain to any part of the Derivative Works, in at least one of the
+following places: within a NOTICE text file distributed as part of the
+Derivative Works; within the Source form or documentation, if provided along
+with the Derivative Works; or, within a display generated by the Derivative
+Works, if and wherever such third-party notices normally appear. The contents of
+the NOTICE file are for informational purposes only and do not modify the
+License. You may add Your own attribution notices within Derivative Works that
+You distribute, alongside or as an addendum to the NOTICE text from the Work,
+provided that such additional attribution notices cannot be construed as
+modifying the License.
+
+You may add Your own copyright statement to Your modifications and may provide
+additional or different license terms and conditions for use, reproduction, or
+distribution of Your modifications, or for any such Derivative Works as a whole,
+provided Your use, reproduction, and distribution of the Work otherwise complies
+with the conditions stated in this License.
+
+#### 5. Submission of Contributions
+
+Unless You explicitly state otherwise, any Contribution intentionally submitted
+for inclusion in the Work by You to the Licensor shall be under the terms and
+conditions of this License, without any additional terms or conditions.
+Notwithstanding the above, nothing herein shall supersede or modify the terms of
+any separate license agreement you may have executed with Licensor regarding
+such Contributions.
+
+#### 6. Trademarks
+
+This License does not grant permission to use the trade names, trademarks,
+service marks, or product names of the Licensor, except as required for
+reasonable and customary use in describing the origin of the Work and
+reproducing the content of the NOTICE file.
+
+#### 7. Disclaimer of Warranty
+
+Unless required by applicable law or agreed to in writing, Licensor provides the
+Work (and each Contributor provides its Contributions) on an “AS IS” BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied,
+including, without limitation, any warranties or conditions of TITLE,
+NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are
+solely responsible for determining the appropriateness of using or
+redistributing the Work and assume any risks associated with Your exercise of
+permissions under this License.
+
+#### 8. Limitation of Liability
+
+In no event and under no legal theory, whether in tort (including negligence),
+contract, or otherwise, unless required by applicable law (such as deliberate
+and grossly negligent acts) or agreed to in writing, shall any Contributor be
+liable to You for damages, including any direct, indirect, special, incidental,
+or consequential damages of any character arising as a result of this License or
+out of the use or inability to use the Work (including but not limited to
+damages for loss of goodwill, work stoppage, computer failure or malfunction, or
+any and all other commercial damages or losses), even if such Contributor has
+been advised of the possibility of such damages.
+
+#### 9. Accepting Warranty or Additional Liability
+
+While redistributing the Work or Derivative Works thereof, You may choose to
+offer, and charge a fee for, acceptance of support, warranty, indemnity, or
+other liability obligations and/or rights consistent with this License. However,
+in accepting such obligations, You may act only on Your own behalf and on Your
+sole responsibility, not on behalf of any other Contributor, and only if You
+agree to indemnify, defend, and hold each Contributor harmless for any liability
+incurred by, or claims asserted against, such Contributor by reason of your
+accepting any such warranty or additional liability.
+
+_END OF TERMS AND CONDITIONS_
+
+### APPENDIX: How to apply the Apache License to your work
+
+To apply the Apache License to your work, attach the following boilerplate
+notice, with the fields enclosed by brackets `[]` replaced with your own
+identifying information. (Don't include the brackets!) The text should be
+enclosed in the appropriate comment syntax for the file format. We also
+recommend that a file or class name and description of purpose be included on
+the same “printed page” as the copyright notice for easier identification within
+third-party archives.
+
+    Copyright [yyyy] [name of copyright owner]
+    
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+    
+      http://www.apache.org/licenses/LICENSE-2.0
+    
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.

xm_slurm/packaging/docker/__init__.py

@@ -1,75 +0,0 @@
-import dataclasses
-import functools
-from typing import Sequence
-
-from absl import flags
-from xmanager import xm
-
-from xm_slurm.executables import Dockerfile, DockerImage, ImageURI, RemoteImage
-from xm_slurm.executors import SlurmSpec
-from xm_slurm.packaging import registry
-from xm_slurm.packaging.docker.abc import DockerClient
-
-FLAGS = flags.FLAGS
-REMOTE_BUILD = flags.DEFINE_enum(
-    "xm_builder", "local", ["local", "gcp", "azure"], "Remote build provider."
-)
-
-IndexedContainer = registry.IndexedContainer
-
-
-@functools.cache
-def docker_client() -> DockerClient:
-    match REMOTE_BUILD.value:
-        case "local":
-            from xm_slurm.packaging.docker.local import LocalDockerClient
-
-            return LocalDockerClient()
-        case "gcp":
-            from xm_slurm.packaging.docker.cloud import GoogleCloudRemoteDockerClient
-
-            return GoogleCloudRemoteDockerClient()
-        case "azure":
-            raise NotImplementedError("Azure remote build is not yet supported.")
-        case _:
-            raise ValueError(f"Unknown remote build provider: {REMOTE_BUILD.value}")
-
-
-@registry.register(Dockerfile)
-def _(
-    targets: Sequence[IndexedContainer[xm.Packageable]],
-) -> list[IndexedContainer[RemoteImage]]:
-    return docker_client().bake(targets=targets)
-
-
-@registry.register(DockerImage)
-def _(
-    targets: Sequence[IndexedContainer[xm.Packageable]],
-) -> list[IndexedContainer[RemoteImage]]:
-    """Build Docker images, this is essentially a passthrough."""
-    images = []
-    client = docker_client()
-    for target in targets:
-        assert isinstance(target.value.executable_spec, DockerImage)
-        assert isinstance(target.value.executor_spec, SlurmSpec)
-        if target.value.executor_spec.tag is not None:
-            raise ValueError(
-                "Executable `DockerImage` should not be tagged via `SlurmSpec`. "
-                "The image URI is provided by the `DockerImage` itself."
-            )
-
-        uri = ImageURI(target.value.executable_spec.image)
-        images.append(
-            dataclasses.replace(
-                target,
-                value=RemoteImage(  # type: ignore
-                    image=str(uri),
-                    workdir=target.value.executable_spec.workdir,
-                    args=target.value.args,
-                    env_vars=target.value.env_vars,
-                    credentials=client.credentials(hostname=uri.domain),
-                ),
-            )
-        )
-
-    return images

xm_slurm/packaging/docker/abc.py

@@ -1,112 +0,0 @@
-import abc
-import collections.abc
-import dataclasses
-import functools
-import os
-from typing import Literal, Mapping, Protocol, Sequence
-
-import jinja2 as j2
-from xmanager import xm
-
-from xm_slurm.executables import RemoteImage, RemoteRepositoryCredentials
-from xm_slurm.packaging.registry import IndexedContainer
-
-
-class DockerCommandProtocol(Protocol):
-    def to_args(self) -> xm.SequentialArgs: ...
-
-
-@dataclasses.dataclass(frozen=True, kw_only=True)
-class DockerBakeCommand(DockerCommandProtocol):
-    targets: str | Sequence[str] | None = None
-    builder: str | None = None
-    files: str | os.PathLike[str] | Sequence[os.PathLike[str] | str] | None = None
-    load: bool = False
-    cache: bool = True
-    print: bool = False
-    pull: bool = False
-    push: bool = False
-    metadata_file: str | os.PathLike[str] | None = None
-    progress: Literal["auto", "plain", "tty"] = "auto"
-    set: Mapping[str, str] | None = None
-
-    def to_args(self) -> xm.SequentialArgs:
-        files = self.files
-        if files is None:
-            files = []
-        if not isinstance(files, collections.abc.Sequence):
-            files = [files]
-
-        targets = self.targets
-        if targets is None:
-            targets = []
-        elif not isinstance(targets, collections.abc.Sequence):
-            targets = [targets]
-
-        return xm.merge_args(
-            ["buildx", "bake"],
-            [f"--progress={self.progress}"],
-            [f"--builder={self.builder}"] if self.builder else [],
-            [f"--metadata-file={self.metadata_file}"] if self.metadata_file else [],
-            ["--print"] if self.print else [],
-            ["--push"] if self.push else [],
-            ["--pull"] if self.pull else [],
-            ["--load"] if self.load else [],
-            ["--no-cache"] if not self.cache else [],
-            [f"--file={file}" for file in files],
-            [f"--set={key}={value}" for key, value in self.set.items()] if self.set else [],
-            targets,
-        )
-
-
-@dataclasses.dataclass(frozen=True, kw_only=True)
-class DockerPullCommand(DockerCommandProtocol):
-    image: str
-
-    def to_args(self) -> xm.SequentialArgs:
-        return xm.merge_args(["pull", self.image])
-
-
-@dataclasses.dataclass(frozen=True, kw_only=True)
-class DockerLoginCommand(DockerCommandProtocol):
-    server: str
-    username: str
-    password: str | None = None
-    password_stdin: bool = False
-
-    def __post_init__(self):
-        if self.password is None and not self.password_stdin:
-            raise ValueError("Either password or password_stdin must be set")
-        if self.password is not None and self.password_stdin:
-            raise ValueError("Only one of password or password_stdin must be set")
-
-    def to_args(self) -> xm.SequentialArgs:
-        return xm.merge_args(
-            ["login", "--username", self.username],
-            ["--password", self.password] if self.password else [],
-            ["--password-stdin"] if self.password_stdin else [],
-            [self.server],
-        )
-
-
-@dataclasses.dataclass(frozen=True, kw_only=True)
-class DockerVersionCommand(DockerCommandProtocol):
-    def to_args(self) -> xm.SequentialArgs:
-        return xm.merge_args(["buildx", "version"])
-
-
-class DockerClient(abc.ABC):
-    @functools.cached_property
-    def _bake_template(self) -> j2.Template:
-        template_loader = j2.PackageLoader("xm_slurm", "templates/docker")
-        template_env = j2.Environment(loader=template_loader, trim_blocks=True, lstrip_blocks=False)
-
-        return template_env.get_template("docker-bake.hcl.j2")
-
-    @abc.abstractmethod
-    def credentials(self, *, hostname: str) -> RemoteRepositoryCredentials | None: ...
-
-    @abc.abstractmethod
-    def bake(
-        self, *, targets: Sequence[IndexedContainer[xm.Packageable]]
-    ) -> list[IndexedContainer[RemoteImage]]: ...