diff --git a/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/INSTALLER b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/METADATA b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..0cf2d721b265bf32a2ff733b34f8d2915c2c1aa6 --- /dev/null +++ b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/METADATA @@ -0,0 +1,100 @@ +Metadata-Version: 2.4 +Name: absl-py +Version: 2.4.0 +Summary: Abseil Python Common Libraries, see https://github.com/abseil/abseil-py. +Project-URL: Changelog, https://github.com/abseil/abseil-py/blob/main/CHANGELOG.md +Project-URL: Documentation, https://abseil.io/docs/python/ +Project-URL: Issues, https://github.com/abseil/abseil-py/issues +Project-URL: Source, https://github.com/abseil/abseil-py +Author: The Abseil Authors +License-Expression: Apache-2.0 +License-File: AUTHORS +License-File: LICENSE +Classifier: Intended Audience :: Developers +Classifier: License :: OSI Approved :: Apache Software License +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Requires-Python: >=3.10 +Description-Content-Type: text/markdown + +[![Package version](https://img.shields.io/pypi/v/absl-py)](https://pypi.org/project/absl-py) +[![Supported Python versions](https://img.shields.io/pypi/pyversions/absl-py.svg?style=flat-square)](https://pypi.org/project/absl-py) +[![License](https://img.shields.io/github/license/abseil/abseil-py)](https://github.com/abseil/abseil-py/blob/main/LICENSE) +[![Build Status](https://github.com/abseil/abseil-py/actions/workflows/test.yml/badge.svg?branch=main)](https://github.com/abseil/abseil-py/actions) +[![Overall downloads](https://pepy.tech/badge/absl-py)](https://pepy.tech/project/absl-py) +[![Last month downloads](https://pepy.tech/badge/absl-py/month)](https://pepy.tech/project/absl-py) + +# Abseil Python Common Libraries + +This repository is a collection of Python library code for building Python +applications. The code is collected from Google's own Python code base, and has +been extensively tested and used in production. + +## Features + +* Simple application startup +* Distributed commandline flags system +* Custom logging module with additional features +* Testing utilities + +## Getting Started + +### Installation + +To install the package, simply run: + +```bash +pip install absl-py +``` + +Or install from source: + +```bash +pip install . +``` + +### Running Tests + +To run Abseil tests, you can clone the git repo and run +[bazel](https://bazel.build/): + +```bash +git clone https://github.com/abseil/abseil-py.git +cd abseil-py +bazel test absl/... +``` + +Please also validate the type annotations against the latest mypy: + +```bash +pip install mypy +mypy absl +``` + +### Example Code + +Please refer to +[smoke_tests/sample_app.py](https://github.com/abseil/abseil-py/blob/main/smoke_tests/sample_app.py) +as an example to get started. + +## Documentation + +See the [Abseil Python Developer Guide](https://abseil.io/docs/python/). + +## Future Releases + +The current repository includes an initial set of libraries for early adoption. +More components and interoperability with Abseil C++ Common Libraries +will come in future releases. + +## License + +The Abseil Python library is licensed under the terms of the Apache +license. See [LICENSE](LICENSE) for more information. diff --git a/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/RECORD b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..116f11d49e7ffd0ab0ccd6c6e5d9dff102bc25d4 --- /dev/null +++ b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/RECORD @@ -0,0 +1,53 @@ +absl/__init__.py,sha256=tQJRn_5rbQ5gWoXdkGAV-fJSGFCOGH6fFBI1NxLx5ho,607 +absl/__pycache__/__init__.cpython-312.pyc,, +absl/__pycache__/app.cpython-312.pyc,, +absl/__pycache__/command_name.cpython-312.pyc,, +absl/app.py,sha256=zC4tAKcK054F7kSrR7pzSCEDqNcJKLVb2AkvjIlmvl8,17065 +absl/app.pyi,sha256=YTcQS7o1DBFyFNgMpdykmvP3g4OljSwp3BuX9sSSLIM,1814 +absl/command_name.py,sha256=FgWUMHmlX0yQxEuMPXxFxn5ayWWZSLB0cq8Vx361TpU,2283 +absl/flags/__init__.py,sha256=n_uLeSK-15_1DKLb0wXivmt6A8xvADDC87D1yoJdqVU,7665 +absl/flags/__pycache__/__init__.cpython-312.pyc,, +absl/flags/__pycache__/_argument_parser.cpython-312.pyc,, +absl/flags/__pycache__/_defines.cpython-312.pyc,, +absl/flags/__pycache__/_exceptions.cpython-312.pyc,, +absl/flags/__pycache__/_flag.cpython-312.pyc,, +absl/flags/__pycache__/_flagvalues.cpython-312.pyc,, +absl/flags/__pycache__/_helpers.cpython-312.pyc,, +absl/flags/__pycache__/_validators.cpython-312.pyc,, +absl/flags/__pycache__/_validators_classes.cpython-312.pyc,, +absl/flags/__pycache__/argparse_flags.cpython-312.pyc,, +absl/flags/_argument_parser.py,sha256=XFIDxu1I8NjgBZhUdzSLUwDHhXFvOC-SbI8FsVmWZ_0,20531 +absl/flags/_defines.py,sha256=P3DLHYAaSgXSBlRgBN6QOS6oRcNUYNtzeapL5EQJ5LA,52499 +absl/flags/_exceptions.py,sha256=FZzlzhvkjqPImTxXqbS1pSPYKr_TvtOd5ellvoiVLDI,3619 +absl/flags/_flag.py,sha256=fOv_KC7t1e-BoabTP3iYHR1dUkTf2yhgfAWonjhByEQ,19891 +absl/flags/_flagvalues.py,sha256=QqCMTxHIJdwqYwlYhZ76nuXzsJQtF20RTzq-x0S_tPg,55319 +absl/flags/_helpers.py,sha256=0cWFV4XFLkcHWu-du6es8sCsWl26k1dZ6fqCPsWZ68k,13301 +absl/flags/_validators.py,sha256=VcsJtZzohliNxsI974NECYpeozD8rswHNHXggrQ4BLo,14140 +absl/flags/_validators_classes.py,sha256=PGUWzO7v3wPOHb9leIKKzry3q-pPeKCoMB_O7prLdnY,6093 +absl/flags/argparse_flags.py,sha256=usJudgMpy3P6Vvq7-LmJNa2Rj3ygHM3hwDTGd1mbAzc,14386 +absl/logging/__init__.py,sha256=3quWtll8Cb949CYs-VOC1AnxdBqClX8ab55VB28knVs,43883 +absl/logging/__init__.pyi,sha256=7yyq5rAJfkBFRaD8iqOTsLW-xhMAIDJ5ub4r3Has04U,5984 +absl/logging/__pycache__/__init__.cpython-312.pyc,, +absl/logging/__pycache__/converter.cpython-312.pyc,, +absl/logging/converter.py,sha256=6eBymfv9UNkog0BGat4HPWlxC_oSqvHcQ46jnSdtaMg,6323 +absl/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +absl/testing/__init__.py,sha256=7cM57swk2T1Hc5wxmt-JpcaR6xfdPJyL_lyRqgODvuM,584 +absl/testing/__pycache__/__init__.cpython-312.pyc,, +absl/testing/__pycache__/_bazelize_command.cpython-312.pyc,, +absl/testing/__pycache__/_pretty_print_reporter.cpython-312.pyc,, +absl/testing/__pycache__/absltest.cpython-312.pyc,, +absl/testing/__pycache__/flagsaver.cpython-312.pyc,, +absl/testing/__pycache__/parameterized.cpython-312.pyc,, +absl/testing/__pycache__/xml_reporter.cpython-312.pyc,, +absl/testing/_bazelize_command.py,sha256=qpioV02ln2sBBJ9kdlHgNpKk8_wxdz2hJGKbG6EWZMI,2287 +absl/testing/_pretty_print_reporter.py,sha256=PZh9NXSXBbXDi0FOk-BOmpse8LXa92Er16tgyBRogMs,3065 +absl/testing/absltest.py,sha256=CkUgJ54S4LJ1RwjNISv4pEW_N9kZBrNDpnzQ1YGPOlc,104586 +absl/testing/flagsaver.py,sha256=BeftLpNxJctYtd041tECgCSglQVFKSsQ4KBFT9D_m5k,13291 +absl/testing/parameterized.py,sha256=TpTlWTUXjikGeUDE45AgubvnmHsPbFNj-wUtAwT-e6E,27817 +absl/testing/xml_reporter.py,sha256=Q0keNHVOh3L9zXwNIzWCjXcWYlCM9ZJxdjMaYyHF3dw,21424 +absl_py-2.4.0.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +absl_py-2.4.0.dist-info/METADATA,sha256=Eh8VboTinee1rj1-UDeBjw_KiicKBxJBi6H0zB1Sq0g,3283 +absl_py-2.4.0.dist-info/RECORD,, +absl_py-2.4.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87 +absl_py-2.4.0.dist-info/licenses/AUTHORS,sha256=YoLudsylaQg7W5mLn4FroQMuEnuNx8RpQrhkd_xvv6U,296 +absl_py-2.4.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358 diff --git a/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/WHEEL b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..ae8ec1bdaa94d726ceb907542d76cbd5d38cafcd --- /dev/null +++ b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/WHEEL @@ -0,0 +1,4 @@ +Wheel-Version: 1.0 +Generator: hatchling 1.28.0 +Root-Is-Purelib: true +Tag: py3-none-any diff --git a/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/licenses/AUTHORS b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/licenses/AUTHORS new file mode 100644 index 0000000000000000000000000000000000000000..23b11ada16bb8e69695cf52e5994784d98054e0d --- /dev/null +++ b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/licenses/AUTHORS @@ -0,0 +1,7 @@ +# This is the list of Abseil authors for copyright purposes. +# +# This does not necessarily list everyone who has contributed code, since in +# some cases, their employer may be the copyright holder. To see the full list +# of contributors, see the revision history in source control. + +Google Inc. diff --git a/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/licenses/LICENSE b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..d645695673349e3947e8e5ae42332d0ac3164cd7 --- /dev/null +++ b/lib/python3.12/site-packages/absl_py-2.4.0.dist-info/licenses/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + 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. diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..40df66499bb87e45daf41d7517cd413deaa8f5f8 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/__main__.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/__main__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ec04d83f3b25002f328f89714f43e214c721a404 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/__main__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/_compat.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/_compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5eeb27b6b3e3a8ab3488eade454327158cfd2682 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/_compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/applications.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/applications.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..088bcd52fffc0eb283ca1daedfb56377fa088521 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/applications.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/background.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/background.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a99b973d3c011a7e2d86769fe5756c7285ae7bc3 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/background.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/cli.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/cli.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..102891f9dfed71280ddfd25a73c81eba67a67557 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/cli.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/concurrency.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/concurrency.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c2a9aab4fe57ce4f1041323cba96e6874fdbedc Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/concurrency.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/datastructures.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/datastructures.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4f9c75a86c084601a332d87057d8c14fbc74eaef Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/datastructures.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/encoders.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/encoders.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8425c04d7d1102dc89b76b4cfb33416eff64c56e Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/encoders.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/exception_handlers.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/exception_handlers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5597c34264230dd9bdac51aab3bcbf2214833ca4 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/exception_handlers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/exceptions.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/exceptions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..03628f7cb38903ff58cde8845f520f9320c16e98 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/exceptions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/logger.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d1fb8e039b8730b0c4d97abadda84b0fb6c3c3d2 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/params.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/params.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80a8dc40001f70bbbba02cd09a57ded9e6bc4305 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/params.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/requests.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/requests.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7a762aaf841aa2d35a983854d612a7b1959006c2 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/requests.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/responses.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/responses.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d46ade6d8c4c301b3db8443b6d87a0df2dc337b3 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/responses.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/routing.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/routing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4b670c49732c5601dc87c00c81f6c7e9c9276f44 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/routing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/staticfiles.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/staticfiles.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e8660496dc6e8c2436be1df93164c49f7ea3f08 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/staticfiles.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/templating.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/templating.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5c633a90034238e52863995186a62175f476d98 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/templating.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/testclient.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/testclient.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b2b665f599c32ddb57b3fa999a41e9610510c1e9 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/testclient.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1144846eed88a29ac624bb1e5984e055a71a1c9d Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83fea0fe5671cdbbc0b7a049dcb9480698dcb7c4 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/__pycache__/websockets.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/__pycache__/websockets.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa84ae0db342178e11e46bb8cf7195cad247f39b Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/__pycache__/websockets.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/dependencies/__init__.py b/lib/python3.12/site-packages/fastapi/dependencies/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46eb78ffbda866f8c32554588272521f5be888b6 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/models.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/models.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a8316cf9804e43e0f39c9c887c039df4564b3da Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/models.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8af38c3b167b3cc9f97c50713a598df1487ed9b4 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/dependencies/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/dependencies/models.py b/lib/python3.12/site-packages/fastapi/dependencies/models.py new file mode 100644 index 0000000000000000000000000000000000000000..418c117259aa31de5c9d9bf6a4e5cfcbb5dab2d8 --- /dev/null +++ b/lib/python3.12/site-packages/fastapi/dependencies/models.py @@ -0,0 +1,37 @@ +from dataclasses import dataclass, field +from typing import Any, Callable, List, Optional, Sequence, Tuple + +from fastapi._compat import ModelField +from fastapi.security.base import SecurityBase + + +@dataclass +class SecurityRequirement: + security_scheme: SecurityBase + scopes: Optional[Sequence[str]] = None + + +@dataclass +class Dependant: + path_params: List[ModelField] = field(default_factory=list) + query_params: List[ModelField] = field(default_factory=list) + header_params: List[ModelField] = field(default_factory=list) + cookie_params: List[ModelField] = field(default_factory=list) + body_params: List[ModelField] = field(default_factory=list) + dependencies: List["Dependant"] = field(default_factory=list) + security_requirements: List[SecurityRequirement] = field(default_factory=list) + name: Optional[str] = None + call: Optional[Callable[..., Any]] = None + request_param_name: Optional[str] = None + websocket_param_name: Optional[str] = None + http_connection_param_name: Optional[str] = None + response_param_name: Optional[str] = None + background_tasks_param_name: Optional[str] = None + security_scopes_param_name: Optional[str] = None + security_scopes: Optional[List[str]] = None + use_cache: bool = True + path: Optional[str] = None + cache_key: Tuple[Optional[Callable[..., Any]], Tuple[str, ...]] = field(init=False) + + def __post_init__(self) -> None: + self.cache_key = (self.call, tuple(sorted(set(self.security_scopes or [])))) diff --git a/lib/python3.12/site-packages/fastapi/dependencies/utils.py b/lib/python3.12/site-packages/fastapi/dependencies/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..84dfa4d0306a6a885004cbc5cc6f5a835fd46955 --- /dev/null +++ b/lib/python3.12/site-packages/fastapi/dependencies/utils.py @@ -0,0 +1,980 @@ +import inspect +from contextlib import AsyncExitStack, contextmanager +from copy import copy, deepcopy +from dataclasses import dataclass +from typing import ( + Any, + Callable, + Coroutine, + Dict, + ForwardRef, + List, + Mapping, + Optional, + Sequence, + Tuple, + Type, + Union, + cast, +) + +import anyio +from fastapi import params +from fastapi._compat import ( + PYDANTIC_V2, + ErrorWrapper, + ModelField, + RequiredParam, + Undefined, + _regenerate_error_with_loc, + copy_field_info, + create_body_model, + evaluate_forwardref, + field_annotation_is_scalar, + get_annotation_from_field_info, + get_cached_model_fields, + get_missing_field_error, + is_bytes_field, + is_bytes_sequence_field, + is_scalar_field, + is_scalar_sequence_field, + is_sequence_field, + is_uploadfile_or_nonable_uploadfile_annotation, + is_uploadfile_sequence_annotation, + lenient_issubclass, + sequence_types, + serialize_sequence_value, + value_is_sequence, +) +from fastapi.background import BackgroundTasks +from fastapi.concurrency import ( + asynccontextmanager, + contextmanager_in_threadpool, +) +from fastapi.dependencies.models import Dependant, SecurityRequirement +from fastapi.logger import logger +from fastapi.security.base import SecurityBase +from fastapi.security.oauth2 import OAuth2, SecurityScopes +from fastapi.security.open_id_connect_url import OpenIdConnect +from fastapi.utils import create_model_field, get_path_param_names +from pydantic import BaseModel +from pydantic.fields import FieldInfo +from starlette.background import BackgroundTasks as StarletteBackgroundTasks +from starlette.concurrency import run_in_threadpool +from starlette.datastructures import ( + FormData, + Headers, + ImmutableMultiDict, + QueryParams, + UploadFile, +) +from starlette.requests import HTTPConnection, Request +from starlette.responses import Response +from starlette.websockets import WebSocket +from typing_extensions import Annotated, get_args, get_origin + +multipart_not_installed_error = ( + 'Form data requires "python-multipart" to be installed. \n' + 'You can install "python-multipart" with: \n\n' + "pip install python-multipart\n" +) +multipart_incorrect_install_error = ( + 'Form data requires "python-multipart" to be installed. ' + 'It seems you installed "multipart" instead. \n' + 'You can remove "multipart" with: \n\n' + "pip uninstall multipart\n\n" + 'And then install "python-multipart" with: \n\n' + "pip install python-multipart\n" +) + + +def ensure_multipart_is_installed() -> None: + try: + from python_multipart import __version__ + + # Import an attribute that can be mocked/deleted in testing + assert __version__ > "0.0.12" + except (ImportError, AssertionError): + try: + # __version__ is available in both multiparts, and can be mocked + from multipart import __version__ # type: ignore[no-redef,import-untyped] + + assert __version__ + try: + # parse_options_header is only available in the right multipart + from multipart.multipart import ( # type: ignore[import-untyped] + parse_options_header, + ) + + assert parse_options_header + except ImportError: + logger.error(multipart_incorrect_install_error) + raise RuntimeError(multipart_incorrect_install_error) from None + except ImportError: + logger.error(multipart_not_installed_error) + raise RuntimeError(multipart_not_installed_error) from None + + +def get_param_sub_dependant( + *, + param_name: str, + depends: params.Depends, + path: str, + security_scopes: Optional[List[str]] = None, +) -> Dependant: + assert depends.dependency + return get_sub_dependant( + depends=depends, + dependency=depends.dependency, + path=path, + name=param_name, + security_scopes=security_scopes, + ) + + +def get_parameterless_sub_dependant(*, depends: params.Depends, path: str) -> Dependant: + assert callable(depends.dependency), ( + "A parameter-less dependency must have a callable dependency" + ) + return get_sub_dependant(depends=depends, dependency=depends.dependency, path=path) + + +def get_sub_dependant( + *, + depends: params.Depends, + dependency: Callable[..., Any], + path: str, + name: Optional[str] = None, + security_scopes: Optional[List[str]] = None, +) -> Dependant: + security_requirement = None + security_scopes = security_scopes or [] + if isinstance(depends, params.Security): + dependency_scopes = depends.scopes + security_scopes.extend(dependency_scopes) + if isinstance(dependency, SecurityBase): + use_scopes: List[str] = [] + if isinstance(dependency, (OAuth2, OpenIdConnect)): + use_scopes = security_scopes + security_requirement = SecurityRequirement( + security_scheme=dependency, scopes=use_scopes + ) + sub_dependant = get_dependant( + path=path, + call=dependency, + name=name, + security_scopes=security_scopes, + use_cache=depends.use_cache, + ) + if security_requirement: + sub_dependant.security_requirements.append(security_requirement) + return sub_dependant + + +CacheKey = Tuple[Optional[Callable[..., Any]], Tuple[str, ...]] + + +def get_flat_dependant( + dependant: Dependant, + *, + skip_repeats: bool = False, + visited: Optional[List[CacheKey]] = None, +) -> Dependant: + if visited is None: + visited = [] + visited.append(dependant.cache_key) + + flat_dependant = Dependant( + path_params=dependant.path_params.copy(), + query_params=dependant.query_params.copy(), + header_params=dependant.header_params.copy(), + cookie_params=dependant.cookie_params.copy(), + body_params=dependant.body_params.copy(), + security_requirements=dependant.security_requirements.copy(), + use_cache=dependant.use_cache, + path=dependant.path, + ) + for sub_dependant in dependant.dependencies: + if skip_repeats and sub_dependant.cache_key in visited: + continue + flat_sub = get_flat_dependant( + sub_dependant, skip_repeats=skip_repeats, visited=visited + ) + flat_dependant.path_params.extend(flat_sub.path_params) + flat_dependant.query_params.extend(flat_sub.query_params) + flat_dependant.header_params.extend(flat_sub.header_params) + flat_dependant.cookie_params.extend(flat_sub.cookie_params) + flat_dependant.body_params.extend(flat_sub.body_params) + flat_dependant.security_requirements.extend(flat_sub.security_requirements) + return flat_dependant + + +def _get_flat_fields_from_params(fields: List[ModelField]) -> List[ModelField]: + if not fields: + return fields + first_field = fields[0] + if len(fields) == 1 and lenient_issubclass(first_field.type_, BaseModel): + fields_to_extract = get_cached_model_fields(first_field.type_) + return fields_to_extract + return fields + + +def get_flat_params(dependant: Dependant) -> List[ModelField]: + flat_dependant = get_flat_dependant(dependant, skip_repeats=True) + path_params = _get_flat_fields_from_params(flat_dependant.path_params) + query_params = _get_flat_fields_from_params(flat_dependant.query_params) + header_params = _get_flat_fields_from_params(flat_dependant.header_params) + cookie_params = _get_flat_fields_from_params(flat_dependant.cookie_params) + return path_params + query_params + header_params + cookie_params + + +def get_typed_signature(call: Callable[..., Any]) -> inspect.Signature: + signature = inspect.signature(call) + globalns = getattr(call, "__globals__", {}) + typed_params = [ + inspect.Parameter( + name=param.name, + kind=param.kind, + default=param.default, + annotation=get_typed_annotation(param.annotation, globalns), + ) + for param in signature.parameters.values() + ] + typed_signature = inspect.Signature(typed_params) + return typed_signature + + +def get_typed_annotation(annotation: Any, globalns: Dict[str, Any]) -> Any: + if isinstance(annotation, str): + annotation = ForwardRef(annotation) + annotation = evaluate_forwardref(annotation, globalns, globalns) + return annotation + + +def get_typed_return_annotation(call: Callable[..., Any]) -> Any: + signature = inspect.signature(call) + annotation = signature.return_annotation + + if annotation is inspect.Signature.empty: + return None + + globalns = getattr(call, "__globals__", {}) + return get_typed_annotation(annotation, globalns) + + +def get_dependant( + *, + path: str, + call: Callable[..., Any], + name: Optional[str] = None, + security_scopes: Optional[List[str]] = None, + use_cache: bool = True, +) -> Dependant: + path_param_names = get_path_param_names(path) + endpoint_signature = get_typed_signature(call) + signature_params = endpoint_signature.parameters + dependant = Dependant( + call=call, + name=name, + path=path, + security_scopes=security_scopes, + use_cache=use_cache, + ) + for param_name, param in signature_params.items(): + is_path_param = param_name in path_param_names + param_details = analyze_param( + param_name=param_name, + annotation=param.annotation, + value=param.default, + is_path_param=is_path_param, + ) + if param_details.depends is not None: + sub_dependant = get_param_sub_dependant( + param_name=param_name, + depends=param_details.depends, + path=path, + security_scopes=security_scopes, + ) + dependant.dependencies.append(sub_dependant) + continue + if add_non_field_param_to_dependency( + param_name=param_name, + type_annotation=param_details.type_annotation, + dependant=dependant, + ): + assert param_details.field is None, ( + f"Cannot specify multiple FastAPI annotations for {param_name!r}" + ) + continue + assert param_details.field is not None + if isinstance(param_details.field.field_info, params.Body): + dependant.body_params.append(param_details.field) + else: + add_param_to_fields(field=param_details.field, dependant=dependant) + return dependant + + +def add_non_field_param_to_dependency( + *, param_name: str, type_annotation: Any, dependant: Dependant +) -> Optional[bool]: + if lenient_issubclass(type_annotation, Request): + dependant.request_param_name = param_name + return True + elif lenient_issubclass(type_annotation, WebSocket): + dependant.websocket_param_name = param_name + return True + elif lenient_issubclass(type_annotation, HTTPConnection): + dependant.http_connection_param_name = param_name + return True + elif lenient_issubclass(type_annotation, Response): + dependant.response_param_name = param_name + return True + elif lenient_issubclass(type_annotation, StarletteBackgroundTasks): + dependant.background_tasks_param_name = param_name + return True + elif lenient_issubclass(type_annotation, SecurityScopes): + dependant.security_scopes_param_name = param_name + return True + return None + + +@dataclass +class ParamDetails: + type_annotation: Any + depends: Optional[params.Depends] + field: Optional[ModelField] + + +def analyze_param( + *, + param_name: str, + annotation: Any, + value: Any, + is_path_param: bool, +) -> ParamDetails: + field_info = None + depends = None + type_annotation: Any = Any + use_annotation: Any = Any + if annotation is not inspect.Signature.empty: + use_annotation = annotation + type_annotation = annotation + # Extract Annotated info + if get_origin(use_annotation) is Annotated: + annotated_args = get_args(annotation) + type_annotation = annotated_args[0] + fastapi_annotations = [ + arg + for arg in annotated_args[1:] + if isinstance(arg, (FieldInfo, params.Depends)) + ] + fastapi_specific_annotations = [ + arg + for arg in fastapi_annotations + if isinstance(arg, (params.Param, params.Body, params.Depends)) + ] + if fastapi_specific_annotations: + fastapi_annotation: Union[FieldInfo, params.Depends, None] = ( + fastapi_specific_annotations[-1] + ) + else: + fastapi_annotation = None + # Set default for Annotated FieldInfo + if isinstance(fastapi_annotation, FieldInfo): + # Copy `field_info` because we mutate `field_info.default` below. + field_info = copy_field_info( + field_info=fastapi_annotation, annotation=use_annotation + ) + assert ( + field_info.default is Undefined or field_info.default is RequiredParam + ), ( + f"`{field_info.__class__.__name__}` default value cannot be set in" + f" `Annotated` for {param_name!r}. Set the default value with `=` instead." + ) + if value is not inspect.Signature.empty: + assert not is_path_param, "Path parameters cannot have default values" + field_info.default = value + else: + field_info.default = RequiredParam + # Get Annotated Depends + elif isinstance(fastapi_annotation, params.Depends): + depends = fastapi_annotation + # Get Depends from default value + if isinstance(value, params.Depends): + assert depends is None, ( + "Cannot specify `Depends` in `Annotated` and default value" + f" together for {param_name!r}" + ) + assert field_info is None, ( + "Cannot specify a FastAPI annotation in `Annotated` and `Depends` as a" + f" default value together for {param_name!r}" + ) + depends = value + # Get FieldInfo from default value + elif isinstance(value, FieldInfo): + assert field_info is None, ( + "Cannot specify FastAPI annotations in `Annotated` and default value" + f" together for {param_name!r}" + ) + field_info = value + if PYDANTIC_V2: + field_info.annotation = type_annotation + + # Get Depends from type annotation + if depends is not None and depends.dependency is None: + # Copy `depends` before mutating it + depends = copy(depends) + depends.dependency = type_annotation + + # Handle non-param type annotations like Request + if lenient_issubclass( + type_annotation, + ( + Request, + WebSocket, + HTTPConnection, + Response, + StarletteBackgroundTasks, + SecurityScopes, + ), + ): + assert depends is None, f"Cannot specify `Depends` for type {type_annotation!r}" + assert field_info is None, ( + f"Cannot specify FastAPI annotation for type {type_annotation!r}" + ) + # Handle default assignations, neither field_info nor depends was not found in Annotated nor default value + elif field_info is None and depends is None: + default_value = value if value is not inspect.Signature.empty else RequiredParam + if is_path_param: + # We might check here that `default_value is RequiredParam`, but the fact is that the same + # parameter might sometimes be a path parameter and sometimes not. See + # `tests/test_infer_param_optionality.py` for an example. + field_info = params.Path(annotation=use_annotation) + elif is_uploadfile_or_nonable_uploadfile_annotation( + type_annotation + ) or is_uploadfile_sequence_annotation(type_annotation): + field_info = params.File(annotation=use_annotation, default=default_value) + elif not field_annotation_is_scalar(annotation=type_annotation): + field_info = params.Body(annotation=use_annotation, default=default_value) + else: + field_info = params.Query(annotation=use_annotation, default=default_value) + + field = None + # It's a field_info, not a dependency + if field_info is not None: + # Handle field_info.in_ + if is_path_param: + assert isinstance(field_info, params.Path), ( + f"Cannot use `{field_info.__class__.__name__}` for path param" + f" {param_name!r}" + ) + elif ( + isinstance(field_info, params.Param) + and getattr(field_info, "in_", None) is None + ): + field_info.in_ = params.ParamTypes.query + use_annotation_from_field_info = get_annotation_from_field_info( + use_annotation, + field_info, + param_name, + ) + if isinstance(field_info, params.Form): + ensure_multipart_is_installed() + if not field_info.alias and getattr(field_info, "convert_underscores", None): + alias = param_name.replace("_", "-") + else: + alias = field_info.alias or param_name + field_info.alias = alias + field = create_model_field( + name=param_name, + type_=use_annotation_from_field_info, + default=field_info.default, + alias=alias, + required=field_info.default in (RequiredParam, Undefined), + field_info=field_info, + ) + if is_path_param: + assert is_scalar_field(field=field), ( + "Path params must be of one of the supported types" + ) + elif isinstance(field_info, params.Query): + assert ( + is_scalar_field(field) + or is_scalar_sequence_field(field) + or ( + lenient_issubclass(field.type_, BaseModel) + # For Pydantic v1 + and getattr(field, "shape", 1) == 1 + ) + ) + + return ParamDetails(type_annotation=type_annotation, depends=depends, field=field) + + +def add_param_to_fields(*, field: ModelField, dependant: Dependant) -> None: + field_info = field.field_info + field_info_in = getattr(field_info, "in_", None) + if field_info_in == params.ParamTypes.path: + dependant.path_params.append(field) + elif field_info_in == params.ParamTypes.query: + dependant.query_params.append(field) + elif field_info_in == params.ParamTypes.header: + dependant.header_params.append(field) + else: + assert field_info_in == params.ParamTypes.cookie, ( + f"non-body parameters must be in path, query, header or cookie: {field.name}" + ) + dependant.cookie_params.append(field) + + +def is_coroutine_callable(call: Callable[..., Any]) -> bool: + if inspect.isroutine(call): + return inspect.iscoroutinefunction(call) + if inspect.isclass(call): + return False + dunder_call = getattr(call, "__call__", None) # noqa: B004 + return inspect.iscoroutinefunction(dunder_call) + + +def is_async_gen_callable(call: Callable[..., Any]) -> bool: + if inspect.isasyncgenfunction(call): + return True + dunder_call = getattr(call, "__call__", None) # noqa: B004 + return inspect.isasyncgenfunction(dunder_call) + + +def is_gen_callable(call: Callable[..., Any]) -> bool: + if inspect.isgeneratorfunction(call): + return True + dunder_call = getattr(call, "__call__", None) # noqa: B004 + return inspect.isgeneratorfunction(dunder_call) + + +async def solve_generator( + *, call: Callable[..., Any], stack: AsyncExitStack, sub_values: Dict[str, Any] +) -> Any: + if is_gen_callable(call): + cm = contextmanager_in_threadpool(contextmanager(call)(**sub_values)) + elif is_async_gen_callable(call): + cm = asynccontextmanager(call)(**sub_values) + return await stack.enter_async_context(cm) + + +@dataclass +class SolvedDependency: + values: Dict[str, Any] + errors: List[Any] + background_tasks: Optional[StarletteBackgroundTasks] + response: Response + dependency_cache: Dict[Tuple[Callable[..., Any], Tuple[str]], Any] + + +async def solve_dependencies( + *, + request: Union[Request, WebSocket], + dependant: Dependant, + body: Optional[Union[Dict[str, Any], FormData]] = None, + background_tasks: Optional[StarletteBackgroundTasks] = None, + response: Optional[Response] = None, + dependency_overrides_provider: Optional[Any] = None, + dependency_cache: Optional[Dict[Tuple[Callable[..., Any], Tuple[str]], Any]] = None, + async_exit_stack: AsyncExitStack, + embed_body_fields: bool, +) -> SolvedDependency: + values: Dict[str, Any] = {} + errors: List[Any] = [] + if response is None: + response = Response() + del response.headers["content-length"] + response.status_code = None # type: ignore + dependency_cache = dependency_cache or {} + sub_dependant: Dependant + for sub_dependant in dependant.dependencies: + sub_dependant.call = cast(Callable[..., Any], sub_dependant.call) + sub_dependant.cache_key = cast( + Tuple[Callable[..., Any], Tuple[str]], sub_dependant.cache_key + ) + call = sub_dependant.call + use_sub_dependant = sub_dependant + if ( + dependency_overrides_provider + and dependency_overrides_provider.dependency_overrides + ): + original_call = sub_dependant.call + call = getattr( + dependency_overrides_provider, "dependency_overrides", {} + ).get(original_call, original_call) + use_path: str = sub_dependant.path # type: ignore + use_sub_dependant = get_dependant( + path=use_path, + call=call, + name=sub_dependant.name, + security_scopes=sub_dependant.security_scopes, + ) + + solved_result = await solve_dependencies( + request=request, + dependant=use_sub_dependant, + body=body, + background_tasks=background_tasks, + response=response, + dependency_overrides_provider=dependency_overrides_provider, + dependency_cache=dependency_cache, + async_exit_stack=async_exit_stack, + embed_body_fields=embed_body_fields, + ) + background_tasks = solved_result.background_tasks + dependency_cache.update(solved_result.dependency_cache) + if solved_result.errors: + errors.extend(solved_result.errors) + continue + if sub_dependant.use_cache and sub_dependant.cache_key in dependency_cache: + solved = dependency_cache[sub_dependant.cache_key] + elif is_gen_callable(call) or is_async_gen_callable(call): + solved = await solve_generator( + call=call, stack=async_exit_stack, sub_values=solved_result.values + ) + elif is_coroutine_callable(call): + solved = await call(**solved_result.values) + else: + solved = await run_in_threadpool(call, **solved_result.values) + if sub_dependant.name is not None: + values[sub_dependant.name] = solved + if sub_dependant.cache_key not in dependency_cache: + dependency_cache[sub_dependant.cache_key] = solved + path_values, path_errors = request_params_to_args( + dependant.path_params, request.path_params + ) + query_values, query_errors = request_params_to_args( + dependant.query_params, request.query_params + ) + header_values, header_errors = request_params_to_args( + dependant.header_params, request.headers + ) + cookie_values, cookie_errors = request_params_to_args( + dependant.cookie_params, request.cookies + ) + values.update(path_values) + values.update(query_values) + values.update(header_values) + values.update(cookie_values) + errors += path_errors + query_errors + header_errors + cookie_errors + if dependant.body_params: + ( + body_values, + body_errors, + ) = await request_body_to_args( # body_params checked above + body_fields=dependant.body_params, + received_body=body, + embed_body_fields=embed_body_fields, + ) + values.update(body_values) + errors.extend(body_errors) + if dependant.http_connection_param_name: + values[dependant.http_connection_param_name] = request + if dependant.request_param_name and isinstance(request, Request): + values[dependant.request_param_name] = request + elif dependant.websocket_param_name and isinstance(request, WebSocket): + values[dependant.websocket_param_name] = request + if dependant.background_tasks_param_name: + if background_tasks is None: + background_tasks = BackgroundTasks() + values[dependant.background_tasks_param_name] = background_tasks + if dependant.response_param_name: + values[dependant.response_param_name] = response + if dependant.security_scopes_param_name: + values[dependant.security_scopes_param_name] = SecurityScopes( + scopes=dependant.security_scopes + ) + return SolvedDependency( + values=values, + errors=errors, + background_tasks=background_tasks, + response=response, + dependency_cache=dependency_cache, + ) + + +def _validate_value_with_model_field( + *, field: ModelField, value: Any, values: Dict[str, Any], loc: Tuple[str, ...] +) -> Tuple[Any, List[Any]]: + if value is None: + if field.required: + return None, [get_missing_field_error(loc=loc)] + else: + return deepcopy(field.default), [] + v_, errors_ = field.validate(value, values, loc=loc) + if isinstance(errors_, ErrorWrapper): + return None, [errors_] + elif isinstance(errors_, list): + new_errors = _regenerate_error_with_loc(errors=errors_, loc_prefix=()) + return None, new_errors + else: + return v_, [] + + +def _get_multidict_value( + field: ModelField, values: Mapping[str, Any], alias: Union[str, None] = None +) -> Any: + alias = alias or field.alias + if is_sequence_field(field) and isinstance(values, (ImmutableMultiDict, Headers)): + value = values.getlist(alias) + else: + value = values.get(alias, None) + if ( + value is None + or ( + isinstance(field.field_info, params.Form) + and isinstance(value, str) # For type checks + and value == "" + ) + or (is_sequence_field(field) and len(value) == 0) + ): + if field.required: + return + else: + return deepcopy(field.default) + return value + + +def request_params_to_args( + fields: Sequence[ModelField], + received_params: Union[Mapping[str, Any], QueryParams, Headers], +) -> Tuple[Dict[str, Any], List[Any]]: + values: Dict[str, Any] = {} + errors: List[Dict[str, Any]] = [] + + if not fields: + return values, errors + + first_field = fields[0] + fields_to_extract = fields + single_not_embedded_field = False + default_convert_underscores = True + if len(fields) == 1 and lenient_issubclass(first_field.type_, BaseModel): + fields_to_extract = get_cached_model_fields(first_field.type_) + single_not_embedded_field = True + # If headers are in a Pydantic model, the way to disable convert_underscores + # would be with Header(convert_underscores=False) at the Pydantic model level + default_convert_underscores = getattr( + first_field.field_info, "convert_underscores", True + ) + + params_to_process: Dict[str, Any] = {} + + processed_keys = set() + + for field in fields_to_extract: + alias = None + if isinstance(received_params, Headers): + # Handle fields extracted from a Pydantic Model for a header, each field + # doesn't have a FieldInfo of type Header with the default convert_underscores=True + convert_underscores = getattr( + field.field_info, "convert_underscores", default_convert_underscores + ) + if convert_underscores: + alias = ( + field.alias + if field.alias != field.name + else field.name.replace("_", "-") + ) + value = _get_multidict_value(field, received_params, alias=alias) + if value is not None: + params_to_process[field.name] = value + processed_keys.add(alias or field.alias) + processed_keys.add(field.name) + + for key, value in received_params.items(): + if key not in processed_keys: + params_to_process[key] = value + + if single_not_embedded_field: + field_info = first_field.field_info + assert isinstance(field_info, params.Param), ( + "Params must be subclasses of Param" + ) + loc: Tuple[str, ...] = (field_info.in_.value,) + v_, errors_ = _validate_value_with_model_field( + field=first_field, value=params_to_process, values=values, loc=loc + ) + return {first_field.name: v_}, errors_ + + for field in fields: + value = _get_multidict_value(field, received_params) + field_info = field.field_info + assert isinstance(field_info, params.Param), ( + "Params must be subclasses of Param" + ) + loc = (field_info.in_.value, field.alias) + v_, errors_ = _validate_value_with_model_field( + field=field, value=value, values=values, loc=loc + ) + if errors_: + errors.extend(errors_) + else: + values[field.name] = v_ + return values, errors + + +def _should_embed_body_fields(fields: List[ModelField]) -> bool: + if not fields: + return False + # More than one dependency could have the same field, it would show up as multiple + # fields but it's the same one, so count them by name + body_param_names_set = {field.name for field in fields} + # A top level field has to be a single field, not multiple + if len(body_param_names_set) > 1: + return True + first_field = fields[0] + # If it explicitly specifies it is embedded, it has to be embedded + if getattr(first_field.field_info, "embed", None): + return True + # If it's a Form (or File) field, it has to be a BaseModel to be top level + # otherwise it has to be embedded, so that the key value pair can be extracted + if isinstance(first_field.field_info, params.Form) and not lenient_issubclass( + first_field.type_, BaseModel + ): + return True + return False + + +async def _extract_form_body( + body_fields: List[ModelField], + received_body: FormData, +) -> Dict[str, Any]: + values = {} + first_field = body_fields[0] + first_field_info = first_field.field_info + + for field in body_fields: + value = _get_multidict_value(field, received_body) + if ( + isinstance(first_field_info, params.File) + and is_bytes_field(field) + and isinstance(value, UploadFile) + ): + value = await value.read() + elif ( + is_bytes_sequence_field(field) + and isinstance(first_field_info, params.File) + and value_is_sequence(value) + ): + # For types + assert isinstance(value, sequence_types) # type: ignore[arg-type] + results: List[Union[bytes, str]] = [] + + async def process_fn( + fn: Callable[[], Coroutine[Any, Any, Any]], + ) -> None: + result = await fn() + results.append(result) # noqa: B023 + + async with anyio.create_task_group() as tg: + for sub_value in value: + tg.start_soon(process_fn, sub_value.read) + value = serialize_sequence_value(field=field, value=results) + if value is not None: + values[field.alias] = value + for key, value in received_body.items(): + if key not in values: + values[key] = value + return values + + +async def request_body_to_args( + body_fields: List[ModelField], + received_body: Optional[Union[Dict[str, Any], FormData]], + embed_body_fields: bool, +) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]: + values: Dict[str, Any] = {} + errors: List[Dict[str, Any]] = [] + assert body_fields, "request_body_to_args() should be called with fields" + single_not_embedded_field = len(body_fields) == 1 and not embed_body_fields + first_field = body_fields[0] + body_to_process = received_body + + fields_to_extract: List[ModelField] = body_fields + + if single_not_embedded_field and lenient_issubclass(first_field.type_, BaseModel): + fields_to_extract = get_cached_model_fields(first_field.type_) + + if isinstance(received_body, FormData): + body_to_process = await _extract_form_body(fields_to_extract, received_body) + + if single_not_embedded_field: + loc: Tuple[str, ...] = ("body",) + v_, errors_ = _validate_value_with_model_field( + field=first_field, value=body_to_process, values=values, loc=loc + ) + return {first_field.name: v_}, errors_ + for field in body_fields: + loc = ("body", field.alias) + value: Optional[Any] = None + if body_to_process is not None: + try: + value = body_to_process.get(field.alias) + # If the received body is a list, not a dict + except AttributeError: + errors.append(get_missing_field_error(loc)) + continue + v_, errors_ = _validate_value_with_model_field( + field=field, value=value, values=values, loc=loc + ) + if errors_: + errors.extend(errors_) + else: + values[field.name] = v_ + return values, errors + + +def get_body_field( + *, flat_dependant: Dependant, name: str, embed_body_fields: bool +) -> Optional[ModelField]: + """ + Get a ModelField representing the request body for a path operation, combining + all body parameters into a single field if necessary. + + Used to check if it's form data (with `isinstance(body_field, params.Form)`) + or JSON and to generate the JSON Schema for a request body. + + This is **not** used to validate/parse the request body, that's done with each + individual body parameter. + """ + if not flat_dependant.body_params: + return None + first_param = flat_dependant.body_params[0] + if not embed_body_fields: + return first_param + model_name = "Body_" + name + BodyModel = create_body_model( + fields=flat_dependant.body_params, model_name=model_name + ) + required = any(True for f in flat_dependant.body_params if f.required) + BodyFieldInfo_kwargs: Dict[str, Any] = { + "annotation": BodyModel, + "alias": "body", + } + if not required: + BodyFieldInfo_kwargs["default"] = None + if any(isinstance(f.field_info, params.File) for f in flat_dependant.body_params): + BodyFieldInfo: Type[params.Body] = params.File + elif any(isinstance(f.field_info, params.Form) for f in flat_dependant.body_params): + BodyFieldInfo = params.Form + else: + BodyFieldInfo = params.Body + + body_param_media_types = [ + f.field_info.media_type + for f in flat_dependant.body_params + if isinstance(f.field_info, params.Body) + ] + if len(set(body_param_media_types)) == 1: + BodyFieldInfo_kwargs["media_type"] = body_param_media_types[0] + final_field = create_model_field( + name="body", + type_=BodyModel, + required=required, + alias="body", + field_info=BodyFieldInfo(**BodyFieldInfo_kwargs), + ) + return final_field diff --git a/lib/python3.12/site-packages/fastapi/middleware/__init__.py b/lib/python3.12/site-packages/fastapi/middleware/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..620296d5ad6ca2cc49eb5d0dc140bcbc3204e9b4 --- /dev/null +++ b/lib/python3.12/site-packages/fastapi/middleware/__init__.py @@ -0,0 +1 @@ +from starlette.middleware import Middleware as Middleware diff --git a/lib/python3.12/site-packages/fastapi/middleware/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66dbbb22b3bfc76627bdf986b64cc3489ca78b72 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/middleware/__pycache__/cors.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/cors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a12709efb8b1b0fc2caeb9c51f8d690969f33bc8 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/cors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/middleware/__pycache__/gzip.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/gzip.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35258ff29d2c1e90cbe3d18b9d5cef81aeb6efd1 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/gzip.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/middleware/__pycache__/httpsredirect.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/httpsredirect.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ae440e39fd676a3e454ec6b737fd57cde1e596b7 Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/httpsredirect.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/middleware/__pycache__/trustedhost.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/trustedhost.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3dba0ed5a87b4e2be1252e1a6da0f65191ba38a Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/trustedhost.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/middleware/__pycache__/wsgi.cpython-312.pyc b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/wsgi.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9a92226ada313960e3d5a7f814da0195efc2b17f Binary files /dev/null and b/lib/python3.12/site-packages/fastapi/middleware/__pycache__/wsgi.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/fastapi/middleware/cors.py b/lib/python3.12/site-packages/fastapi/middleware/cors.py new file mode 100644 index 0000000000000000000000000000000000000000..8dfaad0dbb3ff5300cccb2023748cd30f54bc920 --- /dev/null +++ b/lib/python3.12/site-packages/fastapi/middleware/cors.py @@ -0,0 +1 @@ +from starlette.middleware.cors import CORSMiddleware as CORSMiddleware # noqa diff --git a/lib/python3.12/site-packages/fastapi/middleware/gzip.py b/lib/python3.12/site-packages/fastapi/middleware/gzip.py new file mode 100644 index 0000000000000000000000000000000000000000..bbeb2cc7861a735d6cd5c0e29aeb6dbf8457023a --- /dev/null +++ b/lib/python3.12/site-packages/fastapi/middleware/gzip.py @@ -0,0 +1 @@ +from starlette.middleware.gzip import GZipMiddleware as GZipMiddleware # noqa diff --git a/lib/python3.12/site-packages/fastapi/middleware/httpsredirect.py b/lib/python3.12/site-packages/fastapi/middleware/httpsredirect.py new file mode 100644 index 0000000000000000000000000000000000000000..b7a3d8e078574e87dc6e345d621f5a596c3bdc1e --- /dev/null +++ b/lib/python3.12/site-packages/fastapi/middleware/httpsredirect.py @@ -0,0 +1,3 @@ +from starlette.middleware.httpsredirect import ( # noqa + HTTPSRedirectMiddleware as HTTPSRedirectMiddleware, +) diff --git a/lib/python3.12/site-packages/fastapi/middleware/trustedhost.py b/lib/python3.12/site-packages/fastapi/middleware/trustedhost.py new file mode 100644 index 0000000000000000000000000000000000000000..08d7e035315677856fd2cd0be2044689b57619bf --- /dev/null +++ b/lib/python3.12/site-packages/fastapi/middleware/trustedhost.py @@ -0,0 +1,3 @@ +from starlette.middleware.trustedhost import ( # noqa + TrustedHostMiddleware as TrustedHostMiddleware, +) diff --git a/lib/python3.12/site-packages/fastapi/middleware/wsgi.py b/lib/python3.12/site-packages/fastapi/middleware/wsgi.py new file mode 100644 index 0000000000000000000000000000000000000000..c4c6a797d2675e1c13b028be977c64a822fb649b --- /dev/null +++ b/lib/python3.12/site-packages/fastapi/middleware/wsgi.py @@ -0,0 +1 @@ +from starlette.middleware.wsgi import WSGIMiddleware as WSGIMiddleware # noqa diff --git a/lib/python3.12/site-packages/git/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85f31fd5e4010d65ce887044e11fdeb3c3719882 Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/cmd.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/cmd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2f7a38a4f71763bea845abba89eba70d0ee9a2cb Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/cmd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/compat.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37ddc724631e1062bfd5c2c82dd0f77ebd11197a Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9c98307f297b3bd83b34c9386d5d76aef79a2da Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/db.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/db.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d74cee9388df1be88dfdd08b38fafd131f12b077 Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/db.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/diff.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/diff.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce2171b5bfb6747a82e381fa1bb4fd086736d460 Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/diff.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/exc.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/exc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..198f7b054f0c0e1700c63173d76021cff3d14b3e Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/exc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/remote.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/remote.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..75a3898af4d51d15eb115126e7ced571793abcef Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/remote.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..198014a444dab46a30e9dd9e87d859bdd593a7dd Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/git/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a9ca8e55e445e14df3648e97c6fdd79b811433d Binary files /dev/null and b/lib/python3.12/site-packages/git/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/index/__init__.py b/lib/python3.12/site-packages/git/index/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ba48110fddcdf549b6e61f115d28b137a57cc219 --- /dev/null +++ b/lib/python3.12/site-packages/git/index/__init__.py @@ -0,0 +1,16 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Initialize the index package.""" + +__all__ = [ + "BaseIndexEntry", + "BlobFilter", + "CheckoutError", + "IndexEntry", + "IndexFile", + "StageType", +] + +from .base import CheckoutError, IndexFile +from .typ import BaseIndexEntry, BlobFilter, IndexEntry, StageType diff --git a/lib/python3.12/site-packages/git/index/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/git/index/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d559f22266db66bacdaafb646bcd627fb790abf2 Binary files /dev/null and b/lib/python3.12/site-packages/git/index/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/index/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/git/index/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4cdfe420026a428e87e491a52a897ead0445b10e Binary files /dev/null and b/lib/python3.12/site-packages/git/index/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/index/__pycache__/fun.cpython-312.pyc b/lib/python3.12/site-packages/git/index/__pycache__/fun.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..15f4e448c5d9b32b31642d4a5ab010516071ac38 Binary files /dev/null and b/lib/python3.12/site-packages/git/index/__pycache__/fun.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/index/__pycache__/typ.cpython-312.pyc b/lib/python3.12/site-packages/git/index/__pycache__/typ.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2a68e9aa21189bb47fe48673a519ec2b8a171e3 Binary files /dev/null and b/lib/python3.12/site-packages/git/index/__pycache__/typ.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/index/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/git/index/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6de7938efca917bb7434577d4b3504edea5e072f Binary files /dev/null and b/lib/python3.12/site-packages/git/index/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/index/base.py b/lib/python3.12/site-packages/git/index/base.py new file mode 100644 index 0000000000000000000000000000000000000000..93de7933c3109c7d8989c67688f8705fe2e46462 --- /dev/null +++ b/lib/python3.12/site-packages/git/index/base.py @@ -0,0 +1,1532 @@ +# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors +# +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Module containing :class:`IndexFile`, an Index implementation facilitating all kinds +of index manipulations such as querying and merging.""" + +__all__ = ["IndexFile", "CheckoutError", "StageType"] + +import contextlib +import datetime +import glob +from io import BytesIO +import os +import os.path as osp +from stat import S_ISLNK +import subprocess +import sys +import tempfile + +from gitdb.base import IStream +from gitdb.db import MemoryDB + +from git.compat import defenc, force_bytes +import git.diff as git_diff +from git.exc import CheckoutError, GitCommandError, GitError, InvalidGitRepositoryError +from git.objects import Blob, Commit, Object, Submodule, Tree +from git.objects.util import Serializable +from git.util import ( + Actor, + LazyMixin, + LockedFD, + join_path_native, + file_contents_ro, + to_native_path_linux, + unbare_repo, + to_bin_sha, +) + +from .fun import ( + S_IFGITLINK, + aggressive_tree_merge, + entry_key, + read_cache, + run_commit_hook, + stat_mode_to_index_mode, + write_cache, + write_tree_from_cache, +) +from .typ import BaseIndexEntry, IndexEntry, StageType +from .util import TemporaryFileSwap, post_clear_cache, default_index, git_working_dir + +# typing ----------------------------------------------------------------------------- + +from typing import ( + Any, + BinaryIO, + Callable, + Dict, + Generator, + IO, + Iterable, + Iterator, + List, + NoReturn, + Sequence, + TYPE_CHECKING, + Tuple, + Union, +) + +from git.types import Literal, PathLike + +if TYPE_CHECKING: + from subprocess import Popen + + from git.refs.reference import Reference + from git.repo import Repo + + +Treeish = Union[Tree, Commit, str, bytes] + +# ------------------------------------------------------------------------------------ + + +@contextlib.contextmanager +def _named_temporary_file_for_subprocess(directory: PathLike) -> Generator[str, None, None]: + """Create a named temporary file git subprocesses can open, deleting it afterward. + + :param directory: + The directory in which the file is created. + + :return: + A context manager object that creates the file and provides its name on entry, + and deletes it on exit. + """ + if sys.platform == "win32": + fd, name = tempfile.mkstemp(dir=directory) + os.close(fd) + try: + yield name + finally: + os.remove(name) + else: + with tempfile.NamedTemporaryFile(dir=directory) as ctx: + yield ctx.name + + +class IndexFile(LazyMixin, git_diff.Diffable, Serializable): + """An Index that can be manipulated using a native implementation in order to save + git command function calls wherever possible. + + This provides custom merging facilities allowing to merge without actually changing + your index or your working tree. This way you can perform your own test merges based + on the index only without having to deal with the working copy. This is useful in + case of partial working trees. + + Entries: + + The index contains an entries dict whose keys are tuples of type + :class:`~git.index.typ.IndexEntry` to facilitate access. + + You may read the entries dict or manipulate it using IndexEntry instance, i.e.:: + + index.entries[index.entry_key(index_entry_instance)] = index_entry_instance + + Make sure you use :meth:`index.write() ` once you are done manipulating the + index directly before operating on it using the git command. + """ + + __slots__ = ("repo", "version", "entries", "_extension_data", "_file_path") + + _VERSION = 2 + """The latest version we support.""" + + S_IFGITLINK = S_IFGITLINK + """Flags for a submodule.""" + + def __init__(self, repo: "Repo", file_path: Union[PathLike, None] = None) -> None: + """Initialize this Index instance, optionally from the given `file_path`. + + If no `file_path` is given, we will be created from the current index file. + + If a stream is not given, the stream will be initialized from the current + repository's index on demand. + """ + self.repo = repo + self.version = self._VERSION + self._extension_data = b"" + self._file_path: PathLike = file_path or self._index_path() + + def _set_cache_(self, attr: str) -> None: + if attr == "entries": + try: + fd = os.open(self._file_path, os.O_RDONLY) + except OSError: + # In new repositories, there may be no index, which means we are empty. + self.entries: Dict[Tuple[PathLike, StageType], IndexEntry] = {} + return + # END exception handling + + try: + stream = file_contents_ro(fd, stream=True, allow_mmap=True) + finally: + os.close(fd) + + self._deserialize(stream) + else: + super()._set_cache_(attr) + + def _index_path(self) -> PathLike: + if self.repo.git_dir: + return join_path_native(self.repo.git_dir, "index") + else: + raise GitCommandError("No git directory given to join index path") + + @property + def path(self) -> PathLike: + """:return: Path to the index file we are representing""" + return self._file_path + + def _delete_entries_cache(self) -> None: + """Safely clear the entries cache so it can be recreated.""" + try: + del self.entries + except AttributeError: + # It failed in Python 2.6.5 with AttributeError. + # FIXME: Look into whether we can just remove this except clause now. + pass + # END exception handling + + # { Serializable Interface + + def _deserialize(self, stream: IO) -> "IndexFile": + """Initialize this instance with index values read from the given stream.""" + self.version, self.entries, self._extension_data, _conten_sha = read_cache(stream) + return self + + def _entries_sorted(self) -> List[IndexEntry]: + """:return: List of entries, in a sorted fashion, first by path, then by stage""" + return sorted(self.entries.values(), key=lambda e: (e.path, e.stage)) + + def _serialize(self, stream: IO, ignore_extension_data: bool = False) -> "IndexFile": + entries = self._entries_sorted() + extension_data = self._extension_data # type: Union[None, bytes] + if ignore_extension_data: + extension_data = None + write_cache(entries, stream, extension_data) + return self + + # } END serializable interface + + def write( + self, + file_path: Union[None, PathLike] = None, + ignore_extension_data: bool = False, + ) -> None: + """Write the current state to our file path or to the given one. + + :param file_path: + If ``None``, we will write to our stored file path from which we have been + initialized. Otherwise we write to the given file path. Please note that + this will change the `file_path` of this index to the one you gave. + + :param ignore_extension_data: + If ``True``, the TREE type extension data read in the index will not be + written to disk. NOTE that no extension data is actually written. Use this + if you have altered the index and would like to use + :manpage:`git-write-tree(1)` afterwards to create a tree representing your + written changes. If this data is present in the written index, + :manpage:`git-write-tree(1)` will instead write the stored/cached tree. + Alternatively, use :meth:`write_tree` to handle this case automatically. + """ + # Make sure we have our entries read before getting a write lock. + # Otherwise it would be done when streaming. + # This can happen if one doesn't change the index, but writes it right away. + self.entries # noqa: B018 + lfd = LockedFD(file_path or self._file_path) + stream = lfd.open(write=True, stream=True) + + try: + self._serialize(stream, ignore_extension_data) + except BaseException: + lfd.rollback() + raise + + lfd.commit() + + # Make sure we represent what we have written. + if file_path is not None: + self._file_path = file_path + + @post_clear_cache + @default_index + def merge_tree(self, rhs: Treeish, base: Union[None, Treeish] = None) -> "IndexFile": + """Merge the given `rhs` treeish into the current index, possibly taking + a common base treeish into account. + + As opposed to the :func:`from_tree` method, this allows you to use an already + existing tree as the left side of the merge. + + :param rhs: + Treeish reference pointing to the 'other' side of the merge. + + :param base: + Optional treeish reference pointing to the common base of `rhs` and this + index which equals lhs. + + :return: + self (containing the merge and possibly unmerged entries in case of + conflicts) + + :raise git.exc.GitCommandError: + If there is a merge conflict. The error will be raised at the first + conflicting path. If you want to have proper merge resolution to be done by + yourself, you have to commit the changed index (or make a valid tree from + it) and retry with a three-way :meth:`index.from_tree ` call. + """ + # -i : ignore working tree status + # --aggressive : handle more merge cases + # -m : do an actual merge + args: List[Union[Treeish, str]] = ["--aggressive", "-i", "-m"] + if base is not None: + args.append(base) + args.append(rhs) + + self.repo.git.read_tree(args) + return self + + @classmethod + def new(cls, repo: "Repo", *tree_sha: Union[str, Tree]) -> "IndexFile": + """Merge the given treeish revisions into a new index which is returned. + + This method behaves like ``git-read-tree --aggressive`` when doing the merge. + + :param repo: + The repository treeish are located in. + + :param tree_sha: + 20 byte or 40 byte tree sha or tree objects. + + :return: + New :class:`IndexFile` instance. Its path will be undefined. + If you intend to write such a merged Index, supply an alternate + ``file_path`` to its :meth:`write` method. + """ + tree_sha_bytes: List[bytes] = [to_bin_sha(str(t)) for t in tree_sha] + base_entries = aggressive_tree_merge(repo.odb, tree_sha_bytes) + + inst = cls(repo) + # Convert to entries dict. + entries: Dict[Tuple[PathLike, int], IndexEntry] = dict( + zip( + ((e.path, e.stage) for e in base_entries), + (IndexEntry.from_base(e) for e in base_entries), + ) + ) + + inst.entries = entries + return inst + + @classmethod + def from_tree(cls, repo: "Repo", *treeish: Treeish, **kwargs: Any) -> "IndexFile": + R"""Merge the given treeish revisions into a new index which is returned. + The original index will remain unaltered. + + :param repo: + The repository treeish are located in. + + :param treeish: + One, two or three :class:`~git.objects.tree.Tree` objects, + :class:`~git.objects.commit.Commit`\s or 40 byte hexshas. + + The result changes according to the amount of trees: + + 1. If 1 Tree is given, it will just be read into a new index. + 2. If 2 Trees are given, they will be merged into a new index using a two + way merge algorithm. Tree 1 is the 'current' tree, tree 2 is the 'other' + one. It behaves like a fast-forward. + 3. If 3 Trees are given, a 3-way merge will be performed with the first tree + being the common ancestor of tree 2 and tree 3. Tree 2 is the 'current' + tree, tree 3 is the 'other' one. + + :param kwargs: + Additional arguments passed to :manpage:`git-read-tree(1)`. + + :return: + New :class:`IndexFile` instance. It will point to a temporary index location + which does not exist anymore. If you intend to write such a merged Index, + supply an alternate ``file_path`` to its :meth:`write` method. + + :note: + In the three-way merge case, ``--aggressive`` will be specified to + automatically resolve more cases in a commonly correct manner. Specify + ``trivial=True`` as a keyword argument to override that. + + As the underlying :manpage:`git-read-tree(1)` command takes into account the + current index, it will be temporarily moved out of the way to prevent any + unexpected interference. + """ + if len(treeish) == 0 or len(treeish) > 3: + raise ValueError("Please specify between 1 and 3 treeish, got %i" % len(treeish)) + + arg_list: List[Union[Treeish, str]] = [] + # Ignore that the working tree and index possibly are out of date. + if len(treeish) > 1: + # Drop unmerged entries when reading our index and merging. + arg_list.append("--reset") + # Handle non-trivial cases the way a real merge does. + arg_list.append("--aggressive") + # END merge handling + + # Create the temporary file in the .git directory to be sure renaming + # works - /tmp/ directories could be on another device. + with _named_temporary_file_for_subprocess(repo.git_dir) as tmp_index: + arg_list.append("--index-output=%s" % tmp_index) + arg_list.extend(treeish) + + # Move the current index out of the way - otherwise the merge may fail as it + # considers existing entries. Moving it essentially clears the index. + # Unfortunately there is no 'soft' way to do it. + # The TemporaryFileSwap ensures the original file gets put back. + with TemporaryFileSwap(join_path_native(repo.git_dir, "index")): + repo.git.read_tree(*arg_list, **kwargs) + index = cls(repo, tmp_index) + index.entries # noqa: B018 # Force it to read the file as we will delete the temp-file. + return index + # END index merge handling + + # UTILITIES + + @unbare_repo + def _iter_expand_paths(self: "IndexFile", paths: Sequence[PathLike]) -> Iterator[PathLike]: + """Expand the directories in list of paths to the corresponding paths + accordingly. + + :note: + git will add items multiple times even if a glob overlapped with manually + specified paths or if paths where specified multiple times - we respect that + and do not prune. + """ + + def raise_exc(e: Exception) -> NoReturn: + raise e + + r = str(self.repo.working_tree_dir) + rs = r + os.sep + for path in paths: + abs_path = os.fspath(path) + if not osp.isabs(abs_path): + abs_path = osp.join(r, path) + # END make absolute path + + try: + st = os.lstat(abs_path) # Handles non-symlinks as well. + except OSError: + # The lstat call may fail as the path may contain globs as well. + pass + else: + if S_ISLNK(st.st_mode): + yield abs_path.replace(rs, "") + continue + # END check symlink + + # If the path is not already pointing to an existing file, resolve globs if possible. + if not os.path.exists(abs_path) and ("?" in abs_path or "*" in abs_path or "[" in abs_path): + resolved_paths = glob.glob(abs_path) + # not abs_path in resolved_paths: + # A glob() resolving to the same path we are feeding it with is a + # glob() that failed to resolve. If we continued calling ourselves + # we'd endlessly recurse. If the condition below evaluates to true + # then we are likely dealing with a file whose name contains wildcard + # characters. + if abs_path not in resolved_paths: + for f in self._iter_expand_paths(glob.glob(abs_path)): + yield str(f).replace(rs, "") + continue + # END glob handling + try: + for root, _dirs, files in os.walk(abs_path, onerror=raise_exc): + for rela_file in files: + # Add relative paths only. + yield osp.join(root.replace(rs, ""), rela_file) + # END for each file in subdir + # END for each subdirectory + except OSError: + # It was a file or something that could not be iterated. + yield abs_path.replace(rs, "") + # END path exception handling + # END for each path + + def _write_path_to_stdin( + self, + proc: "Popen", + filepath: PathLike, + item: PathLike, + fmakeexc: Callable[..., GitError], + fprogress: Callable[[PathLike, bool, PathLike], None], + read_from_stdout: bool = True, + ) -> Union[None, str]: + """Write path to ``proc.stdin`` and make sure it processes the item, including + progress. + + :return: + stdout string + + :param read_from_stdout: + If ``True``, ``proc.stdout`` will be read after the item was sent to stdin. + In that case, it will return ``None``. + + :note: + There is a bug in :manpage:`git-update-index(1)` that prevents it from + sending reports just in time. This is why we have a version that tries to + read stdout and one which doesn't. In fact, the stdout is not important as + the piped-in files are processed anyway and just in time. + + :note: + Newlines are essential here, git's behaviour is somewhat inconsistent on + this depending on the version, hence we try our best to deal with newlines + carefully. Usually the last newline will not be sent, instead we will close + stdin to break the pipe. + """ + fprogress(filepath, False, item) + rval: Union[None, str] = None + + if proc.stdin is not None: + try: + proc.stdin.write(("%s\n" % filepath).encode(defenc)) + except IOError as e: + # Pipe broke, usually because some error happened. + raise fmakeexc() from e + # END write exception handling + proc.stdin.flush() + + if read_from_stdout and proc.stdout is not None: + rval = proc.stdout.readline().strip() + fprogress(filepath, True, item) + return rval + + def iter_blobs( + self, predicate: Callable[[Tuple[StageType, Blob]], bool] = lambda t: True + ) -> Iterator[Tuple[StageType, Blob]]: + """ + :return: + Iterator yielding tuples of :class:`~git.objects.blob.Blob` objects and + stages, tuple(stage, Blob). + + :param predicate: + Function(t) returning ``True`` if tuple(stage, Blob) should be yielded by + the iterator. A default filter, the :class:`~git.index.typ.BlobFilter`, allows you + to yield blobs only if they match a given list of paths. + """ + for entry in self.entries.values(): + blob = entry.to_blob(self.repo) + blob.size = entry.size + output = (entry.stage, blob) + if predicate(output): + yield output + # END for each entry + + def unmerged_blobs(self) -> Dict[PathLike, List[Tuple[StageType, Blob]]]: + """ + :return: + Dict(path : list(tuple(stage, Blob, ...))), being a dictionary associating a + path in the index with a list containing sorted stage/blob pairs. + + :note: + Blobs that have been removed in one side simply do not exist in the given + stage. That is, a file removed on the 'other' branch whose entries are at + stage 3 will not have a stage 3 entry. + """ + + def is_unmerged_blob(t: Tuple[StageType, Blob]) -> bool: + return t[0] != 0 + + path_map: Dict[PathLike, List[Tuple[StageType, Blob]]] = {} + for stage, blob in self.iter_blobs(is_unmerged_blob): + path_map.setdefault(blob.path, []).append((stage, blob)) + # END for each unmerged blob + for line in path_map.values(): + line.sort() + + return path_map + + @classmethod + def entry_key(cls, *entry: Union[BaseIndexEntry, PathLike, StageType]) -> Tuple[PathLike, StageType]: + return entry_key(*entry) + + def resolve_blobs(self, iter_blobs: Iterator[Blob]) -> "IndexFile": + """Resolve the blobs given in blob iterator. + + This will effectively remove the index entries of the respective path at all + non-null stages and add the given blob as new stage null blob. + + For each path there may only be one blob, otherwise a :exc:`ValueError` will be + raised claiming the path is already at stage 0. + + :raise ValueError: + If one of the blobs already existed at stage 0. + + :return: + self + + :note: + You will have to write the index manually once you are done, i.e. + ``index.resolve_blobs(blobs).write()``. + """ + for blob in iter_blobs: + stage_null_key = (blob.path, 0) + if stage_null_key in self.entries: + raise ValueError("Path %r already exists at stage 0" % str(blob.path)) + # END assert blob is not stage 0 already + + # Delete all possible stages. + for stage in (1, 2, 3): + try: + del self.entries[(blob.path, stage)] + except KeyError: + pass + # END ignore key errors + # END for each possible stage + + self.entries[stage_null_key] = IndexEntry.from_blob(blob) + # END for each blob + + return self + + def update(self) -> "IndexFile": + """Reread the contents of our index file, discarding all cached information + we might have. + + :note: + This is a possibly dangerous operations as it will discard your changes to + :attr:`index.entries `. + + :return: + self + """ + self._delete_entries_cache() + # Allows to lazily reread on demand. + return self + + def write_tree(self) -> Tree: + """Write this index to a corresponding :class:`~git.objects.tree.Tree` object + into the repository's object database and return it. + + :return: + :class:`~git.objects.tree.Tree` object representing this index. + + :note: + The tree will be written even if one or more objects the tree refers to does + not yet exist in the object database. This could happen if you added entries + to the index directly. + + :raise ValueError: + If there are no entries in the cache. + + :raise git.exc.UnmergedEntriesError: + """ + # We obtain no lock as we just flush our contents to disk as tree. + # If we are a new index, the entries access will load our data accordingly. + mdb = MemoryDB() + entries = self._entries_sorted() + binsha, tree_items = write_tree_from_cache(entries, mdb, slice(0, len(entries))) + + # Copy changed trees only. + mdb.stream_copy(mdb.sha_iter(), self.repo.odb) + + # Note: Additional deserialization could be saved if write_tree_from_cache would + # return sorted tree entries. + root_tree = Tree(self.repo, binsha, path="") + root_tree._cache = tree_items + return root_tree + + def _process_diff_args( + self, + args: List[Union[PathLike, "git_diff.Diffable"]], + ) -> List[Union[PathLike, "git_diff.Diffable"]]: + try: + args.pop(args.index(self)) + except IndexError: + pass + # END remove self + return args + + def _to_relative_path(self, path: PathLike) -> PathLike: + """ + :return: + Version of path relative to our git directory or raise :exc:`ValueError` if + it is not within our git directory. + + :raise ValueError: + """ + if not osp.isabs(path): + return path + if self.repo.bare: + raise InvalidGitRepositoryError("require non-bare repository") + if not osp.normpath(path).startswith(str(self.repo.working_tree_dir)): + raise ValueError("Absolute path %r is not in git repository at %r" % (path, self.repo.working_tree_dir)) + result = os.path.relpath(path, self.repo.working_tree_dir) + if os.fspath(path).endswith(os.sep) and not result.endswith(os.sep): + result += os.sep + return result + + def _preprocess_add_items( + self, items: Union[PathLike, Sequence[Union[PathLike, Blob, BaseIndexEntry, "Submodule"]]] + ) -> Tuple[List[PathLike], List[BaseIndexEntry]]: + """Split the items into two lists of path strings and BaseEntries.""" + paths = [] + entries = [] + # if it is a string put in list + if isinstance(items, (str, os.PathLike)): + items = [items] + + for item in items: + if isinstance(item, (str, os.PathLike)): + paths.append(self._to_relative_path(item)) + elif isinstance(item, (Blob, Submodule)): + entries.append(BaseIndexEntry.from_blob(item)) + elif isinstance(item, BaseIndexEntry): + entries.append(item) + else: + raise TypeError("Invalid Type: %r" % item) + # END for each item + return paths, entries + + def _store_path(self, filepath: PathLike, fprogress: Callable) -> BaseIndexEntry: + """Store file at filepath in the database and return the base index entry. + + :note: + This needs the :func:`~git.index.util.git_working_dir` decorator active! + This must be ensured in the calling code. + """ + st = os.lstat(filepath) # Handles non-symlinks as well. + + if S_ISLNK(st.st_mode): + # In PY3, readlink is a string, but we need bytes. + # In PY2, it was just OS encoded bytes, we assumed UTF-8. + def open_stream() -> BinaryIO: + return BytesIO(force_bytes(os.readlink(filepath), encoding=defenc)) + else: + + def open_stream() -> BinaryIO: + return open(filepath, "rb") + + with open_stream() as stream: + fprogress(filepath, False, filepath) + istream = self.repo.odb.store(IStream(Blob.type, st.st_size, stream)) + fprogress(filepath, True, filepath) + return BaseIndexEntry( + ( + stat_mode_to_index_mode(st.st_mode), + istream.binsha, + 0, + to_native_path_linux(filepath), + ) + ) + + @unbare_repo + @git_working_dir + def _entries_for_paths( + self, + paths: List[str], + path_rewriter: Union[Callable, None], + fprogress: Callable, + entries: List[BaseIndexEntry], + ) -> List[BaseIndexEntry]: + entries_added: List[BaseIndexEntry] = [] + if path_rewriter: + for path in paths: + if osp.isabs(path): + abspath = path + gitrelative_path = path[len(str(self.repo.working_tree_dir)) + 1 :] + else: + gitrelative_path = path + if self.repo.working_tree_dir: + abspath = osp.join(self.repo.working_tree_dir, gitrelative_path) + # END obtain relative and absolute paths + + blob = Blob( + self.repo, + Blob.NULL_BIN_SHA, + stat_mode_to_index_mode(os.stat(abspath).st_mode), + to_native_path_linux(gitrelative_path), + ) + # TODO: variable undefined + entries.append(BaseIndexEntry.from_blob(blob)) + # END for each path + del paths[:] + # END rewrite paths + + # HANDLE PATHS + assert len(entries_added) == 0 + for filepath in self._iter_expand_paths(paths): + entries_added.append(self._store_path(filepath, fprogress)) + # END for each filepath + # END path handling + return entries_added + + def add( + self, + items: Union[PathLike, Sequence[Union[PathLike, Blob, BaseIndexEntry, "Submodule"]]], + force: bool = True, + fprogress: Callable = lambda *args: None, + path_rewriter: Union[Callable[..., PathLike], None] = None, + write: bool = True, + write_extension_data: bool = False, + ) -> List[BaseIndexEntry]: + R"""Add files from the working tree, specific blobs, or + :class:`~git.index.typ.BaseIndexEntry`\s to the index. + + :param items: + Multiple types of items are supported, types can be mixed within one call. + Different types imply a different handling. File paths may generally be + relative or absolute. + + - path string + + Strings denote a relative or absolute path into the repository pointing + to an existing file, e.g., ``CHANGES``, ``lib/myfile.ext``, + ``/home/gitrepo/lib/myfile.ext``. + + Absolute paths must start with working tree directory of this index's + repository to be considered valid. For example, if it was initialized + with a non-normalized path, like ``/root/repo/../repo``, absolute paths + to be added must start with ``/root/repo/../repo``. + + Paths provided like this must exist. When added, they will be written + into the object database. + + PathStrings may contain globs, such as ``lib/__init__*``. Or they can be + directories like ``lib``, which will add all the files within the + directory and subdirectories. + + This equals a straight :manpage:`git-add(1)`. + + They are added at stage 0. + + - :class:`~git.objects.blob.Blob` or + :class:`~git.objects.submodule.base.Submodule` object + + Blobs are added as they are assuming a valid mode is set. + + The file they refer to may or may not exist in the file system, but must + be a path relative to our repository. + + If their sha is null (40*0), their path must exist in the file system + relative to the git repository as an object will be created from the + data at the path. + + The handling now very much equals the way string paths are processed, + except that the mode you have set will be kept. This allows you to + create symlinks by settings the mode respectively and writing the target + of the symlink directly into the file. This equals a default Linux + symlink which is not dereferenced automatically, except that it can be + created on filesystems not supporting it as well. + + Please note that globs or directories are not allowed in + :class:`~git.objects.blob.Blob` objects. + + They are added at stage 0. + + - :class:`~git.index.typ.BaseIndexEntry` or type + + Handling equals the one of :class:`~git.objects.blob.Blob` objects, but + the stage may be explicitly set. Please note that Index Entries require + binary sha's. + + :param force: + **CURRENTLY INEFFECTIVE** + If ``True``, otherwise ignored or excluded files will be added anyway. As + opposed to the :manpage:`git-add(1)` command, we enable this flag by default + as the API user usually wants the item to be added even though they might be + excluded. + + :param fprogress: + Function with signature ``f(path, done=False, item=item)`` called for each + path to be added, one time once it is about to be added where ``done=False`` + and once after it was added where ``done=True``. + + ``item`` is set to the actual item we handle, either a path or a + :class:`~git.index.typ.BaseIndexEntry`. + + Please note that the processed path is not guaranteed to be present in the + index already as the index is currently being processed. + + :param path_rewriter: + Function, with signature ``(string) func(BaseIndexEntry)``, returning a path + for each passed entry which is the path to be actually recorded for the + object created from :attr:`entry.path `. + This allows you to write an index which is not identical to the layout of + the actual files on your hard-disk. If not ``None`` and `items` contain + plain paths, these paths will be converted to Entries beforehand and passed + to the path_rewriter. Please note that ``entry.path`` is relative to the git + repository. + + :param write: + If ``True``, the index will be written once it was altered. Otherwise the + changes only exist in memory and are not available to git commands. + + :param write_extension_data: + If ``True``, extension data will be written back to the index. This can lead + to issues in case it is containing the 'TREE' extension, which will cause + the :manpage:`git-commit(1)` command to write an old tree, instead of a new + one representing the now changed index. + + This doesn't matter if you use :meth:`IndexFile.commit`, which ignores the + 'TREE' extension altogether. You should set it to ``True`` if you intend to + use :meth:`IndexFile.commit` exclusively while maintaining support for + third-party extensions. Besides that, you can usually safely ignore the + built-in extensions when using GitPython on repositories that are not + handled manually at all. + + All current built-in extensions are listed here: + https://git-scm.com/docs/index-format + + :return: + List of :class:`~git.index.typ.BaseIndexEntry`\s representing the entries + just actually added. + + :raise OSError: + If a supplied path did not exist. Please note that + :class:`~git.index.typ.BaseIndexEntry` objects that do not have a null sha + will be added even if their paths do not exist. + """ + # Sort the entries into strings and Entries. + # Blobs are converted to entries automatically. + # Paths can be git-added. For everything else we use git-update-index. + paths, entries = self._preprocess_add_items(items) + entries_added: List[BaseIndexEntry] = [] + # This code needs a working tree, so we try not to run it unless required. + # That way, we are OK on a bare repository as well. + # If there are no paths, the rewriter has nothing to do either. + if paths: + entries_added.extend(self._entries_for_paths(paths, path_rewriter, fprogress, entries)) + + # HANDLE ENTRIES + if entries: + null_mode_entries = [e for e in entries if e.mode == 0] + if null_mode_entries: + raise ValueError( + "At least one Entry has a null-mode - please use index.remove to remove files for clarity" + ) + # END null mode should be remove + + # HANDLE ENTRY OBJECT CREATION + # Create objects if required, otherwise go with the existing shas. + null_entries_indices = [i for i, e in enumerate(entries) if e.binsha == Object.NULL_BIN_SHA] + if null_entries_indices: + + @git_working_dir + def handle_null_entries(self: "IndexFile") -> None: + for ei in null_entries_indices: + null_entry = entries[ei] + new_entry = self._store_path(null_entry.path, fprogress) + + # Update null entry. + entries[ei] = BaseIndexEntry( + ( + null_entry.mode, + new_entry.binsha, + null_entry.stage, + null_entry.path, + ) + ) + # END for each entry index + + # END closure + + handle_null_entries(self) + # END null_entry handling + + # REWRITE PATHS + # If we have to rewrite the entries, do so now, after we have generated all + # object sha's. + if path_rewriter: + for i, e in enumerate(entries): + entries[i] = BaseIndexEntry((e.mode, e.binsha, e.stage, path_rewriter(e))) + # END for each entry + # END handle path rewriting + + # Just go through the remaining entries and provide progress info. + for i, entry in enumerate(entries): + progress_sent = i in null_entries_indices + if not progress_sent: + fprogress(entry.path, False, entry) + fprogress(entry.path, True, entry) + # END handle progress + # END for each entry + entries_added.extend(entries) + # END if there are base entries + + # FINALIZE + # Add the new entries to this instance. + for entry in entries_added: + self.entries[(entry.path, 0)] = IndexEntry.from_base(entry) + + if write: + self.write(ignore_extension_data=not write_extension_data) + # END handle write + + return entries_added + + def _items_to_rela_paths( + self, + items: Union[PathLike, Sequence[Union[PathLike, BaseIndexEntry, Blob, Submodule]]], + ) -> List[PathLike]: + """Returns a list of repo-relative paths from the given items which + may be absolute or relative paths, entries or blobs.""" + paths = [] + # If string, put in list. + if isinstance(items, (str, os.PathLike)): + items = [items] + + for item in items: + if isinstance(item, (BaseIndexEntry, (Blob, Submodule))): + paths.append(self._to_relative_path(item.path)) + elif isinstance(item, (str, os.PathLike)): + paths.append(self._to_relative_path(item)) + else: + raise TypeError("Invalid item type: %r" % item) + # END for each item + return paths + + @post_clear_cache + @default_index + def remove( + self, + items: Union[PathLike, Sequence[Union[PathLike, Blob, BaseIndexEntry, "Submodule"]]], + working_tree: bool = False, + **kwargs: Any, + ) -> List[str]: + R"""Remove the given items from the index and optionally from the working tree + as well. + + :param items: + Multiple types of items are supported which may be be freely mixed. + + - path string + + Remove the given path at all stages. If it is a directory, you must + specify the ``r=True`` keyword argument to remove all file entries below + it. If absolute paths are given, they will be converted to a path + relative to the git repository directory containing the working tree + + The path string may include globs, such as ``*.c``. + + - :class:`~git.objects.blob.Blob` object + + Only the path portion is used in this case. + + - :class:`~git.index.typ.BaseIndexEntry` or compatible type + + The only relevant information here is the path. The stage is ignored. + + :param working_tree: + If ``True``, the entry will also be removed from the working tree, + physically removing the respective file. This may fail if there are + uncommitted changes in it. + + :param kwargs: + Additional keyword arguments to be passed to :manpage:`git-rm(1)`, such as + ``r`` to allow recursive removal. + + :return: + List(path_string, ...) list of repository relative paths that have been + removed effectively. + + This is interesting to know in case you have provided a directory or globs. + Paths are relative to the repository. + """ + args = [] + if not working_tree: + args.append("--cached") + args.append("--") + + # Preprocess paths. + paths = list(map(os.fspath, self._items_to_rela_paths(items))) # type: ignore[arg-type] + removed_paths = self.repo.git.rm(args, paths, **kwargs).splitlines() + + # Process output to gain proper paths. + # rm 'path' + return [p[4:-1] for p in removed_paths] + + @post_clear_cache + @default_index + def move( + self, + items: Union[PathLike, Sequence[Union[PathLike, Blob, BaseIndexEntry, "Submodule"]]], + skip_errors: bool = False, + **kwargs: Any, + ) -> List[Tuple[str, str]]: + """Rename/move the items, whereas the last item is considered the destination of + the move operation. + + If the destination is a file, the first item (of two) must be a file as well. + + If the destination is a directory, it may be preceded by one or more directories + or files. + + The working tree will be affected in non-bare repositories. + + :param items: + Multiple types of items are supported, please see the :meth:`remove` method + for reference. + + :param skip_errors: + If ``True``, errors such as ones resulting from missing source files will be + skipped. + + :param kwargs: + Additional arguments you would like to pass to :manpage:`git-mv(1)`, such as + ``dry_run`` or ``force``. + + :return: + List(tuple(source_path_string, destination_path_string), ...) + + A list of pairs, containing the source file moved as well as its actual + destination. Relative to the repository root. + + :raise ValueError: + If only one item was given. + + :raise git.exc.GitCommandError: + If git could not handle your request. + """ + args = [] + if skip_errors: + args.append("-k") + + paths = self._items_to_rela_paths(items) + if len(paths) < 2: + raise ValueError("Please provide at least one source and one destination of the move operation") + + was_dry_run = kwargs.pop("dry_run", kwargs.pop("n", None)) + kwargs["dry_run"] = True + + # First execute rename in dry run so the command tells us what it actually does + # (for later output). + out = [] + mvlines = self.repo.git.mv(args, paths, **kwargs).splitlines() + + # Parse result - first 0:n/2 lines are 'checking ', the remaining ones are the + # 'renaming' ones which we parse. + for ln in range(int(len(mvlines) / 2), len(mvlines)): + tokens = mvlines[ln].split(" to ") + assert len(tokens) == 2, "Too many tokens in %s" % mvlines[ln] + + # [0] = Renaming x + # [1] = y + out.append((tokens[0][9:], tokens[1])) + # END for each line to parse + + # Either prepare for the real run, or output the dry-run result. + if was_dry_run: + return out + # END handle dry run + + # Now apply the actual operation. + kwargs.pop("dry_run") + self.repo.git.mv(args, paths, **kwargs) + + return out + + def commit( + self, + message: str, + parent_commits: Union[List[Commit], None] = None, + head: bool = True, + author: Union[None, Actor] = None, + committer: Union[None, Actor] = None, + author_date: Union[datetime.datetime, str, None] = None, + commit_date: Union[datetime.datetime, str, None] = None, + skip_hooks: bool = False, + ) -> Commit: + """Commit the current default index file, creating a + :class:`~git.objects.commit.Commit` object. + + For more information on the arguments, see + :meth:`Commit.create_from_tree `. + + :note: + If you have manually altered the :attr:`entries` member of this instance, + don't forget to :meth:`write` your changes to disk beforehand. + + :note: + Passing ``skip_hooks=True`` is the equivalent of using ``-n`` or + ``--no-verify`` on the command line. + + :return: + :class:`~git.objects.commit.Commit` object representing the new commit + """ + if not skip_hooks: + run_commit_hook("pre-commit", self) + + self._write_commit_editmsg(message) + run_commit_hook("commit-msg", self, self._commit_editmsg_filepath()) + message = self._read_commit_editmsg() + self._remove_commit_editmsg() + tree = self.write_tree() + rval = Commit.create_from_tree( + self.repo, + tree, + message, + parent_commits, + head, + author=author, + committer=committer, + author_date=author_date, + commit_date=commit_date, + ) + if not skip_hooks: + run_commit_hook("post-commit", self) + return rval + + def _write_commit_editmsg(self, message: str) -> None: + with open(self._commit_editmsg_filepath(), "wb") as commit_editmsg_file: + commit_editmsg_file.write(message.encode(defenc)) + + def _remove_commit_editmsg(self) -> None: + os.remove(self._commit_editmsg_filepath()) + + def _read_commit_editmsg(self) -> str: + with open(self._commit_editmsg_filepath(), "rb") as commit_editmsg_file: + return commit_editmsg_file.read().decode(defenc) + + def _commit_editmsg_filepath(self) -> str: + return osp.join(self.repo.common_dir, "COMMIT_EDITMSG") + + def _flush_stdin_and_wait(cls, proc: "Popen[bytes]", ignore_stdout: bool = False) -> bytes: + stdin_IO = proc.stdin + if stdin_IO: + stdin_IO.flush() + stdin_IO.close() + + stdout = b"" + if not ignore_stdout and proc.stdout: + stdout = proc.stdout.read() + + if proc.stdout: + proc.stdout.close() + proc.wait() + return stdout + + @default_index + def checkout( + self, + paths: Union[None, Iterable[PathLike]] = None, + force: bool = False, + fprogress: Callable = lambda *args: None, + **kwargs: Any, + ) -> Union[None, Iterator[PathLike], Sequence[PathLike]]: + """Check out the given paths or all files from the version known to the index + into the working tree. + + :note: + Be sure you have written pending changes using the :meth:`write` method in + case you have altered the entries dictionary directly. + + :param paths: + If ``None``, all paths in the index will be checked out. + Otherwise an iterable of relative or absolute paths or a single path + pointing to files or directories in the index is expected. + + :param force: + If ``True``, existing files will be overwritten even if they contain local + modifications. + If ``False``, these will trigger a :exc:`~git.exc.CheckoutError`. + + :param fprogress: + See :meth:`IndexFile.add` for signature and explanation. + + The provided progress information will contain ``None`` as path and item if + no explicit paths are given. Otherwise progress information will be send + prior and after a file has been checked out. + + :param kwargs: + Additional arguments to be passed to :manpage:`git-checkout-index(1)`. + + :return: + Iterable yielding paths to files which have been checked out and are + guaranteed to match the version stored in the index. + + :raise git.exc.CheckoutError: + * If at least one file failed to be checked out. This is a summary, hence it + will checkout as many files as it can anyway. + * If one of files or directories do not exist in the index (as opposed to + the original git command, which ignores them). + + :raise git.exc.GitCommandError: + If error lines could not be parsed - this truly is an exceptional state. + + :note: + The checkout is limited to checking out the files in the index. Files which + are not in the index anymore and exist in the working tree will not be + deleted. This behaviour is fundamentally different to ``head.checkout``, + i.e. if you want :manpage:`git-checkout(1)`-like behaviour, use + ``head.checkout`` instead of ``index.checkout``. + """ + args = ["--index"] + if force: + args.append("--force") + + failed_files = [] + failed_reasons = [] + unknown_lines = [] + + def handle_stderr(proc: "Popen[bytes]", iter_checked_out_files: Iterable[PathLike]) -> None: + stderr_IO = proc.stderr + if not stderr_IO: + return # Return early if stderr empty. + + stderr_bytes = stderr_IO.read() + # line contents: + stderr = stderr_bytes.decode(defenc) + # git-checkout-index: this already exists + endings = ( + " already exists", + " is not in the cache", + " does not exist at stage", + " is unmerged", + ) + for line in stderr.splitlines(): + if not line.startswith("git checkout-index: ") and not line.startswith("git-checkout-index: "): + is_a_dir = " is a directory" + unlink_issue = "unable to unlink old '" + already_exists_issue = " already exists, no checkout" # created by entry.c:checkout_entry(...) + if line.endswith(is_a_dir): + failed_files.append(line[: -len(is_a_dir)]) + failed_reasons.append(is_a_dir) + elif line.startswith(unlink_issue): + failed_files.append(line[len(unlink_issue) : line.rfind("'")]) + failed_reasons.append(unlink_issue) + elif line.endswith(already_exists_issue): + failed_files.append(line[: -len(already_exists_issue)]) + failed_reasons.append(already_exists_issue) + else: + unknown_lines.append(line) + continue + # END special lines parsing + + for e in endings: + if line.endswith(e): + failed_files.append(line[20 : -len(e)]) + failed_reasons.append(e) + break + # END if ending matches + # END for each possible ending + # END for each line + if unknown_lines: + raise GitCommandError(("git-checkout-index",), 128, stderr) + if failed_files: + valid_files = list(set(iter_checked_out_files) - set(failed_files)) + raise CheckoutError( + "Some files could not be checked out from the index due to local modifications", + failed_files, + valid_files, + failed_reasons, + ) + + # END stderr handler + + if paths is None: + args.append("--all") + kwargs["as_process"] = 1 + fprogress(None, False, None) + proc = self.repo.git.checkout_index(*args, **kwargs) + proc.wait() + fprogress(None, True, None) + rval_iter = (e.path for e in self.entries.values()) + handle_stderr(proc, rval_iter) + return rval_iter + else: + if isinstance(paths, str): + paths = [paths] + + # Make sure we have our entries loaded before we start checkout_index, which + # will hold a lock on it. We try to get the lock as well during our entries + # initialization. + self.entries # noqa: B018 + + args.append("--stdin") + kwargs["as_process"] = True + kwargs["istream"] = subprocess.PIPE + proc = self.repo.git.checkout_index(args, **kwargs) + + # FIXME: Reading from GIL! + def make_exc() -> GitCommandError: + return GitCommandError(("git-checkout-index", *args), 128, proc.stderr.read()) + + checked_out_files: List[PathLike] = [] + + for path in paths: + co_path = to_native_path_linux(self._to_relative_path(path)) + # If the item is not in the index, it could be a directory. + path_is_directory = False + + try: + self.entries[(co_path, 0)] + except KeyError: + folder = co_path + if not folder.endswith("/"): + folder += "/" + for entry in self.entries.values(): + if os.fspath(entry.path).startswith(folder): + p = entry.path + self._write_path_to_stdin(proc, p, p, make_exc, fprogress, read_from_stdout=False) + checked_out_files.append(p) + path_is_directory = True + # END if entry is in directory + # END for each entry + # END path exception handlnig + + if not path_is_directory: + self._write_path_to_stdin(proc, co_path, path, make_exc, fprogress, read_from_stdout=False) + checked_out_files.append(co_path) + # END path is a file + # END for each path + try: + self._flush_stdin_and_wait(proc, ignore_stdout=True) + except GitCommandError: + # Without parsing stdout we don't know what failed. + raise CheckoutError( # noqa: B904 + "Some files could not be checked out from the index, probably because they didn't exist.", + failed_files, + [], + failed_reasons, + ) + + handle_stderr(proc, checked_out_files) + return checked_out_files + # END paths handling + + @default_index + def reset( + self, + commit: Union[Commit, "Reference", str] = "HEAD", + working_tree: bool = False, + paths: Union[None, Iterable[PathLike]] = None, + head: bool = False, + **kwargs: Any, + ) -> "IndexFile": + """Reset the index to reflect the tree at the given commit. This will not adjust + our HEAD reference by default, as opposed to + :meth:`HEAD.reset `. + + :param commit: + Revision, :class:`~git.refs.reference.Reference` or + :class:`~git.objects.commit.Commit` specifying the commit we should + represent. + + If you want to specify a tree only, use :meth:`IndexFile.from_tree` and + overwrite the default index. + + :param working_tree: + If ``True``, the files in the working tree will reflect the changed index. + If ``False``, the working tree will not be touched. + Please note that changes to the working copy will be discarded without + warning! + + :param head: + If ``True``, the head will be set to the given commit. This is ``False`` by + default, but if ``True``, this method behaves like + :meth:`HEAD.reset `. + + :param paths: + If given as an iterable of absolute or repository-relative paths, only these + will be reset to their state at the given commit-ish. + The paths need to exist at the commit, otherwise an exception will be + raised. + + :param kwargs: + Additional keyword arguments passed to :manpage:`git-reset(1)`. + + :note: + :meth:`IndexFile.reset`, as opposed to + :meth:`HEAD.reset `, will not delete any files in + order to maintain a consistent working tree. Instead, it will just check out + the files according to their state in the index. + If you want :manpage:`git-reset(1)`-like behaviour, use + :meth:`HEAD.reset ` instead. + + :return: + self + """ + # What we actually want to do is to merge the tree into our existing index, + # which is what git-read-tree does. + new_inst = type(self).from_tree(self.repo, commit) + if not paths: + self.entries = new_inst.entries + else: + nie = new_inst.entries + for path in paths: + path = self._to_relative_path(path) + try: + key = entry_key(path, 0) + self.entries[key] = nie[key] + except KeyError: + # If key is not in theirs, it mustn't be in ours. + try: + del self.entries[key] + except KeyError: + pass + # END handle deletion keyerror + # END handle keyerror + # END for each path + # END handle paths + self.write() + + if working_tree: + self.checkout(paths=paths, force=True) + # END handle working tree + + if head: + self.repo.head.set_commit(self.repo.commit(commit), logmsg="%s: Updating HEAD" % commit) + # END handle head change + + return self + + # FIXME: This is documented to accept the same parameters as Diffable.diff, but this + # does not handle NULL_TREE for `other`. (The suppressed mypy error is about this.) + def diff( + self, + other: Union[ # type: ignore[override] + Literal[git_diff.DiffConstants.INDEX], + "Tree", + "Commit", + str, + None, + ] = git_diff.INDEX, + paths: Union[PathLike, List[PathLike], Tuple[PathLike, ...], None] = None, + create_patch: bool = False, + **kwargs: Any, + ) -> git_diff.DiffIndex[git_diff.Diff]: + """Diff this index against the working copy or a :class:`~git.objects.tree.Tree` + or :class:`~git.objects.commit.Commit` object. + + For documentation of the parameters and return values, see + :meth:`Diffable.diff `. + + :note: + Will only work with indices that represent the default git index as they + have not been initialized with a stream. + """ + # Only run if we are the default repository index. + if self._file_path != self._index_path(): + raise AssertionError("Cannot call %r on indices that do not represent the default git index" % self.diff()) + # Index against index is always empty. + if other is self.INDEX: + return git_diff.DiffIndex() + + # Index against anything but None is a reverse diff with the respective item. + # Handle existing -R flags properly. + # Transform strings to the object so that we can call diff on it. + if isinstance(other, str): + other = self.repo.rev_parse(other) + # END object conversion + + if isinstance(other, Object): # For Tree or Commit. + # Invert the existing R flag. + cur_val = kwargs.get("R", False) + kwargs["R"] = not cur_val + return other.diff(self.INDEX, paths, create_patch, **kwargs) + # END diff against other item handling + + # If other is not None here, something is wrong. + if other is not None: + raise ValueError("other must be None, Diffable.INDEX, a Tree or Commit, was %r" % other) + + # Diff against working copy - can be handled by superclass natively. + return super().diff(other, paths, create_patch, **kwargs) diff --git a/lib/python3.12/site-packages/git/index/fun.py b/lib/python3.12/site-packages/git/index/fun.py new file mode 100644 index 0000000000000000000000000000000000000000..629c19b1e16ea40ae7b30bcad70ffbf2b9ef7265 --- /dev/null +++ b/lib/python3.12/site-packages/git/index/fun.py @@ -0,0 +1,471 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Standalone functions to accompany the index implementation and make it more +versatile.""" + +__all__ = [ + "write_cache", + "read_cache", + "write_tree_from_cache", + "entry_key", + "stat_mode_to_index_mode", + "S_IFGITLINK", + "run_commit_hook", + "hook_path", +] + +from io import BytesIO +import os +import os.path as osp +from pathlib import Path +from stat import S_IFDIR, S_IFLNK, S_IFMT, S_IFREG, S_ISDIR, S_ISLNK, S_IXUSR +import subprocess +import sys + +from gitdb.base import IStream +from gitdb.typ import str_tree_type + +from git.cmd import handle_process_output, safer_popen +from git.compat import defenc, force_bytes, force_text, safe_decode +from git.exc import HookExecutionError, UnmergedEntriesError +from git.objects.fun import ( + traverse_tree_recursive, + traverse_trees_recursive, + tree_to_stream, +) +from git.util import IndexFileSHA1Writer, finalize_process + +from .typ import CE_EXTENDED, BaseIndexEntry, IndexEntry, CE_NAMEMASK, CE_STAGESHIFT +from .util import pack, unpack + +# typing ----------------------------------------------------------------------------- + +from typing import Dict, IO, List, Sequence, TYPE_CHECKING, Tuple, Type, Union, cast + +from git.types import PathLike + +if TYPE_CHECKING: + from git.db import GitCmdObjectDB + from git.objects.tree import TreeCacheTup + + from .base import IndexFile + +# ------------------------------------------------------------------------------------ + +S_IFGITLINK = S_IFLNK | S_IFDIR +"""Flags for a submodule.""" + +CE_NAMEMASK_INV = ~CE_NAMEMASK + + +def hook_path(name: str, git_dir: PathLike) -> str: + """:return: path to the given named hook in the given git repository directory""" + return osp.join(git_dir, "hooks", name) + + +def _has_file_extension(path: str) -> str: + return osp.splitext(path)[1] + + +def run_commit_hook(name: str, index: "IndexFile", *args: str) -> None: + """Run the commit hook of the given name. Silently ignore hooks that do not exist. + + :param name: + Name of hook, like ``pre-commit``. + + :param index: + :class:`~git.index.base.IndexFile` instance. + + :param args: + Arguments passed to hook file. + + :raise git.exc.HookExecutionError: + """ + hp = hook_path(name, index.repo.git_dir) + if not os.access(hp, os.X_OK): + return + + env = os.environ.copy() + env["GIT_INDEX_FILE"] = safe_decode(os.fspath(index.path)) + env["GIT_EDITOR"] = ":" + cmd = [hp] + try: + if sys.platform == "win32" and not _has_file_extension(hp): + # Windows only uses extensions to determine how to open files + # (doesn't understand shebangs). Try using bash to run the hook. + relative_hp = Path(hp).relative_to(index.repo.working_dir).as_posix() + cmd = ["bash.exe", relative_hp] + + process = safer_popen( + cmd + list(args), + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + cwd=index.repo.working_dir, + ) + except Exception as ex: + raise HookExecutionError(hp, ex) from ex + else: + stdout_list: List[str] = [] + stderr_list: List[str] = [] + handle_process_output(process, stdout_list.append, stderr_list.append, finalize_process) + stdout = "".join(stdout_list) + stderr = "".join(stderr_list) + if process.returncode != 0: + stdout = force_text(stdout, defenc) + stderr = force_text(stderr, defenc) + raise HookExecutionError(hp, process.returncode, stderr, stdout) + # END handle return code + + +def stat_mode_to_index_mode(mode: int) -> int: + """Convert the given mode from a stat call to the corresponding index mode and + return it.""" + if S_ISLNK(mode): # symlinks + return S_IFLNK + if S_ISDIR(mode) or S_IFMT(mode) == S_IFGITLINK: # submodules + return S_IFGITLINK + return S_IFREG | (mode & S_IXUSR and 0o755 or 0o644) # blobs with or without executable bit + + +def write_cache( + entries: Sequence[Union[BaseIndexEntry, "IndexEntry"]], + stream: IO[bytes], + extension_data: Union[None, bytes] = None, + ShaStreamCls: Type[IndexFileSHA1Writer] = IndexFileSHA1Writer, +) -> None: + """Write the cache represented by entries to a stream. + + :param entries: + **Sorted** list of entries. + + :param stream: + Stream to wrap into the AdapterStreamCls - it is used for final output. + + :param ShaStreamCls: + Type to use when writing to the stream. It produces a sha while writing to it, + before the data is passed on to the wrapped stream. + + :param extension_data: + Any kind of data to write as a trailer, it must begin a 4 byte identifier, + followed by its size (4 bytes). + """ + # Wrap the stream into a compatible writer. + stream_sha = ShaStreamCls(stream) + + tell = stream_sha.tell + write = stream_sha.write + + # Header + version = 3 if any(entry.extended_flags for entry in entries) else 2 + write(b"DIRC") + write(pack(">LL", version, len(entries))) + + # Body + for entry in entries: + beginoffset = tell() + write(entry.ctime_bytes) # ctime + write(entry.mtime_bytes) # mtime + path_str = str(entry.path) + path: bytes = force_bytes(path_str, encoding=defenc) + plen = len(path) & CE_NAMEMASK # Path length + assert plen == len(path), "Path %s too long to fit into index" % entry.path + flags = plen | (entry.flags & CE_NAMEMASK_INV) # Clear possible previous values. + if entry.extended_flags: + flags |= CE_EXTENDED + write( + pack( + ">LLLLLL20sH", + entry.dev, + entry.inode, + entry.mode, + entry.uid, + entry.gid, + entry.size, + entry.binsha, + flags, + ) + ) + if entry.extended_flags: + write(pack(">H", entry.extended_flags)) + write(path) + real_size = (tell() - beginoffset + 8) & ~7 + write(b"\0" * ((beginoffset + real_size) - tell())) + # END for each entry + + # Write previously cached extensions data. + if extension_data is not None: + stream_sha.write(extension_data) + + # Write the sha over the content. + stream_sha.write_sha() + + +def read_header(stream: IO[bytes]) -> Tuple[int, int]: + """Return tuple(version_long, num_entries) from the given stream.""" + type_id = stream.read(4) + if type_id != b"DIRC": + raise AssertionError("Invalid index file header: %r" % type_id) + unpacked = cast(Tuple[int, int], unpack(">LL", stream.read(4 * 2))) + version, num_entries = unpacked + + assert version in (1, 2, 3), "Unsupported git index version %i, only 1, 2, and 3 are supported" % version + return version, num_entries + + +def entry_key(*entry: Union[BaseIndexEntry, PathLike, int]) -> Tuple[PathLike, int]: + """ + :return: + Key suitable to be used for the + :attr:`index.entries ` dictionary. + + :param entry: + One instance of type BaseIndexEntry or the path and the stage. + """ + + # def is_entry_key_tup(entry_key: Tuple) -> TypeGuard[Tuple[PathLike, int]]: + # return isinstance(entry_key, tuple) and len(entry_key) == 2 + + if len(entry) == 1: + entry_first = entry[0] + assert isinstance(entry_first, BaseIndexEntry) + return (entry_first.path, entry_first.stage) + else: + # assert is_entry_key_tup(entry) + entry = cast(Tuple[PathLike, int], entry) + return entry + # END handle entry + + +def read_cache( + stream: IO[bytes], +) -> Tuple[int, Dict[Tuple[PathLike, int], "IndexEntry"], bytes, bytes]: + """Read a cache file from the given stream. + + :return: + tuple(version, entries_dict, extension_data, content_sha) + + * *version* is the integer version number. + * *entries_dict* is a dictionary which maps IndexEntry instances to a path at a + stage. + * *extension_data* is ``""`` or 4 bytes of type + 4 bytes of size + size bytes. + * *content_sha* is a 20 byte sha on all cache file contents. + """ + version, num_entries = read_header(stream) + count = 0 + entries: Dict[Tuple[PathLike, int], "IndexEntry"] = {} + + read = stream.read + tell = stream.tell + while count < num_entries: + beginoffset = tell() + ctime = unpack(">8s", read(8))[0] + mtime = unpack(">8s", read(8))[0] + (dev, ino, mode, uid, gid, size, sha, flags) = unpack(">LLLLLL20sH", read(20 + 4 * 6 + 2)) + extended_flags = 0 + if flags & CE_EXTENDED: + extended_flags = unpack(">H", read(2))[0] + path_size = flags & CE_NAMEMASK + path = read(path_size).decode(defenc) + + real_size = (tell() - beginoffset + 8) & ~7 + read((beginoffset + real_size) - tell()) + entry = IndexEntry((mode, sha, flags, path, ctime, mtime, dev, ino, uid, gid, size, extended_flags)) + # entry_key would be the method to use, but we save the effort. + entries[(path, entry.stage)] = entry + count += 1 + # END for each entry + + # The footer contains extension data and a sha on the content so far. + # Keep the extension footer,and verify we have a sha in the end. + # Extension data format is: + # 4 bytes ID + # 4 bytes length of chunk + # Repeated 0 - N times + extension_data = stream.read(~0) + assert len(extension_data) > 19, ( + "Index Footer was not at least a sha on content as it was only %i bytes in size" % len(extension_data) + ) + + content_sha = extension_data[-20:] + + # Truncate the sha in the end as we will dynamically create it anyway. + extension_data = extension_data[:-20] + + return (version, entries, extension_data, content_sha) + + +def write_tree_from_cache( + entries: List[IndexEntry], odb: "GitCmdObjectDB", sl: slice, si: int = 0 +) -> Tuple[bytes, List["TreeCacheTup"]]: + R"""Create a tree from the given sorted list of entries and put the respective + trees into the given object database. + + :param entries: + **Sorted** list of :class:`~git.index.typ.IndexEntry`\s. + + :param odb: + Object database to store the trees in. + + :param si: + Start index at which we should start creating subtrees. + + :param sl: + Slice indicating the range we should process on the entries list. + + :return: + tuple(binsha, list(tree_entry, ...)) + + A tuple of a sha and a list of tree entries being a tuple of hexsha, mode, name. + """ + tree_items: List["TreeCacheTup"] = [] + + ci = sl.start + end = sl.stop + while ci < end: + entry = entries[ci] + if entry.stage != 0: + raise UnmergedEntriesError(entry) + # END abort on unmerged + ci += 1 + rbound = entry.path.find("/", si) + if rbound == -1: + # It's not a tree. + tree_items.append((entry.binsha, entry.mode, entry.path[si:])) + else: + # Find common base range. + base = entry.path[si:rbound] + xi = ci + while xi < end: + oentry = entries[xi] + orbound = oentry.path.find("/", si) + if orbound == -1 or oentry.path[si:orbound] != base: + break + # END abort on base mismatch + xi += 1 + # END find common base + + # Enter recursion. + # ci - 1 as we want to count our current item as well. + sha, _tree_entry_list = write_tree_from_cache(entries, odb, slice(ci - 1, xi), rbound + 1) + tree_items.append((sha, S_IFDIR, base)) + + # Skip ahead. + ci = xi + # END handle bounds + # END for each entry + + # Finally create the tree. + sio = BytesIO() + tree_to_stream(tree_items, sio.write) # Writes to stream as bytes, but doesn't change tree_items. + sio.seek(0) + + istream = odb.store(IStream(str_tree_type, len(sio.getvalue()), sio)) + return (istream.binsha, tree_items) + + +def _tree_entry_to_baseindexentry(tree_entry: "TreeCacheTup", stage: int) -> BaseIndexEntry: + return BaseIndexEntry((tree_entry[1], tree_entry[0], stage << CE_STAGESHIFT, tree_entry[2])) + + +def aggressive_tree_merge(odb: "GitCmdObjectDB", tree_shas: Sequence[bytes]) -> List[BaseIndexEntry]: + R""" + :return: + List of :class:`~git.index.typ.BaseIndexEntry`\s representing the aggressive + merge of the given trees. All valid entries are on stage 0, whereas the + conflicting ones are left on stage 1, 2 or 3, whereas stage 1 corresponds to the + common ancestor tree, 2 to our tree and 3 to 'their' tree. + + :param tree_shas: + 1, 2 or 3 trees as identified by their binary 20 byte shas. If 1 or two, the + entries will effectively correspond to the last given tree. If 3 are given, a 3 + way merge is performed. + """ + out: List[BaseIndexEntry] = [] + + # One and two way is the same for us, as we don't have to handle an existing + # index, instrea + if len(tree_shas) in (1, 2): + for entry in traverse_tree_recursive(odb, tree_shas[-1], ""): + out.append(_tree_entry_to_baseindexentry(entry, 0)) + # END for each entry + return out + # END handle single tree + + if len(tree_shas) > 3: + raise ValueError("Cannot handle %i trees at once" % len(tree_shas)) + + # Three trees. + for base, ours, theirs in traverse_trees_recursive(odb, tree_shas, ""): + if base is not None: + # Base version exists. + if ours is not None: + # Ours exists. + if theirs is not None: + # It exists in all branches. Ff it was changed in both + # its a conflict. Otherwise, we take the changed version. + # This should be the most common branch, so it comes first. + if (base[0] != ours[0] and base[0] != theirs[0] and ours[0] != theirs[0]) or ( + base[1] != ours[1] and base[1] != theirs[1] and ours[1] != theirs[1] + ): + # Changed by both. + out.append(_tree_entry_to_baseindexentry(base, 1)) + out.append(_tree_entry_to_baseindexentry(ours, 2)) + out.append(_tree_entry_to_baseindexentry(theirs, 3)) + elif base[0] != ours[0] or base[1] != ours[1]: + # Only we changed it. + out.append(_tree_entry_to_baseindexentry(ours, 0)) + else: + # Either nobody changed it, or they did. In either + # case, use theirs. + out.append(_tree_entry_to_baseindexentry(theirs, 0)) + # END handle modification + else: + if ours[0] != base[0] or ours[1] != base[1]: + # They deleted it, we changed it, conflict. + out.append(_tree_entry_to_baseindexentry(base, 1)) + out.append(_tree_entry_to_baseindexentry(ours, 2)) + # else: + # # We didn't change it, ignore. + # pass + # END handle our change + # END handle theirs + else: + if theirs is None: + # Deleted in both, its fine - it's out. + pass + else: + if theirs[0] != base[0] or theirs[1] != base[1]: + # Deleted in ours, changed theirs, conflict. + out.append(_tree_entry_to_baseindexentry(base, 1)) + out.append(_tree_entry_to_baseindexentry(theirs, 3)) + # END theirs changed + # else: + # # Theirs didn't change. + # pass + # END handle theirs + # END handle ours + else: + # All three can't be None. + if ours is None: + # Added in their branch. + assert theirs is not None + out.append(_tree_entry_to_baseindexentry(theirs, 0)) + elif theirs is None: + # Added in our branch. + out.append(_tree_entry_to_baseindexentry(ours, 0)) + else: + # Both have it, except for the base, see whether it changed. + if ours[0] != theirs[0] or ours[1] != theirs[1]: + out.append(_tree_entry_to_baseindexentry(ours, 2)) + out.append(_tree_entry_to_baseindexentry(theirs, 3)) + else: + # It was added the same in both. + out.append(_tree_entry_to_baseindexentry(ours, 0)) + # END handle two items + # END handle heads + # END handle base exists + # END for each entries tuple + + return out diff --git a/lib/python3.12/site-packages/git/index/typ.py b/lib/python3.12/site-packages/git/index/typ.py new file mode 100644 index 0000000000000000000000000000000000000000..927633a9f6504b21b95bee3923973b58c5ab6218 --- /dev/null +++ b/lib/python3.12/site-packages/git/index/typ.py @@ -0,0 +1,215 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Additional types used by the index.""" + +__all__ = ["BlobFilter", "BaseIndexEntry", "IndexEntry", "StageType"] + +from binascii import b2a_hex +from pathlib import Path + +from git.objects import Blob + +from .util import pack, unpack + +# typing ---------------------------------------------------------------------- + +from typing import NamedTuple, Sequence, TYPE_CHECKING, Tuple, Union, cast + +from git.types import PathLike + +if TYPE_CHECKING: + from git.repo import Repo + +StageType = int + +# --------------------------------------------------------------------------------- + +# { Invariants +CE_NAMEMASK = 0x0FFF +CE_STAGEMASK = 0x3000 +CE_EXTENDED = 0x4000 +CE_VALID = 0x8000 +CE_STAGESHIFT = 12 + +CE_EXT_SKIP_WORKTREE = 0x4000 +CE_EXT_INTENT_TO_ADD = 0x2000 + +# } END invariants + + +class BlobFilter: + """Predicate to be used by + :meth:`IndexFile.iter_blobs ` allowing to + filter only return blobs which match the given list of directories or files. + + The given paths are given relative to the repository. + """ + + __slots__ = ("paths",) + + def __init__(self, paths: Sequence[PathLike]) -> None: + """ + :param paths: + Tuple or list of paths which are either pointing to directories or to files + relative to the current repository. + """ + self.paths = paths + + def __call__(self, stage_blob: Tuple[StageType, Blob]) -> bool: + blob_pathlike: PathLike = stage_blob[1].path + blob_path: Path = blob_pathlike if isinstance(blob_pathlike, Path) else Path(blob_pathlike) + for pathlike in self.paths: + path: Path = pathlike if isinstance(pathlike, Path) else Path(pathlike) + # TODO: Change to use `PosixPath.is_relative_to` once Python 3.8 is no + # longer supported. + filter_parts = path.parts + blob_parts = blob_path.parts + if len(filter_parts) > len(blob_parts): + continue + if all(i == j for i, j in zip(filter_parts, blob_parts)): + return True + return False + + +class BaseIndexEntryHelper(NamedTuple): + """Typed named tuple to provide named attribute access for :class:`BaseIndexEntry`. + + This is needed to allow overriding ``__new__`` in child class to preserve backwards + compatibility. + """ + + mode: int + binsha: bytes + flags: int + path: PathLike + ctime_bytes: bytes = pack(">LL", 0, 0) + mtime_bytes: bytes = pack(">LL", 0, 0) + dev: int = 0 + inode: int = 0 + uid: int = 0 + gid: int = 0 + size: int = 0 + # version 3 extended flags, only when (flags & CE_EXTENDED) is set + extended_flags: int = 0 + + +class BaseIndexEntry(BaseIndexEntryHelper): + R"""Small brother of an index entry which can be created to describe changes + done to the index in which case plenty of additional information is not required. + + As the first 4 data members match exactly to the :class:`IndexEntry` type, methods + expecting a :class:`BaseIndexEntry` can also handle full :class:`IndexEntry`\s even + if they use numeric indices for performance reasons. + """ + + def __new__( + cls, + inp_tuple: Union[ + Tuple[int, bytes, int, PathLike], + Tuple[int, bytes, int, PathLike, bytes, bytes, int, int, int, int, int, int], + ], + ) -> "BaseIndexEntry": + """Override ``__new__`` to allow construction from a tuple for backwards + compatibility.""" + return super().__new__(cls, *inp_tuple) + + def __str__(self) -> str: + return "%o %s %i\t%s" % (self.mode, self.hexsha, self.stage, self.path) + + def __repr__(self) -> str: + return "(%o, %s, %i, %s)" % (self.mode, self.hexsha, self.stage, self.path) + + @property + def hexsha(self) -> str: + """hex version of our sha""" + return b2a_hex(self.binsha).decode("ascii") + + @property + def stage(self) -> int: + """Stage of the entry, either: + + * 0 = default stage + * 1 = stage before a merge or common ancestor entry in case of a 3 way merge + * 2 = stage of entries from the 'left' side of the merge + * 3 = stage of entries from the 'right' side of the merge + + :note: + For more information, see :manpage:`git-read-tree(1)`. + """ + return (self.flags & CE_STAGEMASK) >> CE_STAGESHIFT + + @property + def skip_worktree(self) -> bool: + return (self.extended_flags & CE_EXT_SKIP_WORKTREE) > 0 + + @property + def intent_to_add(self) -> bool: + return (self.extended_flags & CE_EXT_INTENT_TO_ADD) > 0 + + @classmethod + def from_blob(cls, blob: Blob, stage: int = 0) -> "BaseIndexEntry": + """:return: Fully equipped BaseIndexEntry at the given stage""" + return cls((blob.mode, blob.binsha, stage << CE_STAGESHIFT, blob.path)) + + def to_blob(self, repo: "Repo") -> Blob: + """:return: Blob using the information of this index entry""" + return Blob(repo, self.binsha, self.mode, self.path) + + +class IndexEntry(BaseIndexEntry): + """Allows convenient access to index entry data as defined in + :class:`BaseIndexEntry` without completely unpacking it. + + Attributes usually accessed often are cached in the tuple whereas others are + unpacked on demand. + + See the properties for a mapping between names and tuple indices. + """ + + @property + def ctime(self) -> Tuple[int, int]: + """ + :return: + Tuple(int_time_seconds_since_epoch, int_nano_seconds) of the + file's creation time + """ + return cast(Tuple[int, int], unpack(">LL", self.ctime_bytes)) + + @property + def mtime(self) -> Tuple[int, int]: + """See :attr:`ctime` property, but returns modification time.""" + return cast(Tuple[int, int], unpack(">LL", self.mtime_bytes)) + + @classmethod + def from_base(cls, base: "BaseIndexEntry") -> "IndexEntry": + """ + :return: + Minimal entry as created from the given :class:`BaseIndexEntry` instance. + Missing values will be set to null-like values. + + :param base: + Instance of type :class:`BaseIndexEntry`. + """ + time = pack(">LL", 0, 0) + return IndexEntry((base.mode, base.binsha, base.flags, base.path, time, time, 0, 0, 0, 0, 0)) # type: ignore[arg-type] + + @classmethod + def from_blob(cls, blob: Blob, stage: int = 0) -> "IndexEntry": + """:return: Minimal entry resembling the given blob object""" + time = pack(">LL", 0, 0) + return IndexEntry( + ( + blob.mode, + blob.binsha, + stage << CE_STAGESHIFT, + blob.path, + time, + time, + 0, + 0, + 0, + 0, + blob.size, + ) # type: ignore[arg-type] + ) diff --git a/lib/python3.12/site-packages/git/index/util.py b/lib/python3.12/site-packages/git/index/util.py new file mode 100644 index 0000000000000000000000000000000000000000..982a5afb7d48559e284944654c009a0e9da801b9 --- /dev/null +++ b/lib/python3.12/site-packages/git/index/util.py @@ -0,0 +1,121 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Index utilities.""" + +__all__ = ["TemporaryFileSwap", "post_clear_cache", "default_index", "git_working_dir"] + +import contextlib +from functools import wraps +import os +import os.path as osp +import struct +import tempfile +from types import TracebackType + +# typing ---------------------------------------------------------------------- + +from typing import Any, Callable, TYPE_CHECKING, Optional, Type, cast + +from git.types import Literal, PathLike, _T + +if TYPE_CHECKING: + from git.index import IndexFile + +# --------------------------------------------------------------------------------- + +# { Aliases +pack = struct.pack +unpack = struct.unpack +# } END aliases + + +class TemporaryFileSwap: + """Utility class moving a file to a temporary location within the same directory and + moving it back on to where on object deletion.""" + + __slots__ = ("file_path", "tmp_file_path") + + def __init__(self, file_path: PathLike) -> None: + self.file_path = file_path + dirname, basename = osp.split(file_path) + fd, self.tmp_file_path = tempfile.mkstemp(prefix=basename, dir=dirname) + os.close(fd) + with contextlib.suppress(OSError): # It may be that the source does not exist. + os.replace(self.file_path, self.tmp_file_path) + + def __enter__(self) -> "TemporaryFileSwap": + return self + + def __exit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> Literal[False]: + if osp.isfile(self.tmp_file_path): + os.replace(self.tmp_file_path, self.file_path) + return False + + +# { Decorators + + +def post_clear_cache(func: Callable[..., _T]) -> Callable[..., _T]: + """Decorator for functions that alter the index using the git command. + + When a git command alters the index, this invalidates our possibly existing entries + dictionary, which is why it must be deleted to allow it to be lazily reread later. + """ + + @wraps(func) + def post_clear_cache_if_not_raised(self: "IndexFile", *args: Any, **kwargs: Any) -> _T: + rval = func(self, *args, **kwargs) + self._delete_entries_cache() + return rval + + # END wrapper method + + return post_clear_cache_if_not_raised + + +def default_index(func: Callable[..., _T]) -> Callable[..., _T]: + """Decorator ensuring the wrapped method may only run if we are the default + repository index. + + This is as we rely on git commands that operate on that index only. + """ + + @wraps(func) + def check_default_index(self: "IndexFile", *args: Any, **kwargs: Any) -> _T: + if self._file_path != self._index_path(): + raise AssertionError( + "Cannot call %r on indices that do not represent the default git index" % func.__name__ + ) + return func(self, *args, **kwargs) + + # END wrapper method + + return check_default_index + + +def git_working_dir(func: Callable[..., _T]) -> Callable[..., _T]: + """Decorator which changes the current working dir to the one of the git + repository in order to ensure relative paths are handled correctly.""" + + @wraps(func) + def set_git_working_dir(self: "IndexFile", *args: Any, **kwargs: Any) -> _T: + cur_wd = os.getcwd() + os.chdir(cast(PathLike, self.repo.working_tree_dir)) + try: + return func(self, *args, **kwargs) + finally: + os.chdir(cur_wd) + # END handle working dir + + # END wrapper + + return set_git_working_dir + + +# } END decorators diff --git a/lib/python3.12/site-packages/git/objects/__init__.py b/lib/python3.12/site-packages/git/objects/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4447ca50d321963e10b50178253a2d3edf5e707a --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/__init__.py @@ -0,0 +1,25 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Import all submodules' main classes into the package space.""" + +__all__ = [ + "IndexObject", + "Object", + "Blob", + "Commit", + "Submodule", + "UpdateProgress", + "RootModule", + "RootUpdateProgress", + "TagObject", + "Tree", + "TreeModifier", +] + +from .base import IndexObject, Object +from .blob import Blob +from .commit import Commit +from .submodule import RootModule, RootUpdateProgress, Submodule, UpdateProgress +from .tag import TagObject +from .tree import Tree, TreeModifier diff --git a/lib/python3.12/site-packages/git/objects/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c16a142fe17b0fa77867b33f7cbfa9e709641ab2 Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4ff22463b861e651449f3f10fd8cc2c0ea4af849 Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/__pycache__/blob.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/__pycache__/blob.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..81cf16b586aa1e379afe9ec28a5ba2ea0b12aa18 Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/__pycache__/blob.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/__pycache__/commit.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/__pycache__/commit.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fe8560077f483717846d5d92085abd9c4d000e0d Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/__pycache__/commit.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/__pycache__/fun.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/__pycache__/fun.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7ccbe84980c9e1e7419d3b1ad7ae6034427a98e Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/__pycache__/fun.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/__pycache__/tag.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/__pycache__/tag.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9282e44451da760ab7c46b5ca054ed1a392ae60d Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/__pycache__/tag.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/__pycache__/tree.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/__pycache__/tree.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18a1501c256ab1fbdd14881ed4b58c9699fc62bf Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/__pycache__/tree.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..555a29b5b475e39d524b7d67f893021f05e386fc Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/base.py b/lib/python3.12/site-packages/git/objects/base.py new file mode 100644 index 0000000000000000000000000000000000000000..faf600c6b1be2abd307e76b2cf31e0cd2211cdbd --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/base.py @@ -0,0 +1,301 @@ +# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors +# +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["Object", "IndexObject"] + +import os.path as osp + +import gitdb.typ as dbtyp + +from git.exc import WorkTreeRepositoryUnsupported +from git.util import LazyMixin, bin_to_hex, join_path_native, stream_copy + +from .util import get_object_type_by_name + +# typing ------------------------------------------------------------------ + +from typing import Any, TYPE_CHECKING, Union + +from git.types import AnyGitObject, GitObjectTypeString, PathLike + +if TYPE_CHECKING: + from gitdb.base import OStream + + from git.refs.reference import Reference + from git.repo import Repo + + from .blob import Blob + from .submodule.base import Submodule + from .tree import Tree + +IndexObjUnion = Union["Tree", "Blob", "Submodule"] + +# -------------------------------------------------------------------------- + + +class Object(LazyMixin): + """Base class for classes representing git object types. + + The following four leaf classes represent specific kinds of git objects: + + * :class:`Blob ` + * :class:`Tree ` + * :class:`Commit ` + * :class:`TagObject ` + + See :manpage:`gitglossary(7)` on: + + * "object": https://git-scm.com/docs/gitglossary#def_object + * "object type": https://git-scm.com/docs/gitglossary#def_object_type + * "blob": https://git-scm.com/docs/gitglossary#def_blob_object + * "tree object": https://git-scm.com/docs/gitglossary#def_tree_object + * "commit object": https://git-scm.com/docs/gitglossary#def_commit_object + * "tag object": https://git-scm.com/docs/gitglossary#def_tag_object + + :note: + See the :class:`~git.types.AnyGitObject` union type of the four leaf subclasses + that represent actual git object types. + + :note: + :class:`~git.objects.submodule.base.Submodule` is defined under the hierarchy + rooted at this :class:`Object` class, even though submodules are not really a + type of git object. (This also applies to its + :class:`~git.objects.submodule.root.RootModule` subclass.) + + :note: + This :class:`Object` class should not be confused with :class:`object` (the root + of the class hierarchy in Python). + """ + + NULL_HEX_SHA = "0" * 40 + NULL_BIN_SHA = b"\0" * 20 + + TYPES = ( + dbtyp.str_blob_type, + dbtyp.str_tree_type, + dbtyp.str_commit_type, + dbtyp.str_tag_type, + ) + + __slots__ = ("repo", "binsha", "size") + + type: Union[GitObjectTypeString, None] = None + """String identifying (a concrete :class:`Object` subtype for) a git object type. + + The subtypes that this may name correspond to the kinds of git objects that exist, + i.e., the objects that may be present in a git repository. + + :note: + Most subclasses represent specific types of git objects and override this class + attribute accordingly. This attribute is ``None`` in the :class:`Object` base + class, as well as the :class:`IndexObject` intermediate subclass, but never + ``None`` in concrete leaf subclasses representing specific git object types. + + :note: + See also :class:`~git.types.GitObjectTypeString`. + """ + + def __init__(self, repo: "Repo", binsha: bytes) -> None: + """Initialize an object by identifying it by its binary sha. + + All keyword arguments will be set on demand if ``None``. + + :param repo: + Repository this object is located in. + + :param binsha: + 20 byte SHA1 + """ + super().__init__() + self.repo = repo + self.binsha = binsha + assert len(binsha) == 20, "Require 20 byte binary sha, got %r, len = %i" % ( + binsha, + len(binsha), + ) + + @classmethod + def new(cls, repo: "Repo", id: Union[str, "Reference"]) -> AnyGitObject: + """ + :return: + New :class:`Object` instance of a type appropriate to the object type behind + `id`. The id of the newly created object will be a binsha even though the + input id may have been a :class:`~git.refs.reference.Reference` or rev-spec. + + :param id: + :class:`~git.refs.reference.Reference`, rev-spec, or hexsha. + + :note: + This cannot be a ``__new__`` method as it would always call :meth:`__init__` + with the input id which is not necessarily a binsha. + """ + return repo.rev_parse(str(id)) + + @classmethod + def new_from_sha(cls, repo: "Repo", sha1: bytes) -> AnyGitObject: + """ + :return: + New object instance of a type appropriate to represent the given binary sha1 + + :param sha1: + 20 byte binary sha1. + """ + if sha1 == cls.NULL_BIN_SHA: + # The NULL binsha is always the root commit. + return get_object_type_by_name(b"commit")(repo, sha1) + # END handle special case + oinfo = repo.odb.info(sha1) + inst = get_object_type_by_name(oinfo.type)(repo, oinfo.binsha) + inst.size = oinfo.size + return inst + + def _set_cache_(self, attr: str) -> None: + """Retrieve object information.""" + if attr == "size": + oinfo = self.repo.odb.info(self.binsha) + self.size = oinfo.size # type: int + else: + super()._set_cache_(attr) + + def __eq__(self, other: Any) -> bool: + """:return: ``True`` if the objects have the same SHA1""" + if not hasattr(other, "binsha"): + return False + return self.binsha == other.binsha + + def __ne__(self, other: Any) -> bool: + """:return: ``True`` if the objects do not have the same SHA1""" + if not hasattr(other, "binsha"): + return True + return self.binsha != other.binsha + + def __hash__(self) -> int: + """:return: Hash of our id allowing objects to be used in dicts and sets""" + return hash(self.binsha) + + def __str__(self) -> str: + """:return: String of our SHA1 as understood by all git commands""" + return self.hexsha + + def __repr__(self) -> str: + """:return: String with pythonic representation of our object""" + return '' % (self.__class__.__name__, self.hexsha) + + @property + def hexsha(self) -> str: + """:return: 40 byte hex version of our 20 byte binary sha""" + # b2a_hex produces bytes. + return bin_to_hex(self.binsha).decode("ascii") + + @property + def data_stream(self) -> "OStream": + """ + :return: + File-object compatible stream to the uncompressed raw data of the object + + :note: + Returned streams must be read in order. + """ + return self.repo.odb.stream(self.binsha) + + def stream_data(self, ostream: "OStream") -> "Object": + """Write our data directly to the given output stream. + + :param ostream: + File-object compatible stream object. + + :return: + self + """ + istream = self.repo.odb.stream(self.binsha) + stream_copy(istream, ostream) + return self + + +class IndexObject(Object): + """Base for all objects that can be part of the index file. + + The classes representing git object types that can be part of the index file are + :class:`~git.objects.tree.Tree` and :class:`~git.objects.blob.Blob`. In addition, + :class:`~git.objects.submodule.base.Submodule`, which is not really a git object + type but can be part of an index file, is also a subclass. + """ + + __slots__ = ("path", "mode") + + # For compatibility with iterable lists. + _id_attribute_ = "path" + + def __init__( + self, + repo: "Repo", + binsha: bytes, + mode: Union[None, int] = None, + path: Union[None, PathLike] = None, + ) -> None: + """Initialize a newly instanced :class:`IndexObject`. + + :param repo: + The :class:`~git.repo.base.Repo` we are located in. + + :param binsha: + 20 byte sha1. + + :param mode: + The stat-compatible file mode as :class:`int`. + Use the :mod:`stat` module to evaluate the information. + + :param path: + The path to the file in the file system, relative to the git repository + root, like ``file.ext`` or ``folder/other.ext``. + + :note: + Path may not be set if the index object has been created directly, as it + cannot be retrieved without knowing the parent tree. + """ + super().__init__(repo, binsha) + if mode is not None: + self.mode = mode + if path is not None: + self.path = path + + def __hash__(self) -> int: + """ + :return: + Hash of our path as index items are uniquely identifiable by path, not by + their data! + """ + return hash(self.path) + + def _set_cache_(self, attr: str) -> None: + if attr in IndexObject.__slots__: + # They cannot be retrieved later on (not without searching for them). + raise AttributeError( + "Attribute '%s' unset: path and mode attributes must have been set during %s object creation" + % (attr, type(self).__name__) + ) + else: + super()._set_cache_(attr) + # END handle slot attribute + + @property + def name(self) -> str: + """:return: Name portion of the path, effectively being the basename""" + return osp.basename(self.path) + + @property + def abspath(self) -> PathLike: + R""" + :return: + Absolute path to this index object in the file system (as opposed to the + :attr:`path` field which is a path relative to the git repository). + + The returned path will be native to the system and contains ``\`` on + Windows. + """ + if self.repo.working_tree_dir is not None: + return join_path_native(self.repo.working_tree_dir, self.path) + else: + raise WorkTreeRepositoryUnsupported("working_tree_dir was None or empty") diff --git a/lib/python3.12/site-packages/git/objects/blob.py b/lib/python3.12/site-packages/git/objects/blob.py new file mode 100644 index 0000000000000000000000000000000000000000..f7d49c9ccf6d85e3e3a1d4b9cc6889489c3133c3 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/blob.py @@ -0,0 +1,49 @@ +# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors +# +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["Blob"] + +from mimetypes import guess_type +import os +import sys + +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +from . import base + + +class Blob(base.IndexObject): + """A Blob encapsulates a git blob object. + + See :manpage:`gitglossary(7)` on "blob": + https://git-scm.com/docs/gitglossary#def_blob_object + """ + + DEFAULT_MIME_TYPE = "text/plain" + type: Literal["blob"] = "blob" + + # Valid blob modes + executable_mode = 0o100755 + file_mode = 0o100644 + link_mode = 0o120000 + + __slots__ = () + + @property + def mime_type(self) -> str: + """ + :return: + String describing the mime type of this file (based on the filename) + + :note: + Defaults to ``text/plain`` in case the actual file type is unknown. + """ + guesses = None + if self.path: + guesses = guess_type(os.fspath(self.path)) + return guesses and guesses[0] or self.DEFAULT_MIME_TYPE diff --git a/lib/python3.12/site-packages/git/objects/commit.py b/lib/python3.12/site-packages/git/objects/commit.py new file mode 100644 index 0000000000000000000000000000000000000000..8c51254a26dad9ce3f2ad93745f23cc8d07add57 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/commit.py @@ -0,0 +1,909 @@ +# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors +# +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["Commit"] + +from collections import defaultdict +import datetime +from io import BytesIO +import logging +import os +import re +from subprocess import Popen, PIPE +import sys +from time import altzone, daylight, localtime, time, timezone +import warnings + +from gitdb import IStream + +from git.cmd import Git +from git.diff import Diffable +from git.util import Actor, Stats, finalize_process, hex_to_bin + +from . import base +from .tree import Tree +from .util import ( + Serializable, + TraversableIterableObj, + altz_to_utctz_str, + from_timestamp, + parse_actor_and_date, + parse_date, +) + +# typing ------------------------------------------------------------------ + +from typing import ( + Any, + Dict, + IO, + Iterator, + List, + Sequence, + Tuple, + TYPE_CHECKING, + Union, + cast, +) + +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +from git.types import PathLike + +if TYPE_CHECKING: + from git.refs import SymbolicReference + from git.repo import Repo + +# ------------------------------------------------------------------------ + +_logger = logging.getLogger(__name__) + + +class Commit(base.Object, TraversableIterableObj, Diffable, Serializable): + """Wraps a git commit object. + + See :manpage:`gitglossary(7)` on "commit object": + https://git-scm.com/docs/gitglossary#def_commit_object + + :note: + This class will act lazily on some of its attributes and will query the value on + demand only if it involves calling the git binary. + """ + + # ENVIRONMENT VARIABLES + # Read when creating new commits. + env_author_date = "GIT_AUTHOR_DATE" + env_committer_date = "GIT_COMMITTER_DATE" + + # CONFIGURATION KEYS + conf_encoding = "i18n.commitencoding" + + # INVARIANTS + default_encoding = "UTF-8" + + type: Literal["commit"] = "commit" + + __slots__ = ( + "tree", + "author", + "authored_date", + "author_tz_offset", + "committer", + "committed_date", + "committer_tz_offset", + "message", + "parents", + "encoding", + "gpgsig", + ) + + _id_attribute_ = "hexsha" + + parents: Sequence["Commit"] + + def __init__( + self, + repo: "Repo", + binsha: bytes, + tree: Union[Tree, None] = None, + author: Union[Actor, None] = None, + authored_date: Union[int, None] = None, + author_tz_offset: Union[None, float] = None, + committer: Union[Actor, None] = None, + committed_date: Union[int, None] = None, + committer_tz_offset: Union[None, float] = None, + message: Union[str, bytes, None] = None, + parents: Union[Sequence["Commit"], None] = None, + encoding: Union[str, None] = None, + gpgsig: Union[str, None] = None, + ) -> None: + """Instantiate a new :class:`Commit`. All keyword arguments taking ``None`` as + default will be implicitly set on first query. + + :param binsha: + 20 byte sha1. + + :param tree: + A :class:`~git.objects.tree.Tree` object. + + :param author: + The author :class:`~git.util.Actor` object. + + :param authored_date: int_seconds_since_epoch + The authored DateTime - use :func:`time.gmtime` to convert it into a + different format. + + :param author_tz_offset: int_seconds_west_of_utc + The timezone that the `authored_date` is in. + + :param committer: + The committer string, as an :class:`~git.util.Actor` object. + + :param committed_date: int_seconds_since_epoch + The committed DateTime - use :func:`time.gmtime` to convert it into a + different format. + + :param committer_tz_offset: int_seconds_west_of_utc + The timezone that the `committed_date` is in. + + :param message: string + The commit message. + + :param encoding: string + Encoding of the message, defaults to UTF-8. + + :param parents: + List or tuple of :class:`Commit` objects which are our parent(s) in the + commit dependency graph. + + :return: + :class:`Commit` + + :note: + Timezone information is in the same format and in the same sign as what + :func:`time.altzone` returns. The sign is inverted compared to git's UTC + timezone. + """ + super().__init__(repo, binsha) + self.binsha = binsha + if tree is not None: + assert isinstance(tree, Tree), "Tree needs to be a Tree instance, was %s" % type(tree) + if tree is not None: + self.tree = tree + if author is not None: + self.author = author + if authored_date is not None: + self.authored_date = authored_date + if author_tz_offset is not None: + self.author_tz_offset = author_tz_offset + if committer is not None: + self.committer = committer + if committed_date is not None: + self.committed_date = committed_date + if committer_tz_offset is not None: + self.committer_tz_offset = committer_tz_offset + if message is not None: + self.message = message + if parents is not None: + self.parents = parents + if encoding is not None: + self.encoding = encoding + if gpgsig is not None: + self.gpgsig = gpgsig + + @classmethod + def _get_intermediate_items(cls, commit: "Commit") -> Tuple["Commit", ...]: + return tuple(commit.parents) + + @classmethod + def _calculate_sha_(cls, repo: "Repo", commit: "Commit") -> bytes: + """Calculate the sha of a commit. + + :param repo: + :class:`~git.repo.base.Repo` object the commit should be part of. + + :param commit: + :class:`Commit` object for which to generate the sha. + """ + + stream = BytesIO() + commit._serialize(stream) + streamlen = stream.tell() + stream.seek(0) + + istream = repo.odb.store(IStream(cls.type, streamlen, stream)) + return istream.binsha + + def replace(self, **kwargs: Any) -> "Commit": + """Create new commit object from an existing commit object. + + Any values provided as keyword arguments will replace the corresponding + attribute in the new object. + """ + + attrs = {k: getattr(self, k) for k in self.__slots__} + + for attrname in kwargs: + if attrname not in self.__slots__: + raise ValueError("invalid attribute name") + + attrs.update(kwargs) + new_commit = self.__class__(self.repo, self.NULL_BIN_SHA, **attrs) + new_commit.binsha = self._calculate_sha_(self.repo, new_commit) + + return new_commit + + def _set_cache_(self, attr: str) -> None: + if attr in Commit.__slots__: + # Read the data in a chunk, its faster - then provide a file wrapper. + _binsha, _typename, self.size, stream = self.repo.odb.stream(self.binsha) + self._deserialize(BytesIO(stream.read())) + else: + super()._set_cache_(attr) + # END handle attrs + + @property + def authored_datetime(self) -> datetime.datetime: + return from_timestamp(self.authored_date, self.author_tz_offset) + + @property + def committed_datetime(self) -> datetime.datetime: + return from_timestamp(self.committed_date, self.committer_tz_offset) + + @property + def summary(self) -> Union[str, bytes]: + """:return: First line of the commit message""" + if isinstance(self.message, str): + return self.message.split("\n", 1)[0] + else: + return self.message.split(b"\n", 1)[0] + + def count(self, paths: Union[PathLike, Sequence[PathLike]] = "", **kwargs: Any) -> int: + """Count the number of commits reachable from this commit. + + :param paths: + An optional path or a list of paths restricting the return value to commits + actually containing the paths. + + :param kwargs: + Additional options to be passed to :manpage:`git-rev-list(1)`. They must not + alter the output style of the command, or parsing will yield incorrect + results. + + :return: + An int defining the number of reachable commits + """ + # Yes, it makes a difference whether empty paths are given or not in our case as + # the empty paths version will ignore merge commits for some reason. + if paths: + return len(self.repo.git.rev_list(self.hexsha, "--", paths, **kwargs).splitlines()) + return len(self.repo.git.rev_list(self.hexsha, **kwargs).splitlines()) + + @property + def name_rev(self) -> str: + """ + :return: + String describing the commits hex sha based on the closest + :class:`~git.refs.reference.Reference`. + + :note: + Mostly useful for UI purposes. + """ + return self.repo.git.name_rev(self) + + @classmethod + def iter_items( + cls, + repo: "Repo", + rev: Union[str, "Commit", "SymbolicReference"], + paths: Union[PathLike, Sequence[PathLike]] = "", + **kwargs: Any, + ) -> Iterator["Commit"]: + R"""Find all commits matching the given criteria. + + :param repo: + The :class:`~git.repo.base.Repo`. + + :param rev: + Revision specifier. See :manpage:`git-rev-parse(1)` for viable options. + + :param paths: + An optional path or list of paths. If set only :class:`Commit`\s that + include the path or paths will be considered. + + :param kwargs: + Optional keyword arguments to :manpage:`git-rev-list(1)` where: + + * ``max_count`` is the maximum number of commits to fetch. + * ``skip`` is the number of commits to skip. + * ``since`` selects all commits since some date, e.g. ``"1970-01-01"``. + + :return: + Iterator yielding :class:`Commit` items. + """ + if "pretty" in kwargs: + raise ValueError("--pretty cannot be used as parsing expects single sha's only") + # END handle pretty + + # Use -- in all cases, to prevent possibility of ambiguous arguments. + # See https://github.com/gitpython-developers/GitPython/issues/264. + + args_list: List[PathLike] = ["--"] + + if paths: + paths_tup: Tuple[PathLike, ...] + if isinstance(paths, (str, os.PathLike)): + paths_tup = (paths,) + else: + paths_tup = tuple(paths) + + args_list.extend(paths_tup) + # END if paths + + proc = repo.git.rev_list(rev, args_list, as_process=True, **kwargs) + return cls._iter_from_process_or_stream(repo, proc) + + def iter_parents(self, paths: Union[PathLike, Sequence[PathLike]] = "", **kwargs: Any) -> Iterator["Commit"]: + R"""Iterate *all* parents of this commit. + + :param paths: + Optional path or list of paths limiting the :class:`Commit`\s to those that + contain at least one of the paths. + + :param kwargs: + All arguments allowed by :manpage:`git-rev-list(1)`. + + :return: + Iterator yielding :class:`Commit` objects which are parents of ``self`` + """ + # skip ourselves + skip = kwargs.get("skip", 1) + if skip == 0: # skip ourselves + skip = 1 + kwargs["skip"] = skip + + return self.iter_items(self.repo, self, paths, **kwargs) + + @property + def stats(self) -> Stats: + """Create a git stat from changes between this commit and its first parent + or from all changes done if this is the very first commit. + + :return: + :class:`Stats` + """ + + def process_lines(lines: List[str]) -> str: + text = "" + for file_info, line in zip(lines, lines[len(lines) // 2 :]): + change_type = file_info.split("\t")[0][-1] + (insertions, deletions, filename) = line.split("\t") + text += "%s\t%s\t%s\t%s\n" % (change_type, insertions, deletions, filename) + return text + + if not self.parents: + lines = self.repo.git.diff_tree( + self.hexsha, "--", numstat=True, no_renames=True, root=True, raw=True + ).splitlines()[1:] + text = process_lines(lines) + else: + lines = self.repo.git.diff( + self.parents[0].hexsha, self.hexsha, "--", numstat=True, no_renames=True, raw=True + ).splitlines() + text = process_lines(lines) + return Stats._list_from_string(self.repo, text) + + @property + def trailers(self) -> Dict[str, str]: + """Deprecated. Get the trailers of the message as a dictionary. + + :note: + This property is deprecated, please use either :attr:`trailers_list` or + :attr:`trailers_dict`. + + :return: + Dictionary containing whitespace stripped trailer information. + Only contains the latest instance of each trailer key. + """ + warnings.warn( + "Commit.trailers is deprecated, use Commit.trailers_list or Commit.trailers_dict instead", + DeprecationWarning, + stacklevel=2, + ) + return {k: v[0] for k, v in self.trailers_dict.items()} + + @property + def trailers_list(self) -> List[Tuple[str, str]]: + """Get the trailers of the message as a list. + + Git messages can contain trailer information that are similar to :rfc:`822` + e-mail headers. See :manpage:`git-interpret-trailers(1)`. + + This function calls ``git interpret-trailers --parse`` onto the message to + extract the trailer information, returns the raw trailer data as a list. + + Valid message with trailer:: + + Subject line + + some body information + + another information + + key1: value1.1 + key1: value1.2 + key2 : value 2 with inner spaces + + Returned list will look like this:: + + [ + ("key1", "value1.1"), + ("key1", "value1.2"), + ("key2", "value 2 with inner spaces"), + ] + + :return: + List containing key-value tuples of whitespace stripped trailer information. + """ + cmd = ["git", "interpret-trailers", "--parse"] + proc: Git.AutoInterrupt = self.repo.git.execute( # type: ignore[call-overload] + cmd, + as_process=True, + istream=PIPE, + ) + trailer: str = proc.communicate(str(self.message).encode())[0].decode("utf8") + trailer = trailer.strip() + + if not trailer: + return [] + + trailer_list = [] + for t in trailer.split("\n"): + key, val = t.split(":", 1) + trailer_list.append((key.strip(), val.strip())) + + return trailer_list + + @property + def trailers_dict(self) -> Dict[str, List[str]]: + """Get the trailers of the message as a dictionary. + + Git messages can contain trailer information that are similar to :rfc:`822` + e-mail headers. See :manpage:`git-interpret-trailers(1)`. + + This function calls ``git interpret-trailers --parse`` onto the message to + extract the trailer information. The key value pairs are stripped of leading and + trailing whitespaces before they get saved into a dictionary. + + Valid message with trailer:: + + Subject line + + some body information + + another information + + key1: value1.1 + key1: value1.2 + key2 : value 2 with inner spaces + + Returned dictionary will look like this:: + + { + "key1": ["value1.1", "value1.2"], + "key2": ["value 2 with inner spaces"], + } + + + :return: + Dictionary containing whitespace stripped trailer information, mapping + trailer keys to a list of their corresponding values. + """ + d = defaultdict(list) + for key, val in self.trailers_list: + d[key].append(val) + return dict(d) + + @classmethod + def _iter_from_process_or_stream(cls, repo: "Repo", proc_or_stream: Union[Popen, IO]) -> Iterator["Commit"]: + """Parse out commit information into a list of :class:`Commit` objects. + + We expect one line per commit, and parse the actual commit information directly + from our lighting fast object database. + + :param proc: + :manpage:`git-rev-list(1)` process instance - one sha per line. + + :return: + Iterator supplying :class:`Commit` objects + """ + + # def is_proc(inp) -> TypeGuard[Popen]: + # return hasattr(proc_or_stream, 'wait') and not hasattr(proc_or_stream, 'readline') + + # def is_stream(inp) -> TypeGuard[IO]: + # return hasattr(proc_or_stream, 'readline') + + if hasattr(proc_or_stream, "wait"): + proc_or_stream = cast(Popen, proc_or_stream) + if proc_or_stream.stdout is not None: + stream = proc_or_stream.stdout + elif hasattr(proc_or_stream, "readline"): + proc_or_stream = cast(IO, proc_or_stream) # type: ignore[redundant-cast] + stream = proc_or_stream + + readline = stream.readline + while True: + line = readline() + if not line: + break + hexsha = line.strip() + if len(hexsha) > 40: + # Split additional information, as returned by bisect for instance. + hexsha, _ = line.split(None, 1) + # END handle extra info + + assert len(hexsha) == 40, "Invalid line: %s" % hexsha + yield cls(repo, hex_to_bin(hexsha)) + # END for each line in stream + + # TODO: Review this - it seems process handling got a bit out of control due to + # many developers trying to fix the open file handles issue. + if hasattr(proc_or_stream, "wait"): + proc_or_stream = cast(Popen, proc_or_stream) + finalize_process(proc_or_stream) + + @classmethod + def create_from_tree( + cls, + repo: "Repo", + tree: Union[Tree, str], + message: str, + parent_commits: Union[None, List["Commit"]] = None, + head: bool = False, + author: Union[None, Actor] = None, + committer: Union[None, Actor] = None, + author_date: Union[None, str, datetime.datetime] = None, + commit_date: Union[None, str, datetime.datetime] = None, + ) -> "Commit": + """Commit the given tree, creating a :class:`Commit` object. + + :param repo: + :class:`~git.repo.base.Repo` object the commit should be part of. + + :param tree: + :class:`~git.objects.tree.Tree` object or hex or bin sha. + The tree of the new commit. + + :param message: + Commit message. It may be an empty string if no message is provided. It will + be converted to a string, in any case. + + :param parent_commits: + Optional :class:`Commit` objects to use as parents for the new commit. If + empty list, the commit will have no parents at all and become a root commit. + If ``None``, the current head commit will be the parent of the new commit + object. + + :param head: + If ``True``, the HEAD will be advanced to the new commit automatically. + Otherwise the HEAD will remain pointing on the previous commit. This could + lead to undesired results when diffing files. + + :param author: + The name of the author, optional. + If unset, the repository configuration is used to obtain this value. + + :param committer: + The name of the committer, optional. + If unset, the repository configuration is used to obtain this value. + + :param author_date: + The timestamp for the author field. + + :param commit_date: + The timestamp for the committer field. + + :return: + :class:`Commit` object representing the new commit. + + :note: + Additional information about the committer and author are taken from the + environment or from the git configuration. See :manpage:`git-commit-tree(1)` + for more information. + """ + if parent_commits is None: + try: + parent_commits = [repo.head.commit] + except ValueError: + # Empty repositories have no head commit. + parent_commits = [] + # END handle parent commits + else: + for p in parent_commits: + if not isinstance(p, cls): + raise ValueError(f"Parent commit '{p!r}' must be of type {cls}") + # END check parent commit types + # END if parent commits are unset + + # Retrieve all additional information, create a commit object, and serialize it. + # Generally: + # * Environment variables override configuration values. + # * Sensible defaults are set according to the git documentation. + + # COMMITTER AND AUTHOR INFO + cr = repo.config_reader() + env = os.environ + + committer = committer or Actor.committer(cr) + author = author or Actor.author(cr) + + # PARSE THE DATES + unix_time = int(time()) + is_dst = daylight and localtime().tm_isdst > 0 + offset = altzone if is_dst else timezone + + author_date_str = env.get(cls.env_author_date, "") + if author_date: + author_time, author_offset = parse_date(author_date) + elif author_date_str: + author_time, author_offset = parse_date(author_date_str) + else: + author_time, author_offset = unix_time, offset + # END set author time + + committer_date_str = env.get(cls.env_committer_date, "") + if commit_date: + committer_time, committer_offset = parse_date(commit_date) + elif committer_date_str: + committer_time, committer_offset = parse_date(committer_date_str) + else: + committer_time, committer_offset = unix_time, offset + # END set committer time + + # Assume UTF-8 encoding. + enc_section, enc_option = cls.conf_encoding.split(".") + conf_encoding = cr.get_value(enc_section, enc_option, cls.default_encoding) + if not isinstance(conf_encoding, str): + raise TypeError("conf_encoding could not be coerced to str") + + # If the tree is no object, make sure we create one - otherwise the created + # commit object is invalid. + if isinstance(tree, str): + tree = repo.tree(tree) + # END tree conversion + + # CREATE NEW COMMIT + new_commit = cls( + repo, + cls.NULL_BIN_SHA, + tree, + author, + author_time, + author_offset, + committer, + committer_time, + committer_offset, + message, + parent_commits, + conf_encoding, + ) + + new_commit.binsha = cls._calculate_sha_(repo, new_commit) + + if head: + # Need late import here, importing git at the very beginning throws as + # well... + import git.refs + + try: + repo.head.set_commit(new_commit, logmsg=message) + except ValueError: + # head is not yet set to the ref our HEAD points to. + # Happens on first commit. + master = git.refs.Head.create( + repo, + repo.head.ref, + new_commit, + logmsg="commit (initial): %s" % message, + ) + repo.head.set_reference(master, logmsg="commit: Switching to %s" % master) + # END handle empty repositories + # END advance head handling + + return new_commit + + # { Serializable Implementation + + def _serialize(self, stream: BytesIO) -> "Commit": + write = stream.write + write(("tree %s\n" % self.tree).encode("ascii")) + for p in self.parents: + write(("parent %s\n" % p).encode("ascii")) + + a = self.author + aname = a.name + c = self.committer + fmt = "%s %s <%s> %s %s\n" + write( + ( + fmt + % ( + "author", + aname, + a.email, + self.authored_date, + altz_to_utctz_str(self.author_tz_offset), + ) + ).encode(self.encoding) + ) + + # Encode committer. + aname = c.name + write( + ( + fmt + % ( + "committer", + aname, + c.email, + self.committed_date, + altz_to_utctz_str(self.committer_tz_offset), + ) + ).encode(self.encoding) + ) + + if self.encoding != self.default_encoding: + write(("encoding %s\n" % self.encoding).encode("ascii")) + + try: + if self.__getattribute__("gpgsig"): + write(b"gpgsig") + for sigline in self.gpgsig.rstrip("\n").split("\n"): + write((" " + sigline + "\n").encode("ascii")) + except AttributeError: + pass + + write(b"\n") + + # Write plain bytes, be sure its encoded according to our encoding. + if isinstance(self.message, str): + write(self.message.encode(self.encoding)) + else: + write(self.message) + # END handle encoding + return self + + def _deserialize(self, stream: BytesIO) -> "Commit": + readline = stream.readline + self.tree = Tree(self.repo, hex_to_bin(readline().split()[1]), Tree.tree_id << 12, "") + + self.parents = [] + next_line = None + while True: + parent_line = readline() + if not parent_line.startswith(b"parent"): + next_line = parent_line + break + # END abort reading parents + self.parents.append(type(self)(self.repo, hex_to_bin(parent_line.split()[-1].decode("ascii")))) + # END for each parent line + self.parents = tuple(self.parents) + + # We don't know actual author encoding before we have parsed it, so keep the + # lines around. + author_line = next_line + committer_line = readline() + + # We might run into one or more mergetag blocks, skip those for now. + next_line = readline() + while next_line.startswith(b"mergetag "): + next_line = readline() + while next_line.startswith(b" "): + next_line = readline() + # END skip mergetags + + # Now we can have the encoding line, or an empty line followed by the optional + # message. + self.encoding = self.default_encoding + self.gpgsig = "" + + # Read headers. + enc = next_line + buf = enc.strip() + while buf: + if buf[0:10] == b"encoding ": + self.encoding = buf[buf.find(b" ") + 1 :].decode(self.encoding, "ignore") + elif buf[0:7] == b"gpgsig ": + sig = buf[buf.find(b" ") + 1 :] + b"\n" + is_next_header = False + while True: + sigbuf = readline() + if not sigbuf: + break + if sigbuf[0:1] != b" ": + buf = sigbuf.strip() + is_next_header = True + break + sig += sigbuf[1:] + # END read all signature + self.gpgsig = sig.rstrip(b"\n").decode(self.encoding, "ignore") + if is_next_header: + continue + buf = readline().strip() + + # Decode the author's name. + try: + ( + self.author, + self.authored_date, + self.author_tz_offset, + ) = parse_actor_and_date(author_line.decode(self.encoding, "replace")) + except UnicodeDecodeError: + _logger.error( + "Failed to decode author line '%s' using encoding %s", + author_line, + self.encoding, + exc_info=True, + ) + + try: + ( + self.committer, + self.committed_date, + self.committer_tz_offset, + ) = parse_actor_and_date(committer_line.decode(self.encoding, "replace")) + except UnicodeDecodeError: + _logger.error( + "Failed to decode committer line '%s' using encoding %s", + committer_line, + self.encoding, + exc_info=True, + ) + # END handle author's encoding + + # A stream from our data simply gives us the plain message. + # The end of our message stream is marked with a newline that we strip. + self.message = stream.read() + try: + self.message = self.message.decode(self.encoding, "replace") + except UnicodeDecodeError: + _logger.error( + "Failed to decode message '%s' using encoding %s", + self.message, + self.encoding, + exc_info=True, + ) + # END exception handling + + return self + + # } END serializable implementation + + @property + def co_authors(self) -> List[Actor]: + """Search the commit message for any co-authors of this commit. + + Details on co-authors: + https://github.blog/2018-01-29-commit-together-with-co-authors/ + + :return: + List of co-authors for this commit (as :class:`~git.util.Actor` objects). + """ + co_authors = [] + + if self.message: + results = re.findall( + r"^Co-authored-by: (.*) <(.*?)>$", + str(self.message), + re.MULTILINE, + ) + for author in results: + co_authors.append(Actor(*author)) + + return co_authors diff --git a/lib/python3.12/site-packages/git/objects/fun.py b/lib/python3.12/site-packages/git/objects/fun.py new file mode 100644 index 0000000000000000000000000000000000000000..fe57da13ae4599368b9ebb396c2be1dfdb3d6a83 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/fun.py @@ -0,0 +1,281 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Functions that are supposed to be as fast as possible.""" + +__all__ = [ + "tree_to_stream", + "tree_entries_from_data", + "traverse_trees_recursive", + "traverse_tree_recursive", +] + +from stat import S_ISDIR + +from git.compat import safe_decode, defenc + +# typing ---------------------------------------------- + +from typing import ( + Callable, + List, + MutableSequence, + Sequence, + Tuple, + TYPE_CHECKING, + Union, + overload, +) + +if TYPE_CHECKING: + from _typeshed import ReadableBuffer + + from git import GitCmdObjectDB + +EntryTup = Tuple[bytes, int, str] # Same as TreeCacheTup in tree.py. +EntryTupOrNone = Union[EntryTup, None] + +# --------------------------------------------------- + + +def tree_to_stream(entries: Sequence[EntryTup], write: Callable[["ReadableBuffer"], Union[int, None]]) -> None: + """Write the given list of entries into a stream using its ``write`` method. + + :param entries: + **Sorted** list of tuples with (binsha, mode, name). + + :param write: + A ``write`` method which takes a data string. + """ + ord_zero = ord("0") + bit_mask = 7 # 3 bits set. + + for binsha, mode, name in entries: + mode_str = b"" + for i in range(6): + mode_str = bytes([((mode >> (i * 3)) & bit_mask) + ord_zero]) + mode_str + # END for each 8 octal value + + # git slices away the first octal if it's zero. + if mode_str[0] == ord_zero: + mode_str = mode_str[1:] + # END save a byte + + # Here it comes: If the name is actually unicode, the replacement below will not + # work as the binsha is not part of the ascii unicode encoding - hence we must + # convert to an UTF-8 string for it to work properly. According to my tests, + # this is exactly what git does, that is it just takes the input literally, + # which appears to be UTF-8 on linux. + if isinstance(name, str): + name_bytes = name.encode(defenc) + else: + name_bytes = name # type: ignore[unreachable] # check runtime types - is always str? + write(b"".join((mode_str, b" ", name_bytes, b"\0", binsha))) + # END for each item + + +def tree_entries_from_data(data: bytes) -> List[EntryTup]: + """Read the binary representation of a tree and returns tuples of + :class:`~git.objects.tree.Tree` items. + + :param data: + Data block with tree data (as bytes). + + :return: + list(tuple(binsha, mode, tree_relative_path), ...) + """ + ord_zero = ord("0") + space_ord = ord(" ") + len_data = len(data) + i = 0 + out = [] + while i < len_data: + mode = 0 + + # Read Mode + # Some git versions truncate the leading 0, some don't. + # The type will be extracted from the mode later. + while data[i] != space_ord: + # Move existing mode integer up one level being 3 bits and add the actual + # ordinal value of the character. + mode = (mode << 3) + (data[i] - ord_zero) + i += 1 + # END while reading mode + + # Byte is space now, skip it. + i += 1 + + # Parse name, it is NULL separated. + + ns = i + while data[i] != 0: + i += 1 + # END while not reached NULL + + # Default encoding for strings in git is UTF-8. + # Only use the respective unicode object if the byte stream was encoded. + name_bytes = data[ns:i] + name = safe_decode(name_bytes) + + # Byte is NULL, get next 20. + i += 1 + sha = data[i : i + 20] + i = i + 20 + out.append((sha, mode, name)) + # END for each byte in data stream + return out + + +def _find_by_name(tree_data: MutableSequence[EntryTupOrNone], name: str, is_dir: bool, start_at: int) -> EntryTupOrNone: + """Return data entry matching the given name and tree mode or ``None``. + + Before the item is returned, the respective data item is set None in the `tree_data` + list to mark it done. + """ + + try: + item = tree_data[start_at] + if item and item[2] == name and S_ISDIR(item[1]) == is_dir: + tree_data[start_at] = None + return item + except IndexError: + pass + # END exception handling + for index, item in enumerate(tree_data): + if item and item[2] == name and S_ISDIR(item[1]) == is_dir: + tree_data[index] = None + return item + # END if item matches + # END for each item + return None + + +@overload +def _to_full_path(item: None, path_prefix: str) -> None: ... + + +@overload +def _to_full_path(item: EntryTup, path_prefix: str) -> EntryTup: ... + + +def _to_full_path(item: EntryTupOrNone, path_prefix: str) -> EntryTupOrNone: + """Rebuild entry with given path prefix.""" + if not item: + return item + return (item[0], item[1], path_prefix + item[2]) + + +def traverse_trees_recursive( + odb: "GitCmdObjectDB", tree_shas: Sequence[Union[bytes, None]], path_prefix: str +) -> List[Tuple[EntryTupOrNone, ...]]: + """ + :return: + List of list with entries according to the given binary tree-shas. + + The result is encoded in a list + of n tuple|None per blob/commit, (n == len(tree_shas)), where: + + * [0] == 20 byte sha + * [1] == mode as int + * [2] == path relative to working tree root + + The entry tuple is ``None`` if the respective blob/commit did not exist in the + given tree. + + :param tree_shas: + Iterable of shas pointing to trees. All trees must be on the same level. + A tree-sha may be ``None``, in which case ``None``. + + :param path_prefix: + A prefix to be added to the returned paths on this level. + Set it ``""`` for the first iteration. + + :note: + The ordering of the returned items will be partially lost. + """ + trees_data: List[List[EntryTupOrNone]] = [] + + nt = len(tree_shas) + for tree_sha in tree_shas: + if tree_sha is None: + data: List[EntryTupOrNone] = [] + else: + # Make new list for typing as list invariant. + data = list(tree_entries_from_data(odb.stream(tree_sha).read())) + # END handle muted trees + trees_data.append(data) + # END for each sha to get data for + + out: List[Tuple[EntryTupOrNone, ...]] = [] + + # Find all matching entries and recursively process them together if the match is a + # tree. If the match is a non-tree item, put it into the result. + # Processed items will be set None. + for ti, tree_data in enumerate(trees_data): + for ii, item in enumerate(tree_data): + if not item: + continue + # END skip already done items + entries: List[EntryTupOrNone] + entries = [None for _ in range(nt)] + entries[ti] = item + _sha, mode, name = item + is_dir = S_ISDIR(mode) # Type mode bits + + # Find this item in all other tree data items. + # Wrap around, but stop one before our current index, hence ti+nt, not + # ti+1+nt. + for tio in range(ti + 1, ti + nt): + tio = tio % nt + entries[tio] = _find_by_name(trees_data[tio], name, is_dir, ii) + + # END for each other item data + # If we are a directory, enter recursion. + if is_dir: + out.extend( + traverse_trees_recursive( + odb, + [((ei and ei[0]) or None) for ei in entries], + path_prefix + name + "/", + ) + ) + else: + out.append(tuple(_to_full_path(e, path_prefix) for e in entries)) + + # END handle recursion + # Finally mark it done. + tree_data[ii] = None + # END for each item + + # We are done with one tree, set all its data empty. + del tree_data[:] + # END for each tree_data chunk + return out + + +def traverse_tree_recursive(odb: "GitCmdObjectDB", tree_sha: bytes, path_prefix: str) -> List[EntryTup]: + """ + :return: + List of entries of the tree pointed to by the binary `tree_sha`. + + An entry has the following format: + + * [0] 20 byte sha + * [1] mode as int + * [2] path relative to the repository + + :param path_prefix: + Prefix to prepend to the front of all returned paths. + """ + entries = [] + data = tree_entries_from_data(odb.stream(tree_sha).read()) + + # Unpacking/packing is faster than accessing individual items. + for sha, mode, name in data: + if S_ISDIR(mode): + entries.extend(traverse_tree_recursive(odb, sha, path_prefix + name + "/")) + else: + entries.append((sha, mode, path_prefix + name)) + # END for each item + + return entries diff --git a/lib/python3.12/site-packages/git/objects/submodule/__init__.py b/lib/python3.12/site-packages/git/objects/submodule/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c0604e76f42559e1189a2c2114d5e638a7e808a4 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/submodule/__init__.py @@ -0,0 +1,7 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["Submodule", "UpdateProgress", "RootModule", "RootUpdateProgress"] + +from .base import Submodule, UpdateProgress +from .root import RootModule, RootUpdateProgress diff --git a/lib/python3.12/site-packages/git/objects/submodule/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/submodule/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd2122d8bffe0e1014aa7ba99234f0f31644f94d Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/submodule/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/submodule/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/submodule/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..35aeeb258d2c5543973da092015223386fdf8adc Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/submodule/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/submodule/__pycache__/root.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/submodule/__pycache__/root.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..209f701d6d30349f244c4c9551711f09ae402382 Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/submodule/__pycache__/root.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/submodule/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/git/objects/submodule/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..04e5a4b6f40718661d1c39609a0155ec0302d631 Binary files /dev/null and b/lib/python3.12/site-packages/git/objects/submodule/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/objects/submodule/base.py b/lib/python3.12/site-packages/git/objects/submodule/base.py new file mode 100644 index 0000000000000000000000000000000000000000..d183672dbd3067664e38fe8f0c3527ba13e7ced1 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/submodule/base.py @@ -0,0 +1,1642 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["Submodule", "UpdateProgress"] + +import gc +from io import BytesIO +import logging +import os +import os.path as osp +import stat +import sys +import uuid +import urllib + +import git +from git.cmd import Git +from git.compat import defenc +from git.config import GitConfigParser, SectionConstraint, cp +from git.exc import ( + BadName, + InvalidGitRepositoryError, + NoSuchPathError, + RepositoryDirtyError, +) +from git.objects.base import IndexObject, Object +from git.objects.util import TraversableIterableObj +from git.util import ( + IterableList, + RemoteProgress, + join_path_native, + rmtree, + to_native_path_linux, + unbare_repo, +) + +from .util import ( + SubmoduleConfigParser, + find_first_remote_branch, + mkhead, + sm_name, + sm_section, +) + +# typing ---------------------------------------------------------------------- + +from typing import ( + Any, + Callable, + Dict, + Iterator, + Mapping, + Sequence, + TYPE_CHECKING, + Union, + cast, +) + +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +from git.types import Commit_ish, PathLike, TBD + +if TYPE_CHECKING: + from git.index import IndexFile + from git.objects.commit import Commit + from git.refs import Head, RemoteReference + from git.repo import Repo + +# ----------------------------------------------------------------------------- + +_logger = logging.getLogger(__name__) + + +class UpdateProgress(RemoteProgress): + """Class providing detailed progress information to the caller who should + derive from it and implement the + :meth:`update(...) ` message.""" + + CLONE, FETCH, UPDWKTREE = [1 << x for x in range(RemoteProgress._num_op_codes, RemoteProgress._num_op_codes + 3)] + _num_op_codes: int = RemoteProgress._num_op_codes + 3 + + __slots__ = () + + +BEGIN = UpdateProgress.BEGIN +END = UpdateProgress.END +CLONE = UpdateProgress.CLONE +FETCH = UpdateProgress.FETCH +UPDWKTREE = UpdateProgress.UPDWKTREE + + +# IndexObject comes via the util module. It's a 'hacky' fix thanks to Python's import +# mechanism, which causes plenty of trouble if the only reason for packages and modules +# is refactoring - subpackages shouldn't depend on parent packages. +class Submodule(IndexObject, TraversableIterableObj): + """Implements access to a git submodule. They are special in that their sha + represents a commit in the submodule's repository which is to be checked out + at the path of this instance. + + The submodule type does not have a string type associated with it, as it exists + solely as a marker in the tree and index. + + All methods work in bare and non-bare repositories. + """ + + _id_attribute_ = "name" + k_modules_file = ".gitmodules" + k_head_option = "branch" + k_head_default = "master" + k_default_mode = stat.S_IFDIR | stat.S_IFLNK + """Submodule flags. Submodules are directories with link-status.""" + + type: Literal["submodule"] = "submodule" # type: ignore[assignment] + """This is a bogus type string for base class compatibility.""" + + __slots__ = ("_parent_commit", "_url", "_branch_path", "_name", "__weakref__") + + _cache_attrs = ("path", "_url", "_branch_path") + + def __init__( + self, + repo: "Repo", + binsha: bytes, + mode: Union[int, None] = None, + path: Union[PathLike, None] = None, + name: Union[str, None] = None, + parent_commit: Union["Commit", None] = None, + url: Union[str, None] = None, + branch_path: Union[PathLike, None] = None, + ) -> None: + """Initialize this instance with its attributes. + + We only document the parameters that differ from + :class:`~git.objects.base.IndexObject`. + + :param repo: + Our parent repository. + + :param binsha: + Binary sha referring to a commit in the remote repository. + See the `url` parameter. + + :param parent_commit: + The :class:`~git.objects.commit.Commit` whose tree is supposed to contain + the ``.gitmodules`` blob, or ``None`` to always point to the most recent + commit. See :meth:`set_parent_commit` for details. + + :param url: + The URL to the remote repository which is the submodule. + + :param branch_path: + Full repository-relative path to ref to checkout when cloning the remote + repository. + """ + super().__init__(repo, binsha, mode, path) + self.size = 0 + self._parent_commit = parent_commit + if url is not None: + self._url = url + if branch_path is not None: + self._branch_path = branch_path + if name is not None: + self._name = name + + def _set_cache_(self, attr: str) -> None: + if attr in ("path", "_url", "_branch_path"): + reader: SectionConstraint = self.config_reader() + # Default submodule values. + try: + self.path = reader.get("path") + except cp.NoSectionError as e: + if self.repo.working_tree_dir is not None: + raise ValueError( + "This submodule instance does not exist anymore in '%s' file" + % osp.join(self.repo.working_tree_dir, ".gitmodules") + ) from e + + self._url = reader.get("url") + # GitPython extension values - optional. + self._branch_path = reader.get_value(self.k_head_option, git.Head.to_full_path(self.k_head_default)) + elif attr == "_name": + raise AttributeError("Cannot retrieve the name of a submodule if it was not set initially") + else: + super()._set_cache_(attr) + # END handle attribute name + + @classmethod + def _get_intermediate_items(cls, item: "Submodule") -> IterableList["Submodule"]: + """:return: All the submodules of our module repository""" + try: + return cls.list_items(item.module()) + except InvalidGitRepositoryError: + return IterableList("") + # END handle intermediate items + + @classmethod + def _need_gitfile_submodules(cls, git: Git) -> bool: + return git.version_info[:3] >= (1, 7, 5) + + def __eq__(self, other: Any) -> bool: + """Compare with another submodule.""" + # We may only compare by name as this should be the ID they are hashed with. + # Otherwise this type wouldn't be hashable. + # return self.path == other.path and self.url == other.url and super().__eq__(other) + return self._name == other._name + + def __ne__(self, other: object) -> bool: + """Compare with another submodule for inequality.""" + return not (self == other) + + def __hash__(self) -> int: + """Hash this instance using its logical id, not the sha.""" + return hash(self._name) + + def __str__(self) -> str: + return self._name + + def __repr__(self) -> str: + return "git.%s(name=%s, path=%s, url=%s, branch_path=%s)" % ( + type(self).__name__, + self._name, + self.path, + self.url, + self.branch_path, + ) + + @classmethod + def _config_parser( + cls, repo: "Repo", parent_commit: Union["Commit", None], read_only: bool + ) -> SubmoduleConfigParser: + """ + :return: + Config parser constrained to our submodule in read or write mode + + :raise IOError: + If the ``.gitmodules`` file cannot be found, either locally or in the + repository at the given parent commit. Otherwise the exception would be + delayed until the first access of the config parser. + """ + parent_matches_head = True + if parent_commit is not None: + try: + parent_matches_head = repo.head.commit == parent_commit + except ValueError: + # We are most likely in an empty repository, so the HEAD doesn't point + # to a valid ref. + pass + # END handle parent_commit + fp_module: Union[str, BytesIO] + if not repo.bare and parent_matches_head and repo.working_tree_dir: + fp_module = osp.join(repo.working_tree_dir, cls.k_modules_file) + else: + assert parent_commit is not None, "need valid parent_commit in bare repositories" + try: + fp_module = cls._sio_modules(parent_commit) + except KeyError as e: + raise IOError( + "Could not find %s file in the tree of parent commit %s" % (cls.k_modules_file, parent_commit) + ) from e + # END handle exceptions + # END handle non-bare working tree + + if not read_only and (repo.bare or not parent_matches_head): + raise ValueError("Cannot write blobs of 'historical' submodule configurations") + # END handle writes of historical submodules + + return SubmoduleConfigParser(fp_module, read_only=read_only) + + def _clear_cache(self) -> None: + """Clear the possibly changed values.""" + for name in self._cache_attrs: + try: + delattr(self, name) + except AttributeError: + pass + # END try attr deletion + # END for each name to delete + + @classmethod + def _sio_modules(cls, parent_commit: "Commit") -> BytesIO: + """ + :return: + Configuration file as :class:`~io.BytesIO` - we only access it through the + respective blob's data + """ + sio = BytesIO(parent_commit.tree[cls.k_modules_file].data_stream.read()) + sio.name = cls.k_modules_file + return sio + + def _config_parser_constrained(self, read_only: bool) -> SectionConstraint: + """:return: Config parser constrained to our submodule in read or write mode""" + try: + pc = self.parent_commit + except ValueError: + pc = None + # END handle empty parent repository + parser = self._config_parser(self.repo, pc, read_only) + parser.set_submodule(self) + return SectionConstraint(parser, sm_section(self.name)) + + @classmethod + def _module_abspath(cls, parent_repo: "Repo", path: PathLike, name: str) -> PathLike: + if cls._need_gitfile_submodules(parent_repo.git): + return osp.join(parent_repo.git_dir, "modules", name) + if parent_repo.working_tree_dir: + return osp.join(parent_repo.working_tree_dir, path) + raise NotADirectoryError() + + @classmethod + def _clone_repo( + cls, + repo: "Repo", + url: str, + path: PathLike, + name: str, + allow_unsafe_options: bool = False, + allow_unsafe_protocols: bool = False, + **kwargs: Any, + ) -> "Repo": + """ + :return: + :class:`~git.repo.base.Repo` instance of newly cloned repository. + + :param repo: + Our parent repository. + + :param url: + URL to clone from. + + :param path: + Repository-relative path to the submodule checkout location. + + :param name: + Canonical name of the submodule. + + :param allow_unsafe_protocols: + Allow unsafe protocols to be used, like ``ext``. + + :param allow_unsafe_options: + Allow unsafe options to be used, like ``--upload-pack``. + + :param kwargs: + Additional arguments given to :manpage:`git-clone(1)`. + """ + module_abspath = cls._module_abspath(repo, path, name) + module_checkout_path = module_abspath + if cls._need_gitfile_submodules(repo.git): + kwargs["separate_git_dir"] = module_abspath + module_abspath_dir = osp.dirname(module_abspath) + if not osp.isdir(module_abspath_dir): + os.makedirs(module_abspath_dir) + module_checkout_path = osp.join(repo.working_tree_dir, path) # type: ignore[arg-type] + + if url.startswith("../"): + remote_name = cast("RemoteReference", repo.active_branch.tracking_branch()).remote_name + repo_remote_url = repo.remote(remote_name).url + url = os.path.join(repo_remote_url, url) + + clone = git.Repo.clone_from( + url, + module_checkout_path, + allow_unsafe_options=allow_unsafe_options, + allow_unsafe_protocols=allow_unsafe_protocols, + **kwargs, + ) + if cls._need_gitfile_submodules(repo.git): + cls._write_git_file_and_module_config(module_checkout_path, module_abspath) + + return clone + + @classmethod + def _to_relative_path(cls, parent_repo: "Repo", path: PathLike) -> PathLike: + """:return: A path guaranteed to be relative to the given parent repository + + :raise ValueError: + If path is not contained in the parent repository's working tree. + """ + path = to_native_path_linux(path) + if path.endswith("/"): + path = path[:-1] + # END handle trailing slash + + if osp.isabs(path) and parent_repo.working_tree_dir: + working_tree_linux = to_native_path_linux(parent_repo.working_tree_dir) + if not path.startswith(working_tree_linux): + raise ValueError( + "Submodule checkout path '%s' needs to be within the parents repository at '%s'" + % (working_tree_linux, path) + ) + path = path[len(working_tree_linux.rstrip("/")) + 1 :] + if not path: + raise ValueError("Absolute submodule path '%s' didn't yield a valid relative path" % path) + # END verify converted relative path makes sense + # END convert to a relative path + + return path + + @classmethod + def _write_git_file_and_module_config(cls, working_tree_dir: PathLike, module_abspath: PathLike) -> None: + """Write a ``.git`` file containing a (preferably) relative path to the actual + git module repository. + + It is an error if the `module_abspath` cannot be made into a relative path, + relative to the `working_tree_dir`. + + :note: + This will overwrite existing files! + + :note: + As we rewrite both the git file as well as the module configuration, we + might fail on the configuration and will not roll back changes done to the + git file. This should be a non-issue, but may easily be fixed if it becomes + one. + + :param working_tree_dir: + Directory to write the ``.git`` file into. + + :param module_abspath: + Absolute path to the bare repository. + """ + git_file = osp.join(working_tree_dir, ".git") + rela_path = osp.relpath(module_abspath, start=working_tree_dir) + if sys.platform == "win32" and osp.isfile(git_file): + os.remove(git_file) + with open(git_file, "wb") as fp: + fp.write(("gitdir: %s" % rela_path).encode(defenc)) + + with GitConfigParser(osp.join(module_abspath, "config"), read_only=False, merge_includes=False) as writer: + writer.set_value( + "core", + "worktree", + to_native_path_linux(osp.relpath(working_tree_dir, start=module_abspath)), + ) + + # { Edit Interface + + @classmethod + def add( + cls, + repo: "Repo", + name: str, + path: PathLike, + url: Union[str, None] = None, + branch: Union[str, None] = None, + no_checkout: bool = False, + depth: Union[int, None] = None, + env: Union[Mapping[str, str], None] = None, + clone_multi_options: Union[Sequence[TBD], None] = None, + allow_unsafe_options: bool = False, + allow_unsafe_protocols: bool = False, + ) -> "Submodule": + """Add a new submodule to the given repository. This will alter the index as + well as the ``.gitmodules`` file, but will not create a new commit. If the + submodule already exists, no matter if the configuration differs from the one + provided, the existing submodule will be returned. + + :param repo: + Repository instance which should receive the submodule. + + :param name: + The name/identifier for the submodule. + + :param path: + Repository-relative or absolute path at which the submodule should be + located. + It will be created as required during the repository initialization. + + :param url: + ``git clone ...``-compatible URL. See :manpage:`git-clone(1)` for more + information. If ``None``, the repository is assumed to exist, and the URL of + the first remote is taken instead. This is useful if you want to make an + existing repository a submodule of another one. + + :param branch: + Name of branch at which the submodule should (later) be checked out. The + given branch must exist in the remote repository, and will be checked out + locally as a tracking branch. + It will only be written into the configuration if it not ``None``, which is + when the checked out branch will be the one the remote HEAD pointed to. + The result you get in these situation is somewhat fuzzy, and it is + recommended to specify at least ``master`` here. + Examples are ``master`` or ``feature/new``. + + :param no_checkout: + If ``True``, and if the repository has to be cloned manually, no checkout + will be performed. + + :param depth: + Create a shallow clone with a history truncated to the specified number of + commits. + + :param env: + Optional dictionary containing the desired environment variables. + + Note: Provided variables will be used to update the execution environment + for ``git``. If some variable is not specified in `env` and is defined in + attr:`os.environ`, the value from attr:`os.environ` will be used. If you + want to unset some variable, consider providing an empty string as its + value. + + :param clone_multi_options: + A list of clone options. Please see + :meth:`Repo.clone ` for details. + + :param allow_unsafe_protocols: + Allow unsafe protocols to be used, like ``ext``. + + :param allow_unsafe_options: + Allow unsafe options to be used, like ``--upload-pack``. + + :return: + The newly created :class:`Submodule` instance. + + :note: + Works atomically, such that no change will be done if, for example, the + repository update fails. + """ + if repo.bare: + raise InvalidGitRepositoryError("Cannot add submodules to bare repositories") + # END handle bare repos + + path = cls._to_relative_path(repo, path) + + # Ensure we never put backslashes into the URL, as might happen on Windows. + if url is not None: + url = to_native_path_linux(url) + # END ensure URL correctness + + # INSTANTIATE INTERMEDIATE SM + sm = cls( + repo, + cls.NULL_BIN_SHA, + cls.k_default_mode, + path, + name, + url="invalid-temporary", + ) + if sm.exists(): + # Reretrieve submodule from tree. + try: + sm = repo.head.commit.tree[os.fspath(path)] + sm._name = name + return sm + except KeyError: + # Could only be in index. + index = repo.index + entry = index.entries[index.entry_key(path, 0)] + sm.binsha = entry.binsha + return sm + # END handle exceptions + # END handle existing + + # fake-repo - we only need the functionality on the branch instance. + br = git.Head(repo, git.Head.to_full_path(str(branch) or cls.k_head_default)) + has_module = sm.module_exists() + branch_is_default = branch is None + if has_module and url is not None: + if url not in [r.url for r in sm.module().remotes]: + raise ValueError( + "Specified URL '%s' does not match any remote url of the repository at '%s'" % (url, sm.abspath) + ) + # END check url + # END verify urls match + + mrepo: Union[Repo, None] = None + + if url is None: + if not has_module: + raise ValueError("A URL was not given and a repository did not exist at %s" % path) + # END check url + mrepo = sm.module() + # assert isinstance(mrepo, git.Repo) + urls = [r.url for r in mrepo.remotes] + if not urls: + raise ValueError("Didn't find any remote url in repository at %s" % sm.abspath) + # END verify we have url + url = urls[0] + else: + # Clone new repo. + kwargs: Dict[str, Union[bool, int, str, Sequence[TBD]]] = {"n": no_checkout} + if not branch_is_default: + kwargs["b"] = br.name + # END setup checkout-branch + + if depth: + if isinstance(depth, int): + kwargs["depth"] = depth + else: + raise ValueError("depth should be an integer") + if clone_multi_options: + kwargs["multi_options"] = clone_multi_options + + # _clone_repo(cls, repo, url, path, name, **kwargs): + mrepo = cls._clone_repo( + repo, + url, + path, + name, + env=env, + allow_unsafe_options=allow_unsafe_options, + allow_unsafe_protocols=allow_unsafe_protocols, + **kwargs, + ) + # END verify url + + ## See #525 for ensuring git URLs in config-files are valid under Windows. + url = Git.polish_url(url) + + # It's important to add the URL to the parent config, to let `git submodule` know. + # Otherwise there is a '-' character in front of the submodule listing: + # a38efa84daef914e4de58d1905a500d8d14aaf45 mymodule (v0.9.0-1-ga38efa8) + # -a38efa84daef914e4de58d1905a500d8d14aaf45 submodules/intermediate/one + writer: Union[GitConfigParser, SectionConstraint] + + with sm.repo.config_writer() as writer: + writer.set_value(sm_section(name), "url", url) + + # Update configuration and index. + index = sm.repo.index + with sm.config_writer(index=index, write=False) as writer: + writer.set_value("url", url) + writer.set_value("path", path) + + sm._url = url + if not branch_is_default: + # Store full path. + writer.set_value(cls.k_head_option, br.path) + sm._branch_path = br.path + + # We deliberately assume that our head matches our index! + if mrepo: + sm.binsha = mrepo.head.commit.binsha + index.add([sm], write=True) + + return sm + + def update( + self, + recursive: bool = False, + init: bool = True, + to_latest_revision: bool = False, + progress: Union["UpdateProgress", None] = None, + dry_run: bool = False, + force: bool = False, + keep_going: bool = False, + env: Union[Mapping[str, str], None] = None, + clone_multi_options: Union[Sequence[TBD], None] = None, + allow_unsafe_options: bool = False, + allow_unsafe_protocols: bool = False, + ) -> "Submodule": + """Update the repository of this submodule to point to the checkout we point at + with the binsha of this instance. + + :param recursive: + If ``True``, we will operate recursively and update child modules as well. + + :param init: + If ``True``, the module repository will be cloned into place if necessary. + + :param to_latest_revision: + If ``True``, the submodule's sha will be ignored during checkout. Instead, + the remote will be fetched, and the local tracking branch updated. This only + works if we have a local tracking branch, which is the case if the remote + repository had a master branch, or if the ``branch`` option was specified + for this submodule and the branch existed remotely. + + :param progress: + :class:`UpdateProgress` instance, or ``None`` if no progress should be + shown. + + :param dry_run: + If ``True``, the operation will only be simulated, but not performed. + All performed operations are read-only. + + :param force: + If ``True``, we may reset heads even if the repository in question is dirty. + Additionally we will be allowed to set a tracking branch which is ahead of + its remote branch back into the past or the location of the remote branch. + This will essentially 'forget' commits. + + If ``False``, local tracking branches that are in the future of their + respective remote branches will simply not be moved. + + :param keep_going: + If ``True``, we will ignore but log all errors, and keep going recursively. + Unless `dry_run` is set as well, `keep_going` could cause + subsequent/inherited errors you wouldn't see otherwise. + In conjunction with `dry_run`, it can be useful to anticipate all errors + when updating submodules. + + :param env: + Optional dictionary containing the desired environment variables. + + Note: Provided variables will be used to update the execution environment + for ``git``. If some variable is not specified in `env` and is defined in + attr:`os.environ`, value from attr:`os.environ` will be used. + + If you want to unset some variable, consider providing the empty string as + its value. + + :param clone_multi_options: + List of :manpage:`git-clone(1)` options. + Please see :meth:`Repo.clone ` for details. + They only take effect with the `init` option. + + :param allow_unsafe_protocols: + Allow unsafe protocols to be used, like ``ext``. + + :param allow_unsafe_options: + Allow unsafe options to be used, like ``--upload-pack``. + + :note: + Does nothing in bare repositories. + + :note: + This method is definitely not atomic if `recursive` is ``True``. + + :return: + self + """ + if self.repo.bare: + return self + # END pass in bare mode + + if progress is None: + progress = UpdateProgress() + # END handle progress + prefix = "" + if dry_run: + prefix = "DRY-RUN: " + # END handle prefix + + # To keep things plausible in dry-run mode. + if dry_run: + mrepo = None + # END init mrepo + + try: + # ENSURE REPO IS PRESENT AND UP-TO-DATE + ####################################### + try: + mrepo = self.module() + rmts = mrepo.remotes + len_rmts = len(rmts) + for i, remote in enumerate(rmts): + op = FETCH + if i == 0: + op |= BEGIN + # END handle start + + progress.update( + op, + i, + len_rmts, + prefix + "Fetching remote %s of submodule %r" % (remote, self.name), + ) + # =============================== + if not dry_run: + remote.fetch(progress=progress) + # END handle dry-run + # =============================== + if i == len_rmts - 1: + op |= END + # END handle end + progress.update( + op, + i, + len_rmts, + prefix + "Done fetching remote of submodule %r" % self.name, + ) + # END fetch new data + except InvalidGitRepositoryError: + mrepo = None + if not init: + return self + # END early abort if init is not allowed + + # There is no git-repository yet - but delete empty paths. + checkout_module_abspath = self.abspath + if not dry_run and osp.isdir(checkout_module_abspath): + try: + os.rmdir(checkout_module_abspath) + except OSError as e: + raise OSError( + "Module directory at %r does already exist and is non-empty" % checkout_module_abspath + ) from e + # END handle OSError + # END handle directory removal + + # Don't check it out at first - nonetheless it will create a local + # branch according to the remote-HEAD if possible. + progress.update( + BEGIN | CLONE, + 0, + 1, + prefix + + "Cloning url '%s' to '%s' in submodule %r" % (self.url, checkout_module_abspath, self.name), + ) + if not dry_run: + if self.url.startswith("."): + url = urllib.parse.urljoin(self.repo.remotes.origin.url + "/", self.url) + else: + url = self.url + mrepo = self._clone_repo( + self.repo, + url, + self.path, + self.name, + n=True, + env=env, + multi_options=clone_multi_options, + allow_unsafe_options=allow_unsafe_options, + allow_unsafe_protocols=allow_unsafe_protocols, + ) + # END handle dry-run + progress.update( + END | CLONE, + 0, + 1, + prefix + "Done cloning to %s" % checkout_module_abspath, + ) + + if not dry_run: + # See whether we have a valid branch to check out. + try: + mrepo = cast("Repo", mrepo) + # Find a remote which has our branch - we try to be flexible. + remote_branch = find_first_remote_branch(mrepo.remotes, self.branch_name) + local_branch = mkhead(mrepo, self.branch_path) + + # Have a valid branch, but no checkout - make sure we can figure + # that out by marking the commit with a null_sha. + local_branch.set_object(Object(mrepo, self.NULL_BIN_SHA)) + # END initial checkout + branch creation + + # Make sure HEAD is not detached. + mrepo.head.set_reference( + local_branch, + logmsg="submodule: attaching head to %s" % local_branch, + ) + mrepo.head.reference.set_tracking_branch(remote_branch) + except (IndexError, InvalidGitRepositoryError): + _logger.warning("Failed to checkout tracking branch %s", self.branch_path) + # END handle tracking branch + + # NOTE: Have to write the repo config file as well, otherwise the + # default implementation will be offended and not update the + # repository. Maybe this is a good way to ensure it doesn't get into + # our way, but we want to stay backwards compatible too... It's so + # redundant! + with self.repo.config_writer() as writer: + writer.set_value(sm_section(self.name), "url", self.url) + # END handle dry_run + # END handle initialization + + # DETERMINE SHAS TO CHECK OUT + ############################# + binsha = self.binsha + hexsha = self.hexsha + if mrepo is not None: + # mrepo is only set if we are not in dry-run mode or if the module + # existed. + is_detached = mrepo.head.is_detached + # END handle dry_run + + if mrepo is not None and to_latest_revision: + msg_base = "Cannot update to latest revision in repository at %r as " % mrepo.working_dir + if not is_detached: + rref = mrepo.head.reference.tracking_branch() + if rref is not None: + rcommit = rref.commit + binsha = rcommit.binsha + hexsha = rcommit.hexsha + else: + _logger.error( + "%s a tracking branch was not set for local branch '%s'", + msg_base, + mrepo.head.reference, + ) + # END handle remote ref + else: + _logger.error("%s there was no local tracking branch", msg_base) + # END handle detached head + # END handle to_latest_revision option + + # Update the working tree. + # Handles dry_run. + if mrepo is not None and mrepo.head.commit.binsha != binsha: + # We must ensure that our destination sha (the one to point to) is in + # the future of our current head. Otherwise, we will reset changes that + # might have been done on the submodule, but were not yet pushed. We + # also handle the case that history has been rewritten, leaving no + # merge-base. In that case we behave conservatively, protecting possible + # changes the user had done. + may_reset = True + if mrepo.head.commit.binsha != self.NULL_BIN_SHA: + base_commit = mrepo.merge_base(mrepo.head.commit, hexsha) + if len(base_commit) == 0 or (base_commit[0] is not None and base_commit[0].hexsha == hexsha): + if force: + msg = "Will force checkout or reset on local branch that is possibly in the future of" + msg += " the commit it will be checked out to, effectively 'forgetting' new commits" + _logger.debug(msg) + else: + msg = "Skipping %s on branch '%s' of submodule repo '%s' as it contains un-pushed commits" + msg %= ( + is_detached and "checkout" or "reset", + mrepo.head, + mrepo, + ) + _logger.info(msg) + may_reset = False + # END handle force + # END handle if we are in the future + + if may_reset and not force and mrepo.is_dirty(index=True, working_tree=True, untracked_files=True): + raise RepositoryDirtyError(mrepo, "Cannot reset a dirty repository") + # END handle force and dirty state + # END handle empty repo + + # END verify future/past + progress.update( + BEGIN | UPDWKTREE, + 0, + 1, + prefix + + "Updating working tree at %s for submodule %r to revision %s" % (self.path, self.name, hexsha), + ) + + if not dry_run and may_reset: + if is_detached: + # NOTE: For now we force. The user is not supposed to change + # detached submodules anyway. Maybe at some point this becomes + # an option, to properly handle user modifications - see below + # for future options regarding rebase and merge. + mrepo.git.checkout(hexsha, force=force) + else: + mrepo.head.reset(hexsha, index=True, working_tree=True) + # END handle checkout + # If we may reset/checkout. + progress.update( + END | UPDWKTREE, + 0, + 1, + prefix + "Done updating working tree for submodule %r" % self.name, + ) + # END update to new commit only if needed + except Exception as err: + if not keep_going: + raise + _logger.error(str(err)) + # END handle keep_going + + # HANDLE RECURSION + ################## + if recursive: + # In dry_run mode, the module might not exist. + if mrepo is not None: + for submodule in self.iter_items(self.module()): + submodule.update( + recursive, + init, + to_latest_revision, + progress=progress, + dry_run=dry_run, + force=force, + keep_going=keep_going, + ) + # END handle recursive update + # END handle dry run + # END for each submodule + + return self + + @unbare_repo + def move(self, module_path: PathLike, configuration: bool = True, module: bool = True) -> "Submodule": + """Move the submodule to a another module path. This involves physically moving + the repository at our current path, changing the configuration, as well as + adjusting our index entry accordingly. + + :param module_path: + The path to which to move our module in the parent repository's working + tree, given as repository-relative or absolute path. Intermediate + directories will be created accordingly. If the path already exists, it must + be empty. Trailing (back)slashes are removed automatically. + + :param configuration: + If ``True``, the configuration will be adjusted to let the submodule point + to the given path. + + :param module: + If ``True``, the repository managed by this submodule will be moved as well. + If ``False``, we don't move the submodule's checkout, which may leave the + parent repository in an inconsistent state. + + :return: + self + + :raise ValueError: + If the module path existed and was not empty, or was a file. + + :note: + Currently the method is not atomic, and it could leave the repository in an + inconsistent state if a sub-step fails for some reason. + """ + if module + configuration < 1: + raise ValueError("You must specify to move at least the module or the configuration of the submodule") + # END handle input + + module_checkout_path = self._to_relative_path(self.repo, module_path) + + # VERIFY DESTINATION + if module_checkout_path == self.path: + return self + # END handle no change + + module_checkout_abspath = join_path_native(str(self.repo.working_tree_dir), module_checkout_path) + if osp.isfile(module_checkout_abspath): + raise ValueError("Cannot move repository onto a file: %s" % module_checkout_abspath) + # END handle target files + + index = self.repo.index + tekey = index.entry_key(module_checkout_path, 0) + # if the target item already exists, fail + if configuration and tekey in index.entries: + raise ValueError("Index entry for target path did already exist") + # END handle index key already there + + # Remove existing destination. + if module: + if osp.exists(module_checkout_abspath): + if len(os.listdir(module_checkout_abspath)): + raise ValueError("Destination module directory was not empty") + # END handle non-emptiness + + if osp.islink(module_checkout_abspath): + os.remove(module_checkout_abspath) + else: + os.rmdir(module_checkout_abspath) + # END handle link + else: + # Recreate parent directories. + # NOTE: renames() does that now. + pass + # END handle existence + # END handle module + + # Move the module into place if possible. + cur_path = self.abspath + renamed_module = False + if module and osp.exists(cur_path): + os.renames(cur_path, module_checkout_abspath) + renamed_module = True + + if osp.isfile(osp.join(module_checkout_abspath, ".git")): + module_abspath = self._module_abspath(self.repo, self.path, self.name) + self._write_git_file_and_module_config(module_checkout_abspath, module_abspath) + # END handle git file rewrite + # END move physical module + + # Rename the index entry - we have to manipulate the index directly as git-mv + # cannot be used on submodules... yeah. + previous_sm_path = self.path + try: + if configuration: + try: + ekey = index.entry_key(self.path, 0) + entry = index.entries[ekey] + del index.entries[ekey] + nentry = git.IndexEntry(entry[:3] + (module_checkout_path,) + entry[4:]) + index.entries[tekey] = nentry + except KeyError as e: + raise InvalidGitRepositoryError("Submodule's entry at %r did not exist" % (self.path)) from e + # END handle submodule doesn't exist + + # Update configuration. + with self.config_writer(index=index) as writer: # Auto-write. + writer.set_value("path", module_checkout_path) + self.path = module_checkout_path + # END handle configuration flag + except Exception: + if renamed_module: + os.renames(module_checkout_abspath, cur_path) + # END undo module renaming + raise + # END handle undo rename + + # Auto-rename submodule if its name was 'default', that is, the checkout + # directory. + if previous_sm_path == self.name: + self.rename(module_checkout_path) + + return self + + @unbare_repo + def remove( + self, + module: bool = True, + force: bool = False, + configuration: bool = True, + dry_run: bool = False, + ) -> "Submodule": + """Remove this submodule from the repository. This will remove our entry + from the ``.gitmodules`` file and the entry in the ``.git/config`` file. + + :param module: + If ``True``, the checked out module we point to will be deleted as well. If + that module is currently on a commit outside any branch in the remote, or if + it is ahead of its tracking branch, or if there are modified or untracked + files in its working tree, then the removal will fail. In case the removal + of the repository fails for these reasons, the submodule status will not + have been altered. + + If this submodule has child modules of its own, these will be deleted prior + to touching the direct submodule. + + :param force: + Enforces the deletion of the module even though it contains modifications. + This basically enforces a brute-force file system based deletion. + + :param configuration: + If ``True``, the submodule is deleted from the configuration, otherwise it + isn't. Although this should be enabled most of the time, this flag enables + you to safely delete the repository of your submodule. + + :param dry_run: + If ``True``, we will not actually do anything, but throw the errors we would + usually throw. + + :return: + self + + :note: + Doesn't work in bare repositories. + + :note: + Doesn't work atomically, as failure to remove any part of the submodule will + leave an inconsistent state. + + :raise git.exc.InvalidGitRepositoryError: + Thrown if the repository cannot be deleted. + + :raise OSError: + If directories or files could not be removed. + """ + if not (module or configuration): + raise ValueError("Need to specify to delete at least the module, or the configuration") + # END handle parameters + + # Recursively remove children of this submodule. + nc = 0 + for csm in self.children(): + nc += 1 + csm.remove(module, force, configuration, dry_run) + del csm + + if configuration and not dry_run and nc > 0: + # Ensure we don't leave the parent repository in a dirty state, and commit + # our changes. It's important for recursive, unforced, deletions to work as + # expected. + self.module().index.commit("Removed at least one of child-modules of '%s'" % self.name) + # END handle recursion + + # DELETE REPOSITORY WORKING TREE + ################################ + if module and self.module_exists(): + mod = self.module() + git_dir = mod.git_dir + if force: + # Take the fast lane and just delete everything in our module path. + # TODO: If we run into permission problems, we have a highly + # inconsistent state. Delete the .git folders last, start with the + # submodules first. + mp = self.abspath + method: Union[None, Callable[[PathLike], None]] = None + if osp.islink(mp): + method = os.remove + elif osp.isdir(mp): + method = rmtree + elif osp.exists(mp): + raise AssertionError("Cannot forcibly delete repository as it was neither a link, nor a directory") + # END handle brutal deletion + if not dry_run: + assert method + method(mp) + # END apply deletion method + else: + # Verify we may delete our module. + if mod.is_dirty(index=True, working_tree=True, untracked_files=True): + raise InvalidGitRepositoryError( + "Cannot delete module at %s with any modifications, unless force is specified" + % mod.working_tree_dir + ) + # END check for dirt + + # Figure out whether we have new commits compared to the remotes. + # NOTE: If the user pulled all the time, the remote heads might not have + # been updated, so commits coming from the remote look as if they come + # from us. But we stay strictly read-only and don't fetch beforehand. + for remote in mod.remotes: + num_branches_with_new_commits = 0 + rrefs = remote.refs + for rref in rrefs: + num_branches_with_new_commits += len(mod.git.cherry(rref)) != 0 + # END for each remote ref + # Not a single remote branch contained all our commits. + if len(rrefs) and num_branches_with_new_commits == len(rrefs): + raise InvalidGitRepositoryError( + "Cannot delete module at %s as there are new commits" % mod.working_tree_dir + ) + # END handle new commits + # We have to manually delete some references to allow resources to + # be cleaned up immediately when we are done with them, because + # Python's scoping is no more granular than the whole function (loop + # bodies are not scopes). When the objects stay alive longer, they + # can keep handles open. On Windows, this is a problem. + if len(rrefs): + del rref # skipcq: PYL-W0631 + # END handle remotes + del rrefs + del remote + # END for each remote + + # Finally delete our own submodule. + if not dry_run: + self._clear_cache() + wtd = mod.working_tree_dir + del mod # Release file-handles (Windows). + gc.collect() + rmtree(str(wtd)) + # END delete tree if possible + # END handle force + + if not dry_run and osp.isdir(git_dir): + self._clear_cache() + rmtree(git_dir) + # END handle separate bare repository + # END handle module deletion + + # Void our data so as not to delay invalid access. + if not dry_run: + self._clear_cache() + + # DELETE CONFIGURATION + ###################### + if configuration and not dry_run: + # First the index-entry. + parent_index = self.repo.index + try: + del parent_index.entries[parent_index.entry_key(self.path, 0)] + except KeyError: + pass + # END delete entry + parent_index.write() + + # Now git config - we need the config intact, otherwise we can't query + # information anymore. + + with self.repo.config_writer() as gcp_writer: + gcp_writer.remove_section(sm_section(self.name)) + + with self.config_writer() as sc_writer: + sc_writer.remove_section() + # END delete configuration + + return self + + def set_parent_commit(self, commit: Union[Commit_ish, str, None], check: bool = True) -> "Submodule": + """Set this instance to use the given commit whose tree is supposed to + contain the ``.gitmodules`` blob. + + :param commit: + Commit-ish reference pointing at the root tree, or ``None`` to always point + to the most recent commit. + + :param check: + If ``True``, relatively expensive checks will be performed to verify + validity of the submodule. + + :raise ValueError: + If the commit's tree didn't contain the ``.gitmodules`` blob. + + :raise ValueError: + If the parent commit didn't store this submodule under the current path. + + :return: + self + """ + if commit is None: + self._parent_commit = None + return self + # END handle None + pcommit = self.repo.commit(commit) + pctree = pcommit.tree + if self.k_modules_file not in pctree: + raise ValueError("Tree of commit %s did not contain the %s file" % (commit, self.k_modules_file)) + # END handle exceptions + + prev_pc = self._parent_commit + self._parent_commit = pcommit + + if check: + parser = self._config_parser(self.repo, self._parent_commit, read_only=True) + if not parser.has_section(sm_section(self.name)): + self._parent_commit = prev_pc + raise ValueError("Submodule at path %r did not exist in parent commit %s" % (self.path, commit)) + # END handle submodule did not exist + # END handle checking mode + + # Update our sha, it could have changed. + # If check is False, we might see a parent-commit that doesn't even contain the + # submodule anymore. in that case, mark our sha as being NULL. + try: + self.binsha = pctree[str(self.path)].binsha + except KeyError: + self.binsha = self.NULL_BIN_SHA + + self._clear_cache() + return self + + @unbare_repo + def config_writer( + self, index: Union["IndexFile", None] = None, write: bool = True + ) -> SectionConstraint["SubmoduleConfigParser"]: + """ + :return: + A config writer instance allowing you to read and write the data belonging + to this submodule into the ``.gitmodules`` file. + + :param index: + If not ``None``, an :class:`~git.index.base.IndexFile` instance which should + be written. Defaults to the index of the :class:`Submodule`'s parent + repository. + + :param write: + If ``True``, the index will be written each time a configuration value changes. + + :note: + The parameters allow for a more efficient writing of the index, as you can + pass in a modified index on your own, prevent automatic writing, and write + yourself once the whole operation is complete. + + :raise ValueError: + If trying to get a writer on a parent_commit which does not match the + current head commit. + + :raise IOError: + If the ``.gitmodules`` file/blob could not be read. + """ + writer = self._config_parser_constrained(read_only=False) + if index is not None: + writer.config._index = index + writer.config._auto_write = write + return writer + + @unbare_repo + def rename(self, new_name: str) -> "Submodule": + """Rename this submodule. + + :note: + This method takes care of renaming the submodule in various places, such as: + + * ``$parent_git_dir / config`` + * ``$working_tree_dir / .gitmodules`` + * (git >= v1.8.0: move submodule repository to new name) + + As ``.gitmodules`` will be changed, you would need to make a commit afterwards. + The changed ``.gitmodules`` file will already be added to the index. + + :return: + This :class:`Submodule` instance + """ + if self.name == new_name: + return self + + # .git/config + with self.repo.config_writer() as pw: + # As we ourselves didn't write anything about submodules into the parent + # .git/config, we will not require it to exist, and just ignore missing + # entries. + if pw.has_section(sm_section(self.name)): + pw.rename_section(sm_section(self.name), sm_section(new_name)) + + # .gitmodules + with self.config_writer(write=True).config as cw: + cw.rename_section(sm_section(self.name), sm_section(new_name)) + + self._name = new_name + + # .git/modules + mod = self.module() + if mod.has_separate_working_tree(): + destination_module_abspath = self._module_abspath(self.repo, self.path, new_name) + source_dir = mod.git_dir + # Let's be sure the submodule name is not so obviously tied to a directory. + if str(destination_module_abspath).startswith(str(mod.git_dir)): + tmp_dir = self._module_abspath(self.repo, self.path, str(uuid.uuid4())) + os.renames(source_dir, tmp_dir) + source_dir = tmp_dir + # END handle self-containment + os.renames(source_dir, destination_module_abspath) + if mod.working_tree_dir: + self._write_git_file_and_module_config(mod.working_tree_dir, destination_module_abspath) + # END move separate git repository + + return self + + # } END edit interface + + # { Query Interface + + @unbare_repo + def module(self) -> "Repo": + """ + :return: + :class:`~git.repo.base.Repo` instance initialized from the repository at our + submodule path + + :raise git.exc.InvalidGitRepositoryError: + If a repository was not available. + This could also mean that it was not yet initialized. + """ + module_checkout_abspath = self.abspath + try: + repo = git.Repo(module_checkout_abspath) + if repo != self.repo: + return repo + # END handle repo uninitialized + except (InvalidGitRepositoryError, NoSuchPathError) as e: + raise InvalidGitRepositoryError("No valid repository at %s" % module_checkout_abspath) from e + else: + raise InvalidGitRepositoryError("Repository at %r was not yet checked out" % module_checkout_abspath) + # END handle exceptions + + def module_exists(self) -> bool: + """ + :return: + ``True`` if our module exists and is a valid git repository. + See the :meth:`module` method. + """ + try: + self.module() + return True + except Exception: + return False + # END handle exception + + def exists(self) -> bool: + """ + :return: + ``True`` if the submodule exists, ``False`` otherwise. + Please note that a submodule may exist (in the ``.gitmodules`` file) even + though its module doesn't exist on disk. + """ + # Keep attributes for later, and restore them if we have no valid data. + # This way we do not actually alter the state of the object. + loc = locals() + for attr in self._cache_attrs: + try: + if hasattr(self, attr): + loc[attr] = getattr(self, attr) + # END if we have the attribute cache + except (cp.NoSectionError, ValueError): + # On PY3, this can happen apparently... don't know why this doesn't + # happen on PY2. + pass + # END for each attr + self._clear_cache() + + try: + try: + self.path # noqa: B018 + return True + except Exception: + return False + # END handle exceptions + finally: + for attr in self._cache_attrs: + if attr in loc: + setattr(self, attr, loc[attr]) + # END if we have a cache + # END reapply each attribute + # END handle object state consistency + + @property + def branch(self) -> "Head": + """ + :return: + The branch instance that we are to checkout + + :raise git.exc.InvalidGitRepositoryError: + If our module is not yet checked out. + """ + return mkhead(self.module(), self._branch_path) + + @property + def branch_path(self) -> PathLike: + """ + :return: + Full repository-relative path as string to the branch we would checkout from + the remote and track + """ + return self._branch_path + + @property + def branch_name(self) -> str: + """ + :return: + The name of the branch, which is the shortest possible branch name + """ + # Use an instance method, for this we create a temporary Head instance which + # uses a repository that is available at least (it makes no difference). + return git.Head(self.repo, self._branch_path).name + + @property + def url(self) -> str: + """:return: The url to the repository our submodule's repository refers to""" + return self._url + + @property + def parent_commit(self) -> "Commit": + """ + :return: + :class:`~git.objects.commit.Commit` instance with the tree containing the + ``.gitmodules`` file + + :note: + Will always point to the current head's commit if it was not set explicitly. + """ + if self._parent_commit is None: + return self.repo.commit() + return self._parent_commit + + @property + def name(self) -> str: + """ + :return: + The name of this submodule. It is used to identify it within the + ``.gitmodules`` file. + + :note: + By default, this is the name is the path at which to find the submodule, but + in GitPython it should be a unique identifier similar to the identifiers + used for remotes, which allows to change the path of the submodule easily. + """ + return self._name + + def config_reader(self) -> SectionConstraint[SubmoduleConfigParser]: + """ + :return: + ConfigReader instance which allows you to query the configuration values of + this submodule, as provided by the ``.gitmodules`` file. + + :note: + The config reader will actually read the data directly from the repository + and thus does not need nor care about your working tree. + + :note: + Should be cached by the caller and only kept as long as needed. + + :raise IOError: + If the ``.gitmodules`` file/blob could not be read. + """ + return self._config_parser_constrained(read_only=True) + + def children(self) -> IterableList["Submodule"]: + """ + :return: + IterableList(Submodule, ...) An iterable list of :class:`Submodule` + instances which are children of this submodule or 0 if the submodule is not + checked out. + """ + return self._get_intermediate_items(self) + + # } END query interface + + # { Iterable Interface + + @classmethod + def iter_items( + cls, + repo: "Repo", + parent_commit: Union[Commit_ish, str] = "HEAD", + *args: Any, + **kwargs: Any, + ) -> Iterator["Submodule"]: + """ + :return: + Iterator yielding :class:`Submodule` instances available in the given + repository + """ + try: + pc = repo.commit(parent_commit) # Parent commit instance + parser = cls._config_parser(repo, pc, read_only=True) + except (IOError, BadName): + return + # END handle empty iterator + + for sms in parser.sections(): + n = sm_name(sms) + p = parser.get(sms, "path") + u = parser.get(sms, "url") + b = cls.k_head_default + if parser.has_option(sms, cls.k_head_option): + b = str(parser.get(sms, cls.k_head_option)) + # END handle optional information + + # Get the binsha. + index = repo.index + try: + rt = pc.tree # Root tree + sm = rt[p] + except KeyError: + # Try the index, maybe it was just added. + try: + entry = index.entries[index.entry_key(p, 0)] + sm = Submodule(repo, entry.binsha, entry.mode, entry.path) + except KeyError: + # The submodule doesn't exist, probably it wasn't removed from the + # .gitmodules file. + continue + # END handle keyerror + # END handle critical error + + # Make sure we are looking at a submodule object. + if type(sm) is not git.objects.submodule.base.Submodule: + continue + + # Fill in remaining info - saves time as it doesn't have to be parsed again. + sm._name = n + if pc != repo.commit(): + sm._parent_commit = pc + # END set only if not most recent! + sm._branch_path = git.Head.to_full_path(b) + sm._url = u + + yield sm + # END for each section + + # } END iterable interface diff --git a/lib/python3.12/site-packages/git/objects/submodule/root.py b/lib/python3.12/site-packages/git/objects/submodule/root.py new file mode 100644 index 0000000000000000000000000000000000000000..d93193fa3c92201967af27ea382490eeb01fb95d --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/submodule/root.py @@ -0,0 +1,467 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["RootModule", "RootUpdateProgress"] + +import logging + +import git +from git.exc import InvalidGitRepositoryError + +from .base import Submodule, UpdateProgress +from .util import find_first_remote_branch + +# typing ------------------------------------------------------------------- + +from typing import TYPE_CHECKING, Union + +from git.types import Commit_ish + +if TYPE_CHECKING: + from git.repo import Repo + from git.util import IterableList + +# ---------------------------------------------------------------------------- + +_logger = logging.getLogger(__name__) + + +class RootUpdateProgress(UpdateProgress): + """Utility class which adds more opcodes to + :class:`~git.objects.submodule.base.UpdateProgress`.""" + + REMOVE, PATHCHANGE, BRANCHCHANGE, URLCHANGE = [ + 1 << x for x in range(UpdateProgress._num_op_codes, UpdateProgress._num_op_codes + 4) + ] + _num_op_codes = UpdateProgress._num_op_codes + 4 + + __slots__ = () + + +BEGIN = RootUpdateProgress.BEGIN +END = RootUpdateProgress.END +REMOVE = RootUpdateProgress.REMOVE +BRANCHCHANGE = RootUpdateProgress.BRANCHCHANGE +URLCHANGE = RootUpdateProgress.URLCHANGE +PATHCHANGE = RootUpdateProgress.PATHCHANGE + + +class RootModule(Submodule): + """A (virtual) root of all submodules in the given repository. + + This can be used to more easily traverse all submodules of the + superproject (master repository). + """ + + __slots__ = () + + k_root_name = "__ROOT__" + + def __init__(self, repo: "Repo") -> None: + # repo, binsha, mode=None, path=None, name = None, parent_commit=None, url=None, ref=None) + super().__init__( + repo, + binsha=self.NULL_BIN_SHA, + mode=self.k_default_mode, + path="", + name=self.k_root_name, + parent_commit=repo.head.commit, + url="", + branch_path=git.Head.to_full_path(self.k_head_default), + ) + + def _clear_cache(self) -> None: + """May not do anything.""" + pass + + # { Interface + + def update( # type: ignore[override] + self, + previous_commit: Union[Commit_ish, str, None] = None, + recursive: bool = True, + force_remove: bool = False, + init: bool = True, + to_latest_revision: bool = False, + progress: Union[None, "RootUpdateProgress"] = None, + dry_run: bool = False, + force_reset: bool = False, + keep_going: bool = False, + ) -> "RootModule": + """Update the submodules of this repository to the current HEAD commit. + + This method behaves smartly by determining changes of the path of a submodule's + repository, next to changes to the to-be-checked-out commit or the branch to be + checked out. This works if the submodule's ID does not change. + + Additionally it will detect addition and removal of submodules, which will be + handled gracefully. + + :param previous_commit: + If set to a commit-ish, the commit we should use as the previous commit the + HEAD pointed to before it was set to the commit it points to now. + If ``None``, it defaults to ``HEAD@{1}`` otherwise. + + :param recursive: + If ``True``, the children of submodules will be updated as well using the + same technique. + + :param force_remove: + If submodules have been deleted, they will be forcibly removed. Otherwise + the update may fail if a submodule's repository cannot be deleted as changes + have been made to it. + (See :meth:`Submodule.update ` + for more information.) + + :param init: + If we encounter a new module which would need to be initialized, then do it. + + :param to_latest_revision: + If ``True``, instead of checking out the revision pointed to by this + submodule's sha, the checked out tracking branch will be merged with the + latest remote branch fetched from the repository's origin. + + Unless `force_reset` is specified, a local tracking branch will never be + reset into its past, therefore the remote branch must be in the future for + this to have an effect. + + :param force_reset: + If ``True``, submodules may checkout or reset their branch even if the + repository has pending changes that would be overwritten, or if the local + tracking branch is in the future of the remote tracking branch and would be + reset into its past. + + :param progress: + :class:`RootUpdateProgress` instance, or ``None`` if no progress should be + sent. + + :param dry_run: + If ``True``, operations will not actually be performed. Progress messages + will change accordingly to indicate the WOULD DO state of the operation. + + :param keep_going: + If ``True``, we will ignore but log all errors, and keep going recursively. + Unless `dry_run` is set as well, `keep_going` could cause + subsequent/inherited errors you wouldn't see otherwise. + In conjunction with `dry_run`, this can be useful to anticipate all errors + when updating submodules. + + :return: + self + """ + if self.repo.bare: + raise InvalidGitRepositoryError("Cannot update submodules in bare repositories") + # END handle bare + + if progress is None: + progress = RootUpdateProgress() + # END ensure progress is set + + prefix = "" + if dry_run: + prefix = "DRY-RUN: " + + repo = self.repo + + try: + # SETUP BASE COMMIT + ################### + cur_commit = repo.head.commit + if previous_commit is None: + try: + previous_commit = repo.commit(repo.head.log_entry(-1).oldhexsha) + if previous_commit.binsha == previous_commit.NULL_BIN_SHA: + raise IndexError + # END handle initial commit + except IndexError: + # In new repositories, there is no previous commit. + previous_commit = cur_commit + # END exception handling + else: + previous_commit = repo.commit(previous_commit) # Obtain commit object. + # END handle previous commit + + psms: "IterableList[Submodule]" = self.list_items(repo, parent_commit=previous_commit) + sms: "IterableList[Submodule]" = self.list_items(repo) + spsms = set(psms) + ssms = set(sms) + + # HANDLE REMOVALS + ################### + rrsm = spsms - ssms + len_rrsm = len(rrsm) + + for i, rsm in enumerate(rrsm): + op = REMOVE + if i == 0: + op |= BEGIN + # END handle begin + + # Fake it into thinking its at the current commit to allow deletion + # of previous module. Trigger the cache to be updated before that. + progress.update( + op, + i, + len_rrsm, + prefix + "Removing submodule %r at %s" % (rsm.name, rsm.abspath), + ) + rsm._parent_commit = repo.head.commit + rsm.remove( + configuration=False, + module=True, + force=force_remove, + dry_run=dry_run, + ) + + if i == len_rrsm - 1: + op |= END + # END handle end + progress.update(op, i, len_rrsm, prefix + "Done removing submodule %r" % rsm.name) + # END for each removed submodule + + # HANDLE PATH RENAMES + ##################### + # URL changes + branch changes. + csms = spsms & ssms + len_csms = len(csms) + for i, csm in enumerate(csms): + psm: "Submodule" = psms[csm.name] + sm: "Submodule" = sms[csm.name] + + # PATH CHANGES + ############## + if sm.path != psm.path and psm.module_exists(): + progress.update( + BEGIN | PATHCHANGE, + i, + len_csms, + prefix + "Moving repository of submodule %r from %s to %s" % (sm.name, psm.abspath, sm.abspath), + ) + # Move the module to the new path. + if not dry_run: + psm.move(sm.path, module=True, configuration=False) + # END handle dry_run + progress.update( + END | PATHCHANGE, + i, + len_csms, + prefix + "Done moving repository of submodule %r" % sm.name, + ) + # END handle path changes + + if sm.module_exists(): + # HANDLE URL CHANGE + ################### + if sm.url != psm.url: + # Add the new remote, remove the old one. + # This way, if the url just changes, the commits will not have + # to be re-retrieved. + nn = "__new_origin__" + smm = sm.module() + rmts = smm.remotes + + # Don't do anything if we already have the url we search in + # place. + if len([r for r in rmts if r.url == sm.url]) == 0: + progress.update( + BEGIN | URLCHANGE, + i, + len_csms, + prefix + "Changing url of submodule %r from %s to %s" % (sm.name, psm.url, sm.url), + ) + + if not dry_run: + assert nn not in [r.name for r in rmts] + smr = smm.create_remote(nn, sm.url) + smr.fetch(progress=progress) + + # If we have a tracking branch, it should be available + # in the new remote as well. + if len([r for r in smr.refs if r.remote_head == sm.branch_name]) == 0: + raise ValueError( + "Submodule branch named %r was not available in new submodule remote at %r" + % (sm.branch_name, sm.url) + ) + # END head is not detached + + # Now delete the changed one. + rmt_for_deletion = None + for remote in rmts: + if remote.url == psm.url: + rmt_for_deletion = remote + break + # END if urls match + # END for each remote + + # If we didn't find a matching remote, but have exactly + # one, we can safely use this one. + if rmt_for_deletion is None: + if len(rmts) == 1: + rmt_for_deletion = rmts[0] + else: + # If we have not found any remote with the + # original URL we may not have a name. This is a + # special case, and its okay to fail here. + # Alternatively we could just generate a unique + # name and leave all existing ones in place. + raise InvalidGitRepositoryError( + "Couldn't find original remote-repo at url %r" % psm.url + ) + # END handle one single remote + # END handle check we found a remote + + orig_name = rmt_for_deletion.name + smm.delete_remote(rmt_for_deletion) + # NOTE: Currently we leave tags from the deleted remotes + # as well as separate tracking branches in the possibly + # totally changed repository (someone could have changed + # the url to another project). At some point, one might + # want to clean it up, but the danger is high to remove + # stuff the user has added explicitly. + + # Rename the new remote back to what it was. + smr.rename(orig_name) + + # Early on, we verified that the our current tracking + # branch exists in the remote. Now we have to ensure + # that the sha we point to is still contained in the new + # remote tracking branch. + smsha = sm.binsha + found = False + rref = smr.refs[self.branch_name] + for c in rref.commit.traverse(): + if c.binsha == smsha: + found = True + break + # END traverse all commits in search for sha + # END for each commit + + if not found: + # Adjust our internal binsha to use the one of the + # remote this way, it will be checked out in the + # next step. This will change the submodule relative + # to us, so the user will be able to commit the + # change easily. + _logger.warning( + "Current sha %s was not contained in the tracking\ + branch at the new remote, setting it the the remote's tracking branch", + sm.hexsha, + ) + sm.binsha = rref.commit.binsha + # END reset binsha + + # NOTE: All checkout is performed by the base + # implementation of update. + # END handle dry_run + progress.update( + END | URLCHANGE, + i, + len_csms, + prefix + "Done adjusting url of submodule %r" % (sm.name), + ) + # END skip remote handling if new url already exists in module + # END handle url + + # HANDLE PATH CHANGES + ##################### + if sm.branch_path != psm.branch_path: + # Finally, create a new tracking branch which tracks the new + # remote branch. + progress.update( + BEGIN | BRANCHCHANGE, + i, + len_csms, + prefix + + "Changing branch of submodule %r from %s to %s" + % (sm.name, psm.branch_path, sm.branch_path), + ) + if not dry_run: + smm = sm.module() + smmr = smm.remotes + # As the branch might not exist yet, we will have to fetch + # all remotes to be sure... + for remote in smmr: + remote.fetch(progress=progress) + # END for each remote + + try: + tbr = git.Head.create( + smm, + sm.branch_name, + logmsg="branch: Created from HEAD", + ) + except OSError: + # ...or reuse the existing one. + tbr = git.Head(smm, sm.branch_path) + # END ensure tracking branch exists + + tbr.set_tracking_branch(find_first_remote_branch(smmr, sm.branch_name)) + # NOTE: All head-resetting is done in the base + # implementation of update but we will have to checkout the + # new branch here. As it still points to the currently + # checked out commit, we don't do any harm. + # As we don't want to update working-tree or index, changing + # the ref is all there is to do. + smm.head.reference = tbr + # END handle dry_run + + progress.update( + END | BRANCHCHANGE, + i, + len_csms, + prefix + "Done changing branch of submodule %r" % sm.name, + ) + # END handle branch + # END handle + # END for each common submodule + except Exception as err: + if not keep_going: + raise + _logger.error(str(err)) + # END handle keep_going + + # FINALLY UPDATE ALL ACTUAL SUBMODULES + ###################################### + for sm in sms: + # Update the submodule using the default method. + sm.update( + recursive=False, + init=init, + to_latest_revision=to_latest_revision, + progress=progress, + dry_run=dry_run, + force=force_reset, + keep_going=keep_going, + ) + + # Update recursively depth first - question is which inconsistent state will + # be better in case it fails somewhere. Defective branch or defective depth. + # The RootSubmodule type will never process itself, which was done in the + # previous expression. + if recursive: + # The module would exist by now if we are not in dry_run mode. + if sm.module_exists(): + type(self)(sm.module()).update( + recursive=True, + force_remove=force_remove, + init=init, + to_latest_revision=to_latest_revision, + progress=progress, + dry_run=dry_run, + force_reset=force_reset, + keep_going=keep_going, + ) + # END handle dry_run + # END handle recursive + # END for each submodule to update + + return self + + def module(self) -> "Repo": + """:return: The actual repository containing the submodules""" + return self.repo + + # } END interface + + +# } END classes diff --git a/lib/python3.12/site-packages/git/objects/submodule/util.py b/lib/python3.12/site-packages/git/objects/submodule/util.py new file mode 100644 index 0000000000000000000000000000000000000000..c021510d88a47a267f530ce4455b682fe8dee0a6 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/submodule/util.py @@ -0,0 +1,121 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = [ + "sm_section", + "sm_name", + "mkhead", + "find_first_remote_branch", + "SubmoduleConfigParser", +] + +from io import BytesIO +import weakref + +import git +from git.config import GitConfigParser +from git.exc import InvalidGitRepositoryError + +# typing ----------------------------------------------------------------------- + +from typing import Any, Sequence, TYPE_CHECKING, Union + +from git.types import PathLike + +if TYPE_CHECKING: + from weakref import ReferenceType + + from git.refs import Head, RemoteReference + from git.remote import Remote + from git.repo import Repo + + from .base import Submodule + +# { Utilities + + +def sm_section(name: str) -> str: + """:return: Section title used in ``.gitmodules`` configuration file""" + return f'submodule "{name}"' + + +def sm_name(section: str) -> str: + """:return: Name of the submodule as parsed from the section name""" + section = section.strip() + return section[11:-1] + + +def mkhead(repo: "Repo", path: PathLike) -> "Head": + """:return: New branch/head instance""" + return git.Head(repo, git.Head.to_full_path(path)) + + +def find_first_remote_branch(remotes: Sequence["Remote"], branch_name: str) -> "RemoteReference": + """Find the remote branch matching the name of the given branch or raise + :exc:`~git.exc.InvalidGitRepositoryError`.""" + for remote in remotes: + try: + return remote.refs[branch_name] + except IndexError: + continue + # END exception handling + # END for remote + raise InvalidGitRepositoryError("Didn't find remote branch '%r' in any of the given remotes" % branch_name) + + +# } END utilities + +# { Classes + + +class SubmoduleConfigParser(GitConfigParser): + """Catches calls to :meth:`~git.config.GitConfigParser.write`, and updates the + ``.gitmodules`` blob in the index with the new data, if we have written into a + stream. + + Otherwise it would add the local file to the index to make it correspond with the + working tree. Additionally, the cache must be cleared. + + Please note that no mutating method will work in bare mode. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + self._smref: Union["ReferenceType[Submodule]", None] = None + self._index = None + self._auto_write = True + super().__init__(*args, **kwargs) + + # { Interface + def set_submodule(self, submodule: "Submodule") -> None: + """Set this instance's submodule. It must be called before the first write + operation begins.""" + self._smref = weakref.ref(submodule) + + def flush_to_index(self) -> None: + """Flush changes in our configuration file to the index.""" + assert self._smref is not None + # Should always have a file here. + assert not isinstance(self._file_or_files, BytesIO) + + sm = self._smref() + if sm is not None: + index = self._index + if index is None: + index = sm.repo.index + # END handle index + index.add([sm.k_modules_file], write=self._auto_write) + sm._clear_cache() + # END handle weakref + + # } END interface + + # { Overridden Methods + def write(self) -> None: # type: ignore[override] + rval: None = super().write() + self.flush_to_index() + return rval + + # END overridden methods + + +# } END classes diff --git a/lib/python3.12/site-packages/git/objects/tag.py b/lib/python3.12/site-packages/git/objects/tag.py new file mode 100644 index 0000000000000000000000000000000000000000..88671d3165a206500fd9574437cd4c21a25cb5f6 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/tag.py @@ -0,0 +1,140 @@ +# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors +# +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Provides an :class:`~git.objects.base.Object`-based type for annotated tags. + +This defines the :class:`TagObject` class, which represents annotated tags. +For lightweight tags, see the :mod:`git.refs.tag` module. +""" + +__all__ = ["TagObject"] + +import sys + +from git.compat import defenc +from git.util import Actor, hex_to_bin + +from . import base +from .util import get_object_type_by_name, parse_actor_and_date + +# typing ---------------------------------------------- + +from typing import List, TYPE_CHECKING, Union + +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +if TYPE_CHECKING: + from git.repo import Repo + + from .blob import Blob + from .commit import Commit + from .tree import Tree + +# --------------------------------------------------- + + +class TagObject(base.Object): + """Annotated (i.e. non-lightweight) tag carrying additional information about an + object we are pointing to. + + See :manpage:`gitglossary(7)` on "tag object": + https://git-scm.com/docs/gitglossary#def_tag_object + """ + + type: Literal["tag"] = "tag" + + __slots__ = ( + "object", + "tag", + "tagger", + "tagged_date", + "tagger_tz_offset", + "message", + ) + + def __init__( + self, + repo: "Repo", + binsha: bytes, + object: Union[None, base.Object] = None, + tag: Union[None, str] = None, + tagger: Union[None, Actor] = None, + tagged_date: Union[int, None] = None, + tagger_tz_offset: Union[int, None] = None, + message: Union[str, None] = None, + ) -> None: # @ReservedAssignment + """Initialize a tag object with additional data. + + :param repo: + Repository this object is located in. + + :param binsha: + 20 byte SHA1. + + :param object: + :class:`~git.objects.base.Object` instance of object we are pointing to. + + :param tag: + Name of this tag. + + :param tagger: + :class:`~git.util.Actor` identifying the tagger. + + :param tagged_date: int_seconds_since_epoch + The DateTime of the tag creation. + Use :func:`time.gmtime` to convert it into a different format. + + :param tagger_tz_offset: int_seconds_west_of_utc + The timezone that the `tagged_date` is in, in a format similar to + :attr:`time.altzone`. + """ + super().__init__(repo, binsha) + if object is not None: + self.object: Union["Commit", "Blob", "Tree", "TagObject"] = object + if tag is not None: + self.tag = tag + if tagger is not None: + self.tagger = tagger + if tagged_date is not None: + self.tagged_date = tagged_date + if tagger_tz_offset is not None: + self.tagger_tz_offset = tagger_tz_offset + if message is not None: + self.message = message + + def _set_cache_(self, attr: str) -> None: + """Cache all our attributes at once.""" + if attr in TagObject.__slots__: + ostream = self.repo.odb.stream(self.binsha) + lines: List[str] = ostream.read().decode(defenc, "replace").splitlines() + + _obj, hexsha = lines[0].split(" ") + _type_token, type_name = lines[1].split(" ") + object_type = get_object_type_by_name(type_name.encode("ascii")) + self.object = object_type(self.repo, hex_to_bin(hexsha)) + + self.tag = lines[2][4:] # tag + + if len(lines) > 3: + tagger_info = lines[3] # tagger + ( + self.tagger, + self.tagged_date, + self.tagger_tz_offset, + ) = parse_actor_and_date(tagger_info) + + # Line 4 empty - it could mark the beginning of the next header. + # In case there really is no message, it would not exist. + # Otherwise a newline separates header from message. + if len(lines) > 5: + self.message = "\n".join(lines[5:]) + else: + self.message = "" + # END check our attributes + else: + super()._set_cache_(attr) diff --git a/lib/python3.12/site-packages/git/objects/tree.py b/lib/python3.12/site-packages/git/objects/tree.py new file mode 100644 index 0000000000000000000000000000000000000000..a3d611c8003500b31ddfb066cbdfa74d223d64f9 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/tree.py @@ -0,0 +1,418 @@ +# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors +# +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["TreeModifier", "Tree"] + +import os +import sys + +import git.diff as git_diff +from git.util import IterableList, join_path, to_bin_sha + +from . import util +from .base import IndexObjUnion, IndexObject +from .blob import Blob +from .fun import tree_entries_from_data, tree_to_stream +from .submodule.base import Submodule + +# typing ------------------------------------------------- + +from typing import ( + Any, + Callable, + Dict, + Iterable, + Iterator, + List, + Tuple, + TYPE_CHECKING, + Type, + Union, + cast, +) + +if sys.version_info >= (3, 8): + from typing import Literal +else: + from typing_extensions import Literal + +from git.types import PathLike + +if TYPE_CHECKING: + from io import BytesIO + + from git.repo import Repo + +TreeCacheTup = Tuple[bytes, int, str] + +TraversedTreeTup = Union[Tuple[Union["Tree", None], IndexObjUnion, Tuple["Submodule", "Submodule"]]] + +# -------------------------------------------------------- + + +def cmp(a: str, b: str) -> int: + return (a > b) - (a < b) + + +class TreeModifier: + """A utility class providing methods to alter the underlying cache in a list-like + fashion. + + Once all adjustments are complete, the :attr:`_cache`, which really is a reference + to the cache of a tree, will be sorted. This ensures it will be in a serializable + state. + """ + + __slots__ = ("_cache",) + + def __init__(self, cache: List[TreeCacheTup]) -> None: + self._cache = cache + + def _index_by_name(self, name: str) -> int: + """:return: index of an item with name, or -1 if not found""" + for i, t in enumerate(self._cache): + if t[2] == name: + return i + # END found item + # END for each item in cache + return -1 + + # { Interface + def set_done(self) -> "TreeModifier": + """Call this method once you are done modifying the tree information. + + This may be called several times, but be aware that each call will cause a sort + operation. + + :return: + self + """ + self._cache.sort(key=lambda x: (x[2] + "/") if x[1] == Tree.tree_id << 12 else x[2]) + return self + + # } END interface + + # { Mutators + def add(self, sha: bytes, mode: int, name: str, force: bool = False) -> "TreeModifier": + """Add the given item to the tree. + + If an item with the given name already exists, nothing will be done, but a + :exc:`ValueError` will be raised if the sha and mode of the existing item do not + match the one you add, unless `force` is ``True``. + + :param sha: + The 20 or 40 byte sha of the item to add. + + :param mode: + :class:`int` representing the stat-compatible mode of the item. + + :param force: + If ``True``, an item with your name and information will overwrite any + existing item with the same name, no matter which information it has. + + :return: + self + """ + if "/" in name: + raise ValueError("Name must not contain '/' characters") + if (mode >> 12) not in Tree._map_id_to_type: + raise ValueError("Invalid object type according to mode %o" % mode) + + sha = to_bin_sha(sha) + index = self._index_by_name(name) + + item = (sha, mode, name) + + if index == -1: + self._cache.append(item) + else: + if force: + self._cache[index] = item + else: + ex_item = self._cache[index] + if ex_item[0] != sha or ex_item[1] != mode: + raise ValueError("Item %r existed with different properties" % name) + # END handle mismatch + # END handle force + # END handle name exists + return self + + def add_unchecked(self, binsha: bytes, mode: int, name: str) -> None: + """Add the given item to the tree. Its correctness is assumed, so it is the + caller's responsibility to ensure that the input is correct. + + For more information on the parameters, see :meth:`add`. + + :param binsha: + 20 byte binary sha. + """ + assert isinstance(binsha, bytes) and isinstance(mode, int) and isinstance(name, str) + tree_cache = (binsha, mode, name) + + self._cache.append(tree_cache) + + def __delitem__(self, name: str) -> None: + """Delete an item with the given name if it exists.""" + index = self._index_by_name(name) + if index > -1: + del self._cache[index] + + # } END mutators + + +class Tree(IndexObject, git_diff.Diffable, util.Traversable, util.Serializable): + R"""Tree objects represent an ordered list of :class:`~git.objects.blob.Blob`\s and + other :class:`Tree`\s. + + See :manpage:`gitglossary(7)` on "tree object": + https://git-scm.com/docs/gitglossary#def_tree_object + + Subscripting is supported, as with a list or dict: + + * Access a specific blob using the ``tree["filename"]`` notation. + * You may likewise access by index, like ``blob = tree[0]``. + """ + + type: Literal["tree"] = "tree" + + __slots__ = ("_cache",) + + # Actual integer IDs for comparison. + commit_id = 0o16 # Equals stat.S_IFDIR | stat.S_IFLNK - a directory link. + blob_id = 0o10 + symlink_id = 0o12 + tree_id = 0o04 + + _map_id_to_type: Dict[int, Type[IndexObjUnion]] = { + commit_id: Submodule, + blob_id: Blob, + symlink_id: Blob, + # Tree ID added once Tree is defined. + } + + def __init__( + self, + repo: "Repo", + binsha: bytes, + mode: int = tree_id << 12, + path: Union[PathLike, None] = None, + ): + super().__init__(repo, binsha, mode, path) + + @classmethod + def _get_intermediate_items( + cls, + index_object: IndexObjUnion, + ) -> Union[Tuple["Tree", ...], Tuple[()]]: + if index_object.type == "tree": + return tuple(index_object._iter_convert_to_object(index_object._cache)) + return () + + def _set_cache_(self, attr: str) -> None: + if attr == "_cache": + # Set the data when we need it. + ostream = self.repo.odb.stream(self.binsha) + self._cache: List[TreeCacheTup] = tree_entries_from_data(ostream.read()) + else: + super()._set_cache_(attr) + # END handle attribute + + def _iter_convert_to_object(self, iterable: Iterable[TreeCacheTup]) -> Iterator[IndexObjUnion]: + """Iterable yields tuples of (binsha, mode, name), which will be converted to + the respective object representation. + """ + for binsha, mode, name in iterable: + path = join_path(self.path, name) + try: + yield self._map_id_to_type[mode >> 12](self.repo, binsha, mode, path) + except KeyError as e: + raise TypeError("Unknown mode %o found in tree data for path '%s'" % (mode, path)) from e + # END for each item + + def join(self, file: PathLike) -> IndexObjUnion: + """Find the named object in this tree's contents. + + :return: + :class:`~git.objects.blob.Blob`, :class:`Tree`, or + :class:`~git.objects.submodule.base.Submodule` + + :raise KeyError: + If the given file or tree does not exist in this tree. + """ + msg = "Blob or Tree named %r not found" + file = os.fspath(file) + if "/" in file: + tree = self + item = self + tokens = file.split("/") + for i, token in enumerate(tokens): + item = tree[token] + if item.type == "tree": + tree = item + else: + # Safety assertion - blobs are at the end of the path. + if i != len(tokens) - 1: + raise KeyError(msg % file) + return item + # END handle item type + # END for each token of split path + if item == self: + raise KeyError(msg % file) + return item + else: + for info in self._cache: + if info[2] == file: # [2] == name + return self._map_id_to_type[info[1] >> 12]( + self.repo, info[0], info[1], join_path(self.path, info[2]) + ) + # END for each obj + raise KeyError(msg % file) + # END handle long paths + + def __truediv__(self, file: PathLike) -> IndexObjUnion: + """The ``/`` operator is another syntax for joining. + + See :meth:`join` for details. + """ + return self.join(file) + + @property + def trees(self) -> List["Tree"]: + """:return: list(Tree, ...) List of trees directly below this tree""" + return [i for i in self if i.type == "tree"] + + @property + def blobs(self) -> List[Blob]: + """:return: list(Blob, ...) List of blobs directly below this tree""" + return [i for i in self if i.type == "blob"] + + @property + def cache(self) -> TreeModifier: + """ + :return: + An object allowing modification of the internal cache. This can be used to + change the tree's contents. When done, make sure you call + :meth:`~TreeModifier.set_done` on the tree modifier, or serialization + behaviour will be incorrect. + + :note: + See :class:`TreeModifier` for more information on how to alter the cache. + """ + return TreeModifier(self._cache) + + def traverse( + self, + predicate: Callable[[Union[IndexObjUnion, TraversedTreeTup], int], bool] = lambda i, d: True, + prune: Callable[[Union[IndexObjUnion, TraversedTreeTup], int], bool] = lambda i, d: False, + depth: int = -1, + branch_first: bool = True, + visit_once: bool = False, + ignore_self: int = 1, + as_edge: bool = False, + ) -> Union[Iterator[IndexObjUnion], Iterator[TraversedTreeTup]]: + """For documentation, see + `Traversable._traverse() `. + + Trees are set to ``visit_once = False`` to gain more performance in the + traversal. + """ + + # # To typecheck instead of using cast. + # import itertools + # def is_tree_traversed(inp: Tuple) -> TypeGuard[Tuple[Iterator[Union['Tree', 'Blob', 'Submodule']]]]: + # return all(isinstance(x, (Blob, Tree, Submodule)) for x in inp[1]) + + # ret = super().traverse(predicate, prune, depth, branch_first, visit_once, ignore_self) + # ret_tup = itertools.tee(ret, 2) + # assert is_tree_traversed(ret_tup), f"Type is {[type(x) for x in list(ret_tup[0])]}" + # return ret_tup[0] + + return cast( + Union[Iterator[IndexObjUnion], Iterator[TraversedTreeTup]], + super()._traverse( + predicate, # type: ignore[arg-type] + prune, # type: ignore[arg-type] + depth, + branch_first, + visit_once, + ignore_self, + ), + ) + + def list_traverse(self, *args: Any, **kwargs: Any) -> IterableList[IndexObjUnion]: + """ + :return: + :class:`~git.util.IterableList` with the results of the traversal as + produced by :meth:`traverse` + + Tree -> IterableList[Union[Submodule, Tree, Blob]] + """ + return super()._list_traverse(*args, **kwargs) + + # List protocol + + def __getslice__(self, i: int, j: int) -> List[IndexObjUnion]: + return list(self._iter_convert_to_object(self._cache[i:j])) + + def __iter__(self) -> Iterator[IndexObjUnion]: + return self._iter_convert_to_object(self._cache) + + def __len__(self) -> int: + return len(self._cache) + + def __getitem__(self, item: Union[str, int, slice]) -> IndexObjUnion: + if isinstance(item, int): + info = self._cache[item] + return self._map_id_to_type[info[1] >> 12](self.repo, info[0], info[1], join_path(self.path, info[2])) + + if isinstance(item, str): + # compatibility + return self.join(item) + # END index is basestring + + raise TypeError("Invalid index type: %r" % item) + + def __contains__(self, item: Union[IndexObjUnion, PathLike]) -> bool: + if isinstance(item, IndexObject): + for info in self._cache: + if item.binsha == info[0]: + return True + # END compare sha + # END for each entry + # END handle item is index object + # compatibility + + # Treat item as repo-relative path. + else: + path = self.path + for info in self._cache: + if item == join_path(path, info[2]): + return True + # END for each item + return False + + def __reversed__(self) -> Iterator[IndexObjUnion]: + return reversed(self._iter_convert_to_object(self._cache)) # type: ignore[call-overload] + + def _serialize(self, stream: "BytesIO") -> "Tree": + """Serialize this tree into the stream. Assumes sorted tree data. + + :note: + We will assume our tree data to be in a sorted state. If this is not the + case, serialization will not generate a correct tree representation as these + are assumed to be sorted by algorithms. + """ + tree_to_stream(self._cache, stream.write) + return self + + def _deserialize(self, stream: "BytesIO") -> "Tree": + self._cache = tree_entries_from_data(stream.read()) + return self + + +# END tree + +# Finalize map definition. +Tree._map_id_to_type[Tree.tree_id] = Tree diff --git a/lib/python3.12/site-packages/git/objects/util.py b/lib/python3.12/site-packages/git/objects/util.py new file mode 100644 index 0000000000000000000000000000000000000000..a68d701f5142534993bca5e466ed993912d574d0 --- /dev/null +++ b/lib/python3.12/site-packages/git/objects/util.py @@ -0,0 +1,700 @@ +# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors +# +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Utility functions for working with git objects.""" + +__all__ = [ + "get_object_type_by_name", + "parse_date", + "parse_actor_and_date", + "ProcessStreamAdapter", + "Traversable", + "altz_to_utctz_str", + "utctz_to_altz", + "verify_utctz", + "Actor", + "tzoffset", + "utc", +] + +from abc import ABC, abstractmethod +import calendar +from collections import deque +from datetime import datetime, timedelta, tzinfo +import re +from string import digits +import time +import warnings + +from git.util import Actor, IterableList, IterableObj + +# typing ------------------------------------------------------------ + +from typing import ( + Any, + Callable, + Deque, + Iterator, + NamedTuple, + Sequence, + TYPE_CHECKING, + Tuple, + Type, + TypeVar, + Union, + cast, + overload, +) + +from git.types import Has_id_attribute, Literal + +if TYPE_CHECKING: + from io import BytesIO, StringIO + from subprocess import Popen + + from git.types import Protocol, runtime_checkable + + from .blob import Blob + from .commit import Commit + from .submodule.base import Submodule + from .tag import TagObject + from .tree import TraversedTreeTup, Tree +else: + Protocol = ABC + + def runtime_checkable(f): + return f + + +class TraverseNT(NamedTuple): + depth: int + item: Union["Traversable", "Blob"] + src: Union["Traversable", None] + + +T_TIobj = TypeVar("T_TIobj", bound="TraversableIterableObj") # For TraversableIterableObj.traverse() + +TraversedTup = Union[ + Tuple[Union["Traversable", None], "Traversable"], # For Commit, Submodule. + "TraversedTreeTup", # For Tree.traverse(). +] + +# -------------------------------------------------------------------- + +ZERO = timedelta(0) + +# { Functions + + +def mode_str_to_int(modestr: Union[bytes, str]) -> int: + """Convert mode bits from an octal mode string to an integer mode for git. + + :param modestr: + String like ``755`` or ``644`` or ``100644`` - only the last 6 chars will be + used. + + :return: + String identifying a mode compatible to the mode methods ids of the :mod:`stat` + module regarding the rwx permissions for user, group and other, special flags + and file system flags, such as whether it is a symlink. + """ + mode = 0 + for iteration, char in enumerate(reversed(modestr[-6:])): + char = cast(Union[str, int], char) + mode += int(char) << iteration * 3 + # END for each char + return mode + + +def get_object_type_by_name( + object_type_name: bytes, +) -> Union[Type["Commit"], Type["TagObject"], Type["Tree"], Type["Blob"]]: + """Retrieve the Python class GitPython uses to represent a kind of Git object. + + :return: + A type suitable to handle the given as `object_type_name`. + This type can be called create new instances. + + :param object_type_name: + Member of :attr:`Object.TYPES `. + + :raise ValueError: + If `object_type_name` is unknown. + """ + if object_type_name == b"commit": + from . import commit + + return commit.Commit + elif object_type_name == b"tag": + from . import tag + + return tag.TagObject + elif object_type_name == b"blob": + from . import blob + + return blob.Blob + elif object_type_name == b"tree": + from . import tree + + return tree.Tree + else: + raise ValueError("Cannot handle unknown object type: %s" % object_type_name.decode()) + + +def utctz_to_altz(utctz: str) -> int: + """Convert a git timezone offset into a timezone offset west of UTC in seconds + (compatible with :attr:`time.altzone`). + + :param utctz: + git utc timezone string, e.g. +0200 + """ + int_utctz = int(utctz) + seconds = (abs(int_utctz) // 100) * 3600 + (abs(int_utctz) % 100) * 60 + return seconds if int_utctz < 0 else -seconds + + +def altz_to_utctz_str(altz: float) -> str: + """Convert a timezone offset west of UTC in seconds into a Git timezone offset + string. + + :param altz: + Timezone offset in seconds west of UTC. + """ + hours = abs(altz) // 3600 + minutes = (abs(altz) % 3600) // 60 + sign = "-" if altz >= 60 else "+" + return "{}{:02}{:02}".format(sign, hours, minutes) + + +def verify_utctz(offset: str) -> str: + """ + :raise ValueError: + If `offset` is incorrect. + + :return: + `offset` + """ + fmt_exc = ValueError("Invalid timezone offset format: %s" % offset) + if len(offset) != 5: + raise fmt_exc + if offset[0] not in "+-": + raise fmt_exc + if offset[1] not in digits or offset[2] not in digits or offset[3] not in digits or offset[4] not in digits: + raise fmt_exc + # END for each char + return offset + + +class tzoffset(tzinfo): + def __init__(self, secs_west_of_utc: float, name: Union[None, str] = None) -> None: + self._offset = timedelta(seconds=-secs_west_of_utc) + self._name = name or "fixed" + + def __reduce__(self) -> Tuple[Type["tzoffset"], Tuple[float, str]]: + return tzoffset, (-self._offset.total_seconds(), self._name) + + def utcoffset(self, dt: Union[datetime, None]) -> timedelta: + return self._offset + + def tzname(self, dt: Union[datetime, None]) -> str: + return self._name + + def dst(self, dt: Union[datetime, None]) -> timedelta: + return ZERO + + +utc = tzoffset(0, "UTC") + + +def from_timestamp(timestamp: float, tz_offset: float) -> datetime: + """Convert a `timestamp` + `tz_offset` into an aware :class:`~datetime.datetime` + instance.""" + utc_dt = datetime.fromtimestamp(timestamp, utc) + try: + local_dt = utc_dt.astimezone(tzoffset(tz_offset)) + return local_dt + except ValueError: + return utc_dt + + +def parse_date(string_date: Union[str, datetime]) -> Tuple[int, int]: + """Parse the given date as one of the following: + + * Aware datetime instance + * Git internal format: timestamp offset + * :rfc:`2822`: ``Thu, 07 Apr 2005 22:13:13 +0200`` + * ISO 8601: ``2005-04-07T22:13:13`` - The ``T`` can be a space as well. + + :return: + Tuple(int(timestamp_UTC), int(offset)), both in seconds since epoch + + :raise ValueError: + If the format could not be understood. + + :note: + Date can also be ``YYYY.MM.DD``, ``MM/DD/YYYY`` and ``DD.MM.YYYY``. + """ + if isinstance(string_date, datetime): + if string_date.tzinfo: + utcoffset = cast(timedelta, string_date.utcoffset()) # typeguard, if tzinfoand is not None + offset = -int(utcoffset.total_seconds()) + return int(string_date.astimezone(utc).timestamp()), offset + else: + raise ValueError(f"string_date datetime object without tzinfo, {string_date}") + + # Git time + try: + if string_date.count(" ") == 1 and string_date.rfind(":") == -1: + timestamp, offset_str = string_date.split() + if timestamp.startswith("@"): + timestamp = timestamp[1:] + timestamp_int = int(timestamp) + return timestamp_int, utctz_to_altz(verify_utctz(offset_str)) + else: + offset_str = "+0000" # Local time by default. + if string_date[-5] in "-+": + offset_str = verify_utctz(string_date[-5:]) + string_date = string_date[:-6] # skip space as well + # END split timezone info + offset = utctz_to_altz(offset_str) + + # Now figure out the date and time portion - split time. + date_formats = [] + splitter = -1 + if "," in string_date: + date_formats.append("%a, %d %b %Y") + splitter = string_date.rfind(" ") + else: + # ISO plus additional + date_formats.append("%Y-%m-%d") + date_formats.append("%Y.%m.%d") + date_formats.append("%m/%d/%Y") + date_formats.append("%d.%m.%Y") + + splitter = string_date.rfind("T") + if splitter == -1: + splitter = string_date.rfind(" ") + # END handle 'T' and ' ' + # END handle RFC or ISO + + assert splitter > -1 + + # Split date and time. + time_part = string_date[splitter + 1 :] # Skip space. + date_part = string_date[:splitter] + + # Parse time. + tstruct = time.strptime(time_part, "%H:%M:%S") + + for fmt in date_formats: + try: + dtstruct = time.strptime(date_part, fmt) + utctime = calendar.timegm( + ( + dtstruct.tm_year, + dtstruct.tm_mon, + dtstruct.tm_mday, + tstruct.tm_hour, + tstruct.tm_min, + tstruct.tm_sec, + dtstruct.tm_wday, + dtstruct.tm_yday, + tstruct.tm_isdst, + ) + ) + return int(utctime), offset + except ValueError: + continue + # END exception handling + # END for each fmt + + # Still here ? fail. + raise ValueError("no format matched") + # END handle format + except Exception as e: + raise ValueError(f"Unsupported date format or type: {string_date}, type={type(string_date)}") from e + # END handle exceptions + + +# Precompiled regexes +_re_actor_epoch = re.compile(r"^.+? (.*) (\d+) ([+-]\d+).*$") +_re_only_actor = re.compile(r"^.+? (.*)$") + + +def parse_actor_and_date(line: str) -> Tuple[Actor, int, int]: + """Parse out the actor (author or committer) info from a line like:: + + author Tom Preston-Werner 1191999972 -0700 + + :return: + [Actor, int_seconds_since_epoch, int_timezone_offset] + """ + actor, epoch, offset = "", "0", "0" + m = _re_actor_epoch.search(line) + if m: + actor, epoch, offset = m.groups() + else: + m = _re_only_actor.search(line) + actor = m.group(1) if m else line or "" + return (Actor._from_string(actor), int(epoch), utctz_to_altz(offset)) + + +# } END functions + + +# { Classes + + +class ProcessStreamAdapter: + """Class wiring all calls to the contained Process instance. + + Use this type to hide the underlying process to provide access only to a specified + stream. The process is usually wrapped into an :class:`~git.cmd.Git.AutoInterrupt` + class to kill it if the instance goes out of scope. + """ + + __slots__ = ("_proc", "_stream") + + def __init__(self, process: "Popen", stream_name: str) -> None: + self._proc = process + self._stream: StringIO = getattr(process, stream_name) # guessed type + + def __getattr__(self, attr: str) -> Any: + return getattr(self._stream, attr) + + +@runtime_checkable +class Traversable(Protocol): + """Simple interface to perform depth-first or breadth-first traversals in one + direction. + + Subclasses only need to implement one function. + + Instances of the subclass must be hashable. + + Defined subclasses: + + * :class:`Commit ` + * :class:`Tree ` + * :class:`Submodule ` + """ + + __slots__ = () + + @classmethod + @abstractmethod + def _get_intermediate_items(cls, item: Any) -> Sequence["Traversable"]: + """ + :return: + Tuple of items connected to the given item. + Must be implemented in subclass. + + class Commit:: (cls, Commit) -> Tuple[Commit, ...] + class Submodule:: (cls, Submodule) -> Iterablelist[Submodule] + class Tree:: (cls, Tree) -> Tuple[Tree, ...] + """ + raise NotImplementedError("To be implemented in subclass") + + @abstractmethod + def list_traverse(self, *args: Any, **kwargs: Any) -> Any: + """Traverse self and collect all items found. + + Calling this directly on the abstract base class, including via a ``super()`` + proxy, is deprecated. Only overridden implementations should be called. + """ + warnings.warn( + "list_traverse() method should only be called from subclasses." + " Calling from Traversable abstract class will raise NotImplementedError in 4.0.0." + " The concrete subclasses in GitPython itself are 'Commit', 'RootModule', 'Submodule', and 'Tree'.", + DeprecationWarning, + stacklevel=2, + ) + return self._list_traverse(*args, **kwargs) + + def _list_traverse( + self, as_edge: bool = False, *args: Any, **kwargs: Any + ) -> IterableList[Union["Commit", "Submodule", "Tree", "Blob"]]: + """Traverse self and collect all items found. + + :return: + :class:`~git.util.IterableList` with the results of the traversal as + produced by :meth:`traverse`:: + + Commit -> IterableList[Commit] + Submodule -> IterableList[Submodule] + Tree -> IterableList[Union[Submodule, Tree, Blob]] + """ + # Commit and Submodule have id.__attribute__ as IterableObj. + # Tree has id.__attribute__ inherited from IndexObject. + if isinstance(self, Has_id_attribute): + id = self._id_attribute_ + else: + # Shouldn't reach here, unless Traversable subclass created with no + # _id_attribute_. + id = "" + # Could add _id_attribute_ to Traversable, or make all Traversable also + # Iterable? + + if not as_edge: + out: IterableList[Union["Commit", "Submodule", "Tree", "Blob"]] = IterableList(id) + out.extend(self.traverse(as_edge=as_edge, *args, **kwargs)) # noqa: B026 + return out + # Overloads in subclasses (mypy doesn't allow typing self: subclass). + # Union[IterableList['Commit'], IterableList['Submodule'], IterableList[Union['Submodule', 'Tree', 'Blob']]] + else: + # Raise DeprecationWarning, it doesn't make sense to use this. + out_list: IterableList = IterableList(self.traverse(*args, **kwargs)) + return out_list + + @abstractmethod + def traverse(self, *args: Any, **kwargs: Any) -> Any: + """Iterator yielding items found when traversing self. + + Calling this directly on the abstract base class, including via a ``super()`` + proxy, is deprecated. Only overridden implementations should be called. + """ + warnings.warn( + "traverse() method should only be called from subclasses." + " Calling from Traversable abstract class will raise NotImplementedError in 4.0.0." + " The concrete subclasses in GitPython itself are 'Commit', 'RootModule', 'Submodule', and 'Tree'.", + DeprecationWarning, + stacklevel=2, + ) + return self._traverse(*args, **kwargs) + + def _traverse( + self, + predicate: Callable[[Union["Traversable", "Blob", TraversedTup], int], bool] = lambda i, d: True, + prune: Callable[[Union["Traversable", "Blob", TraversedTup], int], bool] = lambda i, d: False, + depth: int = -1, + branch_first: bool = True, + visit_once: bool = True, + ignore_self: int = 1, + as_edge: bool = False, + ) -> Union[Iterator[Union["Traversable", "Blob"]], Iterator[TraversedTup]]: + """Iterator yielding items found when traversing `self`. + + :param predicate: + A function ``f(i,d)`` that returns ``False`` if item i at depth ``d`` should + not be included in the result. + + :param prune: + A function ``f(i,d)`` that returns ``True`` if the search should stop at + item ``i`` at depth ``d``. Item ``i`` will not be returned. + + :param depth: + Defines at which level the iteration should not go deeper if -1. There is no + limit if 0, you would effectively only get `self`, the root of the + iteration. If 1, you would only get the first level of + predecessors/successors. + + :param branch_first: + If ``True``, items will be returned branch first, otherwise depth first. + + :param visit_once: + If ``True``, items will only be returned once, although they might be + encountered several times. Loops are prevented that way. + + :param ignore_self: + If ``True``, `self` will be ignored and automatically pruned from the + result. Otherwise it will be the first item to be returned. If `as_edge` is + ``True``, the source of the first edge is ``None``. + + :param as_edge: + If ``True``, return a pair of items, first being the source, second the + destination, i.e. tuple(src, dest) with the edge spanning from source to + destination. + + :return: + Iterator yielding items found when traversing `self`:: + + Commit -> Iterator[Union[Commit, Tuple[Commit, Commit]] Submodule -> + Iterator[Submodule, Tuple[Submodule, Submodule]] Tree -> + Iterator[Union[Blob, Tree, Submodule, + Tuple[Union[Submodule, Tree], Union[Blob, Tree, + Submodule]]] + + ignore_self=True is_edge=True -> Iterator[item] ignore_self=True + is_edge=False --> Iterator[item] ignore_self=False is_edge=True -> + Iterator[item] | Iterator[Tuple[src, item]] ignore_self=False + is_edge=False -> Iterator[Tuple[src, item]] + """ + + visited = set() + stack: Deque[TraverseNT] = deque() + stack.append(TraverseNT(0, self, None)) # self is always depth level 0. + + def addToStack( + stack: Deque[TraverseNT], + src_item: "Traversable", + branch_first: bool, + depth: int, + ) -> None: + lst = self._get_intermediate_items(item) + if not lst: # Empty list + return + if branch_first: + stack.extendleft(TraverseNT(depth, i, src_item) for i in lst) + else: + reviter = (TraverseNT(depth, lst[i], src_item) for i in range(len(lst) - 1, -1, -1)) + stack.extend(reviter) + + # END addToStack local method + + while stack: + d, item, src = stack.pop() # Depth of item, item, item_source + + if visit_once and item in visited: + continue + + if visit_once: + visited.add(item) + + rval: Union[TraversedTup, "Traversable", "Blob"] + if as_edge: + # If as_edge return (src, item) unless rrc is None + # (e.g. for first item). + rval = (src, item) + else: + rval = item + + if prune(rval, d): + continue + + skipStartItem = ignore_self and (item is self) + if not skipStartItem and predicate(rval, d): + yield rval + + # Only continue to next level if this is appropriate! + next_d = d + 1 + if depth > -1 and next_d > depth: + continue + + addToStack(stack, item, branch_first, next_d) + # END for each item on work stack + + +@runtime_checkable +class Serializable(Protocol): + """Defines methods to serialize and deserialize objects from and into a data + stream.""" + + __slots__ = () + + # @abstractmethod + def _serialize(self, stream: "BytesIO") -> "Serializable": + """Serialize the data of this object into the given data stream. + + :note: + A serialized object would :meth:`_deserialize` into the same object. + + :param stream: + A file-like object. + + :return: + self + """ + raise NotImplementedError("To be implemented in subclass") + + # @abstractmethod + def _deserialize(self, stream: "BytesIO") -> "Serializable": + """Deserialize all information regarding this object from the stream. + + :param stream: + A file-like object. + + :return: + self + """ + raise NotImplementedError("To be implemented in subclass") + + +class TraversableIterableObj(IterableObj, Traversable): + __slots__ = () + + TIobj_tuple = Tuple[Union[T_TIobj, None], T_TIobj] + + def list_traverse(self: T_TIobj, *args: Any, **kwargs: Any) -> IterableList[T_TIobj]: + return super()._list_traverse(*args, **kwargs) + + @overload + def traverse(self: T_TIobj) -> Iterator[T_TIobj]: ... + + @overload + def traverse( + self: T_TIobj, + predicate: Callable[[Union[T_TIobj, Tuple[Union[T_TIobj, None], T_TIobj]], int], bool], + prune: Callable[[Union[T_TIobj, Tuple[Union[T_TIobj, None], T_TIobj]], int], bool], + depth: int, + branch_first: bool, + visit_once: bool, + ignore_self: Literal[True], + as_edge: Literal[False], + ) -> Iterator[T_TIobj]: ... + + @overload + def traverse( + self: T_TIobj, + predicate: Callable[[Union[T_TIobj, Tuple[Union[T_TIobj, None], T_TIobj]], int], bool], + prune: Callable[[Union[T_TIobj, Tuple[Union[T_TIobj, None], T_TIobj]], int], bool], + depth: int, + branch_first: bool, + visit_once: bool, + ignore_self: Literal[False], + as_edge: Literal[True], + ) -> Iterator[Tuple[Union[T_TIobj, None], T_TIobj]]: ... + + @overload + def traverse( + self: T_TIobj, + predicate: Callable[[Union[T_TIobj, TIobj_tuple], int], bool], + prune: Callable[[Union[T_TIobj, TIobj_tuple], int], bool], + depth: int, + branch_first: bool, + visit_once: bool, + ignore_self: Literal[True], + as_edge: Literal[True], + ) -> Iterator[Tuple[T_TIobj, T_TIobj]]: ... + + def traverse( + self: T_TIobj, + predicate: Callable[[Union[T_TIobj, TIobj_tuple], int], bool] = lambda i, d: True, + prune: Callable[[Union[T_TIobj, TIobj_tuple], int], bool] = lambda i, d: False, + depth: int = -1, + branch_first: bool = True, + visit_once: bool = True, + ignore_self: int = 1, + as_edge: bool = False, + ) -> Union[Iterator[T_TIobj], Iterator[Tuple[T_TIobj, T_TIobj]], Iterator[TIobj_tuple]]: + """For documentation, see :meth:`Traversable._traverse`.""" + + ## To typecheck instead of using cast: + # + # import itertools + # from git.types import TypeGuard + # def is_commit_traversed(inp: Tuple) -> TypeGuard[Tuple[Iterator[Tuple['Commit', 'Commit']]]]: + # for x in inp[1]: + # if not isinstance(x, tuple) and len(x) != 2: + # if all(isinstance(inner, Commit) for inner in x): + # continue + # return True + # + # ret = super(Commit, self).traverse(predicate, prune, depth, branch_first, visit_once, ignore_self, as_edge) + # ret_tup = itertools.tee(ret, 2) + # assert is_commit_traversed(ret_tup), f"{[type(x) for x in list(ret_tup[0])]}" + # return ret_tup[0] + + return cast( + Union[Iterator[T_TIobj], Iterator[Tuple[Union[None, T_TIobj], T_TIobj]]], + super()._traverse( + predicate, # type: ignore[arg-type] + prune, # type: ignore[arg-type] + depth, + branch_first, + visit_once, + ignore_self, + as_edge, + ), + ) diff --git a/lib/python3.12/site-packages/git/refs/__init__.py b/lib/python3.12/site-packages/git/refs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d6157e6f37422c85b3f2c62479e68da888a8f598 --- /dev/null +++ b/lib/python3.12/site-packages/git/refs/__init__.py @@ -0,0 +1,21 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = [ + "HEAD", + "Head", + "RefLog", + "RefLogEntry", + "Reference", + "RemoteReference", + "SymbolicReference", + "Tag", + "TagReference", +] + +from .head import HEAD, Head +from .log import RefLog, RefLogEntry +from .reference import Reference +from .remote import RemoteReference +from .symbolic import SymbolicReference +from .tag import Tag, TagReference diff --git a/lib/python3.12/site-packages/git/refs/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/git/refs/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f2a74ea51f1b3adbd6f9890ea59953ada683563a Binary files /dev/null and b/lib/python3.12/site-packages/git/refs/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/refs/__pycache__/head.cpython-312.pyc b/lib/python3.12/site-packages/git/refs/__pycache__/head.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f1562200ccecf3fec07241953836df0d01d03bcc Binary files /dev/null and b/lib/python3.12/site-packages/git/refs/__pycache__/head.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/refs/__pycache__/log.cpython-312.pyc b/lib/python3.12/site-packages/git/refs/__pycache__/log.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..75acddaee4c06a0342893deaf8b1fa4889cc09bb Binary files /dev/null and b/lib/python3.12/site-packages/git/refs/__pycache__/log.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/refs/__pycache__/reference.cpython-312.pyc b/lib/python3.12/site-packages/git/refs/__pycache__/reference.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a640e8ff089cb2b5b6e0d6d858f1cb9120c56e7 Binary files /dev/null and b/lib/python3.12/site-packages/git/refs/__pycache__/reference.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/refs/__pycache__/remote.cpython-312.pyc b/lib/python3.12/site-packages/git/refs/__pycache__/remote.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0576ed1d7bc1ad15af854d2fe5415328d0accddf Binary files /dev/null and b/lib/python3.12/site-packages/git/refs/__pycache__/remote.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/refs/__pycache__/symbolic.cpython-312.pyc b/lib/python3.12/site-packages/git/refs/__pycache__/symbolic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64c244b2a452594fe18df82268542183f4643684 Binary files /dev/null and b/lib/python3.12/site-packages/git/refs/__pycache__/symbolic.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/refs/__pycache__/tag.cpython-312.pyc b/lib/python3.12/site-packages/git/refs/__pycache__/tag.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da67fe2fb55546bf18282b36039437dc617acdd3 Binary files /dev/null and b/lib/python3.12/site-packages/git/refs/__pycache__/tag.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/refs/head.py b/lib/python3.12/site-packages/git/refs/head.py new file mode 100644 index 0000000000000000000000000000000000000000..3c43993e7ed3ff615ae0febd7934f448f1bce94c --- /dev/null +++ b/lib/python3.12/site-packages/git/refs/head.py @@ -0,0 +1,300 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Some ref-based objects. + +Note the distinction between the :class:`HEAD` and :class:`Head` classes. +""" + +__all__ = ["HEAD", "Head"] + +from git.config import GitConfigParser, SectionConstraint +from git.exc import GitCommandError +from git.util import join_path + +from .reference import Reference +from .symbolic import SymbolicReference + +# typing --------------------------------------------------- + +from typing import Any, Sequence, TYPE_CHECKING, Union + +from git.types import Commit_ish, PathLike + +if TYPE_CHECKING: + from git.refs import RemoteReference + from git.repo import Repo + +# ------------------------------------------------------------------- + + +def strip_quotes(string: str) -> str: + if string.startswith('"') and string.endswith('"'): + return string[1:-1] + return string + + +class HEAD(SymbolicReference): + """Special case of a :class:`~git.refs.symbolic.SymbolicReference` representing the + repository's HEAD reference.""" + + _HEAD_NAME = "HEAD" + _ORIG_HEAD_NAME = "ORIG_HEAD" + + __slots__ = () + + def __init__(self, repo: "Repo", path: PathLike = _HEAD_NAME) -> None: + if path != self._HEAD_NAME: + raise ValueError("HEAD instance must point to %r, got %r" % (self._HEAD_NAME, path)) + super().__init__(repo, path) + + def orig_head(self) -> SymbolicReference: + """ + :return: + :class:`~git.refs.symbolic.SymbolicReference` pointing at the ORIG_HEAD, + which is maintained to contain the previous value of HEAD. + """ + return SymbolicReference(self.repo, self._ORIG_HEAD_NAME) + + def reset( + self, + commit: Union[Commit_ish, SymbolicReference, str] = "HEAD", + index: bool = True, + working_tree: bool = False, + paths: Union[PathLike, Sequence[PathLike], None] = None, + **kwargs: Any, + ) -> "HEAD": + """Reset our HEAD to the given commit optionally synchronizing the index and + working tree. The reference we refer to will be set to commit as well. + + :param commit: + :class:`~git.objects.commit.Commit`, :class:`~git.refs.reference.Reference`, + or string identifying a revision we should reset HEAD to. + + :param index: + If ``True``, the index will be set to match the given commit. + Otherwise it will not be touched. + + :param working_tree: + If ``True``, the working tree will be forcefully adjusted to match the given + commit, possibly overwriting uncommitted changes without warning. + If `working_tree` is ``True``, `index` must be ``True`` as well. + + :param paths: + Single path or list of paths relative to the git root directory + that are to be reset. This allows to partially reset individual files. + + :param kwargs: + Additional arguments passed to :manpage:`git-reset(1)`. + + :return: + self + """ + mode: Union[str, None] + mode = "--soft" + if index: + mode = "--mixed" + + # Explicit "--mixed" when passing paths is deprecated since git 1.5.4. + # See https://github.com/gitpython-developers/GitPython/discussions/1876. + if paths: + mode = None + # END special case + # END handle index + + if working_tree: + mode = "--hard" + if not index: + raise ValueError("Cannot reset the working tree if the index is not reset as well") + + # END working tree handling + + try: + self.repo.git.reset(mode, commit, "--", paths, **kwargs) + except GitCommandError as e: + # git nowadays may use 1 as status to indicate there are still unstaged + # modifications after the reset. + if e.status != 1: + raise + # END handle exception + + return self + + +class Head(Reference): + """A Head is a named reference to a :class:`~git.objects.commit.Commit`. Every Head + instance contains a name and a :class:`~git.objects.commit.Commit` object. + + Examples:: + + >>> repo = Repo("/path/to/repo") + >>> head = repo.heads[0] + + >>> head.name + 'master' + + >>> head.commit + + + >>> head.commit.hexsha + '1c09f116cbc2cb4100fb6935bb162daa4723f455' + """ + + _common_path_default = "refs/heads" + k_config_remote = "remote" + k_config_remote_ref = "merge" # Branch to merge from remote. + + @classmethod + def delete(cls, repo: "Repo", *heads: "Union[Head, str]", force: bool = False, **kwargs: Any) -> None: # type: ignore[override] + """Delete the given heads. + + :param force: + If ``True``, the heads will be deleted even if they are not yet merged into + the main development stream. Default ``False``. + """ + flag = "-d" + if force: + flag = "-D" + repo.git.branch(flag, *heads) + + def set_tracking_branch(self, remote_reference: Union["RemoteReference", None]) -> "Head": + """Configure this branch to track the given remote reference. This will + alter this branch's configuration accordingly. + + :param remote_reference: + The remote reference to track or None to untrack any references. + + :return: + self + """ + from .remote import RemoteReference + + if remote_reference is not None and not isinstance(remote_reference, RemoteReference): + raise ValueError("Incorrect parameter type: %r" % remote_reference) + # END handle type + + with self.config_writer() as writer: + if remote_reference is None: + writer.remove_option(self.k_config_remote) + writer.remove_option(self.k_config_remote_ref) + if len(writer.options()) == 0: + writer.remove_section() + else: + writer.set_value(self.k_config_remote, remote_reference.remote_name) + writer.set_value( + self.k_config_remote_ref, + Head.to_full_path(remote_reference.remote_head), + ) + + return self + + def tracking_branch(self) -> Union["RemoteReference", None]: + """ + :return: + The remote reference we are tracking, or ``None`` if we are not a tracking + branch. + """ + from .remote import RemoteReference + + reader = self.config_reader() + if reader.has_option(self.k_config_remote) and reader.has_option(self.k_config_remote_ref): + ref = Head( + self.repo, + Head.to_full_path(strip_quotes(reader.get_value(self.k_config_remote_ref))), + ) + remote_refpath = RemoteReference.to_full_path(join_path(reader.get_value(self.k_config_remote), ref.name)) + return RemoteReference(self.repo, remote_refpath) + # END handle have tracking branch + + # We are not a tracking branch. + return None + + def rename(self, new_path: PathLike, force: bool = False) -> "Head": + """Rename self to a new path. + + :param new_path: + Either a simple name or a path, e.g. ``new_name`` or ``features/new_name``. + The prefix ``refs/heads`` is implied. + + :param force: + If ``True``, the rename will succeed even if a head with the target name + already exists. + + :return: + self + + :note: + Respects the ref log, as git commands are used. + """ + flag = "-m" + if force: + flag = "-M" + + self.repo.git.branch(flag, self, new_path) + self.path = "%s/%s" % (self._common_path_default, new_path) + return self + + def checkout(self, force: bool = False, **kwargs: Any) -> Union["HEAD", "Head"]: + """Check out this head by setting the HEAD to this reference, by updating the + index to reflect the tree we point to and by updating the working tree to + reflect the latest index. + + The command will fail if changed working tree files would be overwritten. + + :param force: + If ``True``, changes to the index and the working tree will be discarded. + If ``False``, :exc:`~git.exc.GitCommandError` will be raised in that + situation. + + :param kwargs: + Additional keyword arguments to be passed to git checkout, e.g. + ``b="new_branch"`` to create a new branch at the given spot. + + :return: + The active branch after the checkout operation, usually self unless a new + branch has been created. + If there is no active branch, as the HEAD is now detached, the HEAD + reference will be returned instead. + + :note: + By default it is only allowed to checkout heads - everything else will leave + the HEAD detached which is allowed and possible, but remains a special state + that some tools might not be able to handle. + """ + kwargs["f"] = force + if kwargs["f"] is False: + kwargs.pop("f") + + self.repo.git.checkout(self, **kwargs) + if self.repo.head.is_detached: + return self.repo.head + else: + return self.repo.active_branch + + # { Configuration + def _config_parser(self, read_only: bool) -> SectionConstraint[GitConfigParser]: + if read_only: + parser = self.repo.config_reader() + else: + parser = self.repo.config_writer() + # END handle parser instance + + return SectionConstraint(parser, 'branch "%s"' % self.name) + + def config_reader(self) -> SectionConstraint[GitConfigParser]: + """ + :return: + A configuration parser instance constrained to only read this instance's + values. + """ + return self._config_parser(read_only=True) + + def config_writer(self) -> SectionConstraint[GitConfigParser]: + """ + :return: + A configuration writer instance with read-and write access to options of + this head. + """ + return self._config_parser(read_only=False) + + # } END configuration diff --git a/lib/python3.12/site-packages/git/refs/log.py b/lib/python3.12/site-packages/git/refs/log.py new file mode 100644 index 0000000000000000000000000000000000000000..4751cff99b82887b0bd7c898bca3fd6a827ec50f --- /dev/null +++ b/lib/python3.12/site-packages/git/refs/log.py @@ -0,0 +1,399 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["RefLog", "RefLogEntry"] + +from mmap import mmap +import os.path as osp +import re +import time as _time + +from git.compat import defenc +from git.objects.util import ( + Serializable, + altz_to_utctz_str, + parse_date, +) +from git.util import ( + Actor, + LockedFD, + LockFile, + assure_directory_exists, + bin_to_hex, + file_contents_ro_filepath, + to_native_path, +) + +# typing ------------------------------------------------------------------ + +from typing import Iterator, List, Tuple, TYPE_CHECKING, Union + +from git.types import PathLike + +if TYPE_CHECKING: + from io import BytesIO + + from git.config import GitConfigParser, SectionConstraint + from git.refs import SymbolicReference + +# ------------------------------------------------------------------------------ + + +class RefLogEntry(Tuple[str, str, Actor, Tuple[int, int], str]): + """Named tuple allowing easy access to the revlog data fields.""" + + _re_hexsha_only = re.compile(r"^[0-9A-Fa-f]{40}$") + + __slots__ = () + + def __repr__(self) -> str: + """Representation of ourselves in git reflog format.""" + return self.format() + + def format(self) -> str: + """:return: A string suitable to be placed in a reflog file.""" + act = self.actor + time = self.time + return "{} {} {} <{}> {!s} {}\t{}\n".format( + self.oldhexsha, + self.newhexsha, + act.name, + act.email, + time[0], + altz_to_utctz_str(time[1]), + self.message, + ) + + @property + def oldhexsha(self) -> str: + """The hexsha to the commit the ref pointed to before the change.""" + return self[0] + + @property + def newhexsha(self) -> str: + """The hexsha to the commit the ref now points to, after the change.""" + return self[1] + + @property + def actor(self) -> Actor: + """Actor instance, providing access.""" + return self[2] + + @property + def time(self) -> Tuple[int, int]: + """Time as tuple: + + * [0] = ``int(time)`` + * [1] = ``int(timezone_offset)`` in :attr:`time.altzone` format + """ + return self[3] + + @property + def message(self) -> str: + """Message describing the operation that acted on the reference.""" + return self[4] + + @classmethod + def new( + cls, + oldhexsha: str, + newhexsha: str, + actor: Actor, + time: int, + tz_offset: int, + message: str, + ) -> "RefLogEntry": # skipcq: PYL-W0621 + """:return: New instance of a :class:`RefLogEntry`""" + if not isinstance(actor, Actor): + raise ValueError("Need actor instance, got %s" % actor) + # END check types + return RefLogEntry((oldhexsha, newhexsha, actor, (time, tz_offset), message)) + + @classmethod + def from_line(cls, line: bytes) -> "RefLogEntry": + """:return: New :class:`RefLogEntry` instance from the given revlog line. + + :param line: + Line bytes without trailing newline + + :raise ValueError: + If `line` could not be parsed. + """ + line_str = line.decode(defenc) + fields = line_str.split("\t", 1) + if len(fields) == 1: + info, msg = fields[0], None + elif len(fields) == 2: + info, msg = fields + else: + raise ValueError("Line must have up to two TAB-separated fields. Got %s" % repr(line_str)) + # END handle first split + + oldhexsha = info[:40] + newhexsha = info[41:81] + for hexsha in (oldhexsha, newhexsha): + if not cls._re_hexsha_only.match(hexsha): + raise ValueError("Invalid hexsha: %r" % (hexsha,)) + # END if hexsha re doesn't match + # END for each hexsha + + email_end = info.find(">", 82) + if email_end == -1: + raise ValueError("Missing token: >") + # END handle missing end brace + + actor = Actor._from_string(info[82 : email_end + 1]) + time, tz_offset = parse_date(info[email_end + 2 :]) # skipcq: PYL-W0621 + + return RefLogEntry((oldhexsha, newhexsha, actor, (time, tz_offset), msg)) # type: ignore [arg-type] + + +class RefLog(List[RefLogEntry], Serializable): + R"""A reflog contains :class:`RefLogEntry`\s, each of which defines a certain state + of the head in question. Custom query methods allow to retrieve log entries by date + or by other criteria. + + Reflog entries are ordered. The first added entry is first in the list. The last + entry, i.e. the last change of the head or reference, is last in the list. + """ + + __slots__ = ("_path",) + + def __new__(cls, filepath: Union[PathLike, None] = None) -> "RefLog": + inst = super().__new__(cls) + return inst + + def __init__(self, filepath: Union[PathLike, None] = None) -> None: + """Initialize this instance with an optional filepath, from which we will + initialize our data. The path is also used to write changes back using the + :meth:`write` method.""" + self._path = filepath + if filepath is not None: + self._read_from_file() + # END handle filepath + + def _read_from_file(self) -> None: + try: + fmap = file_contents_ro_filepath(self._path, stream=True, allow_mmap=True) + except OSError: + # It is possible and allowed that the file doesn't exist! + return + # END handle invalid log + + try: + self._deserialize(fmap) + finally: + fmap.close() + # END handle closing of handle + + # { Interface + + @classmethod + def from_file(cls, filepath: PathLike) -> "RefLog": + """ + :return: + A new :class:`RefLog` instance containing all entries from the reflog at the + given `filepath`. + + :param filepath: + Path to reflog. + + :raise ValueError: + If the file could not be read or was corrupted in some way. + """ + return cls(filepath) + + @classmethod + def path(cls, ref: "SymbolicReference") -> str: + """ + :return: + String to absolute path at which the reflog of the given ref instance would + be found. The path is not guaranteed to point to a valid file though. + + :param ref: + :class:`~git.refs.symbolic.SymbolicReference` instance + """ + return osp.join(ref.repo.git_dir, "logs", to_native_path(ref.path)) + + @classmethod + def iter_entries(cls, stream: Union[str, "BytesIO", mmap]) -> Iterator[RefLogEntry]: + """ + :return: + Iterator yielding :class:`RefLogEntry` instances, one for each line read + from the given stream. + + :param stream: + File-like object containing the revlog in its native format or string + instance pointing to a file to read. + """ + new_entry = RefLogEntry.from_line + if isinstance(stream, str): + # Default args return mmap since Python 3. + _stream = file_contents_ro_filepath(stream) + assert isinstance(_stream, mmap) + else: + _stream = stream + # END handle stream type + while True: + line = _stream.readline() + if not line: + return + yield new_entry(line.strip()) + # END endless loop + + @classmethod + def entry_at(cls, filepath: PathLike, index: int) -> "RefLogEntry": + """ + :return: + :class:`RefLogEntry` at the given index. + + :param filepath: + Full path to the index file from which to read the entry. + + :param index: + Python list compatible index, i.e. it may be negative to specify an entry + counted from the end of the list. + + :raise IndexError: + If the entry didn't exist. + + :note: + This method is faster as it only parses the entry at index, skipping all + other lines. Nonetheless, the whole file has to be read if the index is + negative. + """ + with open(filepath, "rb") as fp: + if index < 0: + return RefLogEntry.from_line(fp.readlines()[index].strip()) + # Read until index is reached. + + for i in range(index + 1): + line = fp.readline() + if not line: + raise IndexError(f"Index file ended at line {i + 1}, before given index was reached") + # END abort on eof + # END handle runup + + return RefLogEntry.from_line(line.strip()) + # END handle index + + def to_file(self, filepath: PathLike) -> None: + """Write the contents of the reflog instance to a file at the given filepath. + + :param filepath: + Path to file. Parent directories are assumed to exist. + """ + lfd = LockedFD(filepath) + assure_directory_exists(filepath, is_file=True) + + fp = lfd.open(write=True, stream=True) + try: + self._serialize(fp) + lfd.commit() + except BaseException: + lfd.rollback() + raise + # END handle change + + @classmethod + def append_entry( + cls, + config_reader: Union[Actor, "GitConfigParser", "SectionConstraint", None], + filepath: PathLike, + oldbinsha: bytes, + newbinsha: bytes, + message: str, + write: bool = True, + ) -> "RefLogEntry": + """Append a new log entry to the revlog at filepath. + + :param config_reader: + Configuration reader of the repository - used to obtain user information. + May also be an :class:`~git.util.Actor` instance identifying the committer + directly or ``None``. + + :param filepath: + Full path to the log file. + + :param oldbinsha: + Binary sha of the previous commit. + + :param newbinsha: + Binary sha of the current commit. + + :param message: + Message describing the change to the reference. + + :param write: + If ``True``, the changes will be written right away. + Otherwise the change will not be written. + + :return: + :class:`RefLogEntry` objects which was appended to the log. + + :note: + As we are append-only, concurrent access is not a problem as we do not + interfere with readers. + """ + + if len(oldbinsha) != 20 or len(newbinsha) != 20: + raise ValueError("Shas need to be given in binary format") + # END handle sha type + assure_directory_exists(filepath, is_file=True) + first_line = message.split("\n")[0] + if isinstance(config_reader, Actor): + committer = config_reader # mypy thinks this is Actor | Gitconfigparser, but why? + else: + committer = Actor.committer(config_reader) + entry = RefLogEntry( + ( + bin_to_hex(oldbinsha).decode("ascii"), + bin_to_hex(newbinsha).decode("ascii"), + committer, + (int(_time.time()), _time.altzone), + first_line, + ) + ) + + if write: + lf = LockFile(filepath) + lf._obtain_lock_or_raise() + fd = open(filepath, "ab") + try: + fd.write(entry.format().encode(defenc)) + finally: + fd.close() + lf._release_lock() + # END handle write operation + return entry + + def write(self) -> "RefLog": + """Write this instance's data to the file we are originating from. + + :return: + self + """ + if self._path is None: + raise ValueError("Instance was not initialized with a path, use to_file(...) instead") + # END assert path + self.to_file(self._path) + return self + + # } END interface + + # { Serializable Interface + + def _serialize(self, stream: "BytesIO") -> "RefLog": + write = stream.write + + # Write all entries. + for e in self: + write(e.format().encode(defenc)) + # END for each entry + return self + + def _deserialize(self, stream: "BytesIO") -> "RefLog": + self.extend(self.iter_entries(stream)) + return self + + # } END serializable interface diff --git a/lib/python3.12/site-packages/git/refs/reference.py b/lib/python3.12/site-packages/git/refs/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..0c4327225df105be2123e3acb6fb1ed1a8797e1b --- /dev/null +++ b/lib/python3.12/site-packages/git/refs/reference.py @@ -0,0 +1,177 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["Reference"] + +import os +from git.util import IterableObj, LazyMixin + +from .symbolic import SymbolicReference, T_References + +# typing ------------------------------------------------------------------ + +from typing import Any, Callable, Iterator, TYPE_CHECKING, Type, Union + +from git.types import AnyGitObject, PathLike, _T + +if TYPE_CHECKING: + from git.repo import Repo + +# ------------------------------------------------------------------------------ + +# { Utilities + + +def require_remote_ref_path(func: Callable[..., _T]) -> Callable[..., _T]: + """A decorator raising :exc:`ValueError` if we are not a valid remote, based on the + path.""" + + def wrapper(self: T_References, *args: Any) -> _T: + if not self.is_remote(): + raise ValueError("ref path does not point to a remote reference: %s" % self.path) + return func(self, *args) + + # END wrapper + wrapper.__name__ = func.__name__ + return wrapper + + +# } END utilities + + +class Reference(SymbolicReference, LazyMixin, IterableObj): + """A named reference to any object. + + Subclasses may apply restrictions though, e.g., a :class:`~git.refs.head.Head` can + only point to commits. + """ + + __slots__ = () + + _points_to_commits_only = False + _resolve_ref_on_create = True + _common_path_default = "refs" + + def __init__(self, repo: "Repo", path: PathLike, check_path: bool = True) -> None: + """Initialize this instance. + + :param repo: + Our parent repository. + + :param path: + Path relative to the ``.git/`` directory pointing to the ref in question, + e.g. ``refs/heads/master``. + + :param check_path: + If ``False``, you can provide any path. + Otherwise the path must start with the default path prefix of this type. + """ + if check_path and not os.fspath(path).startswith(self._common_path_default + "/"): + raise ValueError(f"Cannot instantiate {self.__class__.__name__!r} from path {path}") + self.path: str # SymbolicReference converts to string at the moment. + super().__init__(repo, path) + + def __str__(self) -> str: + return self.name + + # { Interface + + # @ReservedAssignment + def set_object( + self, + object: Union[AnyGitObject, "SymbolicReference", str], + logmsg: Union[str, None] = None, + ) -> "Reference": + """Special version which checks if the head-log needs an update as well. + + :return: + self + """ + oldbinsha = None + if logmsg is not None: + head = self.repo.head + if not head.is_detached and head.ref == self: + oldbinsha = self.commit.binsha + # END handle commit retrieval + # END handle message is set + + super().set_object(object, logmsg) + + if oldbinsha is not None: + # From refs/files-backend.c in git-source: + # /* + # * Special hack: If a branch is updated directly and HEAD + # * points to it (may happen on the remote side of a push + # * for example) then logically the HEAD reflog should be + # * updated too. + # * A generic solution implies reverse symref information, + # * but finding all symrefs pointing to the given branch + # * would be rather costly for this rare event (the direct + # * update of a branch) to be worth it. So let's cheat and + # * check with HEAD only which should cover 99% of all usage + # * scenarios (even 100% of the default ones). + # */ + self.repo.head.log_append(oldbinsha, logmsg) + # END check if the head + + return self + + # NOTE: No need to overwrite properties, as the will only work without a the log. + + @property + def name(self) -> str: + """ + :return: + (shortest) Name of this reference - it may contain path components + """ + # The first two path tokens can be removed as they are + # refs/heads or refs/tags or refs/remotes. + tokens = self.path.split("/") + if len(tokens) < 3: + return self.path # could be refs/HEAD + return "/".join(tokens[2:]) + + @classmethod + def iter_items( + cls: Type[T_References], + repo: "Repo", + common_path: Union[PathLike, None] = None, + *args: Any, + **kwargs: Any, + ) -> Iterator[T_References]: + """Equivalent to + :meth:`SymbolicReference.iter_items `, + but will return non-detached references as well.""" + return cls._iter_items(repo, common_path) + + # } END interface + + # { Remote Interface + + @property + @require_remote_ref_path + def remote_name(self) -> str: + """ + :return: + Name of the remote we are a reference of, such as ``origin`` for a reference + named ``origin/master``. + """ + tokens = self.path.split("/") + # /refs/remotes// + return tokens[2] + + @property + @require_remote_ref_path + def remote_head(self) -> str: + """ + :return: + Name of the remote head itself, e.g. ``master``. + + :note: + The returned name is usually not qualified enough to uniquely identify a + branch. + """ + tokens = self.path.split("/") + return "/".join(tokens[3:]) + + # } END remote interface diff --git a/lib/python3.12/site-packages/git/refs/remote.py b/lib/python3.12/site-packages/git/refs/remote.py new file mode 100644 index 0000000000000000000000000000000000000000..b4f4f7b366282695cf2676de246758edd631b1ca --- /dev/null +++ b/lib/python3.12/site-packages/git/refs/remote.py @@ -0,0 +1,79 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Module implementing a remote object allowing easy access to git remotes.""" + +__all__ = ["RemoteReference"] + +import os + +from git.util import join_path + +from .head import Head + +# typing ------------------------------------------------------------------ + +from typing import Any, Iterator, NoReturn, TYPE_CHECKING, Union + +from git.types import PathLike + +if TYPE_CHECKING: + from git.remote import Remote + from git.repo import Repo + +# ------------------------------------------------------------------------------ + + +class RemoteReference(Head): + """A reference pointing to a remote head.""" + + _common_path_default = Head._remote_common_path_default + + @classmethod + def iter_items( + cls, + repo: "Repo", + common_path: Union[PathLike, None] = None, + remote: Union["Remote", None] = None, + *args: Any, + **kwargs: Any, + ) -> Iterator["RemoteReference"]: + """Iterate remote references, and if given, constrain them to the given remote.""" + common_path = common_path or cls._common_path_default + if remote is not None: + common_path = join_path(common_path, str(remote)) + # END handle remote constraint + # super is Reference + return super().iter_items(repo, common_path) + + # The Head implementation of delete also accepts strs, but this implementation does + # not. mypy doesn't have a way of representing tightening the types of arguments in + # subclasses and recommends Any or "type: ignore". + # (See: https://github.com/python/typing/issues/241) + @classmethod + def delete(cls, repo: "Repo", *refs: "RemoteReference", **kwargs: Any) -> None: # type: ignore[override] + """Delete the given remote references. + + :note: + `kwargs` are given for comparability with the base class method as we + should not narrow the signature. + """ + repo.git.branch("-d", "-r", *refs) + # The official deletion method will ignore remote symbolic refs - these are + # generally ignored in the refs/ folder. We don't though and delete remainders + # manually. + for ref in refs: + try: + os.remove(os.path.join(repo.common_dir, ref.path)) + except OSError: + pass + try: + os.remove(os.path.join(repo.git_dir, ref.path)) + except OSError: + pass + # END for each ref + + @classmethod + def create(cls, *args: Any, **kwargs: Any) -> NoReturn: + """Raise :exc:`TypeError`. Defined so the ``create`` method is disabled.""" + raise TypeError("Cannot explicitly create remote references") diff --git a/lib/python3.12/site-packages/git/refs/symbolic.py b/lib/python3.12/site-packages/git/refs/symbolic.py new file mode 100644 index 0000000000000000000000000000000000000000..99af4f57cabac79f179df06239d41ee4fde32874 --- /dev/null +++ b/lib/python3.12/site-packages/git/refs/symbolic.py @@ -0,0 +1,934 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +__all__ = ["SymbolicReference"] + +import os +from pathlib import Path + +from gitdb.exc import BadName, BadObject + +from git.compat import defenc +from git.objects.base import Object +from git.objects.commit import Commit +from git.refs.log import RefLog +from git.util import ( + LockedFD, + assure_directory_exists, + hex_to_bin, + join_path, + join_path_native, + to_native_path_linux, +) + +# typing ------------------------------------------------------------------ + +from typing import ( + Any, + Iterator, + List, + TYPE_CHECKING, + Tuple, + Type, + TypeVar, + Union, + cast, +) + +from git.types import AnyGitObject, PathLike + +if TYPE_CHECKING: + from git.config import GitConfigParser + from git.objects.commit import Actor + from git.refs.log import RefLogEntry + from git.refs.reference import Reference + from git.repo import Repo + + +T_References = TypeVar("T_References", bound="SymbolicReference") + +# ------------------------------------------------------------------------------ + + +def _git_dir(repo: "Repo", path: Union[PathLike, None]) -> PathLike: + """Find the git dir that is appropriate for the path.""" + name = f"{path}" + if name in ["HEAD", "ORIG_HEAD", "FETCH_HEAD", "index", "logs"]: + return repo.git_dir + return repo.common_dir + + +class SymbolicReference: + """Special case of a reference that is symbolic. + + This does not point to a specific commit, but to another + :class:`~git.refs.head.Head`, which itself specifies a commit. + + A typical example for a symbolic reference is :class:`~git.refs.head.HEAD`. + """ + + __slots__ = ("repo", "path") + + _resolve_ref_on_create = False + _points_to_commits_only = True + _common_path_default = "" + _remote_common_path_default = "refs/remotes" + _id_attribute_ = "name" + + def __init__(self, repo: "Repo", path: PathLike, check_path: bool = False) -> None: + self.repo = repo + self.path: PathLike = path + + def __str__(self) -> str: + return os.fspath(self.path) + + def __repr__(self) -> str: + return '' % (self.__class__.__name__, self.path) + + def __eq__(self, other: object) -> bool: + if hasattr(other, "path"): + other = cast(SymbolicReference, other) + return self.path == other.path + return False + + def __ne__(self, other: object) -> bool: + return not (self == other) + + def __hash__(self) -> int: + return hash(self.path) + + @property + def name(self) -> str: + """ + :return: + In case of symbolic references, the shortest assumable name is the path + itself. + """ + return os.fspath(self.path) + + @property + def abspath(self) -> PathLike: + return join_path_native(_git_dir(self.repo, self.path), self.path) + + @classmethod + def _get_packed_refs_path(cls, repo: "Repo") -> str: + return os.path.join(repo.common_dir, "packed-refs") + + @classmethod + def _iter_packed_refs(cls, repo: "Repo") -> Iterator[Tuple[str, str]]: + """Return an iterator yielding pairs of sha1/path pairs (as strings) for the + corresponding refs. + + :note: + The packed refs file will be kept open as long as we iterate. + """ + try: + with open(cls._get_packed_refs_path(repo), "rt", encoding="UTF-8") as fp: + for line in fp: + line = line.strip() + if not line: + continue + if line.startswith("#"): + # "# pack-refs with: peeled fully-peeled sorted" + # the git source code shows "peeled", + # "fully-peeled" and "sorted" as the keywords + # that can go on this line, as per comments in git file + # refs/packed-backend.c + # I looked at master on 2017-10-11, + # commit 111ef79afe, after tag v2.15.0-rc1 + # from repo https://github.com/git/git.git + if line.startswith("# pack-refs with:") and "peeled" not in line: + raise TypeError("PackingType of packed-Refs not understood: %r" % line) + # END abort if we do not understand the packing scheme + continue + # END parse comment + + # Skip dereferenced tag object entries - previous line was actual + # tag reference for it. + if line[0] == "^": + continue + + yield cast(Tuple[str, str], tuple(line.split(" ", 1))) + # END for each line + except OSError: + return None + # END no packed-refs file handling + + @classmethod + def dereference_recursive(cls, repo: "Repo", ref_path: Union[PathLike, None]) -> str: + """ + :return: + hexsha stored in the reference at the given `ref_path`, recursively + dereferencing all intermediate references as required + + :param repo: + The repository containing the reference at `ref_path`. + """ + + while True: + hexsha, ref_path = cls._get_ref_info(repo, ref_path) + if hexsha is not None: + return hexsha + # END recursive dereferencing + + @staticmethod + def _check_ref_name_valid(ref_path: PathLike) -> None: + """Check a ref name for validity. + + This is based on the rules described in :manpage:`git-check-ref-format(1)`. + """ + previous: Union[str, None] = None + one_before_previous: Union[str, None] = None + for c in os.fspath(ref_path): + if c in " ~^:?*[\\": + raise ValueError( + f"Invalid reference '{ref_path}': references cannot contain spaces, tildes (~), carets (^)," + f" colons (:), question marks (?), asterisks (*), open brackets ([) or backslashes (\\)" + ) + elif c == ".": + if previous is None or previous == "/": + raise ValueError( + f"Invalid reference '{ref_path}': references cannot start with a period (.) or contain '/.'" + ) + elif previous == ".": + raise ValueError(f"Invalid reference '{ref_path}': references cannot contain '..'") + elif c == "/": + if previous == "/": + raise ValueError(f"Invalid reference '{ref_path}': references cannot contain '//'") + elif previous is None: + raise ValueError( + f"Invalid reference '{ref_path}': references cannot start with forward slashes '/'" + ) + elif c == "{" and previous == "@": + raise ValueError(f"Invalid reference '{ref_path}': references cannot contain '@{{'") + elif ord(c) < 32 or ord(c) == 127: + raise ValueError(f"Invalid reference '{ref_path}': references cannot contain ASCII control characters") + + one_before_previous = previous + previous = c + + if previous == ".": + raise ValueError(f"Invalid reference '{ref_path}': references cannot end with a period (.)") + elif previous == "/": + raise ValueError(f"Invalid reference '{ref_path}': references cannot end with a forward slash (/)") + elif previous == "@" and one_before_previous is None: + raise ValueError(f"Invalid reference '{ref_path}': references cannot be '@'") + elif any(component.endswith(".lock") for component in Path(ref_path).parts): + raise ValueError( + f"Invalid reference '{ref_path}': references cannot have slash-separated components that end with" + " '.lock'" + ) + + @classmethod + def _get_ref_info_helper( + cls, repo: "Repo", ref_path: Union[PathLike, None] + ) -> Union[Tuple[str, None], Tuple[None, str]]: + """ + :return: + *(str(sha), str(target_ref_path))*, where: + + * *sha* is of the file at rela_path points to if available, or ``None``. + * *target_ref_path* is the reference we point to, or ``None``. + """ + if ref_path: + cls._check_ref_name_valid(ref_path) + + tokens: Union[None, List[str], Tuple[str, str]] = None + repodir = _git_dir(repo, ref_path) + try: + with open(os.path.join(repodir, ref_path), "rt", encoding="UTF-8") as fp: # type: ignore[arg-type] + value = fp.read().rstrip() + # Don't only split on spaces, but on whitespace, which allows to parse lines like: + # 60b64ef992065e2600bfef6187a97f92398a9144 branch 'master' of git-server:/path/to/repo + tokens = value.split() + assert len(tokens) != 0 + except OSError: + # Probably we are just packed. Find our entry in the packed refs file. + # NOTE: We are not a symbolic ref if we are in a packed file, as these + # are excluded explicitly. + for sha, path in cls._iter_packed_refs(repo): + if path != ref_path: + continue + # sha will be used. + tokens = sha, path + break + # END for each packed ref + # END handle packed refs + if tokens is None: + raise ValueError("Reference at %r does not exist" % ref_path) + + # Is it a reference? + if tokens[0] == "ref:": + return (None, tokens[1]) + + # It's a commit. + if repo.re_hexsha_only.match(tokens[0]): + return (tokens[0], None) + + raise ValueError("Failed to parse reference information from %r" % ref_path) + + @classmethod + def _get_ref_info(cls, repo: "Repo", ref_path: Union[PathLike, None]) -> Union[Tuple[str, None], Tuple[None, str]]: + """ + :return: + *(str(sha), str(target_ref_path))*, where: + + * *sha* is of the file at rela_path points to if available, or ``None``. + * *target_ref_path* is the reference we point to, or ``None``. + """ + return cls._get_ref_info_helper(repo, ref_path) + + def _get_object(self) -> AnyGitObject: + """ + :return: + The object our ref currently refers to. Refs can be cached, they will always + point to the actual object as it gets re-created on each query. + """ + # We have to be dynamic here as we may be a tag which can point to anything. + # Our path will be resolved to the hexsha which will be used accordingly. + return Object.new_from_sha(self.repo, hex_to_bin(self.dereference_recursive(self.repo, self.path))) + + def _get_commit(self) -> "Commit": + """ + :return: + :class:`~git.objects.commit.Commit` object we point to. This works for + detached and non-detached :class:`SymbolicReference` instances. The symbolic + reference will be dereferenced recursively. + """ + obj = self._get_object() + if obj.type == "tag": + obj = obj.object + # END dereference tag + + if obj.type != Commit.type: + raise TypeError("Symbolic Reference pointed to object %r, commit was required" % obj) + # END handle type + return obj + + def set_commit( + self, + commit: Union[Commit, "SymbolicReference", str], + logmsg: Union[str, None] = None, + ) -> "SymbolicReference": + """Like :meth:`set_object`, but restricts the type of object to be a + :class:`~git.objects.commit.Commit`. + + :raise ValueError: + If `commit` is not a :class:`~git.objects.commit.Commit` object, nor does it + point to a commit. + + :return: + self + """ + # Check the type - assume the best if it is a base-string. + invalid_type = False + if isinstance(commit, Object): + invalid_type = commit.type != Commit.type + elif isinstance(commit, SymbolicReference): + invalid_type = commit.object.type != Commit.type + else: + try: + invalid_type = self.repo.rev_parse(commit).type != Commit.type + except (BadObject, BadName) as e: + raise ValueError("Invalid object: %s" % commit) from e + # END handle exception + # END verify type + + if invalid_type: + raise ValueError("Need commit, got %r" % commit) + # END handle raise + + # We leave strings to the rev-parse method below. + self.set_object(commit, logmsg) + + return self + + def set_object( + self, + object: Union[AnyGitObject, "SymbolicReference", str], + logmsg: Union[str, None] = None, + ) -> "SymbolicReference": + """Set the object we point to, possibly dereference our symbolic reference + first. If the reference does not exist, it will be created. + + :param object: + A refspec, a :class:`SymbolicReference` or an + :class:`~git.objects.base.Object` instance. + + * :class:`SymbolicReference` instances will be dereferenced beforehand to + obtain the git object they point to. + * :class:`~git.objects.base.Object` instances must represent git objects + (:class:`~git.types.AnyGitObject`). + + :param logmsg: + If not ``None``, the message will be used in the reflog entry to be written. + Otherwise the reflog is not altered. + + :note: + Plain :class:`SymbolicReference` instances may not actually point to objects + by convention. + + :return: + self + """ + if isinstance(object, SymbolicReference): + object = object.object # @ReservedAssignment + # END resolve references + + is_detached = True + try: + is_detached = self.is_detached + except ValueError: + pass + # END handle non-existing ones + + if is_detached: + return self.set_reference(object, logmsg) + + # set the commit on our reference + return self._get_reference().set_object(object, logmsg) + + @property + def commit(self) -> "Commit": + """Query or set commits directly""" + return self._get_commit() + + @commit.setter + def commit(self, commit: Union[Commit, "SymbolicReference", str]) -> "SymbolicReference": + return self.set_commit(commit) + + @property + def object(self) -> AnyGitObject: + """Return the object our ref currently refers to""" + return self._get_object() + + @object.setter + def object(self, object: Union[AnyGitObject, "SymbolicReference", str]) -> "SymbolicReference": + return self.set_object(object) + + def _get_reference(self) -> "Reference": + """ + :return: + :class:`~git.refs.reference.Reference` object we point to + + :raise TypeError: + If this symbolic reference is detached, hence it doesn't point to a + reference, but to a commit. + """ + sha, target_ref_path = self._get_ref_info(self.repo, self.path) + if target_ref_path is None: + raise TypeError("%s is a detached symbolic reference as it points to %r" % (self, sha)) + return cast("Reference", self.from_path(self.repo, target_ref_path)) + + def set_reference( + self, + ref: Union[AnyGitObject, "SymbolicReference", str], + logmsg: Union[str, None] = None, + ) -> "SymbolicReference": + """Set ourselves to the given `ref`. + + It will stay a symbol if the `ref` is a :class:`~git.refs.reference.Reference`. + + Otherwise a git object, specified as a :class:`~git.objects.base.Object` + instance or refspec, is assumed. If it is valid, this reference will be set to + it, which effectively detaches the reference if it was a purely symbolic one. + + :param ref: + A :class:`SymbolicReference` instance, an :class:`~git.objects.base.Object` + instance (specifically an :class:`~git.types.AnyGitObject`), or a refspec + string. Only if the ref is a :class:`SymbolicReference` instance, we will + point to it. Everything else is dereferenced to obtain the actual object. + + :param logmsg: + If set to a string, the message will be used in the reflog. + Otherwise, a reflog entry is not written for the changed reference. + The previous commit of the entry will be the commit we point to now. + + See also: :meth:`log_append` + + :return: + self + + :note: + This symbolic reference will not be dereferenced. For that, see + :meth:`set_object`. + """ + write_value = None + obj = None + if isinstance(ref, SymbolicReference): + write_value = "ref: %s" % ref.path + elif isinstance(ref, Object): + obj = ref + write_value = ref.hexsha + elif isinstance(ref, str): + try: + obj = self.repo.rev_parse(ref + "^{}") # Optionally dereference tags. + write_value = obj.hexsha + except (BadObject, BadName) as e: + raise ValueError("Could not extract object from %s" % ref) from e + # END end try string + else: + raise ValueError("Unrecognized Value: %r" % ref) + # END try commit attribute + + # typecheck + if obj is not None and self._points_to_commits_only and obj.type != Commit.type: + raise TypeError("Require commit, got %r" % obj) + # END verify type + + oldbinsha: bytes = b"" + if logmsg is not None: + try: + oldbinsha = self.commit.binsha + except ValueError: + oldbinsha = Commit.NULL_BIN_SHA + # END handle non-existing + # END retrieve old hexsha + + fpath = self.abspath + assure_directory_exists(fpath, is_file=True) + + lfd = LockedFD(fpath) + fd = lfd.open(write=True, stream=True) + try: + fd.write(write_value.encode("utf-8") + b"\n") + lfd.commit() + except BaseException: + lfd.rollback() + raise + # Adjust the reflog + if logmsg is not None: + self.log_append(oldbinsha, logmsg) + + return self + + # Aliased reference + @property + def reference(self) -> "Reference": + return self._get_reference() + + @reference.setter + def reference(self, ref: Union[AnyGitObject, "SymbolicReference", str]) -> "SymbolicReference": + return self.set_reference(ref) + + ref = reference + + def is_valid(self) -> bool: + """ + :return: + ``True`` if the reference is valid, hence it can be read and points to a + valid object or reference. + """ + try: + self.object # noqa: B018 + except (OSError, ValueError): + return False + else: + return True + + @property + def is_detached(self) -> bool: + """ + :return: + ``True`` if we are a detached reference, hence we point to a specific commit + instead to another reference. + """ + try: + self.ref # noqa: B018 + return False + except TypeError: + return True + + def log(self) -> "RefLog": + """ + :return: + :class:`~git.refs.log.RefLog` for this reference. + Its last entry reflects the latest change applied to this reference. + + :note: + As the log is parsed every time, its recommended to cache it for use instead + of calling this method repeatedly. It should be considered read-only. + """ + return RefLog.from_file(RefLog.path(self)) + + def log_append( + self, + oldbinsha: bytes, + message: Union[str, None], + newbinsha: Union[bytes, None] = None, + ) -> "RefLogEntry": + """Append a logentry to the logfile of this ref. + + :param oldbinsha: + Binary sha this ref used to point to. + + :param message: + A message describing the change. + + :param newbinsha: + The sha the ref points to now. If None, our current commit sha will be used. + + :return: + The added :class:`~git.refs.log.RefLogEntry` instance. + """ + # NOTE: We use the committer of the currently active commit - this should be + # correct to allow overriding the committer on a per-commit level. + # See https://github.com/gitpython-developers/GitPython/pull/146. + try: + committer_or_reader: Union["Actor", "GitConfigParser"] = self.commit.committer + except ValueError: + committer_or_reader = self.repo.config_reader() + # END handle newly cloned repositories + if newbinsha is None: + newbinsha = self.commit.binsha + + if message is None: + message = "" + + return RefLog.append_entry(committer_or_reader, RefLog.path(self), oldbinsha, newbinsha, message) + + def log_entry(self, index: int) -> "RefLogEntry": + """ + :return: + :class:`~git.refs.log.RefLogEntry` at the given index + + :param index: + Python list compatible positive or negative index. + + :note: + This method must read part of the reflog during execution, hence it should + be used sparingly, or only if you need just one index. In that case, it will + be faster than the :meth:`log` method. + """ + return RefLog.entry_at(RefLog.path(self), index) + + @classmethod + def to_full_path(cls, path: Union[PathLike, "SymbolicReference"]) -> PathLike: + """ + :return: + String with a full repository-relative path which can be used to initialize + a :class:`~git.refs.reference.Reference` instance, for instance by using + :meth:`Reference.from_path `. + """ + if isinstance(path, SymbolicReference): + path = path.path + full_ref_path = path + if not cls._common_path_default: + return full_ref_path + if not os.fspath(path).startswith(cls._common_path_default + "/"): + full_ref_path = "%s/%s" % (cls._common_path_default, path) + return full_ref_path + + @classmethod + def delete(cls, repo: "Repo", path: PathLike) -> None: + """Delete the reference at the given path. + + :param repo: + Repository to delete the reference from. + + :param path: + Short or full path pointing to the reference, e.g. ``refs/myreference`` or + just ``myreference``, hence ``refs/`` is implied. + Alternatively the symbolic reference to be deleted. + """ + full_ref_path = cls.to_full_path(path) + abs_path = os.path.join(repo.common_dir, full_ref_path) + if os.path.exists(abs_path): + os.remove(abs_path) + else: + # Check packed refs. + pack_file_path = cls._get_packed_refs_path(repo) + try: + with open(pack_file_path, "rb") as reader: + new_lines = [] + made_change = False + dropped_last_line = False + for line_bytes in reader: + line = line_bytes.decode(defenc) + _, _, line_ref = line.partition(" ") + line_ref = line_ref.strip() + # Keep line if it is a comment or if the ref to delete is not in + # the line. + # If we deleted the last line and this one is a tag-reference + # object, we drop it as well. + if (line.startswith("#") or full_ref_path != line_ref) and ( + not dropped_last_line or dropped_last_line and not line.startswith("^") + ): + new_lines.append(line) + dropped_last_line = False + continue + # END skip comments and lines without our path + + # Drop this line. + made_change = True + dropped_last_line = True + + # Write the new lines. + if made_change: + # Binary writing is required, otherwise Windows will open the file + # in text mode and change LF to CRLF! + with open(pack_file_path, "wb") as fd: + fd.writelines(line.encode(defenc) for line in new_lines) + + except OSError: + pass # It didn't exist at all. + + # Delete the reflog. + reflog_path = RefLog.path(cls(repo, full_ref_path)) + if os.path.isfile(reflog_path): + os.remove(reflog_path) + # END remove reflog + + @classmethod + def _create( + cls: Type[T_References], + repo: "Repo", + path: PathLike, + resolve: bool, + reference: Union["SymbolicReference", str], + force: bool, + logmsg: Union[str, None] = None, + ) -> T_References: + """Internal method used to create a new symbolic reference. + + If `resolve` is ``False``, the reference will be taken as is, creating a proper + symbolic reference. Otherwise it will be resolved to the corresponding object + and a detached symbolic reference will be created instead. + """ + git_dir = _git_dir(repo, path) + full_ref_path = cls.to_full_path(path) + abs_ref_path = os.path.join(git_dir, full_ref_path) + + # Figure out target data. + target = reference + if resolve: + target = repo.rev_parse(str(reference)) + + if not force and os.path.isfile(abs_ref_path): + target_data = str(target) + if isinstance(target, SymbolicReference): + target_data = os.fspath(target.path) + if not resolve: + target_data = "ref: " + target_data + with open(abs_ref_path, "rb") as fd: + existing_data = fd.read().decode(defenc).strip() + if existing_data != target_data: + raise OSError( + "Reference at %r does already exist, pointing to %r, requested was %r" + % (full_ref_path, existing_data, target_data) + ) + # END no force handling + + ref = cls(repo, full_ref_path) + ref.set_reference(target, logmsg) + return ref + + @classmethod + def create( + cls: Type[T_References], + repo: "Repo", + path: PathLike, + reference: Union["SymbolicReference", str] = "HEAD", + logmsg: Union[str, None] = None, + force: bool = False, + **kwargs: Any, + ) -> T_References: + """Create a new symbolic reference: a reference pointing to another reference. + + :param repo: + Repository to create the reference in. + + :param path: + Full path at which the new symbolic reference is supposed to be created at, + e.g. ``NEW_HEAD`` or ``symrefs/my_new_symref``. + + :param reference: + The reference which the new symbolic reference should point to. + If it is a commit-ish, the symbolic ref will be detached. + + :param force: + If ``True``, force creation even if a symbolic reference with that name + already exists. Raise :exc:`OSError` otherwise. + + :param logmsg: + If not ``None``, the message to append to the reflog. + If ``None``, no reflog entry is written. + + :return: + Newly created symbolic reference + + :raise OSError: + If a (Symbolic)Reference with the same name but different contents already + exists. + + :note: + This does not alter the current HEAD, index or working tree. + """ + return cls._create(repo, path, cls._resolve_ref_on_create, reference, force, logmsg) + + def rename(self, new_path: PathLike, force: bool = False) -> "SymbolicReference": + """Rename self to a new path. + + :param new_path: + Either a simple name or a full path, e.g. ``new_name`` or + ``features/new_name``. + The prefix ``refs/`` is implied for references and will be set as needed. + In case this is a symbolic ref, there is no implied prefix. + + :param force: + If ``True``, the rename will succeed even if a head with the target name + already exists. It will be overwritten in that case. + + :return: + self + + :raise OSError: + If a file at path but with different contents already exists. + """ + new_path = self.to_full_path(new_path) + if self.path == new_path: + return self + + new_abs_path = os.path.join(_git_dir(self.repo, new_path), new_path) + cur_abs_path = os.path.join(_git_dir(self.repo, self.path), self.path) + if os.path.isfile(new_abs_path): + if not force: + # If they point to the same file, it's not an error. + with open(new_abs_path, "rb") as fd1: + f1 = fd1.read().strip() + with open(cur_abs_path, "rb") as fd2: + f2 = fd2.read().strip() + if f1 != f2: + raise OSError("File at path %r already exists" % new_abs_path) + # else: We could remove ourselves and use the other one, but... + # ...for clarity, we just continue as usual. + # END not force handling + os.remove(new_abs_path) + # END handle existing target file + + dname = os.path.dirname(new_abs_path) + if not os.path.isdir(dname): + os.makedirs(dname) + # END create directory + + os.rename(cur_abs_path, new_abs_path) + self.path = new_path + + return self + + @classmethod + def _iter_items( + cls: Type[T_References], repo: "Repo", common_path: Union[PathLike, None] = None + ) -> Iterator[T_References]: + if common_path is None: + common_path = cls._common_path_default + rela_paths = set() + + # Walk loose refs. + # Currently we do not follow links. + for root, dirs, files in os.walk(join_path_native(repo.common_dir, common_path)): + if "refs" not in root.split(os.sep): # Skip non-refs subfolders. + refs_id = [d for d in dirs if d == "refs"] + if refs_id: + dirs[0:] = ["refs"] + # END prune non-refs folders + + for f in files: + if f == "packed-refs": + continue + abs_path = to_native_path_linux(join_path(root, f)) + rela_paths.add(abs_path.replace(to_native_path_linux(repo.common_dir) + "/", "")) + # END for each file in root directory + # END for each directory to walk + + # Read packed refs. + for _sha, rela_path in cls._iter_packed_refs(repo): + if rela_path.startswith(os.fspath(common_path)): + rela_paths.add(rela_path) + # END relative path matches common path + # END packed refs reading + + # Yield paths in sorted order. + for path in sorted(rela_paths): + try: + yield cls.from_path(repo, path) + except ValueError: + continue + # END for each sorted relative refpath + + @classmethod + def iter_items( + cls: Type[T_References], + repo: "Repo", + common_path: Union[PathLike, None] = None, + *args: Any, + **kwargs: Any, + ) -> Iterator[T_References]: + """Find all refs in the repository. + + :param repo: + The :class:`~git.repo.base.Repo`. + + :param common_path: + Optional keyword argument to the path which is to be shared by all returned + Ref objects. + Defaults to class specific portion if ``None``, ensuring that only refs + suitable for the actual class are returned. + + :return: + A list of :class:`SymbolicReference`, each guaranteed to be a symbolic ref + which is not detached and pointing to a valid ref. + + The list is lexicographically sorted. The returned objects are instances of + concrete subclasses, such as :class:`~git.refs.head.Head` or + :class:`~git.refs.tag.TagReference`. + """ + return (r for r in cls._iter_items(repo, common_path) if r.__class__ is SymbolicReference or not r.is_detached) + + @classmethod + def from_path(cls: Type[T_References], repo: "Repo", path: PathLike) -> T_References: + """Make a symbolic reference from a path. + + :param path: + Full ``.git``-directory-relative path name to the Reference to instantiate. + + :note: + Use :meth:`to_full_path` if you only have a partial path of a known + Reference type. + + :return: + Instance of type :class:`~git.refs.reference.Reference`, + :class:`~git.refs.head.Head`, or :class:`~git.refs.tag.Tag`, depending on + the given path. + """ + if not path: + raise ValueError("Cannot create Reference from %r" % path) + + # Names like HEAD are inserted after the refs module is imported - we have an + # import dependency cycle and don't want to import these names in-function. + from . import HEAD, Head, RemoteReference, TagReference, Reference + + for ref_type in ( + HEAD, + Head, + RemoteReference, + TagReference, + Reference, + SymbolicReference, + ): + try: + instance = cast(T_References, ref_type(repo, path)) + if instance.__class__ is SymbolicReference and instance.is_detached: + raise ValueError("SymbolicRef was detached, we drop it") + else: + return instance + + except ValueError: + pass + # END exception handling + # END for each type to try + raise ValueError("Could not find reference type suitable to handle path %r" % path) + + def is_remote(self) -> bool: + """:return: True if this symbolic reference points to a remote branch""" + return os.fspath(self.path).startswith(self._remote_common_path_default + "/") diff --git a/lib/python3.12/site-packages/git/refs/tag.py b/lib/python3.12/site-packages/git/refs/tag.py new file mode 100644 index 0000000000000000000000000000000000000000..4525b09cb24bd3ce7ac62a08591c4b31774cdf4b --- /dev/null +++ b/lib/python3.12/site-packages/git/refs/tag.py @@ -0,0 +1,155 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Provides a :class:`~git.refs.reference.Reference`-based type for lightweight tags. + +This defines the :class:`TagReference` class (and its alias :class:`Tag`), which +represents lightweight tags. For annotated tags (which are git objects), see the +:mod:`git.objects.tag` module. +""" + +__all__ = ["TagReference", "Tag"] + +from .reference import Reference + +# typing ------------------------------------------------------------------ + +from typing import Any, TYPE_CHECKING, Type, Union + +from git.types import AnyGitObject, PathLike + +if TYPE_CHECKING: + from git.objects import Commit, TagObject + from git.refs import SymbolicReference + from git.repo import Repo + +# ------------------------------------------------------------------------------ + + +class TagReference(Reference): + """A lightweight tag reference which either points to a commit, a tag object or any + other object. In the latter case additional information, like the signature or the + tag-creator, is available. + + This tag object will always point to a commit object, but may carry additional + information in a tag object:: + + tagref = TagReference.list_items(repo)[0] + print(tagref.commit.message) + if tagref.tag is not None: + print(tagref.tag.message) + """ + + __slots__ = () + + _common_default = "tags" + _common_path_default = Reference._common_path_default + "/" + _common_default + + @property # type: ignore[misc] + def commit(self) -> "Commit": # LazyMixin has unrelated commit method + """:return: Commit object the tag ref points to + + :raise ValueError: + If the tag points to a tree or blob. + """ + obj = self.object + while obj.type != "commit": + if obj.type == "tag": + # It is a tag object which carries the commit as an object - we can point to anything. + obj = obj.object + else: + raise ValueError( + ( + "Cannot resolve commit as tag %s points to a %s object - " + + "use the `.object` property instead to access it" + ) + % (self, obj.type) + ) + return obj + + @property + def tag(self) -> Union["TagObject", None]: + """ + :return: + Tag object this tag ref points to, or ``None`` in case we are a lightweight + tag + """ + obj = self.object + if obj.type == "tag": + return obj + return None + + # Make object read-only. It should be reasonably hard to adjust an existing tag. + @property # type: ignore[misc] + def object(self) -> AnyGitObject: + return Reference._get_object(self) + + @classmethod + def create( + cls: Type["TagReference"], + repo: "Repo", + path: PathLike, + reference: Union[str, "SymbolicReference"] = "HEAD", + logmsg: Union[str, None] = None, + force: bool = False, + **kwargs: Any, + ) -> "TagReference": + """Create a new tag reference. + + :param repo: + The :class:`~git.repo.base.Repo` to create the tag in. + + :param path: + The name of the tag, e.g. ``1.0`` or ``releases/1.0``. + The prefix ``refs/tags`` is implied. + + :param reference: + A reference to the :class:`~git.objects.base.Object` you want to tag. + The referenced object can be a commit, tree, or blob. + + :param logmsg: + If not ``None``, the message will be used in your tag object. This will also + create an additional tag object that allows to obtain that information, + e.g.:: + + tagref.tag.message + + :param message: + Synonym for the `logmsg` parameter. Included for backwards compatibility. + `logmsg` takes precedence if both are passed. + + :param force: + If ``True``, force creation of a tag even though that tag already exists. + + :param kwargs: + Additional keyword arguments to be passed to :manpage:`git-tag(1)`. + + :return: + A new :class:`TagReference`. + """ + if "ref" in kwargs and kwargs["ref"]: + reference = kwargs["ref"] + + if "message" in kwargs and kwargs["message"]: + kwargs["m"] = kwargs["message"] + del kwargs["message"] + + if logmsg: + kwargs["m"] = logmsg + + if force: + kwargs["f"] = True + + args = (path, reference) + + repo.git.tag(*args, **kwargs) + return TagReference(repo, "%s/%s" % (cls._common_path_default, path)) + + @classmethod + def delete(cls, repo: "Repo", *tags: "TagReference") -> None: # type: ignore[override] + """Delete the given existing tag or tags.""" + repo.git.tag("-d", *tags) + + +# Provide an alias. +Tag = TagReference diff --git a/lib/python3.12/site-packages/git/repo/__init__.py b/lib/python3.12/site-packages/git/repo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..66319ef95f5c06956660127d6cb131fa99d3b780 --- /dev/null +++ b/lib/python3.12/site-packages/git/repo/__init__.py @@ -0,0 +1,8 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""Initialize the repo package.""" + +__all__ = ["Repo"] + +from .base import Repo diff --git a/lib/python3.12/site-packages/git/repo/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/git/repo/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f217cfba074a274dd6dc1118a83d22753464ecd Binary files /dev/null and b/lib/python3.12/site-packages/git/repo/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/repo/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/git/repo/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2334f4c5b4b12e4412449fade4b48f93c001bf59 Binary files /dev/null and b/lib/python3.12/site-packages/git/repo/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/repo/__pycache__/fun.cpython-312.pyc b/lib/python3.12/site-packages/git/repo/__pycache__/fun.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8b77c1817e2974de60ee729a7a75afe58bf9b96f Binary files /dev/null and b/lib/python3.12/site-packages/git/repo/__pycache__/fun.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/git/repo/base.py b/lib/python3.12/site-packages/git/repo/base.py new file mode 100644 index 0000000000000000000000000000000000000000..1f543cc574d35dd3287452d7f3c66bf8460c6dfa --- /dev/null +++ b/lib/python3.12/site-packages/git/repo/base.py @@ -0,0 +1,1633 @@ +# Copyright (C) 2008, 2009 Michael Trier (mtrier@gmail.com) and contributors +# +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +from __future__ import annotations + +__all__ = ["Repo"] + +import gc +import logging +import os +import os.path as osp +from pathlib import Path +import re +import shlex +import sys +import warnings + +import gitdb +from gitdb.db.loose import LooseObjectDB +from gitdb.exc import BadObject + +from git.cmd import Git, handle_process_output +from git.compat import defenc, safe_decode +from git.config import GitConfigParser +from git.db import GitCmdObjectDB +from git.exc import ( + GitCommandError, + InvalidGitRepositoryError, + NoSuchPathError, +) +from git.index import IndexFile +from git.objects import Submodule, RootModule, Commit +from git.refs import HEAD, Head, Reference, TagReference +from git.remote import Remote, add_progress, to_progress_instance +from git.util import ( + Actor, + cygpath, + expand_path, + finalize_process, + hex_to_bin, + remove_password_if_present, +) + +from .fun import ( + find_submodule_git_dir, + find_worktree_git_dir, + is_git_dir, + rev_parse, + touch, +) + +# typing ------------------------------------------------------ + +from git.types import ( + CallableProgress, + Commit_ish, + Lit_config_levels, + PathLike, + TBD, + Tree_ish, + assert_never, +) +from typing import ( + Any, + BinaryIO, + Callable, + Dict, + Iterator, + List, + Mapping, + NamedTuple, + Optional, + Sequence, + TYPE_CHECKING, + TextIO, + Tuple, + Type, + Union, + cast, +) + +from git.types import ConfigLevels_Tup, TypedDict + +if TYPE_CHECKING: + from git.objects import Tree + from git.objects.submodule.base import UpdateProgress + from git.refs.symbolic import SymbolicReference + from git.remote import RemoteProgress + from git.util import IterableList + +# ----------------------------------------------------------- + +_logger = logging.getLogger(__name__) + + +class BlameEntry(NamedTuple): + commit: Dict[str, Commit] + linenos: range + orig_path: Optional[str] + orig_linenos: range + + +class Repo: + """Represents a git repository and allows you to query references, create commit + information, generate diffs, create and clone repositories, and query the log. + + The following attributes are worth using: + + * :attr:`working_dir` is the working directory of the git command, which is the + working tree directory if available or the ``.git`` directory in case of bare + repositories. + + * :attr:`working_tree_dir` is the working tree directory, but will return ``None`` + if we are a bare repository. + + * :attr:`git_dir` is the ``.git`` repository directory, which is always set. + """ + + DAEMON_EXPORT_FILE = "git-daemon-export-ok" + + # Must exist, or __del__ will fail in case we raise on `__init__()`. + git = cast("Git", None) + + working_dir: PathLike + """The working directory of the git command.""" + + # stored as string for easier processing, but annotated as path for clearer intention + _working_tree_dir: Optional[PathLike] = None + + git_dir: PathLike + """The ``.git`` repository directory.""" + + _common_dir: PathLike = "" + + # Precompiled regex + re_whitespace = re.compile(r"\s+") + re_hexsha_only = re.compile(r"^[0-9A-Fa-f]{40}$") + re_hexsha_shortened = re.compile(r"^[0-9A-Fa-f]{4,40}$") + re_envvars = re.compile(r"(\$(\{\s?)?[a-zA-Z_]\w*(\}\s?)?|%\s?[a-zA-Z_]\w*\s?%)") + re_author_committer_start = re.compile(r"^(author|committer)") + re_tab_full_line = re.compile(r"^\t(.*)$") + + unsafe_git_clone_options = [ + # Executes arbitrary commands: + "--upload-pack", + "-u", + # Can override configuration variables that execute arbitrary commands: + "--config", + "-c", + ] + """Options to :manpage:`git-clone(1)` that allow arbitrary commands to be executed. + + The ``--upload-pack``/``-u`` option allows users to execute arbitrary commands + directly: + https://git-scm.com/docs/git-clone#Documentation/git-clone.txt---upload-packltupload-packgt + + The ``--config``/``-c`` option allows users to override configuration variables like + ``protocol.allow`` and ``core.gitProxy`` to execute arbitrary commands: + https://git-scm.com/docs/git-clone#Documentation/git-clone.txt---configltkeygtltvaluegt + """ + + # Invariants + config_level: ConfigLevels_Tup = ("system", "user", "global", "repository") + """Represents the configuration level of a configuration file.""" + + # Subclass configuration + GitCommandWrapperType = Git + """Subclasses may easily bring in their own custom types by placing a constructor or + type here.""" + + def __init__( + self, + path: Optional[PathLike] = None, + odbt: Type[LooseObjectDB] = GitCmdObjectDB, + search_parent_directories: bool = False, + expand_vars: bool = True, + ) -> None: + R"""Create a new :class:`Repo` instance. + + :param path: + The path to either the worktree directory or the .git directory itself:: + + repo = Repo("/Users/mtrier/Development/git-python") + repo = Repo("/Users/mtrier/Development/git-python.git") + repo = Repo("~/Development/git-python.git") + repo = Repo("$REPOSITORIES/Development/git-python.git") + repo = Repo(R"C:\Users\mtrier\Development\git-python\.git") + + - In *Cygwin*, `path` may be a ``cygdrive/...`` prefixed path. + - If `path` is ``None`` or an empty string, :envvar:`GIT_DIR` is used. If + that environment variable is absent or empty, the current directory is + used. + + :param odbt: + Object DataBase type - a type which is constructed by providing the + directory containing the database objects, i.e. ``.git/objects``. It will be + used to access all object data. + + :param search_parent_directories: + If ``True``, all parent directories will be searched for a valid repo as + well. + + Please note that this was the default behaviour in older versions of + GitPython, which is considered a bug though. + + :raise git.exc.InvalidGitRepositoryError: + + :raise git.exc.NoSuchPathError: + + :return: + :class:`Repo` + """ + + epath = path or os.getenv("GIT_DIR") + if not epath: + epath = os.getcwd() + epath = os.fspath(epath) + if Git.is_cygwin(): + # Given how the tests are written, this seems more likely to catch Cygwin + # git used from Windows than Windows git used from Cygwin. Therefore + # changing to Cygwin-style paths is the relevant operation. + epath = cygpath(epath) + + if expand_vars and re.search(self.re_envvars, epath): + warnings.warn( + "The use of environment variables in paths is deprecated" + + "\nfor security reasons and may be removed in the future!!", + stacklevel=1, + ) + epath = expand_path(epath, expand_vars) + if epath is not None: + if not os.path.exists(epath): + raise NoSuchPathError(epath) + + # Walk up the path to find the `.git` dir. + curpath = epath + git_dir = None + while curpath: + # ABOUT osp.NORMPATH + # It's important to normalize the paths, as submodules will otherwise + # initialize their repo instances with paths that depend on path-portions + # that will not exist after being removed. It's just cleaner. + if is_git_dir(curpath): + git_dir = curpath + # from man git-config : core.worktree + # Set the path to the root of the working tree. If GIT_COMMON_DIR + # environment variable is set, core.worktree is ignored and not used for + # determining the root of working tree. This can be overridden by the + # GIT_WORK_TREE environment variable. The value can be an absolute path + # or relative to the path to the .git directory, which is either + # specified by GIT_DIR, or automatically discovered. If GIT_DIR is + # specified but none of GIT_WORK_TREE and core.worktree is specified, + # the current working directory is regarded as the top level of your + # working tree. + self._working_tree_dir = os.path.dirname(git_dir) + if os.environ.get("GIT_COMMON_DIR") is None: + gitconf = self._config_reader("repository", git_dir) + if gitconf.has_option("core", "worktree"): + self._working_tree_dir = gitconf.get("core", "worktree") + if "GIT_WORK_TREE" in os.environ: + self._working_tree_dir = os.getenv("GIT_WORK_TREE") + break + + dotgit = osp.join(curpath, ".git") + sm_gitpath = find_submodule_git_dir(dotgit) + if sm_gitpath is not None: + git_dir = osp.normpath(sm_gitpath) + + sm_gitpath = find_submodule_git_dir(dotgit) + if sm_gitpath is None: + sm_gitpath = find_worktree_git_dir(dotgit) + + if sm_gitpath is not None: + git_dir = expand_path(sm_gitpath, expand_vars) + self._working_tree_dir = curpath + break + + if not search_parent_directories: + break + curpath, tail = osp.split(curpath) + if not tail: + break + # END while curpath + + if git_dir is None: + raise InvalidGitRepositoryError(epath) + self.git_dir = git_dir + + self._bare = False + try: + self._bare = self.config_reader("repository").getboolean("core", "bare") + except Exception: + # Let's not assume the option exists, although it should. + pass + + try: + common_dir = (Path(self.git_dir) / "commondir").read_text().splitlines()[0].strip() + self._common_dir = osp.join(self.git_dir, common_dir) + except OSError: + self._common_dir = "" + + # Adjust the working directory in case we are actually bare - we didn't know + # that in the first place. + if self._bare: + self._working_tree_dir = None + # END working dir handling + + self.working_dir: PathLike = self._working_tree_dir or self.common_dir + self.git = self.GitCommandWrapperType(self.working_dir) + + # Special handling, in special times. + rootpath = osp.join(self.common_dir, "objects") + if issubclass(odbt, GitCmdObjectDB): + self.odb = odbt(rootpath, self.git) + else: + self.odb = odbt(rootpath) + + def __enter__(self) -> "Repo": + return self + + def __exit__(self, *args: Any) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + def close(self) -> None: + if self.git: + self.git.clear_cache() + # Tempfiles objects on Windows are holding references to open files until + # they are collected by the garbage collector, thus preventing deletion. + # TODO: Find these references and ensure they are closed and deleted + # synchronously rather than forcing a gc collection. + if sys.platform == "win32": + gc.collect() + gitdb.util.mman.collect() + if sys.platform == "win32": + gc.collect() + + def __eq__(self, rhs: object) -> bool: + if isinstance(rhs, Repo): + return self.git_dir == rhs.git_dir + return False + + def __ne__(self, rhs: object) -> bool: + return not self.__eq__(rhs) + + def __hash__(self) -> int: + return hash(self.git_dir) + + @property + def description(self) -> str: + """The project's description""" + filename = osp.join(self.git_dir, "description") + with open(filename, "rb") as fp: + return fp.read().rstrip().decode(defenc) + + @description.setter + def description(self, descr: str) -> None: + filename = osp.join(self.git_dir, "description") + with open(filename, "wb") as fp: + fp.write((descr + "\n").encode(defenc)) + + @property + def working_tree_dir(self) -> Optional[PathLike]: + """ + :return: + The working tree directory of our git repository. + If this is a bare repository, ``None`` is returned. + """ + return self._working_tree_dir + + @property + def common_dir(self) -> PathLike: + """ + :return: + The git dir that holds everything except possibly HEAD, FETCH_HEAD, + ORIG_HEAD, COMMIT_EDITMSG, index, and logs/. + """ + return self._common_dir or self.git_dir + + @property + def bare(self) -> bool: + """:return: ``True`` if the repository is bare""" + return self._bare + + @property + def heads(self) -> "IterableList[Head]": + """A list of :class:`~git.refs.head.Head` objects representing the branch heads + in this repo. + + :return: + ``git.IterableList(Head, ...)`` + """ + return Head.list_items(self) + + @property + def branches(self) -> "IterableList[Head]": + """Alias for heads. + A list of :class:`~git.refs.head.Head` objects representing the branch heads + in this repo. + + :return: + ``git.IterableList(Head, ...)`` + """ + return self.heads + + @property + def references(self) -> "IterableList[Reference]": + """A list of :class:`~git.refs.reference.Reference` objects representing tags, + heads and remote references. + + :return: + ``git.IterableList(Reference, ...)`` + """ + return Reference.list_items(self) + + @property + def refs(self) -> "IterableList[Reference]": + """Alias for references. + A list of :class:`~git.refs.reference.Reference` objects representing tags, + heads and remote references. + + :return: + ``git.IterableList(Reference, ...)`` + """ + return self.references + + @property + def index(self) -> "IndexFile": + """ + :return: + A :class:`~git.index.base.IndexFile` representing this repository's index. + + :note: + This property can be expensive, as the returned + :class:`~git.index.base.IndexFile` will be reinitialized. + It is recommended to reuse the object. + """ + return IndexFile(self) + + @property + def head(self) -> "HEAD": + """ + :return: + :class:`~git.refs.head.HEAD` object pointing to the current head reference + """ + return HEAD(self, "HEAD") + + @property + def remotes(self) -> "IterableList[Remote]": + """A list of :class:`~git.remote.Remote` objects allowing to access and + manipulate remotes. + + :return: + ``git.IterableList(Remote, ...)`` + """ + return Remote.list_items(self) + + def remote(self, name: str = "origin") -> "Remote": + """:return: The remote with the specified name + + :raise ValueError: + If no remote with such a name exists. + """ + r = Remote(self, name) + if not r.exists(): + raise ValueError("Remote named '%s' didn't exist" % name) + return r + + # { Submodules + + @property + def submodules(self) -> "IterableList[Submodule]": + """ + :return: + git.IterableList(Submodule, ...) of direct submodules available from the + current head + """ + return Submodule.list_items(self) + + def submodule(self, name: str) -> "Submodule": + """:return: The submodule with the given name + + :raise ValueError: + If no such submodule exists. + """ + try: + return self.submodules[name] + except IndexError as e: + raise ValueError("Didn't find submodule named %r" % name) from e + # END exception handling + + def create_submodule(self, *args: Any, **kwargs: Any) -> Submodule: + """Create a new submodule. + + :note: + For a description of the applicable parameters, see the documentation of + :meth:`Submodule.add `. + + :return: + The created submodule. + """ + return Submodule.add(self, *args, **kwargs) + + def iter_submodules(self, *args: Any, **kwargs: Any) -> Iterator[Submodule]: + """An iterator yielding Submodule instances. + + See the :class:`~git.objects.util.Traversable` interface for a description of `args` + and `kwargs`. + + :return: + Iterator + """ + return RootModule(self).traverse(*args, **kwargs) + + def submodule_update(self, *args: Any, **kwargs: Any) -> RootModule: + """Update the submodules, keeping the repository consistent as it will + take the previous state into consideration. + + :note: + For more information, please see the documentation of + :meth:`RootModule.update `. + """ + return RootModule(self).update(*args, **kwargs) + + # }END submodules + + @property + def tags(self) -> "IterableList[TagReference]": + """A list of :class:`~git.refs.tag.TagReference` objects that are available in + this repo. + + :return: + ``git.IterableList(TagReference, ...)`` + """ + return TagReference.list_items(self) + + def tag(self, path: PathLike) -> TagReference: + """ + :return: + :class:`~git.refs.tag.TagReference` object, reference pointing to a + :class:`~git.objects.commit.Commit` or tag + + :param path: + Path to the tag reference, e.g. ``0.1.5`` or ``tags/0.1.5``. + """ + full_path = self._to_full_tag_path(path) + return TagReference(self, full_path) + + @staticmethod + def _to_full_tag_path(path: PathLike) -> str: + path_str = str(path) + if path_str.startswith(TagReference._common_path_default + "/"): + return path_str + if path_str.startswith(TagReference._common_default + "/"): + return Reference._common_path_default + "/" + path_str + else: + return TagReference._common_path_default + "/" + path_str + + def create_head( + self, + path: PathLike, + commit: Union["SymbolicReference", "str"] = "HEAD", + force: bool = False, + logmsg: Optional[str] = None, + ) -> "Head": + """Create a new head within the repository. + + :note: + For more documentation, please see the + :meth:`Head.create ` method. + + :return: + Newly created :class:`~git.refs.head.Head` Reference. + """ + return Head.create(self, path, commit, logmsg, force) + + def delete_head(self, *heads: "Union[str, Head]", **kwargs: Any) -> None: + """Delete the given heads. + + :param kwargs: + Additional keyword arguments to be passed to :manpage:`git-branch(1)`. + """ + return Head.delete(self, *heads, **kwargs) + + def create_tag( + self, + path: PathLike, + ref: Union[str, "SymbolicReference"] = "HEAD", + message: Optional[str] = None, + force: bool = False, + **kwargs: Any, + ) -> TagReference: + """Create a new tag reference. + + :note: + For more documentation, please see the + :meth:`TagReference.create ` method. + + :return: + :class:`~git.refs.tag.TagReference` object + """ + return TagReference.create(self, path, ref, message, force, **kwargs) + + def delete_tag(self, *tags: TagReference) -> None: + """Delete the given tag references.""" + return TagReference.delete(self, *tags) + + def create_remote(self, name: str, url: str, **kwargs: Any) -> Remote: + """Create a new remote. + + For more information, please see the documentation of the + :meth:`Remote.create ` method. + + :return: + :class:`~git.remote.Remote` reference + """ + return Remote.create(self, name, url, **kwargs) + + def delete_remote(self, remote: "Remote") -> str: + """Delete the given remote.""" + return Remote.remove(self, remote) + + def _get_config_path(self, config_level: Lit_config_levels, git_dir: Optional[PathLike] = None) -> str: + if git_dir is None: + git_dir = self.git_dir + # We do not support an absolute path of the gitconfig on Windows. + # Use the global config instead. + if sys.platform == "win32" and config_level == "system": + config_level = "global" + + if config_level == "system": + return "/etc/gitconfig" + elif config_level == "user": + config_home = os.environ.get("XDG_CONFIG_HOME") or osp.join(os.environ.get("HOME", "~"), ".config") + return osp.normpath(osp.expanduser(osp.join(config_home, "git", "config"))) + elif config_level == "global": + return osp.normpath(osp.expanduser("~/.gitconfig")) + elif config_level == "repository": + repo_dir = self._common_dir or git_dir + if not repo_dir: + raise NotADirectoryError + else: + return osp.normpath(osp.join(repo_dir, "config")) + else: + assert_never( # type: ignore[unreachable] + config_level, + ValueError(f"Invalid configuration level: {config_level!r}"), + ) + + def config_reader( + self, + config_level: Optional[Lit_config_levels] = None, + ) -> GitConfigParser: + """ + :return: + :class:`~git.config.GitConfigParser` allowing to read the full git + configuration, but not to write it. + + The configuration will include values from the system, user and repository + configuration files. + + :param config_level: + For possible values, see the :meth:`config_writer` method. If ``None``, all + applicable levels will be used. Specify a level in case you know which file + you wish to read to prevent reading multiple files. + + :note: + On Windows, system configuration cannot currently be read as the path is + unknown, instead the global path will be used. + """ + return self._config_reader(config_level=config_level) + + def _config_reader( + self, + config_level: Optional[Lit_config_levels] = None, + git_dir: Optional[PathLike] = None, + ) -> GitConfigParser: + if config_level is None: + files = [self._get_config_path(f, git_dir) for f in self.config_level if f] + else: + files = [self._get_config_path(config_level, git_dir)] + return GitConfigParser(files, read_only=True, repo=self) + + def config_writer(self, config_level: Lit_config_levels = "repository") -> GitConfigParser: + """ + :return: + A :class:`~git.config.GitConfigParser` allowing to write values of the + specified configuration file level. Config writers should be retrieved, used + to change the configuration, and written right away as they will lock the + configuration file in question and prevent other's to write it. + + :param config_level: + One of the following values: + + * ``"system"`` = system wide configuration file + * ``"global"`` = user level configuration file + * ``"`repository"`` = configuration file for this repository only + """ + return GitConfigParser(self._get_config_path(config_level), read_only=False, repo=self, merge_includes=False) + + def commit(self, rev: Union[str, Commit_ish, None] = None) -> Commit: + """The :class:`~git.objects.commit.Commit` object for the specified revision. + + :param rev: + Revision specifier, see :manpage:`git-rev-parse(1)` for viable options. + + :return: + :class:`~git.objects.commit.Commit` + """ + if rev is None: + return self.head.commit + return self.rev_parse(str(rev) + "^0") + + def iter_trees(self, *args: Any, **kwargs: Any) -> Iterator["Tree"]: + """:return: Iterator yielding :class:`~git.objects.tree.Tree` objects + + :note: + Accepts all arguments known to the :meth:`iter_commits` method. + """ + return (c.tree for c in self.iter_commits(*args, **kwargs)) + + def tree(self, rev: Union[Tree_ish, str, None] = None) -> "Tree": + """The :class:`~git.objects.tree.Tree` object for the given tree-ish revision. + + Examples:: + + repo.tree(repo.heads[0]) + + :param rev: + A revision pointing to a Treeish (being a commit or tree). + + :return: + :class:`~git.objects.tree.Tree` + + :note: + If you need a non-root level tree, find it by iterating the root tree. + Otherwise it cannot know about its path relative to the repository root and + subsequent operations might have unexpected results. + """ + if rev is None: + return self.head.commit.tree + return self.rev_parse(str(rev) + "^{tree}") + + def iter_commits( + self, + rev: Union[str, Commit, "SymbolicReference", None] = None, + paths: Union[PathLike, Sequence[PathLike]] = "", + **kwargs: Any, + ) -> Iterator[Commit]: + """An iterator of :class:`~git.objects.commit.Commit` objects representing the + history of a given ref/commit. + + :param rev: + Revision specifier, see :manpage:`git-rev-parse(1)` for viable options. + If ``None``, the active branch will be used. + + :param paths: + An optional path or a list of paths. If set, only commits that include the + path or paths will be returned. + + :param kwargs: + Arguments to be passed to :manpage:`git-rev-list(1)`. + Common ones are ``max_count`` and ``skip``. + + :note: + To receive only commits between two named revisions, use the + ``"revA...revB"`` revision specifier. + + :return: + Iterator of :class:`~git.objects.commit.Commit` objects + """ + if rev is None: + rev = self.head.commit + + return Commit.iter_items(self, rev, paths, **kwargs) + + def merge_base(self, *rev: TBD, **kwargs: Any) -> List[Commit]: + R"""Find the closest common ancestor for the given revision + (:class:`~git.objects.commit.Commit`\s, :class:`~git.refs.tag.Tag`\s, + :class:`~git.refs.reference.Reference`\s, etc.). + + :param rev: + At least two revs to find the common ancestor for. + + :param kwargs: + Additional arguments to be passed to the ``repo.git.merge_base()`` command + which does all the work. + + :return: + A list of :class:`~git.objects.commit.Commit` objects. If ``--all`` was + not passed as a keyword argument, the list will have at max one + :class:`~git.objects.commit.Commit`, or is empty if no common merge base + exists. + + :raise ValueError: + If fewer than two revisions are provided. + """ + if len(rev) < 2: + raise ValueError("Please specify at least two revs, got only %i" % len(rev)) + # END handle input + + res: List[Commit] = [] + try: + lines: List[str] = self.git.merge_base(*rev, **kwargs).splitlines() + except GitCommandError as err: + if err.status == 128: + raise + # END handle invalid rev + # Status code 1 is returned if there is no merge-base. + # (See: https://github.com/git/git/blob/v2.44.0/builtin/merge-base.c#L19) + return res + # END exception handling + + for line in lines: + res.append(self.commit(line)) + # END for each merge-base + + return res + + def is_ancestor(self, ancestor_rev: Commit, rev: Commit) -> bool: + """Check if a commit is an ancestor of another. + + :param ancestor_rev: + Rev which should be an ancestor. + + :param rev: + Rev to test against `ancestor_rev`. + + :return: + ``True`` if `ancestor_rev` is an ancestor to `rev`. + """ + try: + self.git.merge_base(ancestor_rev, rev, is_ancestor=True) + except GitCommandError as err: + if err.status == 1: + return False + raise + return True + + def is_valid_object(self, sha: str, object_type: Union[str, None] = None) -> bool: + try: + complete_sha = self.odb.partial_to_complete_sha_hex(sha) + object_info = self.odb.info(complete_sha) + if object_type: + if object_info.type == object_type.encode(): + return True + else: + _logger.debug( + "Commit hash points to an object of type '%s'. Requested were objects of type '%s'", + object_info.type.decode(), + object_type, + ) + return False + else: + return True + except BadObject: + _logger.debug("Commit hash is invalid.") + return False + + def _get_daemon_export(self) -> bool: + if self.git_dir: + filename = osp.join(self.git_dir, self.DAEMON_EXPORT_FILE) + return osp.exists(filename) + + def _set_daemon_export(self, value: object) -> None: + if self.git_dir: + filename = osp.join(self.git_dir, self.DAEMON_EXPORT_FILE) + fileexists = osp.exists(filename) + if value and not fileexists: + touch(filename) + elif not value and fileexists: + os.unlink(filename) + + @property + def daemon_export(self) -> bool: + """If True, git-daemon may export this repository""" + return self._get_daemon_export() + + @daemon_export.setter + def daemon_export(self, value: object) -> None: + self._set_daemon_export(value) + + def _get_alternates(self) -> List[str]: + """The list of alternates for this repo from which objects can be retrieved. + + :return: + List of strings being pathnames of alternates + """ + if self.git_dir: + alternates_path = osp.join(self.git_dir, "objects", "info", "alternates") + + if osp.exists(alternates_path): + with open(alternates_path, "rb") as f: + alts = f.read().decode(defenc) + return alts.strip().splitlines() + return [] + + def _set_alternates(self, alts: List[str]) -> None: + """Set the alternates. + + :param alts: + The array of string paths representing the alternates at which git should + look for objects, i.e. ``/home/user/repo/.git/objects``. + + :raise git.exc.NoSuchPathError: + + :note: + The method does not check for the existence of the paths in `alts`, as the + caller is responsible. + """ + alternates_path = osp.join(self.common_dir, "objects", "info", "alternates") + if not alts: + if osp.isfile(alternates_path): + os.remove(alternates_path) + else: + with open(alternates_path, "wb") as f: + f.write("\n".join(alts).encode(defenc)) + + @property + def alternates(self) -> List[str]: + """Retrieve a list of alternates paths or set a list paths to be used as alternates""" + return self._get_alternates() + + @alternates.setter + def alternates(self, alts: List[str]) -> None: + self._set_alternates(alts) + + def is_dirty( + self, + index: bool = True, + working_tree: bool = True, + untracked_files: bool = False, + submodules: bool = True, + path: Optional[PathLike] = None, + ) -> bool: + """ + :return: + ``True`` if the repository is considered dirty. By default it will react + like a :manpage:`git-status(1)` without untracked files, hence it is dirty + if the index or the working copy have changes. + """ + if self._bare: + # Bare repositories with no associated working directory are + # always considered to be clean. + return False + + # Start from the one which is fastest to evaluate. + default_args = ["--abbrev=40", "--full-index", "--raw"] + if not submodules: + default_args.append("--ignore-submodules") + if path: + default_args.extend(["--", os.fspath(path)]) + if index: + # diff index against HEAD. + if osp.isfile(self.index.path) and len(self.git.diff("--cached", *default_args)): + return True + # END index handling + if working_tree: + # diff index against working tree. + if len(self.git.diff(*default_args)): + return True + # END working tree handling + if untracked_files: + if len(self._get_untracked_files(path, ignore_submodules=not submodules)): + return True + # END untracked files + return False + + @property + def untracked_files(self) -> List[str]: + """ + :return: + list(str,...) + + Files currently untracked as they have not been staged yet. Paths are + relative to the current working directory of the git command. + + :note: + Ignored files will not appear here, i.e. files mentioned in ``.gitignore``. + + :note: + This property is expensive, as no cache is involved. To process the result, + please consider caching it yourself. + """ + return self._get_untracked_files() + + def _get_untracked_files(self, *args: Any, **kwargs: Any) -> List[str]: + # Make sure we get all files, not only untracked directories. + proc = self.git.status(*args, porcelain=True, untracked_files=True, as_process=True, **kwargs) + # Untracked files prefix in porcelain mode + prefix = "?? " + untracked_files = [] + for line in proc.stdout: + line = line.decode(defenc) + if not line.startswith(prefix): + continue + filename = line[len(prefix) :].rstrip("\n") + # Special characters are escaped + if filename[0] == filename[-1] == '"': + filename = filename[1:-1] + # WHATEVER ... it's a mess, but works for me + filename = filename.encode("ascii").decode("unicode_escape").encode("latin1").decode(defenc) + untracked_files.append(filename) + finalize_process(proc) + return untracked_files + + def ignored(self, *paths: PathLike) -> List[str]: + """Checks if paths are ignored via ``.gitignore``. + + This does so using the :manpage:`git-check-ignore(1)` method. + + :param paths: + List of paths to check whether they are ignored or not. + + :return: + Subset of those paths which are ignored + """ + try: + proc: str = self.git.check_ignore(*paths) + except GitCommandError as err: + if err.status == 1: + # If return code is 1, this means none of the items in *paths are + # ignored by Git, so return an empty list. + return [] + else: + # Raise the exception on all other return codes. + raise + + return proc.replace("\\\\", "\\").replace('"', "").split("\n") + + @property + def active_branch(self) -> Head: + """The name of the currently active branch. + + :raise TypeError: + If HEAD is detached. + + :return: + :class:`~git.refs.head.Head` to the active branch + """ + # reveal_type(self.head.reference) # => Reference + return self.head.reference + + def blame_incremental(self, rev: str | HEAD | None, file: str, **kwargs: Any) -> Iterator["BlameEntry"]: + """Iterator for blame information for the given file at the given revision. + + Unlike :meth:`blame`, this does not return the actual file's contents, only a + stream of :class:`BlameEntry` tuples. + + :param rev: + Revision specifier. If ``None``, the blame will include all the latest + uncommitted changes. Otherwise, anything successfully parsed by + :manpage:`git-rev-parse(1)` is a valid option. + + :return: + Lazy iterator of :class:`BlameEntry` tuples, where the commit indicates the + commit to blame for the line, and range indicates a span of line numbers in + the resulting file. + + If you combine all line number ranges outputted by this command, you should get + a continuous range spanning all line numbers in the file. + """ + + data: bytes = self.git.blame(rev, "--", file, p=True, incremental=True, stdout_as_string=False, **kwargs) + commits: Dict[bytes, Commit] = {} + + stream = (line for line in data.split(b"\n") if line) + while True: + try: + # When exhausted, causes a StopIteration, terminating this function. + line = next(stream) + except StopIteration: + return + split_line = line.split() + hexsha, orig_lineno_b, lineno_b, num_lines_b = split_line + lineno = int(lineno_b) + num_lines = int(num_lines_b) + orig_lineno = int(orig_lineno_b) + if hexsha not in commits: + # Now read the next few lines and build up a dict of properties for this + # commit. + props: Dict[bytes, bytes] = {} + while True: + try: + line = next(stream) + except StopIteration: + return + if line == b"boundary": + # "boundary" indicates a root commit and occurs instead of the + # "previous" tag. + continue + + tag, value = line.split(b" ", 1) + props[tag] = value + if tag == b"filename": + # "filename" formally terminates the entry for --incremental. + orig_filename = value + break + + c = Commit( + self, + hex_to_bin(hexsha), + author=Actor( + safe_decode(props[b"author"]), + safe_decode(props[b"author-mail"].lstrip(b"<").rstrip(b">")), + ), + authored_date=int(props[b"author-time"]), + committer=Actor( + safe_decode(props[b"committer"]), + safe_decode(props[b"committer-mail"].lstrip(b"<").rstrip(b">")), + ), + committed_date=int(props[b"committer-time"]), + ) + commits[hexsha] = c + else: + # Discard all lines until we find "filename" which is guaranteed to be + # the last line. + while True: + try: + # Will fail if we reach the EOF unexpectedly. + line = next(stream) + except StopIteration: + return + tag, value = line.split(b" ", 1) + if tag == b"filename": + orig_filename = value + break + + yield BlameEntry( + commits[hexsha], + range(lineno, lineno + num_lines), + safe_decode(orig_filename), + range(orig_lineno, orig_lineno + num_lines), + ) + + def blame( + self, + rev: Union[str, HEAD, None], + file: str, + incremental: bool = False, + rev_opts: Optional[List[str]] = None, + **kwargs: Any, + ) -> List[List[Commit | List[str | bytes] | None]] | Iterator[BlameEntry] | None: + """The blame information for the given file at the given revision. + + :param rev: + Revision specifier. If ``None``, the blame will include all the latest + uncommitted changes. Otherwise, anything successfully parsed by + :manpage:`git-rev-parse(1)` is a valid option. + + :return: + list: [git.Commit, list: []] + + A list of lists associating a :class:`~git.objects.commit.Commit` object + with a list of lines that changed within the given commit. The + :class:`~git.objects.commit.Commit` objects will be given in order of + appearance. + """ + if incremental: + return self.blame_incremental(rev, file, **kwargs) + rev_opts = rev_opts or [] + data: bytes = self.git.blame(rev, *rev_opts, "--", file, p=True, stdout_as_string=False, **kwargs) + commits: Dict[str, Commit] = {} + blames: List[List[Commit | List[str | bytes] | None]] = [] + + class InfoTD(TypedDict, total=False): + sha: str + id: str + filename: str + summary: str + author: str + author_email: str + author_date: int + committer: str + committer_email: str + committer_date: int + + info: InfoTD = {} + + keepends = True + for line_bytes in data.splitlines(keepends): + try: + line_str = line_bytes.rstrip().decode(defenc) + except UnicodeDecodeError: + firstpart = "" + parts = [] + is_binary = True + else: + # As we don't have an idea when the binary data ends, as it could + # contain multiple newlines in the process. So we rely on being able to + # decode to tell us what it is. This can absolutely fail even on text + # files, but even if it does, we should be fine treating it as binary + # instead. + parts = self.re_whitespace.split(line_str, 1) + firstpart = parts[0] + is_binary = False + # END handle decode of line + + if self.re_hexsha_only.search(firstpart): + # handles + # 634396b2f541a9f2d58b00be1a07f0c358b999b3 1 1 7 - indicates blame-data start + # 634396b2f541a9f2d58b00be1a07f0c358b999b3 2 2 - indicates + # another line of blame with the same data + digits = parts[-1].split(" ") + if len(digits) == 3: + info = {"id": firstpart} + blames.append([None, []]) + elif info["id"] != firstpart: + info = {"id": firstpart} + blames.append([commits.get(firstpart), []]) + # END blame data initialization + else: + m = self.re_author_committer_start.search(firstpart) + if m: + # handles: + # author Tom Preston-Werner + # author-mail + # author-time 1192271832 + # author-tz -0700 + # committer Tom Preston-Werner + # committer-mail + # committer-time 1192271832 + # committer-tz -0700 - IGNORED BY US + role = m.group(0) + if role == "author": + if firstpart.endswith("-mail"): + info["author_email"] = parts[-1] + elif firstpart.endswith("-time"): + info["author_date"] = int(parts[-1]) + elif role == firstpart: + info["author"] = parts[-1] + elif role == "committer": + if firstpart.endswith("-mail"): + info["committer_email"] = parts[-1] + elif firstpart.endswith("-time"): + info["committer_date"] = int(parts[-1]) + elif role == firstpart: + info["committer"] = parts[-1] + # END distinguish mail,time,name + else: + # handle + # filename lib/grit.rb + # summary add Blob + # + if firstpart.startswith("filename"): + info["filename"] = parts[-1] + elif firstpart.startswith("summary"): + info["summary"] = parts[-1] + elif firstpart == "": + if info: + sha = info["id"] + c = commits.get(sha) + if c is None: + c = Commit( + self, + hex_to_bin(sha), + author=Actor._from_string(f"{info['author']} {info['author_email']}"), + authored_date=info["author_date"], + committer=Actor._from_string(f"{info['committer']} {info['committer_email']}"), + committed_date=info["committer_date"], + ) + commits[sha] = c + blames[-1][0] = c + # END if commit objects needs initial creation + + if blames[-1][1] is not None: + line: str | bytes + if not is_binary: + if line_str and line_str[0] == "\t": + line_str = line_str[1:] + line = line_str + else: + line = line_bytes + # NOTE: We are actually parsing lines out of binary + # data, which can lead to the binary being split up + # along the newline separator. We will append this + # to the blame we are currently looking at, even + # though it should be concatenated with the last + # line we have seen. + blames[-1][1].append(line) + + info = {"id": sha} + # END if we collected commit info + # END distinguish filename,summary,rest + # END distinguish author|committer vs filename,summary,rest + # END distinguish hexsha vs other information + return blames + + @classmethod + def init( + cls, + path: Union[PathLike, None] = None, + mkdir: bool = True, + odbt: Type[GitCmdObjectDB] = GitCmdObjectDB, + expand_vars: bool = True, + **kwargs: Any, + ) -> "Repo": + """Initialize a git repository at the given path if specified. + + :param path: + The full path to the repo (traditionally ends with ``/.git``). Or + ``None``, in which case the repository will be created in the current + working directory. + + :param mkdir: + If specified, will create the repository directory if it doesn't already + exist. Creates the directory with a mode=0755. + Only effective if a path is explicitly given. + + :param odbt: + Object DataBase type - a type which is constructed by providing the + directory containing the database objects, i.e. ``.git/objects``. It will be + used to access all object data. + + :param expand_vars: + If specified, environment variables will not be escaped. This can lead to + information disclosure, allowing attackers to access the contents of + environment variables. + + :param kwargs: + Keyword arguments serving as additional options to the + :manpage:`git-init(1)` command. + + :return: + :class:`Repo` (the newly created repo) + """ + if path: + path = expand_path(path, expand_vars) + if mkdir and path and not osp.exists(path): + os.makedirs(path, 0o755) + + # git command automatically chdir into the directory + git = cls.GitCommandWrapperType(path) + git.init(**kwargs) + return cls(path, odbt=odbt) + + @classmethod + def _clone( + cls, + git: "Git", + url: PathLike, + path: PathLike, + odb_default_type: Type[GitCmdObjectDB], + progress: Union["RemoteProgress", "UpdateProgress", Callable[..., "RemoteProgress"], None] = None, + multi_options: Optional[List[str]] = None, + allow_unsafe_protocols: bool = False, + allow_unsafe_options: bool = False, + **kwargs: Any, + ) -> "Repo": + odbt = kwargs.pop("odbt", odb_default_type) + + # url may be a path and this has no effect if it is a string + url = os.fspath(url) + path = os.fspath(path) + + ## A bug win cygwin's Git, when `--bare` or `--separate-git-dir` + # it prepends the cwd or(?) the `url` into the `path, so:: + # git clone --bare /cygwin/d/foo.git C:\\Work + # becomes:: + # git clone --bare /cygwin/d/foo.git /cygwin/d/C:\\Work + # + clone_path = Git.polish_url(path) if Git.is_cygwin() and "bare" in kwargs else path + sep_dir = kwargs.get("separate_git_dir") + if sep_dir: + kwargs["separate_git_dir"] = Git.polish_url(sep_dir) + multi = None + if multi_options: + multi = shlex.split(" ".join(multi_options)) + + if not allow_unsafe_protocols: + Git.check_unsafe_protocols(url) + if not allow_unsafe_options: + Git.check_unsafe_options(options=list(kwargs.keys()), unsafe_options=cls.unsafe_git_clone_options) + if not allow_unsafe_options and multi_options: + Git.check_unsafe_options(options=multi_options, unsafe_options=cls.unsafe_git_clone_options) + + proc = git.clone( + multi, + "--", + Git.polish_url(url), + clone_path, + with_extended_output=True, + as_process=True, + v=True, + universal_newlines=True, + **add_progress(kwargs, git, progress), + ) + if progress: + handle_process_output( + proc, + None, + to_progress_instance(progress).new_message_handler(), + finalize_process, + decode_streams=False, + ) + else: + (stdout, stderr) = proc.communicate() + cmdline = getattr(proc, "args", "") + cmdline = remove_password_if_present(cmdline) + + _logger.debug("Cmd(%s)'s unused stdout: %s", cmdline, stdout) + finalize_process(proc, stderr=stderr) + + # Our git command could have a different working dir than our actual + # environment, hence we prepend its working dir if required. + if not osp.isabs(path): + path = osp.join(git._working_dir, path) if git._working_dir is not None else path + + repo = cls(path, odbt=odbt) + + # Retain env values that were passed to _clone(). + repo.git.update_environment(**git.environment()) + + # Adjust remotes - there may be operating systems which use backslashes, These + # might be given as initial paths, but when handling the config file that + # contains the remote from which we were clones, git stops liking it as it will + # escape the backslashes. Hence we undo the escaping just to be sure. + if repo.remotes: + with repo.remotes[0].config_writer as writer: + writer.set_value("url", Git.polish_url(repo.remotes[0].url)) + # END handle remote repo + return repo + + def clone( + self, + path: PathLike, + progress: Optional[CallableProgress] = None, + multi_options: Optional[List[str]] = None, + allow_unsafe_protocols: bool = False, + allow_unsafe_options: bool = False, + **kwargs: Any, + ) -> "Repo": + """Create a clone from this repository. + + :param path: + The full path of the new repo (traditionally ends with ``./.git``). + + :param progress: + See :meth:`Remote.push `. + + :param multi_options: + A list of :manpage:`git-clone(1)` options that can be provided multiple + times. + + One option per list item which is passed exactly as specified to clone. + For example:: + + [ + "--config core.filemode=false", + "--config core.ignorecase", + "--recurse-submodule=repo1_path", + "--recurse-submodule=repo2_path", + ] + + :param allow_unsafe_protocols: + Allow unsafe protocols to be used, like ``ext``. + + :param allow_unsafe_options: + Allow unsafe options to be used, like ``--upload-pack``. + + :param kwargs: + * ``odbt`` = ObjectDatabase Type, allowing to determine the object database + implementation used by the returned :class:`Repo` instance. + * All remaining keyword arguments are given to the :manpage:`git-clone(1)` + command. + + :return: + :class:`Repo` (the newly cloned repo) + """ + return self._clone( + self.git, + self.common_dir, + path, + type(self.odb), + progress, # type: ignore[arg-type] + multi_options, + allow_unsafe_protocols=allow_unsafe_protocols, + allow_unsafe_options=allow_unsafe_options, + **kwargs, + ) + + @classmethod + def clone_from( + cls, + url: PathLike, + to_path: PathLike, + progress: CallableProgress = None, + env: Optional[Mapping[str, str]] = None, + multi_options: Optional[List[str]] = None, + allow_unsafe_protocols: bool = False, + allow_unsafe_options: bool = False, + **kwargs: Any, + ) -> "Repo": + """Create a clone from the given URL. + + :param url: + Valid git url, see: https://git-scm.com/docs/git-clone#URLS + + :param to_path: + Path to which the repository should be cloned to. + + :param progress: + See :meth:`Remote.push `. + + :param env: + Optional dictionary containing the desired environment variables. + + Note: Provided variables will be used to update the execution environment + for ``git``. If some variable is not specified in `env` and is defined in + :attr:`os.environ`, value from :attr:`os.environ` will be used. If you want + to unset some variable, consider providing empty string as its value. + + :param multi_options: + See the :meth:`clone` method. + + :param allow_unsafe_protocols: + Allow unsafe protocols to be used, like ``ext``. + + :param allow_unsafe_options: + Allow unsafe options to be used, like ``--upload-pack``. + + :param kwargs: + See the :meth:`clone` method. + + :return: + :class:`Repo` instance pointing to the cloned directory. + """ + git = cls.GitCommandWrapperType(os.getcwd()) + if env is not None: + git.update_environment(**env) + return cls._clone( + git, + url, + to_path, + GitCmdObjectDB, + progress, # type: ignore[arg-type] + multi_options, + allow_unsafe_protocols=allow_unsafe_protocols, + allow_unsafe_options=allow_unsafe_options, + **kwargs, + ) + + def archive( + self, + ostream: Union[TextIO, BinaryIO], + treeish: Optional[str] = None, + prefix: Optional[str] = None, + **kwargs: Any, + ) -> Repo: + """Archive the tree at the given revision. + + :param ostream: + File-compatible stream object to which the archive will be written as bytes. + + :param treeish: + The treeish name/id, defaults to active branch. + + :param prefix: + The optional prefix to prepend to each filename in the archive. + + :param kwargs: + Additional arguments passed to :manpage:`git-archive(1)`: + + * Use the ``format`` argument to define the kind of format. Use specialized + ostreams to write any format supported by Python. + * You may specify the special ``path`` keyword, which may either be a + repository-relative path to a directory or file to place into the archive, + or a list or tuple of multiple paths. + + :raise git.exc.GitCommandError: + If something went wrong. + + :return: + self + """ + if treeish is None: + treeish = self.head.commit + if prefix and "prefix" not in kwargs: + kwargs["prefix"] = prefix + kwargs["output_stream"] = ostream + path = kwargs.pop("path", []) + path = cast(Union[PathLike, List[PathLike], Tuple[PathLike, ...]], path) + if not isinstance(path, (tuple, list)): + path = [path] + # END ensure paths is list (or tuple) + self.git.archive("--", treeish, *path, **kwargs) + return self + + def has_separate_working_tree(self) -> bool: + """ + :return: + True if our :attr:`git_dir` is not at the root of our + :attr:`working_tree_dir`, but a ``.git`` file with a platform-agnostic + symbolic link. Our :attr:`git_dir` will be wherever the ``.git`` file points + to. + + :note: + Bare repositories will always return ``False`` here. + """ + if self.bare: + return False + if self.working_tree_dir: + return osp.isfile(osp.join(self.working_tree_dir, ".git")) + else: + return False # Or raise Error? + + rev_parse = rev_parse + + def __repr__(self) -> str: + clazz = self.__class__ + return "<%s.%s %r>" % (clazz.__module__, clazz.__name__, self.git_dir) + + def currently_rebasing_on(self) -> Commit | None: + """ + :return: + The commit which is currently being replayed while rebasing. + + ``None`` if we are not currently rebasing. + """ + if self.git_dir: + rebase_head_file = osp.join(self.git_dir, "REBASE_HEAD") + if not osp.isfile(rebase_head_file): + return None + with open(rebase_head_file, "rt") as f: + content = f.readline().strip() + return self.commit(content) diff --git a/lib/python3.12/site-packages/git/repo/fun.py b/lib/python3.12/site-packages/git/repo/fun.py new file mode 100644 index 0000000000000000000000000000000000000000..3f00e60eac9c100ffc8526e6bb1ed6ec981b720f --- /dev/null +++ b/lib/python3.12/site-packages/git/repo/fun.py @@ -0,0 +1,425 @@ +# This module is part of GitPython and is released under the +# 3-Clause BSD License: https://opensource.org/license/bsd-3-clause/ + +"""General repository-related functions.""" + +from __future__ import annotations + +__all__ = [ + "rev_parse", + "is_git_dir", + "touch", + "find_submodule_git_dir", + "name_to_object", + "short_to_long", + "deref_tag", + "to_commit", + "find_worktree_git_dir", +] + +import os +import os.path as osp +from pathlib import Path +import stat +from string import digits + +from gitdb.exc import BadName, BadObject + +from git.cmd import Git +from git.exc import WorkTreeRepositoryUnsupported +from git.objects import Object +from git.refs import SymbolicReference +from git.util import cygpath, bin_to_hex, hex_to_bin + +# Typing ---------------------------------------------------------------------- + +from typing import Optional, TYPE_CHECKING, Union, cast, overload + +from git.types import AnyGitObject, Literal, PathLike + +if TYPE_CHECKING: + from git.db import GitCmdObjectDB + from git.objects import Commit, TagObject + from git.refs.reference import Reference + from git.refs.tag import Tag + + from .base import Repo + +# ---------------------------------------------------------------------------- + + +def touch(filename: str) -> str: + with open(filename, "ab"): + pass + return filename + + +def is_git_dir(d: PathLike) -> bool: + """This is taken from the git setup.c:is_git_directory function. + + :raise git.exc.WorkTreeRepositoryUnsupported: + If it sees a worktree directory. It's quite hacky to do that here, but at least + clearly indicates that we don't support it. There is the unlikely danger to + throw if we see directories which just look like a worktree dir, but are none. + """ + if osp.isdir(d): + if (osp.isdir(osp.join(d, "objects")) or "GIT_OBJECT_DIRECTORY" in os.environ) and osp.isdir( + osp.join(d, "refs") + ): + headref = osp.join(d, "HEAD") + return osp.isfile(headref) or (osp.islink(headref) and os.readlink(headref).startswith("refs")) + elif ( + osp.isfile(osp.join(d, "gitdir")) + and osp.isfile(osp.join(d, "commondir")) + and osp.isfile(osp.join(d, "gitfile")) + ): + raise WorkTreeRepositoryUnsupported(d) + return False + + +def find_worktree_git_dir(dotgit: PathLike) -> Optional[str]: + """Search for a gitdir for this worktree.""" + try: + statbuf = os.stat(dotgit) + except OSError: + return None + if not stat.S_ISREG(statbuf.st_mode): + return None + + try: + lines = Path(dotgit).read_text().splitlines() + for key, value in [line.strip().split(": ") for line in lines]: + if key == "gitdir": + return value + except ValueError: + pass + return None + + +def find_submodule_git_dir(d: PathLike) -> Optional[PathLike]: + """Search for a submodule repo.""" + if is_git_dir(d): + return d + + try: + with open(d) as fp: + content = fp.read().rstrip() + except IOError: + # It's probably not a file. + pass + else: + if content.startswith("gitdir: "): + path = content[8:] + + if Git.is_cygwin(): + # Cygwin creates submodules prefixed with `/cygdrive/...`. + # Cygwin git understands Cygwin paths much better than Windows ones. + # Also the Cygwin tests are assuming Cygwin paths. + path = cygpath(path) + if not osp.isabs(path): + path = osp.normpath(osp.join(osp.dirname(d), path)) + return find_submodule_git_dir(path) + # END handle exception + return None + + +def short_to_long(odb: "GitCmdObjectDB", hexsha: str) -> Optional[bytes]: + """ + :return: + Long hexadecimal sha1 from the given less than 40 byte hexsha, or ``None`` if no + candidate could be found. + + :param hexsha: + hexsha with less than 40 bytes. + """ + try: + return bin_to_hex(odb.partial_to_complete_sha_hex(hexsha)) + except BadObject: + return None + # END exception handling + + +@overload +def name_to_object(repo: "Repo", name: str, return_ref: Literal[False] = ...) -> AnyGitObject: ... + + +@overload +def name_to_object(repo: "Repo", name: str, return_ref: Literal[True]) -> Union[AnyGitObject, SymbolicReference]: ... + + +def name_to_object(repo: "Repo", name: str, return_ref: bool = False) -> Union[AnyGitObject, SymbolicReference]: + """ + :return: + Object specified by the given name - hexshas (short and long) as well as + references are supported. + + :param return_ref: + If ``True``, and name specifies a reference, we will return the reference + instead of the object. Otherwise it will raise :exc:`~gitdb.exc.BadObject` or + :exc:`~gitdb.exc.BadName`. + """ + hexsha: Union[None, str, bytes] = None + + # Is it a hexsha? Try the most common ones, which is 7 to 40. + if repo.re_hexsha_shortened.match(name): + if len(name) != 40: + # Find long sha for short sha. + hexsha = short_to_long(repo.odb, name) + else: + hexsha = name + # END handle short shas + # END find sha if it matches + + # If we couldn't find an object for what seemed to be a short hexsha, try to find it + # as reference anyway, it could be named 'aaa' for instance. + if hexsha is None: + for base in ( + "%s", + "refs/%s", + "refs/tags/%s", + "refs/heads/%s", + "refs/remotes/%s", + "refs/remotes/%s/HEAD", + ): + try: + hexsha = SymbolicReference.dereference_recursive(repo, base % name) + if return_ref: + return SymbolicReference(repo, base % name) + # END handle symbolic ref + break + except ValueError: + pass + # END for each base + # END handle hexsha + + # Didn't find any ref, this is an error. + if return_ref: + raise BadObject("Couldn't find reference named %r" % name) + # END handle return ref + + # Tried everything ? fail. + if hexsha is None: + raise BadName(name) + # END assert hexsha was found + + return Object.new_from_sha(repo, hex_to_bin(hexsha)) + + +def deref_tag(tag: "Tag") -> AnyGitObject: + """Recursively dereference a tag and return the resulting object.""" + while True: + try: + tag = tag.object + except AttributeError: + break + # END dereference tag + return tag + + +def to_commit(obj: Object) -> "Commit": + """Convert the given object to a commit if possible and return it.""" + if obj.type == "tag": + obj = deref_tag(obj) + + if obj.type != "commit": + raise ValueError("Cannot convert object %r to type commit" % obj) + # END verify type + return obj + + +def rev_parse(repo: "Repo", rev: str) -> AnyGitObject: + """Parse a revision string. Like :manpage:`git-rev-parse(1)`. + + :return: + `~git.objects.base.Object` at the given revision. + + This may be any type of git object: + + * :class:`Commit ` + * :class:`TagObject ` + * :class:`Tree ` + * :class:`Blob ` + + :param rev: + :manpage:`git-rev-parse(1)`-compatible revision specification as string. + Please see :manpage:`git-rev-parse(1)` for details. + + :raise gitdb.exc.BadObject: + If the given revision could not be found. + + :raise ValueError: + If `rev` couldn't be parsed. + + :raise IndexError: + If an invalid reflog index is specified. + """ + # Are we in colon search mode? + if rev.startswith(":/"): + # Colon search mode + raise NotImplementedError("commit by message search (regex)") + # END handle search + + obj: Optional[AnyGitObject] = None + ref = None + output_type = "commit" + start = 0 + parsed_to = 0 + lr = len(rev) + while start < lr: + if rev[start] not in "^~:@": + start += 1 + continue + # END handle start + + token = rev[start] + + if obj is None: + # token is a rev name. + if start == 0: + ref = repo.head.ref + else: + if token == "@": + ref = cast("Reference", name_to_object(repo, rev[:start], return_ref=True)) + else: + obj = name_to_object(repo, rev[:start]) + # END handle token + # END handle refname + else: + if ref is not None: + obj = ref.commit + # END handle ref + # END initialize obj on first token + + start += 1 + + # Try to parse {type}. + if start < lr and rev[start] == "{": + end = rev.find("}", start) + if end == -1: + raise ValueError("Missing closing brace to define type in %s" % rev) + output_type = rev[start + 1 : end] # Exclude brace. + + # Handle type. + if output_type == "commit": + obj = cast("TagObject", obj) + if obj and obj.type == "tag": + obj = deref_tag(obj) + else: + # Cannot do anything for non-tags. + pass + # END handle tag + elif output_type == "tree": + try: + obj = cast(AnyGitObject, obj) + obj = to_commit(obj).tree + except (AttributeError, ValueError): + pass # Error raised later. + # END exception handling + elif output_type in ("", "blob"): + obj = cast("TagObject", obj) + if obj and obj.type == "tag": + obj = deref_tag(obj) + else: + # Cannot do anything for non-tags. + pass + # END handle tag + elif token == "@": + # try single int + assert ref is not None, "Require Reference to access reflog" + revlog_index = None + try: + # Transform reversed index into the format of our revlog. + revlog_index = -(int(output_type) + 1) + except ValueError as e: + # TODO: Try to parse the other date options, using parse_date maybe. + raise NotImplementedError("Support for additional @{...} modes not implemented") from e + # END handle revlog index + + try: + entry = ref.log_entry(revlog_index) + except IndexError as e: + raise IndexError("Invalid revlog index: %i" % revlog_index) from e + # END handle index out of bound + + obj = Object.new_from_sha(repo, hex_to_bin(entry.newhexsha)) + + # Make it pass the following checks. + output_type = "" + else: + raise ValueError("Invalid output type: %s ( in %s )" % (output_type, rev)) + # END handle output type + + # Empty output types don't require any specific type, its just about + # dereferencing tags. + if output_type and obj and obj.type != output_type: + raise ValueError("Could not accommodate requested object type %r, got %s" % (output_type, obj.type)) + # END verify output type + + start = end + 1 # Skip brace. + parsed_to = start + continue + # END parse type + + # Try to parse a number. + num = 0 + if token != ":": + found_digit = False + while start < lr: + if rev[start] in digits: + num = num * 10 + int(rev[start]) + start += 1 + found_digit = True + else: + break + # END handle number + # END number parse loop + + # No explicit number given, 1 is the default. It could be 0 though. + if not found_digit: + num = 1 + # END set default num + # END number parsing only if non-blob mode + + parsed_to = start + # Handle hierarchy walk. + try: + obj = cast(AnyGitObject, obj) + if token == "~": + obj = to_commit(obj) + for _ in range(num): + obj = obj.parents[0] + # END for each history item to walk + elif token == "^": + obj = to_commit(obj) + # Must be n'th parent. + if num: + obj = obj.parents[num - 1] + elif token == ":": + if obj.type != "tree": + obj = obj.tree + # END get tree type + obj = obj[rev[start:]] + parsed_to = lr + else: + raise ValueError("Invalid token: %r" % token) + # END end handle tag + except (IndexError, AttributeError) as e: + raise BadName( + f"Invalid revision spec '{rev}' - not enough parent commits to reach '{token}{int(num)}'" + ) from e + # END exception handling + # END parse loop + + # Still no obj? It's probably a simple name. + if obj is None: + obj = name_to_object(repo, rev) + parsed_to = lr + # END handle simple name + + if obj is None: + raise ValueError("Revision specifier could not be parsed: %s" % rev) + + if parsed_to != lr: + raise ValueError("Didn't consume complete rev spec %s, consumed part: %s" % (rev, rev[:parsed_to])) + + return obj diff --git a/lib/python3.12/site-packages/matplotlib_inline/__init__.py b/lib/python3.12/site-packages/matplotlib_inline/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d275175ecd970d64aa94db01b833fd4ebf7ce1ea --- /dev/null +++ b/lib/python3.12/site-packages/matplotlib_inline/__init__.py @@ -0,0 +1,8 @@ +from . import backend_inline, config # noqa + +__version__ = "0.2.1" + +# we can't ''.join(...) otherwise finding the version number at build time requires +# import which introduces IPython and matplotlib at build time, and thus circular +# dependencies. +version_info = tuple(int(s) for s in __version__.split(".")[:3]) diff --git a/lib/python3.12/site-packages/matplotlib_inline/__pycache__/config.cpython-312.pyc b/lib/python3.12/site-packages/matplotlib_inline/__pycache__/config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..535e2cde23802caa6a02ce7fb9b08417f2792f83 Binary files /dev/null and b/lib/python3.12/site-packages/matplotlib_inline/__pycache__/config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/matplotlib_inline/backend_inline.py b/lib/python3.12/site-packages/matplotlib_inline/backend_inline.py new file mode 100644 index 0000000000000000000000000000000000000000..a798c562a7fa6c07d32e6e0dd742d0e6c711b699 --- /dev/null +++ b/lib/python3.12/site-packages/matplotlib_inline/backend_inline.py @@ -0,0 +1,336 @@ +"""A matplotlib backend for publishing figures via display_data""" + +# Copyright (c) IPython Development Team. +# Distributed under the terms of the BSD 3-Clause License. + +import matplotlib +from IPython.core.getipython import get_ipython +from IPython.core.interactiveshell import InteractiveShell +from IPython.core.pylabtools import select_figure_formats +from IPython.display import display +from matplotlib import colors +from matplotlib._pylab_helpers import Gcf +from matplotlib.backends import backend_agg +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure + +from .config import InlineBackend + + +def new_figure_manager(num, *args, FigureClass=Figure, **kwargs): + """ + Return a new figure manager for a new figure instance. + + This function is part of the API expected by Matplotlib backends. + """ + return new_figure_manager_given_figure(num, FigureClass(*args, **kwargs)) + + +def new_figure_manager_given_figure(num, figure): + """ + Return a new figure manager for a given figure instance. + + This function is part of the API expected by Matplotlib backends. + """ + manager = backend_agg.new_figure_manager_given_figure(num, figure) + + # Hack: matplotlib FigureManager objects in interacive backends (at least + # in some of them) monkeypatch the figure object and add a .show() method + # to it. This applies the same monkeypatch in order to support user code + # that might expect `.show()` to be part of the official API of figure + # objects. For further reference: + # https://github.com/ipython/ipython/issues/1612 + # https://github.com/matplotlib/matplotlib/issues/835 + + if not hasattr(figure, "show"): + # Queue up `figure` for display + figure.show = lambda *a: display( + figure, metadata=_fetch_figure_metadata(figure) + ) + + # If matplotlib was manually set to non-interactive mode, this function + # should be a no-op (otherwise we'll generate duplicate plots, since a user + # who set ioff() manually expects to make separate draw/show calls). + if not matplotlib.is_interactive(): + return manager + + # ensure current figure will be drawn, and each subsequent call + # of draw_if_interactive() moves the active figure to ensure it is + # drawn last + try: + show._to_draw.remove(figure) + except ValueError: + # ensure it only appears in the draw list once + pass + # Queue up the figure for drawing in next show() call + show._to_draw.append(figure) + show._draw_called = True + + return manager + + +def show(close=None, block=None): + """Show all figures as SVG/PNG payloads sent to the IPython clients. + + Parameters + ---------- + close : bool, optional + If true, a ``plt.close('all')`` call is automatically issued after + sending all the figures. If this is set, the figures will entirely + removed from the internal list of figures. + block : Not used. + The `block` parameter is a Matplotlib experimental parameter. + We accept it in the function signature for compatibility with other + backends. + """ + if close is None: + close = InlineBackend.instance().close_figures + try: + for figure_manager in Gcf.get_all_fig_managers(): + display( + figure_manager.canvas.figure, + metadata=_fetch_figure_metadata(figure_manager.canvas.figure), + ) + finally: + show._to_draw = [] + # only call close('all') if any to close + # close triggers gc.collect, which can be slow + if close and Gcf.get_all_fig_managers(): + matplotlib.pyplot.close("all") + + +# This flag will be reset by draw_if_interactive when called +show._draw_called = False # type: ignore[attr-defined] +# list of figures to draw when flush_figures is called +show._to_draw = [] # type: ignore[attr-defined] + + +def flush_figures(): + """Send all figures that changed + + This is meant to be called automatically and will call show() if, during + prior code execution, there had been any calls to draw_if_interactive. + + This function is meant to be used as a post_execute callback in IPython, + so user-caused errors are handled with showtraceback() instead of being + allowed to raise. If this function is not called from within IPython, + then these exceptions will raise. + """ + if not show._draw_called: + return + + try: + if InlineBackend.instance().close_figures: + # ignore the tracking, just draw and close all figures + try: + return show(True) + except Exception as e: + # safely show traceback if in IPython, else raise + ip = get_ipython() + if ip is None: + raise e + else: + ip.showtraceback() + return + + # exclude any figures that were closed: + active = set([fm.canvas.figure for fm in Gcf.get_all_fig_managers()]) + for fig in [fig for fig in show._to_draw if fig in active]: + try: + display(fig, metadata=_fetch_figure_metadata(fig)) + except Exception as e: + # safely show traceback if in IPython, else raise + ip = get_ipython() + if ip is None: + raise e + else: + ip.showtraceback() + return + finally: + # clear flags for next round + show._to_draw = [] + show._draw_called = False + + +# Changes to matplotlib in version 1.2 requires a mpl backend to supply a default +# figurecanvas. This is set here to a Agg canvas +# See https://github.com/matplotlib/matplotlib/pull/1125 +FigureCanvas = FigureCanvasAgg + + +def configure_inline_support(shell, backend): + """Configure an IPython shell object for matplotlib use. + + Parameters + ---------- + shell : InteractiveShell instance + + backend : matplotlib backend + """ + # If using our svg payload backend, register the post-execution + # function that will pick up the results for display. This can only be + # done with access to the real shell object. + + cfg = InlineBackend.instance(parent=shell) + cfg.shell = shell + if cfg not in shell.configurables: + shell.configurables.append(cfg) + + if backend in ("inline", "module://matplotlib_inline.backend_inline"): + shell.events.register("post_execute", flush_figures) + + # Save rcParams that will be overwrittern + shell._saved_rcParams = {} + for k in cfg.rc: + shell._saved_rcParams[k] = matplotlib.rcParams[k] + # load inline_rc + matplotlib.rcParams.update(cfg.rc) + new_backend_name = "inline" + else: + try: + shell.events.unregister("post_execute", flush_figures) + except ValueError: + pass + if hasattr(shell, "_saved_rcParams"): + matplotlib.rcParams.update(shell._saved_rcParams) + del shell._saved_rcParams + new_backend_name = "other" + + # only enable the formats once -> don't change the enabled formats (which the user may + # has changed) when getting another "%matplotlib inline" call. + # See https://github.com/ipython/ipykernel/issues/29 + cur_backend = getattr(configure_inline_support, "current_backend", "unset") + if new_backend_name != cur_backend: + # Setup the default figure format + select_figure_formats(shell, cfg.figure_formats, **cfg.print_figure_kwargs) + configure_inline_support.current_backend = new_backend_name + + +def _enable_matplotlib_integration(): + """Enable extra IPython matplotlib integration when we are loaded as the matplotlib backend.""" + ip = get_ipython() + + import matplotlib + + if matplotlib.__version_info__ >= (3, 10): + backend = matplotlib.get_backend(auto_select=False) + else: + backend = matplotlib.rcParams._get("backend") + + if ip and backend in ("inline", "module://matplotlib_inline.backend_inline"): + from IPython.core.pylabtools import activate_matplotlib + + try: + activate_matplotlib(backend) + configure_inline_support(ip, backend) + except (ImportError, AttributeError): + # bugs may cause a circular import on Python 2 + def configure_once(*args): + activate_matplotlib(backend) + configure_inline_support(ip, backend) + ip.events.unregister("post_run_cell", configure_once) + + ip.events.register("post_run_cell", configure_once) + + +_enable_matplotlib_integration() + + +def _fetch_figure_metadata(fig): + """Get some metadata to help with displaying a figure.""" + # determine if a background is needed for legibility + if _is_transparent(fig.get_facecolor()): + # the background is transparent + ticksLight = _is_light( + [ + label.get_color() + for axes in fig.axes + for axis in (axes.xaxis, axes.yaxis) + for label in axis.get_ticklabels() + ] + ) + if ticksLight.size and (ticksLight == ticksLight[0]).all(): + # there are one or more tick labels, all with the same lightness + return {"needs_background": "dark" if ticksLight[0] else "light"} + + return None + + +def _is_light(color): + """Determines if a color (or each of a sequence of colors) is light (as + opposed to dark). Based on ITU BT.601 luminance formula (see + https://stackoverflow.com/a/596241).""" + rgbaArr = colors.to_rgba_array(color) + return rgbaArr[:, :3].dot((0.299, 0.587, 0.114)) > 0.5 + + +def _is_transparent(color): + """Determine transparency from alpha.""" + rgba = colors.to_rgba(color) + return rgba[3] < 0.5 + + +def set_matplotlib_formats(*formats, **kwargs): + """Select figure formats for the inline backend. Optionally pass quality for JPEG. + + For example, this enables PNG and JPEG output with a JPEG quality of 90%:: + + In [1]: set_matplotlib_formats('png', 'jpeg', + pil_kwargs={'quality': 90}) + + To set this in your notebook by `%config` magic:: + + In [1]: %config InlineBackend.figure_formats = {'png', 'jpeg'} + %config InlineBackend.print_figure_kwargs = \\ + {'pil_kwargs': {'quality' : 90}} + + To set this in your config files use the following:: + + c.InlineBackend.figure_formats = {'png', 'jpeg'} + c.InlineBackend.print_figure_kwargs.update({ + 'pil_kwargs': {'quality' : 90} + }) + + Parameters + ---------- + *formats : strs + One or more figure formats to enable: 'png', 'retina', 'jpeg', 'svg', 'pdf'. + **kwargs + Keyword args will be relayed to ``figure.canvas.print_figure``. + + In addition, see the docstrings of `plt.savefig()`, + `matplotlib.figure.Figure.savefig()`, `PIL.Image.Image.save()` and + :ref:`Pillow Image file formats `. + """ + # build kwargs, starting with InlineBackend config + cfg = InlineBackend.instance() + kw = {} + kw.update(cfg.print_figure_kwargs) + kw.update(**kwargs) + shell = InteractiveShell.instance() + select_figure_formats(shell, formats, **kw) + + +def set_matplotlib_close(close=True): + """Set whether the inline backend closes all figures automatically or not. + + By default, the inline backend used in the IPython Notebook will close all + matplotlib figures automatically after each cell is run. This means that + plots in different cells won't interfere. Sometimes, you may want to make + a plot in one cell and then refine it in later cells. This can be accomplished + by:: + + In [1]: set_matplotlib_close(False) + + To set this in your config files use the following:: + + c.InlineBackend.close_figures = False + + Parameters + ---------- + close : bool + Should all matplotlib figures be automatically closed after each cell is + run? + """ + cfg = InlineBackend.instance() + cfg.close_figures = close diff --git a/lib/python3.12/site-packages/matplotlib_inline/config.py b/lib/python3.12/site-packages/matplotlib_inline/config.py new file mode 100644 index 0000000000000000000000000000000000000000..ce9be032fbd3af9563facf42d5b950990b03317e --- /dev/null +++ b/lib/python3.12/site-packages/matplotlib_inline/config.py @@ -0,0 +1,109 @@ +"""Configurable for configuring the IPython inline backend + +This module does not import anything from matplotlib. +""" + +# Copyright (c) IPython Development Team. +# Distributed under the terms of the BSD 3-Clause License. + +from traitlets import Bool, Dict, Instance, Set, TraitError, Unicode +from traitlets.config.configurable import SingletonConfigurable + + +# Configurable for inline backend options +def pil_available(): + """Test if PIL/Pillow is available""" + out = False + try: + from PIL import Image # noqa + + out = True + except ImportError: + pass + return out + + +# Inherit from InlineBackendConfig for deprecation purposes +class InlineBackendConfig(SingletonConfigurable): + pass + + +class InlineBackend(InlineBackendConfig): + """An object to store configuration of the inline backend.""" + + # While we are deprecating overriding matplotlib defaults out of the + # box, this structure should remain here (empty) for API compatibility + # and the use of other tools that may need it. Specifically Spyder takes + # advantage of it. + # See https://github.com/ipython/ipython/issues/10383 for details. + rc = Dict( + {}, + help="""Dict to manage matplotlib configuration defaults in the inline + backend. As of v0.1.4 IPython/Jupyter do not override defaults out of + the box, but third-party tools may use it to manage rc data. To change + personal defaults for matplotlib, use matplotlib's configuration + tools, or customize this class in your `ipython_config.py` file for + IPython/Jupyter-specific usage.""", + ).tag(config=True) + + figure_formats = Set( + {"png"}, + help="""A set of figure formats to enable: 'png', + 'retina', 'jpeg', 'svg', 'pdf'.""", + ).tag(config=True) + + def _update_figure_formatters(self): + if self.shell is not None: + from IPython.core.pylabtools import select_figure_formats + + select_figure_formats( + self.shell, self.figure_formats, **self.print_figure_kwargs + ) + + def _figure_formats_changed(self, name, old, new): + if "jpg" in new or "jpeg" in new: + if not pil_available(): + raise TraitError("Requires PIL/Pillow for JPG figures") + self._update_figure_formatters() + + figure_format = Unicode( + help="""The figure format to enable (deprecated + use `figure_formats` instead)""" + ).tag(config=True) + + def _figure_format_changed(self, name, old, new): + if new: + self.figure_formats = {new} + + print_figure_kwargs = Dict( + {"bbox_inches": "tight"}, + help="""Extra kwargs to be passed to fig.canvas.print_figure. + + Logical examples include: bbox_inches, pil_kwargs, etc. In addition, + see the docstrings of `set_matplotlib_formats`. + """, + ).tag(config=True) + _print_figure_kwargs_changed = _update_figure_formatters + + close_figures = Bool( + True, + help="""Close all figures at the end of each cell. + + When True, ensures that each cell starts with no active figures, but it + also means that one must keep track of references in order to edit or + redraw figures in subsequent cells. This mode is ideal for the notebook, + where residual plots from other cells might be surprising. + + When False, one must call figure() to create new figures. This means + that gcf() and getfigs() can reference figures created in other cells, + and the active figure can continue to be edited with pylab/pyplot + methods that reference the current active figure. This mode facilitates + iterative editing of figures, and behaves most consistently with + other matplotlib backends, but figure barriers between cells must + be explicit. + """, + ).tag(config=True) + + shell = Instance( + "IPython.core.interactiveshell.InteractiveShellABC", allow_none=True + ) diff --git a/lib/python3.12/site-packages/matplotlib_inline/py.typed b/lib/python3.12/site-packages/matplotlib_inline/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/pure_eval/__init__.py b/lib/python3.12/site-packages/pure_eval/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ec4ba7350307f19e2ab021cc77fa598941f30414 --- /dev/null +++ b/lib/python3.12/site-packages/pure_eval/__init__.py @@ -0,0 +1,17 @@ +from .core import Evaluator, CannotEval, group_expressions, is_expression_interesting +from .my_getattr_static import getattr_static + +try: + from .version import __version__ +except ImportError: + # version.py is auto-generated with the git tag when building + __version__ = "???" + +__all__ = [ + "Evaluator", + "CannotEval", + "group_expressions", + "is_expression_interesting", + "getattr_static", + "__version__", +] diff --git a/lib/python3.12/site-packages/pure_eval/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/pure_eval/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24f017e6c414c8c1637c3afaf41cde56780aa252 Binary files /dev/null and b/lib/python3.12/site-packages/pure_eval/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pure_eval/__pycache__/core.cpython-312.pyc b/lib/python3.12/site-packages/pure_eval/__pycache__/core.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..22196ac97ed02df656675bbda2b1aa98d4961843 Binary files /dev/null and b/lib/python3.12/site-packages/pure_eval/__pycache__/core.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pure_eval/__pycache__/my_getattr_static.cpython-312.pyc b/lib/python3.12/site-packages/pure_eval/__pycache__/my_getattr_static.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d3a3d9a3f37acaaf0f4b66a7acd997273a31c2fe Binary files /dev/null and b/lib/python3.12/site-packages/pure_eval/__pycache__/my_getattr_static.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pure_eval/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/pure_eval/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..194eebc0621bdcfccfb8565709fb5b0c952ad8b9 Binary files /dev/null and b/lib/python3.12/site-packages/pure_eval/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pure_eval/__pycache__/version.cpython-312.pyc b/lib/python3.12/site-packages/pure_eval/__pycache__/version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..10b24aec401b043774d38e301c135ffddd9d11a0 Binary files /dev/null and b/lib/python3.12/site-packages/pure_eval/__pycache__/version.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/pure_eval/core.py b/lib/python3.12/site-packages/pure_eval/core.py new file mode 100644 index 0000000000000000000000000000000000000000..748f0518d4f941c6edb1cd95bc4273ded8f14455 --- /dev/null +++ b/lib/python3.12/site-packages/pure_eval/core.py @@ -0,0 +1,449 @@ +import ast +import builtins +import operator +from collections import ChainMap, OrderedDict, deque +from contextlib import suppress +from types import FrameType +from typing import Any, Tuple, Iterable, List, Mapping, Dict, Union, Set + +from pure_eval.my_getattr_static import getattr_static +from pure_eval.utils import ( + CannotEval, + has_ast_name, + copy_ast_without_context, + is_standard_types, + of_standard_types, + is_any, + of_type, + ensure_dict, +) + + +class Evaluator: + def __init__(self, names: Mapping[str, Any]): + """ + Construct a new evaluator with the given variable names. + This is a low level API, typically you will use `Evaluator.from_frame(frame)`. + + :param names: a mapping from variable names to their values. + """ + + self.names = names + self._cache = {} # type: Dict[ast.expr, Any] + + @classmethod + def from_frame(cls, frame: FrameType) -> 'Evaluator': + """ + Construct an Evaluator that can look up variables from the given frame. + + :param frame: a frame object, e.g. from a traceback or `inspect.currentframe().f_back`. + """ + + return cls(ChainMap( + ensure_dict(frame.f_locals), + ensure_dict(frame.f_globals), + ensure_dict(frame.f_builtins), + )) + + def __getitem__(self, node: ast.expr) -> Any: + """ + Find the value of the given node. + If it cannot be evaluated safely, this raises `CannotEval`. + The result is cached either way. + + :param node: an AST expression to evaluate + :return: the value of the node + """ + + if not isinstance(node, ast.expr): + raise TypeError("node should be an ast.expr, not {!r}".format(type(node).__name__)) + + with suppress(KeyError): + result = self._cache[node] + if result is CannotEval: + raise CannotEval + else: + return result + + try: + self._cache[node] = result = self._handle(node) + return result + except CannotEval: + self._cache[node] = CannotEval + raise + + def _handle(self, node: ast.expr) -> Any: + """ + This is where the evaluation happens. + Users should use `__getitem__`, i.e. `evaluator[node]`, + as it provides caching. + + :param node: an AST expression to evaluate + :return: the value of the node + """ + + with suppress(Exception): + return ast.literal_eval(node) + + if isinstance(node, ast.Name): + try: + return self.names[node.id] + except KeyError: + raise CannotEval + elif isinstance(node, ast.Attribute): + value = self[node.value] + attr = node.attr + return getattr_static(value, attr) + elif isinstance(node, ast.Subscript): + return self._handle_subscript(node) + elif isinstance(node, (ast.List, ast.Tuple, ast.Set, ast.Dict)): + return self._handle_container(node) + elif isinstance(node, ast.UnaryOp): + return self._handle_unary(node) + elif isinstance(node, ast.BinOp): + return self._handle_binop(node) + elif isinstance(node, ast.BoolOp): + return self._handle_boolop(node) + elif isinstance(node, ast.Compare): + return self._handle_compare(node) + elif isinstance(node, ast.Call): + return self._handle_call(node) + raise CannotEval + + def _handle_call(self, node): + if node.keywords: + raise CannotEval + func = self[node.func] + args = [self[arg] for arg in node.args] + + if ( + is_any( + func, + slice, + int, + range, + round, + complex, + list, + tuple, + abs, + hex, + bin, + oct, + bool, + ord, + float, + len, + chr, + ) + or len(args) == 0 + and is_any(func, set, dict, str, frozenset, bytes, bytearray, object) + or len(args) >= 2 + and is_any(func, str, divmod, bytes, bytearray, pow) + ): + args = [ + of_standard_types(arg, check_dict_values=False, deep=False) + for arg in args + ] + try: + return func(*args) + except Exception as e: + raise CannotEval from e + + if len(args) == 1: + arg = args[0] + if is_any(func, id, type): + try: + return func(arg) + except Exception as e: + raise CannotEval from e + if is_any(func, all, any, sum): + of_type(arg, tuple, frozenset, list, set, dict, OrderedDict, deque) + for x in arg: + of_standard_types(x, check_dict_values=False, deep=False) + try: + return func(arg) + except Exception as e: + raise CannotEval from e + + if is_any( + func, sorted, min, max, hash, set, dict, ascii, str, repr, frozenset + ): + of_standard_types(arg, check_dict_values=True, deep=True) + try: + return func(arg) + except Exception as e: + raise CannotEval from e + raise CannotEval + + def _handle_compare(self, node): + left = self[node.left] + result = True + + for op, right in zip(node.ops, node.comparators): + right = self[right] + + op_type = type(op) + op_func = { + ast.Eq: operator.eq, + ast.NotEq: operator.ne, + ast.Lt: operator.lt, + ast.LtE: operator.le, + ast.Gt: operator.gt, + ast.GtE: operator.ge, + ast.Is: operator.is_, + ast.IsNot: operator.is_not, + ast.In: (lambda a, b: a in b), + ast.NotIn: (lambda a, b: a not in b), + }[op_type] + + if op_type not in (ast.Is, ast.IsNot): + of_standard_types(left, check_dict_values=False, deep=True) + of_standard_types(right, check_dict_values=False, deep=True) + + try: + result = op_func(left, right) + except Exception as e: + raise CannotEval from e + if not result: + return result + left = right + + return result + + def _handle_boolop(self, node): + left = of_standard_types( + self[node.values[0]], check_dict_values=False, deep=False + ) + + for right in node.values[1:]: + # We need short circuiting so that the whole operation can be evaluated + # even if the right operand can't + if isinstance(node.op, ast.Or): + left = left or of_standard_types( + self[right], check_dict_values=False, deep=False + ) + else: + assert isinstance(node.op, ast.And) + left = left and of_standard_types( + self[right], check_dict_values=False, deep=False + ) + return left + + def _handle_binop(self, node): + op_type = type(node.op) + op = { + ast.Add: operator.add, + ast.Sub: operator.sub, + ast.Mult: operator.mul, + ast.Div: operator.truediv, + ast.FloorDiv: operator.floordiv, + ast.Mod: operator.mod, + ast.Pow: operator.pow, + ast.LShift: operator.lshift, + ast.RShift: operator.rshift, + ast.BitOr: operator.or_, + ast.BitXor: operator.xor, + ast.BitAnd: operator.and_, + }.get(op_type) + if not op: + raise CannotEval + left = self[node.left] + hash_type = is_any(type(left), set, frozenset, dict, OrderedDict) + left = of_standard_types(left, check_dict_values=False, deep=hash_type) + formatting = type(left) in (str, bytes) and op_type == ast.Mod + + right = of_standard_types( + self[node.right], + check_dict_values=formatting, + deep=formatting or hash_type, + ) + try: + return op(left, right) + except Exception as e: + raise CannotEval from e + + def _handle_unary(self, node: ast.UnaryOp): + value = of_standard_types( + self[node.operand], check_dict_values=False, deep=False + ) + op_type = type(node.op) + op = { + ast.USub: operator.neg, + ast.UAdd: operator.pos, + ast.Not: operator.not_, + ast.Invert: operator.invert, + }[op_type] + try: + return op(value) + except Exception as e: + raise CannotEval from e + + def _handle_subscript(self, node): + value = self[node.value] + of_standard_types( + value, check_dict_values=False, deep=is_any(type(value), dict, OrderedDict) + ) + index = node.slice + if isinstance(index, ast.Slice): + index = slice( + *[ + None if p is None else self[p] + for p in [index.lower, index.upper, index.step] + ] + ) + elif isinstance(index, ast.ExtSlice): + raise CannotEval + else: + if isinstance(index, ast.Index): + index = index.value + index = self[index] + of_standard_types(index, check_dict_values=False, deep=True) + + try: + return value[index] + except Exception: + raise CannotEval + + def _handle_container( + self, + node: Union[ast.List, ast.Tuple, ast.Set, ast.Dict] + ) -> Union[List, Tuple, Set, Dict]: + """Handle container nodes, including List, Set, Tuple and Dict""" + if isinstance(node, ast.Dict): + elts = node.keys + if None in elts: # ** unpacking inside {}, not yet supported + raise CannotEval + else: + elts = node.elts + elts = [self[elt] for elt in elts] + if isinstance(node, ast.List): + return elts + if isinstance(node, ast.Tuple): + return tuple(elts) + + # Set and Dict + if not all( + is_standard_types(elt, check_dict_values=False, deep=True) for elt in elts + ): + raise CannotEval + + if isinstance(node, ast.Set): + try: + return set(elts) + except TypeError: + raise CannotEval + + assert isinstance(node, ast.Dict) + + pairs = [(elt, self[val]) for elt, val in zip(elts, node.values)] + try: + return dict(pairs) + except TypeError: + raise CannotEval + + def find_expressions(self, root: ast.AST) -> Iterable[Tuple[ast.expr, Any]]: + """ + Find all expressions in the given tree that can be safely evaluated. + This is a low level API, typically you will use `interesting_expressions_grouped`. + + :param root: any AST node + :return: generator of pairs (tuples) of expression nodes and their corresponding values. + """ + + for node in ast.walk(root): + if not isinstance(node, ast.expr): + continue + + try: + value = self[node] + except CannotEval: + continue + + yield node, value + + def interesting_expressions_grouped(self, root: ast.AST) -> List[Tuple[List[ast.expr], Any]]: + """ + Find all interesting expressions in the given tree that can be safely evaluated, + grouping equivalent nodes together. + + For more control and details, see: + - Evaluator.find_expressions + - is_expression_interesting + - group_expressions + + :param root: any AST node + :return: A list of pairs (tuples) containing: + - A list of equivalent AST expressions + - The value of the first expression node + (which should be the same for all nodes, unless threads are involved) + """ + + return group_expressions( + pair + for pair in self.find_expressions(root) + if is_expression_interesting(*pair) + ) + + +def is_expression_interesting(node: ast.expr, value: Any) -> bool: + """ + Determines if an expression is potentially interesting, at least in my opinion. + Returns False for the following expressions whose value is generally obvious: + - Literals (e.g. 123, 'abc', [1, 2, 3], {'a': (), 'b': ([1, 2], [3])}) + - Variables or attributes whose name is equal to the value's __name__. + For example, a function `def foo(): ...` is not interesting when referred to + as `foo` as it usually would, but `bar` can be interesting if `bar is foo`. + Similarly the method `self.foo` is not interesting. + - Builtins (e.g. `len`) referred to by their usual name. + + This is a low level API, typically you will use `interesting_expressions_grouped`. + + :param node: an AST expression + :param value: the value of the node + :return: a boolean: True if the expression is interesting, False otherwise + """ + + with suppress(ValueError): + ast.literal_eval(node) + return False + + # TODO exclude inner modules, e.g. numpy.random.__name__ == 'numpy.random' != 'random' + # TODO exclude common module abbreviations, e.g. numpy as np, pandas as pd + if has_ast_name(value, node): + return False + + if ( + isinstance(node, ast.Name) + and getattr(builtins, node.id, object()) is value + ): + return False + + return True + + +def group_expressions(expressions: Iterable[Tuple[ast.expr, Any]]) -> List[Tuple[List[ast.expr], Any]]: + """ + Organise expression nodes and their values such that equivalent nodes are together. + Two nodes are considered equivalent if they have the same structure, + ignoring context (Load, Store, or Delete) and location (lineno, col_offset). + For example, this will group together the same variable name mentioned multiple times in an expression. + + This will not check the values of the nodes. Equivalent nodes should have the same values, + unless threads are involved. + + This is a low level API, typically you will use `interesting_expressions_grouped`. + + :param expressions: pairs of AST expressions and their values, as obtained from + `Evaluator.find_expressions`, or `(node, evaluator[node])`. + :return: A list of pairs (tuples) containing: + - A list of equivalent AST expressions + - The value of the first expression node + (which should be the same for all nodes, unless threads are involved) + """ + + result = {} + for node, value in expressions: + dump = ast.dump(copy_ast_without_context(node)) + result.setdefault(dump, ([], value))[0].append(node) + return list(result.values()) diff --git a/lib/python3.12/site-packages/pure_eval/my_getattr_static.py b/lib/python3.12/site-packages/pure_eval/my_getattr_static.py new file mode 100644 index 0000000000000000000000000000000000000000..5490b90adec97a64955ef6ffdc35879cffc60603 --- /dev/null +++ b/lib/python3.12/site-packages/pure_eval/my_getattr_static.py @@ -0,0 +1,140 @@ +import types + +from pure_eval.utils import of_type, CannotEval + +_sentinel = object() + + +def _static_getmro(klass): + return type.__dict__['__mro__'].__get__(klass) + + +def _check_instance(obj, attr): + instance_dict = {} + try: + instance_dict = object.__getattribute__(obj, "__dict__") + except AttributeError: + pass + return dict.get(instance_dict, attr, _sentinel) + + +def _check_class(klass, attr): + for entry in _static_getmro(klass): + if _shadowed_dict(type(entry)) is _sentinel: + try: + return entry.__dict__[attr] + except KeyError: + pass + else: + break + return _sentinel + + +def _is_type(obj): + try: + _static_getmro(obj) + except TypeError: + return False + return True + + +def _shadowed_dict(klass): + dict_attr = type.__dict__["__dict__"] + for entry in _static_getmro(klass): + try: + class_dict = dict_attr.__get__(entry)["__dict__"] + except KeyError: + pass + else: + if not (type(class_dict) is types.GetSetDescriptorType and + class_dict.__name__ == "__dict__" and + class_dict.__objclass__ is entry): + return class_dict + return _sentinel + + +def getattr_static(obj, attr): + """Retrieve attributes without triggering dynamic lookup via the + descriptor protocol, __getattr__ or __getattribute__. + + Note: this function may not be able to retrieve all attributes + that getattr can fetch (like dynamically created attributes) + and may find attributes that getattr can't (like descriptors + that raise AttributeError). It can also return descriptor objects + instead of instance members in some cases. See the + documentation for details. + """ + instance_result = _sentinel + if not _is_type(obj): + klass = type(obj) + dict_attr = _shadowed_dict(klass) + if (dict_attr is _sentinel or + type(dict_attr) is types.MemberDescriptorType): + instance_result = _check_instance(obj, attr) + else: + raise CannotEval + else: + klass = obj + + klass_result = _check_class(klass, attr) + + if instance_result is not _sentinel and klass_result is not _sentinel: + if _check_class(type(klass_result), "__get__") is not _sentinel and ( + _check_class(type(klass_result), "__set__") is not _sentinel + or _check_class(type(klass_result), "__delete__") is not _sentinel + ): + return _resolve_descriptor(klass_result, obj, klass) + + if instance_result is not _sentinel: + return instance_result + if klass_result is not _sentinel: + get = _check_class(type(klass_result), '__get__') + if get is _sentinel: + return klass_result + else: + if obj is klass: + instance = None + else: + instance = obj + return _resolve_descriptor(klass_result, instance, klass) + + if obj is klass: + # for types we check the metaclass too + for entry in _static_getmro(type(klass)): + if _shadowed_dict(type(entry)) is _sentinel: + try: + result = entry.__dict__[attr] + get = _check_class(type(result), '__get__') + if get is not _sentinel: + raise CannotEval + return result + except KeyError: + pass + raise CannotEval + + +class _foo: + __slots__ = ['foo'] + method = lambda: 0 + + +slot_descriptor = _foo.foo +wrapper_descriptor = str.__dict__['__add__'] +method_descriptor = str.__dict__['startswith'] +user_method_descriptor = _foo.__dict__['method'] + +safe_descriptors_raw = [ + slot_descriptor, + wrapper_descriptor, + method_descriptor, + user_method_descriptor, +] + +safe_descriptor_types = list(map(type, safe_descriptors_raw)) + + +def _resolve_descriptor(d, instance, owner): + try: + return type(of_type(d, *safe_descriptor_types)).__get__(d, instance, owner) + except AttributeError as e: + raise CannotEval from e diff --git a/lib/python3.12/site-packages/pure_eval/py.typed b/lib/python3.12/site-packages/pure_eval/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..298c64a9041faa49413d4025444a28acc4a1f44c --- /dev/null +++ b/lib/python3.12/site-packages/pure_eval/py.typed @@ -0,0 +1 @@ +# Marker file for PEP 561. The pure_eval package uses inline types. diff --git a/lib/python3.12/site-packages/pure_eval/utils.py b/lib/python3.12/site-packages/pure_eval/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..19ead6585060f191ecc7ec070f21c7323eac72ee --- /dev/null +++ b/lib/python3.12/site-packages/pure_eval/utils.py @@ -0,0 +1,206 @@ +from collections import OrderedDict, deque +from datetime import date, time, datetime +from decimal import Decimal +from fractions import Fraction +import ast +import enum +import typing + + +class CannotEval(Exception): + def __repr__(self): + return self.__class__.__name__ + + __str__ = __repr__ + + +def is_any(x, *args): + return any( + x is arg + for arg in args + ) + + +def of_type(x, *types): + if is_any(type(x), *types): + return x + else: + raise CannotEval + + +def of_standard_types(x, *, check_dict_values: bool, deep: bool): + if is_standard_types(x, check_dict_values=check_dict_values, deep=deep): + return x + else: + raise CannotEval + + +def is_standard_types(x, *, check_dict_values: bool, deep: bool): + try: + return _is_standard_types_deep(x, check_dict_values, deep)[0] + except RecursionError: + return False + + +def _is_standard_types_deep(x, check_dict_values: bool, deep: bool): + typ = type(x) + if is_any( + typ, + str, + int, + bool, + float, + bytes, + complex, + date, + time, + datetime, + Fraction, + Decimal, + type(None), + object, + ): + return True, 0 + + if is_any(typ, tuple, frozenset, list, set, dict, OrderedDict, deque, slice): + if typ in [slice]: + length = 0 + else: + length = len(x) + assert isinstance(deep, bool) + if not deep: + return True, length + + if check_dict_values and typ in (dict, OrderedDict): + items = (v for pair in x.items() for v in pair) + elif typ is slice: + items = [x.start, x.stop, x.step] + else: + items = x + for item in items: + if length > 100000: + return False, length + is_standard, item_length = _is_standard_types_deep( + item, check_dict_values, deep + ) + if not is_standard: + return False, length + length += item_length + return True, length + + return False, 0 + + +class _E(enum.Enum): + pass + + +class _C: + def foo(self): pass # pragma: nocover + + def bar(self): pass # pragma: nocover + + @classmethod + def cm(cls): pass # pragma: nocover + + @staticmethod + def sm(): pass # pragma: nocover + + +safe_name_samples = { + "len": len, + "append": list.append, + "__add__": list.__add__, + "insert": [].insert, + "__mul__": [].__mul__, + "fromkeys": dict.__dict__['fromkeys'], + "is_any": is_any, + "__repr__": CannotEval.__repr__, + "foo": _C().foo, + "bar": _C.bar, + "cm": _C.cm, + "sm": _C.sm, + "ast": ast, + "CannotEval": CannotEval, + "_E": _E, +} + +typing_annotation_samples = { + name: getattr(typing, name) + for name in "List Dict Tuple Set Callable Mapping".split() +} + +safe_name_types = tuple({ + type(f) + for f in safe_name_samples.values() +}) + + +typing_annotation_types = tuple({ + type(f) + for f in typing_annotation_samples.values() +}) + + +def eq_checking_types(a, b): + return type(a) is type(b) and a == b + + +def ast_name(node): + if isinstance(node, ast.Name): + return node.id + elif isinstance(node, ast.Attribute): + return node.attr + else: + return None + + +def safe_name(value): + typ = type(value) + if is_any(typ, *safe_name_types): + return value.__name__ + elif value is typing.Optional: + return "Optional" + elif value is typing.Union: + return "Union" + elif is_any(typ, *typing_annotation_types): + return getattr(value, "__name__", None) or getattr(value, "_name", None) + else: + return None + + +def has_ast_name(value, node): + value_name = safe_name(value) + if type(value_name) is not str: + return False + return eq_checking_types(ast_name(node), value_name) + + +def copy_ast_without_context(x): + if isinstance(x, ast.AST): + kwargs = { + field: copy_ast_without_context(getattr(x, field)) + for field in x._fields + if field != 'ctx' + if hasattr(x, field) + } + a = type(x)(**kwargs) + if hasattr(a, 'ctx'): + # Python 3.13.0b2+ defaults to Load when we don't pass ctx + # https://github.com/python/cpython/pull/118871 + del a.ctx + return a + elif isinstance(x, list): + return list(map(copy_ast_without_context, x)) + else: + return x + + +def ensure_dict(x): + """ + Handles invalid non-dict inputs + """ + try: + return dict(x) + except Exception: + return {} diff --git a/lib/python3.12/site-packages/pure_eval/version.py b/lib/python3.12/site-packages/pure_eval/version.py new file mode 100644 index 0000000000000000000000000000000000000000..db9cf742db5ea82c8ef72e791395e73d5736ee2e --- /dev/null +++ b/lib/python3.12/site-packages/pure_eval/version.py @@ -0,0 +1 @@ +__version__ = '0.2.3' \ No newline at end of file diff --git a/lib/python3.12/site-packages/ray/__init__.py b/lib/python3.12/site-packages/ray/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..372ec27bcee1ce9e183307f25d70bb3567a33fcb --- /dev/null +++ b/lib/python3.12/site-packages/ray/__init__.py @@ -0,0 +1,304 @@ +# isort: skip_file +from ray._private import log # isort: skip # noqa: F401 +import logging +import os +import sys +from typing import TYPE_CHECKING + +log.generate_logging_config() +logger = logging.getLogger(__name__) + + +def _configure_system(): + import os + import platform + import sys + + """Wraps system configuration to avoid 'leaking' variables into ray.""" + + # Sanity check pickle5 if it has been installed. + if "pickle5" in sys.modules: + if sys.version_info >= (3, 8): + logger.warning( + "Package pickle5 becomes unnecessary in Python 3.8 and above. " + "Its presence may confuse libraries including Ray. " + "Please uninstall the package." + ) + + import importlib.metadata + + try: + version_str = importlib.metadata.version("pickle5") + version = tuple(int(n) for n in version_str.split(".")) + if version < (0, 0, 10): + logger.warning( + "Although not used by Ray, a version of pickle5 that leaks memory " + "is found in the environment. Please run 'pip install pickle5 -U' " + "to upgrade." + ) + except importlib.metadata.PackageNotFoundError: + logger.warning( + "You are using the 'pickle5' module, but " + "the exact version is unknown (possibly carried as " + "an internal component by another module). Please " + "make sure you are using pickle5 >= 0.0.10 because " + "previous versions may leak memory." + ) + + # Importing psutil. Must be before ray._raylet is + # initialized. + thirdparty_files = os.path.join( + os.path.abspath(os.path.dirname(__file__)), "thirdparty_files" + ) + sys.path.insert(0, thirdparty_files) + + if ( + platform.system() == "Linux" + and "Microsoft".lower() in platform.release().lower() + ): + from ray._private import compat # noqa: E402 + + compat.patch_psutil() + + # Expose ray ABI symbols which may be dependent by other shared + # libraries such as _streaming.so. See BUILD.bazel:_raylet + python_shared_lib_suffix = ".so" if sys.platform != "win32" else ".pyd" + so_path = os.path.join( + os.path.dirname(__file__), "_raylet" + python_shared_lib_suffix + ) + if os.path.exists(so_path): + import ctypes + from ctypes import CDLL + + CDLL(so_path, ctypes.RTLD_GLOBAL) + + +_configure_system() +# Delete configuration function. +del _configure_system + +from ray import _version # noqa: E402 + +__commit__ = _version.commit +__version__ = _version.version + +import ray._raylet # noqa: E402 + +from ray._raylet import ( # noqa: E402,F401 + ActorClassID, + ActorID, + NodeID, + Config as _Config, + JobID, + WorkerID, + FunctionID, + ObjectID, + ObjectRef, + ObjectRefGenerator, + DynamicObjectRefGenerator, + TaskID, + UniqueID, + Language, + PlacementGroupID, + ClusterID, +) + +_config = _Config() + +from ray._private.state import ( # noqa: E402,F401 + nodes, + timeline, + cluster_resources, + available_resources, +) +from ray._private.worker import ( # noqa: E402,F401 + LOCAL_MODE, + SCRIPT_MODE, + WORKER_MODE, + RESTORE_WORKER_MODE, + SPILL_WORKER_MODE, + cancel, + get, + get_actor, + get_gpu_ids, + init, + is_initialized, + put, + kill, + remote, + shutdown, + wait, +) + +from ray._private.ray_logging.logging_config import LoggingConfig # noqa: E402 + +# We import ray.actor because some code is run in actor.py which initializes +# some functions in the worker. +import ray.actor # noqa: E402,F401 +from ray.actor import method # noqa: E402,F401 + +# TODO(qwang): We should remove this exporting in Ray2.0. +from ray.cross_language import java_function, java_actor_class # noqa: E402,F401 +from ray.runtime_context import get_runtime_context # noqa: E402,F401 +from ray import internal # noqa: E402,F401 +from ray import util # noqa: E402,F401 +from ray import _private # noqa: E402,F401 + +# We import ClientBuilder so that modules can inherit from `ray.ClientBuilder`. +from ray.client_builder import client, ClientBuilder # noqa: E402,F401 + + +class _DeprecationWrapper: + def __init__(self, name, real_worker): + self._name = name + self._real_worker = real_worker + self._warned = set() + + def __getattr__(self, attr): + value = getattr(self._real_worker, attr) + if attr not in self._warned: + self._warned.add(attr) + logger.warning( + f"DeprecationWarning: `ray.{self._name}.{attr}` is a private " + "attribute and access will be removed in a future Ray version." + ) + return value + + +# TODO(ekl) remove this entirely after 3rd party libraries are all migrated. +worker = _DeprecationWrapper("worker", ray._private.worker) +ray_constants = _DeprecationWrapper("ray_constants", ray._private.ray_constants) +serialization = _DeprecationWrapper("serialization", ray._private.serialization) +state = _DeprecationWrapper("state", ray._private.state) + + +# Pulic Ray APIs +__all__ = [ + "__version__", + "_config", + "get_runtime_context", + "autoscaler", + "available_resources", + "cancel", + "client", + "ClientBuilder", + "cluster_resources", + "get", + "get_actor", + "get_gpu_ids", + "init", + "is_initialized", + "java_actor_class", + "java_function", + "cpp_function", + "kill", + "Language", + "method", + "nodes", + "put", + "remote", + "shutdown", + "show_in_dashboard", + "timeline", + "wait", + "LOCAL_MODE", + "SCRIPT_MODE", + "WORKER_MODE", + "LoggingConfig", +] + +# Public APIs that should automatically trigger ray.init(). +AUTO_INIT_APIS = { + "cancel", + "get", + "get_actor", + "get_gpu_ids", + "kill", + "put", + "wait", + "get_runtime_context", +} + +# Public APIs that should not automatically trigger ray.init(). +NON_AUTO_INIT_APIS = { + "ClientBuilder", + "LOCAL_MODE", + "Language", + "SCRIPT_MODE", + "WORKER_MODE", + "__version__", + "_config", + "autoscaler", + "available_resources", + "client", + "cluster_resources", + "cpp_function", + "init", + "is_initialized", + "java_actor_class", + "java_function", + "method", + "nodes", + "remote", + "show_in_dashboard", + "shutdown", + "timeline", + "LoggingConfig", +} + +assert set(__all__) == AUTO_INIT_APIS | NON_AUTO_INIT_APIS +from ray._private.auto_init_hook import wrap_auto_init_for_all_apis # noqa: E402 + +wrap_auto_init_for_all_apis(AUTO_INIT_APIS) +del wrap_auto_init_for_all_apis + +# Subpackages +__all__ += [ + "actor", + "autoscaler", + "data", + "internal", + "util", + "widgets", + "workflow", +] + +# ID types +__all__ += [ + "ActorClassID", + "ActorID", + "NodeID", + "JobID", + "WorkerID", + "FunctionID", + "ObjectID", + "ObjectRef", + "ObjectRefGenerator", + "DynamicObjectRefGenerator", + "TaskID", + "UniqueID", + "PlacementGroupID", +] + + +# Delay importing of expensive, isolated subpackages. Note that for proper type +# checking support these imports must be kept in sync between type checking and +# runtime behavior. +if TYPE_CHECKING: + from ray import autoscaler + from ray import data + from ray import workflow +else: + + def __getattr__(name: str): + import importlib + + if name in ["data", "workflow", "autoscaler"]: + return importlib.import_module("." + name, __name__) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +del os +del logging +del sys +del TYPE_CHECKING diff --git a/lib/python3.12/site-packages/ray/_raylet.pxd b/lib/python3.12/site-packages/ray/_raylet.pxd new file mode 100644 index 0000000000000000000000000000000000000000..d87de681a3600dd47e9cef50104e2df2ed13b896 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_raylet.pxd @@ -0,0 +1,185 @@ +# cython: profile=False +# distutils: language = c++ +# cython: embedsignature = True +# cython: language_level = 3 + +from cpython.pystate cimport PyThreadState_Get + +from libc.stdint cimport ( + int64_t, +) +from libcpp cimport bool as c_bool +from libcpp.string cimport string as c_string +from libcpp.vector cimport vector as c_vector +from libcpp.unordered_map cimport unordered_map +from libcpp.memory cimport ( + shared_ptr, + unique_ptr +) +from libcpp.pair cimport pair as c_pair +from libcpp.utility cimport pair +from ray.includes.optional cimport ( + optional, + nullopt, + make_optional, +) + +from ray.includes.common cimport ( + CBuffer, + CRayObject, + CAddress, + CConcurrencyGroup, + CSchedulingStrategy, + CLabelMatchExpressions, + CTensorTransport, +) +from ray.includes.libcoreworker cimport ( + ActorHandleSharedPtr, + CActorHandle, + CFiberEvent, +) + +from ray.includes.unique_ids cimport ( + CObjectID, + CActorID, + CTaskID, +) +from ray.includes.function_descriptor cimport ( + CFunctionDescriptor, +) + +cdef extern from *: + """ + #if __OPTIMIZE__ && __OPTIMIZE__ == 1 + #undef __OPTIMIZE__ + int __OPTIMIZE__ = 1; + #define __OPTIMIZE__ 1 + #elif defined(BAZEL_OPT) + // For compilers that don't define __OPTIMIZE__ + int __OPTIMIZE__ = 1; + #else + int __OPTIMIZE__ = 0; + #endif + """ + int __OPTIMIZE__ + +cdef extern from "Python.h": + # Note(simon): This is used to configure asyncio actor stack size. + # Cython made PyThreadState an opaque types. Saying that if the user wants + # specific attributes, they can be declared manually. + + # You can find the cpython definition in Include/cpython/pystate.h#L59 + ctypedef struct CPyThreadState "PyThreadState": + int recursion_limit + int recursion_remaining + int c_recursion_remaining + + # From Include/ceveal.h#67 + int Py_GetRecursionLimit() + void Py_SetRecursionLimit(int) + +# Note that `functional.pxd` in the Cython repository supports only a limited subset of +# . Therefore, `from libcpp.functional cimport function` is not enough, and we +# still need to expose some functions here. +cdef extern from "" namespace "std" nogil: + T bind[T, Args](T callable, Args args) + # Reference: https://github.com/scipy/scipy/blob/6b56162fa6880b0182faea44af88d6a1587f35a8/scipy/stats/_qmc_cy.pyx#L31-L34 + cdef cppclass reference_wrapper[T]: + pass + cdef reference_wrapper[T] ref[T](T&) + +cdef class Buffer: + cdef: + shared_ptr[CBuffer] buffer + Py_ssize_t shape + Py_ssize_t strides + + @staticmethod + cdef make(const shared_ptr[CBuffer]& buffer) + +cdef class BaseID: + # To avoid the error of "Python int too large to convert to C ssize_t", + # here `cdef size_t` is required. + cdef size_t hash(self) + +cdef class ObjectRef(BaseID): + cdef: + CObjectID data + c_string owner_addr + # Flag indicating whether or not this object ref was added to the set + # of active IDs in the core worker so we know whether we should clean + # it up. + c_bool in_core_worker + c_string call_site_data + int tensor_transport_val + + cdef CObjectID native(self) + + cdef CTensorTransport c_tensor_transport(self) + +cdef class ActorID(BaseID): + cdef CActorID data + + cdef CActorID native(self) + + cdef size_t hash(self) + + +cdef class CoreWorker: + cdef: + c_bool is_driver + object async_thread + object async_event_loop + object job_config + object current_runtime_env + c_bool is_local_mode + + object cgname_to_eventloop_dict + object eventloop_for_default_cg + object thread_for_default_cg + object fd_to_cgname_dict + object _task_id_to_future_lock + dict _task_id_to_future + object event_loop_executor + object _gc_thread + + cdef unique_ptr[CAddress] _convert_python_address(self, address=*) + cdef store_task_output( + self, serialized_object, + const CObjectID &return_id, + const CObjectID &generator_id, + size_t data_size, shared_ptr[CBuffer] &metadata, const c_vector[CObjectID] + &contained_id, const CAddress &caller_address, + int64_t *task_output_inlined_bytes, + shared_ptr[CRayObject] *return_ptr) + cdef store_task_outputs( + self, + worker, outputs, + const CAddress &caller_address, + c_vector[c_pair[CObjectID, shared_ptr[CRayObject]]] *returns, + ref_generator_id=*, # CObjectID + CTensorTransport c_tensor_transport=*, + ) + cdef make_actor_handle(self, ActorHandleSharedPtr c_actor_handle, + c_bool weak_ref) + cdef c_function_descriptors_to_python( + self, const c_vector[CFunctionDescriptor] &c_function_descriptors) + cdef initialize_eventloops_for_actor_concurrency_group( + self, const c_vector[CConcurrencyGroup] &c_defined_concurrency_groups) + cdef python_scheduling_strategy_to_c( + self, python_scheduling_strategy, + CSchedulingStrategy *c_scheduling_strategy) + cdef python_label_match_expressions_to_c( + self, python_expressions, + CLabelMatchExpressions *c_expressions) + cdef CObjectID allocate_dynamic_return_id_for_generator( + self, + const CAddress &owner_address, + const CTaskID &task_id, + return_size, + generator_index, + is_async_actor) + +cdef class FunctionDescriptor: + cdef: + CFunctionDescriptor descriptor diff --git a/lib/python3.12/site-packages/ray/_raylet.pyi b/lib/python3.12/site-packages/ray/_raylet.pyi new file mode 100644 index 0000000000000000000000000000000000000000..fff69c451b67f0b4d38e3d9c41bfce96ef8adee9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_raylet.pyi @@ -0,0 +1,37 @@ +from ray.includes.object_ref import ObjectRef, _set_future_helper +from ray.includes.unique_ids import ( + ActorClassID, + ActorID, + BaseID, + ClusterID, + FunctionID, + JobID, + NodeID, + ObjectID, + PlacementGroupID, + TaskID, + UniqueID, + WorkerID, + check_id, +) + +__all__ = [ + # ray.includes.unique_ids + "ActorClassID", + "ActorID", + "BaseID", + "ClusterID", + "FunctionID", + "JobID", + "NodeID", + "ObjectID", + "PlacementGroupID", + "TaskID", + "UniqueID", + "WorkerID", + "check_id", + + # ray.includes.object_ref + "_set_future_helper", + "ObjectRef", +] diff --git a/lib/python3.12/site-packages/ray/_version.py b/lib/python3.12/site-packages/ray/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..15c91cf88a104aa450b36729158b744aa5f701a5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/_version.py @@ -0,0 +1,6 @@ +# Replaced with the current commit when building the wheels. +commit = "4ebdc0abe5e5a551625fe7f87053c7e668a6ff74" +version = "2.52.1" + +if __name__ == "__main__": + print("%s %s" % (version, commit)) diff --git a/lib/python3.12/site-packages/ray/actor.py b/lib/python3.12/site-packages/ray/actor.py new file mode 100644 index 0000000000000000000000000000000000000000..b33c8c0f575ddc241ddfcd850eb000d064296750 --- /dev/null +++ b/lib/python3.12/site-packages/ray/actor.py @@ -0,0 +1,2483 @@ +import inspect +import logging +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Generic, + List, + Literal, + Optional, + Tuple, + TypeVar, + Union, + overload, +) + +try: + from typing import Concatenate, ParamSpec +except ImportError: + from typing_extensions import Concatenate, ParamSpec + +import ray._common.signature as signature +import ray._private.ray_constants as ray_constants +import ray._raylet +from ray import ActorClassID, Language, ObjectRef, cross_language +from ray._common import ray_option_utils +from ray._common.ray_constants import DEFAULT_MAX_CONCURRENCY_ASYNC +from ray._common.ray_option_utils import _warn_if_using_deprecated_placement_group +from ray._private.async_compat import has_async_methods +from ray._private.auto_init_hook import wrap_auto_init +from ray._private.client_mode_hook import ( + client_mode_convert_actor, + client_mode_hook, + client_mode_should_convert, +) +from ray._private.custom_types import ( + TensorTransportEnum, +) +from ray._private.inspect_util import ( + is_class_method, + is_function_or_method, + is_static_method, +) +from ray._private.utils import get_runtime_env_info, parse_runtime_env_for_task_or_actor +from ray._raylet import ( + STREAMING_GENERATOR_RETURN, + ObjectRefGenerator, + PythonFunctionDescriptor, + raise_sys_exit_with_custom_error_message, +) +from ray.exceptions import AsyncioActorExit +from ray.util.annotations import DeveloperAPI, PublicAPI +from ray.util.placement_group import _configure_placement_group_based_on_context +from ray.util.scheduling_strategies import ( + PlacementGroupSchedulingStrategy, + SchedulingStrategyT, +) +from ray.util.tracing.tracing_helper import ( + _inject_tracing_into_class, + _tracing_actor_creation, + _tracing_actor_method_invocation, +) + +if TYPE_CHECKING: + pass + +logger = logging.getLogger(__name__) + +# Hook to call with (actor, resources, strategy) on each local actor creation. +_actor_launch_hook = None + +# TypeVar for generic ActorHandle +T = TypeVar("T") + +# return type of ActorClass[T].remote() +ActorProxy = Union["ActorHandle[T]", type[T]] + +_Ret = TypeVar("_Ret") +_P = ParamSpec("_P") +_T0 = TypeVar("_T0") +_T1 = TypeVar("_T1") +_T2 = TypeVar("_T2") +_T3 = TypeVar("_T3") +_T4 = TypeVar("_T4") +_T5 = TypeVar("_T5") +_T6 = TypeVar("_T6") +_T7 = TypeVar("_T7") +_T8 = TypeVar("_T8") +_T9 = TypeVar("_T9") + + +class _RemoteMethodNoArgs(Generic[_Ret]): + def remote(self) -> "ObjectRef[_Ret]": + ... + + def bind(self) -> Any: + ... + + +class _RemoteMethod0(Generic[_Ret, _T0]): + def remote(self, __arg0: "Union[_T0, ObjectRef[_T0]]") -> "ObjectRef[_Ret]": + ... + + def bind(self, __arg0: _T0) -> Any: + ... + + +class _RemoteMethod1(Generic[_Ret, _T0, _T1]): + def remote( + self, __arg0: "Union[_T0, ObjectRef[_T0]]", __arg1: "Union[_T1, ObjectRef[_T1]]" + ) -> "ObjectRef[_Ret]": + ... + + def bind(self, __arg0: _T0, __arg1: _T1) -> Any: + ... + + +class _RemoteMethod2(Generic[_Ret, _T0, _T1, _T2]): + def remote( + self, + __arg0: "Union[_T0, ObjectRef[_T0]]", + __arg1: "Union[_T1, ObjectRef[_T1]]", + __arg2: "Union[_T2, ObjectRef[_T2]]", + ) -> "ObjectRef[_Ret]": + ... + + def bind(self, __arg0: _T0, __arg1: _T1, __arg2: _T2) -> Any: + ... + + +class _RemoteMethod3(Generic[_Ret, _T0, _T1, _T2, _T3]): + def remote( + self, + __arg0: "Union[_T0, ObjectRef[_T0]]", + __arg1: "Union[_T1, ObjectRef[_T1]]", + __arg2: "Union[_T2, ObjectRef[_T2]]", + __arg3: "Union[_T3, ObjectRef[_T3]]", + ) -> "ObjectRef[_Ret]": + ... + + def bind(self, __arg0: _T0, __arg1: _T1, __arg2: _T2, __arg3: _T3) -> Any: + ... + + +class _RemoteMethod4(Generic[_Ret, _T0, _T1, _T2, _T3, _T4]): + def remote( + self, + __arg0: "Union[_T0, ObjectRef[_T0]]", + __arg1: "Union[_T1, ObjectRef[_T1]]", + __arg2: "Union[_T2, ObjectRef[_T2]]", + __arg3: "Union[_T3, ObjectRef[_T3]]", + __arg4: "Union[_T4, ObjectRef[_T4]]", + ) -> "ObjectRef[_Ret]": + ... + + def bind( + self, __arg0: _T0, __arg1: _T1, __arg2: _T2, __arg3: _T3, __arg4: _T4 + ) -> Any: + ... + + +class _RemoteMethod5(Generic[_Ret, _T0, _T1, _T2, _T3, _T4, _T5]): + def remote( + self, + __arg0: "Union[_T0, ObjectRef[_T0]]", + __arg1: "Union[_T1, ObjectRef[_T1]]", + __arg2: "Union[_T2, ObjectRef[_T2]]", + __arg3: "Union[_T3, ObjectRef[_T3]]", + __arg4: "Union[_T4, ObjectRef[_T4]]", + __arg5: "Union[_T5, ObjectRef[_T5]]", + ) -> "ObjectRef[_Ret]": + ... + + def bind( + self, + __arg0: _T0, + __arg1: _T1, + __arg2: _T2, + __arg3: _T3, + __arg4: _T4, + __arg5: _T5, + ) -> Any: + ... + + +class _RemoteMethod6(Generic[_Ret, _T0, _T1, _T2, _T3, _T4, _T5, _T6]): + def remote( + self, + __arg0: "Union[_T0, ObjectRef[_T0]]", + __arg1: "Union[_T1, ObjectRef[_T1]]", + __arg2: "Union[_T2, ObjectRef[_T2]]", + __arg3: "Union[_T3, ObjectRef[_T3]]", + __arg4: "Union[_T4, ObjectRef[_T4]]", + __arg5: "Union[_T5, ObjectRef[_T5]]", + __arg6: "Union[_T6, ObjectRef[_T6]]", + ) -> "ObjectRef[_Ret]": + ... + + def bind( + self, + __arg0: _T0, + __arg1: _T1, + __arg2: _T2, + __arg3: _T3, + __arg4: _T4, + __arg5: _T5, + __arg6: _T6, + ) -> Any: + ... + + +class _RemoteMethod7(Generic[_Ret, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7]): + def remote( + self, + __arg0: "Union[_T0, ObjectRef[_T0]]", + __arg1: "Union[_T1, ObjectRef[_T1]]", + __arg2: "Union[_T2, ObjectRef[_T2]]", + __arg3: "Union[_T3, ObjectRef[_T3]]", + __arg4: "Union[_T4, ObjectRef[_T4]]", + __arg5: "Union[_T5, ObjectRef[_T5]]", + __arg6: "Union[_T6, ObjectRef[_T6]]", + __arg7: "Union[_T7, ObjectRef[_T7]]", + ) -> "ObjectRef[_Ret]": + ... + + def bind( + self, + __arg0: _T0, + __arg1: _T1, + __arg2: _T2, + __arg3: _T3, + __arg4: _T4, + __arg5: _T5, + __arg6: _T6, + __arg7: _T7, + ) -> Any: + ... + + +class _RemoteMethod8(Generic[_Ret, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, _T8]): + def remote( + self, + __arg0: "Union[_T0, ObjectRef[_T0]]", + __arg1: "Union[_T1, ObjectRef[_T1]]", + __arg2: "Union[_T2, ObjectRef[_T2]]", + __arg3: "Union[_T3, ObjectRef[_T3]]", + __arg4: "Union[_T4, ObjectRef[_T4]]", + __arg5: "Union[_T5, ObjectRef[_T5]]", + __arg6: "Union[_T6, ObjectRef[_T6]]", + __arg7: "Union[_T7, ObjectRef[_T7]]", + __arg8: "Union[_T8, ObjectRef[_T8]]", + ) -> "ObjectRef[_Ret]": + ... + + def bind( + self, + __arg0: _T0, + __arg1: _T1, + __arg2: _T2, + __arg3: _T3, + __arg4: _T4, + __arg5: _T5, + __arg6: _T6, + __arg7: _T7, + __arg8: _T8, + ) -> Any: + ... + + +class _RemoteMethod9(Generic[_Ret, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, _T8, _T9]): + def remote( + self, + __arg0: "Union[_T0, ObjectRef[_T0]]", + __arg1: "Union[_T1, ObjectRef[_T1]]", + __arg2: "Union[_T2, ObjectRef[_T2]]", + __arg3: "Union[_T3, ObjectRef[_T3]]", + __arg4: "Union[_T4, ObjectRef[_T4]]", + __arg5: "Union[_T5, ObjectRef[_T5]]", + __arg6: "Union[_T6, ObjectRef[_T6]]", + __arg7: "Union[_T7, ObjectRef[_T7]]", + __arg8: "Union[_T8, ObjectRef[_T8]]", + __arg9: "Union[_T9, ObjectRef[_T9]]", + ) -> "ObjectRef[_Ret]": + ... + + def bind( + self, + __arg0: _T0, + __arg1: _T1, + __arg2: _T2, + __arg3: _T3, + __arg4: _T4, + __arg5: _T5, + __arg6: _T6, + __arg7: _T7, + __arg8: _T8, + __arg9: _T9, + ) -> Any: + ... + + +@overload +def method( + __method: Callable[[Any, _T0], _Ret], +) -> _RemoteMethod0[_Ret, _T0]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1], _Ret], +) -> _RemoteMethod1[_Ret, _T0, _T1]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1, _T2], _Ret], +) -> _RemoteMethod2[_Ret, _T0, _T1, _T2]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1, _T2, _T3], _Ret], +) -> _RemoteMethod3[_Ret, _T0, _T1, _T2, _T3]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1, _T2, _T3, _T4], _Ret], +) -> _RemoteMethod4[_Ret, _T0, _T1, _T2, _T3, _T4]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1, _T2, _T3, _T4, _T5], _Ret], +) -> _RemoteMethod5[_Ret, _T0, _T1, _T2, _T3, _T4, _T5]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1, _T2, _T3, _T4, _T5, _T6], _Ret], +) -> _RemoteMethod6[_Ret, _T0, _T1, _T2, _T3, _T4, _T5, _T6]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7], _Ret], +) -> _RemoteMethod7[_Ret, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, _T8], _Ret], +) -> _RemoteMethod8[_Ret, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, _T8]: + ... + + +@overload +def method( + __method: Callable[[Any, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, _T8, _T9], _Ret], +) -> _RemoteMethod9[_Ret, _T0, _T1, _T2, _T3, _T4, _T5, _T6, _T7, _T8, _T9]: + ... + + +@overload +def method( + __method: Callable[[Any], _Ret], +) -> _RemoteMethodNoArgs[_Ret]: + ... + + +@overload +def method( + *, + num_returns: Optional[Union[int, Literal["streaming"]]] = None, + concurrency_group: Optional[str] = None, + max_task_retries: Optional[int] = None, + retry_exceptions: Optional[Union[bool, list, tuple]] = None, + _generator_backpressure_num_objects: Optional[int] = None, + enable_task_events: Optional[bool] = None, + tensor_transport: Optional[TensorTransportEnum] = None, +) -> Callable[[Callable[Concatenate[Any, _P], _Ret]], Any]: + ... + + +@PublicAPI +@client_mode_hook +def method(*args, **kwargs): + """Annotate an actor method. + + .. code-block:: python + + @ray.remote + class Foo: + @ray.method(num_returns=2) + def bar(self): + return 1, 2 + + f = Foo.remote() + + _, _ = f.bar.remote() + + Args: + num_returns: The number of object refs that should be returned by + invocations of this actor method. The default value is 1 for a + normal actor task and "streaming" for an actor generator task (a + function that yields objects instead of returning them). + max_task_retries: How many times to retry an actor task if the task + fails due to a runtime error, e.g., the actor has died. The + default value is 0. If set to -1, the system will retry the + failed task until the task succeeds, or the actor has reached + its max_restarts limit. If set to `n > 0`, the system will retry + the failed task up to n times, after which the task will throw a + `RayActorError` exception upon :obj:`ray.get`. Note that Python + exceptions may trigger retries + *only if* `retry_exceptions` is set for the method, in that case + when `max_task_retries` runs out the task will rethrow the + exception from the task. You can override this number with the + method's `max_task_retries` option in `@ray.method` decorator or + in `.option()`. + retry_exceptions: Boolean of whether to retry all Python + exceptions, or a list of allowlist exceptions to retry. The default + value is False (only retry tasks upon system failures and if + max_task_retries is set) + concurrency_group: The name of the concurrency group + to use for the actor method. By default, the actor is + single-threaded and runs all actor tasks on the same thread. + See :ref:`Defining Concurrency Groups `. + tensor_transport: [Alpha] The tensor transport protocol to + use for the actor method. The valid values are "OBJECT_STORE" + (default), "NCCL", "GLOO", or "NIXL" (case-insensitive). If a + non-object store transport is specified, Ray will store a + *reference* instead of a copy of any torch.Tensors found inside + values returned by this task, and the tensors will be sent directly + to other tasks using the specified transport. NCCL and GLOO + transports require first creating a collective with the involved + actors using + :func:`ray.experimental.collective.create_collective_group`. + See :ref:`Ray Direct Transport (RDT) ` for more + details. + """ + valid_kwargs = [ + "num_returns", + "concurrency_group", + "max_task_retries", + "retry_exceptions", + "_generator_backpressure_num_objects", + "enable_task_events", + "tensor_transport", + ] + + def annotate_method(method: Callable[_P, _Ret]): + if "num_returns" in kwargs: + method.__ray_num_returns__ = kwargs["num_returns"] + if "max_task_retries" in kwargs: + method.__ray_max_task_retries__ = kwargs["max_task_retries"] + if "retry_exceptions" in kwargs: + method.__ray_retry_exceptions__ = kwargs["retry_exceptions"] + if "concurrency_group" in kwargs: + method.__ray_concurrency_group__ = kwargs["concurrency_group"] + if "_generator_backpressure_num_objects" in kwargs: + method.__ray_generator_backpressure_num_objects__ = kwargs[ + "_generator_backpressure_num_objects" + ] + if "enable_task_events" in kwargs and kwargs["enable_task_events"] is not None: + method.__ray_enable_task_events__ = kwargs["enable_task_events"] + if "tensor_transport" in kwargs: + method.__ray_tensor_transport__ = TensorTransportEnum.from_str( + kwargs["tensor_transport"] + ) + return method + + # Check if decorator is called without parentheses (args[0] would be the function) + if len(args) == 1 and callable(args[0]) and len(kwargs) == 0: + # Called as @ray.method (without parentheses) + return annotate_method(args[0]) + + # Called as @ray.method() or @ray.method(options...) + error_string = ( + "The @ray.method decorator must be applied using no arguments or at " + f"least one of the arguments in the list {valid_kwargs}, for example " + "'@ray.method(num_returns=2)'." + ) + assert len(args) == 0, error_string + for key in kwargs: + key_error_string = ( + f"Unexpected keyword argument to @ray.method: '{key}'. The " + f"supported keyword arguments are {valid_kwargs}" + ) + assert key in valid_kwargs, key_error_string + + return annotate_method + + +class _ActorMethodMetadata: + """A container for the metadata required to invoke an actor method. + + This class intentionally does *not* hold a reference to the `ActorHandle`, as that causes + a circular reference that delays `ActorHandle` destruction until the Python GC runs. + + Instead, it can be used as a factory to lazily generate `ActorMethod` instances that can + be used to submit actor tasks for this method. + """ + + def __init__( + self, + method_name: str, + num_returns: Optional[Union[int, Literal["streaming"]]], + max_task_retries: int, + retry_exceptions: Union[bool, list, tuple], + is_generator: bool, + generator_backpressure_num_objects: int, + enable_task_events: bool, + decorator: Optional[Any] = None, + signature: Optional[List[inspect.Parameter]] = None, + tensor_transport: Optional[TensorTransportEnum] = None, + ): + """Initialize an _ActorMethodMetadata. + + Args: + method_name: The name of the actor method. + num_returns: The default number of return values that the method + invocation should return. If None is given, it uses + DEFAULT_ACTOR_METHOD_NUM_RETURN_VALS for a normal actor task + and "streaming" for a generator task (when `is_generator` is True). + max_task_retries: Number of retries on method failure. + retry_exceptions: Boolean or list/tuple of exceptions to retry. + is_generator: True if the method is a generator. + generator_backpressure_num_objects: Generator-only config for backpressure. + enable_task_events: True if task events are enabled for this method. + decorator: Optional decorator for the method invocation. + signature: The signature of the actor method. + tensor_transport: The tensor transport protocol to use for the actor method. + """ + self._method_name = method_name + + # Default case. + if num_returns is None: + if is_generator: + num_returns = "streaming" + else: + num_returns = ray_constants.DEFAULT_ACTOR_METHOD_NUM_RETURN_VALS + self._num_returns = num_returns + self._max_task_retries = max_task_retries + self._retry_exceptions = retry_exceptions + self._is_generator = is_generator + self._generator_backpressure_num_objects = generator_backpressure_num_objects + self._enable_task_events = enable_task_events + self._decorator = decorator + self._signature = signature + self._tensor_transport = tensor_transport + + def bind(self, actor_handle: "ActorHandle") -> "ActorMethod": + """ + Produce a bound ActorMethod that holds a strong reference to actor_handle. + """ + return ActorMethod( + actor_handle, + self._method_name, + self._num_returns, + self._max_task_retries, + self._retry_exceptions, + self._is_generator, + self._generator_backpressure_num_objects, + self._enable_task_events, + decorator=self._decorator, + signature=self._signature, + tensor_transport=self._tensor_transport, + ) + + +# Create objects to wrap method invocations. This is done so that we can +# invoke methods with actor.method.remote() instead of actor.method(). +@PublicAPI +class ActorMethod: + """A class used to invoke an actor method. + + Note: This class should not be instantiated directly. Instead, it should + only be used as a return value from the `@ray.method` decorator. + """ + + def __init__( + self, + actor, + method_name, + num_returns: Optional[Union[int, Literal["streaming"]]], + max_task_retries: int, + retry_exceptions: Union[bool, list, tuple], + is_generator: bool, + generator_backpressure_num_objects: int, + enable_task_events: bool, + decorator=None, + signature: Optional[List[inspect.Parameter]] = None, + tensor_transport: Optional[TensorTransportEnum] = None, + ): + """Initialize an ActorMethod. + + Args: + actor: The actor instance this method belongs to. + method_name: The name of the actor method. + num_returns: The default number of return values that the method + invocation should return. If None is given, it uses + DEFAULT_ACTOR_METHOD_NUM_RETURN_VALS for a normal actor task + and "streaming" for a generator task (when `is_generator` is True). + max_task_retries: Number of retries on method failure. + retry_exceptions: Boolean of whether you want to retry all user-raised + exceptions, or a list of allowlist exceptions to retry. + is_generator: True if a given method is a Python generator. + generator_backpressure_num_objects: Generator-only config. + If a number of unconsumed objects reach this threshold, + a actor task stop pausing. + enable_task_events: True if task events is enabled, i.e., task events from + the actor should be reported. Defaults to True. + decorator: An optional decorator that should be applied to the actor + method invocation. + signature: The signature of the actor method. It is None only when cross + language feature is used. + tensor_transport: The tensor transport protocol to use for the actor method. + """ + self._actor = actor + self._method_name = method_name + self._num_returns = num_returns + + # Default case. + if self._num_returns is None: + if is_generator: + self._num_returns = "streaming" + else: + self._num_returns = ray_constants.DEFAULT_ACTOR_METHOD_NUM_RETURN_VALS + + self._max_task_retries = max_task_retries + self._retry_exceptions = retry_exceptions + self._is_generator = is_generator + self._generator_backpressure_num_objects = generator_backpressure_num_objects + self._enable_task_events = enable_task_events + self._signature = signature + # This is a decorator that is used to wrap the function invocation (as + # opposed to the function execution). The decorator must return a + # function that takes in two arguments ("args" and "kwargs"). In most + # cases, it should call the function that was passed into the decorator + # and return the resulting ObjectRefs. + self._decorator = decorator + + # If the task call doesn't specify a tensor transport option, use `_tensor_transport` + # as the default transport for this actor method. + if tensor_transport is None: + tensor_transport = TensorTransportEnum.OBJECT_STORE + self._tensor_transport = tensor_transport + + def __call__(self, *args, **kwargs): + raise TypeError( + "Actor methods cannot be called directly. Instead " + f"of running 'object.{self._method_name}()', try " + f"'object.{self._method_name}.remote()'." + ) + + @DeveloperAPI + def bind(self, *args, **kwargs): + """ + Bind arguments to the actor method for Ray DAG building. + + This method generates and returns an intermediate representation (IR) + node that indicates the actor method will be called with the given + arguments at execution time. + + This method is used in both :ref:`Ray DAG ` and + :ref:`Ray Compiled Graph ` for building a DAG. + """ + return self._bind(args, kwargs) + + def remote(self, *args, **kwargs): + return self._remote(args, kwargs) + + def options(self, **options): + """Convenience method for executing an actor method call with options. + + Same arguments as func._remote(), but returns a wrapped function + that a non-underscore .remote() can be called on. + + Examples: + # The following two calls are equivalent. + >>> actor.my_method._remote(args=[x, y], name="foo", num_returns=2) + >>> actor.my_method.options(name="foo", num_returns=2).remote(x, y) + """ + + func_cls = self + + tensor_transport = options.get("tensor_transport", None) + if tensor_transport is not None: + options["tensor_transport"] = TensorTransportEnum.from_str(tensor_transport) + + class FuncWrapper: + def remote(self, *args, **kwargs): + return func_cls._remote(args=args, kwargs=kwargs, **options) + + @DeveloperAPI + def bind(self, *args, **kwargs): + return func_cls._bind(args=args, kwargs=kwargs, **options) + + return FuncWrapper() + + @wrap_auto_init + @_tracing_actor_method_invocation + def _bind( + self, + args=None, + kwargs=None, + name="", + num_returns=None, + concurrency_group=None, + _generator_backpressure_num_objects=None, + ) -> Union["ray.dag.ClassMethodNode", Tuple["ray.dag.ClassMethodNode", ...]]: + from ray.dag.class_node import ( + BIND_INDEX_KEY, + IS_CLASS_METHOD_OUTPUT_KEY, + PARENT_CLASS_NODE_KEY, + PREV_CLASS_METHOD_CALL_KEY, + ClassMethodNode, + ) + + # TODO(sang): unify option passing + options = { + "name": name, + "num_returns": num_returns, + "concurrency_group": concurrency_group, + "_generator_backpressure_num_objects": _generator_backpressure_num_objects, + } + + actor = self._actor + if actor is None: + # Ref is GC'ed. It happens when the actor handle is GC'ed + # when bind is called. + raise RuntimeError("Lost reference to actor") + + other_args_to_resolve = { + PARENT_CLASS_NODE_KEY: actor, + PREV_CLASS_METHOD_CALL_KEY: None, + BIND_INDEX_KEY: actor._ray_dag_bind_index, + } + actor._ray_dag_bind_index += 1 + + assert ( + self._signature is not None + ), "self._signature should be set for .bind API." + try: + signature.validate_args(self._signature, args, kwargs) + except TypeError as e: + signature_copy = self._signature.copy() + if len(signature_copy) > 0 and signature_copy[-1].name == "_ray_trace_ctx": + # Remove the trace context arg for readability. + signature_copy.pop(-1) + signature_copy = inspect.Signature(parameters=signature_copy) + raise TypeError( + f"{str(e)}. The function `{self._method_name}` has a signature " + f"`{signature_copy}`, but the given arguments to `bind` doesn't " + f"match. args: {args}. kwargs: {kwargs}." + ) from None + + node = ClassMethodNode( + self._method_name, + args, + kwargs, + options, + other_args_to_resolve=other_args_to_resolve, + ) + + if node.num_returns > 1: + output_nodes: List[ClassMethodNode] = [] + for i in range(node.num_returns): + output_node = ClassMethodNode( + f"return_idx_{i}", + (node, i), + dict(), + dict(), + {IS_CLASS_METHOD_OUTPUT_KEY: True, PARENT_CLASS_NODE_KEY: actor}, + ) + output_nodes.append(output_node) + return tuple(output_nodes) + else: + return node + + @wrap_auto_init + @_tracing_actor_method_invocation + def _remote( + self, + args=None, + kwargs=None, + name="", + num_returns=None, + max_task_retries=None, + retry_exceptions=None, + concurrency_group=None, + _generator_backpressure_num_objects=None, + enable_task_events=None, + tensor_transport: Optional[TensorTransportEnum] = None, + ): + if num_returns is None: + num_returns = self._num_returns + if max_task_retries is None: + max_task_retries = self._max_task_retries + if max_task_retries is None: + max_task_retries = 0 + if retry_exceptions is None: + retry_exceptions = self._retry_exceptions + if enable_task_events is None: + enable_task_events = self._enable_task_events + if _generator_backpressure_num_objects is None: + _generator_backpressure_num_objects = ( + self._generator_backpressure_num_objects + ) + + if tensor_transport is None: + tensor_transport = self._tensor_transport + + if tensor_transport != TensorTransportEnum.OBJECT_STORE: + if num_returns != 1: + raise ValueError( + f"Currently, methods with tensor_transport={tensor_transport.name} only support 1 return value. " + "Please make sure the actor method is decorated with `@ray.method(num_returns=1)` (the default)." + ) + if not self._actor._ray_enable_tensor_transport: + raise ValueError( + f'Currently, methods with .options(tensor_transport="{tensor_transport.name}") are not supported when enable_tensor_transport=False. ' + "Please set @ray.remote(enable_tensor_transport=True) on the actor class definition." + ) + gpu_object_manager = ray._private.worker.global_worker.gpu_object_manager + if not gpu_object_manager.actor_has_tensor_transport( + self._actor, tensor_transport + ): + raise ValueError( + f'{self._actor} does not have tensor transport {tensor_transport.name} available. If using a collective-based transport ("nccl" or "gloo"), please create a communicator with ' + "`ray.experimental.collective.create_collective_group` " + "before calling actor tasks with non-default tensor_transport." + ) + + args = args or [] + kwargs = kwargs or {} + + def invocation(args, kwargs): + dst_actor = self._actor + if dst_actor is None: + # See https://github.com/ray-project/ray/issues/6265 for more details. + raise RuntimeError( + "Lost reference to actor. Actor handles must be stored as variables, e.g. `actor = MyActor.remote()` before calling methods." + ) + + gpu_object_manager = ray._private.worker.global_worker.gpu_object_manager + gpu_object_manager.trigger_out_of_band_tensor_transfer(dst_actor, args) + + return dst_actor._actor_method_call( + self._method_name, + args=args, + kwargs=kwargs, + name=name, + num_returns=num_returns, + max_task_retries=max_task_retries, + retry_exceptions=retry_exceptions, + concurrency_group_name=concurrency_group, + generator_backpressure_num_objects=( + _generator_backpressure_num_objects + ), + enable_task_events=enable_task_events, + tensor_transport=tensor_transport, + ) + + # Apply the decorator if there is one. + if self._decorator is not None: + invocation = self._decorator(invocation) + + object_refs = invocation(args, kwargs) + if tensor_transport != TensorTransportEnum.OBJECT_STORE: + # Currently, we only support transfer tensor out-of-band when + # num_returns is 1. + assert isinstance(object_refs, ObjectRef) + object_ref = object_refs + gpu_object_manager = ray._private.worker.global_worker.gpu_object_manager + gpu_object_manager.add_gpu_object_ref( + object_ref, self._actor, tensor_transport + ) + + return object_refs + + def __getstate__(self): + return { + "actor": self._actor, + "method_name": self._method_name, + "num_returns": self._num_returns, + "max_task_retries": self._max_task_retries, + "retry_exceptions": self._retry_exceptions, + "decorator": self._decorator, + "is_generator": self._is_generator, + "generator_backpressure_num_objects": self._generator_backpressure_num_objects, # noqa + "enable_task_events": self._enable_task_events, + "_tensor_transport": self._tensor_transport, + } + + def __setstate__(self, state): + self.__init__( + state["actor"], + state["method_name"], + state["num_returns"], + state["max_task_retries"], + state["retry_exceptions"], + state["is_generator"], + state["generator_backpressure_num_objects"], + state["enable_task_events"], + state["decorator"], + state["_tensor_transport"], + ) + + +class _ActorClassMethodMetadata(object): + """Metadata for all methods in an actor class. This data can be cached. + + Attributes: + methods: The actor methods. + decorators: Optional decorators that should be applied to the + method invocation function before invoking the actor methods. These + can be set by attaching the attribute + "__ray_invocation_decorator__" to the actor method. + signatures: The signatures of the methods. + num_returns: The default number of return values for + each actor method. + max_task_retries: Number of retries on method failure. + retry_exceptions: Boolean of whether you want to retry all user-raised + exceptions, or a list of allowlist exceptions to retry, for each method. + enable_task_events: True if tracing is enabled, i.e., task events from + the actor should be reported. Defaults to True. + """ + + _cache = {} # This cache will be cleared in ray._private.worker.disconnect() + + def __init__(self): + class_name = type(self).__name__ + raise TypeError( + f"{class_name} can not be constructed directly, " + f"instead of running '{class_name}()', " + f"try '{class_name}.create()'" + ) + + @classmethod + def reset_cache(cls): + cls._cache.clear() + + @classmethod + def create( + cls, + modified_class, + actor_creation_function_descriptor, + ): + # Try to create an instance from cache. + cached_meta = cls._cache.get(actor_creation_function_descriptor) + if cached_meta is not None: + return cached_meta + + # Create an instance without __init__ called. + self = cls.__new__(cls) + + actor_methods = inspect.getmembers(modified_class, is_function_or_method) + self.methods = dict(actor_methods) + + # Extract the signatures of each of the methods. This will be used + # to catch some errors if the methods are called with inappropriate + # arguments. + self.decorators = {} + self.signatures = {} + self.num_returns = {} + self.max_task_retries = {} + self.retry_exceptions = {} + self.method_is_generator = {} + self.enable_task_events = {} + self.generator_backpressure_num_objects = {} + self.concurrency_group_for_methods = {} + self.method_name_to_tensor_transport: Dict[str, TensorTransportEnum] = {} + + # Check whether any actor methods specify a non-default tensor transport. + self.has_tensor_transport_methods = any( + getattr( + method, "__ray_tensor_transport__", TensorTransportEnum.OBJECT_STORE + ) + != TensorTransportEnum.OBJECT_STORE + for _, method in actor_methods + ) + + for method_name, method in actor_methods: + # Whether or not this method requires binding of its first + # argument. For class and static methods, we do not want to bind + # the first argument, but we do for instance methods + method = inspect.unwrap(method) + is_bound = is_class_method(method) or is_static_method( + modified_class, method_name + ) + + # Print a warning message if the method signature is not + # supported. We don't raise an exception because if the actor + # inherits from a class that has a method whose signature we + # don't support, there may not be much the user can do about it. + self.signatures[method_name] = signature.extract_signature( + method, ignore_first=not is_bound + ) + # Set the default number of return values for this method. + if hasattr(method, "__ray_num_returns__"): + self.num_returns[method_name] = method.__ray_num_returns__ + else: + self.num_returns[method_name] = None + + # Only contains entries from `@ray.method(max_task_retries=...)` + # Ray may not populate the others with max_task_retries here because you may + # have set in `actor.method.options(max_task_retries=...)`. So Ray always + # stores max_task_retries both from the method and from the actor, and + # favors the former. + if hasattr(method, "__ray_max_task_retries__"): + self.max_task_retries[method_name] = method.__ray_max_task_retries__ + + if hasattr(method, "__ray_retry_exceptions__"): + self.retry_exceptions[method_name] = method.__ray_retry_exceptions__ + + if hasattr(method, "__ray_invocation_decorator__"): + self.decorators[method_name] = method.__ray_invocation_decorator__ + + if hasattr(method, "__ray_concurrency_group__"): + self.concurrency_group_for_methods[ + method_name + ] = method.__ray_concurrency_group__ + + if hasattr(method, "__ray_enable_task_events__"): + self.enable_task_events[method_name] = method.__ray_enable_task_events__ + + is_generator = inspect.isgeneratorfunction( + method + ) or inspect.isasyncgenfunction(method) + self.method_is_generator[method_name] = is_generator + + if hasattr(method, "__ray_generator_backpressure_num_objects__"): + self.generator_backpressure_num_objects[ + method_name + ] = method.__ray_generator_backpressure_num_objects__ + + if hasattr(method, "__ray_tensor_transport__"): + self.method_name_to_tensor_transport[ + method_name + ] = method.__ray_tensor_transport__ + + # Update cache. + cls._cache[actor_creation_function_descriptor] = self + return self + + +class _ActorClassMetadata: + """Metadata for an actor class. + + Attributes: + language: The actor language, e.g. Python, Java. + modified_class: The original class that was decorated (with some + additional methods added like __ray_terminate__). + actor_creation_function_descriptor: The function descriptor for + the actor creation task. + class_id: The ID of this actor class. + method_meta: The actor method metadata. + class_name: The name of this class. + num_cpus: The default number of CPUs required by the actor creation + task. + num_gpus: The default number of GPUs required by the actor creation + task. + memory: The heap memory quota for this actor. + resources: The default resources required by the actor creation task. + label_selector: The labels required for the node on which this actor + can be scheduled on. The label selector consist of key-value pairs, where the keys + are label names and the value are expressions consisting of an operator with label + values or just a value to indicate equality. + fallback_strategy: If specified, expresses soft constraints through a list of decorator + options to fall back on when scheduling on a node. Decorator options are evaluated + together during scheduling. The first satisfied dict of options is used. Currently + only `label_selector` is a supported option. + accelerator_type: The specified type of accelerator required for the + node on which this actor runs. + See :ref:`accelerator types `. + runtime_env: The runtime environment for this actor. + scheduling_strategy: Strategy about how to schedule this actor. + last_export_cluster_and_job: A pair of the last exported cluster + and job to help us to know whether this function was exported. + This is an imperfect mechanism used to determine if we need to + export the remote function again. It is imperfect in the sense that + the actor class definition could be exported multiple times by + different workers. + enable_tensor_transport: Whether to enable out-of-band tensor transport + for this actor. + """ + + def __init__( + self, + language, + modified_class, + actor_creation_function_descriptor, + class_id, + method_meta, + max_restarts, + max_task_retries, + num_cpus, + num_gpus, + memory, + object_store_memory, + resources, + label_selector, + fallback_strategy, + accelerator_type, + runtime_env, + concurrency_groups, + scheduling_strategy: SchedulingStrategyT, + enable_tensor_transport: bool, + ): + self.language = language + self.modified_class = modified_class + self.actor_creation_function_descriptor = actor_creation_function_descriptor + self.method_meta = method_meta + self.class_name = actor_creation_function_descriptor.class_name + self.is_cross_language = language != Language.PYTHON + self.class_id = class_id + self.max_restarts = max_restarts + self.max_task_retries = max_task_retries + self.num_cpus = num_cpus + self.num_gpus = num_gpus + self.memory = memory + self.object_store_memory = object_store_memory + self.resources = resources + self.label_selector = label_selector + self.fallback_strategy = fallback_strategy + self.accelerator_type = accelerator_type + self.runtime_env = runtime_env + self.concurrency_groups = concurrency_groups + self.scheduling_strategy = scheduling_strategy + self.last_export_cluster_and_job = None + self.enable_tensor_transport = enable_tensor_transport + + +@PublicAPI +class ActorClassInheritanceException(TypeError): + pass + + +def _process_option_dict(actor_options, has_tensor_transport_methods): + _filled_options = {} + arg_names = set(inspect.getfullargspec(_ActorClassMetadata.__init__)[0]) + for k, v in ray_option_utils.actor_options.items(): + if k in arg_names: + _filled_options[k] = actor_options.get(k, v.default_value) + _filled_options["runtime_env"] = parse_runtime_env_for_task_or_actor( + _filled_options["runtime_env"] + ) + # If any actor method has a non-default tensor transport, automatically + # enable tensor transport, unless it was explicitly set to False by the + # user. + if has_tensor_transport_methods: + if _filled_options["enable_tensor_transport"] is False: + raise ValueError( + "Actor class has methods with @ray.method(tensor_transport=...) decorator but @ray.remote(enable_tensor_transport=False). " + "Either set enable_tensor_transport=True or remove the @ray.method(tensor_transport=...) decorator from the methods." + ) + _filled_options["enable_tensor_transport"] = True + + # Ray GPU objects requires a background thread for data transfer. However, + # currently by default the background thread will be blocked if the main + # thread does not yield. For now, we explicitly create the background thread + # if `@ray.remote(enable_tensor_transport=True)` or if any methods are + # decorated with `@ray.method(tensor_transport=...)` and a non-default + # tensor transport. This forces Ray to execute all tasks on background + # threads instead of the main thread. + # TODO(swang): Remove this code once + # https://github.com/ray-project/ray/issues/54639 is fixed. + enable_tensor_transport = _filled_options.get("enable_tensor_transport", False) + if enable_tensor_transport: + if _filled_options.get("concurrency_groups", None) is None: + _filled_options["concurrency_groups"] = {} + _filled_options["concurrency_groups"]["_ray_system"] = 1 + _filled_options["concurrency_groups"]["_ray_system_error"] = 1 + + return _filled_options + + +@PublicAPI +class ActorClass(Generic[T]): + """An actor class. + + This is a decorated class. It can be used to create actors. + + Attributes: + __ray_metadata__: Contains metadata for the actor. + """ + + def __init__(cls, name, bases, attr): + """Prevents users from directly inheriting from an ActorClass. + + This will be called when a class is defined with an ActorClass object + as one of its base classes. To intentionally construct an ActorClass, + use the '_ray_from_modified_class' classmethod. + + Raises: + ActorClassInheritanceException: When ActorClass is inherited. + AssertionError: If ActorClassInheritanceException is not raised i.e., + conditions for raising it are not met in any + iteration of the loop. + TypeError: In all other cases. + """ + for base in bases: + if isinstance(base, ActorClass): + raise ActorClassInheritanceException( + f"Attempted to define subclass '{name}' of actor " + f"class '{base.__ray_metadata__.class_name}'. " + "Inheriting from actor classes is " + "not currently supported. You can instead " + "inherit from a non-actor base class and make " + "the derived class an actor class (with " + "@ray.remote)." + ) + + # This shouldn't be reached because one of the base classes must be + # an actor class if this was meant to be subclassed. + assert False, ( + "ActorClass.__init__ should not be called. Please use " + "the @ray.remote decorator instead." + ) + + def __call__(self, *args, **kwargs): + """Prevents users from directly instantiating an ActorClass. + + This will be called instead of __init__ when 'ActorClass()' is executed + because an is an object rather than a metaobject. To properly + instantiated a remote actor, use 'ActorClass.remote()'. + + Raises: + Exception: Always. + """ + raise TypeError( + "Actors cannot be instantiated directly. " + f"Instead of '{self.__ray_metadata__.class_name}()', " + f"use '{self.__ray_metadata__.class_name}.remote()'." + ) + + @classmethod + def _ray_from_modified_class( + cls, + modified_class, + class_id, + actor_options, + ): + for attribute in [ + "remote", + "_remote", + "_ray_from_modified_class", + "_ray_from_function_descriptor", + ]: + if hasattr(modified_class, attribute): + logger.warning( + "Creating an actor from class " + f"{modified_class.__name__} overwrites " + f"attribute {attribute} of that class" + ) + + # Make sure the actor class we are constructing inherits from the + # original class so it retains all class properties. + class DerivedActorClass(cls, modified_class): + def __init__(self, *args, **kwargs): + try: + cls.__init__(self, *args, **kwargs) + except Exception as e: + # Delegate call to modified_class.__init__ only + # if the exception raised by cls.__init__ is + # TypeError and not ActorClassInheritanceException(TypeError). + # In all other cases proceed with raise e. + if isinstance(e, TypeError) and not isinstance( + e, ActorClassInheritanceException + ): + modified_class.__init__(self, *args, **kwargs) + else: + raise e + + name = f"ActorClass({modified_class.__name__})" + DerivedActorClass.__module__ = modified_class.__module__ + DerivedActorClass.__name__ = name + DerivedActorClass.__qualname__ = name + # Construct the base object. + self = DerivedActorClass.__new__(DerivedActorClass) + # Actor creation function descriptor. + actor_creation_function_descriptor = PythonFunctionDescriptor.from_class( + modified_class.__ray_actor_class__ + ) + + actor_method_meta = _ActorClassMethodMetadata.create( + modified_class, + actor_creation_function_descriptor, + ) + self.__ray_metadata__ = _ActorClassMetadata( + Language.PYTHON, + modified_class, + actor_creation_function_descriptor, + class_id, + actor_method_meta, + **_process_option_dict( + actor_options, actor_method_meta.has_tensor_transport_methods + ), + ) + self._default_options = actor_options + if "runtime_env" in self._default_options: + self._default_options["runtime_env"] = self.__ray_metadata__.runtime_env + + return self + + @classmethod + def _ray_from_function_descriptor( + cls, + language, + actor_creation_function_descriptor, + actor_options, + ): + self = ActorClass.__new__(ActorClass) + modified_class = None + actor_method_meta = _ActorClassMethodMetadata.create( + modified_class, + actor_creation_function_descriptor, + ) + self.__ray_metadata__ = _ActorClassMetadata( + language, + modified_class, + actor_creation_function_descriptor, + None, + actor_method_meta, + **_process_option_dict( + actor_options, actor_method_meta.has_tensor_transport_methods + ), + ) + self._default_options = actor_options + if "runtime_env" in self._default_options: + self._default_options["runtime_env"] = self.__ray_metadata__.runtime_env + return self + + def remote(self, *args, **kwargs) -> ActorProxy[T]: + """Create an actor. + + Args: + args: These arguments are forwarded directly to the actor + constructor. + kwargs: These arguments are forwarded directly to the actor + constructor. + + Returns: + A handle to the newly created actor. + """ + return self._remote(args=args, kwargs=kwargs, **self._default_options) + + def options(self, **actor_options) -> "ActorClass[T]": + """Configures and overrides the actor instantiation parameters. + + The arguments are the same as those that can be passed + to :obj:`ray.remote`. + + Args: + num_cpus: The quantity of CPU cores to reserve + for this task or for the lifetime of the actor. + num_gpus: The quantity of GPUs to reserve + for this task or for the lifetime of the actor. + resources (Dict[str, float]): The quantity of various custom resources + to reserve for this task or for the lifetime of the actor. + This is a dictionary mapping strings (resource names) to floats. + label_selector (Dict[str, str]): If specified, requires that the actor run + on a node which meets the specified label conditions (equals, in, not in, etc.). + fallback_strategy (List[Dict[str, Any]]): If specified, expresses soft constraints + through a list of decorator options to fall back on when scheduling on a node. + accelerator_type: If specified, requires that the task or actor run + on a node with the specified type of accelerator. + See :ref:`accelerator types `. + memory: The heap memory request in bytes for this task/actor, + rounded down to the nearest integer. + object_store_memory: The object store memory request for actors only. + max_restarts: This specifies the maximum + number of times that the actor should be restarted when it dies + unexpectedly. The minimum valid value is 0 (default), + which indicates that the actor doesn't need to be restarted. + A value of -1 indicates that an actor should be restarted + indefinitely. + max_task_retries: How many times to retry an actor task if the task + fails due to a runtime error, e.g., the actor has died. The + default value is 0. If set to -1, the system will retry the + failed task until the task succeeds, or the actor has reached + its max_restarts limit. If set to `n > 0`, the system will retry + the failed task up to n times, after which the task will throw a + `RayActorError` exception upon :obj:`ray.get`. Note that Python + exceptions may trigger retries + *only if* `retry_exceptions` is set for the method, in that case + when `max_task_retries` runs out the task will rethrow the + exception from the task. You can override this number with the + method's `max_task_retries` option in `@ray.method` decorator or + in `.option()`. + max_pending_calls: Set the max number of pending calls + allowed on the actor handle. When this value is exceeded, + PendingCallsLimitExceeded will be raised for further tasks. + Note that this limit is counted per handle. -1 means that the + number of pending calls is unlimited. + max_concurrency: The max number of concurrent calls to allow for + this actor. This only works with direct actor calls. The max + concurrency defaults to 1 for threaded execution, and 1000 for + asyncio execution. Note that the execution order is not + guaranteed when max_concurrency > 1. + allow_out_of_order_execution: Only for *actors*. Whether Ray executes actor + tasks out of order. If you're using multi-threaded + (``max_concurrency > 1``) or async actors, you can't set this to False. + Defaults to True if you're using multi-threaded or async actors, and + False otherwise. Actor task retries are always executed out of order. + name: The globally unique name for the actor, which can be used + to retrieve the actor via ray.get_actor(name) as long as the + actor is still alive. + namespace: Override the namespace to use for the actor. By default, + actors are created in an anonymous namespace. The actor can + be retrieved via ray.get_actor(name=name, namespace=namespace). + lifetime: Either `None`, which defaults to the actor will fate + share with its creator and will be deleted once its refcount + drops to zero, or "detached", which means the actor will live + as a global object independent of the creator. + runtime_env (Dict[str, Any]): Specifies the runtime environment for + this actor or task and its children. See + :ref:`runtime-environments` for detailed documentation. + scheduling_strategy: Strategy about how to + schedule a remote function or actor. Possible values are + None: ray will figure out the scheduling strategy to use, it + will either be the PlacementGroupSchedulingStrategy using parent's + placement group if parent has one and has + placement_group_capture_child_tasks set to true, + or "DEFAULT"; + "DEFAULT": default hybrid scheduling; + "SPREAD": best effort spread scheduling; + `PlacementGroupSchedulingStrategy`: + placement group based scheduling; + `NodeAffinitySchedulingStrategy`: + node id based affinity scheduling. + enable_task_events: True if tracing is enabled, i.e., task events from + the actor should be reported. Defaults to True. + + Examples: + + .. code-block:: python + + @ray.remote(num_cpus=2, resources={"CustomResource": 1}) + class Foo: + def method(self): + return 1 + # Class Bar will require 1 cpu instead of 2. + # It will also require no custom resources. + Bar = Foo.options(num_cpus=1, resources=None) + """ + + actor_cls = self + + # override original options + default_options = self._default_options.copy() + # "concurrency_groups" could not be used in ".options()", + # we should remove it before merging options from '@ray.remote'. + default_options.pop("concurrency_groups", None) + updated_options = ray_option_utils.update_options( + default_options, actor_options + ) + ray_option_utils.validate_actor_options(updated_options, in_options=True) + + # only update runtime_env when ".options()" specifies new runtime_env + if "runtime_env" in actor_options: + updated_options["runtime_env"] = parse_runtime_env_for_task_or_actor( + updated_options["runtime_env"] + ) + + class ActorOptionWrapper: + def remote(self, *args, **kwargs): + return actor_cls._remote(args=args, kwargs=kwargs, **updated_options) + + @DeveloperAPI + def bind(self, *args, **kwargs): + """ + For Ray DAG building that creates static graph from decorated + class or functions. + """ + from ray.dag.class_node import ClassNode + + return ClassNode( + actor_cls.__ray_metadata__.modified_class, + args, + kwargs, + updated_options, + ) + + return ActorOptionWrapper() + + @wrap_auto_init + @_tracing_actor_creation + def _remote(self, args=None, kwargs=None, **actor_options) -> ActorProxy[T]: + """Create an actor. + + This method allows more flexibility than the remote method because + resource requirements can be specified and override the defaults in the + decorator. + + Args: + args: The arguments to forward to the actor constructor. + kwargs: The keyword arguments to forward to the actor constructor. + **actor_options: Keyword arguments for configuring the actor options. + See ``ActorClass.options`` for more details. + + Returns: + A handle to the newly created actor. + """ + name = actor_options.get("name") + namespace = actor_options.get("namespace") + if name is not None: + if not isinstance(name, str): + raise TypeError(f"name must be None or a string, got: '{type(name)}'.") + elif name == "": + raise ValueError("Actor name cannot be an empty string.") + if namespace is not None: + ray._private.utils.validate_namespace(namespace) + + # Handle the get-or-create case. + if actor_options.get("get_if_exists"): + try: + return ray.get_actor(name, namespace=namespace) + except ValueError: + # Attempt to create it (may race with other attempts). + updated_options = actor_options.copy() + updated_options["get_if_exists"] = False # prevent infinite loop + try: + return self._remote(args, kwargs, **updated_options) + except ValueError: + # We lost the creation race, ignore. + pass + return ray.get_actor(name, namespace=namespace) + + # We pop the "concurrency_groups" coming from "@ray.remote" here. We no longer + # need it in "_remote()". + actor_options.pop("concurrency_groups", None) + + if args is None: + args = [] + if kwargs is None: + kwargs = {} + meta = self.__ray_metadata__ + is_asyncio = has_async_methods(meta.modified_class) + + if actor_options.get("max_concurrency") is None: + actor_options["max_concurrency"] = ( + DEFAULT_MAX_CONCURRENCY_ASYNC + if is_asyncio + else ray_constants.DEFAULT_MAX_CONCURRENCY_THREADED + ) + + if client_mode_should_convert(): + return client_mode_convert_actor(self, args, kwargs, **actor_options) + + # fill actor required options + for k, v in ray_option_utils.actor_options.items(): + actor_options[k] = actor_options.get(k, v.default_value) + # "concurrency_groups" already takes effects and should not apply again. + # Remove the default value here. + actor_options.pop("concurrency_groups", None) + + # TODO(suquark): cleanup these fields + max_concurrency = actor_options["max_concurrency"] + lifetime = actor_options["lifetime"] + runtime_env = actor_options["runtime_env"] + placement_group = actor_options["placement_group"] + placement_group_bundle_index = actor_options["placement_group_bundle_index"] + placement_group_capture_child_tasks = actor_options[ + "placement_group_capture_child_tasks" + ] + scheduling_strategy = actor_options["scheduling_strategy"] + max_restarts = actor_options["max_restarts"] + max_task_retries = actor_options["max_task_retries"] + max_pending_calls = actor_options["max_pending_calls"] + + # Override enable_task_events to default for actor if not specified (i.e. None) + enable_task_events = actor_options.get("enable_task_events") + + if scheduling_strategy is None or not isinstance( + scheduling_strategy, PlacementGroupSchedulingStrategy + ): + _warn_if_using_deprecated_placement_group(actor_options, 3) + + worker = ray._private.worker.global_worker + worker.check_connected() + + if worker.mode != ray._private.worker.WORKER_MODE: + from ray._common.usage import usage_lib + + usage_lib.record_library_usage("core") + + # Check whether the name is already taken. + # TODO(edoakes): this check has a race condition because two drivers + # could pass the check and then create the same named actor. We should + # instead check this when we create the actor, but that's currently an + # async call. + if name is not None: + try: + ray.get_actor(name, namespace=namespace) + except ValueError: # Name is not taken. + pass + else: + raise ValueError( + f"The name {name} (namespace={namespace}) is already " + "taken. Please use " + "a different name or get the existing actor using " + f"ray.get_actor('{name}', namespace='{namespace}')" + ) + + if lifetime is None: + detached = None + elif lifetime == "detached": + detached = True + elif lifetime == "non_detached": + detached = False + else: + raise ValueError( + "actor `lifetime` argument must be one of 'detached', " + "'non_detached' and 'None'." + ) + + # LOCAL_MODE cannot handle cross_language + if worker.mode == ray.LOCAL_MODE: + assert ( + not meta.is_cross_language + ), "Cross language ActorClass cannot be executed locally." + + # Export the actor. + if not meta.is_cross_language and ( + meta.last_export_cluster_and_job != worker.current_cluster_and_job + ): + # If this actor class was not exported in this cluster and job, + # we need to export this function again, because current GCS + # doesn't have it. + + # After serialize / deserialize modified class, the __module__ + # of modified class will be ray.cloudpickle.cloudpickle. + # So, here pass actor_creation_function_descriptor to make + # sure export actor class correct. + worker.function_actor_manager.export_actor_class( + meta.modified_class, + meta.actor_creation_function_descriptor, + meta.method_meta.methods.keys(), + ) + meta.last_export_cluster_and_job = worker.current_cluster_and_job + + resources = ray._common.utils.resources_from_ray_options(actor_options) + # Set the actor's default resources if not already set. First three + # conditions are to check that no resources were specified in the + # decorator. Last three conditions are to check that no resources were + # specified when _remote() was called. + # TODO(suquark): In the original code, memory is not considered as resources, + # when deciding the default CPUs. It is strange, but we keep the original + # semantics in case that it breaks user applications & tests. + if not set(resources.keys()).difference({"memory", "object_store_memory"}): + # In the default case, actors acquire no resources for + # their lifetime, and actor methods will require 1 CPU. + resources.setdefault("CPU", ray_constants.DEFAULT_ACTOR_CREATION_CPU_SIMPLE) + actor_method_cpu = ray_constants.DEFAULT_ACTOR_METHOD_CPU_SIMPLE + else: + # If any resources are specified (here or in decorator), then + # all resources are acquired for the actor's lifetime and no + # resources are associated with methods. + resources.setdefault( + "CPU", ray_constants.DEFAULT_ACTOR_CREATION_CPU_SPECIFIED + ) + actor_method_cpu = ray_constants.DEFAULT_ACTOR_METHOD_CPU_SPECIFIED + + # If the actor methods require CPU resources, then set the required + # placement resources. If actor_placement_resources is empty, then + # the required placement resources will be the same as resources. + actor_placement_resources = {} + assert actor_method_cpu in [0, 1] + if actor_method_cpu == 1: + actor_placement_resources = resources.copy() + actor_placement_resources["CPU"] += 1 + if meta.is_cross_language: + creation_args = cross_language._format_args(worker, args, kwargs) + else: + function_signature = meta.method_meta.signatures["__init__"] + creation_args = signature.flatten_args(function_signature, args, kwargs) + + use_placement_group = scheduling_strategy is not None and isinstance( + scheduling_strategy, PlacementGroupSchedulingStrategy + ) + is_restartable = max_restarts > 0 or max_restarts == -1 + if use_placement_group and detached and is_restartable: + # TODO(kevin85421): Checking `max_restarts > 0` is because Ray Serve currently schedules detached actors with + # placement groups. Adding the check avoids printing this warning for all Ray Serve applications. In the future, + # we should consider raising an error instead of a warning, but this is a breaking change. + logger.warning( + "Scheduling a restartable detached actor with a placement group is not recommended " + "because Ray will kill the actor when the placement group is removed and the actor will " + "not be able to be restarted." + ) + + if scheduling_strategy is None or isinstance( + scheduling_strategy, PlacementGroupSchedulingStrategy + ): + # TODO(jjyao) Clean this up once the + # placement_group option is removed. + # We should also consider pushing this logic down to c++ + # so that it can be reused by all languages. + if isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy): + placement_group = scheduling_strategy.placement_group + placement_group_bundle_index = ( + scheduling_strategy.placement_group_bundle_index + ) + placement_group_capture_child_tasks = ( + scheduling_strategy.placement_group_capture_child_tasks + ) + + if placement_group_capture_child_tasks is None: + placement_group_capture_child_tasks = ( + worker.should_capture_child_tasks_in_placement_group + ) + placement_group = _configure_placement_group_based_on_context( + placement_group_capture_child_tasks, + placement_group_bundle_index, + resources, + actor_placement_resources, + meta.class_name, + placement_group=placement_group, + ) + if not placement_group.is_empty: + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group, + placement_group_bundle_index, + placement_group_capture_child_tasks, + ) + else: + scheduling_strategy = "DEFAULT" + + serialized_runtime_env_info = None + if runtime_env is not None: + serialized_runtime_env_info = get_runtime_env_info( + runtime_env, + is_job_runtime_env=False, + serialize=True, + ) + + concurrency_groups_dict = {} + if meta.concurrency_groups is None: + meta.concurrency_groups = [] + for cg_name in meta.concurrency_groups: + concurrency_groups_dict[cg_name] = { + "name": cg_name, + "max_concurrency": meta.concurrency_groups[cg_name], + "function_descriptors": [], + } + + # Update methods + for method_name in meta.method_meta.concurrency_group_for_methods: + cg_name = meta.method_meta.concurrency_group_for_methods[method_name] + assert cg_name in concurrency_groups_dict + + module_name = meta.actor_creation_function_descriptor.module_name + class_name = meta.actor_creation_function_descriptor.class_name + concurrency_groups_dict[cg_name]["function_descriptors"].append( + PythonFunctionDescriptor(module_name, method_name, class_name) + ) + + # Update the creation descriptor based on number of arguments + if meta.is_cross_language: + func_name = "" + if meta.language == Language.CPP: + func_name = meta.actor_creation_function_descriptor.function_name + meta.actor_creation_function_descriptor = ( + cross_language._get_function_descriptor_for_actor_method( + meta.language, + meta.actor_creation_function_descriptor, + func_name, + str(len(args) + len(kwargs)), + ) + ) + + allow_out_of_order_execution = actor_options.get("allow_out_of_order_execution") + + # If the actor is async or multi-threaded, default to out-of-order execution. + if allow_out_of_order_execution is None: + allow_out_of_order_execution = is_asyncio or max_concurrency > 1 + + if is_asyncio and not allow_out_of_order_execution: + raise ValueError( + "If you're using async actors, Ray can't execute actor tasks in order. " + "Set `allow_out_of_order_execution=True` to allow out-of-order " + "execution." + ) + + elif max_concurrency > 1 and not allow_out_of_order_execution: + raise ValueError( + "If you're using multi-threaded actors, Ray can't execute actor tasks " + "in order. Set `allow_out_of_order_execution=True` to allow " + "out-of-order execution." + ) + + actor_id = worker.core_worker.create_actor( + meta.language, + meta.actor_creation_function_descriptor, + creation_args, + max_restarts, + max_task_retries, + resources, + actor_placement_resources, + max_concurrency, + detached, + name if name is not None else "", + namespace if namespace is not None else "", + is_asyncio, + # Store actor_method_cpu in actor handle's extension data. + extension_data=str(actor_method_cpu), + serialized_runtime_env_info=serialized_runtime_env_info or "{}", + concurrency_groups_dict=concurrency_groups_dict or dict(), + max_pending_calls=max_pending_calls, + scheduling_strategy=scheduling_strategy, + enable_task_events=enable_task_events, + labels=actor_options.get("_labels"), + label_selector=actor_options.get("label_selector"), + fallback_strategy=actor_options.get("fallback_strategy"), + allow_out_of_order_execution=allow_out_of_order_execution, + enable_tensor_transport=meta.enable_tensor_transport, + ) + + if _actor_launch_hook: + _actor_launch_hook( + meta.actor_creation_function_descriptor, resources, scheduling_strategy + ) + + actor_handle = ActorHandle( + meta.language, + actor_id, + max_task_retries, + enable_task_events, + meta.method_meta.method_is_generator, + meta.method_meta.decorators, + meta.method_meta.signatures, + meta.method_meta.num_returns, + meta.method_meta.max_task_retries, + meta.method_meta.retry_exceptions, + meta.method_meta.generator_backpressure_num_objects, + meta.method_meta.enable_task_events, + meta.enable_tensor_transport, + meta.method_meta.method_name_to_tensor_transport, + actor_method_cpu, + meta.actor_creation_function_descriptor, + worker.current_cluster_and_job, + original_handle=True, + allow_out_of_order_execution=allow_out_of_order_execution, + ) + + return actor_handle + + @DeveloperAPI + def bind(self, *args, **kwargs): + """ + For Ray DAG building that creates static graph from decorated + class or functions. + """ + from ray.dag.class_node import ClassNode + + return ClassNode( + self.__ray_metadata__.modified_class, args, kwargs, self._default_options + ) + + +@PublicAPI +class ActorHandle(Generic[T]): + """A handle to an actor. + + The fields in this class are prefixed with _ray_ to hide them from the user + and to avoid collision with actor method names. + + An ActorHandle can be created in three ways. First, by calling .remote() on + an ActorClass. Second, by passing an actor handle into a task (forking the + ActorHandle). Third, by directly serializing the ActorHandle (e.g., with + cloudpickle). + + Attributes: + _ray_actor_language: The actor language. + _ray_actor_id: Actor ID. + _ray_enable_task_events: The default value of whether task events is + enabled, i.e., task events from the actor should be reported. + _ray_method_is_generator: Map of method name -> if it is a generator + method. + _ray_method_decorators: Optional decorators for the function + invocation. This can be used to change the behavior on the + invocation side, whereas a regular decorator can be used to change + the behavior on the execution side. + _ray_method_signatures: The signatures of the actor methods. + _ray_method_max_task_retries: Max number of retries on method failure. + _ray_method_num_returns: The default number of return values for + each method. + _ray_method_retry_exceptions: The default value of boolean of whether you want + to retry all user-raised exceptions, or a list of allowlist exceptions to + retry. + _ray_method_generator_backpressure_num_objects: Generator-only + config. The max number of objects to generate before it + starts pausing a generator. + _ray_method_enable_task_events: The value of whether task + tracing is enabled for the actor methods. This overrides the + actor's default value (`_ray_enable_task_events`). + _ray_method_name_to_tensor_transport: A dictionary mapping method names to their tensor transport protocol settings. + The valid values are OBJECT_STORE (default), NCCL, or GLOO, and they are case-insensitive. + _ray_actor_method_cpus: The number of CPUs required by actor methods. + _ray_original_handle: True if this is the original actor handle for a + given actor. If this is true, then the actor will be destroyed when + this handle goes out of scope. + _ray_weak_ref: True means that this handle does not count towards the + distributed ref count for the actor, i.e. the actor may be GCed + while this handle is still in scope. This is set to True if the + handle was created by getting an actor by name or by getting the + self handle. It is set to False if this is the original handle or + if it was created by passing the original handle through task args + and returns. + _ray_is_cross_language: Whether this actor is cross language. + _ray_actor_creation_function_descriptor: The function descriptor + of the actor creation task. + _ray_allow_out_of_order_execution: Whether the actor can execute tasks out of order. + _ray_enable_tensor_transport: Whether tensor transport is enabled for this actor. + """ + + def __init__( + self, + language, + actor_id, + max_task_retries: Optional[int], + enable_task_events: bool, + method_is_generator: Dict[str, bool], + method_decorators, + method_signatures, + method_num_returns: Dict[str, Union[int, Literal["streaming"]]], + method_max_task_retries: Dict[str, int], + method_retry_exceptions: Dict[str, Union[bool, list, tuple]], + method_generator_backpressure_num_objects: Dict[str, int], + method_enable_task_events: Dict[str, bool], + enable_tensor_transport: bool, + method_name_to_tensor_transport: Dict[str, TensorTransportEnum], + actor_method_cpus: int, + actor_creation_function_descriptor, + cluster_and_job, + original_handle=False, + weak_ref: bool = False, + allow_out_of_order_execution: Optional[bool] = None, + ): + """Initialize an ActorHandle. + + Args: + language: The actor language. + actor_id: The ID of the actor. + max_task_retries: The maximum number of times to retry a task when it fails. + enable_task_events: Whether task events should be enabled for this actor. + method_is_generator: Dictionary mapping method names to whether they are generator methods. + method_decorators: Dictionary mapping method names to their decorators. + method_signatures: Dictionary mapping method names to their signatures. + method_num_returns: Dictionary mapping method names to their number of return values. + method_max_task_retries: Dictionary mapping method names to their maximum task retries. + method_retry_exceptions: Dictionary mapping method names to their retry exception settings. + method_generator_backpressure_num_objects: Dictionary mapping method names to their generator backpressure settings. + method_enable_task_events: Dictionary mapping method names to whether task events are enabled. + enable_tensor_transport: Whether tensor transport is enabled for + this actor. If True, then methods can be called with + .options(tensor_transport=...) to specify a non-default tensor + transport. + method_name_to_tensor_transport: Dictionary mapping method names to their tensor transport settings. + actor_method_cpus: The number of CPUs required by actor methods. + actor_creation_function_descriptor: The function descriptor for actor creation. + cluster_and_job: The cluster and job information. + original_handle: Whether this is the original actor handle. + weak_ref: Whether this is a weak reference to the actor. + allow_out_of_order_execution: Whether the actor can execute tasks out of order. + """ + self._ray_actor_language = language + self._ray_actor_id = actor_id + self._ray_max_task_retries = max_task_retries + self._ray_original_handle = original_handle + self._ray_weak_ref = weak_ref + self._ray_enable_task_events = enable_task_events + self._ray_allow_out_of_order_execution = allow_out_of_order_execution + + self._ray_method_is_generator = method_is_generator + self._ray_method_decorators = method_decorators + self._ray_method_signatures = method_signatures + self._ray_method_num_returns = method_num_returns + self._ray_method_max_task_retries = method_max_task_retries + self._ray_method_retry_exceptions = method_retry_exceptions + self._ray_method_generator_backpressure_num_objects = ( + method_generator_backpressure_num_objects + ) + self._ray_method_enable_task_events = method_enable_task_events + self._ray_enable_tensor_transport = enable_tensor_transport + self._ray_method_name_to_tensor_transport = method_name_to_tensor_transport + self._ray_actor_method_cpus = actor_method_cpus + self._ray_cluster_and_job = cluster_and_job + self._ray_is_cross_language = language != Language.PYTHON + self._ray_actor_creation_function_descriptor = ( + actor_creation_function_descriptor + ) + self._ray_function_descriptor = {} + # This is incremented each time `bind()` is called on an actor handle + # (in Ray DAGs), therefore capturing the bind order of the actor methods. + # TODO: this does not work properly if the caller has two copies of the + # same actor handle, and needs to be fixed. + self._ray_dag_bind_index = 0 + + if not self._ray_is_cross_language: + assert isinstance( + actor_creation_function_descriptor, PythonFunctionDescriptor + ) + module_name = actor_creation_function_descriptor.module_name + class_name = actor_creation_function_descriptor.class_name + for method_name in self._ray_method_signatures.keys(): + function_descriptor = PythonFunctionDescriptor( + module_name, method_name, class_name + ) + self._ray_function_descriptor[method_name] = function_descriptor + + # Build an _ActorMethodMetadata per method to cache expensive parsing logic. + # The _ActorMethodMetadata doesn't take a reference to this ActorHandle to avoid a circular reference. + # Instead, we will lazily bind this ActorHandle to the _ActorMethodMetadata when a method is invoked. + self._method_shells = {} + for method_name, method_signature in self._ray_method_signatures.items(): + self._method_shells[method_name] = _ActorMethodMetadata( + method_name=method_name, + num_returns=self._ray_method_num_returns.get(method_name, None), + max_task_retries=self._ray_method_max_task_retries.get( + method_name, self._ray_max_task_retries + ) + or 0, + retry_exceptions=self._ray_method_retry_exceptions.get(method_name), + is_generator=self._ray_method_is_generator.get(method_name), + generator_backpressure_num_objects=self._ray_method_generator_backpressure_num_objects.get( + method_name + ), + enable_task_events=self._ray_method_enable_task_events.get( + method_name, self._ray_enable_task_events + ), + decorator=self._ray_method_decorators.get(method_name), + signature=method_signature, + tensor_transport=self._ray_method_name_to_tensor_transport.get( + method_name + ), + ) + + def __del__(self): + # Weak references don't count towards the distributed ref count, so no + # need to decrement the ref count. + if self._ray_weak_ref: + return + + try: + # Mark that this actor handle has gone out of scope. Once all actor + # handles are out of scope, the actor will exit. + if ray._private.worker: + worker = ray._private.worker.global_worker + if worker.connected and hasattr(worker, "core_worker"): + worker.core_worker.remove_actor_handle_reference(self._ray_actor_id) + except AttributeError: + # Suppress the attribute error which is caused by + # python destruction ordering issue. + # It only happen when python exits. + pass + + def _actor_method_call( + self, + method_name: str, + args: List[Any] = None, + kwargs: Dict[str, Any] = None, + name: str = "", + num_returns: Optional[Union[int, Literal["streaming"]]] = None, + max_task_retries: int = None, + retry_exceptions: Union[bool, list, tuple] = None, + concurrency_group_name: Optional[str] = None, + generator_backpressure_num_objects: Optional[int] = None, + enable_task_events: Optional[bool] = None, + tensor_transport: Optional[TensorTransportEnum] = None, + ): + """Method execution stub for an actor handle. + + This is the function that executes when + `actor.method_name.remote(*args, **kwargs)` is called. Instead of + executing locally, the method is packaged as a task and scheduled + to the remote actor instance. + + Args: + method_name: The name of the actor method to execute. + args: A list of arguments for the actor method. + kwargs: A dictionary of keyword arguments for the actor method. + name: The name to give the actor method call task. + num_returns: The number of return values for the method. + max_task_retries: Number of retries when method fails. + retry_exceptions: Boolean of whether you want to retry all user-raised + exceptions, or a list of allowlist exceptions to retry. + concurrency_group_name: The name of the concurrency group to use. + generator_backpressure_num_objects: The number of objects to generate + before applying backpressure. + enable_task_events: True if tracing is enabled, i.e., task events from + the actor should be reported. + tensor_transport: The tensor transport protocol to use for the actor method. + + Returns: + object_refs: A list of object refs returned by the remote actor + method. + """ + worker = ray._private.worker.global_worker + + args = args or [] + kwargs = kwargs or {} + if self._ray_is_cross_language: + list_args = cross_language._format_args(worker, args, kwargs) + function_descriptor = cross_language._get_function_descriptor_for_actor_method( # noqa: E501 + self._ray_actor_language, + self._ray_actor_creation_function_descriptor, + method_name, + # The signature for xlang should be "{length_of_arguments}" to handle + # overloaded methods. + signature=str(len(args) + len(kwargs)), + ) + else: + function_signature = self._ray_method_signatures[method_name] + + if not args and not kwargs and not function_signature: + list_args = [] + else: + list_args = signature.flatten_args(function_signature, args, kwargs) + function_descriptor = self._ray_function_descriptor[method_name] + + if worker.mode == ray.LOCAL_MODE: + assert ( + not self._ray_is_cross_language + ), "Cross language remote actor method cannot be executed locally." + + if num_returns == "dynamic": + num_returns = -1 + elif num_returns == "streaming": + # TODO(sang): This is a temporary private API. + # Remove it when we migrate to the streaming generator. + num_returns = ray._raylet.STREAMING_GENERATOR_RETURN + + retry_exception_allowlist = None + if retry_exceptions is None: + retry_exceptions = False + elif isinstance(retry_exceptions, (list, tuple)): + retry_exception_allowlist = tuple(retry_exceptions) + retry_exceptions = True + assert isinstance( + retry_exceptions, bool + ), "retry_exceptions can either be \ + boolean or list/tuple of exception types." + + if generator_backpressure_num_objects is None: + generator_backpressure_num_objects = -1 + + object_refs = worker.core_worker.submit_actor_task( + self._ray_actor_language, + self._ray_actor_id, + function_descriptor, + list_args, + name, + num_returns, + max_task_retries, + retry_exceptions, + retry_exception_allowlist, + self._ray_actor_method_cpus, + concurrency_group_name if concurrency_group_name is not None else b"", + generator_backpressure_num_objects, + enable_task_events, + tensor_transport.value, + ) + + if num_returns == STREAMING_GENERATOR_RETURN: + # Streaming generator will return a single ref + # that is for the generator task. + assert len(object_refs) == 1 + generator_ref = object_refs[0] + return ObjectRefGenerator(generator_ref, worker) + if len(object_refs) == 1: + object_refs = object_refs[0] + elif len(object_refs) == 0: + object_refs = None + + return object_refs + + def __getattr__(self, item: str) -> Any: + """Handle dynamic attribute access for actor methods. + + This method is called when accessing attributes that don't exist as direct + instance attributes. It's the core mechanism for actor method invocation. + + For Python actors (99% of cases): + - We use strict validation: only methods in _method_shells are allowed + - This prevents typos and provides clear error messages + - Returns a bound ActorMethod created from the cached _ActorMethodMetadata + + For cross-language actors: + - We can't validate method names client-side (the target language defines them) + - We allow arbitrary method calls to pass through + - Some Python-specific methods like `__ray_terminate__` are blocked with warnings + + Args: + item: The attribute/method name being accessed + + Returns: + ActorMethod: A bound method ready for .remote() calls + + Raises: + AttributeError: For Python actors when accessing non-existent methods + """ + # If this name matches a remote method, bind and return it. + if item in self._method_shells: + return self._method_shells[item].bind(self) + + if not self._ray_is_cross_language: + raise AttributeError( + f"'{type(self).__name__}' object has " f"no attribute '{item}'" + ) + if item in ["__ray_terminate__"]: + + class FakeActorMethod(object): + def __call__(self, *args, **kwargs): + raise TypeError( + "Actor methods cannot be called directly. Instead " + "of running 'object.{}()', try 'object.{}.remote()'.".format( + item, item + ) + ) + + def remote(self, *args, **kwargs): + logger.warning( + f"Actor method {item} is not supported by cross language." + ) + + return FakeActorMethod() + + return ActorMethod( + self, # actor + item, # method_name + ray_constants.DEFAULT_ACTOR_METHOD_NUM_RETURN_VALS, + 0, # max_task_retries + False, # retry_exceptions + False, # is_generator + self._ray_method_generator_backpressure_num_objects.get(item, -1), + self._ray_enable_task_events, # enable_task_events + # Currently, cross-lang actor method not support decorator + decorator=None, + signature=None, + ) + + # Make tab completion work. + def __dir__(self): + return self._ray_method_signatures.keys() + + def __repr__(self): + return ( + "Actor(" + f"{self._ray_actor_creation_function_descriptor.class_name}, " + f"{self._actor_id.hex()})" + ) + + def __hash__(self): + return hash(self._actor_id) + + def __eq__(self, __value): + return hash(self) == hash(__value) + + @property + def _actor_id(self): + return self._ray_actor_id + + def _get_local_state(self): + """Get the local actor state. + + NOTE: this method only returns accurate actor state + after a first actor method call is made against + this actor handle due to https://github.com/ray-project/ray/pull/24600. + + Returns: + ActorTableData.ActorState or None if the state is unknown. + """ + worker = ray._private.worker.global_worker + worker.check_connected() + return worker.core_worker.get_local_actor_state(self._ray_actor_id) + + def _serialization_helper(self): + """This is defined in order to make pickling work. + + Returns: + A dictionary of the information needed to reconstruct the object. + """ + worker = ray._private.worker.global_worker + worker.check_connected() + + if hasattr(worker, "core_worker"): + # Non-local mode + state = worker.core_worker.serialize_actor_handle(self._ray_actor_id) + else: + # Local mode + state = ( + { + "actor_language": self._ray_actor_language, + "actor_id": self._ray_actor_id, + "max_task_retries": self._ray_max_task_retries, + "enable_task_events": self._enable_task_events, + "method_is_generator": self._ray_method_is_generator, + "method_decorators": self._ray_method_decorators, + "method_signatures": self._ray_method_signatures, + "method_num_returns": self._ray_method_num_returns, + "method_max_task_retries": self._ray_method_max_task_retries, + "method_retry_exceptions": self._ray_method_retry_exceptions, + "method_generator_backpressure_num_objects": ( + self._ray_method_generator_backpressure_num_objects + ), + "method_enable_task_events": self._ray_method_enable_task_events, + "enable_tensor_transport": self._ray_enable_tensor_transport, + "method_name_to_tensor_transport": self._ray_method_name_to_tensor_transport, + "actor_method_cpus": self._ray_actor_method_cpus, + "actor_creation_function_descriptor": self._ray_actor_creation_function_descriptor, # noqa: E501 + }, + None, + ) + + return (*state, self._ray_weak_ref) + + @classmethod + def _deserialization_helper(cls, state, weak_ref: bool, outer_object_ref=None): + """This is defined in order to make pickling work. + + Args: + state: The serialized state of the actor handle. + outer_object_ref: The ObjectRef that the serialized actor handle + was contained in, if any. This is used for counting references + to the actor handle. + weak_ref: Whether this was serialized from an actor handle with a + weak ref to the actor. + + """ + worker = ray._private.worker.global_worker + worker.check_connected() + + if hasattr(worker, "core_worker"): + # Non-local mode + return worker.core_worker.deserialize_and_register_actor_handle( + state, + outer_object_ref, + weak_ref, + ) + else: + # Local mode + assert worker.current_cluster_and_job == state["current_cluster_and_job"] + return cls( + # TODO(swang): Accessing the worker's current task ID is not + # thread-safe. + state["actor_language"], + state["actor_id"], + state["max_task_retries"], + state["enable_task_events"], + state["method_is_generator"], + state["method_decorators"], + state["method_signatures"], + state["method_num_returns"], + state["method_max_task_retries"], + state["method_retry_exceptions"], + state["method_generator_backpressure_num_objects"], + state["method_enable_task_events"], + state["enable_tensor_transport"], + state["method_name_to_tensor_transport"], + state["actor_method_cpus"], + state["actor_creation_function_descriptor"], + state["current_cluster_and_job"], + ) + + def __reduce__(self): + """This code path is used by pickling but not by Ray forking.""" + (serialized, _, weak_ref) = self._serialization_helper() + # There is no outer object ref when the actor handle is + # deserialized out-of-band using pickle. + return ActorHandle._deserialization_helper, (serialized, weak_ref, None) + + +def _modify_class(cls): + # cls has been modified. + if hasattr(cls, "__ray_actor_class__"): + return cls + + # Give an error if cls is an old-style class. + if not issubclass(cls, object): + raise TypeError( + "The @ray.remote decorator cannot be applied to old-style " + "classes. In Python 2, you must declare the class with " + "'class ClassName(object):' instead of 'class ClassName:'." + ) + + # Modify the class to have additional default methods. + class Class(cls): + __ray_actor_class__ = cls # The original actor class + + def __ray_ready__(self): + return True + + def __ray_call__(self, fn, *args, **kwargs): + return fn(self, *args, **kwargs) + + def __ray_terminate__(self): + worker = ray._private.worker.global_worker + if worker.mode != ray.LOCAL_MODE: + ray.actor.exit_actor() + + Class.__module__ = cls.__module__ + Class.__name__ = cls.__name__ + + if not is_function_or_method(getattr(Class, "__init__", None)): + # Add __init__ if it does not exist. + # Actor creation will be executed with __init__ together. + + # Assign an __init__ function will avoid many checks later on. + def __init__(self): + pass + + Class.__init__ = __init__ + + return Class + + +def _make_actor(cls, actor_options): + Class = _modify_class(cls) + _inject_tracing_into_class(Class) + + if "max_restarts" in actor_options: + if actor_options["max_restarts"] != -1: # -1 represents infinite restart + # Make sure we don't pass too big of an int to C++, causing + # an overflow. + actor_options["max_restarts"] = min( + actor_options["max_restarts"], ray_constants.MAX_INT64_VALUE + ) + + return ActorClass._ray_from_modified_class( + Class, + ActorClassID.from_random(), + actor_options, + ) + + +@PublicAPI +def exit_actor(): + """Intentionally exit the current actor. + + This API can be used only inside an actor. Use ray.kill + API if you'd like to kill an actor using actor handle. + + When this API is called, an exception is raised and the actor + will exit immediately. For asyncio actors, there may be a short + delay before the actor exits if the API is called from a background + task. + Any queued methods will fail. Any ``atexit`` + handlers installed in the actor will be run. + + Raises: + TypeError: An exception is raised if this is a driver or this + worker is not an actor. + """ + worker = ray._private.worker.global_worker + if worker.mode == ray.WORKER_MODE and not worker.actor_id.is_nil(): + worker.core_worker.set_current_actor_should_exit() + # In asyncio actor mode, we can't raise SystemExit because it will just + # quit the asycnio event loop thread, not the main thread. Instead, we + # raise a custom error to the main thread to tell it to exit. + if worker.core_worker.current_actor_is_asyncio(): + raise AsyncioActorExit() + + # Set a flag to indicate this is an intentional actor exit. This + # reduces log verbosity. + raise_sys_exit_with_custom_error_message("exit_actor() is called.") + else: + raise TypeError( + "exit_actor API is called on a non-actor worker, " + f"{worker.mode}. Call this API inside an actor methods" + "if you'd like to exit the actor gracefully." + ) diff --git a/lib/python3.12/site-packages/ray/client_builder.py b/lib/python3.12/site-packages/ray/client_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..48d7f4cd07181d6063635c2ec63326a72132ebfd --- /dev/null +++ b/lib/python3.12/site-packages/ray/client_builder.py @@ -0,0 +1,376 @@ +import importlib +import inspect +import json +import logging +import os +import sys +import warnings +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple + +import ray.util.client_connect +from ray._private.ray_constants import ( + RAY_ADDRESS_ENVIRONMENT_VARIABLE, + RAY_NAMESPACE_ENVIRONMENT_VARIABLE, + RAY_RUNTIME_ENV_ENVIRONMENT_VARIABLE, +) +from ray._private.utils import get_ray_client_dependency_error, split_address +from ray._private.worker import BaseContext, init as ray_driver_init +from ray.job_config import JobConfig +from ray.util.annotations import Deprecated, PublicAPI + +logger = logging.getLogger(__name__) + +CLIENT_DOCS_URL = ( + "https://docs.ray.io/en/latest/cluster/running-applications/" + "job-submission/ray-client.html" +) + + +@dataclass +@PublicAPI +class ClientContext(BaseContext): + """ + Basic context manager for a ClientBuilder connection. + + """ + + dashboard_url: Optional[str] + python_version: str + ray_version: str + ray_commit: str + _num_clients: int + _context_to_restore: Optional[ray.util.client.RayAPIStub] + + def __enter__(self) -> "ClientContext": + self._swap_context() + return self + + def __exit__(self, *exc) -> None: + self._disconnect_with_context(False) + self._swap_context() + + def disconnect(self) -> None: + self._swap_context() + self._disconnect_with_context(True) + self._swap_context() + + def _swap_context(self): + if self._context_to_restore is not None: + self._context_to_restore = ray.util.client.ray.set_context( + self._context_to_restore + ) + + def _disconnect_with_context(self, force_disconnect: bool) -> None: + """ + Disconnect Ray. If it's a ray client and created with `allow_multiple`, + it will do nothing. For other cases this either disconnects from the + remote Client Server or shuts the current driver down. + """ + if ray.util.client.ray.is_connected(): + if ray.util.client.ray.is_default() or force_disconnect: + # This is the only client connection + ray.util.client_connect.disconnect() + elif ray._private.worker.global_worker.node is None: + # Already disconnected. + return + elif ray._private.worker.global_worker.node.is_head(): + logger.debug( + "The current Ray Cluster is scoped to this process. " + "Disconnecting is not possible as it will shutdown the " + "cluster." + ) + else: + # This is only a driver connected to an existing cluster. + ray.shutdown() + + +@Deprecated +class ClientBuilder: + """ + Builder for a Ray Client connection. This class can be subclassed by + custom builder classes to modify connection behavior to include additional + features or altered semantics. One example is the ``_LocalClientBuilder``. + """ + + def __init__(self, address: Optional[str]) -> None: + if get_ray_client_dependency_error() is not None: + raise ValueError( + "Ray Client requires pip package `ray[client]`. " + "If you installed the minimal Ray (e.g. `pip install ray`), " + "please reinstall by executing `pip install ray[client]`." + ) + self.address = address + self._job_config = JobConfig() + self._remote_init_kwargs = {} + # Whether to allow connections to multiple clusters" + # " (allow_multiple=True). + self._allow_multiple_connections = False + self._credentials = None + self._metadata = None + # Set to False if ClientBuilder is being constructed by internal + # methods + self._deprecation_warn_enabled = True + + def env(self, env: Dict[str, Any]) -> "ClientBuilder": + """ + Set an environment for the session. + Args: + env (Dict[st, Any]): A runtime environment to use for this + connection. See :ref:`runtime-environments` for what values are + accepted in this dict. + """ + self._job_config.set_runtime_env(env) + return self + + def namespace(self, namespace: str) -> "ClientBuilder": + """ + Sets the namespace for the session. + Args: + namespace: Namespace to use. + """ + self._job_config.set_ray_namespace(namespace) + return self + + def connect(self) -> ClientContext: + """ + Begin a connection to the address passed in via ray.client(...). + + Returns: + ClientInfo: Dataclass with information about the setting. This + includes the server's version of Python & Ray as well as the + dashboard_url. + """ + if self._deprecation_warn_enabled: + self._client_deprecation_warn() + # Fill runtime env/namespace from environment if not already set. + # Should be done *after* the deprecation warning, since warning will + # check if those values are already set. + self._fill_defaults_from_env() + + # If it has already connected to the cluster with allow_multiple=True, + # connect to the default one is not allowed. + # But if it has connected to the default one, connect to other clients + # with allow_multiple=True is allowed + default_cli_connected = ray.util.client.ray.is_connected() + has_cli_connected = ray.util.client.num_connected_contexts() > 0 + if ( + not self._allow_multiple_connections + and not default_cli_connected + and has_cli_connected + ): + raise ValueError( + "The client has already connected to the cluster " + "with allow_multiple=True. Please set allow_multiple=True" + " to proceed" + ) + + old_ray_cxt = None + if self._allow_multiple_connections: + old_ray_cxt = ray.util.client.ray.set_context(None) + + client_info_dict = ray.util.client_connect.connect( + self.address, + job_config=self._job_config, + _credentials=self._credentials, + ray_init_kwargs=self._remote_init_kwargs, + metadata=self._metadata, + ) + + dashboard_url = ray.util.client.ray._get_dashboard_url() + + cxt = ClientContext( + dashboard_url=dashboard_url, + python_version=client_info_dict["python_version"], + ray_version=client_info_dict["ray_version"], + ray_commit=client_info_dict["ray_commit"], + _num_clients=client_info_dict["num_clients"], + _context_to_restore=ray.util.client.ray.get_context(), + ) + if self._allow_multiple_connections: + ray.util.client.ray.set_context(old_ray_cxt) + return cxt + + def _fill_defaults_from_env(self): + # Check environment variables for default values + namespace_env_var = os.environ.get(RAY_NAMESPACE_ENVIRONMENT_VARIABLE) + if namespace_env_var and self._job_config.ray_namespace is None: + self.namespace(namespace_env_var) + + runtime_env_var = os.environ.get(RAY_RUNTIME_ENV_ENVIRONMENT_VARIABLE) + if runtime_env_var and self._job_config.runtime_env is None: + self.env(json.loads(runtime_env_var)) + + def _init_args(self, **kwargs) -> "ClientBuilder": + """ + When a client builder is constructed through ray.init, for example + `ray.init(ray://..., namespace=...)`, all of the + arguments passed into ray.init with non-default values are passed + again into this method. Custom client builders can override this method + to do their own handling/validation of arguments. + """ + # Use namespace and runtime_env from ray.init call + if kwargs.get("namespace") is not None: + self.namespace(kwargs["namespace"]) + del kwargs["namespace"] + if kwargs.get("runtime_env") is not None: + self.env(kwargs["runtime_env"]) + del kwargs["runtime_env"] + + if kwargs.get("allow_multiple") is True: + self._allow_multiple_connections = True + del kwargs["allow_multiple"] + + if "_credentials" in kwargs.keys(): + self._credentials = kwargs["_credentials"] + del kwargs["_credentials"] + + if "_metadata" in kwargs.keys(): + self._metadata = kwargs["_metadata"] + del kwargs["_metadata"] + + if kwargs: + expected_sig = inspect.signature(ray_driver_init) + extra_args = set(kwargs.keys()).difference(expected_sig.parameters.keys()) + if len(extra_args) > 0: + raise RuntimeError( + "Got unexpected kwargs: {}".format(", ".join(extra_args)) + ) + self._remote_init_kwargs = kwargs + unknown = ", ".join(kwargs) + logger.info( + "Passing the following kwargs to ray.init() " + f"on the server: {unknown}" + ) + return self + + def _client_deprecation_warn(self) -> None: + """ + Generates a warning for user's if this ClientBuilder instance was + created directly or through ray.client, instead of relying on + internal methods (ray.init, or auto init) + """ + namespace = self._job_config.ray_namespace + runtime_env = self._job_config.runtime_env + replacement_args = [] + if self.address: + if isinstance(self, _LocalClientBuilder): + # Address might be set for LocalClientBuilder if ray.client() + # is called while ray_current_cluster is set + # (see _get_builder_from_address). In this case, + # leave off the ray:// so the user attaches the driver directly + replacement_args.append(f'"{self.address}"') + else: + replacement_args.append(f'"ray://{self.address}"') + if namespace: + replacement_args.append(f'namespace="{namespace}"') + if runtime_env: + # Use a placeholder here, since the real runtime_env would be + # difficult to read if formatted in directly + replacement_args.append("runtime_env=") + args_str = ", ".join(replacement_args) + replacement_call = f"ray.init({args_str})" + + # Note: stack level is set to 3 since we want the warning to reach the + # call to ray.client(...).connect(). The intervening frames are + # connect() -> client_deprecation_warn() -> warnings.warn() + # https://docs.python.org/3/library/warnings.html#available-functions + warnings.warn( + "Starting a connection through `ray.client` will be deprecated " + "in future ray versions in favor of `ray.init`. See the docs for " + f"more details: {CLIENT_DOCS_URL}. You can replace your call to " + "`ray.client().connect()` with the following:\n" + f" {replacement_call}\n", + DeprecationWarning, + stacklevel=3, + ) + + +class _LocalClientBuilder(ClientBuilder): + def connect(self) -> ClientContext: + """ + Begin a connection to the address passed in via ray.client(...) + """ + if self._deprecation_warn_enabled: + self._client_deprecation_warn() + # Fill runtime env/namespace from environment if not already set. + # Should be done *after* the deprecation warning, since warning will + # check if those values are already set. + self._fill_defaults_from_env() + + connection_dict = ray.init(address=self.address, job_config=self._job_config) + return ClientContext( + dashboard_url=connection_dict["webui_url"], + python_version="{}.{}.{}".format( + sys.version_info[0], sys.version_info[1], sys.version_info[2] + ), + ray_version=ray.__version__, + ray_commit=ray.__commit__, + _num_clients=1, + _context_to_restore=None, + ) + + +def _split_address(address: str) -> Tuple[str, str]: + """ + Splits address into a module string (scheme) and an inner_address. + + If the scheme is not present, then "ray://" is prepended to the address. + """ + if "://" not in address: + address = "ray://" + address + return split_address(address) + + +def _get_builder_from_address(address: Optional[str]) -> ClientBuilder: + if address == "local": + return _LocalClientBuilder("local") + if address is None: + # NOTE: This is not placed in `Node::get_temp_dir_path`, because + # this file is accessed before the `Node` object is created. + address = ray._private.services.canonicalize_bootstrap_address(address) + return _LocalClientBuilder(address) + module_string, inner_address = _split_address(address) + try: + module = importlib.import_module(module_string) + except Exception as e: + raise RuntimeError( + f"Module: {module_string} does not exist.\n" + f"This module was parsed from Address: {address}" + ) from e + assert "ClientBuilder" in dir( + module + ), f"Module: {module_string} does not have ClientBuilder." + return module.ClientBuilder(inner_address) + + +@Deprecated +def client( + address: Optional[str] = None, _deprecation_warn_enabled: bool = True +) -> ClientBuilder: + """ + Creates a ClientBuilder based on the provided address. The address can be + of the following forms: + + * None: Connects to or creates a local cluster and connects to it. + * ``"local"``: Creates a new cluster locally and connects to it. + * ``"IP:Port"``: Connects to a Ray Client Server at the given address. + * ``"module://inner_address"``: load module.ClientBuilder & pass + inner_address + + The _deprecation_warn_enabled flag enables deprecation warnings, and is + for internal use only. Set it to False to suppress client deprecation + warnings. + """ + env_address = os.environ.get(RAY_ADDRESS_ENVIRONMENT_VARIABLE) + if env_address and address is None: + logger.debug( + f"Using address ({env_address}) instead of auto-detection " + f"because {RAY_ADDRESS_ENVIRONMENT_VARIABLE} is set." + ) + address = env_address + + builder = _get_builder_from_address(address) + # Disable client deprecation warn when ray.client is used internally + builder._deprecation_warn_enabled = _deprecation_warn_enabled + return builder diff --git a/lib/python3.12/site-packages/ray/cluster_utils.py b/lib/python3.12/site-packages/ray/cluster_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..bb8249ff13af4044fd3d7687f67a617ce2351761 --- /dev/null +++ b/lib/python3.12/site-packages/ray/cluster_utils.py @@ -0,0 +1,415 @@ +import copy +import json +import logging +import os +import subprocess +import tempfile +import time +from typing import Dict, Optional + +import yaml + +import ray +import ray._private.services +from ray._private import ray_constants +from ray._private.client_mode_hook import disable_client_hook +from ray._raylet import GcsClientOptions +from ray.autoscaler._private.fake_multi_node.node_provider import FAKE_HEAD_NODE_ID +from ray.util.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + +cluster_not_supported = os.name == "nt" + + +@DeveloperAPI +class AutoscalingCluster: + """Create a local autoscaling cluster for testing. + + See test_autoscaler_fake_multinode.py for an end-to-end example. + """ + + def __init__( + self, + head_resources: dict, + worker_node_types: dict, + autoscaler_v2: bool = False, + **config_kwargs, + ): + """Create the cluster. + + Args: + head_resources: resources of the head node, including CPU. + worker_node_types: autoscaler node types config for worker nodes. + """ + self._head_resources = head_resources + self._config = self._generate_config( + head_resources, + worker_node_types, + autoscaler_v2=autoscaler_v2, + **config_kwargs, + ) + self._autoscaler_v2 = autoscaler_v2 + + def _generate_config( + self, head_resources, worker_node_types, autoscaler_v2=False, **config_kwargs + ): + base_config = yaml.safe_load( + open( + os.path.join( + os.path.dirname(ray.__file__), + "autoscaler/_private/fake_multi_node/example.yaml", + ) + ) + ) + custom_config = copy.deepcopy(base_config) + custom_config["available_node_types"] = worker_node_types + custom_config["available_node_types"]["ray.head.default"] = { + "resources": head_resources, + "node_config": {}, + "max_workers": 0, + } + + # Autoscaler v2 specific configs + if autoscaler_v2: + custom_config["provider"]["launch_multiple"] = True + custom_config["provider"]["head_node_id"] = FAKE_HEAD_NODE_ID + custom_config.update(config_kwargs) + return custom_config + + def start(self, _system_config=None, override_env: Optional[Dict] = None): + """Start the cluster. + + After this call returns, you can connect to the cluster with + ray.init("auto"). + """ + subprocess.check_call(["ray", "stop", "--force"]) + _, fake_config = tempfile.mkstemp() + with open(fake_config, "w") as f: + f.write(json.dumps(self._config)) + cmd = [ + "ray", + "start", + "--autoscaling-config={}".format(fake_config), + "--head", + ] + if "CPU" in self._head_resources: + cmd.append("--num-cpus={}".format(self._head_resources.pop("CPU"))) + if "GPU" in self._head_resources: + cmd.append("--num-gpus={}".format(self._head_resources.pop("GPU"))) + if "object_store_memory" in self._head_resources: + cmd.append( + "--object-store-memory={}".format( + self._head_resources.pop("object_store_memory") + ) + ) + if self._head_resources: + cmd.append("--resources='{}'".format(json.dumps(self._head_resources))) + if _system_config is not None: + cmd.append( + "--system-config={}".format( + json.dumps(_system_config, separators=(",", ":")) + ) + ) + env = os.environ.copy() + env.update({"AUTOSCALER_UPDATE_INTERVAL_S": "1", "RAY_FAKE_CLUSTER": "1"}) + if self._autoscaler_v2: + # Set the necessary environment variables for autoscaler v2. + env.update( + { + "RAY_enable_autoscaler_v2": "1", + "RAY_CLOUD_INSTANCE_ID": FAKE_HEAD_NODE_ID, + "RAY_OVERRIDE_NODE_ID_FOR_TESTING": FAKE_HEAD_NODE_ID, + } + ) + if override_env: + env.update(override_env) + subprocess.check_call(cmd, env=env) + + def shutdown(self): + """Terminate the cluster.""" + subprocess.check_call(["ray", "stop", "--force"]) + + +@DeveloperAPI +class Cluster: + def __init__( + self, + initialize_head: bool = False, + connect: bool = False, + head_node_args: dict = None, + shutdown_at_exit: bool = True, + ): + """Initializes all services of a Ray cluster. + + Args: + initialize_head: Automatically start a Ray cluster + by initializing the head node. Defaults to False. + connect: If `initialize_head=True` and `connect=True`, + ray.init will be called with the address of this cluster + passed in. + head_node_args: Arguments to be passed into + `start_ray_head` via `self.add_node`. + shutdown_at_exit: If True, registers an exit hook + for shutting down all started processes. + """ + if cluster_not_supported: + logger.warning( + "Ray cluster mode is currently experimental and untested on " + "Windows. If you are using it and running into issues please " + "file a report at https://github.com/ray-project/ray/issues." + ) + self.head_node = None + self.worker_nodes = set() + self.redis_address = None + self.connected = False + # Create a new global state accessor for fetching GCS table. + self.global_state = ray._private.state.GlobalState() + self._shutdown_at_exit = shutdown_at_exit + if not initialize_head and connect: + raise RuntimeError("Cannot connect to uninitialized cluster.") + + if initialize_head: + head_node_args = head_node_args or {} + self.add_node(**head_node_args) + if connect: + self.connect() + + @property + def gcs_address(self): + if self.head_node is None: + return None + return self.head_node.gcs_address + + @property + def address(self): + return self.gcs_address + + def connect(self, namespace=None): + """Connect the driver to the cluster.""" + assert self.address is not None + assert not self.connected + output_info = ray.init( + namespace=namespace, + ignore_reinit_error=True, + address=self.address, + _redis_username=self.redis_username, + _redis_password=self.redis_password, + ) + logger.info(output_info) + self.connected = True + + def add_node(self, wait: bool = True, **node_args): + """Adds a node to the local Ray Cluster. + + All nodes are by default started with the following settings: + cleanup=True, + num_cpus=1, + object_store_memory=150 * 1024 * 1024 # 150 MiB + + Args: + wait: Whether to wait until the node is alive. + node_args: Keyword arguments used in `start_ray_head` and + `start_ray_node`. Overrides defaults. + + Returns: + Node object of the added Ray node. + """ + default_kwargs = { + "num_cpus": 1, + "num_gpus": 0, + "object_store_memory": 150 * 1024 * 1024, # 150 MiB + "min_worker_port": 0, + "max_worker_port": 0, + } + ray_params = ray._private.parameter.RayParams(**node_args) + ray_params.update_if_absent(**default_kwargs) + with disable_client_hook(): + if self.head_node is None: + node = ray._private.node.Node( + ray_params, + head=True, + shutdown_at_exit=self._shutdown_at_exit, + spawn_reaper=self._shutdown_at_exit, + ) + self.head_node = node + self.redis_address = self.head_node.redis_address + self.redis_username = node_args.get( + "redis_username", ray_constants.REDIS_DEFAULT_USERNAME + ) + self.redis_password = node_args.get( + "redis_password", ray_constants.REDIS_DEFAULT_PASSWORD + ) + self.webui_url = self.head_node.webui_url + # Init global state accessor when creating head node. + gcs_options = GcsClientOptions.create( + node.gcs_address, + None, + allow_cluster_id_nil=True, + fetch_cluster_id_if_nil=False, + ) + self.global_state._initialize_global_state(gcs_options) + # Write the Ray cluster address for convenience in unit + # testing. ray.init() and ray.init(address="auto") will connect + # to the local cluster. + ray._private.utils.write_ray_address(self.head_node.gcs_address) + else: + ray_params.update_if_absent(redis_address=self.redis_address) + ray_params.update_if_absent(gcs_address=self.gcs_address) + # We only need one log monitor per physical node. + ray_params.update_if_absent(include_log_monitor=False) + # Let grpc pick a port. + ray_params.update_if_absent(node_manager_port=0) + if "dashboard_agent_listen_port" not in node_args: + # Pick a random one to not conflict + # with the head node dashboard agent + ray_params.dashboard_agent_listen_port = None + + node = ray._private.node.Node( + ray_params, + head=False, + shutdown_at_exit=self._shutdown_at_exit, + spawn_reaper=self._shutdown_at_exit, + ) + self.worker_nodes.add(node) + + if wait: + # Wait for the node to appear in the client table. We do this + # so that the nodes appears in the client table in the order + # that the corresponding calls to add_node were made. We do + # this because in the tests we assume that the driver is + # connected to the first node that is added. + self._wait_for_node(node) + + return node + + def remove_node(self, node, allow_graceful=True): + """Kills all processes associated with worker node. + + Args: + node: Worker node of which all associated processes + will be removed. + """ + global_node = ray._private.worker.global_worker.node + if global_node is not None: + if node._raylet_socket_name == global_node._raylet_socket_name: + ray.shutdown() + raise ValueError( + "Removing a node that is connected to this Ray client " + "is not allowed because it will break the driver. " + "You can use the get_other_node utility to avoid removing " + "a node that the Ray client is connected." + ) + + node.destroy_external_storage() + if self.head_node == node: + # We have to wait to prevent the raylet becomes a zombie which will prevent + # worker from exiting + self.head_node.kill_all_processes( + check_alive=False, allow_graceful=allow_graceful, wait=True + ) + self.head_node = None + # TODO(rliaw): Do we need to kill all worker processes? + else: + # We have to wait to prevent the raylet becomes a zombie which will prevent + # worker from exiting + node.kill_all_processes( + check_alive=False, allow_graceful=allow_graceful, wait=True + ) + self.worker_nodes.remove(node) + + assert ( + not node.any_processes_alive() + ), "There are zombie processes left over after killing." + + def _wait_for_node(self, node, timeout: float = 30): + """Wait until this node has appeared in the client table. + + Args: + node (ray._private.node.Node): The node to wait for. + timeout: The amount of time in seconds to wait before raising an + exception. + + Raises: + TimeoutError: An exception is raised if the timeout expires before + the node appears in the client table. + """ + ray._private.services.wait_for_node( + node.gcs_address, + node.plasma_store_socket_name, + timeout, + ) + + def wait_for_nodes(self, timeout: float = 30): + """Waits for correct number of nodes to be registered. + + This will wait until the number of live nodes in the client table + exactly matches the number of "add_node" calls minus the number of + "remove_node" calls that have been made on this cluster. This means + that if a node dies without "remove_node" having been called, this will + raise an exception. + + Args: + timeout: The number of seconds to wait for nodes to join + before failing. + + Raises: + TimeoutError: An exception is raised if we time out while waiting + for nodes to join. + """ + start_time = time.time() + while time.time() - start_time < timeout: + live_clients = self.global_state._live_node_ids() + + expected = len(self.list_all_nodes()) + if len(live_clients) == expected: + logger.debug("All nodes registered as expected.") + return + else: + logger.debug( + f"{len(live_clients)} nodes are currently registered, " + f"but we are expecting {expected}" + ) + time.sleep(0.1) + raise TimeoutError("Timed out while waiting for nodes to join.") + + def list_all_nodes(self): + """Lists all nodes. + + TODO(rliaw): What is the desired behavior if a head node + dies before worker nodes die? + + Returns: + List of all nodes, including the head node. + """ + nodes = list(self.worker_nodes) + if self.head_node: + nodes = [self.head_node] + nodes + return nodes + + def remaining_processes_alive(self): + """Returns a bool indicating whether all processes are alive or not. + + Note that this ignores processes that have been explicitly killed, + e.g., via a command like node.kill_raylet(). + + Returns: + True if all processes are alive and false otherwise. + """ + return all(node.remaining_processes_alive() for node in self.list_all_nodes()) + + def shutdown(self): + """Removes all nodes.""" + + # We create a list here as a copy because `remove_node` + # modifies `self.worker_nodes`. + all_nodes = list(self.worker_nodes) + for node in all_nodes: + self.remove_node(node) + + if self.head_node is not None: + self.remove_node(self.head_node) + # need to reset internal kv since gcs is down + ray.experimental.internal_kv._internal_kv_reset() + # Delete the cluster address. + ray._common.utils.reset_ray_address() diff --git a/lib/python3.12/site-packages/ray/cross_language.py b/lib/python3.12/site-packages/ray/cross_language.py new file mode 100644 index 0000000000000000000000000000000000000000..1539954e56f0f45a9f2fbc89b99e72153db0b603 --- /dev/null +++ b/lib/python3.12/site-packages/ray/cross_language.py @@ -0,0 +1,137 @@ +from __future__ import absolute_import, division, print_function + +from ray import Language +from ray._raylet import CppFunctionDescriptor, JavaFunctionDescriptor +from ray.util.annotations import PublicAPI + +__all__ = [ + "java_function", + "java_actor_class", + "cpp_function", +] + + +@PublicAPI(stability="beta") +def java_function(class_name: str, function_name: str): + """Define a Java function. + + Args: + class_name: Java class name. + function_name: Java function name. + """ + from ray.remote_function import RemoteFunction + + return RemoteFunction( + Language.JAVA, + lambda *args, **kwargs: None, + JavaFunctionDescriptor(class_name, function_name, ""), + {}, + ) + + +@PublicAPI(stability="beta") +def cpp_function(function_name: str): + """Define a Cpp function. + + Args: + function_name: Cpp function name. + """ + from ray.remote_function import RemoteFunction + + return RemoteFunction( + Language.CPP, + lambda *args, **kwargs: None, + CppFunctionDescriptor(function_name, "PYTHON"), + {}, + ) + + +@PublicAPI(stability="beta") +def java_actor_class(class_name: str): + """Define a Java actor class. + + Args: + class_name: Java class name. + """ + from ray.actor import ActorClass + + return ActorClass._ray_from_function_descriptor( + Language.JAVA, + JavaFunctionDescriptor(class_name, "", ""), + {}, + ) + + +@PublicAPI(stability="beta") +def cpp_actor_class(create_function_name: str, class_name: str): + """Define a Cpp actor class. + + Args: + create_function_name: Create cpp class function name. + class_name: Cpp class name. + """ + from ray.actor import ActorClass + + print("create func=", create_function_name, "class_name=", class_name) + return ActorClass._ray_from_function_descriptor( + Language.CPP, + CppFunctionDescriptor(create_function_name, "PYTHON", class_name), + {}, + ) + + +def _format_args(worker, args, kwargs): + """Format args for various languages. + + Args: + worker: The global worker instance. + args: The arguments for cross language. + kwargs: The keyword arguments for cross language. + + Returns: + List of args and kwargs (if supported). + """ + if not worker.load_code_from_local: + raise ValueError( + "Cross language feature needs --load-code-from-local to be set." + ) + if kwargs: + raise TypeError( + f"Cross language remote functions does not support kwargs, " + f"kwargs:{str(kwargs)}." + ) + return args + + +def _get_function_descriptor_for_actor_method( + language: str, actor_creation_function_descriptor, method_name: str, signature: str +): + """Get function descriptor for cross language actor method call. + + Args: + language: Target language. + actor_creation_function_descriptor: + The function signature for actor creation. + method_name: The name of actor method. + signature: The signature for the actor method. When calling Java from Python, + it should be string in the form of "{length_of_args}". + + Returns: + Function descriptor for cross language actor method call. + """ + if language == Language.JAVA: + return JavaFunctionDescriptor( + actor_creation_function_descriptor.class_name, + method_name, + signature, + ) + elif language == Language.CPP: + return CppFunctionDescriptor( + method_name, + "PYTHON", + actor_creation_function_descriptor.class_name, + ) + else: + raise NotImplementedError( + "Cross language remote actor method " f"not support language {language}" + ) diff --git a/lib/python3.12/site-packages/ray/exceptions.py b/lib/python3.12/site-packages/ray/exceptions.py new file mode 100644 index 0000000000000000000000000000000000000000..3b9e8f35f7c6a1edacec8a2900447a95fb73013d --- /dev/null +++ b/lib/python3.12/site-packages/ray/exceptions.py @@ -0,0 +1,991 @@ +import logging +import os +import sys +from traceback import format_exception +from typing import Optional, Union + +import colorama + +import ray._private.ray_constants as ray_constants +import ray.cloudpickle as pickle +from ray._raylet import ActorID, TaskID, WorkerID +from ray.core.generated.common_pb2 import ( + PYTHON, + ActorDiedErrorContext, + Address, + Language, + NodeDeathInfo, + RayException, +) +from ray.util.annotations import DeveloperAPI, PublicAPI + +logger = logging.getLogger(__name__) + + +@PublicAPI +class RayError(Exception): + """Super class of all ray exception types.""" + + def to_bytes(self): + # Extract exc_info from exception object. + exc_info = (type(self), self, self.__traceback__) + formatted_exception_string = "\n".join(format_exception(*exc_info)) + return RayException( + language=PYTHON, + serialized_exception=pickle.dumps(self), + formatted_exception_string=formatted_exception_string, + ).SerializeToString() + + @staticmethod + def from_bytes(b): + ray_exception = RayException() + ray_exception.ParseFromString(b) + return RayError.from_ray_exception(ray_exception) + + @staticmethod + def from_ray_exception(ray_exception): + if ray_exception.language == PYTHON: + try: + return pickle.loads(ray_exception.serialized_exception) + except Exception: + # formatted_exception_string is set in to_bytes() above by calling + # traceback.format_exception() on the original exception. It contains + # the string representation and stack trace of the original error. + original_stacktrace = getattr( + ray_exception, + "formatted_exception_string", + "No formatted exception string available.", + ) + return UnserializableException(original_stacktrace) + else: + return CrossLanguageError(ray_exception) + + +@PublicAPI +class CrossLanguageError(RayError): + """Raised from another language.""" + + def __init__(self, ray_exception): + super().__init__( + "An exception raised from {}:\n{}".format( + Language.Name(ray_exception.language), + ray_exception.formatted_exception_string, + ) + ) + + +@PublicAPI +class TaskCancelledError(RayError): + """Raised when this task is cancelled. + + Args: + task_id: The TaskID of the function that was directly + cancelled. + """ + + def __init__( + self, task_id: Optional[TaskID] = None, error_message: Optional[str] = None + ): + self.task_id = task_id + self.error_message = error_message + + def __str__(self): + msg = "" + if self.task_id: + msg = "Task: " + str(self.task_id) + " was cancelled. " + if self.error_message: + msg += self.error_message + return msg + + +@PublicAPI +class RayTaskError(RayError): + """Indicates that a task threw an exception during execution. + + If a task throws an exception during execution, a RayTaskError is stored in + the object store for each of the task's outputs. When an object is + retrieved from the object store, the Python method that retrieved it checks + to see if the object is a RayTaskError and if it is then an exception is + thrown propagating the error message. + """ + + def __init__( + self, + function_name, + traceback_str, + cause, + proctitle=None, + pid=None, + ip=None, + actor_repr=None, + actor_id=None, + ): + """Initialize a RayTaskError.""" + import ray + + if proctitle: + self.proctitle = proctitle + else: + self.proctitle = ray._raylet.getproctitle() + self.pid = pid or os.getpid() + self.ip = ip or ray.util.get_node_ip_address() + self.function_name = function_name + self.traceback_str = traceback_str + self.actor_repr = actor_repr + self._actor_id = actor_id + self.cause = cause + + try: + pickle.dumps(cause) + except (pickle.PicklingError, TypeError) as e: + err_msg = ( + "The original cause of the RayTaskError" + f" ({self.cause.__class__}) isn't serializable: {e}." + " Overwriting the cause to a RayError." + ) + logger.warning(err_msg) + self.cause = RayError(err_msg) + + # BaseException implements a __reduce__ method that returns + # a tuple with the type and the value of self.args. + # https://stackoverflow.com/a/49715949/2213289 + self.args = (function_name, traceback_str, self.cause, proctitle, pid, ip) + + assert traceback_str is not None + + def make_dual_exception_instance(self) -> "RayTaskError": + """Makes a object instance that inherits from both RayTaskError and the type of + `self.cause`. Raises TypeError if the cause class can't be subclassed""" + # For normal user Exceptions, we subclass from both + # RayTaskError and the user exception. For ExceptionGroup, + # we special handle it because it has a different __new__() + # signature from Exception. + # Ref: https://docs.python.org/3/library/exceptions.html#exception-groups + if sys.version_info >= (3, 11) and isinstance( + self.cause, ExceptionGroup # noqa: F821 + ): + return self._make_exceptiongroup_dual_exception_instance() + return self._make_normal_dual_exception_instance() + + def _make_normal_dual_exception_instance(self) -> "RayTaskError": + cause_cls = self.cause.__class__ + error_msg = str(self) + + class cls(RayTaskError, cause_cls): + def __init__(self, cause): + self.cause = cause + self.args = (cause,) + + def __reduce__(self): + return (cls, self.args) + + def __getattr__(self, name): + return getattr(self.cause, name) + + def __str__(self): + return error_msg + + name = f"RayTaskError({cause_cls.__name__})" + cls.__name__ = name + cls.__qualname__ = name + + return cls(self.cause) + + def _make_exceptiongroup_dual_exception_instance(self) -> "RayTaskError": + cause_cls = self.cause.__class__ + error_msg = str(self) + + class cls(RayTaskError, cause_cls): + def __new__(cls, cause): + self = super().__new__(cls, cause.message, cause.exceptions) + return self + + def __init__(self, cause): + self.cause = cause + self.args = (cause,) + + def __reduce__(self): + return (cls, self.args) + + def __getattr__(self, name): + return getattr(self.cause, name) + + def __str__(self): + return error_msg + + name = f"RayTaskError({cause_cls.__name__})" + cls.__name__ = name + cls.__qualname__ = name + + return cls(self.cause) + + def as_instanceof_cause(self): + """Returns an exception that's an instance of the cause's class. + + The returned exception inherits from both RayTaskError and the + cause class and contains all of the attributes of the cause + exception. + + If the cause class can't be subclassed, issues a warning and returns `self`. + """ + cause_cls = self.cause.__class__ + if issubclass(RayTaskError, cause_cls): + return self # already satisfied + + try: + return self.make_dual_exception_instance() + except TypeError as e: + logger.warning( + f"User exception type {type(self.cause)} in RayTaskError can't" + " be subclassed! This exception is raised as" + " RayTaskError only. You can use `ray_task_error.cause` to" + f" access the user exception. Failure in subclassing: {e}" + ) + return self + + def __str__(self): + """Format a RayTaskError as a string.""" + lines = self.traceback_str.strip().split("\n") + out = [] + code_from_internal_file = False + + # Format tracebacks. + # Python stacktrace consists of + # Traceback...: Indicate the next line will be a traceback. + # File [file_name + line number] + # code + # XError: [message] + # NOTE: For _raylet.pyx (Cython), the code is not always included. + for i, line in enumerate(lines): + # Convert traceback to the readable information. + if line.startswith("Traceback "): + traceback_line = ( + f"{colorama.Fore.CYAN}" + f"{self.proctitle}()" + f"{colorama.Fore.RESET} " + f"(pid={self.pid}, ip={self.ip}" + ) + if self.actor_repr: + traceback_line += ( + f", actor_id={self._actor_id}, repr={self.actor_repr})" + ) + else: + traceback_line += ")" + code_from_internal_file = False + out.append(traceback_line) + elif line.startswith(" File ") and ( + "ray/worker.py" in line + or "ray/_private/" in line + or "ray/util/tracing/" in line + or "ray/_raylet.pyx" in line + ): + # TODO(windows) + # Process the internal file line. + # The file line always starts with 2 space and File. + # https://github.com/python/cpython/blob/0a0a135bae2692d069b18d2d590397fbe0a0d39a/Lib/traceback.py#L421 # noqa + if "ray._raylet.raise_if_dependency_failed" in line: + # It means the current task is failed + # due to the dependency failure. + # Print out an user-friendly + # message to explain that.. + out.append( + " At least one of the input arguments for " + "this task could not be computed:" + ) + if i + 1 < len(lines) and lines[i + 1].startswith(" "): + # If the next line is indented with 2 space, + # that means it contains internal code information. + # For example, + # File [file_name] [line] + # [code] # if the next line is indented, it is code. + # Note there there are 4 spaces in the code line. + code_from_internal_file = True + elif code_from_internal_file: + # If the current line is internal file's code, + # the next line is not code anymore. + code_from_internal_file = False + else: + out.append(line) + return "\n".join(out) + + +@PublicAPI +class LocalRayletDiedError(RayError): + """Indicates that the task's local raylet died.""" + + def __str__(self): + return "The task's local raylet died. Check raylet.out for more information." + + +@PublicAPI +class WorkerCrashedError(RayError): + """Indicates that the worker died unexpectedly while executing a task.""" + + def __str__(self): + return ( + "The worker died unexpectedly while executing this task. " + "Check python-core-worker-*.log files for more information." + ) + + +@PublicAPI +class RayActorError(RayError): + """Indicates that the actor has outages unexpectedly before finishing a task. + + This exception could happen because the actor process is dead, or is unavailable for + the moment. Ray raises subclasses `ActorDiedError` and `ActorUnavailableError` + respectively. + """ + + BASE_ERROR_MSG = "The actor experienced an error before finishing this task." + + def __init__( + self, + actor_id: str = None, + error_msg: str = BASE_ERROR_MSG, + actor_init_failed: bool = False, + preempted: bool = False, + ): + #: The actor ID in hex string. + self.actor_id = actor_id + #: Whether the actor failed in the middle of __init__. + self.error_msg = error_msg + #: The full error message. + self._actor_init_failed = actor_init_failed + #: Whether the actor died because the node was preempted. + self._preempted = preempted + + def __str__(self) -> str: + return self.error_msg + + @property + def preempted(self) -> bool: + return self._preempted + + @property + def actor_init_failed(self) -> bool: + return self._actor_init_failed + + +@DeveloperAPI +class ActorDiedError(RayActorError): + """Indicates that the actor died unexpectedly before finishing a task. + + This exception could happen either because the actor process dies while + executing a task, or because a task is submitted to a dead actor. + + Args: + cause: The cause of the actor error. `RayTaskError` type means + the actor has died because of an exception within `__init__`. + `ActorDiedErrorContext` means the actor has died because of + an unexpected system error. None means the cause isn't known. + Theoretically, this shouldn't happen, + but it's there as a safety check. + """ + + BASE_ERROR_MSG = "The actor died unexpectedly before finishing this task." + + def __init__( + self, cause: Optional[Union[RayTaskError, ActorDiedErrorContext]] = None + ): + """ + Construct a RayActorError by building the arguments. + """ + + actor_id = None + error_msg = ActorDiedError.BASE_ERROR_MSG + actor_init_failed = False + preempted = False + + if not cause: + # Use the defaults above. + pass + elif isinstance(cause, RayTaskError): + actor_init_failed = True + actor_id = cause._actor_id + error_msg = ( + "The actor died because of an error" + " raised in its creation task, " + f"{cause.__str__()}" + ) + else: + # Inidicating system-level actor failures. + assert isinstance(cause, ActorDiedErrorContext) + error_msg_lines = [ActorDiedError.BASE_ERROR_MSG] + error_msg_lines.append(f"\tclass_name: {cause.class_name}") + error_msg_lines.append(f"\tactor_id: {ActorID(cause.actor_id).hex()}") + # Below items are optional fields. + if cause.pid != 0: + error_msg_lines.append(f"\tpid: {cause.pid}") + if cause.name != "": + error_msg_lines.append(f"\tname: {cause.name}") + if cause.ray_namespace != "": + error_msg_lines.append(f"\tnamespace: {cause.ray_namespace}") + if cause.node_ip_address != "": + error_msg_lines.append(f"\tip: {cause.node_ip_address}") + error_msg_lines.append(cause.error_message) + if cause.never_started: + error_msg_lines.append( + "The actor never ran - it was cancelled before it started running." + ) + if ( + cause.node_death_info + and cause.node_death_info.reason + == NodeDeathInfo.AUTOSCALER_DRAIN_PREEMPTED + ): + preempted = True + error_msg = "\n".join(error_msg_lines) + actor_id = ActorID(cause.actor_id).hex() + super().__init__(actor_id, error_msg, actor_init_failed, preempted) + + @staticmethod + def from_task_error(task_error: RayTaskError): + return ActorDiedError(task_error) + + +@DeveloperAPI +class ActorUnavailableError(RayActorError): + """Raised when the actor is temporarily unavailable but may be available later.""" + + def __init__(self, error_message: str, actor_id: Optional[bytes]): + actor_id = ActorID(actor_id).hex() if actor_id is not None else None + error_msg = ( + f"The actor {actor_id} is unavailable: {error_message}. The task may or " + "may not have been executed on the actor." + ) + actor_init_failed = False + preempted = False + + super().__init__(actor_id, error_msg, actor_init_failed, preempted) + + +@PublicAPI +class RaySystemError(RayError): + """Indicates that Ray encountered a system error. + + This exception can be thrown when the raylet is killed. + """ + + def __init__(self, client_exc, traceback_str=None): + self.client_exc = client_exc + self.traceback_str = traceback_str + + def __str__(self): + error_msg = f"System error: {self.client_exc}" + if self.traceback_str: + error_msg += f"\ntraceback: {self.traceback_str}" + return error_msg + + +@PublicAPI +class AuthenticationError(RayError): + """Indicates that an authentication error occurred. + + Most commonly, this is caused by a missing or mismatching token set on the client + (e.g., a Ray CLI command interacting with a remote cluster). + + Only applicable when `RAY_AUTH_MODE` is not set to `disabled`. + """ + + def __init__(self, message: str): + self.message = message + + # Always hide traceback for cleaner output + self.__suppress_context__ = True + super().__init__(message) + + def __str__(self) -> str: + return self.message + ( + ". Ensure that you have `RAY_AUTH_MODE=token` set and the token for the cluster is available as the `RAY_AUTH_TOKEN` environment variable or a local file. " + "For more information, see: https://docs.ray.io/en/latest/ray-security/auth.html" + ) + + +@DeveloperAPI +class UserCodeException(RayError): + """Indicates that an exception occurred while executing user code. + For example, this exception can be used to wrap user code exceptions + from a remote task or actor. The `retry_exceptions` parameter will + still respect the underlying cause of this exception.""" + + pass + + +@PublicAPI +class ObjectStoreFullError(RayError): + """Indicates that the object store is full. + + This is raised if the attempt to store the object fails + because the object store is full even after multiple retries. + """ + + def __str__(self): + return super(ObjectStoreFullError, self).__str__() + ( + "\n" + "The local object store is full of objects that are still in " + "scope and cannot be evicted. Tip: Use the `ray memory` command " + "to list active objects in the cluster." + ) + + +@PublicAPI +class OutOfDiskError(RayError): + """Indicates that the local disk is full. + + This is raised if the attempt to store the object fails + because both the object store and disk are full. + """ + + def __str__(self): + # TODO(scv119): expose more disk usage information and link to a doc. + return super(OutOfDiskError, self).__str__() + ( + "\n" + "The object cannot be created because the local object store" + " is full and the local disk's utilization is over capacity" + " (95% by default)." + "Tip: Use `df` on this node to check disk usage and " + "`ray memory` to check object store memory usage." + ) + + +@PublicAPI +class OutOfMemoryError(RayError): + """Indicates that the node is running out of memory and is close to full. + + This is raised if the node is low on memory and tasks or actors are being + evicted to free up memory. + """ + + # TODO: (clarng) expose the error message string here and format it with proto + def __init__(self, message): + self.message = message + + def __str__(self): + return self.message + + +@PublicAPI +class NodeDiedError(RayError): + """Indicates that the node is either dead or unreachable.""" + + # TODO: (clarng) expose the error message string here and format it with proto + def __init__(self, message): + self.message = message + + def __str__(self): + return self.message + + +@PublicAPI +class ObjectLostError(RayError): + """Indicates that the object is lost from distributed memory, due to + node failure or system error. + + Args: + object_ref_hex: Hex ID of the object. + """ + + def __init__(self, object_ref_hex, owner_address, call_site): + self.object_ref_hex = object_ref_hex + self.owner_address = owner_address + self.call_site = call_site.replace( + ray_constants.CALL_STACK_LINE_DELIMITER, "\n " + ) + + def _base_str(self): + msg = f"Failed to retrieve object {self.object_ref_hex}. " + if self.call_site: + msg += f"The ObjectRef was created at: {self.call_site}" + else: + msg += ( + "To see information about where this ObjectRef was created " + "in Python, set the environment variable " + "RAY_record_ref_creation_sites=1 during `ray start` and " + "`ray.init()`." + ) + return msg + + def __str__(self): + return ( + self._base_str() + + "\n\n" + + ( + f"All copies of {self.object_ref_hex} have been lost due to node " + "failure. Check cluster logs (`/tmp/ray/session_latest/logs`) for " + "more information about the failure." + ) + ) + + +@PublicAPI +class ObjectFetchTimedOutError(ObjectLostError): + """Indicates that an object fetch timed out. + + Args: + object_ref_hex: Hex ID of the object. + """ + + def __str__(self): + return ( + self._base_str() + + "\n\n" + + ( + f"Fetch for object {self.object_ref_hex} timed out because no " + "locations were found for the object. This may indicate a " + "system-level bug." + ) + ) + + +@DeveloperAPI +class RpcError(RayError): + """Indicates an error in the underlying RPC system.""" + + def __init__(self, message, rpc_code=None): + self.message = message + self.rpc_code = rpc_code + + def __str__(self): + return self.message + + +@DeveloperAPI +class ReferenceCountingAssertionError(ObjectLostError, AssertionError): + """Indicates that an object has been deleted while there was still a + reference to it. + + Args: + object_ref_hex: Hex ID of the object. + """ + + def __str__(self): + return ( + self._base_str() + + "\n\n" + + ( + "The object has already been deleted by the reference counting " + "protocol. This should not happen." + ) + ) + + +@DeveloperAPI +class ObjectFreedError(ObjectLostError): + """Indicates that an object was manually freed by the application. + + Attributes: + object_ref_hex: Hex ID of the object. + """ + + def __str__(self): + return ( + self._base_str() + + "\n\n" + + ( + "The object was manually freed using the internal `free` call. " + "Please ensure that `free` is only called once the object is no " + "longer needed." + ) + ) + + +@PublicAPI +class OwnerDiedError(ObjectLostError): + """Indicates that the owner of the object has died while there is still a + reference to the object. + + Args: + object_ref_hex: Hex ID of the object. + """ + + def __str__(self): + log_loc = "`/tmp/ray/session_latest/logs`" + if self.owner_address: + try: + addr = Address() + addr.ParseFromString(self.owner_address) + ip_addr = addr.ip_address + worker_id = WorkerID(addr.worker_id) + log_loc = ( + f"`/tmp/ray/session_latest/logs/*{worker_id.hex()}*`" + f" at IP address {ip_addr}" + ) + except Exception: + # Catch all to make sure we always at least print the default + # message. + pass + + return ( + self._base_str() + + "\n\n" + + ( + "The object's owner has exited. This is the Python " + "worker that first created the ObjectRef via `.remote()` or " + "`ray.put()`. " + f"Check cluster logs ({log_loc}) for more " + "information about the Python worker failure." + ) + ) + + +@PublicAPI +class ObjectReconstructionFailedError(ObjectLostError): + """Indicates that the object cannot be reconstructed. + + Args: + object_ref_hex: Hex ID of the object. + """ + + def __str__(self): + return ( + self._base_str() + + "\n\n" + + ( + "The object cannot be reconstructed " + "because it was created by an actor, ray.put() call, or its " + "ObjectRef was created by a different worker." + ) + ) + + +@PublicAPI +class ObjectReconstructionFailedMaxAttemptsExceededError(ObjectLostError): + """Indicates that the object cannot be reconstructed because the maximum + number of task retries has been exceeded. + + Args: + object_ref_hex: Hex ID of the object. + """ + + def __str__(self): + return ( + self._base_str() + + "\n\n" + + ( + "The object cannot be reconstructed " + "because the maximum number of task retries has been exceeded. " + "To prevent this error, set " + "`@ray.remote(max_retries=)` (default 3)." + ) + ) + + +@PublicAPI +class ObjectReconstructionFailedLineageEvictedError(ObjectLostError): + """Indicates that the object cannot be reconstructed because its lineage + was evicted due to memory pressure. + + Args: + object_ref_hex: Hex ID of the object. + """ + + def __str__(self): + return ( + self._base_str() + + "\n\n" + + ( + "The object cannot be reconstructed because its lineage has been " + "evicted to reduce memory pressure. " + "To prevent this error, set the environment variable " + "RAY_max_lineage_bytes= (default 1GB) during `ray start`." + ) + ) + + +@PublicAPI +class GetTimeoutError(RayError, TimeoutError): + """Indicates that a call to the worker timed out.""" + + pass + + +@PublicAPI +class PlasmaObjectNotAvailable(RayError): + """Called when an object was not available within the given timeout.""" + + pass + + +@PublicAPI +class AsyncioActorExit(RayError): + """Raised when an asyncio actor intentionally exits via exit_actor().""" + + pass + + +@PublicAPI +class RuntimeEnvSetupError(RayError): + """Raised when a runtime environment fails to be set up. + + Args: + error_message: The error message that explains + why runtime env setup has failed. + """ + + def __init__(self, error_message: str = None): + self.error_message = error_message + + def __str__(self): + msgs = ["Failed to set up runtime environment."] + if self.error_message: + msgs.append(self.error_message) + return "\n".join(msgs) + + +@PublicAPI +class TaskPlacementGroupRemoved(RayError): + """Raised when the corresponding placement group was removed.""" + + def __str__(self): + return "The placement group corresponding to this task has been removed." + + +@PublicAPI +class ActorPlacementGroupRemoved(RayError): + """Raised when the corresponding placement group was removed.""" + + def __str__(self): + return "The placement group corresponding to this Actor has been removed." + + +@PublicAPI +class PendingCallsLimitExceeded(RayError): + """Raised when the pending actor calls exceeds `max_pending_calls` option. + + This exception could happen probably because the caller calls the callee + too frequently. + """ + + pass + + +@PublicAPI +class TaskUnschedulableError(RayError): + """Raised when the task cannot be scheduled. + + One example is that the node specified through + NodeAffinitySchedulingStrategy is dead. + """ + + def __init__(self, error_message: str): + self.error_message = error_message + + def __str__(self): + return f"The task is not schedulable: {self.error_message}" + + +@PublicAPI +class ActorUnschedulableError(RayError): + """Raised when the actor cannot be scheduled. + + One example is that the node specified through + NodeAffinitySchedulingStrategy is dead. + """ + + def __init__(self, error_message: str): + self.error_message = error_message + + def __str__(self): + return f"The actor is not schedulable: {self.error_message}" + + +@DeveloperAPI +class ObjectRefStreamEndOfStreamError(RayError): + """Raised by streaming generator tasks when there are no more ObjectRefs to + read. + """ + + pass + + +@DeveloperAPI +class OufOfBandObjectRefSerializationException(RayError): + """Raised when an `ray.ObjectRef` is out of band serialized by + `ray.cloudpickle`. It is an anti pattern. + """ + + pass + + +@PublicAPI(stability="alpha") +class RayChannelError(RaySystemError): + """Indicates that Ray encountered a system error related + to ray.experimental.channel. + """ + + pass + + +@PublicAPI(stability="alpha") +class RayChannelTimeoutError(RayChannelError, TimeoutError): + """Raised when the Compiled Graph channel operation times out.""" + + pass + + +@PublicAPI(stability="alpha") +class RayCgraphCapacityExceeded(RaySystemError): + """Raised when the Compiled Graph channel's buffer is at max capacity""" + + pass + + +@PublicAPI(stability="alpha") +class UnserializableException(RayError): + """Raised when there is an error deserializing a serialized exception. + + This occurs when deserializing (unpickling) a previously serialized exception + fails. In this case, we fall back to raising the string representation of + the original exception along with its stack trace that was captured at the + time of serialization. + + For more details and how to handle this with custom serializers, :ref:`configuring custom exeception serializers ` + + Args: + original_stack_trace: The string representation and stack trace of the + original exception that was captured during serialization. + """ + + def __init__(self, original_stack_trace: str): + self._original_stack_trace = original_stack_trace + + def __str__(self): + return ( + "Failed to deserialize exception. Refer to https://docs.ray.io/en/latest/ray-core/objects/serialization.html#custom-serializers-for-exceptions for more information.\n" + "Original exception:\n" + f"{self._original_stack_trace}" + ) + + +RAY_EXCEPTION_TYPES = [ + PlasmaObjectNotAvailable, + RayError, + RayTaskError, + WorkerCrashedError, + RayActorError, + ObjectStoreFullError, + ObjectLostError, + ObjectFetchTimedOutError, + ReferenceCountingAssertionError, + ObjectReconstructionFailedError, + ObjectReconstructionFailedMaxAttemptsExceededError, + ObjectReconstructionFailedLineageEvictedError, + OwnerDiedError, + GetTimeoutError, + AsyncioActorExit, + RuntimeEnvSetupError, + TaskPlacementGroupRemoved, + ActorPlacementGroupRemoved, + PendingCallsLimitExceeded, + LocalRayletDiedError, + TaskUnschedulableError, + ActorDiedError, + ActorUnschedulableError, + ActorUnavailableError, + RayChannelError, + RayChannelTimeoutError, + OufOfBandObjectRefSerializationException, + RayCgraphCapacityExceeded, + UnserializableException, + AuthenticationError, +] diff --git a/lib/python3.12/site-packages/ray/job_config.py b/lib/python3.12/site-packages/ray/job_config.py new file mode 100644 index 0000000000000000000000000000000000000000..5bb9d72eebb6dcf3783f67113d55d4b599014750 --- /dev/null +++ b/lib/python3.12/site-packages/ray/job_config.py @@ -0,0 +1,253 @@ +import uuid +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +import ray.cloudpickle as pickle +from ray._private.ray_logging.logging_config import LoggingConfig +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.runtime_env import RuntimeEnv + + +@PublicAPI +class JobConfig: + """A class used to store the configurations of a job. + + Examples: + .. testcode:: + :hide: + + import ray + ray.shutdown() + + .. testcode:: + + import ray + from ray.job_config import JobConfig + + ray.init(job_config=JobConfig(default_actor_lifetime="non_detached")) + + Args: + jvm_options: The jvm options for java workers of the job. + code_search_path: A list of directories or jar files that + specify the search path for user code. This will be used as + `CLASSPATH` in Java and `PYTHONPATH` in Python. + See :ref:`Ray cross-language programming ` for more details. + runtime_env: A :ref:`runtime environment ` dictionary. + metadata: An opaque metadata dictionary. + ray_namespace: A :ref:`namespace ` + is a logical grouping of jobs and named actors. + default_actor_lifetime: The default value of actor lifetime, + can be "detached" or "non_detached". + See :ref:`actor lifetimes ` for more details. + """ + + def __init__( + self, + jvm_options: Optional[List[str]] = None, + code_search_path: Optional[List[str]] = None, + runtime_env: Optional[dict] = None, + _client_job: bool = False, + metadata: Optional[dict] = None, + ray_namespace: Optional[str] = None, + default_actor_lifetime: str = "non_detached", + _py_driver_sys_path: Optional[List[str]] = None, + ): + #: The jvm options for java workers of the job. + self.jvm_options = jvm_options or [] + #: A list of directories or jar files that + #: specify the search path for user code. + self.code_search_path = code_search_path or [] + # It's difficult to find the error that caused by the + # code_search_path is a string. So we assert here. + assert isinstance(self.code_search_path, (list, tuple)), ( + f"The type of code search path is incorrect: " f"{type(code_search_path)}" + ) + self._client_job = _client_job + #: An opaque metadata dictionary. + self.metadata = metadata or {} + #: A namespace is a logical grouping of jobs and named actors. + self.ray_namespace = ray_namespace + self.set_runtime_env(runtime_env) + self.set_default_actor_lifetime(default_actor_lifetime) + # A list of directories that specify the search path for python workers. + self._py_driver_sys_path = _py_driver_sys_path or [] + # Python logging configurations that will be passed to Ray tasks/actors. + self.py_logging_config = None + + def set_metadata(self, key: str, value: str) -> None: + """Add key-value pair to the metadata dictionary. + + If the key already exists, the value is overwritten to the new value. + + Examples: + .. testcode:: + + import ray + from ray.job_config import JobConfig + + job_config = JobConfig() + job_config.set_metadata("submitter", "foo") + + Args: + key: The key of the metadata. + value: The value of the metadata. + """ + self.metadata[key] = value + + def _serialize(self) -> str: + """Serialize the struct into protobuf string""" + return self._get_proto_job_config().SerializeToString() + + def set_runtime_env( + self, + runtime_env: Optional[Union[Dict[str, Any], "RuntimeEnv"]], + validate: bool = False, + ) -> None: + """Modify the runtime_env of the JobConfig. + + We don't validate the runtime_env by default here because it may go + through some translation before actually being passed to C++ (e.g., + working_dir translated from a local directory to a URI). + + Args: + runtime_env: A :ref:`runtime environment ` dictionary. + validate: Whether to validate the runtime env. + """ + self.runtime_env = runtime_env if runtime_env is not None else {} + if validate: + self.runtime_env = self._validate_runtime_env() + self._cached_pb = None + + def set_py_logging_config( + self, + logging_config: Optional[LoggingConfig] = None, + ): + """Set the logging configuration for the job. + + The logging configuration will be applied to the root loggers of + all Ray task and actor processes that belong to this job. + + Args: + logging_config: The logging configuration to set. + """ + self.py_logging_config = logging_config + + def set_ray_namespace(self, ray_namespace: str) -> None: + """Set Ray :ref:`namespace `. + + Args: + ray_namespace: The namespace to set. + """ + + if ray_namespace != self.ray_namespace: + self.ray_namespace = ray_namespace + self._cached_pb = None + + def set_default_actor_lifetime(self, default_actor_lifetime: str) -> None: + """Set the default actor lifetime, which can be "detached" or "non_detached". + + See :ref:`actor lifetimes ` for more details. + + Args: + default_actor_lifetime: The default actor lifetime to set. + """ + import ray.core.generated.common_pb2 as common_pb2 + + if default_actor_lifetime == "detached": + self._default_actor_lifetime = common_pb2.JobConfig.ActorLifetime.DETACHED + elif default_actor_lifetime == "non_detached": + self._default_actor_lifetime = ( + common_pb2.JobConfig.ActorLifetime.NON_DETACHED + ) + else: + raise ValueError( + "Default actor lifetime must be one of `detached`, `non_detached`" + ) + + def _validate_runtime_env(self): + # TODO(edoakes): this is really unfortunate, but JobConfig is imported + # all over the place so this causes circular imports. We should remove + # this dependency and pass in a validated runtime_env instead. + from ray.runtime_env import RuntimeEnv + from ray.runtime_env.runtime_env import _validate_no_local_paths + + runtime_env = self.runtime_env + + if not isinstance(runtime_env, RuntimeEnv): + runtime_env = RuntimeEnv(**self.runtime_env) + _validate_no_local_paths(runtime_env) + return runtime_env + + def _get_proto_job_config(self): + """Return the protobuf structure of JobConfig.""" + # TODO(edoakes): this is really unfortunate, but JobConfig is imported + # all over the place so this causes circular imports. We should remove + # this dependency and pass in a validated runtime_env instead. + import ray.core.generated.common_pb2 as common_pb2 + from ray._private.utils import get_runtime_env_info + + if self._cached_pb is None: + pb = common_pb2.JobConfig() + if self.ray_namespace is None: + pb.ray_namespace = str(uuid.uuid4()) + else: + pb.ray_namespace = self.ray_namespace + pb.jvm_options.extend(self.jvm_options) + pb.code_search_path.extend(self.code_search_path) + pb.py_driver_sys_path.extend(self._py_driver_sys_path) + for k, v in self.metadata.items(): + pb.metadata[k] = v + + parsed_env = self._validate_runtime_env() + pb.runtime_env_info.CopyFrom( + get_runtime_env_info( + parsed_env, + is_job_runtime_env=True, + serialize=False, + ) + ) + + if self._default_actor_lifetime is not None: + pb.default_actor_lifetime = self._default_actor_lifetime + if self.py_logging_config: + pb.serialized_py_logging_config = pickle.dumps(self.py_logging_config) + self._cached_pb = pb + + return self._cached_pb + + def _runtime_env_has_working_dir(self): + return self._validate_runtime_env().has_working_dir() + + def _get_serialized_runtime_env(self) -> str: + """Return the JSON-serialized parsed runtime env dict""" + return self._validate_runtime_env().serialize() + + def _get_proto_runtime_env_config(self) -> str: + """Return the JSON-serialized parsed runtime env info""" + return self._get_proto_job_config().runtime_env_info.runtime_env_config + + @classmethod + def from_json(cls, job_config_json): + """Generates a JobConfig object from json. + + Examples: + .. testcode:: + + from ray.job_config import JobConfig + + job_config = JobConfig.from_json( + {"runtime_env": {"working_dir": "uri://abc"}}) + + Args: + job_config_json: The job config json dictionary. + """ + return cls( + jvm_options=job_config_json.get("jvm_options", None), + code_search_path=job_config_json.get("code_search_path", None), + runtime_env=job_config_json.get("runtime_env", None), + metadata=job_config_json.get("metadata", None), + ray_namespace=job_config_json.get("ray_namespace", None), + _client_job=job_config_json.get("client_job", False), + _py_driver_sys_path=job_config_json.get("py_driver_sys_path", None), + ) diff --git a/lib/python3.12/site-packages/ray/job_submission/__init__.py b/lib/python3.12/site-packages/ray/job_submission/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b3a76be1e53557b0a590b7461097254517d1599d --- /dev/null +++ b/lib/python3.12/site-packages/ray/job_submission/__init__.py @@ -0,0 +1,13 @@ +from ray.dashboard.modules.job.common import JobErrorType, JobInfo, JobStatus +from ray.dashboard.modules.job.pydantic_models import DriverInfo, JobDetails, JobType +from ray.dashboard.modules.job.sdk import JobSubmissionClient + +__all__ = [ + "JobSubmissionClient", + "JobStatus", + "JobErrorType", + "JobInfo", + "JobDetails", + "DriverInfo", + "JobType", +] diff --git a/lib/python3.12/site-packages/ray/job_submission/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/job_submission/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..61c3a83f1ca264dc31fe766cb820f6dae24bf4b8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/job_submission/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/nightly-wheels.yaml b/lib/python3.12/site-packages/ray/nightly-wheels.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1648a1d01400826ac652f1cb6fb7cf9a38c6aa7a --- /dev/null +++ b/lib/python3.12/site-packages/ray/nightly-wheels.yaml @@ -0,0 +1,11 @@ +linux: + "3.8": https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp38-cp38-manylinux2014_x86_64.whl + "3.7": https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl + +darwin: + "3.8": https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp38-cp38-macosx_10_15_x86_64.whl + "3.7": https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-macosx_10_15_intel.whl + +win32: + "3.8": https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp38-cp38-win_amd64.whl + "3.7": https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-3.0.0.dev0-cp37-cp37m-win_amd64.whl diff --git a/lib/python3.12/site-packages/ray/py.typed b/lib/python3.12/site-packages/ray/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/remote_function.py b/lib/python3.12/site-packages/ray/remote_function.py new file mode 100644 index 0000000000000000000000000000000000000000..e1b3f9ae7ebba47bf00dd5c6f71c3b9b91a2bcdf --- /dev/null +++ b/lib/python3.12/site-packages/ray/remote_function.py @@ -0,0 +1,538 @@ +import inspect +import logging +import os +import uuid +from functools import wraps +from threading import Lock +from typing import Optional + +import ray._common.signature +from ray import Language, cross_language +from ray._common import ray_option_utils +from ray._common.ray_option_utils import _warn_if_using_deprecated_placement_group +from ray._common.serialization import pickle_dumps +from ray._private.auto_init_hook import wrap_auto_init +from ray._private.client_mode_hook import ( + client_mode_convert_function, + client_mode_should_convert, +) +from ray._private.utils import get_runtime_env_info, parse_runtime_env_for_task_or_actor +from ray._raylet import ( + STREAMING_GENERATOR_RETURN, + ObjectRefGenerator, + PythonFunctionDescriptor, +) +from ray.util.annotations import DeveloperAPI, PublicAPI +from ray.util.placement_group import _configure_placement_group_based_on_context +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy +from ray.util.tracing.tracing_helper import ( + _inject_tracing_into_function, + _tracing_task_invocation, +) + +logger = logging.getLogger(__name__) + + +# Hook to call with (fn, resources, strategy) on each local task submission. +_task_launch_hook = None + + +@PublicAPI +class RemoteFunction: + """A remote function. + + This is a decorated function. It can be used to spawn tasks. + + Attributes: + _language: The target language. + _function: The original function. + _function_descriptor: The function descriptor. This is not defined + until the remote function is first invoked because that is when the + function is pickled, and the pickled function is used to compute + the function descriptor. + _function_name: The module and function name. + _num_cpus: The default number of CPUs to use for invocations of this + remote function. + _num_gpus: The default number of GPUs to use for invocations of this + remote function. + _memory: The heap memory request in bytes for this task/actor, + rounded down to the nearest integer. + _label_selector: The label requirements on a node for scheduling of the task or actor. + _fallback_strategy: Soft constraints of a list of decorator options to fall back on when scheduling on a node. + _resources: The default custom resource requirements for invocations of + this remote function. + _num_returns: The default number of return values for invocations + of this remote function. + _max_calls: The number of times a worker can execute this function + before exiting. + _max_retries: The number of times this task may be retried + on worker failure. + _retry_exceptions: Whether application-level errors should be retried. + This can be a boolean or a list/tuple of exceptions that should be retried. + _runtime_env: The runtime environment for this task. + _decorator: An optional decorator that should be applied to the remote + function invocation (as opposed to the function execution) before + invoking the function. The decorator must return a function that + takes in two arguments ("args" and "kwargs"). In most cases, it + should call the function that was passed into the decorator and + return the resulting ObjectRefs. For an example, see + "test_decorated_function" in "python/ray/tests/test_basic.py". + _function_signature: The function signature. + _last_export_cluster_and_job: A pair of the last exported cluster + and job to help us to know whether this function was exported. + This is an imperfect mechanism used to determine if we need to + export the remote function again. It is imperfect in the sense that + the actor class definition could be exported multiple times by + different workers. + _scheduling_strategy: Strategy about how to schedule + this remote function. + """ + + def __init__( + self, + language, + function, + function_descriptor, + task_options, + ): + if inspect.iscoroutinefunction(function): + raise ValueError( + "'async def' should not be used for remote tasks. You can wrap the " + "async function with `asyncio.run(f())`. See more at:" + "https://docs.ray.io/en/latest/ray-core/actors/async_api.html " + ) + self._default_options = task_options + + # When gpu is used, set the task non-recyclable by default. + # https://github.com/ray-project/ray/issues/29624 for more context. + # Note: Ray task worker process is not being reused when nsight + # profiler is running, as nsight/rocprof-sys generate report + # once the process exit. + num_gpus = self._default_options.get("num_gpus") or 0 + if ( + num_gpus > 0 and self._default_options.get("max_calls", None) is None + ) or any( + [ + s in (self._default_options.get(s) or {}) + for s in ["nsight", "rocprof-sys"] + ] + ): + self._default_options["max_calls"] = 1 + + # TODO(suquark): This is a workaround for class attributes of options. + # They are being used in some other places, mostly tests. Need cleanup later. + # E.g., actors uses "__ray_metadata__" to collect options, we can so something + # similar for remote functions. + for k, v in ray_option_utils.task_options.items(): + setattr(self, "_" + k, task_options.get(k, v.default_value)) + self._runtime_env = parse_runtime_env_for_task_or_actor(self._runtime_env) + if "runtime_env" in self._default_options: + self._default_options["runtime_env"] = self._runtime_env + + # Pre-calculate runtime env info, to avoid re-calculation at `remote` + # invocation. When `remote` call has specified extra `option` field, + # runtime env will be overwritten and re-serialized. + # + # Caveat: To support dynamic runtime envs in + # `func.option(runtime_env={...}).remote()`, we recalculate the serialized + # runtime env info in the `option` call. But it's acceptable since + # pre-calculation here only happens once at `RemoteFunction` initialization. + self._serialized_base_runtime_env_info = "" + if self._runtime_env: + self._serialized_base_runtime_env_info = get_runtime_env_info( + self._runtime_env, + is_job_runtime_env=False, + serialize=True, + ) + + self._language = language + self._is_generator = inspect.isgeneratorfunction(function) + self._function = function + self._function_signature = None + # Guards trace injection to enforce exactly once semantics + self._inject_lock = Lock() + self._function_name = function.__module__ + "." + function.__name__ + self._function_descriptor = function_descriptor + self._is_cross_language = language != Language.PYTHON + self._decorator = getattr(function, "__ray_invocation_decorator__", None) + self._last_export_cluster_and_job = None + self._uuid = uuid.uuid4() + + # Override task.remote's signature and docstring + @wraps(function) + def _remote_proxy(*args, **kwargs): + return self._remote( + serialized_runtime_env_info=self._serialized_base_runtime_env_info, + args=args, + kwargs=kwargs, + **self._default_options, + ) + + self.remote = _remote_proxy + + def __call__(self, *args, **kwargs): + raise TypeError( + "Remote functions cannot be called directly. Instead " + f"of running '{self._function_name}()', " + f"try '{self._function_name}.remote()'." + ) + + # Lock is not picklable + def __getstate__(self): + attrs = self.__dict__.copy() + del attrs["_inject_lock"] + return attrs + + def __setstate__(self, state): + self.__dict__.update(state) + self.__dict__["_inject_lock"] = Lock() + + def options(self, **task_options): + """Configures and overrides the task invocation parameters. + + The arguments are the same as those that can be passed to :obj:`ray.remote`. + Overriding `max_calls` is not supported. + + Args: + num_returns: It specifies the number of object refs returned by + the remote function invocation. + num_cpus: The quantity of CPU cores to reserve + for this task or for the lifetime of the actor. + num_gpus: The quantity of GPUs to reserve + for this task or for the lifetime of the actor. + resources (Dict[str, float]): The quantity of various custom resources + to reserve for this task or for the lifetime of the actor. + This is a dictionary mapping strings (resource names) to floats. + label_selector (Dict[str, str]): If specified, the labels required for the node on + which this actor can be scheduled on. The label selector consist of key-value pairs, + where the keys are label names and the value are expressions consisting of an operator + with label values or just a value to indicate equality. + fallback_strategy (List[Dict[str, Any]]): If specified, expresses soft constraints + through a list of decorator options to fall back on when scheduling on a node. + accelerator_type: If specified, requires that the task or actor run + on a node with the specified type of accelerator. + See :ref:`accelerator types `. + memory: The heap memory request in bytes for this task/actor, + rounded down to the nearest integer. + object_store_memory: The object store memory request for actors only. + max_calls: This specifies the + maximum number of times that a given worker can execute + the given remote function before it must exit + (this can be used to address memory leaks in third-party + libraries or to reclaim resources that cannot easily be + released, e.g., GPU memory that was acquired by TensorFlow). + By default this is infinite for CPU tasks and 1 for GPU tasks + (to force GPU tasks to release resources after finishing). + max_retries: This specifies the maximum number of times that the remote + function should be rerun when the worker process executing it + crashes unexpectedly. The minimum valid value is 0, + the default is 3 (default), and a value of -1 indicates + infinite retries. + runtime_env (Dict[str, Any]): Specifies the runtime environment for + this actor or task and its children. See + :ref:`runtime-environments` for detailed documentation. + retry_exceptions: This specifies whether application-level errors + should be retried up to max_retries times. + scheduling_strategy: Strategy about how to + schedule a remote function or actor. Possible values are + None: ray will figure out the scheduling strategy to use, it + will either be the PlacementGroupSchedulingStrategy using parent's + placement group if parent has one and has + placement_group_capture_child_tasks set to true, + or "DEFAULT"; + "DEFAULT": default hybrid scheduling; + "SPREAD": best effort spread scheduling; + `PlacementGroupSchedulingStrategy`: + placement group based scheduling; + `NodeAffinitySchedulingStrategy`: + node id based affinity scheduling. + enable_task_events: This specifies whether to enable task events for this + task. If set to True, task events such as (task running, finished) + are emitted, and available to Ray Dashboard and State API. + See :ref:`state-api-overview-ref` for more details. + _labels: The key-value labels of a task. + + Examples: + + .. code-block:: python + + @ray.remote(num_gpus=1, max_calls=1, num_returns=2) + def f(): + return 1, 2 + # Task g will require 2 gpus instead of 1. + g = f.options(num_gpus=2) + """ + + func_cls = self + + # override original options + default_options = self._default_options.copy() + # max_calls could not be used in ".options()", we should remove it before + # merging options from '@ray.remote'. + default_options.pop("max_calls", None) + updated_options = ray_option_utils.update_options(default_options, task_options) + ray_option_utils.validate_task_options(updated_options, in_options=True) + + # Only update runtime_env and re-calculate serialized runtime env info when + # ".options()" specifies new runtime_env. + serialized_runtime_env_info = self._serialized_base_runtime_env_info + if "runtime_env" in task_options: + updated_options["runtime_env"] = parse_runtime_env_for_task_or_actor( + updated_options["runtime_env"] + ) + # Re-calculate runtime env info based on updated runtime env. + if updated_options["runtime_env"]: + serialized_runtime_env_info = get_runtime_env_info( + updated_options["runtime_env"], + is_job_runtime_env=False, + serialize=True, + ) + + class FuncWrapper: + def remote(self, *args, **kwargs): + return func_cls._remote( + args=args, + kwargs=kwargs, + serialized_runtime_env_info=serialized_runtime_env_info, + **updated_options, + ) + + @DeveloperAPI + def bind(self, *args, **kwargs): + """ + For Ray DAG building that creates static graph from decorated + class or functions. + """ + from ray.dag.function_node import FunctionNode + + return FunctionNode(func_cls._function, args, kwargs, updated_options) + + return FuncWrapper() + + @wrap_auto_init + @_tracing_task_invocation + def _remote( + self, + args=None, + kwargs=None, + serialized_runtime_env_info: Optional[str] = None, + **task_options, + ): + """Submit the remote function for execution.""" + # We pop the "max_calls" coming from "@ray.remote" here. We no longer need + # it in "_remote()". + task_options.pop("max_calls", None) + if client_mode_should_convert(): + return client_mode_convert_function(self, args, kwargs, **task_options) + + worker = ray._private.worker.global_worker + worker.check_connected() + + if worker.mode != ray._private.worker.WORKER_MODE: + # Only need to record on the driver side + # since workers are created via tasks or actors + # launched from the driver. + from ray._common.usage import usage_lib + + usage_lib.record_library_usage("core") + + # We cannot do this when the function is first defined, because we need + # ray.init() to have been called when this executes + with self._inject_lock: + if self._function_signature is None: + self._function = _inject_tracing_into_function(self._function) + self._function_signature = ray._common.signature.extract_signature( + self._function + ) + + # If this function was not exported in this cluster and job, we need to + # export this function again, because the current GCS doesn't have it. + if ( + not self._is_cross_language + and self._last_export_cluster_and_job != worker.current_cluster_and_job + ): + self._function_descriptor = PythonFunctionDescriptor.from_function( + self._function, self._uuid + ) + # There is an interesting question here. If the remote function is + # used by a subsequent driver (in the same script), should the + # second driver pickle the function again? If yes, then the remote + # function definition can differ in the second driver (e.g., if + # variables in its closure have changed). We probably want the + # behavior of the remote function in the second driver to be + # independent of whether or not the function was invoked by the + # first driver. This is an argument for repickling the function, + # which we do here. + self._pickled_function = pickle_dumps( + self._function, + f"Could not serialize the function {self._function_descriptor.repr}", + ) + + self._last_export_cluster_and_job = worker.current_cluster_and_job + worker.function_actor_manager.export(self) + + kwargs = {} if kwargs is None else kwargs + args = [] if args is None else args + + # fill task required options + for k, v in ray_option_utils.task_options.items(): + if k == "max_retries": + # TODO(swang): We need to override max_retries here because the default + # value gets set at Ray import time. Ideally, we should allow setting + # default values from env vars for other options too. + v.default_value = os.environ.get( + "RAY_TASK_MAX_RETRIES", v.default_value + ) + v.default_value = int(v.default_value) + task_options[k] = task_options.get(k, v.default_value) + # "max_calls" already takes effects and should not apply again. + # Remove the default value here. + task_options.pop("max_calls", None) + + # TODO(suquark): cleanup these fields + name = task_options["name"] + placement_group = task_options["placement_group"] + placement_group_bundle_index = task_options["placement_group_bundle_index"] + placement_group_capture_child_tasks = task_options[ + "placement_group_capture_child_tasks" + ] + scheduling_strategy = task_options["scheduling_strategy"] + + num_returns = task_options["num_returns"] + if num_returns is None: + if self._is_generator: + num_returns = "streaming" + else: + num_returns = 1 + + if num_returns == "dynamic": + num_returns = -1 + elif num_returns == "streaming": + # TODO(sang): This is a temporary private API. + # Remove it when we migrate to the streaming generator. + num_returns = ray._raylet.STREAMING_GENERATOR_RETURN + generator_backpressure_num_objects = task_options[ + "_generator_backpressure_num_objects" + ] + if generator_backpressure_num_objects is None: + generator_backpressure_num_objects = -1 + + max_retries = task_options["max_retries"] + retry_exceptions = task_options["retry_exceptions"] + if isinstance(retry_exceptions, (list, tuple)): + retry_exception_allowlist = tuple(retry_exceptions) + retry_exceptions = True + else: + retry_exception_allowlist = None + + if scheduling_strategy is None or not isinstance( + scheduling_strategy, PlacementGroupSchedulingStrategy + ): + _warn_if_using_deprecated_placement_group(task_options, 4) + + resources = ray._common.utils.resources_from_ray_options(task_options) + + if scheduling_strategy is None or isinstance( + scheduling_strategy, PlacementGroupSchedulingStrategy + ): + if isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy): + placement_group = scheduling_strategy.placement_group + placement_group_bundle_index = ( + scheduling_strategy.placement_group_bundle_index + ) + placement_group_capture_child_tasks = ( + scheduling_strategy.placement_group_capture_child_tasks + ) + + if placement_group_capture_child_tasks is None: + placement_group_capture_child_tasks = ( + worker.should_capture_child_tasks_in_placement_group + ) + placement_group = _configure_placement_group_based_on_context( + placement_group_capture_child_tasks, + placement_group_bundle_index, + resources, + {}, # no placement_resources for tasks + self._function_descriptor.function_name, + placement_group=placement_group, + ) + if not placement_group.is_empty: + scheduling_strategy = PlacementGroupSchedulingStrategy( + placement_group, + placement_group_bundle_index, + placement_group_capture_child_tasks, + ) + else: + scheduling_strategy = "DEFAULT" + + if _task_launch_hook: + _task_launch_hook(self._function_descriptor, resources, scheduling_strategy) + + # Override enable_task_events to default for actor if not specified (i.e. None) + enable_task_events = task_options.get("enable_task_events") + labels = task_options.get("_labels") + label_selector = task_options.get("label_selector") + fallback_strategy = task_options.get("fallback_strategy") + + def invocation(args, kwargs): + if self._is_cross_language: + list_args = cross_language._format_args(worker, args, kwargs) + elif not args and not kwargs and not self._function_signature: + list_args = [] + else: + list_args = ray._common.signature.flatten_args( + self._function_signature, args, kwargs + ) + + if worker.mode == ray._private.worker.LOCAL_MODE: + assert ( + not self._is_cross_language + ), "Cross language remote function cannot be executed locally." + object_refs = worker.core_worker.submit_task( + self._language, + self._function_descriptor, + list_args, + name if name is not None else "", + num_returns, + resources, + max_retries, + retry_exceptions, + retry_exception_allowlist, + scheduling_strategy, + worker.debugger_breakpoint, + serialized_runtime_env_info or "{}", + generator_backpressure_num_objects, + enable_task_events, + labels, + label_selector, + fallback_strategy, + ) + # Reset worker's debug context from the last "remote" command + # (which applies only to this .remote call). + worker.debugger_breakpoint = b"" + if num_returns == STREAMING_GENERATOR_RETURN: + # Streaming generator will return a single ref + # that is for the generator task. + assert len(object_refs) == 1 + generator_ref = object_refs[0] + return ObjectRefGenerator(generator_ref, worker) + if len(object_refs) == 1: + return object_refs[0] + elif len(object_refs) > 1: + return object_refs + + if self._decorator is not None: + invocation = self._decorator(invocation) + + return invocation(args, kwargs) + + @DeveloperAPI + def bind(self, *args, **kwargs): + """ + For Ray DAG building that creates static graph from decorated + class or functions. + """ + + from ray.dag.function_node import FunctionNode + + return FunctionNode(self._function, args, kwargs, self._default_options) diff --git a/lib/python3.12/site-packages/ray/runtime_context.py b/lib/python3.12/site-packages/ray/runtime_context.py new file mode 100644 index 0000000000000000000000000000000000000000..5abbc2471a143533a62e7dd0f1d0b60254734399 --- /dev/null +++ b/lib/python3.12/site-packages/ray/runtime_context.py @@ -0,0 +1,584 @@ +import logging +import threading +from typing import Any, Dict, List, Optional + +import ray._private.worker +from ray._private.client_mode_hook import client_mode_hook +from ray._private.state import actors +from ray._private.utils import parse_pg_formatted_resources_to_original +from ray._raylet import TaskID +from ray.runtime_env import RuntimeEnv +from ray.util.annotations import Deprecated, PublicAPI + +logger = logging.getLogger(__name__) + + +@PublicAPI +class RuntimeContext(object): + """A class used for getting runtime context.""" + + def __init__(self, worker): + assert worker is not None + self.worker = worker + + @Deprecated( + message="Use get_xxx_id() methods to get relevant ids instead", warning=True + ) + def get(self) -> Dict[str, Any]: + """Get a dictionary of the current context. + + Returns: + dict: Dictionary of the current context. + """ + context = { + "job_id": self.job_id, + "node_id": self.node_id, + "namespace": self.namespace, + } + if self.worker.mode == ray._private.worker.WORKER_MODE: + if self.task_id is not None: + context["task_id"] = self.task_id + if self.actor_id is not None: + context["actor_id"] = self.actor_id + + return context + + @property + @Deprecated(message="Use get_job_id() instead", warning=True) + def job_id(self): + """Get current job ID for this worker or driver. + + Job ID is the id of your Ray drivers that create tasks or actors. + + Returns: + If called by a driver, this returns the job ID. If called in + a task, return the job ID of the associated driver. + + """ + job_id = self.worker.current_job_id + assert not job_id.is_nil() + return job_id + + def get_job_id(self) -> str: + """Get current job ID for this worker or driver. + + Job ID is the id of your Ray drivers that create tasks or actors. + + Returns: + If called by a driver, this returns the job ID. If called in + a task, return the job ID of the associated driver. The + job ID will be hex format. + + Raises: + AssertionError: If not called in a driver or worker. Generally, + this means that ray.init() was not called. + """ + assert ( + ray.is_initialized() + ), "Job ID is not available because Ray has not been initialized." + job_id = self.worker.current_job_id + return job_id.hex() + + @property + @Deprecated(message="Use get_node_id() instead", warning=True) + def node_id(self): + """Get the ID for the node that this process is running on. + + This can be called from within a driver, task, or actor. + When called from a driver that is connected to a remote Ray cluster using + Ray Client, this returns the ID of the head node. + + Returns: + A node id for this worker or driver. + """ + node_id = self.worker.current_node_id + assert not node_id.is_nil() + return node_id + + def get_node_id(self) -> str: + """Get the ID for the node that this process is running on. + + This can be called from within a driver, task, or actor. + When called from a driver that is connected to a remote Ray cluster using + Ray Client, this returns the ID of the head node. + + Returns: + A node id in hex format for this worker or driver. + + Raises: + AssertionError: If not called in a driver or worker. Generally, + this means that ray.init() was not called. + """ + assert ( + ray.is_initialized() + ), "Node ID is not available because Ray has not been initialized." + node_id = self.worker.current_node_id + return node_id.hex() + + def get_worker_id(self) -> str: + """Get current worker ID for this worker or driver process. + + Returns: + A worker id in hex format for this worker or driver process. + """ + assert ( + ray.is_initialized() + ), "Worker ID is not available because Ray has not been initialized." + return self.worker.worker_id.hex() + + @property + @Deprecated(message="Use get_task_id() instead", warning=True) + def task_id(self): + """Get current task ID for this worker. + + Task ID is the id of a Ray task. + This shouldn't be used in a driver process. + + Example: + + .. testcode:: + + import ray + + @ray.remote + class Actor: + def ready(self): + return True + + @ray.remote + def f(): + return True + + # All the below code generates different task ids. + # Task ids are available for actor creation. + a = Actor.remote() + # Task ids are available for actor tasks. + a.ready.remote() + # Task ids are available for normal tasks. + f.remote() + + Returns: + The current worker's task id. None if there's no task id. + """ + # only worker mode has task_id + assert ( + self.worker.mode == ray._private.worker.WORKER_MODE + ), f"This method is only available when the process is a\ + worker. Current mode: {self.worker.mode}" + + task_id = self._get_current_task_id() + return task_id if not task_id.is_nil() else None + + def get_task_id(self) -> Optional[str]: + """Get current task ID for this worker. + + Task ID is the id of a Ray task. The ID will be in hex format. + This shouldn't be used in a driver process. + + Example: + + .. testcode:: + + import ray + + @ray.remote + class Actor: + def get_task_id(self): + return ray.get_runtime_context().get_task_id() + + @ray.remote + def get_task_id(): + return ray.get_runtime_context().get_task_id() + + # All the below code generates different task ids. + a = Actor.remote() + # Task ids are available for actor tasks. + print(ray.get(a.get_task_id.remote())) + # Task ids are available for normal tasks. + print(ray.get(get_task_id.remote())) + + .. testoutput:: + :options: +MOCK + + 16310a0f0a45af5c2746a0e6efb235c0962896a201000000 + c2668a65bda616c1ffffffffffffffffffffffff01000000 + + Returns: + The current worker's task id in hex. None if there's no task id. + """ + # only worker mode has task_id + if self.worker.mode != ray._private.worker.WORKER_MODE: + logger.warning( + "This method is only available when the process is a " + f"worker. Current mode: {self.worker.mode}" + ) + return None + task_id = self._get_current_task_id() + return task_id.hex() if not task_id.is_nil() else None + + def _get_current_task_id(self) -> TaskID: + return self.worker.current_task_id + + def get_task_name(self) -> Optional[str]: + """Get current task name for this worker. + + Task name by default is the task's funciton call string. It can also be + specified in options when triggering a task. + + Example: + + .. testcode:: + + import ray + + @ray.remote + class Actor: + def get_task_name(self): + return ray.get_runtime_context().get_task_name() + + @ray.remote + class AsyncActor: + async def get_task_name(self): + return ray.get_runtime_context().get_task_name() + + @ray.remote + def get_task_name(): + return ray.get_runtime_context().get_task_name() + + a = Actor.remote() + b = AsyncActor.remote() + # Task names are available for actor tasks. + print(ray.get(a.get_task_name.remote())) + # Task names are avaiable for async actor tasks. + print(ray.get(b.get_task_name.remote())) + # Task names are available for normal tasks. + # Get default task name + print(ray.get(get_task_name.remote())) + # Get specified task name + print(ray.get(get_task_name.options(name="task_name").remote())) + + .. testoutput:: + :options: +MOCK + + Actor.get_task_name + AsyncActor.get_task_name + get_task_name + task_nams + + Returns: + The current worker's task name + """ + # only worker mode has task_name + if self.worker.mode != ray._private.worker.WORKER_MODE: + logger.warning( + "This method is only available when the process is a " + f"worker. Current mode: {self.worker.mode}" + ) + return None + return self.worker.current_task_name + + def get_task_function_name(self) -> Optional[str]: + """Get current task function name string for this worker. + + Example: + + .. testcode:: + + import ray + + @ray.remote + class Actor: + def get_task_function_name(self): + return ray.get_runtime_context().get_task_function_name() + + @ray.remote + class AsyncActor: + async def get_task_function_name(self): + return ray.get_runtime_context().get_task_function_name() + + @ray.remote + def get_task_function_name(): + return ray.get_runtime_context().get_task_function_name() + + a = Actor.remote() + b = AsyncActor.remote() + # Task functions are available for actor tasks. + print(ray.get(a.get_task_function_name.remote())) + # Task functions are available for async actor tasks. + print(ray.get(b.get_task_function_name.remote())) + # Task functions are available for normal tasks. + print(ray.get(get_task_function_name.remote())) + + .. testoutput:: + :options: +MOCK + + [python modual name].Actor.get_task_function_name + [python modual name].AsyncActor.get_task_function_name + [python modual name].get_task_function_name + + Returns: + The current worker's task function call string + """ + # only worker mode has task_function_name + if self.worker.mode != ray._private.worker.WORKER_MODE: + logger.warning( + "This method is only available when the process is a " + f"worker. Current mode: {self.worker.mode}" + ) + return None + return self.worker.current_task_function_name + + @property + @Deprecated(message="Use get_actor_id() instead", warning=True) + def actor_id(self): + """Get the current actor ID in this worker. + + ID of the actor of the current process. + This shouldn't be used in a driver process. + + Returns: + The current actor id in this worker. None if there's no actor id. + """ + # only worker mode has actor_id + assert ( + self.worker.mode == ray._private.worker.WORKER_MODE + ), f"This method is only available when the process is a\ + worker. Current mode: {self.worker.mode}" + actor_id = self.worker.actor_id + return actor_id if not actor_id.is_nil() else None + + def get_actor_id(self) -> Optional[str]: + """Get the current actor ID in this worker. + + ID of the actor of the current process. + This shouldn't be used in a driver process. + The ID will be in hex format. + + Returns: + The current actor id in hex format in this worker. None if there's no + actor id. + """ + # only worker mode has actor_id + if self.worker.mode != ray._private.worker.WORKER_MODE: + logger.debug( + "This method is only available when the process is a " + f"worker. Current mode: {self.worker.mode}" + ) + return None + actor_id = self.worker.actor_id + return actor_id.hex() if not actor_id.is_nil() else None + + def get_actor_name(self) -> Optional[str]: + """Get the current actor name of this worker. + + This shouldn't be used in a driver process. + The name is in string format. + + Returns: + The current actor name of this worker. + If a current worker is an actor, and + if actor name doesn't exist, it returns an empty string. + If a current worker is not an actor, it returns None. + """ + # only worker mode has actor_id + if self.worker.mode != ray._private.worker.WORKER_MODE: + logger.warning( + "This method is only available when the process is a " + f"worker. Current mode: {self.worker.mode}" + ) + return None + actor_id = self.worker.actor_id + return self.worker.actor_name if not actor_id.is_nil() else None + + @property + def namespace(self): + """Get the current namespace of this worker. + + Returns: + The current namespace of this worker. + """ + return self.worker.namespace + + @property + def was_current_actor_reconstructed(self): + """Check whether this actor has been restarted. + + Returns: + Whether this actor has been ever restarted. + """ + assert ( + not self.actor_id.is_nil() + ), "This method should't be called inside Ray tasks." + actor_info = actors(actor_id=self.actor_id.hex()) + return actor_info and actor_info["NumRestarts"] != 0 + + @property + @Deprecated(message="Use get_placement_group_id() instead", warning=True) + def current_placement_group_id(self): + """Get the current Placement group ID of this worker. + + Returns: + The current placement group id of this worker. + """ + return self.worker.placement_group_id + + def get_placement_group_id(self) -> Optional[str]: + """Get the current Placement group ID of this worker. + + Returns: + The current placement group id in hex format of this worker. + """ + pg_id = self.worker.placement_group_id + return pg_id.hex() if not pg_id.is_nil() else None + + @property + def should_capture_child_tasks_in_placement_group(self): + """Get if the current task should capture parent's placement group. + + This returns True if it is called inside a driver. + + Returns: + Return True if the current task should implicitly + capture the parent placement group. + """ + return self.worker.should_capture_child_tasks_in_placement_group + + def get_assigned_resources(self): + """Get the assigned resources to this worker. + + By default for tasks, this will return {"CPU": 1}. + By default for actors, this will return {}. This is because + actors do not have CPUs assigned to them by default. + + Returns: + A dictionary mapping the name of a resource to a float, where + the float represents the amount of that resource reserved + for this worker. + """ + assert ( + self.worker.mode == ray._private.worker.WORKER_MODE + ), f"This method is only available when the process is a\ + worker. Current mode: {self.worker.mode}" + self.worker.check_connected() + resource_id_map = self.worker.core_worker.resource_ids() + resource_map = { + res: sum(amt for _, amt in mapping) + for res, mapping in resource_id_map.items() + } + result = parse_pg_formatted_resources_to_original(resource_map) + return result + + def get_runtime_env_string(self): + """Get the runtime env string used for the current driver or worker. + + Returns: + The runtime env string currently using by this worker. + """ + return self.worker.runtime_env + + @property + def runtime_env(self): + """Get the runtime env used for the current driver or worker. + + Returns: + The runtime env currently using by this worker. The type of + return value is ray.runtime_env.RuntimeEnv. + """ + + return RuntimeEnv.deserialize(self.get_runtime_env_string()) + + @property + def current_actor(self): + """Get the current actor handle of this actor itself. + + Returns: + The handle of current actor. + """ + worker = self.worker + worker.check_connected() + actor_id = worker.actor_id + if actor_id.is_nil(): + raise RuntimeError("This method is only available in an actor.") + + return worker.core_worker.get_actor_handle(actor_id) + + @property + def gcs_address(self): + """Get the GCS address of the ray cluster. + + Returns: + The GCS address of the cluster. + """ + self.worker.check_connected() + return self.worker.gcs_client.address + + @Deprecated(message="Use get_accelerator_ids() instead", warning=True) + def get_resource_ids(self) -> Dict[str, List[str]]: + return self.get_accelerator_ids() + + def get_accelerator_ids(self) -> Dict[str, List[str]]: + """ + Get the current worker's visible accelerator ids. + + Returns: + A dictionary keyed by the accelerator resource name. The values are a list + of ids `{'GPU': ['0', '1'], 'neuron_cores': ['0', '1'], + 'TPU': ['0', '1']}`. + """ + worker = self.worker + worker.check_connected() + ids_dict: Dict[str, List[str]] = {} + for ( + accelerator_resource_name + ) in ray._private.accelerators.get_all_accelerator_resource_names(): + accelerator_ids = worker.get_accelerator_ids_for_accelerator_resource( + accelerator_resource_name, + f"^{accelerator_resource_name}_group_[0-9A-Za-z]+$", + ) + ids_dict[accelerator_resource_name] = [str(id) for id in accelerator_ids] + return ids_dict + + def get_node_labels(self) -> Dict[str, List[str]]: + """ + Get the node labels of the current worker. + + Returns: + A dictionary of label key-value pairs. + """ + worker = self.worker + worker.check_connected() + + return worker.current_node_labels + + +_runtime_context = None +_runtime_context_lock = threading.Lock() + + +@PublicAPI +@client_mode_hook +def get_runtime_context() -> RuntimeContext: + """Get the runtime context of the current driver/worker. + + The obtained runtime context can be used to get the metadata + of the current driver, task, or actor. + + Example: + + .. testcode:: + + import ray + # Get the job id. + ray.get_runtime_context().get_job_id() + # Get the actor id. + ray.get_runtime_context().get_actor_id() + # Get the task id. + ray.get_runtime_context().get_task_id() + + """ + with _runtime_context_lock: + global _runtime_context + if _runtime_context is None: + _runtime_context = RuntimeContext(ray._private.worker.global_worker) + + return _runtime_context diff --git a/lib/python3.12/site-packages/ray/setup-dev.py b/lib/python3.12/site-packages/ray/setup-dev.py new file mode 100644 index 0000000000000000000000000000000000000000..8f65e4e716ff31d39e63605f68557dd6c8f5f3ea --- /dev/null +++ b/lib/python3.12/site-packages/ray/setup-dev.py @@ -0,0 +1,196 @@ +#!/usr/bin/env python +# ruff: noqa: E402 +"""This script allows you to develop Ray Python code without needing to compile +Ray. +See https://docs.ray.io/en/master/development.html#building-ray-python-only""" + +import os +import sys + +# types.py can conflict with stdlib's types.py in some python versions, +# see https://github.com/python/cpython/issues/101210. +# To avoid import errors, we move the current working dir to the end of sys.path. +this_dir = os.path.dirname(__file__) +if this_dir in sys.path: + sys.path.remove(this_dir) + sys.path.append(this_dir) + +import argparse +import shutil +import subprocess + +import click + +import ray + + +def do_link(package, force=False, skip_list=None, allow_list=None, local_path=None): + if skip_list and package in skip_list: + print(f"Skip creating symbolic link for {package}") + return + if allow_list is not None and package not in allow_list: + print(f"Skip creating symbolic link for {package} (not in allow list)") + return + package_home = os.path.abspath(os.path.join(ray.__file__, f"../{package}")) + # Infer local_path automatically. + if local_path is None: + local_path = f"../{package}" + local_home = os.path.abspath(os.path.join(__file__, local_path)) + # If installed package dir does not exist, continue either way. We'll + # remove it/create a link from there anyways. + if not os.path.isdir(package_home) and not os.path.isfile(package_home): + print(f"{package_home} does not exist. Continuing to link.") + # Make sure the path we are linking to does exist. + assert os.path.exists(local_home), local_home + # Confirm with user. + if not force and not click.confirm( + f"This will replace:\n {package_home}\nwith " f"a symlink to:\n {local_home}", + default=True, + ): + return + + # Windows: Create directory junction. + if os.name == "nt": + try: + shutil.rmtree(package_home) + except FileNotFoundError: + pass + except OSError: + os.remove(package_home) + + # create symlink for directory or file + if os.path.isdir(local_home): + subprocess.check_call( + ["mklink", "/J", package_home, local_home], shell=True + ) + elif os.path.isfile(local_home): + subprocess.check_call( + ["mklink", "/H", package_home, local_home], shell=True + ) + else: + print(f"{local_home} is neither directory nor file. Link failed.") + + # Posix: Use `ln -s` to create softlink. + else: + sudo = [] + if not os.access(os.path.dirname(package_home), os.W_OK): + print("You don't have write permission " f"to {package_home}, using sudo:") + sudo = ["sudo"] + print(f"Creating symbolic link from \n {local_home} to \n {package_home}") + + # Preserve ray/serve/generated + serve_temp_dir = "/tmp/ray/_serve/" + if package == "serve": + # Copy generated folder to a temp dir + generated_folder = os.path.join(package_home, "generated") + if not os.path.exists(serve_temp_dir): + os.makedirs(serve_temp_dir) + subprocess.check_call(["mv", generated_folder, serve_temp_dir]) + + # Create backup of the old directory if it exists + if os.path.exists(package_home): + backup_dir = f"{package_home}.bak" + print(f"Creating backup of {package_home} to {backup_dir}") + subprocess.check_call(sudo + ["cp", "-r", package_home, backup_dir]) + + subprocess.check_call(sudo + ["rm", "-rf", package_home]) + subprocess.check_call(sudo + ["ln", "-s", local_home, package_home]) + + # Move generated folder to local_home + if package == "serve": + tmp_generated_folder = os.path.join(serve_temp_dir, "generated") + package_generated_folder = os.path.join(package_home, "generated") + if not os.path.exists(package_generated_folder): + subprocess.check_call( + ["mv", tmp_generated_folder, package_generated_folder] + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + formatter_class=argparse.RawDescriptionHelpFormatter, description="Setup dev." + ) + parser.add_argument( + "--yes", "-y", action="store_true", help="Don't ask for confirmation." + ) + parser.add_argument( + "--skip", + "-s", + nargs="*", + help="List of folders to skip linking to facilitate workspace dev", + required=False, + ) + parser.add_argument( + "--allow", + "-a", + nargs="*", + help="List of folders to link (only these will be linked)", + required=False, + ) + parser.add_argument( + "--extras", + "-e", + nargs="*", + help="List of extra folders to link to facilitate workspace dev", + required=False, + ) + + args = parser.parse_args() + if args.skip and args.allow: + print("Error: --skip and --allow cannot be used together.") + sys.exit(1) + + if not args.yes: + print("NOTE: Use '-y' to override all python files without confirmation.") + + # Dictionary of packages to link, with optional local_path + packages_to_link = { + "llm": None, + "serve/llm": None, + "data/llm.py": None, + "rllib": "../../../rllib", + "air": None, + "tune": None, + "train": None, + "autoscaler": None, + "cloudpickle": None, + "data": None, + "scripts": None, + "internal": None, + "tests": None, + "experimental": None, + "util": None, + "workflow": None, + "serve": None, + "dag": None, + "widgets": None, + "cluster_utils.py": None, + "_private": None, + "_common": None, + "dashboard": None, + } + + # Link all packages using a for loop + for package, local_path in packages_to_link.items(): + do_link( + package, + force=args.yes, + skip_list=args.skip, + allow_list=args.allow, + local_path=local_path, + ) + + if args.extras is not None: + for package in args.extras: + do_link(package, force=args.yes, skip_list=args.skip, allow_list=args.allow) + + print( + "Created links.\n\nIf you run into issues initializing Ray, please " + "ensure that your local repo and the installed Ray are in sync " + "(pip install -U the latest wheels at " + "https://docs.ray.io/en/master/installation.html, " + "and ensure you are up-to-date on the master branch on git).\n\n" + "Note that you may need to delete the package symlinks when pip " + "installing new Ray versions to prevent pip from overwriting files " + "in your git repo." + ) diff --git a/lib/python3.12/site-packages/ray/types.py b/lib/python3.12/site-packages/ray/types.py new file mode 100644 index 0000000000000000000000000000000000000000..a4733f44e98eb8660449dd46bf0c45f81856d8a6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/types.py @@ -0,0 +1,14 @@ +from typing import Generic, TypeVar + +from ray.util.annotations import PublicAPI + +T = TypeVar("T") + + +# TODO(ekl) this is a dummy generic ref type for documentation purposes only. +# We should try to make the Cython ray.ObjectRef properly generic. +# NOTE(sang): Looks like using Generic in Cython is not currently possible. +# We should update Cython > 3.0 for this. +@PublicAPI +class ObjectRef(Generic[T]): + pass diff --git a/lib/python3.12/site-packages/ray/util/__init__.py b/lib/python3.12/site-packages/ray/util/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..19d58a0dd318acae01568b74d0ad8c171fdf9003 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/__init__.py @@ -0,0 +1,75 @@ +from typing import List + +import ray +from ray._private.auto_init_hook import wrap_auto_init +from ray._private.client_mode_hook import client_mode_hook +from ray._private.services import get_node_instance_id, get_node_ip_address +from ray.util import accelerators, debugpy as ray_debugpy, iter, rpdb as pdb +from ray.util.actor_pool import ActorPool +from ray.util.annotations import PublicAPI +from ray.util.check_serialize import inspect_serializability +from ray.util.client_connect import connect, disconnect +from ray.util.debug import disable_log_once_globally, enable_periodic_logging, log_once +from ray.util.helpers import as_completed, map_unordered +from ray.util.placement_group import ( + get_current_placement_group, + get_placement_group, + placement_group, + placement_group_table, + remove_placement_group, +) +from ray.util.serialization import deregister_serializer, register_serializer + + +@PublicAPI(stability="beta") +@wrap_auto_init +@client_mode_hook +def list_named_actors(all_namespaces: bool = False) -> List[str]: + """List all named actors in the system. + + Actors must have been created with Actor.options(name="name").remote(). + This works for both detached & non-detached actors. + + By default, only actors in the current namespace will be returned + and the returned entries will simply be their name. + + If `all_namespaces` is set to True, all actors in the cluster will be + returned regardless of namespace, and the returned entries will be of the + form {"namespace": namespace, "name": name}. + """ + worker = ray._private.worker.global_worker + worker.check_connected() + + actors = worker.core_worker.list_named_actors(all_namespaces) + if all_namespaces: + return [{"name": name, "namespace": namespace} for namespace, name in actors] + else: + return [name for _, name in actors] + + +__all__ = [ + "accelerators", + "ActorPool", + "as_completed", + "disable_log_once_globally", + "enable_periodic_logging", + "iter", + "log_once", + "pdb", + "placement_group", + "placement_group_table", + "get_placement_group", + "get_current_placement_group", + "get_node_instance_id", + "get_node_ip_address", + "map_unordered", + "remove_placement_group", + "ray_debugpy", + "inspect_serializability", + "collective", + "connect", + "disconnect", + "register_serializer", + "deregister_serializer", + "list_named_actors", +] diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e89529b262ae9cebcecb7e9e26d41c061216c3e5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/actor_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/actor_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..45b64810a9cf1a443d8e634c322f32b49d5222da Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/actor_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/actor_pool.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/actor_pool.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c3324bae54e40f3f011f7bd360e41c72fe6c754 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/actor_pool.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/annotations.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/annotations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..738b46914f4f4481e7c8f2da866554cfbaacef1d Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/annotations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/check_open_ports.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/check_open_ports.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce7adbd71d6ca57c231bda06b4ab4bba91f06261 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/check_open_ports.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/check_serialize.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/check_serialize.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fe86837d20c4ee91867593b670a421e344047b6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/check_serialize.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/client_connect.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/client_connect.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3019e71e833bf56c198c892e8f3cfe6bbed2ed9b Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/client_connect.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..533fd3467f631c199c8538f6587d786b807e271c Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/debug.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/debug.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80500a74ca64a66ae18daeb25e167af2a3e9e3d1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/debug.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/debugpy.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/debugpy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f03e06c7910692f4dd377b1c1e2bcd05a9904aa0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/debugpy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/helpers.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bfe23217013939446c6390fd1d57fac01e424dbf Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/iter.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/iter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56b8def2613c503e5baa66cd9c573d821be68ea8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/iter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/iter_metrics.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/iter_metrics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..608b0c655bc5dc9305366331b3a5c338c18b7f3b Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/iter_metrics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/metrics.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/metrics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9b6f83d488527685cfccb8d0bd32c297143c0ec5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/metrics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/placement_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/placement_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea195cabbfd632c9a02cc0aff9b80ef719f7ff8c Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/placement_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/queue.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/queue.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..381998ce2412ca8ca4373e325e3cd882cf549c12 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/queue.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/rpdb.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/rpdb.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e380f6f8b96dc756d6e8c60b5e9e08bad495f50 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/rpdb.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/scheduling_strategies.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/scheduling_strategies.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14ff980036f9cf7b216ed69471ccf08759402462 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/scheduling_strategies.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/serialization.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/serialization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4459ee121d2e8579e79343855ab9417e4b3c2319 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/serialization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/serialization_addons.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/serialization_addons.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..216b54b0583f4dbdc7d616e3c5f0a5d644dff319 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/serialization_addons.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/timer.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/timer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c6764aabb5dc89e44957c9aa59794b05338e745d Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/timer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/__pycache__/tpu.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/__pycache__/tpu.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba44c2e1f825e59da1d99a29ce4800ebea841991 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/__pycache__/tpu.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/accelerators/__init__.py b/lib/python3.12/site-packages/ray/util/accelerators/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..53d6a501fbaa17324fed2128f181c6f039062b3c --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/accelerators/__init__.py @@ -0,0 +1,78 @@ +import warnings + +from ray.util import tpu +from ray.util.accelerators.accelerators import ( + AMD_INSTINCT_MI100, + AMD_INSTINCT_MI210, + AMD_INSTINCT_MI250, + AMD_RADEON_HD_7900, + AMD_RADEON_R9_200_HD_7900, + AWS_NEURON_CORE, + GOOGLE_TPU_V2, + GOOGLE_TPU_V3, + GOOGLE_TPU_V4, + GOOGLE_TPU_V5LITEPOD, + GOOGLE_TPU_V5P, + GOOGLE_TPU_V6E, + INTEL_GAUDI, + INTEL_MAX_1100, + INTEL_MAX_1550, + NVIDIA_A100, + NVIDIA_H100, + NVIDIA_L4, + NVIDIA_TESLA_A10G, + NVIDIA_TESLA_K80, + NVIDIA_TESLA_P4, + NVIDIA_TESLA_P100, + NVIDIA_TESLA_T4, + NVIDIA_TESLA_V100, + AMD_INSTINCT_MI250x, + AMD_INSTINCT_MI300x, +) + +__all__ = [ + "tpu", + "NVIDIA_TESLA_V100", + "NVIDIA_TESLA_P100", + "NVIDIA_TESLA_T4", + "NVIDIA_TESLA_P4", + "NVIDIA_TESLA_K80", + "NVIDIA_TESLA_A10G", + "NVIDIA_L4", + "NVIDIA_A100", + "NVIDIA_A100_40G", + "NVIDIA_A100_80G", + "NVIDIA_H100", + "INTEL_MAX_1550", + "INTEL_MAX_1100", + "INTEL_GAUDI", + "AMD_INSTINCT_MI100", + "AMD_INSTINCT_MI210", + "AMD_INSTINCT_MI250", + "AMD_INSTINCT_MI250x", + "AMD_INSTINCT_MI300x", + "AMD_RADEON_R9_200_HD_7900", + "AMD_RADEON_HD_7900", + "AWS_NEURON_CORE", + "GOOGLE_TPU_V2", + "GOOGLE_TPU_V3", + "GOOGLE_TPU_V4", + "GOOGLE_TPU_V5P", + "GOOGLE_TPU_V5LITEPOD", + "GOOGLE_TPU_V6E", + # Deprecated + "NVIDIA_TESLA_A100", +] + + +def __getattr__(name: str): + if name == "NVIDIA_TESLA_A100": + from ray.util.annotations import RayDeprecationWarning + + warnings.warn( + "NVIDIA_TESLA_A100 is deprecated, use NVIDIA_A100 instead.", + RayDeprecationWarning, + stacklevel=2, + ) + return NVIDIA_A100 + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/lib/python3.12/site-packages/ray/util/accelerators/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/accelerators/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb07b2a4b58bd336b3dd79055f436446616925cf Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/accelerators/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/accelerators/__pycache__/accelerators.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/accelerators/__pycache__/accelerators.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a2ee91d8a1f46e52851e2fd137728c5b6313d434 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/accelerators/__pycache__/accelerators.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/accelerators/accelerators.py b/lib/python3.12/site-packages/ray/util/accelerators/accelerators.py new file mode 100644 index 0000000000000000000000000000000000000000..b68d0460b538e0a46ae89a4c2f082b34e7edb98b --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/accelerators/accelerators.py @@ -0,0 +1,44 @@ +NVIDIA_TESLA_V100 = "V100" +NVIDIA_TESLA_P100 = "P100" +NVIDIA_TESLA_T4 = "T4" +NVIDIA_TESLA_P4 = "P4" +NVIDIA_TESLA_K80 = "K80" +NVIDIA_TESLA_A10G = "A10G" +NVIDIA_L4 = "L4" +NVIDIA_L40S = "L40S" +NVIDIA_A100 = "A100" +NVIDIA_H100 = "H100" +NVIDIA_H200 = "H200" +NVIDIA_H20 = "H20" +NVIDIA_B200 = "B200" +INTEL_MAX_1550 = "Intel-GPU-Max-1550" +INTEL_MAX_1100 = "Intel-GPU-Max-1100" +INTEL_GAUDI = "Intel-GAUDI" +AMD_INSTINCT_MI100 = "AMD-Instinct-MI100" +AMD_INSTINCT_MI250x = "AMD-Instinct-MI250X" +AMD_INSTINCT_MI250 = "AMD-Instinct-MI250X-MI250" +AMD_INSTINCT_MI210 = "AMD-Instinct-MI210" +AMD_INSTINCT_MI300A = "AMD-Instinct-MI300A" +AMD_INSTINCT_MI300x = "AMD-Instinct-MI300X-OAM" +AMD_INSTINCT_MI300x_HF = "AMD-Instinct-MI300X-HF" +AMD_INSTINCT_MI308x = "AMD-Instinct-MI308X" +AMD_INSTINCT_MI325x = "AMD-Instinct-MI325X-OAM" +AMD_INSTINCT_MI350x = "AMD-Instinct-MI350X-OAM" +AMD_INSTINCT_MI355x = "AMD-Instinct-MI355X-OAM" +AMD_RADEON_R9_200_HD_7900 = "AMD-Radeon-R9-200-HD-7900" +AMD_RADEON_HD_7900 = "AMD-Radeon-HD-7900" +AWS_NEURON_CORE = "aws-neuron-core" +GOOGLE_TPU_V2 = "TPU-V2" +GOOGLE_TPU_V3 = "TPU-V3" +GOOGLE_TPU_V4 = "TPU-V4" +GOOGLE_TPU_V5P = "TPU-V5P" +GOOGLE_TPU_V5LITEPOD = "TPU-V5LITEPOD" +GOOGLE_TPU_V6E = "TPU-V6E" +HUAWEI_NPU_910B = "Ascend910B" +HUAWEI_NPU_910B4 = "Ascend910B4" + +# Use these instead of NVIDIA_A100 if you need a specific accelerator size. Note that +# these labels are not auto-added to nodes, you'll have to add them manually in +# addition to the default A100 label if needed. +NVIDIA_A100_40G = "A100-40G" +NVIDIA_A100_80G = "A100-80G" diff --git a/lib/python3.12/site-packages/ray/util/actor_group.py b/lib/python3.12/site-packages/ray/util/actor_group.py new file mode 100644 index 0000000000000000000000000000000000000000..5cd343f1b17d3ceacc0470e9c588fff2bc83e2cf --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/actor_group.py @@ -0,0 +1,230 @@ +import logging +import weakref +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple, Type, TypeVar + +import ray +from ray._private.utils import get_ray_doc_version +from ray.actor import ActorHandle +from ray.util.annotations import Deprecated + +T = TypeVar("T") +ActorMetadata = TypeVar("ActorMetadata") + +logger = logging.getLogger(__name__) + + +@dataclass +class ActorWrapper: + """Class containing an actor and its metadata.""" + + actor: ActorHandle + metadata: ActorMetadata + + +@dataclass +class ActorConfig: + num_cpus: float + num_gpus: float + resources: Optional[Dict[str, float]] + init_args: Tuple + init_kwargs: Dict + + +class ActorGroupMethod: + def __init__(self, actor_group: "ActorGroup", method_name: str): + self.actor_group = weakref.ref(actor_group) + self._method_name = method_name + + def __call__(self, *args, **kwargs): + raise TypeError( + "ActorGroup methods cannot be called directly. " + "Instead " + f"of running 'object.{self._method_name}()', try " + f"'object.{self._method_name}.remote()'." + ) + + def remote(self, *args, **kwargs): + return [ + getattr(a.actor, self._method_name).remote(*args, **kwargs) + for a in self.actor_group().actors + ] + + +@Deprecated( + message="For stateless/task processing, use ray.util.multiprocessing, see details " + f"in https://docs.ray.io/en/{get_ray_doc_version()}/ray-more-libs/multiprocessing.html. " # noqa: E501 + "For stateful/actor processing such as batch prediction, use " + "Datasets.map_batches(compute=ActorPoolStrategy, ...), see details in " + f"https://docs.ray.io/en/{get_ray_doc_version()}/data/api/dataset.html#ray.data.Dataset.map_batches.", # noqa: E501 + warning=True, +) +class ActorGroup: + """Group of Ray Actors that can execute arbitrary functions. + + ``ActorGroup`` launches Ray actors according to the given + specification. It can then execute arbitrary Python functions in each of + these actors. + + If not enough resources are available to launch the actors, the Ray + cluster will automatically scale up if autoscaling is enabled. + + Args: + actor_cls: The class to use as the remote actors. + num_actors: The number of the provided Ray actors to + launch. Defaults to 1. + num_cpus_per_actor: The number of CPUs to reserve for each + actor. Fractional values are allowed. Defaults to 1. + num_gpus_per_actor: The number of GPUs to reserve for each + actor. Fractional values are allowed. Defaults to 0. + resources_per_actor (Optional[Dict[str, float]]): + Dictionary specifying the resources that will be + requested for each actor in addition to ``num_cpus_per_actor`` + and ``num_gpus_per_actor``. + init_args, init_kwargs: If ``actor_cls`` is provided, + these args will be used for the actor initialization. + + """ + + def __init__( + self, + actor_cls: Type, + num_actors: int = 1, + num_cpus_per_actor: float = 1, + num_gpus_per_actor: float = 0, + resources_per_actor: Optional[Dict[str, float]] = None, + init_args: Optional[Tuple] = None, + init_kwargs: Optional[Dict] = None, + ): + from ray._common.usage.usage_lib import record_library_usage + + record_library_usage("util.ActorGroup") + + if num_actors <= 0: + raise ValueError( + "The provided `num_actors` must be greater " + f"than 0. Received num_actors={num_actors} " + f"instead." + ) + if num_cpus_per_actor < 0 or num_gpus_per_actor < 0: + raise ValueError( + "The number of CPUs and GPUs per actor must " + "not be negative. Received " + f"num_cpus_per_actor={num_cpus_per_actor} and " + f"num_gpus_per_actor={num_gpus_per_actor}." + ) + + self.actors = [] + + self.num_actors = num_actors + + self.actor_config = ActorConfig( + num_cpus=num_cpus_per_actor, + num_gpus=num_gpus_per_actor, + resources=resources_per_actor, + init_args=init_args or (), + init_kwargs=init_kwargs or {}, + ) + + self._remote_cls = ray.remote( + num_cpus=self.actor_config.num_cpus, + num_gpus=self.actor_config.num_gpus, + resources=self.actor_config.resources, + )(actor_cls) + + self.start() + + def __getattr__(self, item): + if len(self.actors) == 0: + raise RuntimeError( + "This ActorGroup has been shutdown. Please start it again." + ) + # Same implementation as actor.py + return ActorGroupMethod(self, item) + + def __len__(self): + return len(self.actors) + + def __getitem__(self, item): + return self.actors[item] + + def start(self): + """Starts all the actors in this actor group.""" + if self.actors and len(self.actors) > 0: + raise RuntimeError( + "The actors have already been started. " + "Please call `shutdown` first if you want to " + "restart them." + ) + + logger.debug(f"Starting {self.num_actors} actors.") + self.add_actors(self.num_actors) + logger.debug(f"{len(self.actors)} actors have successfully started.") + + def shutdown(self, patience_s: float = 5): + """Shutdown all the actors in this actor group. + + Args: + patience_s: Attempt a graceful shutdown + of the actors for this many seconds. Fallback to force kill + if graceful shutdown is not complete after this time. If + this is less than or equal to 0, immediately force kill all + actors. + """ + logger.debug(f"Shutting down {len(self.actors)} actors.") + if patience_s <= 0: + for actor in self.actors: + ray.kill(actor.actor) + else: + done_refs = [w.actor.__ray_terminate__.remote() for w in self.actors] + # Wait for actors to die gracefully. + done, not_done = ray.wait(done_refs, timeout=patience_s) + if not_done: + logger.debug("Graceful termination failed. Falling back to force kill.") + # If all actors are not able to die gracefully, then kill them. + for actor in self.actors: + ray.kill(actor.actor) + + logger.debug("Shutdown successful.") + self.actors = [] + + def remove_actors(self, actor_indexes: List[int]): + """Removes the actors with the specified indexes. + + Args: + actor_indexes (List[int]): The indexes of the actors to remove. + """ + new_actors = [] + for i in range(len(self.actors)): + if i not in actor_indexes: + new_actors.append(self.actors[i]) + self.actors = new_actors + + def add_actors(self, num_actors: int): + """Adds ``num_actors`` to this ActorGroup. + + Args: + num_actors: The number of actors to add. + """ + new_actors = [] + new_actor_metadata = [] + for _ in range(num_actors): + actor = self._remote_cls.remote( + *self.actor_config.init_args, **self.actor_config.init_kwargs + ) + new_actors.append(actor) + if hasattr(actor, "get_actor_metadata"): + new_actor_metadata.append(actor.get_actor_metadata.remote()) + + # Get metadata from all actors. + metadata = ray.get(new_actor_metadata) + + if len(metadata) == 0: + metadata = [None] * len(new_actors) + + for i in range(len(new_actors)): + self.actors.append(ActorWrapper(actor=new_actors[i], metadata=metadata[i])) + + @property + def actor_metadata(self): + return [a.metadata for a in self.actors] diff --git a/lib/python3.12/site-packages/ray/util/annotations.py b/lib/python3.12/site-packages/ray/util/annotations.py new file mode 100644 index 0000000000000000000000000000000000000000..a2e3fc664d55d15bf554b8c46fcf086b311ed338 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/annotations.py @@ -0,0 +1,268 @@ +import inspect +import sys +import warnings +from enum import Enum +from functools import wraps +from typing import Optional + + +class AnnotationType(Enum): + PUBLIC_API = "PublicAPI" + DEVELOPER_API = "DeveloperAPI" + DEPRECATED = "Deprecated" + UNKNOWN = "Unknown" + + +def PublicAPI(*args, **kwargs): + """Annotation for documenting public APIs. + + Public APIs are classes and methods exposed to end users of Ray. + + If ``stability="alpha"``, the API can be used by advanced users who are + tolerant to and expect breaking changes. + + If ``stability="beta"``, the API is still public and can be used by early + users, but are subject to change. + + If ``stability="stable"``, the APIs will remain backwards compatible across + minor Ray releases (e.g., Ray 1.4 -> 1.8). + + For a full definition of the stability levels, please refer to the + :ref:`Ray API Stability definitions `. + + Args: + stability: One of {"stable", "beta", "alpha"}. + api_group: Optional. Used only for doc rendering purpose. APIs in the same group + will be grouped together in the API doc pages. + + Examples: + >>> from ray.util.annotations import PublicAPI + >>> @PublicAPI + ... def func(x): + ... return x + + >>> @PublicAPI(stability="beta") + ... def func(y): + ... return y + """ + if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): + return PublicAPI(stability="stable", api_group="Others")(args[0]) + + if "stability" in kwargs: + stability = kwargs["stability"] + assert stability in ["stable", "beta", "alpha"], stability + else: + stability = "stable" + api_group = kwargs.get("api_group", "Others") + + def wrap(obj): + if stability in ["alpha", "beta"]: + message = ( + f"**PublicAPI ({stability}):** This API is in {stability} " + "and may change before becoming stable." + ) + _append_doc(obj, message=message) + + _mark_annotated(obj, type=AnnotationType.PUBLIC_API, api_group=api_group) + return obj + + return wrap + + +def DeveloperAPI(*args, **kwargs): + """Annotation for documenting developer APIs. + + Developer APIs are lower-level methods explicitly exposed to advanced Ray + users and library developers. Their interfaces may change across minor + Ray releases. + + Examples: + >>> from ray.util.annotations import DeveloperAPI + >>> @DeveloperAPI + ... def func(x): + ... return x + """ + if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): + return DeveloperAPI()(args[0]) + + def wrap(obj): + _append_doc( + obj, + message="**DeveloperAPI:** This API may change across minor Ray releases.", + ) + _mark_annotated(obj, type=AnnotationType.DEVELOPER_API) + return obj + + return wrap + + +class RayDeprecationWarning(DeprecationWarning): + """Specialized Deprecation Warning for fine grained filtering control""" + + pass + + +# By default, print the first occurrence of matching warnings for +# each module where the warning is issued (regardless of line number) +if not sys.warnoptions: + warnings.filterwarnings("module", category=RayDeprecationWarning) + + +def Deprecated(*args, **kwargs): + """Annotation for documenting a deprecated API. + + Deprecated APIs may be removed in future releases of Ray. + + Args: + message: a message to help users understand the reason for the + deprecation, and provide a migration path. + + Examples: + >>> from ray.util.annotations import Deprecated + >>> @Deprecated + ... def func(x): + ... return x + + >>> @Deprecated(message="g() is deprecated because the API is error " + ... "prone. Please call h() instead.") + ... def g(y): + ... return y + """ + if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): + return Deprecated()(args[0]) + + doc_message = ( + "**DEPRECATED**: This API is deprecated and may be removed " + "in future Ray releases." + ) + warning_message = ( + "This API is deprecated and may be removed in future Ray releases. " + "You could suppress this warning by setting env variable " + 'PYTHONWARNINGS="ignore::DeprecationWarning"' + ) + + warning = kwargs.pop("warning", False) + + if "message" in kwargs: + doc_message = doc_message + "\n" + kwargs["message"] + warning_message = warning_message + "\n" + kwargs["message"] + del kwargs["message"] + + if kwargs: + raise ValueError("Unknown kwargs: {}".format(kwargs.keys())) + + def inner(obj): + _append_doc(obj, message=doc_message, directive="warning") + _mark_annotated(obj, type=AnnotationType.DEPRECATED) + + if not warning: + return obj + + if inspect.isclass(obj): + obj_init = obj.__init__ + + def patched_init(*args, **kwargs): + warnings.warn(warning_message, RayDeprecationWarning, stacklevel=2) + return obj_init(*args, **kwargs) + + obj.__init__ = patched_init + return obj + else: + # class method or function. + @wraps(obj) + def wrapper(*args, **kwargs): + warnings.warn(warning_message, RayDeprecationWarning, stacklevel=2) + return obj(*args, **kwargs) + + return wrapper + + return inner + + +def _append_doc(obj, *, message: str, directive: Optional[str] = None) -> str: + if not obj.__doc__: + obj.__doc__ = "" + + obj.__doc__ = obj.__doc__.rstrip() + + indent = _get_indent(obj.__doc__) + obj.__doc__ += "\n\n" + + if directive is not None: + obj.__doc__ += f"{' ' * indent}.. {directive}::\n\n" + + message = message.replace("\n", "\n" + " " * (indent + 4)) + obj.__doc__ += f"{' ' * (indent + 4)}{message}" + else: + message = message.replace("\n", "\n" + " " * (indent + 4)) + obj.__doc__ += f"{' ' * indent}{message}" + obj.__doc__ += f"\n{' ' * indent}" + + +def _get_indent(docstring: str) -> int: + """ + + Example: + >>> def f(): + ... '''Docstring summary.''' + >>> f.__doc__ + 'Docstring summary.' + >>> _get_indent(f.__doc__) + 0 + + >>> def g(foo): + ... '''Docstring summary. + ... + ... Args: + ... foo: Does bar. + ... ''' + >>> g.__doc__ + 'Docstring summary.\\n\\n Args:\\n foo: Does bar.\\n ' + >>> _get_indent(g.__doc__) + 4 + + >>> class A: + ... def h(): + ... '''Docstring summary. + ... + ... Returns: + ... None. + ... ''' + >>> A.h.__doc__ + 'Docstring summary.\\n\\n Returns:\\n None.\\n ' + >>> _get_indent(A.h.__doc__) + 8 + """ + if not docstring: + return 0 + + non_empty_lines = list(filter(bool, docstring.splitlines())) + if len(non_empty_lines) == 1: + # Docstring contains summary only. + return 0 + + # The docstring summary isn't indented, so check the indentation of the second + # non-empty line. + return len(non_empty_lines[1]) - len(non_empty_lines[1].lstrip()) + + +def _mark_annotated( + obj, type: AnnotationType = AnnotationType.UNKNOWN, api_group="Others" +) -> None: + # Set magic token for check_api_annotations linter. + if hasattr(obj, "__name__"): + obj._annotated = obj.__name__ + obj._annotated_type = type + obj._annotated_api_group = api_group + + +def _is_annotated(obj) -> bool: + # Check the magic token exists and applies to this class (not a subclass). + return hasattr(obj, "_annotated") and obj._annotated == obj.__name__ + + +def _get_annotation_type(obj) -> Optional[str]: + if not _is_annotated(obj): + return None + + return obj._annotated_type.value diff --git a/lib/python3.12/site-packages/ray/util/check_open_ports.py b/lib/python3.12/site-packages/ray/util/check_open_ports.py new file mode 100644 index 0000000000000000000000000000000000000000..67f5e1fd87a56d8ff810b406190ae67d47b95048 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/check_open_ports.py @@ -0,0 +1,181 @@ +"""A CLI utility for check open ports in the Ray cluster. + +See https://www.anyscale.com/blog/update-on-ray-cve-2023-48022-new-verification-tooling-available # noqa: E501 +for more details. +""" +import json +import subprocess +import urllib +from typing import List, Tuple + +import click + +import ray +from ray.autoscaler._private.cli_logger import add_click_logging_options, cli_logger +from ray.autoscaler._private.constants import RAY_PROCESSES +from ray.util.annotations import PublicAPI +from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy + +import psutil + + +def _get_ray_ports() -> List[int]: + unique_ports = set() + + process_infos = [] + for proc in psutil.process_iter(["name", "cmdline"]): + try: + process_infos.append((proc, proc.name(), proc.cmdline())) + except psutil.Error: + pass + + for keyword, filter_by_cmd in RAY_PROCESSES: + for candidate in process_infos: + proc, proc_cmd, proc_args = candidate + corpus = proc_cmd if filter_by_cmd else subprocess.list2cmdline(proc_args) + if keyword in corpus: + try: + for connection in proc.connections(): + if connection.status == psutil.CONN_LISTEN: + unique_ports.add(connection.laddr.port) + except psutil.AccessDenied: + cli_logger.info( + "Access denied to process connections for process," + " worker process probably restarted", + proc, + ) + + return sorted(unique_ports) + + +def _check_for_open_ports_from_internet( + service_url: str, ports: List[int] +) -> Tuple[List[int], List[int]]: + request = urllib.request.Request( + method="POST", + url=service_url, + headers={ + "Content-Type": "application/json", + "X-Ray-Open-Port-Check": "1", + }, + data=json.dumps({"ports": ports}).encode("utf-8"), + ) + + response = urllib.request.urlopen(request) + if response.status != 200: + raise RuntimeError( + f"Failed to check with Ray Open Port Service: {response.status}" + ) + response_body = json.load(response) + + publicly_open_ports = response_body.get("open_ports", []) + checked_ports = response_body.get("checked_ports", []) + + return publicly_open_ports, checked_ports + + +def _check_if_exposed_to_internet( + service_url: str, +) -> Tuple[List[int], List[int]]: + return _check_for_open_ports_from_internet(service_url, _get_ray_ports()) + + +def _check_ray_cluster( + service_url: str, +) -> List[Tuple[str, Tuple[List[int], List[int]]]]: + ray.init(ignore_reinit_error=True) + + @ray.remote(num_cpus=0) + def check(node_id, service_url): + return node_id, _check_if_exposed_to_internet(service_url) + + ray_node_ids = [node["NodeID"] for node in ray.nodes() if node["Alive"]] + cli_logger.info( + f"Cluster has {len(ray_node_ids)} node(s)." + " Scheduling tasks on each to check for exposed ports", + ) + + per_node_tasks = { + node_id: ( + check.options( + scheduling_strategy=NodeAffinitySchedulingStrategy( + node_id=node_id, soft=False + ) + ).remote(node_id, service_url) + ) + for node_id in ray_node_ids + } + + results = [] + for node_id, per_node_task in per_node_tasks.items(): + try: + results.append(ray.get(per_node_task)) + except Exception as e: + cli_logger.info(f"Failed to check on node {node_id}: {e}") + + return results + + +@click.command() +@click.option( + "--yes", "-y", is_flag=True, default=False, help="Don't ask for confirmation." +) +@click.option( + "--service-url", + required=False, + type=str, + default="https://ray-open-port-checker.uc.r.appspot.com/open-port-check", + help="The url of service that checks whether submitted ports are open.", +) +@add_click_logging_options +@PublicAPI +def check_open_ports(yes, service_url): + """Check open ports in the local Ray cluster.""" + if not cli_logger.confirm( + yes=yes, + msg=( + "Do you want to check the local Ray cluster" + " for any nodes with ports accessible to the internet?" + ), + _default=True, + ): + cli_logger.info("Exiting without checking as instructed") + return + + cluster_open_ports = _check_ray_cluster(service_url) + + public_nodes = [] + for node_id, (open_ports, checked_ports) in cluster_open_ports: + if open_ports: + cli_logger.info( + f"[🛑] open ports detected open_ports={open_ports!r} node={node_id!r}" + ) + public_nodes.append((node_id, open_ports, checked_ports)) + else: + cli_logger.info( + f"[🟢] No open ports detected " + f"checked_ports={checked_ports!r} node={node_id!r}" + ) + + cli_logger.info("Check complete, results:") + + if public_nodes: + cli_logger.info( + """ +[🛑] An server on the internet was able to open a connection to one of this Ray +cluster's public IP on one of Ray's internal ports. If this is not a false +positive, this is an extremely unsafe configuration for Ray to be running in. +Ray is not meant to be exposed to untrusted clients and will allow them to run +arbitrary code on your machine. + +You should take immediate action to validate this result and if confirmed shut +down your Ray cluster immediately and take appropriate action to remediate its +exposure. Anything either running on this Ray cluster or that this cluster has +had access to could be at risk. + +For guidance on how to operate Ray safely, please review [Ray's security +documentation](https://docs.ray.io/en/latest/ray-security/index.html). +""".strip() + ) + else: + cli_logger.info("[🟢] No open ports detected from any Ray nodes") diff --git a/lib/python3.12/site-packages/ray/util/check_serialize.py b/lib/python3.12/site-packages/ray/util/check_serialize.py new file mode 100644 index 0000000000000000000000000000000000000000..04e1c9633f26dea5b886c853d0fa13017543efc4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/check_serialize.py @@ -0,0 +1,266 @@ +"""A utility for debugging serialization issues.""" +import inspect +from contextlib import contextmanager +from typing import Any, Optional, Set, Tuple + +import colorama + +# Import ray first to use the bundled colorama +import ray # noqa: F401 +import ray.cloudpickle as cp +from ray.util.annotations import DeveloperAPI + + +@contextmanager +def _indent(printer): + printer.level += 1 + yield + printer.level -= 1 + + +class _Printer: + def __init__(self, print_file): + self.level = 0 + self.print_file = print_file + + def indent(self): + return _indent(self) + + def print(self, msg): + indent = " " * self.level + print(indent + msg, file=self.print_file) + + +@DeveloperAPI +class FailureTuple: + """Represents the serialization 'frame'. + + Attributes: + obj: The object that fails serialization. + name: The variable name of the object. + parent: The object that references the `obj`. + """ + + def __init__(self, obj: Any, name: str, parent: Any): + self.obj = obj + self.name = name + self.parent = parent + + def __repr__(self): + return f"FailTuple({self.name} [obj={self.obj}, parent={self.parent}])" + + +def _inspect_func_serialization(base_obj, depth, parent, failure_set, printer): + """Adds the first-found non-serializable element to the failure_set.""" + assert inspect.isfunction(base_obj) + closure = inspect.getclosurevars(base_obj) + found = False + if closure.globals: + printer.print( + f"Detected {len(closure.globals)} global variables. " + "Checking serializability..." + ) + + with printer.indent(): + for name, obj in closure.globals.items(): + serializable, _ = _inspect_serializability( + obj, + name=name, + depth=depth - 1, + parent=parent, + failure_set=failure_set, + printer=printer, + ) + found = found or not serializable + if found: + break + + if closure.nonlocals: + printer.print( + f"Detected {len(closure.nonlocals)} nonlocal variables. " + "Checking serializability..." + ) + with printer.indent(): + for name, obj in closure.nonlocals.items(): + serializable, _ = _inspect_serializability( + obj, + name=name, + depth=depth - 1, + parent=parent, + failure_set=failure_set, + printer=printer, + ) + found = found or not serializable + if found: + break + if not found: + printer.print( + f"WARNING: Did not find non-serializable object in {base_obj}. " + "This may be an oversight." + ) + return found + + +def _inspect_generic_serialization(base_obj, depth, parent, failure_set, printer): + """Adds the first-found non-serializable element to the failure_set.""" + assert not inspect.isfunction(base_obj) + functions = inspect.getmembers(base_obj, predicate=inspect.isfunction) + found = False + with printer.indent(): + for name, obj in functions: + serializable, _ = _inspect_serializability( + obj, + name=name, + depth=depth - 1, + parent=parent, + failure_set=failure_set, + printer=printer, + ) + found = found or not serializable + if found: + break + + with printer.indent(): + members = inspect.getmembers(base_obj) + for name, obj in members: + if name.startswith("__") and name.endswith("__") or inspect.isbuiltin(obj): + continue + serializable, _ = _inspect_serializability( + obj, + name=name, + depth=depth - 1, + parent=parent, + failure_set=failure_set, + printer=printer, + ) + found = found or not serializable + if found: + break + if not found: + printer.print( + f"WARNING: Did not find non-serializable object in {base_obj}. " + "This may be an oversight." + ) + return found + + +@DeveloperAPI +def inspect_serializability( + base_obj: Any, + name: Optional[str] = None, + depth: int = 3, + print_file: Optional[Any] = None, +) -> Tuple[bool, Set[FailureTuple]]: + """Identifies what objects are preventing serialization. + + Args: + base_obj: Object to be serialized. + name: Optional name of string. + depth: Depth of the scope stack to walk through. Defaults to 3. + print_file: file argument that will be passed to print(). + + Returns: + bool: True if serializable. + set[FailureTuple]: Set of unserializable objects. + + .. versionadded:: 1.1.0 + + """ + printer = _Printer(print_file) + return _inspect_serializability(base_obj, name, depth, None, None, printer) + + +def _inspect_serializability( + base_obj, name, depth, parent, failure_set, printer +) -> Tuple[bool, Set[FailureTuple]]: + colorama.init() + top_level = False + declaration = "" + found = False + if failure_set is None: + top_level = True + failure_set = set() + declaration = f"Checking Serializability of {base_obj}" + printer.print("=" * min(len(declaration), 80)) + printer.print(declaration) + printer.print("=" * min(len(declaration), 80)) + + if name is None: + name = str(base_obj) + else: + printer.print(f"Serializing '{name}' {base_obj}...") + try: + cp.dumps(base_obj) + return True, failure_set + except Exception as e: + printer.print( + f"{colorama.Fore.RED}!!! FAIL{colorama.Fore.RESET} " f"serialization: {e}" + ) + found = True + try: + if depth == 0: + failure_set.add(FailureTuple(base_obj, name, parent)) + # Some objects may not be hashable, so we skip adding this to the set. + except Exception: + pass + + if depth <= 0: + return False, failure_set + + # TODO: we only differentiate between 'function' and 'object' + # but we should do a better job of diving into something + # more specific like a Type, Object, etc. + if inspect.isfunction(base_obj): + _inspect_func_serialization( + base_obj, + depth=depth, + parent=base_obj, + failure_set=failure_set, + printer=printer, + ) + else: + _inspect_generic_serialization( + base_obj, + depth=depth, + parent=base_obj, + failure_set=failure_set, + printer=printer, + ) + + if not failure_set: + failure_set.add(FailureTuple(base_obj, name, parent)) + + if top_level: + printer.print("=" * min(len(declaration), 80)) + if not failure_set: + printer.print( + "Nothing failed the inspect_serialization test, though " + "serialization did not succeed." + ) + else: + fail_vars = ( + f"\n\n\t{colorama.Style.BRIGHT}" + + "\n".join(str(k) for k in failure_set) + + f"{colorama.Style.RESET_ALL}\n\n" + ) + printer.print( + f"Variable: {fail_vars}was found to be non-serializable. " + "There may be multiple other undetected variables that were " + "non-serializable. " + ) + printer.print( + "Consider either removing the " + "instantiation/imports of these variables or moving the " + "instantiation into the scope of the function/class. " + ) + printer.print("=" * min(len(declaration), 80)) + printer.print( + "Check https://docs.ray.io/en/master/ray-core/objects/serialization.html#troubleshooting for more information." # noqa + ) + printer.print( + "If you have any suggestions on how to improve " + "this error message, please reach out to the " + "Ray developers on github.com/ray-project/ray/issues/" + ) + printer.print("=" * min(len(declaration), 80)) + return not found, failure_set diff --git a/lib/python3.12/site-packages/ray/util/client/__init__.py b/lib/python3.12/site-packages/ray/util/client/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..97f4bf2802bce0776e322f18bf7699356faaa052 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/__init__.py @@ -0,0 +1,301 @@ +import logging +import os +import threading +from typing import Any, Dict, List, Optional, Tuple + +import ray._private.ray_constants as ray_constants +from ray._private.client_mode_hook import ( + _explicitly_disable_client_mode, + _explicitly_enable_client_mode, +) +from ray._private.ray_logging import setup_logger +from ray._private.utils import check_version_info +from ray.job_config import JobConfig +from ray.util.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + + +class _ClientContext: + def __init__(self): + from ray.util.client.api import _ClientAPI + + self.api = _ClientAPI() + self.client_worker = None + self._server = None + self._connected_with_init = False + self._inside_client_test = False + + def connect( + self, + conn_str: str, + job_config: JobConfig = None, + secure: bool = False, + metadata: List[Tuple[str, str]] = None, + connection_retries: int = 3, + namespace: str = None, + *, + ignore_version: bool = False, + _credentials: Optional["grpc.ChannelCredentials"] = None, # noqa: F821 + ray_init_kwargs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Connect the Ray Client to a server. + + Args: + conn_str: Connection string, in the form "[host]:port" + job_config: The job config of the server. + secure: Whether to use a TLS secured gRPC channel + metadata: gRPC metadata to send on connect + connection_retries: number of connection attempts to make + ignore_version: whether to ignore Python or Ray version mismatches. + This should only be used for debugging purposes. + + Returns: + Dictionary of connection info, e.g., {"num_clients": 1}. + """ + # Delay imports until connect to avoid circular imports. + from ray.util.client.worker import Worker + + if self.client_worker is not None: + if self._connected_with_init: + return + raise Exception("ray.init() called, but ray client is already connected") + if not self._inside_client_test: + # If we're calling a client connect specifically and we're not + # currently in client mode, ensure we are. + _explicitly_enable_client_mode() + if namespace is not None: + job_config = job_config or JobConfig() + job_config.set_ray_namespace(namespace) + + logging_level = ray_constants.LOGGER_LEVEL + logging_format = ray_constants.LOGGER_FORMAT + + if ray_init_kwargs is None: + ray_init_kwargs = {} + + # NOTE(architkulkarni): env_hook is not supported with Ray Client. + ray_init_kwargs["_skip_env_hook"] = True + + if ray_init_kwargs.get("logging_level") is not None: + logging_level = ray_init_kwargs["logging_level"] + if ray_init_kwargs.get("logging_format") is not None: + logging_format = ray_init_kwargs["logging_format"] + + setup_logger(logging_level, logging_format) + + try: + self.client_worker = Worker( + conn_str, + secure=secure, + _credentials=_credentials, + metadata=metadata, + connection_retries=connection_retries, + ) + self.api.worker = self.client_worker + self.client_worker._server_init(job_config, ray_init_kwargs) + conn_info = self.client_worker.connection_info() + self._check_versions(conn_info, ignore_version) + self._register_serializers() + return conn_info + except Exception: + self.disconnect() + raise + + def _register_serializers(self): + """Register the custom serializer addons at the client side. + + The server side should have already registered the serializers via + regular worker's serialization_context mechanism. + """ + import ray.util.serialization_addons + from ray.util.serialization import StandaloneSerializationContext + + ctx = StandaloneSerializationContext() + ray.util.serialization_addons.apply(ctx) + + def _check_versions(self, conn_info: Dict[str, Any], ignore_version: bool) -> None: + # conn_info has "python_version" and "ray_version" so it can be used to compare. + ignore_version = ignore_version or ("RAY_IGNORE_VERSION_MISMATCH" in os.environ) + check_version_info( + conn_info, + "Ray Client", + raise_on_mismatch=not ignore_version, + python_version_match_level="minor", + ) + + def disconnect(self): + """Disconnect the Ray Client.""" + from ray.util.client.api import _ClientAPI + + if self.client_worker is not None: + self.client_worker.close() + self.api = _ClientAPI() + self.client_worker = None + + # remote can be called outside of a connection, which is why it + # exists on the same API layer as connect() itself. + def remote(self, *args, **kwargs): + """remote is the hook stub passed on to replace `ray.remote`. + + This sets up remote functions or actors, as the decorator, + but does not execute them. + + Args: + args: opaque arguments + kwargs: opaque keyword arguments + """ + return self.api.remote(*args, **kwargs) + + def __getattr__(self, key: str): + if self.is_connected(): + return getattr(self.api, key) + elif key in ["is_initialized", "_internal_kv_initialized"]: + # Client is not connected, thus Ray is not considered initialized. + return lambda: False + else: + raise Exception( + "Ray Client is not connected. Please connect by calling `ray.init`." + ) + + def is_connected(self) -> bool: + if self.client_worker is None: + return False + return self.client_worker.is_connected() + + def init(self, *args, **kwargs): + if self._server is not None: + raise Exception("Trying to start two instances of ray via client") + import ray.util.client.server.server as ray_client_server + + server_handle, address_info = ray_client_server.init_and_serve( + "127.0.0.1", 50051, *args, **kwargs + ) + self._server = server_handle.grpc_server + self.connect("127.0.0.1:50051") + self._connected_with_init = True + return address_info + + def shutdown(self, _exiting_interpreter=False): + self.disconnect() + import ray.util.client.server.server as ray_client_server + + if self._server is None: + return + ray_client_server.shutdown_with_server(self._server, _exiting_interpreter) + self._server = None + + +# All connected context will be put here +# This struct will be guarded by a lock for thread safety +_all_contexts = set() +_lock = threading.Lock() + +# This is the default context which is used when allow_multiple is not True +_default_context = _ClientContext() + + +@DeveloperAPI +class RayAPIStub: + """This class stands in as the replacement API for the `import ray` module. + + Much like the ray module, this mostly delegates the work to the + _client_worker. As parts of the ray API are covered, they are piped through + here or on the client worker API. + """ + + def __init__(self): + self._cxt = threading.local() + self._cxt.handler = _default_context + self._inside_client_test = False + + def get_context(self): + try: + return self._cxt.__getattribute__("handler") + except AttributeError: + self._cxt.handler = _default_context + return self._cxt.handler + + def set_context(self, cxt): + old_cxt = self.get_context() + if cxt is None: + self._cxt.handler = _ClientContext() + else: + self._cxt.handler = cxt + return old_cxt + + def is_default(self): + return self.get_context() == _default_context + + def connect(self, *args, **kw_args): + self.get_context()._inside_client_test = self._inside_client_test + conn = self.get_context().connect(*args, **kw_args) + global _lock, _all_contexts + with _lock: + _all_contexts.add(self._cxt.handler) + return conn + + def disconnect(self, *args, **kw_args): + global _lock, _all_contexts, _default_context + with _lock: + if _default_context == self.get_context(): + for cxt in _all_contexts: + cxt.disconnect(*args, **kw_args) + _all_contexts = set() + else: + self.get_context().disconnect(*args, **kw_args) + if self.get_context() in _all_contexts: + _all_contexts.remove(self.get_context()) + if len(_all_contexts) == 0: + _explicitly_disable_client_mode() + + def remote(self, *args, **kwargs): + return self.get_context().remote(*args, **kwargs) + + def __getattr__(self, name): + return self.get_context().__getattr__(name) + + def is_connected(self, *args, **kwargs): + return self.get_context().is_connected(*args, **kwargs) + + def init(self, *args, **kwargs): + ret = self.get_context().init(*args, **kwargs) + global _lock, _all_contexts + with _lock: + _all_contexts.add(self._cxt.handler) + return ret + + def shutdown(self, *args, **kwargs): + global _lock, _all_contexts + with _lock: + if _default_context == self.get_context(): + for cxt in _all_contexts: + cxt.shutdown(*args, **kwargs) + _all_contexts = set() + else: + self.get_context().shutdown(*args, **kwargs) + if self.get_context() in _all_contexts: + _all_contexts.remove(self.get_context()) + if len(_all_contexts) == 0: + _explicitly_disable_client_mode() + + +ray = RayAPIStub() + + +@DeveloperAPI +def num_connected_contexts(): + """Return the number of client connections active.""" + global _lock, _all_contexts + with _lock: + return len(_all_contexts) + + +# Someday we might add methods in this module so that someone who +# tries to `import ray_client as ray` -- as a module, instead of +# `from ray_client import ray` -- as the API stub +# still gets expected functionality. This is the way the ray package +# worked in the past. +# +# This really calls for PEP 562: https://www.python.org/dev/peps/pep-0562/ +# But until Python 3.6 is EOL, here we are. diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d23e06ba906cda4ea58456d0d9bef4cad775f4de Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6cd5a4ab1398825b842d08d623629ebb5fa3a6c2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/client_app.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/client_app.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d00da77f62f7c86215c0aeb93aa8abc6b3f0f08a Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/client_app.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/client_pickler.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/client_pickler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16b6f15831a8f02fd5070e807c02df39175590dc Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/client_pickler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8862f2d0a36e4eea42338bacd7b9f89b29ddbc03 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/dataclient.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/dataclient.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83c2d23a55ad29dea65067bd037034fcc86994bd Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/dataclient.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/logsclient.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/logsclient.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0c83110efdce4678f28710c570bc614e0c8d9a7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/logsclient.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/options.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/options.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a70ae0b8662d292077d2101e72e0018a7e23d30f Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/options.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/ray_client_helpers.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/ray_client_helpers.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0e49096dc8b1a83724c36c70b5c2ffdf6cbdd5e7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/ray_client_helpers.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/runtime_context.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/runtime_context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a10067f59363153bd60a28c3239f734fced7643e Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/runtime_context.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/__pycache__/worker.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/__pycache__/worker.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86f94b9d8e90d16d9f7f63a33f6e6af5cafdf340 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/__pycache__/worker.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/api.py b/lib/python3.12/site-packages/ray/util/client/api.py new file mode 100644 index 0000000000000000000000000000000000000000..f9dd185a5e36af1bd40e9bb0a659fc0d70a714d3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/api.py @@ -0,0 +1,406 @@ +"""This file defines the interface between the ray client worker +and the overall ray module API. +""" +import json +import logging +from concurrent.futures import Future +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union + +from ray._common import ray_option_utils +from ray.util.client.runtime_context import _ClientWorkerPropertyAPI + +if TYPE_CHECKING: + from ray.actor import ActorClass + from ray.core.generated.ray_client_pb2 import DataResponse + from ray.remote_function import RemoteFunction + from ray.util.client.common import ClientActorHandle, ClientObjectRef, ClientStub + +logger = logging.getLogger(__name__) + + +def _as_bytes(value): + if isinstance(value, str): + return value.encode("utf-8") + return value + + +class _ClientAPI: + """The Client-side methods corresponding to the ray API. Delegates + to the Client Worker that contains the connection to the ClientServer. + """ + + def __init__(self, worker=None): + self.worker = worker + + def get(self, vals, *, timeout=None): + """get is the hook stub passed on to replace `ray.get` + + Args: + vals: [Client]ObjectRef or list of these refs to retrieve. + timeout: Optional timeout in milliseconds + """ + return self.worker.get(vals, timeout=timeout) + + def put(self, *args, **kwargs): + """put is the hook stub passed on to replace `ray.put` + + Args: + val: The value to `put`. + args: opaque arguments + kwargs: opaque keyword arguments + """ + return self.worker.put(*args, **kwargs) + + def wait(self, *args, **kwargs): + """wait is the hook stub passed on to replace `ray.wait` + + Args: + args: opaque arguments + kwargs: opaque keyword arguments + """ + return self.worker.wait(*args, **kwargs) + + def remote(self, *args, **kwargs): + """remote is the hook stub passed on to replace `ray.remote`. + + This sets up remote functions or actors, as the decorator, + but does not execute them. + + Args: + args: opaque arguments + kwargs: opaque keyword arguments + """ + # Delayed import to avoid a cyclic import + from ray.util.client.common import remote_decorator + + if len(args) == 1 and len(kwargs) == 0 and callable(args[0]): + # This is the case where the decorator is just @ray.remote. + return remote_decorator(options=None)(args[0]) + assert ( + len(args) == 0 and len(kwargs) > 0 + ), ray_option_utils.remote_args_error_string + return remote_decorator(options=kwargs) + + # TODO(mwtian): consider adding _internal_ prefix to call_remote / + # call_release / call_retain. + def call_remote(self, instance: "ClientStub", *args, **kwargs) -> List[Future]: + """call_remote is called by stub objects to execute them remotely. + + This is used by stub objects in situations where they're called + with .remote, eg, `f.remote()` or `actor_cls.remote()`. + This allows the client stub objects to delegate execution to be + implemented in the most effective way whether it's in the client, + clientserver, or raylet worker. + + Args: + instance: The Client-side stub reference to a remote object + args: opaque arguments + kwargs: opaque keyword arguments + """ + return self.worker.call_remote(instance, *args, **kwargs) + + def call_release(self, id: bytes) -> None: + """Attempts to release an object reference. + + When client references are destructed, they release their reference, + which can opportunistically send a notification through the datachannel + to release the reference being held for that object on the server. + + Args: + id: The id of the reference to release on the server side. + """ + return self.worker.call_release(id) + + def call_retain(self, id: bytes) -> None: + """Attempts to retain a client object reference. + + Increments the reference count on the client side, to prevent + the client worker from attempting to release the server reference. + + Args: + id: The id of the reference to retain on the client side. + """ + return self.worker.call_retain(id) + + def close(self) -> None: + """close cleans up an API connection by closing any channels or + shutting down any servers gracefully. + """ + return self.worker.close() + + def get_actor( + self, name: str, namespace: Optional[str] = None + ) -> "ClientActorHandle": + """Returns a handle to an actor by name. + + Args: + name: The name passed to this actor by + Actor.options(name="name").remote() + """ + return self.worker.get_actor(name, namespace) + + def list_named_actors(self, all_namespaces: bool = False) -> List[str]: + """List all named actors in the system. + + Actors must have been created with Actor.options(name="name").remote(). + This works for both detached & non-detached actors. + + By default, only actors in the current namespace will be returned + and the returned entries will simply be their name. + + If `all_namespaces` is set to True, all actors in the cluster will be + returned regardless of namespace, and the retunred entries will be of + the form '/'. + """ + return self.worker.list_named_actors(all_namespaces) + + def kill(self, actor: "ClientActorHandle", *, no_restart=True): + """kill forcibly stops an actor running in the cluster + + Args: + no_restart: Whether this actor should be restarted if it's a + restartable actor. + """ + return self.worker.terminate_actor(actor, no_restart) + + def cancel(self, obj: "ClientObjectRef", *, force=False, recursive=True): + """Cancels a task on the cluster. + + If the specified task is pending execution, it will not be executed. If + the task is currently executing, the behavior depends on the ``force`` + flag, as per `ray.cancel()` + + Only non-actor tasks can be canceled. Canceled tasks will not be + retried (max_retries will not be respected). + + Args: + object_ref: ObjectRef returned by the task + that should be canceled. + force: Whether to force-kill a running task by killing + the worker that is running the task. + recursive: Whether to try to cancel tasks submitted by + the task specified. + """ + return self.worker.terminate_task(obj, force, recursive) + + # Various metadata methods for the client that are defined in the protocol. + def is_initialized(self) -> bool: + """True if our client is connected, and if the server is initialized. + Returns: + A boolean determining if the client is connected and + server initialized. + """ + return self.worker.is_initialized() + + def nodes(self): + """Get a list of the nodes in the cluster (for debugging only). + + Returns: + Information about the Ray clients in the cluster. + """ + # This should be imported here, otherwise, it will error doc build. + import ray.core.generated.ray_client_pb2 as ray_client_pb2 + + return self.worker.get_cluster_info(ray_client_pb2.ClusterInfoType.NODES) + + def method(self, *args, **kwargs): + """Annotate an actor method + + Args: + num_returns: The number of object refs that should be returned by + invocations of this actor method. + """ + + # NOTE: So this follows the same logic as in ray/actor.py::method() + # The reason to duplicate it here is to simplify the client mode + # redirection logic. As the annotated method gets pickled and sent to + # the server from the client it carries this private variable, it + # activates the same logic on the server side; so there's no need to + # pass anything else. It's inside the class definition that becomes an + # actor. Similar annotations would follow the same way. + valid_kwargs = ["num_returns", "concurrency_group"] + error_string = ( + "The @ray.method decorator must be applied using at least one of " + f"the arguments in the list {valid_kwargs}, for example " + "'@ray.method(num_returns=2)'." + ) + assert len(args) == 0 and len(kwargs) > 0, error_string + for key in kwargs: + key_error_string = ( + f'Unexpected keyword argument to @ray.method: "{key}". The ' + f"supported keyword arguments are {valid_kwargs}" + ) + assert key in valid_kwargs, key_error_string + + def annotate_method(method): + if "num_returns" in kwargs: + method.__ray_num_returns__ = kwargs["num_returns"] + if "concurrency_group" in kwargs: + method.__ray_concurrency_group__ = kwargs["concurrency_group"] + return method + + return annotate_method + + def cluster_resources(self): + """Get the current total cluster resources. + + Note that this information can grow stale as nodes are added to or + removed from the cluster. + + Returns: + A dictionary mapping resource name to the total quantity of that + resource in the cluster. + """ + # This should be imported here, otherwise, it will error doc build. + import ray.core.generated.ray_client_pb2 as ray_client_pb2 + + return self.worker.get_cluster_info( + ray_client_pb2.ClusterInfoType.CLUSTER_RESOURCES + ) + + def available_resources(self): + """Get the current available cluster resources. + + This is different from `cluster_resources` in that this will return + idle (available) resources rather than total resources. + + Note that this information can grow stale as tasks start and finish. + + Returns: + A dictionary mapping resource name to the total quantity of that + resource in the cluster. + """ + # This should be imported here, otherwise, it will error doc build. + import ray.core.generated.ray_client_pb2 as ray_client_pb2 + + return self.worker.get_cluster_info( + ray_client_pb2.ClusterInfoType.AVAILABLE_RESOURCES + ) + + def get_runtime_context(self): + """Return a Ray RuntimeContext describing the state on the server + + Returns: + A RuntimeContext wrapping a client making get_cluster_info calls. + """ + return _ClientWorkerPropertyAPI(self.worker).build_runtime_context() + + # Client process isn't assigned any GPUs. + def get_gpu_ids(self) -> list: + return [] + + def timeline(self, filename: Optional[str] = None) -> Optional[List[Any]]: + logger.warning( + "Timeline will include events from other clients using this server." + ) + # This should be imported here, otherwise, it will error doc build. + import ray.core.generated.ray_client_pb2 as ray_client_pb2 + + all_events = self.worker.get_cluster_info( + ray_client_pb2.ClusterInfoType.TIMELINE + ) + if filename is not None: + with open(filename, "w") as outfile: + json.dump(all_events, outfile) + else: + return all_events + + def _internal_kv_initialized(self) -> bool: + """Hook for internal_kv._internal_kv_initialized.""" + # NOTE(edoakes): the kv is always initialized because we initialize it + # manually in the proxier with a GCS client if Ray hasn't been + # initialized yet. + return True + + def _internal_kv_exists( + self, key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None + ) -> bool: + """Hook for internal_kv._internal_kv_exists.""" + return self.worker.internal_kv_exists( + _as_bytes(key), namespace=_as_bytes(namespace) + ) + + def _internal_kv_get( + self, key: Union[str, bytes], *, namespace: Optional[Union[str, bytes]] = None + ) -> bytes: + """Hook for internal_kv._internal_kv_get.""" + return self.worker.internal_kv_get( + _as_bytes(key), namespace=_as_bytes(namespace) + ) + + def _internal_kv_put( + self, + key: Union[str, bytes], + value: Union[str, bytes], + overwrite: bool = True, + *, + namespace: Optional[Union[str, bytes]] = None, + ) -> bool: + """Hook for internal_kv._internal_kv_put.""" + return self.worker.internal_kv_put( + _as_bytes(key), _as_bytes(value), overwrite, namespace=_as_bytes(namespace) + ) + + def _internal_kv_del( + self, + key: Union[str, bytes], + *, + del_by_prefix: bool = False, + namespace: Optional[Union[str, bytes]] = None, + ) -> int: + """Hook for internal_kv._internal_kv_del.""" + return self.worker.internal_kv_del( + _as_bytes(key), del_by_prefix=del_by_prefix, namespace=_as_bytes(namespace) + ) + + def _internal_kv_list( + self, + prefix: Union[str, bytes], + *, + namespace: Optional[Union[str, bytes]] = None, + ) -> List[bytes]: + """Hook for internal_kv._internal_kv_list.""" + return self.worker.internal_kv_list( + _as_bytes(prefix), namespace=_as_bytes(namespace) + ) + + def _pin_runtime_env_uri(self, uri: str, expiration_s: int) -> None: + """Hook for internal_kv._pin_runtime_env_uri.""" + return self.worker.pin_runtime_env_uri(uri, expiration_s) + + def _convert_actor(self, actor: "ActorClass") -> str: + """Register a ClientActorClass for the ActorClass and return a UUID""" + return self.worker._convert_actor(actor) + + def _convert_function(self, func: "RemoteFunction") -> str: + """Register a ClientRemoteFunc for the ActorClass and return a UUID""" + return self.worker._convert_function(func) + + def _get_converted(self, key: str) -> "ClientStub": + """Given a UUID, return the converted object""" + return self.worker._get_converted(key) + + def _converted_key_exists(self, key: str) -> bool: + """Check if a key UUID is present in the store of converted objects.""" + return self.worker._converted_key_exists(key) + + def __getattr__(self, key: str): + if not key.startswith("_"): + raise NotImplementedError( + "Not available in Ray client: `ray.{}`. This method is only " + "available within Ray remote functions and is not yet " + "implemented in the client API.".format(key) + ) + return self.__getattribute__(key) + + def _register_callback( + self, ref: "ClientObjectRef", callback: Callable[["DataResponse"], None] + ) -> None: + self.worker.register_callback(ref, callback) + + def _get_dashboard_url(self) -> str: + import ray.core.generated.ray_client_pb2 as ray_client_pb2 + + return self.worker.get_cluster_info( + ray_client_pb2.ClusterInfoType.DASHBOARD_URL + ).get("dashboard_url", "") diff --git a/lib/python3.12/site-packages/ray/util/client/client_app.py b/lib/python3.12/site-packages/ray/util/client/client_app.py new file mode 100644 index 0000000000000000000000000000000000000000..612700147f4f4b5dacf8a4b58c674f41389a6329 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/client_app.py @@ -0,0 +1,91 @@ +from typing import Tuple + +from ray.util.client import ray + +ray.connect("localhost:50051") + + +@ray.remote +class HelloActor: + def __init__(self): + self.count = 0 + + def say_hello(self, whom: str) -> Tuple[str, int]: + self.count += 1 + return ("Hello " + whom, self.count) + + +actor = HelloActor.remote() +s, count = ray.get(actor.say_hello.remote("you")) +print(s, count) +assert s == "Hello you" +assert count == 1 +s, count = ray.get(actor.say_hello.remote("world")) +print(s, count) +assert s == "Hello world" +assert count == 2 + + +@ray.remote +def plus2(x): + return x + 2 + + +@ray.remote +def fact(x): + print(x, type(fact)) + if x <= 0: + return 1 + # This hits the "nested tasks" issue + # https://github.com/ray-project/ray/issues/3644 + # So we're on the right track! + return ray.get(fact.remote(x - 1)) * x + + +@ray.remote +def get_nodes(): + return ray.nodes() # Can access the full Ray API in remote methods. + + +print("Cluster nodes", ray.get(get_nodes.remote())) +print(ray.nodes()) + +objectref = ray.put("hello world") + +# `ClientObjectRef(...)` +print(objectref) + +# `hello world` +print(ray.get(objectref)) + +ref2 = plus2.remote(234) +# `ClientObjectRef(...)` +print(ref2) +# `236` +print(ray.get(ref2)) + +ref3 = fact.remote(20) +# `ClientObjectRef(...)` +print(ref3) +# `2432902008176640000` +print(ray.get(ref3)) + +# Reuse the cached ClientRemoteFunc object +ref4 = fact.remote(5) +# `120` +print(ray.get(ref4)) + +ref5 = fact.remote(10) + +print([ref2, ref3, ref4, ref5]) +# should return ref2, ref3, ref4 +res = ray.wait([ref5, ref2, ref3, ref4], num_returns=3) +print(res) +assert [ref2, ref3, ref4] == res[0] +assert [ref5] == res[1] + +# should return ref2, ref3, ref4, ref5 +res = ray.wait([ref2, ref3, ref4, ref5], num_returns=4) +print(res) +assert [ref2, ref3, ref4, ref5] == res[0] +assert [] == res[1] diff --git a/lib/python3.12/site-packages/ray/util/client/client_pickler.py b/lib/python3.12/site-packages/ray/util/client/client_pickler.py new file mode 100644 index 0000000000000000000000000000000000000000..39a025f1efacc289fc07543eadbf56da1275a363 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/client_pickler.py @@ -0,0 +1,175 @@ +"""Implements the client side of the client/server pickling protocol. + +All ray client client/server data transfer happens through this pickling +protocol. The model is as follows: + + * All Client objects (eg ClientObjectRef) always live on the client and + are never represented in the server + * All Ray objects (eg, ray.ObjectRef) always live on the server and are + never returned to the client + * In order to translate between these two references, PickleStub tuples + are generated as persistent ids in the data blobs during the pickling + and unpickling of these objects. + +The PickleStubs have just enough information to find or generate their +associated partner object on either side. + +This also has the advantage of avoiding predefined pickle behavior for ray +objects, which may include ray internal reference counting. + +ClientPickler dumps things from the client into the appropriate stubs +ServerUnpickler loads stubs from the server into their client counterparts. +""" + +import io +import pickle # noqa: F401 +from typing import Any, Dict, NamedTuple, Optional + +import ray.cloudpickle as cloudpickle +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +from ray.util.client import RayAPIStub +from ray.util.client.common import ( + ClientActorClass, + ClientActorHandle, + ClientActorRef, + ClientObjectRef, + ClientRemoteFunc, + ClientRemoteMethod, + InProgressSentinel, + OptionWrapper, +) + + +# NOTE(barakmich): These PickleStubs are really close to +# the data for an execution, with no arguments. Combine the two? +class PickleStub( + NamedTuple( + "PickleStub", + [ + ("type", str), + ("client_id", str), + ("ref_id", bytes), + ("name", Optional[str]), + ("baseline_options", Optional[Dict]), + ], + ) +): + def __reduce__(self): + # PySpark's namedtuple monkey patch breaks compatibility with + # cloudpickle. Thus we revert this patch here if it exists. + return object.__reduce__(self) + + +class ClientPickler(cloudpickle.CloudPickler): + def __init__(self, client_id, *args, **kwargs): + super().__init__(*args, **kwargs) + self.client_id = client_id + + def persistent_id(self, obj): + if isinstance(obj, RayAPIStub): + return PickleStub( + type="Ray", + client_id=self.client_id, + ref_id=b"", + name=None, + baseline_options=None, + ) + elif isinstance(obj, ClientObjectRef): + return PickleStub( + type="Object", + client_id=self.client_id, + ref_id=obj.id, + name=None, + baseline_options=None, + ) + elif isinstance(obj, ClientActorHandle): + return PickleStub( + type="Actor", + client_id=self.client_id, + ref_id=obj._actor_id.id, + name=None, + baseline_options=None, + ) + elif isinstance(obj, ClientRemoteFunc): + if obj._ref is None: + obj._ensure_ref() + if type(obj._ref) is InProgressSentinel: + return PickleStub( + type="RemoteFuncSelfReference", + client_id=self.client_id, + ref_id=obj._client_side_ref.id, + name=None, + baseline_options=None, + ) + return PickleStub( + type="RemoteFunc", + client_id=self.client_id, + ref_id=obj._ref.id, + name=None, + baseline_options=obj._options, + ) + elif isinstance(obj, ClientActorClass): + if obj._ref is None: + obj._ensure_ref() + if type(obj._ref) is InProgressSentinel: + return PickleStub( + type="RemoteActorSelfReference", + client_id=self.client_id, + ref_id=obj._client_side_ref.id, + name=None, + baseline_options=None, + ) + return PickleStub( + type="RemoteActor", + client_id=self.client_id, + ref_id=obj._ref.id, + name=None, + baseline_options=obj._options, + ) + elif isinstance(obj, ClientRemoteMethod): + return PickleStub( + type="RemoteMethod", + client_id=self.client_id, + ref_id=obj._actor_handle.actor_ref.id, + name=obj._method_name, + baseline_options=None, + ) + elif isinstance(obj, OptionWrapper): + raise NotImplementedError("Sending a partial option is unimplemented") + return None + + +class ServerUnpickler(pickle.Unpickler): + def persistent_load(self, pid): + assert isinstance(pid, PickleStub) + if pid.type == "Object": + return ClientObjectRef(pid.ref_id) + elif pid.type == "Actor": + return ClientActorHandle(ClientActorRef(pid.ref_id)) + else: + raise NotImplementedError("Being passed back an unknown stub") + + +def dumps_from_client(obj: Any, client_id: str, protocol=None) -> bytes: + with io.BytesIO() as file: + cp = ClientPickler(client_id, file, protocol=protocol) + cp.dump(obj) + return file.getvalue() + + +def loads_from_server( + data: bytes, *, fix_imports=True, encoding="ASCII", errors="strict" +) -> Any: + if isinstance(data, str): + raise TypeError("Can't load pickle from unicode string") + file = io.BytesIO(data) + return ServerUnpickler( + file, fix_imports=fix_imports, encoding=encoding, errors=errors + ).load() + + +def convert_to_arg(val: Any, client_id: str) -> ray_client_pb2.Arg: + out = ray_client_pb2.Arg() + out.local = ray_client_pb2.Arg.Locality.INTERNED + out.data = dumps_from_client(val, client_id) + return out diff --git a/lib/python3.12/site-packages/ray/util/client/common.py b/lib/python3.12/site-packages/ray/util/client/common.py new file mode 100644 index 0000000000000000000000000000000000000000..80435bc5c4fdb23363f552ff15edc64098ac1204 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/common.py @@ -0,0 +1,954 @@ +import inspect +import logging +import os +import pickle +import threading +import uuid +from collections import OrderedDict +from concurrent.futures import Future +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import grpc + +import ray._raylet as raylet +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc +from ray._common.signature import extract_signature, get_signature +from ray._private import ray_constants +from ray._private.inspect_util import ( + is_class_method, + is_cython, + is_function_or_method, + is_static_method, +) +from ray._private.utils import check_oversized_function +from ray.util.client import ray +from ray.util.client.options import validate_options +from ray.util.common import INT32_MAX + +logger = logging.getLogger(__name__) + +# gRPC status codes that the client shouldn't attempt to recover from +# Resource exhausted: Server is low on resources, or has hit the max number +# of client connections +# Invalid argument: Reserved for application errors +# Not found: Set if the client is attempting to reconnect to a session that +# does not exist +# Failed precondition: Reserverd for application errors +# Aborted: Set when an error is serialized into the details of the context, +# signals that error should be deserialized on the client side +GRPC_UNRECOVERABLE_ERRORS = ( + grpc.StatusCode.RESOURCE_EXHAUSTED, + grpc.StatusCode.INVALID_ARGUMENT, + grpc.StatusCode.NOT_FOUND, + grpc.StatusCode.FAILED_PRECONDITION, + grpc.StatusCode.ABORTED, +) + +# TODO: Instead of just making the max message size large, the right thing to +# do is to split up the bytes representation of serialized data into multiple +# messages and reconstruct them on either end. That said, since clients are +# drivers and really just feed initial things in and final results out, (when +# not going to S3 or similar) then a large limit will suffice for many use +# cases. +# +# Currently, this is 2GiB, the max for a signed int. +GRPC_MAX_MESSAGE_SIZE = (2 * 1024 * 1024 * 1024) - 1 + +# 30 seconds because ELB timeout is 60 seconds +GRPC_KEEPALIVE_TIME_MS = 1000 * 30 + +# Long timeout because we do not want gRPC ending a connection. +GRPC_KEEPALIVE_TIMEOUT_MS = 1000 * 600 + +GRPC_OPTIONS = [ + *ray_constants.GLOBAL_GRPC_OPTIONS, + ("grpc.max_send_message_length", GRPC_MAX_MESSAGE_SIZE), + ("grpc.max_receive_message_length", GRPC_MAX_MESSAGE_SIZE), + ("grpc.keepalive_time_ms", GRPC_KEEPALIVE_TIME_MS), + ("grpc.keepalive_timeout_ms", GRPC_KEEPALIVE_TIMEOUT_MS), + ("grpc.keepalive_permit_without_calls", 1), + # Send an infinite number of pings + ("grpc.http2.max_pings_without_data", 0), + ("grpc.http2.min_ping_interval_without_data_ms", GRPC_KEEPALIVE_TIME_MS - 50), + # Allow many strikes + ("grpc.http2.max_ping_strikes", 0), +] + +CLIENT_SERVER_MAX_THREADS = float(os.getenv("RAY_CLIENT_SERVER_MAX_THREADS", 100)) + +# Large objects are chunked into 5 MiB messages, ref PR #35025 +OBJECT_TRANSFER_CHUNK_SIZE = 5 * 2**20 + +# Warn the user if the object being transferred is larger than 2 GiB +OBJECT_TRANSFER_WARNING_SIZE = 2 * 2**30 + + +class ClientObjectRef(raylet.ObjectRef): + def __init__(self, id: Union[bytes, Future]): + self._mutex = threading.Lock() + self._worker = ray.get_context().client_worker + self._id_future = None + if isinstance(id, bytes): + self._set_id(id) + elif isinstance(id, Future): + self._id_future = id + else: + raise TypeError("Unexpected type for id {}".format(id)) + + def __del__(self): + if self._worker is not None and self._worker.is_connected(): + try: + if not self.is_nil(): + self._worker.call_release(self.id) + except Exception: + logger.info( + "Exception in ObjectRef is ignored in destructor. " + "To receive this exception in application code, call " + "a method on the actor reference before its destructor " + "is run." + ) + + def binary(self): + self._wait_for_id() + return super().binary() + + def hex(self): + self._wait_for_id() + return super().hex() + + def is_nil(self): + self._wait_for_id() + return super().is_nil() + + def __hash__(self): + self._wait_for_id() + return hash(self.id) + + def task_id(self): + self._wait_for_id() + return super().task_id() + + @property + def id(self): + return self.binary() + + def future(self) -> Future: + fut = Future() + + def set_future(data: Any) -> None: + """Schedules a callback to set the exception or result + in the Future.""" + + if isinstance(data, Exception): + fut.set_exception(data) + else: + fut.set_result(data) + + self._on_completed(set_future) + + # Prevent this object ref from being released. + fut.object_ref = self + return fut + + def _on_completed(self, py_callback: Callable[[Any], None]) -> None: + """Register a callback that will be called after Object is ready. + If the ObjectRef is already ready, the callback will be called soon. + The callback should take the result as the only argument. The result + can be an exception object in case of task error. + """ + + def deserialize_obj( + resp: Union[ray_client_pb2.DataResponse, Exception] + ) -> None: + from ray.util.client.client_pickler import loads_from_server + + if isinstance(resp, Exception): + data = resp + elif isinstance(resp, bytearray): + data = loads_from_server(resp) + else: + obj = resp.get + data = None + if not obj.valid: + data = loads_from_server(resp.get.error) + else: + data = loads_from_server(resp.get.data) + + py_callback(data) + + self._worker.register_callback(self, deserialize_obj) + + def _set_id(self, id): + super()._set_id(id) + self._worker.call_retain(id) + + def _wait_for_id(self, timeout=None): + if self._id_future: + with self._mutex: + if self._id_future: + self._set_id(self._id_future.result(timeout=timeout)) + self._id_future = None + + +class ClientActorRef(raylet.ActorID): + def __init__( + self, + id: Union[bytes, Future], + weak_ref: Optional[bool] = False, + ): + self._weak_ref = weak_ref + self._mutex = threading.Lock() + self._worker = ray.get_context().client_worker + if isinstance(id, bytes): + self._set_id(id) + self._id_future = None + elif isinstance(id, Future): + self._id_future = id + else: + raise TypeError("Unexpected type for id {}".format(id)) + + def __del__(self): + if self._weak_ref: + return + + if self._worker is not None and self._worker.is_connected(): + try: + if not self.is_nil(): + self._worker.call_release(self.id) + except Exception: + logger.debug( + "Exception from actor creation is ignored in destructor. " + "To receive this exception in application code, call " + "a method on the actor reference before its destructor " + "is run." + ) + + def binary(self): + self._wait_for_id() + return super().binary() + + def hex(self): + self._wait_for_id() + return super().hex() + + def is_nil(self): + self._wait_for_id() + return super().is_nil() + + def __hash__(self): + self._wait_for_id() + return hash(self.id) + + @property + def id(self): + return self.binary() + + def _set_id(self, id): + super()._set_id(id) + self._worker.call_retain(id) + + def _wait_for_id(self, timeout=None): + if self._id_future: + with self._mutex: + if self._id_future: + self._set_id(self._id_future.result(timeout=timeout)) + self._id_future = None + + +class ClientStub: + pass + + +class ClientRemoteFunc(ClientStub): + """A stub created on the Ray Client to represent a remote + function that can be exectued on the cluster. + + This class is allowed to be passed around between remote functions. + + Args: + _func: The actual function to execute remotely + _name: The original name of the function + _ref: The ClientObjectRef of the pickled code of the function, _func + """ + + def __init__(self, f, options=None): + self._lock = threading.Lock() + self._func = f + self._name = f.__name__ + self._signature = get_signature(f) + self._ref = None + self._client_side_ref = ClientSideRefID.generate_id() + self._options = validate_options(options) + + def __call__(self, *args, **kwargs): + raise TypeError( + "Remote function cannot be called directly. " + f"Use {self._name}.remote method instead" + ) + + def remote(self, *args, **kwargs): + # Check if supplied parameters match the function signature. Same case + # at the other callsites. + self._signature.bind(*args, **kwargs) + return return_refs(ray.call_remote(self, *args, **kwargs)) + + def options(self, **kwargs): + return OptionWrapper(self, kwargs) + + def _remote(self, args=None, kwargs=None, **option_args): + if args is None: + args = [] + if kwargs is None: + kwargs = {} + return self.options(**option_args).remote(*args, **kwargs) + + def __repr__(self): + return "ClientRemoteFunc(%s, %s)" % (self._name, self._ref) + + def _ensure_ref(self): + with self._lock: + if self._ref is None: + # While calling ray.put() on our function, if + # our function is recursive, it will attempt to + # encode the ClientRemoteFunc -- itself -- and + # infinitely recurse on _ensure_ref. + # + # So we set the state of the reference to be an + # in-progress self reference value, which + # the encoding can detect and handle correctly. + self._ref = InProgressSentinel() + data = ray.worker._dumps_from_client(self._func) + # Check pickled size before sending it to server, which is more + # efficient and can be done synchronously inside remote() call. + check_oversized_function(data, self._name, "remote function", None) + self._ref = ray.worker._put_pickled( + data, client_ref_id=self._client_side_ref.id + ) + + def _prepare_client_task(self) -> ray_client_pb2.ClientTask: + self._ensure_ref() + task = ray_client_pb2.ClientTask() + task.type = ray_client_pb2.ClientTask.FUNCTION + task.name = self._name + task.payload_id = self._ref.id + set_task_options(task, self._options, "baseline_options") + return task + + def _num_returns(self) -> int: + if not self._options: + return None + return self._options.get("num_returns") + + +class ClientActorClass(ClientStub): + """A stub created on the Ray Client to represent an actor class. + + It is wrapped by ray.remote and can be executed on the cluster. + + Args: + actor_cls: The actual class to execute remotely + _name: The original name of the class + _ref: The ClientObjectRef of the pickled `actor_cls` + """ + + def __init__(self, actor_cls, options=None): + self.actor_cls = actor_cls + self._lock = threading.Lock() + self._name = actor_cls.__name__ + self._init_signature = inspect.Signature( + parameters=extract_signature(actor_cls.__init__, ignore_first=True) + ) + self._ref = None + self._client_side_ref = ClientSideRefID.generate_id() + self._options = validate_options(options) + + def __call__(self, *args, **kwargs): + raise TypeError( + "Remote actor cannot be instantiated directly. " + f"Use {self._name}.remote() instead" + ) + + def _ensure_ref(self): + with self._lock: + if self._ref is None: + # As before, set the state of the reference to be an + # in-progress self reference value, which + # the encoding can detect and handle correctly. + self._ref = InProgressSentinel() + data = ray.worker._dumps_from_client(self.actor_cls) + # Check pickled size before sending it to server, which is more + # efficient and can be done synchronously inside remote() call. + check_oversized_function(data, self._name, "actor", None) + self._ref = ray.worker._put_pickled( + data, client_ref_id=self._client_side_ref.id + ) + + def remote(self, *args, **kwargs) -> "ClientActorHandle": + self._init_signature.bind(*args, **kwargs) + # Actually instantiate the actor + futures = ray.call_remote(self, *args, **kwargs) + assert len(futures) == 1 + return ClientActorHandle(ClientActorRef(futures[0]), actor_class=self) + + def options(self, **kwargs): + return ActorOptionWrapper(self, kwargs) + + def _remote(self, args=None, kwargs=None, **option_args): + if args is None: + args = [] + if kwargs is None: + kwargs = {} + return self.options(**option_args).remote(*args, **kwargs) + + def __repr__(self): + return "ClientActorClass(%s, %s)" % (self._name, self._ref) + + def __getattr__(self, key): + if key not in self.__dict__: + raise AttributeError("Not a class attribute") + raise NotImplementedError("static methods") + + def _prepare_client_task(self) -> ray_client_pb2.ClientTask: + self._ensure_ref() + task = ray_client_pb2.ClientTask() + task.type = ray_client_pb2.ClientTask.ACTOR + task.name = self._name + task.payload_id = self._ref.id + set_task_options(task, self._options, "baseline_options") + return task + + @staticmethod + def _num_returns() -> int: + return 1 + + +class ClientActorHandle(ClientStub): + """Client-side stub for instantiated actor. + + A stub created on the Ray Client to represent a remote actor that + has been started on the cluster. This class is allowed to be passed + around between remote functions. + + Args: + actor_ref: A reference to the running actor given to the client. This + is a serialized version of the actual handle as an opaque token. + """ + + def __init__( + self, + actor_ref: ClientActorRef, + actor_class: Optional[ClientActorClass] = None, + ): + self.actor_ref = actor_ref + self._dir: Optional[List[str]] = None + if actor_class is not None: + self._method_num_returns = {} + self._method_signatures = {} + for method_name, method_obj in inspect.getmembers( + actor_class.actor_cls, is_function_or_method + ): + self._method_num_returns[method_name] = getattr( + method_obj, "__ray_num_returns__", None + ) + self._method_signatures[method_name] = inspect.Signature( + parameters=extract_signature( + method_obj, + ignore_first=( + not ( + is_class_method(method_obj) + or is_static_method(actor_class.actor_cls, method_name) + ) + ), + ) + ) + else: + self._method_num_returns = None + self._method_signatures = None + + def __dir__(self) -> List[str]: + if self._method_num_returns is not None: + return self._method_num_returns.keys() + if ray.is_connected(): + self._init_class_info() + return self._method_num_returns.keys() + return super().__dir__() + + # For compatibility with core worker ActorHandle._actor_id which returns + # ActorID + @property + def _actor_id(self) -> ClientActorRef: + return self.actor_ref + + def __hash__(self) -> int: + return hash(self._actor_id) + + def __eq__(self, __value) -> bool: + return hash(self) == hash(__value) + + def __getattr__(self, key): + if key == "_method_num_returns": + # We need to explicitly handle this value since it is used below, + # otherwise we may end up infinitely recursing when deserializing. + # This can happen after unpickling an object but before + # _method_num_returns is correctly populated. + raise AttributeError(f"ClientActorRef has no attribute '{key}'") + + if self._method_num_returns is None: + self._init_class_info() + if key not in self._method_signatures: + raise AttributeError(f"ClientActorRef has no attribute '{key}'") + return ClientRemoteMethod( + self, + key, + self._method_num_returns.get(key), + self._method_signatures.get(key), + ) + + def __repr__(self): + return "ClientActorHandle(%s)" % (self.actor_ref.id.hex()) + + def _init_class_info(self): + # TODO: fetch Ray method decorators + @ray.remote(num_cpus=0) + def get_class_info(x): + return x._ray_method_num_returns, x._ray_method_signatures + + self._method_num_returns, method_parameters = ray.get( + get_class_info.remote(self) + ) + + self._method_signatures = {} + for method, parameters in method_parameters.items(): + self._method_signatures[method] = inspect.Signature(parameters=parameters) + + +class ClientRemoteMethod(ClientStub): + """A stub for a method on a remote actor. + + Can be annotated with execution options. + + Args: + actor_handle: A reference to the ClientActorHandle that generated + this method and will have this method called upon it. + method_name: The name of this method + """ + + def __init__( + self, + actor_handle: ClientActorHandle, + method_name: str, + num_returns: int, + signature: inspect.Signature, + ): + self._actor_handle = actor_handle + self._method_name = method_name + self._method_num_returns = num_returns + self._signature = signature + + def __call__(self, *args, **kwargs): + raise TypeError( + "Actor methods cannot be called directly. Instead " + f"of running 'object.{self._method_name}()', try " + f"'object.{self._method_name}.remote()'." + ) + + def remote(self, *args, **kwargs): + self._signature.bind(*args, **kwargs) + return return_refs(ray.call_remote(self, *args, **kwargs)) + + def __repr__(self): + return "ClientRemoteMethod(%s, %s, %s)" % ( + self._method_name, + self._actor_handle, + self._method_num_returns, + ) + + def options(self, **kwargs): + return OptionWrapper(self, kwargs) + + def _remote(self, args=None, kwargs=None, **option_args): + if args is None: + args = [] + if kwargs is None: + kwargs = {} + return self.options(**option_args).remote(*args, **kwargs) + + def _prepare_client_task(self) -> ray_client_pb2.ClientTask: + task = ray_client_pb2.ClientTask() + task.type = ray_client_pb2.ClientTask.METHOD + task.name = self._method_name + task.payload_id = self._actor_handle.actor_ref.id + return task + + def _num_returns(self) -> int: + return self._method_num_returns + + +class OptionWrapper: + def __init__(self, stub: ClientStub, options: Optional[Dict[str, Any]]): + self._remote_stub = stub + self._options = validate_options(options) + + def remote(self, *args, **kwargs): + self._remote_stub._signature.bind(*args, **kwargs) + return return_refs(ray.call_remote(self, *args, **kwargs)) + + def __getattr__(self, key): + return getattr(self._remote_stub, key) + + def _prepare_client_task(self): + task = self._remote_stub._prepare_client_task() + set_task_options(task, self._options) + return task + + def _num_returns(self) -> int: + if self._options: + num = self._options.get("num_returns") + if num is not None: + return num + return self._remote_stub._num_returns() + + +class ActorOptionWrapper(OptionWrapper): + def remote(self, *args, **kwargs): + self._remote_stub._init_signature.bind(*args, **kwargs) + futures = ray.call_remote(self, *args, **kwargs) + assert len(futures) == 1 + actor_class = None + if isinstance(self._remote_stub, ClientActorClass): + actor_class = self._remote_stub + return ClientActorHandle(ClientActorRef(futures[0]), actor_class=actor_class) + + +def set_task_options( + task: ray_client_pb2.ClientTask, + options: Optional[Dict[str, Any]], + field: str = "options", +) -> None: + if options is None: + task.ClearField(field) + return + + getattr(task, field).pickled_options = pickle.dumps(options) + + +def return_refs( + futures: List[Future], +) -> Union[None, ClientObjectRef, List[ClientObjectRef]]: + if not futures: + return None + if len(futures) == 1: + return ClientObjectRef(futures[0]) + return [ClientObjectRef(fut) for fut in futures] + + +class InProgressSentinel: + def __repr__(self) -> str: + return self.__class__.__name__ + + +class ClientSideRefID: + """An ID generated by the client for objects not yet given an ObjectRef""" + + def __init__(self, id: bytes): + assert len(id) != 0 + self.id = id + + @staticmethod + def generate_id() -> "ClientSideRefID": + tid = uuid.uuid4() + return ClientSideRefID(b"\xcc" + tid.bytes) + + +def remote_decorator(options: Optional[Dict[str, Any]]): + def decorator(function_or_class) -> ClientStub: + if inspect.isfunction(function_or_class) or is_cython(function_or_class): + return ClientRemoteFunc(function_or_class, options=options) + elif inspect.isclass(function_or_class): + return ClientActorClass(function_or_class, options=options) + else: + raise TypeError( + "The @ray.remote decorator must be applied to " + "either a function or to a class." + ) + + return decorator + + +@dataclass +class ClientServerHandle: + """Holds the handles to the registered gRPC servicers and their server.""" + + task_servicer: ray_client_pb2_grpc.RayletDriverServicer + data_servicer: ray_client_pb2_grpc.RayletDataStreamerServicer + logs_servicer: ray_client_pb2_grpc.RayletLogStreamerServicer + grpc_server: grpc.Server + + def stop(self, grace: int) -> None: + # The data servicer might be sleeping while waiting for clients to + # reconnect. Signal that they no longer have to sleep and can exit + # immediately, since the RPC server is stopped. + self.grpc_server.stop(grace) + self.data_servicer.stopped.set() + + # Add a hook for all the cases that previously + # expected simply a gRPC server + def __getattr__(self, attr): + return getattr(self.grpc_server, attr) + + +def _get_client_id_from_context(context: Any) -> str: + """ + Get `client_id` from gRPC metadata. If the `client_id` is not present, + this function logs an error and sets the status_code. + """ + metadata = dict(context.invocation_metadata()) + client_id = metadata.get("client_id") or "" + if client_id == "": + logger.error("Client connecting with no client_id") + context.set_code(grpc.StatusCode.FAILED_PRECONDITION) + return client_id + + +def _propagate_error_in_context(e: Exception, context: Any) -> bool: + """ + Encode an error into the context of an RPC response. Returns True + if the error can be recovered from, false otherwise + """ + try: + if isinstance(e, grpc.RpcError): + # RPC error, propagate directly by copying details into context + context.set_code(e.code()) + context.set_details(e.details()) + return e.code() not in GRPC_UNRECOVERABLE_ERRORS + except Exception: + # Extra precaution -- if encoding the RPC directly fails fallback + # to treating it as a regular error + pass + context.set_code(grpc.StatusCode.FAILED_PRECONDITION) + context.set_details(str(e)) + return False + + +def _id_is_newer(id1: int, id2: int) -> bool: + """ + We should only replace cache entries with the responses for newer IDs. + Most of the time newer IDs will be the ones with higher value, except when + the req_id counter rolls over. We check for this case by checking the + distance between the two IDs. If the distance is significant, then it's + likely that the req_id counter rolled over, and the smaller id should + still be used to replace the one in cache. + """ + diff = abs(id2 - id1) + # Int32 max is also the maximum number of simultaneous in-flight requests. + if diff > (INT32_MAX // 2): + # Rollover likely occurred. In this case the smaller ID is newer + return id1 < id2 + return id1 > id2 + + +class ResponseCache: + """ + Cache for blocking method calls. Needed to prevent retried requests from + being applied multiple times on the server, for example when the client + disconnects. This is used to cache requests/responses sent through + unary-unary RPCs to the RayletServicer. + + Note that no clean up logic is used, the last response for each thread + will always be remembered, so at most the cache will hold N entries, + where N is the number of threads on the client side. This relies on the + assumption that a thread will not make a new blocking request until it has + received a response for a previous one, at which point it's safe to + overwrite the old response. + + The high level logic is: + + 1. Before making a call, check the cache for the current thread. + 2. If present in the cache, check the request id of the cached + response. + a. If it matches the current request_id, then the request has been + received before and we shouldn't re-attempt the logic. Wait for + the response to become available in the cache, and then return it + b. If it doesn't match, then this is a new request and we can + proceed with calling the real stub. While the response is still + being generated, temporarily keep (req_id, None) in the cache. + Once the call is finished, update the cache entry with the + new (req_id, response) pair. Notify other threads that may + have been waiting for the response to be prepared. + """ + + def __init__(self): + self.cv = threading.Condition() + self.cache: Dict[int, Tuple[int, Any]] = {} + + def check_cache(self, thread_id: int, request_id: int) -> Optional[Any]: + """ + Check the cache for a given thread, and see if the entry in the cache + matches the current request_id. Returns None if the request_id has + not been seen yet, otherwise returns the cached result. + + Throws an error if the placeholder in the cache doesn't match the + request_id -- this means that a new request evicted the old value in + the cache, and that the RPC for `request_id` is redundant and the + result can be discarded, i.e.: + + 1. Request A is sent (A1) + 2. Channel disconnects + 3. Request A is resent (A2) + 4. A1 is received + 5. A2 is received, waits for A1 to finish + 6. A1 finishes and is sent back to client + 7. Request B is sent + 8. Request B overwrites cache entry + 9. A2 wakes up extremely late, but cache is now invalid + + In practice this is VERY unlikely to happen, but the error can at + least serve as a sanity check or catch invalid request id's. + """ + with self.cv: + if thread_id in self.cache: + cached_request_id, cached_resp = self.cache[thread_id] + if cached_request_id == request_id: + while cached_resp is None: + # The call was started, but the response hasn't yet + # been added to the cache. Let go of the lock and + # wait until the response is ready. + self.cv.wait() + cached_request_id, cached_resp = self.cache[thread_id] + if cached_request_id != request_id: + raise RuntimeError( + "Cached response doesn't match the id of the " + "original request. This might happen if this " + "request was received out of order. The " + "result of the caller is no longer needed. " + f"({request_id} != {cached_request_id})" + ) + return cached_resp + if not _id_is_newer(request_id, cached_request_id): + raise RuntimeError( + "Attempting to replace newer cache entry with older " + "one. This might happen if this request was received " + "out of order. The result of the caller is no " + f"longer needed. ({request_id} != {cached_request_id}" + ) + self.cache[thread_id] = (request_id, None) + return None + + def update_cache(self, thread_id: int, request_id: int, response: Any) -> None: + """ + Inserts `response` into the cache for `request_id`. + """ + with self.cv: + cached_request_id, cached_resp = self.cache[thread_id] + if cached_request_id != request_id or cached_resp is not None: + # The cache was overwritten by a newer requester between + # our call to check_cache and our call to update it. + # This can't happen if the assumption that the cached requests + # are all blocking on the client side, so if you encounter + # this, check if any async requests are being cached. + raise RuntimeError( + "Attempting to update the cache, but placeholder's " + "do not match the current request_id. This might happen " + "if this request was received out of order. The result " + f"of the caller is no longer needed. ({request_id} != " + f"{cached_request_id})" + ) + self.cache[thread_id] = (request_id, response) + self.cv.notify_all() + + +class OrderedResponseCache: + """ + Cache for streaming RPCs, i.e. the DataServicer. Relies on explicit + ack's from the client to determine when it can clean up cache entries. + """ + + def __init__(self): + self.last_received = 0 + self.cv = threading.Condition() + self.cache: Dict[int, Any] = OrderedDict() + + def check_cache(self, req_id: int) -> Optional[Any]: + """ + Check the cache for a given thread, and see if the entry in the cache + matches the current request_id. Returns None if the request_id has + not been seen yet, otherwise returns the cached result. + """ + with self.cv: + if _id_is_newer(self.last_received, req_id) or self.last_received == req_id: + # Request is for an id that has already been cleared from + # cache/acknowledged. + raise RuntimeError( + "Attempting to accesss a cache entry that has already " + "cleaned up. The client has already acknowledged " + f"receiving this response. ({req_id}, " + f"{self.last_received})" + ) + if req_id in self.cache: + cached_resp = self.cache[req_id] + while cached_resp is None: + # The call was started, but the response hasn't yet been + # added to the cache. Let go of the lock and wait until + # the response is ready + self.cv.wait() + if req_id not in self.cache: + raise RuntimeError( + "Cache entry was removed. This likely means that " + "the result of this call is no longer needed." + ) + cached_resp = self.cache[req_id] + return cached_resp + self.cache[req_id] = None + return None + + def update_cache(self, req_id: int, resp: Any) -> None: + """ + Inserts `response` into the cache for `request_id`. + """ + with self.cv: + self.cv.notify_all() + if req_id not in self.cache: + raise RuntimeError( + "Attempting to update the cache, but placeholder is " + "missing. This might happen on a redundant call to " + f"update_cache. ({req_id})" + ) + self.cache[req_id] = resp + + def invalidate(self, e: Exception) -> bool: + """ + Invalidate any partially populated cache entries, replacing their + placeholders with the passed in exception. Useful to prevent a thread + from waiting indefinitely on a failed call. + + Returns True if the cache contains an error, False otherwise + """ + with self.cv: + invalid = False + for req_id in self.cache: + if self.cache[req_id] is None: + self.cache[req_id] = e + if isinstance(self.cache[req_id], Exception): + invalid = True + self.cv.notify_all() + return invalid + + def cleanup(self, last_received: int) -> None: + """ + Cleanup all of the cached requests up to last_received. Assumes that + the cache entries were inserted in ascending order. + """ + with self.cv: + if _id_is_newer(last_received, self.last_received): + self.last_received = last_received + to_remove = [] + for req_id in self.cache: + if _id_is_newer(last_received, req_id) or last_received == req_id: + to_remove.append(req_id) + else: + break + for req_id in to_remove: + del self.cache[req_id] + self.cv.notify_all() diff --git a/lib/python3.12/site-packages/ray/util/client/dataclient.py b/lib/python3.12/site-packages/ray/util/client/dataclient.py new file mode 100644 index 0000000000000000000000000000000000000000..6ef6f29c190beee9a7b5f07733339c4acb0ce5bf --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/dataclient.py @@ -0,0 +1,599 @@ +"""This file implements a threaded stream controller to abstract a data stream +back to the ray clientserver. +""" +import logging +import math +import queue +import threading +import warnings +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union + +import grpc + +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc +from ray.util.client.common import ( + INT32_MAX, + OBJECT_TRANSFER_CHUNK_SIZE, + OBJECT_TRANSFER_WARNING_SIZE, +) +from ray.util.debug import log_once + +if TYPE_CHECKING: + from ray.util.client.worker import Worker + +logger = logging.getLogger(__name__) + +ResponseCallable = Callable[[Union[ray_client_pb2.DataResponse, Exception]], None] + +# Send an acknowledge on every 32nd response received +ACKNOWLEDGE_BATCH_SIZE = 32 + + +def chunk_put(req: ray_client_pb2.DataRequest): + """ + Chunks a put request. Doing this lazily is important for large objects, + since taking slices of bytes objects does a copy. This means if we + immediately materialized every chunk of a large object and inserted them + into the result_queue, we would effectively double the memory needed + on the client to handle the put. + """ + # When accessing a protobuf field, deserialization is performed, which will + # generate a copy. So we need to avoid accessing the `data` field multiple + # times in the loop + request_data = req.put.data + total_size = len(request_data) + assert total_size > 0, "Cannot chunk object with missing data" + if total_size >= OBJECT_TRANSFER_WARNING_SIZE and log_once( + "client_object_put_size_warning" + ): + size_gb = total_size / 2**30 + warnings.warn( + "Ray Client is attempting to send a " + f"{size_gb:.2f} GiB object over the network, which may " + "be slow. Consider serializing the object and using a remote " + "URI to transfer via S3 or Google Cloud Storage instead. " + "Documentation for doing this can be found here: " + "https://docs.ray.io/en/latest/handling-dependencies.html#remote-uris", + UserWarning, + ) + total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE) + for chunk_id in range(0, total_chunks): + start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE + end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE) + chunk = ray_client_pb2.PutRequest( + client_ref_id=req.put.client_ref_id, + data=request_data[start:end], + chunk_id=chunk_id, + total_chunks=total_chunks, + total_size=total_size, + owner_id=req.put.owner_id, + ) + yield ray_client_pb2.DataRequest(req_id=req.req_id, put=chunk) + + +def chunk_task(req: ray_client_pb2.DataRequest): + """ + Chunks a client task. Doing this lazily is important with large arguments, + since taking slices of bytes objects does a copy. This means if we + immediately materialized every chunk of a large argument and inserted them + into the result_queue, we would effectively double the memory needed + on the client to handle the task. + """ + # When accessing a protobuf field, deserialization is performed, which will + # generate a copy. So we need to avoid accessing the `data` field multiple + # times in the loop + request_data = req.task.data + total_size = len(request_data) + assert total_size > 0, "Cannot chunk object with missing data" + total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE) + for chunk_id in range(0, total_chunks): + start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE + end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE) + chunk = ray_client_pb2.ClientTask( + type=req.task.type, + name=req.task.name, + payload_id=req.task.payload_id, + client_id=req.task.client_id, + options=req.task.options, + baseline_options=req.task.baseline_options, + namespace=req.task.namespace, + data=request_data[start:end], + chunk_id=chunk_id, + total_chunks=total_chunks, + ) + yield ray_client_pb2.DataRequest(req_id=req.req_id, task=chunk) + + +class ChunkCollector: + """ + This object collects chunks from async get requests via __call__, and + calls the underlying callback when the object is fully received, or if an + exception while retrieving the object occurs. + + This is not used in synchronous gets (synchronous gets interact with the + raylet servicer directly, not through the datapath). + + __call__ returns true once the underlying call back has been called. + """ + + def __init__(self, callback: ResponseCallable, request: ray_client_pb2.DataRequest): + # Bytearray containing data received so far + self.data = bytearray() + # The callback that will be called once all data is received + self.callback = callback + # The id of the last chunk we've received, or -1 if haven't seen any yet + self.last_seen_chunk = -1 + # The GetRequest that initiated the transfer. start_chunk_id will be + # updated as chunks are received to avoid re-requesting chunks that + # we've already received. + self.request = request + + def __call__(self, response: Union[ray_client_pb2.DataResponse, Exception]) -> bool: + if isinstance(response, Exception): + self.callback(response) + return True + get_resp = response.get + if not get_resp.valid: + self.callback(response) + return True + if get_resp.total_size > OBJECT_TRANSFER_WARNING_SIZE and log_once( + "client_object_transfer_size_warning" + ): + size_gb = get_resp.total_size / 2**30 + warnings.warn( + "Ray Client is attempting to retrieve a " + f"{size_gb:.2f} GiB object over the network, which may " + "be slow. Consider serializing the object to a file and " + "using rsync or S3 instead.", + UserWarning, + ) + chunk_data = get_resp.data + chunk_id = get_resp.chunk_id + if chunk_id == self.last_seen_chunk + 1: + self.data.extend(chunk_data) + self.last_seen_chunk = chunk_id + # If we disconnect partway through, restart the get request + # at the first chunk we haven't seen + self.request.get.start_chunk_id = self.last_seen_chunk + 1 + elif chunk_id > self.last_seen_chunk + 1: + # A chunk was skipped. This shouldn't happen in practice since + # grpc guarantees that chunks will arrive in order. + msg = ( + f"Received chunk {chunk_id} when we expected " + f"{self.last_seen_chunk + 1} for request {response.req_id}" + ) + logger.warning(msg) + self.callback(RuntimeError(msg)) + return True + else: + # We received a chunk that've already seen before. Ignore, since + # it should already be appended to self.data. + logger.debug( + f"Received a repeated chunk {chunk_id} " + f"from request {response.req_id}." + ) + + if get_resp.chunk_id == get_resp.total_chunks - 1: + self.callback(self.data) + return True + else: + # Not done yet + return False + + +class DataClient: + def __init__(self, client_worker: "Worker", client_id: str, metadata: list): + """Initializes a thread-safe datapath over a Ray Client gRPC channel. + + Args: + client_worker: The Ray Client worker that manages this client + client_id: the generated ID representing this client + metadata: metadata to pass to gRPC requests + """ + self.client_worker = client_worker + self._client_id = client_id + self._metadata = metadata + self.data_thread = self._start_datathread() + + # Track outstanding requests to resend in case of disconnection + self.outstanding_requests: Dict[int, Any] = OrderedDict() + + # Serialize access to all mutable internal states: self.request_queue, + # self.ready_data, self.asyncio_waiting_data, + # self._in_shutdown, self._req_id, self.outstanding_requests and + # calling self._next_id() + self.lock = threading.Lock() + + # Waiting for response or shutdown. + self.cv = threading.Condition(lock=self.lock) + + self.request_queue = self._create_queue() + self.ready_data: Dict[int, Any] = {} + # NOTE: Dictionary insertion is guaranteed to complete before lookup + # and/or removal because of synchronization via the request_queue. + self.asyncio_waiting_data: Dict[int, ResponseCallable] = {} + self._in_shutdown = False + self._req_id = 0 + self._last_exception = None + self._acknowledge_counter = 0 + + self.data_thread.start() + + # Must hold self.lock when calling this function. + def _next_id(self) -> int: + assert self.lock.locked() + self._req_id += 1 + if self._req_id > INT32_MAX: + self._req_id = 1 + # Responses that aren't tracked (like opportunistic releases) + # have req_id=0, so make sure we never mint such an id. + assert self._req_id != 0 + return self._req_id + + def _start_datathread(self) -> threading.Thread: + return threading.Thread( + target=self._data_main, + name="ray_client_streaming_rpc", + args=(), + daemon=True, + ) + + # A helper that takes requests from queue. If the request wraps a PutRequest, + # lazily chunks and yields the request. Otherwise, yields the request directly. + def _requests(self): + while True: + req = self.request_queue.get() + if req is None: + # Stop when client signals shutdown. + return + req_type = req.WhichOneof("type") + if req_type == "put": + yield from chunk_put(req) + elif req_type == "task": + yield from chunk_task(req) + else: + yield req + + def _data_main(self) -> None: + reconnecting = False + try: + while not self.client_worker._in_shutdown: + stub = ray_client_pb2_grpc.RayletDataStreamerStub( + self.client_worker.channel + ) + metadata = self._metadata + [("reconnecting", str(reconnecting))] + resp_stream = stub.Datapath( + self._requests(), + metadata=metadata, + wait_for_ready=True, + ) + try: + for response in resp_stream: + self._process_response(response) + return + except grpc.RpcError as e: + reconnecting = self._can_reconnect(e) + if not reconnecting: + self._last_exception = e + return + self._reconnect_channel() + except Exception as e: + self._last_exception = e + finally: + logger.debug("Shutting down data channel.") + self._shutdown() + + def _process_response(self, response: Any) -> None: + """ + Process responses from the data servicer. + """ + if response.req_id == 0: + # This is not being waited for. + logger.debug(f"Got unawaited response {response}") + return + if response.req_id in self.asyncio_waiting_data: + can_remove = True + try: + callback = self.asyncio_waiting_data[response.req_id] + if isinstance(callback, ChunkCollector): + can_remove = callback(response) + elif callback: + callback(response) + if can_remove: + # NOTE: calling del self.asyncio_waiting_data results + # in the destructor of ClientObjectRef running, which + # calls ReleaseObject(). So self.asyncio_waiting_data + # is accessed without holding self.lock. Holding the + # lock shouldn't be necessary either. + del self.asyncio_waiting_data[response.req_id] + except Exception: + logger.exception("Callback error:") + with self.lock: + # Update outstanding requests + if response.req_id in self.outstanding_requests and can_remove: + del self.outstanding_requests[response.req_id] + # Acknowledge response + self._acknowledge(response.req_id) + else: + with self.lock: + self.ready_data[response.req_id] = response + self.cv.notify_all() + + def _can_reconnect(self, e: grpc.RpcError) -> bool: + """ + Processes RPC errors that occur while reading from data stream. + Returns True if the error can be recovered from, False otherwise. + """ + if not self.client_worker._can_reconnect(e): + logger.error("Unrecoverable error in data channel.") + logger.debug(e) + return False + logger.debug("Recoverable error in data channel.") + logger.debug(e) + return True + + def _shutdown(self) -> None: + """ + Shutdown the data channel + """ + with self.lock: + self._in_shutdown = True + self.cv.notify_all() + + callbacks = self.asyncio_waiting_data.values() + self.asyncio_waiting_data = {} + + if self._last_exception: + # Abort async requests with the error. + err = ConnectionError( + "Failed during this or a previous request. Exception that " + f"broke the connection: {self._last_exception}" + ) + else: + err = ConnectionError( + "Request cannot be fulfilled because the data client has " + "disconnected." + ) + for callback in callbacks: + if callback: + callback(err) + # Since self._in_shutdown is set to True, no new item + # will be added to self.asyncio_waiting_data + + def _acknowledge(self, req_id: int) -> None: + """ + Puts an acknowledge request on the request queue periodically. + Lock should be held before calling this. Used when an async or + blocking response is received. + """ + if not self.client_worker._reconnect_enabled: + # Skip ACKs if reconnect isn't enabled + return + assert self.lock.locked() + self._acknowledge_counter += 1 + if self._acknowledge_counter % ACKNOWLEDGE_BATCH_SIZE == 0: + self.request_queue.put( + ray_client_pb2.DataRequest( + acknowledge=ray_client_pb2.AcknowledgeRequest(req_id=req_id) + ) + ) + + def _reconnect_channel(self) -> None: + """ + Attempts to reconnect the gRPC channel and resend outstanding + requests. First, the server is pinged to see if the current channel + still works. If the ping fails, then the current channel is closed + and replaced with a new one. + + Once a working channel is available, a new request queue is made + and filled with any outstanding requests to be resent to the server. + """ + try: + # Ping the server to see if the current channel is reuseable, for + # example if gRPC reconnected the channel on its own or if the + # RPC error was transient and the channel is still open + ping_succeeded = self.client_worker.ping_server(timeout=5) + except grpc.RpcError: + ping_succeeded = False + + if not ping_succeeded: + # Ping failed, try refreshing the data channel + logger.warning( + "Encountered connection issues in the data channel. " + "Attempting to reconnect." + ) + try: + self.client_worker._connect_channel(reconnecting=True) + except ConnectionError: + logger.warning("Failed to reconnect the data channel") + raise + logger.debug("Reconnection succeeded!") + + # Recreate the request queue, and resend outstanding requests + with self.lock: + self.request_queue = self._create_queue() + for request in self.outstanding_requests.values(): + # Resend outstanding requests + self.request_queue.put(request) + + # Use SimpleQueue to avoid deadlocks when appending to queue from __del__() + @staticmethod + def _create_queue(): + return queue.SimpleQueue() + + def close(self) -> None: + thread = None + with self.lock: + self._in_shutdown = True + # Notify blocking operations to fail. + self.cv.notify_all() + # Add sentinel to terminate streaming RPC. + if self.request_queue is not None: + # Intentional shutdown, tell server it can clean up the + # connection immediately and ignore the reconnect grace period. + cleanup_request = ray_client_pb2.DataRequest( + connection_cleanup=ray_client_pb2.ConnectionCleanupRequest() + ) + self.request_queue.put(cleanup_request) + self.request_queue.put(None) + if self.data_thread is not None: + thread = self.data_thread + # Wait until streaming RPCs are done. + if thread is not None: + thread.join() + + def _blocking_send( + self, req: ray_client_pb2.DataRequest + ) -> ray_client_pb2.DataResponse: + with self.lock: + self._check_shutdown() + req_id = self._next_id() + req.req_id = req_id + self.request_queue.put(req) + self.outstanding_requests[req_id] = req + + self.cv.wait_for(lambda: req_id in self.ready_data or self._in_shutdown) + self._check_shutdown() + + data = self.ready_data[req_id] + del self.ready_data[req_id] + del self.outstanding_requests[req_id] + self._acknowledge(req_id) + + return data + + def _async_send( + self, + req: ray_client_pb2.DataRequest, + callback: Optional[ResponseCallable] = None, + ) -> None: + with self.lock: + self._check_shutdown() + req_id = self._next_id() + req.req_id = req_id + self.asyncio_waiting_data[req_id] = callback + self.outstanding_requests[req_id] = req + self.request_queue.put(req) + + # Must hold self.lock when calling this function. + def _check_shutdown(self): + assert self.lock.locked() + if not self._in_shutdown: + return + + self.lock.release() + + # Do not try disconnect() or throw exceptions in self.data_thread. + # Otherwise deadlock can occur. + if threading.current_thread().ident == self.data_thread.ident: + return + + from ray.util import disconnect + + disconnect() + + self.lock.acquire() + + if self._last_exception is not None: + msg = ( + "Request can't be sent because the Ray client has already " + "been disconnected due to an error. Last exception: " + f"{self._last_exception}" + ) + else: + msg = ( + "Request can't be sent because the Ray client has already " + "been disconnected." + ) + + raise ConnectionError(msg) + + def Init( + self, request: ray_client_pb2.InitRequest, context=None + ) -> ray_client_pb2.InitResponse: + datareq = ray_client_pb2.DataRequest( + init=request, + ) + resp = self._blocking_send(datareq) + return resp.init + + def PrepRuntimeEnv( + self, request: ray_client_pb2.PrepRuntimeEnvRequest, context=None + ) -> ray_client_pb2.PrepRuntimeEnvResponse: + datareq = ray_client_pb2.DataRequest( + prep_runtime_env=request, + ) + resp = self._blocking_send(datareq) + return resp.prep_runtime_env + + def ConnectionInfo(self, context=None) -> ray_client_pb2.ConnectionInfoResponse: + datareq = ray_client_pb2.DataRequest( + connection_info=ray_client_pb2.ConnectionInfoRequest() + ) + resp = self._blocking_send(datareq) + return resp.connection_info + + def GetObject( + self, request: ray_client_pb2.GetRequest, context=None + ) -> ray_client_pb2.GetResponse: + datareq = ray_client_pb2.DataRequest( + get=request, + ) + resp = self._blocking_send(datareq) + return resp.get + + def RegisterGetCallback( + self, request: ray_client_pb2.GetRequest, callback: ResponseCallable + ) -> None: + if len(request.ids) != 1: + raise ValueError( + "RegisterGetCallback() must have exactly 1 Object ID. " + f"Actual: {request}" + ) + datareq = ray_client_pb2.DataRequest( + get=request, + ) + collector = ChunkCollector(callback=callback, request=datareq) + self._async_send(datareq, collector) + + # TODO: convert PutObject to async + def PutObject( + self, request: ray_client_pb2.PutRequest, context=None + ) -> ray_client_pb2.PutResponse: + datareq = ray_client_pb2.DataRequest( + put=request, + ) + resp = self._blocking_send(datareq) + return resp.put + + def ReleaseObject( + self, request: ray_client_pb2.ReleaseRequest, context=None + ) -> None: + datareq = ray_client_pb2.DataRequest( + release=request, + ) + self._async_send(datareq) + + def Schedule(self, request: ray_client_pb2.ClientTask, callback: ResponseCallable): + datareq = ray_client_pb2.DataRequest(task=request) + self._async_send(datareq, callback) + + def Terminate( + self, request: ray_client_pb2.TerminateRequest + ) -> ray_client_pb2.TerminateResponse: + req = ray_client_pb2.DataRequest( + terminate=request, + ) + resp = self._blocking_send(req) + return resp.terminate + + def ListNamedActors( + self, request: ray_client_pb2.ClientListNamedActorsRequest + ) -> ray_client_pb2.ClientListNamedActorsResponse: + req = ray_client_pb2.DataRequest( + list_named_actors=request, + ) + resp = self._blocking_send(req) + return resp.list_named_actors diff --git a/lib/python3.12/site-packages/ray/util/client/logsclient.py b/lib/python3.12/site-packages/ray/util/client/logsclient.py new file mode 100644 index 0000000000000000000000000000000000000000..34ad3f9f6ce902ce3038564b3063aac891e8e384 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/logsclient.py @@ -0,0 +1,135 @@ +"""This file implements a threaded stream controller to return logs back from +the ray clientserver. +""" +import logging +import queue +import sys +import threading +import time +from typing import TYPE_CHECKING + +import grpc + +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc +from ray.util.debug import log_once + +if TYPE_CHECKING: + from ray.util.client.worker import Worker + +logger = logging.getLogger(__name__) +# TODO(barakmich): Running a logger in a logger causes loopback. +# The client logger need its own root -- possibly this one. +# For the moment, let's just not propogate beyond this point. +logger.propagate = False + + +class LogstreamClient: + def __init__(self, client_worker: "Worker", metadata: list): + """Initializes a thread-safe log stream over a Ray Client gRPC channel. + + Args: + client_worker: The Ray Client worker that manages this client + metadata: metadata to pass to gRPC requests + """ + self.client_worker = client_worker + self._metadata = metadata + self.request_queue = queue.Queue() + self.log_thread = self._start_logthread() + self.log_thread.start() + self.last_req = None + + def _start_logthread(self) -> threading.Thread: + return threading.Thread(target=self._log_main, args=(), daemon=True) + + def _log_main(self) -> None: + reconnecting = False + while not self.client_worker._in_shutdown: + if reconnecting: + # Refresh queue and retry last request + self.request_queue = queue.Queue() + if self.last_req: + self.request_queue.put(self.last_req) + stub = ray_client_pb2_grpc.RayletLogStreamerStub(self.client_worker.channel) + try: + log_stream = stub.Logstream( + iter(self.request_queue.get, None), metadata=self._metadata + ) + except ValueError: + # Trying to use the stub on a cancelled channel will raise + # ValueError. This should only happen when the data client + # is attempting to reset the connection -- sleep and try + # again. + time.sleep(0.5) + continue + try: + for record in log_stream: + if record.level < 0: + self.stdstream(level=record.level, msg=record.msg) + self.log(level=record.level, msg=record.msg) + return + except grpc.RpcError as e: + reconnecting = self._process_rpc_error(e) + if not reconnecting: + return + + def _process_rpc_error(self, e: grpc.RpcError) -> bool: + """ + Processes RPC errors that occur while reading from data stream. + Returns True if the error can be recovered from, False otherwise. + """ + if self.client_worker._can_reconnect(e): + if log_once("lost_reconnect_logs"): + logger.warning( + "Log channel is reconnecting. Logs produced while " + "the connection was down can be found on the head " + "node of the cluster in " + "`ray_client_server_[port].out`" + ) + logger.debug("Log channel dropped, retrying.") + time.sleep(0.5) + return True + logger.debug("Shutting down log channel.") + if not self.client_worker._in_shutdown: + logger.exception("Unexpected exception:") + return False + + def log(self, level: int, msg: str): + """Log the message from the log stream. + By default, calls logger.log but this can be overridden. + + Args: + level: The loglevel of the received log message + msg: The content of the message + """ + logger.log(level=level, msg=msg) + + def stdstream(self, level: int, msg: str): + """Log the stdout/stderr entry from the log stream. + By default, calls print but this can be overridden. + + Args: + level: The loglevel of the received log message + msg: The content of the message + """ + print_file = sys.stderr if level == -2 else sys.stdout + print(msg, file=print_file, end="") + + def set_logstream_level(self, level: int): + logger.setLevel(level) + req = ray_client_pb2.LogSettingsRequest() + req.enabled = True + req.loglevel = level + self.request_queue.put(req) + self.last_req = req + + def close(self) -> None: + self.request_queue.put(None) + if self.log_thread is not None: + self.log_thread.join() + + def disable_logs(self) -> None: + req = ray_client_pb2.LogSettingsRequest() + req.enabled = False + self.request_queue.put(req) + self.last_req = req diff --git a/lib/python3.12/site-packages/ray/util/client/options.py b/lib/python3.12/site-packages/ray/util/client/options.py new file mode 100644 index 0000000000000000000000000000000000000000..bd0946fa1975d543d7a10dd6fa7b353c271c9417 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/options.py @@ -0,0 +1,45 @@ +from typing import Any, Dict, Optional + +from ray._common import ray_option_utils +from ray.util.placement_group import PlacementGroup, check_placement_group_index +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + +def validate_options(kwargs_dict: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]: + if kwargs_dict is None: + return None + if len(kwargs_dict) == 0: + return None + + out = {} + for k, v in kwargs_dict.items(): + if k not in ray_option_utils.valid_options: + raise ValueError( + f"Invalid option keyword: '{k}'. " + f"{ray_option_utils.remote_args_error_string}" + ) + ray_option_utils.valid_options[k].validate(k, v) + out[k] = v + + # Validate placement setting similar to the logic in ray/actor.py and + # ray/remote_function.py. The difference is that when + # placement_group = default and placement_group_capture_child_tasks + # specified, placement group cannot be resolved at client. So this check + # skips this case and relies on server to enforce any condition. + bundle_index = out.get("placement_group_bundle_index", None) + pg = out.get("placement_group", None) + scheduling_strategy = out.get("scheduling_strategy", None) + if isinstance(scheduling_strategy, PlacementGroupSchedulingStrategy): + pg = scheduling_strategy.placement_group + bundle_index = scheduling_strategy.placement_group_bundle_index + if bundle_index is not None: + if pg is None: + pg = PlacementGroup.empty() + if pg == "default" and ( + out.get("placement_group_capture_child_tasks", None) is None + ): + pg = PlacementGroup.empty() + if isinstance(pg, PlacementGroup): + check_placement_group_index(pg, bundle_index) + + return out diff --git a/lib/python3.12/site-packages/ray/util/client/ray_client_helpers.py b/lib/python3.12/site-packages/ray/util/client/ray_client_helpers.py new file mode 100644 index 0000000000000000000000000000000000000000..1554bd5e1c23b321fb03552bc10169a6a206c075 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/ray_client_helpers.py @@ -0,0 +1,82 @@ +import time +from contextlib import contextmanager +from typing import Any, Dict + +import ray as real_ray +import ray.util.client.server.server as ray_client_server +from ray._private.client_mode_hook import disable_client_hook +from ray.job_config import JobConfig +from ray.util.client import ray + + +@contextmanager +def ray_start_client_server(metadata=None, ray_connect_handler=None, **kwargs): + with ray_start_client_server_pair( + metadata=metadata, ray_connect_handler=ray_connect_handler, **kwargs + ) as pair: + client, server = pair + yield client + + +@contextmanager +def ray_start_client_server_for_address(address): + """ + Starts a Ray client server that initializes drivers at the specified address. + """ + + def connect_handler( + job_config: JobConfig = None, **ray_init_kwargs: Dict[str, Any] + ): + import ray + + with disable_client_hook(): + if not ray.is_initialized(): + return ray.init(address, job_config=job_config, **ray_init_kwargs) + + with ray_start_client_server(ray_connect_handler=connect_handler) as ray: + yield ray + + +@contextmanager +def ray_start_client_server_pair(metadata=None, ray_connect_handler=None, **kwargs): + ray._inside_client_test = True + with disable_client_hook(): + assert not ray.is_initialized() + server = ray_client_server.serve( + "127.0.0.1", 50051, ray_connect_handler=ray_connect_handler + ) + ray.connect("127.0.0.1:50051", metadata=metadata, **kwargs) + try: + yield ray, server + finally: + ray._inside_client_test = False + ray.disconnect() + server.stop(0) + del server + start = time.monotonic() + with disable_client_hook(): + while ray.is_initialized(): + time.sleep(1) + if time.monotonic() - start > 30: + raise RuntimeError("Failed to terminate Ray") + # Allow windows to close processes before moving on + time.sleep(3) + + +@contextmanager +def ray_start_cluster_client_server_pair(address): + ray._inside_client_test = True + + def ray_connect_handler(job_config=None, **ray_init_kwargs): + real_ray.init(address=address) + + server = ray_client_server.serve( + "127.0.0.1", 50051, ray_connect_handler=ray_connect_handler + ) + ray.connect("127.0.0.1:50051") + try: + yield ray, server + finally: + ray._inside_client_test = False + ray.disconnect() + server.stop(0) diff --git a/lib/python3.12/site-packages/ray/util/client/runtime_context.py b/lib/python3.12/site-packages/ray/util/client/runtime_context.py new file mode 100644 index 0000000000000000000000000000000000000000..ea28055361d8b3e8ef8eac12b5aaee5413eb2671 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/runtime_context.py @@ -0,0 +1,65 @@ +from types import SimpleNamespace +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from ray import JobID, NodeID + from ray.runtime_context import RuntimeContext + + +class _ClientWorkerPropertyAPI: + """Emulates the properties of the ray._private.worker object for the client""" + + def __init__(self, worker): + assert worker is not None + self.worker = worker + + def build_runtime_context(self) -> "RuntimeContext": + """Creates a RuntimeContext backed by the properites of this API""" + # Defer the import of RuntimeContext until needed to avoid cycles + from ray.runtime_context import RuntimeContext + + return RuntimeContext(self) + + def _fetch_runtime_context(self): + import ray.core.generated.ray_client_pb2 as ray_client_pb2 + + return self.worker.get_cluster_info( + ray_client_pb2.ClusterInfoType.RUNTIME_CONTEXT + ) + + @property + def mode(self): + from ray._private.worker import SCRIPT_MODE + + return SCRIPT_MODE + + @property + def current_job_id(self) -> "JobID": + from ray import JobID + + return JobID(self._fetch_runtime_context().job_id) + + @property + def current_node_id(self) -> "NodeID": + from ray import NodeID + + return NodeID(self._fetch_runtime_context().node_id) + + @property + def namespace(self) -> str: + return self._fetch_runtime_context().namespace + + @property + def should_capture_child_tasks_in_placement_group(self) -> bool: + return self._fetch_runtime_context().capture_client_tasks + + @property + def runtime_env(self) -> str: + return self._fetch_runtime_context().runtime_env + + def check_connected(self) -> bool: + return self.worker.ping_server() + + @property + def gcs_client(self) -> str: + return SimpleNamespace(address=self._fetch_runtime_context().gcs_address) diff --git a/lib/python3.12/site-packages/ray/util/client/server/__init__.py b/lib/python3.12/site-packages/ray/util/client/server/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..37c7767bb0cf1430118b672b9ace20dac6eded17 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/server/__init__.py @@ -0,0 +1 @@ +from ray.util.client.server.server import serve # noqa diff --git a/lib/python3.12/site-packages/ray/util/client/server/__main__.py b/lib/python3.12/site-packages/ray/util/client/server/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..b582d767cd21e8778ac909d7a0ab8eb8c4f53450 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/server/__main__.py @@ -0,0 +1,4 @@ +if __name__ == "__main__": + from ray.util.client.server.server import main + + main() diff --git a/lib/python3.12/site-packages/ray/util/client/server/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5cc743f93d2e426498f47f77eb8a2bb192b42cdf Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/server/__pycache__/__main__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/__main__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9dfb21f87c4d9660939637e35496b093f62b4aa3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/__main__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/server/__pycache__/dataservicer.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/dataservicer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f2215e047c9986248649a2ebec0e2544621d1ab Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/dataservicer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/server/__pycache__/logservicer.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/logservicer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fd90d980a1a248de5db466f0097cb92116ad794f Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/logservicer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/server/__pycache__/proxier.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/proxier.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fd9c4f575e2726589c5310d57cc9bcb577d98cb Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/proxier.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f2b0f2b29d4a95208a650d50dbb87d371ae863c Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server_pickler.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server_pickler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..de64a53e2e37efb40c9e27866ae66b640ed3e6be Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server_pickler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server_stubs.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server_stubs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dcc66333ecbe710147d126f3329151b00ac2f8fb Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/client/server/__pycache__/server_stubs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/client/server/dataservicer.py b/lib/python3.12/site-packages/ray/util/client/server/dataservicer.py new file mode 100644 index 0000000000000000000000000000000000000000..0e9363ea36404eb35c6ae382b43e1ff2278f812c --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/server/dataservicer.py @@ -0,0 +1,416 @@ +import logging +import sys +import time +from collections import defaultdict +from queue import Queue +from threading import Event, Lock, Thread +from typing import TYPE_CHECKING, Any, Dict, Iterator, Union + +import grpc + +import ray +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc +from ray._private.client_mode_hook import disable_client_hook +from ray.util.client.common import ( + CLIENT_SERVER_MAX_THREADS, + OrderedResponseCache, + _propagate_error_in_context, +) +from ray.util.client.server.server_pickler import loads_from_client +from ray.util.debug import log_once + +if TYPE_CHECKING: + from ray.util.client.server.server import RayletServicer + +logger = logging.getLogger(__name__) + +QUEUE_JOIN_SECONDS = 10 + + +def _get_reconnecting_from_context(context: Any) -> bool: + """ + Get `reconnecting` from gRPC metadata, or False if missing. + """ + metadata = dict(context.invocation_metadata()) + val = metadata.get("reconnecting") + if val is None or val not in ("True", "False"): + logger.error( + f'Client connecting with invalid value for "reconnecting": {val}, ' + "This may be because you have a mismatched client and server " + "version." + ) + return False + return val == "True" + + +def _should_cache(req: ray_client_pb2.DataRequest) -> bool: + """ + Returns True if the response should to the given request should be cached, + false otherwise. At the moment the only requests we do not cache are: + - asynchronous gets: These arrive out of order. Skipping caching here + is fine, since repeating an async get is idempotent + - acks: Repeating acks is idempotent + - clean up requests: Also idempotent, and client has likely already + wrapped up the data connection by this point. + - puts: We should only cache when we receive the final chunk, since + any earlier chunks won't generate a response + - tasks: We should only cache when we receive the final chunk, + since any earlier chunks won't generate a response + """ + req_type = req.WhichOneof("type") + if req_type == "get" and req.get.asynchronous: + return False + if req_type == "put": + return req.put.chunk_id == req.put.total_chunks - 1 + if req_type == "task": + return req.task.chunk_id == req.task.total_chunks - 1 + return req_type not in ("acknowledge", "connection_cleanup") + + +def fill_queue( + grpc_input_generator: Iterator[ray_client_pb2.DataRequest], + output_queue: "Queue[Union[ray_client_pb2.DataRequest, ray_client_pb2.DataResponse]]", # noqa: E501 +) -> None: + """ + Pushes incoming requests to a shared output_queue. + """ + try: + for req in grpc_input_generator: + output_queue.put(req) + except grpc.RpcError as e: + logger.debug( + "closing dataservicer reader thread " + f"grpc error reading request_iterator: {e}" + ) + finally: + # Set the sentinel value for the output_queue + output_queue.put(None) + + +class ChunkCollector: + """ + Helper class for collecting chunks from PutObject or ClientTask messages + """ + + def __init__(self): + self.curr_req_id = None + self.last_seen_chunk_id = -1 + self.data = bytearray() + + def add_chunk( + self, + req: ray_client_pb2.DataRequest, + chunk: Union[ray_client_pb2.PutRequest, ray_client_pb2.ClientTask], + ): + if self.curr_req_id is not None and self.curr_req_id != req.req_id: + raise RuntimeError( + "Expected to receive a chunk from request with id " + f"{self.curr_req_id}, but found {req.req_id} instead." + ) + self.curr_req_id = req.req_id + next_chunk = self.last_seen_chunk_id + 1 + if chunk.chunk_id < next_chunk: + # Repeated chunk, ignore + return + if chunk.chunk_id > next_chunk: + raise RuntimeError( + f"A chunk {chunk.chunk_id} of request {req.req_id} was " + "received out of order." + ) + elif chunk.chunk_id == self.last_seen_chunk_id + 1: + self.data.extend(chunk.data) + self.last_seen_chunk_id = chunk.chunk_id + return chunk.chunk_id + 1 == chunk.total_chunks + + def reset(self): + self.curr_req_id = None + self.last_seen_chunk_id = -1 + self.data = bytearray() + + +class DataServicer(ray_client_pb2_grpc.RayletDataStreamerServicer): + def __init__(self, basic_service: "RayletServicer"): + self.basic_service = basic_service + self.clients_lock = Lock() + self.num_clients = 0 # guarded by self.clients_lock + # dictionary mapping client_id's to the last time they connected + self.client_last_seen: Dict[str, float] = {} + # dictionary mapping client_id's to their reconnect grace periods + self.reconnect_grace_periods: Dict[str, float] = {} + # dictionary mapping client_id's to their response cache + self.response_caches: Dict[str, OrderedResponseCache] = defaultdict( + OrderedResponseCache + ) + # stopped event, useful for signals that the server is shut down + self.stopped = Event() + # Helper for collecting chunks from PutObject calls. Assumes that + # that put requests from different objects aren't interleaved. + self.put_request_chunk_collector = ChunkCollector() + # Helper for collecting chunks from ClientTask calls. Assumes that + # schedule requests from different remote calls aren't interleaved. + self.client_task_chunk_collector = ChunkCollector() + + def Datapath(self, request_iterator, context): + start_time = time.time() + # set to True if client shuts down gracefully + cleanup_requested = False + metadata = dict(context.invocation_metadata()) + client_id = metadata.get("client_id") + if client_id is None: + logger.error("Client connecting with no client_id") + return + logger.debug(f"New data connection from client {client_id}: ") + accepted_connection = self._init(client_id, context, start_time) + response_cache = self.response_caches[client_id] + # Set to False if client requests a reconnect grace period of 0 + reconnect_enabled = True + if not accepted_connection: + return + try: + request_queue = Queue() + queue_filler_thread = Thread( + target=fill_queue, daemon=True, args=(request_iterator, request_queue) + ) + queue_filler_thread.start() + """For non `async get` requests, this loop yields immediately + For `async get` requests, this loop: + 1) does not yield, it just continues + 2) When the result is ready, it yields + """ + for req in iter(request_queue.get, None): + if isinstance(req, ray_client_pb2.DataResponse): + # Early shortcut if this is the result of an async get. + yield req + continue + + assert isinstance(req, ray_client_pb2.DataRequest) + if _should_cache(req) and reconnect_enabled: + cached_resp = response_cache.check_cache(req.req_id) + if isinstance(cached_resp, Exception): + # Cache state is invalid, raise exception + raise cached_resp + if cached_resp is not None: + yield cached_resp + continue + + resp = None + req_type = req.WhichOneof("type") + if req_type == "init": + resp_init = self.basic_service.Init(req.init) + resp = ray_client_pb2.DataResponse( + init=resp_init, + ) + with self.clients_lock: + self.reconnect_grace_periods[ + client_id + ] = req.init.reconnect_grace_period + if req.init.reconnect_grace_period == 0: + reconnect_enabled = False + + elif req_type == "get": + if req.get.asynchronous: + get_resp = self.basic_service._async_get_object( + req.get, client_id, req.req_id, request_queue + ) + if get_resp is None: + # Skip sending a response for this request and + # continue to the next requst. The response for + # this request will be sent when the object is + # ready. + continue + else: + get_resp = self.basic_service._get_object(req.get, client_id) + resp = ray_client_pb2.DataResponse(get=get_resp) + elif req_type == "put": + if not self.put_request_chunk_collector.add_chunk(req, req.put): + # Put request still in progress + continue + put_resp = self.basic_service._put_object( + self.put_request_chunk_collector.data, + req.put.client_ref_id, + client_id, + req.put.owner_id, + ) + self.put_request_chunk_collector.reset() + resp = ray_client_pb2.DataResponse(put=put_resp) + elif req_type == "release": + released = [] + for rel_id in req.release.ids: + rel = self.basic_service.release(client_id, rel_id) + released.append(rel) + resp = ray_client_pb2.DataResponse( + release=ray_client_pb2.ReleaseResponse(ok=released) + ) + elif req_type == "connection_info": + resp = ray_client_pb2.DataResponse( + connection_info=self._build_connection_response() + ) + elif req_type == "prep_runtime_env": + with self.clients_lock: + resp_prep = self.basic_service.PrepRuntimeEnv( + req.prep_runtime_env + ) + resp = ray_client_pb2.DataResponse(prep_runtime_env=resp_prep) + elif req_type == "connection_cleanup": + cleanup_requested = True + cleanup_resp = ray_client_pb2.ConnectionCleanupResponse() + resp = ray_client_pb2.DataResponse(connection_cleanup=cleanup_resp) + elif req_type == "acknowledge": + # Clean up acknowledged cache entries + response_cache.cleanup(req.acknowledge.req_id) + continue + elif req_type == "task": + with self.clients_lock: + task = req.task + if not self.client_task_chunk_collector.add_chunk(req, task): + # Not all serialized arguments have arrived + continue + arglist, kwargs = loads_from_client( + self.client_task_chunk_collector.data, self.basic_service + ) + self.client_task_chunk_collector.reset() + resp_ticket = self.basic_service.Schedule( + req.task, arglist, kwargs, context + ) + resp = ray_client_pb2.DataResponse(task_ticket=resp_ticket) + del arglist + del kwargs + elif req_type == "terminate": + with self.clients_lock: + response = self.basic_service.Terminate(req.terminate, context) + resp = ray_client_pb2.DataResponse(terminate=response) + elif req_type == "list_named_actors": + with self.clients_lock: + response = self.basic_service.ListNamedActors( + req.list_named_actors + ) + resp = ray_client_pb2.DataResponse(list_named_actors=response) + else: + raise Exception( + f"Unreachable code: Request type " + f"{req_type} not handled in Datapath" + ) + resp.req_id = req.req_id + if _should_cache(req) and reconnect_enabled: + response_cache.update_cache(req.req_id, resp) + yield resp + except Exception as e: + logger.exception("Error in data channel:") + recoverable = _propagate_error_in_context(e, context) + invalid_cache = response_cache.invalidate(e) + if not recoverable or invalid_cache: + context.set_code(grpc.StatusCode.FAILED_PRECONDITION) + # Connection isn't recoverable, skip cleanup + cleanup_requested = True + finally: + logger.debug(f"Stream is broken with client {client_id}") + queue_filler_thread.join(QUEUE_JOIN_SECONDS) + if queue_filler_thread.is_alive(): + logger.error( + "Queue filler thread failed to join before timeout: {}".format( + QUEUE_JOIN_SECONDS + ) + ) + cleanup_delay = self.reconnect_grace_periods.get(client_id) + if not cleanup_requested and cleanup_delay is not None: + logger.debug( + "Cleanup wasn't requested, delaying cleanup by" + f"{cleanup_delay} seconds." + ) + # Delay cleanup, since client may attempt a reconnect + # Wait on the "stopped" event in case the grpc server is + # stopped and we can clean up earlier. + self.stopped.wait(timeout=cleanup_delay) + else: + logger.debug("Cleanup was requested, cleaning up immediately.") + with self.clients_lock: + if client_id not in self.client_last_seen: + logger.debug("Connection already cleaned up.") + # Some other connection has already cleaned up this + # this client's session. This can happen if the client + # reconnects and then gracefully shut's down immediately. + return + last_seen = self.client_last_seen[client_id] + if last_seen > start_time: + # The client successfully reconnected and updated + # last seen some time during the grace period + logger.debug("Client reconnected, skipping cleanup") + return + # Either the client shut down gracefully, or the client + # failed to reconnect within the grace period. Clean up + # the connection. + self.basic_service.release_all(client_id) + del self.client_last_seen[client_id] + if client_id in self.reconnect_grace_periods: + del self.reconnect_grace_periods[client_id] + if client_id in self.response_caches: + del self.response_caches[client_id] + self.num_clients -= 1 + logger.debug( + f"Removed client {client_id}, " f"remaining={self.num_clients}" + ) + + # It's important to keep the Ray shutdown + # within this locked context or else Ray could hang. + # NOTE: it is strange to start ray in server.py but shut it + # down here. Consider consolidating ray lifetime management. + with disable_client_hook(): + if self.num_clients == 0: + logger.debug("Shutting down ray.") + ray.shutdown() + + def _init(self, client_id: str, context: Any, start_time: float): + """ + Checks if resources allow for another client. + Returns a boolean indicating if initialization was successful. + """ + with self.clients_lock: + reconnecting = _get_reconnecting_from_context(context) + threshold = int(CLIENT_SERVER_MAX_THREADS / 2) + if self.num_clients >= threshold: + logger.warning( + f"[Data Servicer]: Num clients {self.num_clients} " + f"has reached the threshold {threshold}. " + f"Rejecting client: {client_id}. " + ) + if log_once("client_threshold"): + logger.warning( + "You can configure the client connection " + "threshold by setting the " + "RAY_CLIENT_SERVER_MAX_THREADS env var " + f"(currently set to {CLIENT_SERVER_MAX_THREADS})." + ) + context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED) + return False + if reconnecting and client_id not in self.client_last_seen: + # Client took too long to reconnect, session has been + # cleaned up. + context.set_code(grpc.StatusCode.NOT_FOUND) + context.set_details( + "Attempted to reconnect to a session that has already " + "been cleaned up." + ) + return False + if client_id in self.client_last_seen: + logger.debug(f"Client {client_id} has reconnected.") + else: + self.num_clients += 1 + logger.debug( + f"Accepted data connection from {client_id}. " + f"Total clients: {self.num_clients}" + ) + self.client_last_seen[client_id] = start_time + return True + + def _build_connection_response(self): + with self.clients_lock: + cur_num_clients = self.num_clients + return ray_client_pb2.ConnectionInfoResponse( + num_clients=cur_num_clients, + python_version="{}.{}.{}".format( + sys.version_info[0], sys.version_info[1], sys.version_info[2] + ), + ray_version=ray.__version__, + ray_commit=ray.__commit__, + ) diff --git a/lib/python3.12/site-packages/ray/util/client/server/logservicer.py b/lib/python3.12/site-packages/ray/util/client/server/logservicer.py new file mode 100644 index 0000000000000000000000000000000000000000..764e6c82c65347b17ef9f76735f1d33636e27de9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/server/logservicer.py @@ -0,0 +1,125 @@ +"""This file responds to log stream requests and forwards logs +with its handler. +""" +import io +import logging +import queue +import threading +import uuid + +import grpc + +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc +from ray._private.ray_logging import global_worker_stdstream_dispatcher +from ray._private.worker import print_worker_logs +from ray.util.client.common import CLIENT_SERVER_MAX_THREADS + +logger = logging.getLogger(__name__) + + +class LogstreamHandler(logging.Handler): + def __init__(self, queue, level): + super().__init__() + self.queue = queue + self.level = level + + def emit(self, record: logging.LogRecord): + logdata = ray_client_pb2.LogData() + logdata.msg = record.getMessage() + logdata.level = record.levelno + logdata.name = record.name + self.queue.put(logdata) + + +class StdStreamHandler: + def __init__(self, queue): + self.queue = queue + self.id = str(uuid.uuid4()) + + def handle(self, data): + logdata = ray_client_pb2.LogData() + logdata.level = -2 if data["is_err"] else -1 + logdata.name = "stderr" if data["is_err"] else "stdout" + with io.StringIO() as file: + print_worker_logs(data, file) + logdata.msg = file.getvalue() + self.queue.put(logdata) + + def register_global(self): + global_worker_stdstream_dispatcher.add_handler(self.id, self.handle) + + def unregister_global(self): + global_worker_stdstream_dispatcher.remove_handler(self.id) + + +def log_status_change_thread(log_queue, request_iterator): + std_handler = StdStreamHandler(log_queue) + current_handler = None + root_logger = logging.getLogger("ray") + default_level = root_logger.getEffectiveLevel() + try: + for req in request_iterator: + if current_handler is not None: + root_logger.setLevel(default_level) + root_logger.removeHandler(current_handler) + std_handler.unregister_global() + if not req.enabled: + current_handler = None + continue + current_handler = LogstreamHandler(log_queue, req.loglevel) + std_handler.register_global() + root_logger.addHandler(current_handler) + root_logger.setLevel(req.loglevel) + except grpc.RpcError as e: + logger.debug(f"closing log thread " f"grpc error reading request_iterator: {e}") + finally: + if current_handler is not None: + root_logger.setLevel(default_level) + root_logger.removeHandler(current_handler) + std_handler.unregister_global() + log_queue.put(None) + + +class LogstreamServicer(ray_client_pb2_grpc.RayletLogStreamerServicer): + def __init__(self): + super().__init__() + self.num_clients = 0 + self.client_lock = threading.Lock() + + def Logstream(self, request_iterator, context): + initialized = False + with self.client_lock: + threshold = CLIENT_SERVER_MAX_THREADS / 2 + if self.num_clients + 1 >= threshold: + context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED) + logger.warning( + f"Logstream: Num clients {self.num_clients} has reached " + f"the threshold {threshold}. Rejecting new connection." + ) + return + self.num_clients += 1 + initialized = True + logger.info( + "New logs connection established. " f"Total clients: {self.num_clients}" + ) + log_queue = queue.Queue() + thread = threading.Thread( + target=log_status_change_thread, + args=(log_queue, request_iterator), + daemon=True, + ) + thread.start() + try: + queue_iter = iter(log_queue.get, None) + for record in queue_iter: + if record is None: + break + yield record + except grpc.RpcError as e: + logger.debug(f"Closing log channel: {e}") + finally: + thread.join() + with self.client_lock: + if initialized: + self.num_clients -= 1 diff --git a/lib/python3.12/site-packages/ray/util/client/server/proxier.py b/lib/python3.12/site-packages/ray/util/client/server/proxier.py new file mode 100644 index 0000000000000000000000000000000000000000..613f486f7e995a474c0589988f7ffe78193498ab --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/server/proxier.py @@ -0,0 +1,909 @@ +import atexit +import json +import logging +import socket +import sys +import time +import traceback +import urllib +from concurrent import futures +from dataclasses import dataclass +from itertools import chain +from threading import Event, Lock, RLock, Thread +from typing import Callable, Dict, List, Optional, Tuple + +import grpc + +import ray +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc +import ray.core.generated.runtime_env_agent_pb2 as runtime_env_agent_pb2 +from ray._common.network_utils import build_address, is_ipv6, is_localhost +from ray._private.authentication.http_token_authentication import ( + format_authentication_http_error, + get_auth_headers_if_auth_enabled, +) +from ray._private.client_mode_hook import disable_client_hook +from ray._private.grpc_utils import init_grpc_channel +from ray._private.parameter import RayParams +from ray._private.runtime_env.context import RuntimeEnvContext +from ray._private.services import ProcessInfo, start_ray_client_server +from ray._private.tls_utils import add_port_to_grpc_server +from ray._private.utils import detect_fate_sharing_support +from ray._raylet import GcsClient +from ray.cloudpickle.compat import pickle +from ray.exceptions import AuthenticationError +from ray.job_config import JobConfig +from ray.util.client.common import ( + CLIENT_SERVER_MAX_THREADS, + GRPC_OPTIONS, + ClientServerHandle, + _get_client_id_from_context, + _propagate_error_in_context, +) +from ray.util.client.server.dataservicer import _get_reconnecting_from_context + +# Import psutil after ray so the packaged version is used. +import psutil + +logger = logging.getLogger(__name__) + +CHECK_PROCESS_INTERVAL_S = 30 + +MIN_SPECIFIC_SERVER_PORT = 23000 +MAX_SPECIFIC_SERVER_PORT = 24000 + +CHECK_CHANNEL_TIMEOUT_S = 30 + +LOGSTREAM_RETRIES = 5 +LOGSTREAM_RETRY_INTERVAL_SEC = 2 + + +@dataclass +class SpecificServer: + port: int + process_handle_future: futures.Future + channel: "grpc._channel.Channel" + + def is_ready(self) -> bool: + """Check if the server is ready or not (doesn't block).""" + return self.process_handle_future.done() + + def wait_ready(self, timeout: Optional[float] = None) -> None: + """ + Wait for the server to actually start up. + """ + res = self.process_handle_future.result(timeout=timeout) + if res is None: + # This is only set to none when server creation specifically fails. + raise RuntimeError("Server startup failed.") + + def poll(self) -> Optional[int]: + """Check if the process has exited.""" + try: + proc = self.process_handle_future.result(timeout=0.1) + if proc is not None: + return proc.process.poll() + except futures.TimeoutError: + return + + def kill(self) -> None: + """Try to send a KILL signal to the process.""" + try: + proc = self.process_handle_future.result(timeout=0.1) + if proc is not None: + proc.process.kill() + except futures.TimeoutError: + # Server has not been started yet. + pass + + def set_result(self, proc: Optional[ProcessInfo]) -> None: + """Set the result of the internal future if it is currently unset.""" + if not self.is_ready(): + self.process_handle_future.set_result(proc) + + +def _match_running_client_server(command: List[str]) -> bool: + """ + Detects if the main process in the given command is the RayClient Server. + This works by ensuring that the command is of the form: + -m ray.util.client.server + """ + flattened = " ".join(command) + return "-m ray.util.client.server" in flattened + + +class ProxyManager: + def __init__( + self, + address: Optional[str], + runtime_env_agent_address: str, + *, + session_dir: Optional[str] = None, + redis_username: Optional[str] = None, + redis_password: Optional[str] = None, + runtime_env_agent_port: int = 0, + ): + self.servers: Dict[str, SpecificServer] = dict() + self.server_lock = RLock() + self._address = address + self._redis_username = redis_username + self._redis_password = redis_password + self._free_ports: List[int] = list( + range(MIN_SPECIFIC_SERVER_PORT, MAX_SPECIFIC_SERVER_PORT) + ) + + self._runtime_env_agent_address = runtime_env_agent_address + + self._check_thread = Thread(target=self._check_processes, daemon=True) + self._check_thread.start() + + self.fate_share = bool(detect_fate_sharing_support()) + self._node: Optional[ray._private.node.Node] = None + atexit.register(self._cleanup) + + def _get_unused_port(self, family: int = socket.AF_INET) -> int: + """ + Search for a port in _free_ports that is unused. + """ + with self.server_lock: + num_ports = len(self._free_ports) + for _ in range(num_ports): + port = self._free_ports.pop(0) + s = socket.socket(family, socket.SOCK_STREAM) + try: + s.bind(("", port)) + except OSError: + self._free_ports.append(port) + continue + finally: + s.close() + return port + raise RuntimeError("Unable to succeed in selecting a random port.") + + @property + def address(self) -> str: + """ + Returns the provided Ray bootstrap address, or creates a new cluster. + """ + if self._address: + return self._address + # Start a new, locally scoped cluster. + connection_tuple = ray.init() + self._address = connection_tuple["address"] + self._session_dir = connection_tuple["session_dir"] + return self._address + + @property + def node(self) -> ray._private.node.Node: + """Gets a 'ray.Node' object for this node (the head node). + If it does not already exist, one is created using the bootstrap + address. + """ + if self._node: + return self._node + ray_params = RayParams(gcs_address=self.address) + + self._node = ray._private.node.Node( + ray_params, + head=False, + shutdown_at_exit=False, + spawn_reaper=False, + connect_only=True, + ) + + return self._node + + def create_specific_server(self, client_id: str) -> SpecificServer: + """ + Create, but not start a SpecificServer for a given client. This + method must be called once per client. + """ + with self.server_lock: + assert ( + self.servers.get(client_id) is None + ), f"Server already created for Client: {client_id}" + + host = "127.0.0.1" + port = self._get_unused_port( + socket.AF_INET6 if is_ipv6(host) else socket.AF_INET + ) + + server = SpecificServer( + port=port, + process_handle_future=futures.Future(), + channel=init_grpc_channel( + build_address(host, port), options=GRPC_OPTIONS + ), + ) + self.servers[client_id] = server + return server + + def _create_runtime_env( + self, + serialized_runtime_env: str, + runtime_env_config: str, + specific_server: SpecificServer, + ): + """Increase the runtime_env reference by sending an RPC to the agent. + + Includes retry logic to handle the case when the agent is + temporarily unreachable (e.g., hasn't been started up yet). + """ + logger.info( + f"Increasing runtime env reference for " + f"ray_client_server_{specific_server.port}." + f"Serialized runtime env is {serialized_runtime_env}." + ) + + assert ( + len(self._runtime_env_agent_address) > 0 + ), "runtime_env_agent_address not set" + + create_env_request = runtime_env_agent_pb2.GetOrCreateRuntimeEnvRequest( + serialized_runtime_env=serialized_runtime_env, + runtime_env_config=runtime_env_config, + job_id=f"ray_client_server_{specific_server.port}".encode("utf-8"), + source_process="client_server", + ) + + retries = 0 + max_retries = 5 + wait_time_s = 0.5 + last_exception = None + while retries <= max_retries: + try: + url = urllib.parse.urljoin( + self._runtime_env_agent_address, "/get_or_create_runtime_env" + ) + data = create_env_request.SerializeToString() + headers = {"Content-Type": "application/octet-stream"} + headers.update(**get_auth_headers_if_auth_enabled(headers)) + req = urllib.request.Request( + url, data=data, method="POST", headers=headers + ) + response = urllib.request.urlopen(req, timeout=None) + response_data = response.read() + r = runtime_env_agent_pb2.GetOrCreateRuntimeEnvReply() + r.ParseFromString(response_data) + + if r.status == runtime_env_agent_pb2.AgentRpcStatus.AGENT_RPC_STATUS_OK: + return r.serialized_runtime_env_context + elif ( + r.status + == runtime_env_agent_pb2.AgentRpcStatus.AGENT_RPC_STATUS_FAILED + ): + raise RuntimeError( + "Failed to create runtime_env for Ray client " + f"server, it is caused by:\n{r.error_message}" + ) + else: + assert False, f"Unknown status: {r.status}." + except urllib.error.HTTPError as e: + body = "" + try: + body = e.read().decode("utf-8", "ignore") + except Exception: + body = e.reason if hasattr(e, "reason") else str(e) + + formatted_error = format_authentication_http_error(e.code, body or "") + if formatted_error: + raise AuthenticationError(formatted_error) from e + + # Treat non-auth HTTP errors like URLError (retry with backoff) + last_exception = e + logger.warning( + f"GetOrCreateRuntimeEnv request failed with HTTP {e.code}: {body or e}. " + f"Retrying after {wait_time_s}s. " + f"{max_retries-retries} retries remaining." + ) + + except urllib.error.URLError as e: + last_exception = e + logger.warning( + f"GetOrCreateRuntimeEnv request failed: {e}. " + f"Retrying after {wait_time_s}s. " + f"{max_retries-retries} retries remaining." + ) + + # Exponential backoff. + time.sleep(wait_time_s) + retries += 1 + wait_time_s *= 2 + + raise TimeoutError( + f"GetOrCreateRuntimeEnv request failed after {max_retries} attempts." + f" Last exception: {last_exception}" + ) + + def start_specific_server(self, client_id: str, job_config: JobConfig) -> bool: + """ + Start up a RayClient Server for an incoming client to + communicate with. Returns whether creation was successful. + """ + specific_server = self._get_server_for_client(client_id) + assert specific_server, f"Server has not been created for: {client_id}" + + output, error = self.node.get_log_file_handles( + f"ray_client_server_{specific_server.port}", unique=True + ) + + serialized_runtime_env = job_config._get_serialized_runtime_env() + runtime_env_config = job_config._get_proto_runtime_env_config() + if not serialized_runtime_env or serialized_runtime_env == "{}": + # TODO(edoakes): can we just remove this case and always send it + # to the agent? + serialized_runtime_env_context = RuntimeEnvContext().serialize() + else: + serialized_runtime_env_context = self._create_runtime_env( + serialized_runtime_env=serialized_runtime_env, + runtime_env_config=runtime_env_config, + specific_server=specific_server, + ) + + proc = start_ray_client_server( + self.address, + "127.0.0.1", + specific_server.port, + stdout_file=output, + stderr_file=error, + fate_share=self.fate_share, + server_type="specific-server", + serialized_runtime_env_context=serialized_runtime_env_context, + redis_username=self._redis_username, + redis_password=self._redis_password, + ) + + # Wait for the process being run transitions from the shim process + # to the actual RayClient Server. + pid = proc.process.pid + if sys.platform != "win32": + psutil_proc = psutil.Process(pid) + else: + psutil_proc = None + # Don't use `psutil` on Win32 + while psutil_proc is not None: + if proc.process.poll() is not None: + logger.error(f"SpecificServer startup failed for client: {client_id}") + break + cmd = psutil_proc.cmdline() + if _match_running_client_server(cmd): + break + logger.debug("Waiting for Process to reach the actual client server.") + time.sleep(0.5) + specific_server.set_result(proc) + logger.info( + f"SpecificServer started on port: {specific_server.port} " + f"with PID: {pid} for client: {client_id}" + ) + return proc.process.poll() is None + + def _get_server_for_client(self, client_id: str) -> Optional[SpecificServer]: + with self.server_lock: + client = self.servers.get(client_id) + if client is None: + logger.error(f"Unable to find channel for client: {client_id}") + return client + + def has_channel(self, client_id: str) -> bool: + server = self._get_server_for_client(client_id) + if server is None: + return False + + return server.is_ready() + + def get_channel( + self, + client_id: str, + ) -> Optional["grpc._channel.Channel"]: + """ + Find the gRPC Channel for the given client_id. This will block until + the server process has started. + """ + server = self._get_server_for_client(client_id) + if server is None: + return None + # Wait for the SpecificServer to become ready. + server.wait_ready() + try: + grpc.channel_ready_future(server.channel).result( + timeout=CHECK_CHANNEL_TIMEOUT_S + ) + return server.channel + except grpc.FutureTimeoutError: + logger.exception(f"Timeout waiting for channel for {client_id}") + return None + + def _check_processes(self): + """ + Keeps the internal servers dictionary up-to-date with running servers. + """ + while True: + with self.server_lock: + for client_id, specific_server in list(self.servers.items()): + if specific_server.poll() is not None: + logger.info( + f"Specific server {client_id} is no longer running" + f", freeing its port {specific_server.port}" + ) + del self.servers[client_id] + # Port is available to use again. + self._free_ports.append(specific_server.port) + + time.sleep(CHECK_PROCESS_INTERVAL_S) + + def _cleanup(self) -> None: + """ + Forcibly kill all spawned RayClient Servers. This ensures cleanup + for platforms where fate sharing is not supported. + """ + for server in self.servers.values(): + server.kill() + + +class RayletServicerProxy(ray_client_pb2_grpc.RayletDriverServicer): + def __init__(self, ray_connect_handler: Callable, proxy_manager: ProxyManager): + self.proxy_manager = proxy_manager + self.ray_connect_handler = ray_connect_handler + + def _call_inner_function( + self, request, context, method: str + ) -> Optional[ray_client_pb2_grpc.RayletDriverStub]: + client_id = _get_client_id_from_context(context) + chan = self.proxy_manager.get_channel(client_id) + if not chan: + logger.error(f"Channel for Client: {client_id} not found!") + context.set_code(grpc.StatusCode.NOT_FOUND) + return None + + stub = ray_client_pb2_grpc.RayletDriverStub(chan) + try: + metadata = [("client_id", client_id)] + if context: + metadata = context.invocation_metadata() + return getattr(stub, method)(request, metadata=metadata) + except Exception as e: + # Error while proxying -- propagate the error's context to user + logger.exception(f"Proxying call to {method} failed!") + _propagate_error_in_context(e, context) + + def _has_channel_for_request(self, context): + client_id = _get_client_id_from_context(context) + return self.proxy_manager.has_channel(client_id) + + def Init(self, request, context=None) -> ray_client_pb2.InitResponse: + return self._call_inner_function(request, context, "Init") + + def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse: + """Proxies internal_kv.put. + + This is used by the working_dir code to upload to the GCS before + ray.init is called. In that case (if we don't have a server yet) + we directly make the internal KV call from the proxier. + + Otherwise, we proxy the call to the downstream server as usual. + """ + if self._has_channel_for_request(context): + return self._call_inner_function(request, context, "KVPut") + + with disable_client_hook(): + already_exists = ray.experimental.internal_kv._internal_kv_put( + request.key, request.value, overwrite=request.overwrite + ) + return ray_client_pb2.KVPutResponse(already_exists=already_exists) + + def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse: + """Proxies internal_kv.get. + + This is used by the working_dir code to upload to the GCS before + ray.init is called. In that case (if we don't have a server yet) + we directly make the internal KV call from the proxier. + + Otherwise, we proxy the call to the downstream server as usual. + """ + if self._has_channel_for_request(context): + return self._call_inner_function(request, context, "KVGet") + + with disable_client_hook(): + value = ray.experimental.internal_kv._internal_kv_get(request.key) + return ray_client_pb2.KVGetResponse(value=value) + + def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse: + """Proxies internal_kv.delete. + + This is used by the working_dir code to upload to the GCS before + ray.init is called. In that case (if we don't have a server yet) + we directly make the internal KV call from the proxier. + + Otherwise, we proxy the call to the downstream server as usual. + """ + if self._has_channel_for_request(context): + return self._call_inner_function(request, context, "KVDel") + + with disable_client_hook(): + ray.experimental.internal_kv._internal_kv_del(request.key) + return ray_client_pb2.KVDelResponse() + + def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse: + """Proxies internal_kv.list. + + This is used by the working_dir code to upload to the GCS before + ray.init is called. In that case (if we don't have a server yet) + we directly make the internal KV call from the proxier. + + Otherwise, we proxy the call to the downstream server as usual. + """ + if self._has_channel_for_request(context): + return self._call_inner_function(request, context, "KVList") + + with disable_client_hook(): + keys = ray.experimental.internal_kv._internal_kv_list(request.prefix) + return ray_client_pb2.KVListResponse(keys=keys) + + def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse: + """Proxies internal_kv.exists. + + This is used by the working_dir code to upload to the GCS before + ray.init is called. In that case (if we don't have a server yet) + we directly make the internal KV call from the proxier. + + Otherwise, we proxy the call to the downstream server as usual. + """ + if self._has_channel_for_request(context): + return self._call_inner_function(request, context, "KVExists") + + with disable_client_hook(): + exists = ray.experimental.internal_kv._internal_kv_exists(request.key) + return ray_client_pb2.KVExistsResponse(exists=exists) + + def PinRuntimeEnvURI( + self, request, context=None + ) -> ray_client_pb2.ClientPinRuntimeEnvURIResponse: + """Proxies internal_kv.pin_runtime_env_uri. + + This is used by the working_dir code to upload to the GCS before + ray.init is called. In that case (if we don't have a server yet) + we directly make the internal KV call from the proxier. + + Otherwise, we proxy the call to the downstream server as usual. + """ + if self._has_channel_for_request(context): + return self._call_inner_function(request, context, "PinRuntimeEnvURI") + + with disable_client_hook(): + ray.experimental.internal_kv._pin_runtime_env_uri( + request.uri, expiration_s=request.expiration_s + ) + return ray_client_pb2.ClientPinRuntimeEnvURIResponse() + + def ListNamedActors( + self, request, context=None + ) -> ray_client_pb2.ClientListNamedActorsResponse: + return self._call_inner_function(request, context, "ListNamedActors") + + def ClusterInfo(self, request, context=None) -> ray_client_pb2.ClusterInfoResponse: + + # NOTE: We need to respond to the PING request here to allow the client + # to continue with connecting. + if request.type == ray_client_pb2.ClusterInfoType.PING: + resp = ray_client_pb2.ClusterInfoResponse(json=json.dumps({})) + return resp + return self._call_inner_function(request, context, "ClusterInfo") + + def Terminate(self, req, context=None): + return self._call_inner_function(req, context, "Terminate") + + def GetObject(self, request, context=None): + try: + yield from self._call_inner_function(request, context, "GetObject") + except Exception as e: + # Error while iterating over response from GetObject stream + logger.exception("Proxying call to GetObject failed!") + _propagate_error_in_context(e, context) + + def PutObject( + self, request: ray_client_pb2.PutRequest, context=None + ) -> ray_client_pb2.PutResponse: + return self._call_inner_function(request, context, "PutObject") + + def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse: + return self._call_inner_function(request, context, "WaitObject") + + def Schedule(self, task, context=None) -> ray_client_pb2.ClientTaskTicket: + return self._call_inner_function(task, context, "Schedule") + + +def ray_client_server_env_prep(job_config: JobConfig) -> JobConfig: + return job_config + + +def prepare_runtime_init_req( + init_request: ray_client_pb2.DataRequest, +) -> Tuple[ray_client_pb2.DataRequest, JobConfig]: + """ + Extract JobConfig and possibly mutate InitRequest before it is passed to + the specific RayClient Server. + """ + init_type = init_request.WhichOneof("type") + assert init_type == "init", ( + "Received initial message of type " f"{init_type}, not 'init'." + ) + req = init_request.init + job_config = JobConfig() + if req.job_config: + job_config = pickle.loads(req.job_config) + new_job_config = ray_client_server_env_prep(job_config) + modified_init_req = ray_client_pb2.InitRequest( + job_config=pickle.dumps(new_job_config), + ray_init_kwargs=init_request.init.ray_init_kwargs, + reconnect_grace_period=init_request.init.reconnect_grace_period, + ) + + init_request.init.CopyFrom(modified_init_req) + return (init_request, new_job_config) + + +class RequestIteratorProxy: + def __init__(self, request_iterator): + self.request_iterator = request_iterator + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self.request_iterator) + except grpc.RpcError as e: + # To stop proxying already CANCLLED request stream gracefully, + # we only translate the exact grpc.RpcError to StopIteration, + # not its subsclasses. ex: grpc._Rendezvous + # https://github.com/grpc/grpc/blob/v1.43.0/src/python/grpcio/grpc/_server.py#L353-L354 + # This fixes the https://github.com/ray-project/ray/issues/23865 + if type(e) is not grpc.RpcError: + raise e # re-raise other grpc exceptions + logger.exception( + "Stop iterating cancelled request stream with the following exception:" + ) + raise StopIteration + + +class DataServicerProxy(ray_client_pb2_grpc.RayletDataStreamerServicer): + def __init__(self, proxy_manager: ProxyManager): + self.num_clients = 0 + # dictionary mapping client_id's to the last time they connected + self.clients_last_seen: Dict[str, float] = {} + self.reconnect_grace_periods: Dict[str, float] = {} + self.clients_lock = Lock() + self.proxy_manager = proxy_manager + self.stopped = Event() + + def modify_connection_info_resp( + self, init_resp: ray_client_pb2.DataResponse + ) -> ray_client_pb2.DataResponse: + """ + Modify the `num_clients` returned the ConnectionInfoResponse because + individual SpecificServers only have **one** client. + """ + init_type = init_resp.WhichOneof("type") + if init_type != "connection_info": + return init_resp + modified_resp = ray_client_pb2.DataResponse() + modified_resp.CopyFrom(init_resp) + with self.clients_lock: + modified_resp.connection_info.num_clients = self.num_clients + return modified_resp + + def Datapath(self, request_iterator, context): + request_iterator = RequestIteratorProxy(request_iterator) + cleanup_requested = False + start_time = time.time() + client_id = _get_client_id_from_context(context) + if client_id == "": + return + reconnecting = _get_reconnecting_from_context(context) + + if reconnecting: + with self.clients_lock: + if client_id not in self.clients_last_seen: + # Client took too long to reconnect, session has already + # been cleaned up + context.set_code(grpc.StatusCode.NOT_FOUND) + context.set_details( + "Attempted to reconnect a session that has already " + "been cleaned up" + ) + return + self.clients_last_seen[client_id] = start_time + server = self.proxy_manager._get_server_for_client(client_id) + channel = self.proxy_manager.get_channel(client_id) + # iterator doesn't need modification on reconnect + new_iter = request_iterator + else: + # Create Placeholder *before* reading the first request. + server = self.proxy_manager.create_specific_server(client_id) + with self.clients_lock: + self.clients_last_seen[client_id] = start_time + self.num_clients += 1 + + try: + if not reconnecting: + logger.info(f"New data connection from client {client_id}: ") + init_req = next(request_iterator) + with self.clients_lock: + self.reconnect_grace_periods[ + client_id + ] = init_req.init.reconnect_grace_period + try: + modified_init_req, job_config = prepare_runtime_init_req(init_req) + if not self.proxy_manager.start_specific_server( + client_id, job_config + ): + logger.error( + f"Server startup failed for client: {client_id}, " + f"using JobConfig: {job_config}!" + ) + raise RuntimeError( + "Starting Ray client server failed. See " + f"ray_client_server_{server.port}.err for " + "detailed logs." + ) + channel = self.proxy_manager.get_channel(client_id) + if channel is None: + logger.error(f"Channel not found for {client_id}") + raise RuntimeError( + "Proxy failed to Connect to backend! Check " + "`ray_client_server.err` and " + f"`ray_client_server_{server.port}.err` on the " + "head node of the cluster for the relevant logs. " + "By default these are located at " + "/tmp/ray/session_latest/logs." + ) + except Exception: + init_resp = ray_client_pb2.DataResponse( + init=ray_client_pb2.InitResponse( + ok=False, msg=traceback.format_exc() + ) + ) + init_resp.req_id = init_req.req_id + yield init_resp + return None + + new_iter = chain([modified_init_req], request_iterator) + + stub = ray_client_pb2_grpc.RayletDataStreamerStub(channel) + metadata = [("client_id", client_id), ("reconnecting", str(reconnecting))] + resp_stream = stub.Datapath(new_iter, metadata=metadata) + for resp in resp_stream: + resp_type = resp.WhichOneof("type") + if resp_type == "connection_cleanup": + # Specific server is skipping cleanup, proxier should too + cleanup_requested = True + yield self.modify_connection_info_resp(resp) + except Exception as e: + logger.exception("Proxying Datapath failed!") + # Propogate error through context + recoverable = _propagate_error_in_context(e, context) + if not recoverable: + # Client shouldn't attempt to recover, clean up connection + cleanup_requested = True + finally: + cleanup_delay = self.reconnect_grace_periods.get(client_id) + if not cleanup_requested and cleanup_delay is not None: + # Delay cleanup, since client may attempt a reconnect + # Wait on stopped event in case the server closes and we + # can clean up earlier + self.stopped.wait(timeout=cleanup_delay) + with self.clients_lock: + if client_id not in self.clients_last_seen: + logger.info(f"{client_id} not found. Skipping clean up.") + # Connection has already been cleaned up + return + last_seen = self.clients_last_seen[client_id] + logger.info( + f"{client_id} last started stream at {last_seen}. Current " + f"stream started at {start_time}." + ) + if last_seen > start_time: + logger.info("Client reconnected. Skipping cleanup.") + # Client has reconnected, don't clean up + return + logger.debug(f"Client detached: {client_id}") + self.num_clients -= 1 + del self.clients_last_seen[client_id] + if client_id in self.reconnect_grace_periods: + del self.reconnect_grace_periods[client_id] + server.set_result(None) + + +class LogstreamServicerProxy(ray_client_pb2_grpc.RayletLogStreamerServicer): + def __init__(self, proxy_manager: ProxyManager): + super().__init__() + self.proxy_manager = proxy_manager + + def Logstream(self, request_iterator, context): + request_iterator = RequestIteratorProxy(request_iterator) + client_id = _get_client_id_from_context(context) + if client_id == "": + return + logger.debug(f"New logstream connection from client {client_id}: ") + + channel = None + # We need to retry a few times because the LogClient *may* connect + # Before the DataClient has finished connecting. + for i in range(LOGSTREAM_RETRIES): + channel = self.proxy_manager.get_channel(client_id) + + if channel is not None: + break + logger.warning(f"Retrying Logstream connection. {i+1} attempts failed.") + time.sleep(LOGSTREAM_RETRY_INTERVAL_SEC) + + if channel is None: + context.set_code(grpc.StatusCode.NOT_FOUND) + context.set_details( + "Logstream proxy failed to connect. Channel for client " + f"{client_id} not found." + ) + return None + + stub = ray_client_pb2_grpc.RayletLogStreamerStub(channel) + + resp_stream = stub.Logstream( + request_iterator, metadata=[("client_id", client_id)] + ) + try: + for resp in resp_stream: + yield resp + except Exception: + logger.exception("Proxying Logstream failed!") + + +def serve_proxier( + host: str, + port: int, + gcs_address: Optional[str], + *, + redis_username: Optional[str] = None, + redis_password: Optional[str] = None, + session_dir: Optional[str] = None, + runtime_env_agent_address: Optional[str] = None, +): + # Initialize internal KV to be used to upload and download working_dir + # before calling ray.init within the RayletServicers. + # NOTE(edoakes): redis_address and redis_password should only be None in + # tests. + if gcs_address is not None: + gcs_cli = GcsClient(address=gcs_address) + ray.experimental.internal_kv._initialize_internal_kv(gcs_cli) + + from ray._private.grpc_utils import create_grpc_server_with_interceptors + + server = create_grpc_server_with_interceptors( + max_workers=CLIENT_SERVER_MAX_THREADS, + thread_name_prefix="ray_client_proxier", + options=GRPC_OPTIONS, + asynchronous=False, + ) + proxy_manager = ProxyManager( + gcs_address, + session_dir=session_dir, + redis_username=redis_username, + redis_password=redis_password, + runtime_env_agent_address=runtime_env_agent_address, + ) + task_servicer = RayletServicerProxy(None, proxy_manager) + data_servicer = DataServicerProxy(proxy_manager) + logs_servicer = LogstreamServicerProxy(proxy_manager) + ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server) + ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server) + ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server) + if not is_localhost(host): + add_port_to_grpc_server(server, f"127.0.0.1:{port}") + add_port_to_grpc_server(server, f"{host}:{port}") + server.start() + return ClientServerHandle( + task_servicer=task_servicer, + data_servicer=data_servicer, + logs_servicer=logs_servicer, + grpc_server=server, + ) diff --git a/lib/python3.12/site-packages/ray/util/client/server/server.py b/lib/python3.12/site-packages/ray/util/client/server/server.py new file mode 100644 index 0000000000000000000000000000000000000000..c59ca172079208366803a84911fb79f448afa533 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/server/server.py @@ -0,0 +1,953 @@ +import base64 +import functools +import gc +import inspect +import json +import logging +import math +import pickle +import queue +import threading +import time +from collections import defaultdict +from typing import Any, Callable, Dict, List, Optional, Set, Union + +import grpc + +import ray +import ray._private.state +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc +from ray import cloudpickle +from ray._common.network_utils import build_address, is_localhost +from ray._private import ray_constants +from ray._private.client_mode_hook import disable_client_hook +from ray._private.ray_constants import env_integer +from ray._private.ray_logging import setup_logger +from ray._private.services import canonicalize_bootstrap_address_or_die +from ray._private.tls_utils import add_port_to_grpc_server +from ray._raylet import GcsClient +from ray.job_config import JobConfig +from ray.util.client.common import ( + CLIENT_SERVER_MAX_THREADS, + GRPC_OPTIONS, + OBJECT_TRANSFER_CHUNK_SIZE, + ClientServerHandle, + ResponseCache, +) +from ray.util.client.server.dataservicer import DataServicer +from ray.util.client.server.logservicer import LogstreamServicer +from ray.util.client.server.proxier import serve_proxier +from ray.util.client.server.server_pickler import dumps_from_server, loads_from_client +from ray.util.client.server.server_stubs import current_server + +logger = logging.getLogger(__name__) + +TIMEOUT_FOR_SPECIFIC_SERVER_S = env_integer("TIMEOUT_FOR_SPECIFIC_SERVER_S", 30) + + +def _use_response_cache(func): + """ + Decorator for gRPC stubs. Before calling the real stubs, checks if there's + an existing entry in the caches. If there is, then return the cached + entry. Otherwise, call the real function and use the real cache + """ + + @functools.wraps(func) + def wrapper(self, request, context): + metadata = dict(context.invocation_metadata()) + expected_ids = ("client_id", "thread_id", "req_id") + if any(i not in metadata for i in expected_ids): + # Missing IDs, skip caching and call underlying stub directly + return func(self, request, context) + + # Get relevant IDs to check cache + client_id = metadata["client_id"] + thread_id = metadata["thread_id"] + req_id = int(metadata["req_id"]) + + # Check if response already cached + response_cache = self.response_caches[client_id] + cached_entry = response_cache.check_cache(thread_id, req_id) + if cached_entry is not None: + if isinstance(cached_entry, Exception): + # Original call errored, propogate error + context.set_code(grpc.StatusCode.FAILED_PRECONDITION) + context.set_details(str(cached_entry)) + raise cached_entry + return cached_entry + + try: + # Response wasn't cached, call underlying stub and cache result + resp = func(self, request, context) + except Exception as e: + # Unexpected error in underlying stub -- update cache and + # propagate to user through context + response_cache.update_cache(thread_id, req_id, e) + context.set_code(grpc.StatusCode.FAILED_PRECONDITION) + context.set_details(str(e)) + raise + response_cache.update_cache(thread_id, req_id, resp) + return resp + + return wrapper + + +class RayletServicer(ray_client_pb2_grpc.RayletDriverServicer): + def __init__(self, ray_connect_handler: Callable): + """Construct a raylet service + + Args: + ray_connect_handler: Function to connect to ray cluster + """ + # Stores client_id -> (ref_id -> ObjectRef) + self.object_refs: Dict[str, Dict[bytes, ray.ObjectRef]] = defaultdict(dict) + # Stores client_id -> (client_ref_id -> ref_id (in self.object_refs)) + self.client_side_ref_map: Dict[str, Dict[bytes, bytes]] = defaultdict(dict) + self.function_refs = {} + self.actor_refs: Dict[bytes, ray.ActorHandle] = {} + self.actor_owners: Dict[str, Set[bytes]] = defaultdict(set) + self.registered_actor_classes = {} + self.named_actors = set() + self.state_lock = threading.Lock() + self.ray_connect_handler = ray_connect_handler + self.response_caches: Dict[str, ResponseCache] = defaultdict(ResponseCache) + + def Init( + self, request: ray_client_pb2.InitRequest, context=None + ) -> ray_client_pb2.InitResponse: + if request.job_config: + job_config = pickle.loads(request.job_config) + job_config._client_job = True + else: + job_config = None + current_job_config = None + with disable_client_hook(): + if ray.is_initialized(): + worker = ray._private.worker.global_worker + current_job_config = worker.core_worker.get_job_config() + else: + extra_kwargs = json.loads(request.ray_init_kwargs or "{}") + try: + self.ray_connect_handler(job_config, **extra_kwargs) + except Exception as e: + logger.exception("Running Ray Init failed:") + return ray_client_pb2.InitResponse( + ok=False, + msg=f"Call to `ray.init()` on the server failed with: {e}", + ) + if job_config is None: + return ray_client_pb2.InitResponse(ok=True) + + # NOTE(edoakes): this code should not be necessary anymore because we + # only allow a single client/job per server. There is an existing test + # that tests the behavior of multiple clients with the same job config + # connecting to one server (test_client_init.py::test_num_clients), + # so I'm leaving it here for now. + job_config = job_config._get_proto_job_config() + # If the server has been initialized, we need to compare whether the + # runtime env is compatible. + if current_job_config: + job_uris = set(job_config.runtime_env_info.uris.working_dir_uri) + job_uris.update(job_config.runtime_env_info.uris.py_modules_uris) + current_job_uris = set( + current_job_config.runtime_env_info.uris.working_dir_uri + ) + current_job_uris.update( + current_job_config.runtime_env_info.uris.py_modules_uris + ) + if job_uris != current_job_uris and len(job_uris) > 0: + return ray_client_pb2.InitResponse( + ok=False, + msg="Runtime environment doesn't match " + f"request one {job_config.runtime_env_info.uris} " + f"current one {current_job_config.runtime_env_info.uris}", + ) + return ray_client_pb2.InitResponse(ok=True) + + @_use_response_cache + def KVPut(self, request, context=None) -> ray_client_pb2.KVPutResponse: + try: + with disable_client_hook(): + already_exists = ray.experimental.internal_kv._internal_kv_put( + request.key, + request.value, + overwrite=request.overwrite, + namespace=request.namespace, + ) + except Exception as e: + return_exception_in_context(e, context) + already_exists = False + return ray_client_pb2.KVPutResponse(already_exists=already_exists) + + def KVGet(self, request, context=None) -> ray_client_pb2.KVGetResponse: + try: + with disable_client_hook(): + value = ray.experimental.internal_kv._internal_kv_get( + request.key, namespace=request.namespace + ) + except Exception as e: + return_exception_in_context(e, context) + value = b"" + return ray_client_pb2.KVGetResponse(value=value) + + @_use_response_cache + def KVDel(self, request, context=None) -> ray_client_pb2.KVDelResponse: + try: + with disable_client_hook(): + deleted_num = ray.experimental.internal_kv._internal_kv_del( + request.key, + del_by_prefix=request.del_by_prefix, + namespace=request.namespace, + ) + except Exception as e: + return_exception_in_context(e, context) + deleted_num = 0 + return ray_client_pb2.KVDelResponse(deleted_num=deleted_num) + + def KVList(self, request, context=None) -> ray_client_pb2.KVListResponse: + try: + with disable_client_hook(): + keys = ray.experimental.internal_kv._internal_kv_list( + request.prefix, namespace=request.namespace + ) + except Exception as e: + return_exception_in_context(e, context) + keys = [] + return ray_client_pb2.KVListResponse(keys=keys) + + def KVExists(self, request, context=None) -> ray_client_pb2.KVExistsResponse: + try: + with disable_client_hook(): + exists = ray.experimental.internal_kv._internal_kv_exists( + request.key, namespace=request.namespace + ) + except Exception as e: + return_exception_in_context(e, context) + exists = False + return ray_client_pb2.KVExistsResponse(exists=exists) + + def ListNamedActors( + self, request, context=None + ) -> ray_client_pb2.ClientListNamedActorsResponse: + with disable_client_hook(): + actors = ray.util.list_named_actors(all_namespaces=request.all_namespaces) + + return ray_client_pb2.ClientListNamedActorsResponse( + actors_json=json.dumps(actors) + ) + + def ClusterInfo(self, request, context=None) -> ray_client_pb2.ClusterInfoResponse: + resp = ray_client_pb2.ClusterInfoResponse() + resp.type = request.type + if request.type == ray_client_pb2.ClusterInfoType.CLUSTER_RESOURCES: + with disable_client_hook(): + resources = ray.cluster_resources() + # Normalize resources into floats + # (the function may return values that are ints) + float_resources = {k: float(v) for k, v in resources.items()} + resp.resource_table.CopyFrom( + ray_client_pb2.ClusterInfoResponse.ResourceTable(table=float_resources) + ) + elif request.type == ray_client_pb2.ClusterInfoType.AVAILABLE_RESOURCES: + with disable_client_hook(): + resources = ray.available_resources() + # Normalize resources into floats + # (the function may return values that are ints) + float_resources = {k: float(v) for k, v in resources.items()} + resp.resource_table.CopyFrom( + ray_client_pb2.ClusterInfoResponse.ResourceTable(table=float_resources) + ) + elif request.type == ray_client_pb2.ClusterInfoType.RUNTIME_CONTEXT: + ctx = ray_client_pb2.ClusterInfoResponse.RuntimeContext() + with disable_client_hook(): + rtc = ray.get_runtime_context() + ctx.job_id = ray._common.utils.hex_to_binary(rtc.get_job_id()) + ctx.node_id = ray._common.utils.hex_to_binary(rtc.get_node_id()) + ctx.namespace = rtc.namespace + ctx.capture_client_tasks = ( + rtc.should_capture_child_tasks_in_placement_group + ) + ctx.gcs_address = rtc.gcs_address + ctx.runtime_env = rtc.get_runtime_env_string() + resp.runtime_context.CopyFrom(ctx) + else: + with disable_client_hook(): + resp.json = self._return_debug_cluster_info(request, context) + return resp + + def _return_debug_cluster_info(self, request, context=None) -> str: + """Handle ClusterInfo requests that only return a json blob.""" + data = None + if request.type == ray_client_pb2.ClusterInfoType.NODES: + data = ray.nodes() + elif request.type == ray_client_pb2.ClusterInfoType.IS_INITIALIZED: + data = ray.is_initialized() + elif request.type == ray_client_pb2.ClusterInfoType.TIMELINE: + data = ray.timeline() + elif request.type == ray_client_pb2.ClusterInfoType.PING: + data = {} + elif request.type == ray_client_pb2.ClusterInfoType.DASHBOARD_URL: + data = {"dashboard_url": ray._private.worker.get_dashboard_url()} + else: + raise TypeError("Unsupported cluster info type") + return json.dumps(data) + + def release(self, client_id: str, id: bytes) -> bool: + with self.state_lock: + if client_id in self.object_refs: + if id in self.object_refs[client_id]: + logger.debug(f"Releasing object {id.hex()} for {client_id}") + del self.object_refs[client_id][id] + return True + + if client_id in self.actor_owners: + if id in self.actor_owners[client_id]: + logger.debug(f"Releasing actor {id.hex()} for {client_id}") + self.actor_owners[client_id].remove(id) + if self._can_remove_actor_ref(id): + logger.debug(f"Deleting reference to actor {id.hex()}") + del self.actor_refs[id] + return True + + return False + + def release_all(self, client_id): + with self.state_lock: + self._release_objects(client_id) + self._release_actors(client_id) + # NOTE: Try to actually dereference the object and actor refs. + # Otherwise dereferencing will happen later, which may run concurrently + # with ray.shutdown() and will crash the process. The crash is a bug + # that should be fixed eventually. + gc.collect() + + def _can_remove_actor_ref(self, actor_id_bytes): + no_owner = not any( + actor_id_bytes in actor_list for actor_list in self.actor_owners.values() + ) + return no_owner and actor_id_bytes not in self.named_actors + + def _release_objects(self, client_id): + if client_id not in self.object_refs: + logger.debug(f"Releasing client with no references: {client_id}") + return + count = len(self.object_refs[client_id]) + del self.object_refs[client_id] + if client_id in self.client_side_ref_map: + del self.client_side_ref_map[client_id] + if client_id in self.response_caches: + del self.response_caches[client_id] + logger.debug(f"Released all {count} objects for client {client_id}") + + def _release_actors(self, client_id): + if client_id not in self.actor_owners: + logger.debug(f"Releasing client with no actors: {client_id}") + return + + count = 0 + actors_to_remove = self.actor_owners.pop(client_id) + for id_bytes in actors_to_remove: + count += 1 + if self._can_remove_actor_ref(id_bytes): + logger.debug(f"Deleting reference to actor {id_bytes.hex()}") + del self.actor_refs[id_bytes] + + logger.debug(f"Released all {count} actors for client: {client_id}") + + @_use_response_cache + def Terminate(self, req, context=None): + if req.WhichOneof("terminate_type") == "task_object": + try: + object_ref = self.object_refs[req.client_id][req.task_object.id] + with disable_client_hook(): + ray.cancel( + object_ref, + force=req.task_object.force, + recursive=req.task_object.recursive, + ) + except Exception as e: + return_exception_in_context(e, context) + elif req.WhichOneof("terminate_type") == "actor": + try: + actor_ref = self.actor_refs[req.actor.id] + with disable_client_hook(): + ray.kill(actor_ref, no_restart=req.actor.no_restart) + except Exception as e: + return_exception_in_context(e, context) + else: + raise RuntimeError( + "Client requested termination without providing a valid terminate_type" + ) + return ray_client_pb2.TerminateResponse(ok=True) + + def _async_get_object( + self, + request: ray_client_pb2.GetRequest, + client_id: str, + req_id: int, + result_queue: queue.Queue, + context=None, + ) -> Optional[ray_client_pb2.GetResponse]: + """Attempts to schedule a callback to push the GetResponse to the + main loop when the desired object is ready. If there is some failure + in scheduling, a GetResponse will be immediately returned. + """ + if len(request.ids) != 1: + raise ValueError( + f"Async get() must have exactly 1 Object ID. Actual: {request}" + ) + rid = request.ids[0] + ref = self.object_refs[client_id].get(rid, None) + if not ref: + return ray_client_pb2.GetResponse( + valid=False, + error=cloudpickle.dumps( + ValueError( + f"ClientObjectRef with id {rid} not found for " + f"client {client_id}" + ) + ), + ) + try: + logger.debug("async get: %s" % ref) + with disable_client_hook(): + + def send_get_response(result: Any) -> None: + """Pushes GetResponses to the main DataPath loop to send + to the client. This is called when the object is ready + on the server side.""" + try: + serialized = dumps_from_server(result, client_id, self) + total_size = len(serialized) + assert total_size > 0, "Serialized object cannot be zero bytes" + total_chunks = math.ceil( + total_size / OBJECT_TRANSFER_CHUNK_SIZE + ) + for chunk_id in range(request.start_chunk_id, total_chunks): + start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE + end = min( + total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE + ) + get_resp = ray_client_pb2.GetResponse( + valid=True, + data=serialized[start:end], + chunk_id=chunk_id, + total_chunks=total_chunks, + total_size=total_size, + ) + chunk_resp = ray_client_pb2.DataResponse( + get=get_resp, req_id=req_id + ) + result_queue.put(chunk_resp) + except Exception as exc: + get_resp = ray_client_pb2.GetResponse( + valid=False, error=cloudpickle.dumps(exc) + ) + resp = ray_client_pb2.DataResponse(get=get_resp, req_id=req_id) + result_queue.put(resp) + + ref._on_completed(send_get_response) + return None + + except Exception as e: + return ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e)) + + def GetObject(self, request: ray_client_pb2.GetRequest, context): + metadata = dict(context.invocation_metadata()) + client_id = metadata.get("client_id") + if client_id is None: + yield ray_client_pb2.GetResponse( + valid=False, + error=cloudpickle.dumps( + ValueError("client_id is not specified in request metadata") + ), + ) + else: + yield from self._get_object(request, client_id) + + def _get_object(self, request: ray_client_pb2.GetRequest, client_id: str): + objectrefs = [] + for rid in request.ids: + ref = self.object_refs[client_id].get(rid, None) + if ref: + objectrefs.append(ref) + else: + yield ray_client_pb2.GetResponse( + valid=False, + error=cloudpickle.dumps( + ValueError( + f"ClientObjectRef {rid} is not found for client {client_id}" + ) + ), + ) + return + try: + logger.debug("get: %s" % objectrefs) + with disable_client_hook(): + items = ray.get(objectrefs, timeout=request.timeout) + except Exception as e: + yield ray_client_pb2.GetResponse(valid=False, error=cloudpickle.dumps(e)) + return + serialized = dumps_from_server(items, client_id, self) + total_size = len(serialized) + assert total_size > 0, "Serialized object cannot be zero bytes" + total_chunks = math.ceil(total_size / OBJECT_TRANSFER_CHUNK_SIZE) + for chunk_id in range(request.start_chunk_id, total_chunks): + start = chunk_id * OBJECT_TRANSFER_CHUNK_SIZE + end = min(total_size, (chunk_id + 1) * OBJECT_TRANSFER_CHUNK_SIZE) + yield ray_client_pb2.GetResponse( + valid=True, + data=serialized[start:end], + chunk_id=chunk_id, + total_chunks=total_chunks, + total_size=total_size, + ) + + def PutObject( + self, request: ray_client_pb2.PutRequest, context=None + ) -> ray_client_pb2.PutResponse: + """gRPC entrypoint for unary PutObject""" + return self._put_object( + request.data, request.client_ref_id, "", request.owner_id, context + ) + + def _put_object( + self, + data: Union[bytes, bytearray], + client_ref_id: bytes, + client_id: str, + owner_id: bytes, + context=None, + ): + """Put an object in the cluster with ray.put() via gRPC. + + Args: + data: Pickled data. Can either be bytearray if this is called + from the dataservicer, or bytes if called from PutObject. + client_ref_id: The id associated with this object on the client. + client_id: The client who owns this data, for tracking when to + delete this reference. + owner_id: The owner id of the object. + context: gRPC context. + """ + try: + obj = loads_from_client(data, self) + + if owner_id: + owner = self.actor_refs[owner_id] + else: + owner = None + with disable_client_hook(): + objectref = ray.put(obj, _owner=owner) + except Exception as e: + logger.exception("Put failed:") + return ray_client_pb2.PutResponse( + id=b"", valid=False, error=cloudpickle.dumps(e) + ) + + self.object_refs[client_id][objectref.binary()] = objectref + if len(client_ref_id) > 0: + self.client_side_ref_map[client_id][client_ref_id] = objectref.binary() + logger.debug("put: %s" % objectref) + return ray_client_pb2.PutResponse(id=objectref.binary(), valid=True) + + def WaitObject(self, request, context=None) -> ray_client_pb2.WaitResponse: + object_refs = [] + for rid in request.object_ids: + if rid not in self.object_refs[request.client_id]: + raise Exception( + "Asking for a ref not associated with this client: %s" % str(rid) + ) + object_refs.append(self.object_refs[request.client_id][rid]) + num_returns = request.num_returns + timeout = request.timeout + try: + with disable_client_hook(): + ready_object_refs, remaining_object_refs = ray.wait( + object_refs, + num_returns=num_returns, + timeout=timeout if timeout != -1 else None, + ) + except Exception as e: + # TODO(ameer): improve exception messages. + logger.error(f"Exception {e}") + return ray_client_pb2.WaitResponse(valid=False) + logger.debug( + "wait: %s %s" % (str(ready_object_refs), str(remaining_object_refs)) + ) + ready_object_ids = [ + ready_object_ref.binary() for ready_object_ref in ready_object_refs + ] + remaining_object_ids = [ + remaining_object_ref.binary() + for remaining_object_ref in remaining_object_refs + ] + return ray_client_pb2.WaitResponse( + valid=True, + ready_object_ids=ready_object_ids, + remaining_object_ids=remaining_object_ids, + ) + + def Schedule( + self, + task: ray_client_pb2.ClientTask, + arglist: List[Any], + kwargs: Dict[str, Any], + context=None, + ) -> ray_client_pb2.ClientTaskTicket: + logger.debug( + "schedule: %s %s" + % (task.name, ray_client_pb2.ClientTask.RemoteExecType.Name(task.type)) + ) + try: + with disable_client_hook(): + if task.type == ray_client_pb2.ClientTask.FUNCTION: + result = self._schedule_function(task, arglist, kwargs, context) + elif task.type == ray_client_pb2.ClientTask.ACTOR: + result = self._schedule_actor(task, arglist, kwargs, context) + elif task.type == ray_client_pb2.ClientTask.METHOD: + result = self._schedule_method(task, arglist, kwargs, context) + elif task.type == ray_client_pb2.ClientTask.NAMED_ACTOR: + result = self._schedule_named_actor(task, context) + else: + raise NotImplementedError( + "Unimplemented Schedule task type: %s" + % ray_client_pb2.ClientTask.RemoteExecType.Name(task.type) + ) + result.valid = True + return result + except Exception as e: + logger.debug("Caught schedule exception", exc_info=True) + return ray_client_pb2.ClientTaskTicket( + valid=False, error=cloudpickle.dumps(e) + ) + + def _schedule_method( + self, + task: ray_client_pb2.ClientTask, + arglist: List[Any], + kwargs: Dict[str, Any], + context=None, + ) -> ray_client_pb2.ClientTaskTicket: + actor_handle = self.actor_refs.get(task.payload_id) + if actor_handle is None: + raise Exception("Can't run an actor the server doesn't have a handle for") + method = getattr(actor_handle, task.name) + opts = decode_options(task.options) + if opts is not None: + method = method.options(**opts) + output = method.remote(*arglist, **kwargs) + ids = self.unify_and_track_outputs(output, task.client_id) + return ray_client_pb2.ClientTaskTicket(return_ids=ids) + + def _schedule_actor( + self, + task: ray_client_pb2.ClientTask, + arglist: List[Any], + kwargs: Dict[str, Any], + context=None, + ) -> ray_client_pb2.ClientTaskTicket: + remote_class = self.lookup_or_register_actor( + task.payload_id, task.client_id, decode_options(task.baseline_options) + ) + opts = decode_options(task.options) + if opts is not None: + remote_class = remote_class.options(**opts) + with current_server(self): + actor = remote_class.remote(*arglist, **kwargs) + self.actor_refs[actor._actor_id.binary()] = actor + self.actor_owners[task.client_id].add(actor._actor_id.binary()) + return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()]) + + def _schedule_function( + self, + task: ray_client_pb2.ClientTask, + arglist: List[Any], + kwargs: Dict[str, Any], + context=None, + ) -> ray_client_pb2.ClientTaskTicket: + remote_func = self.lookup_or_register_func( + task.payload_id, task.client_id, decode_options(task.baseline_options) + ) + opts = decode_options(task.options) + if opts is not None: + remote_func = remote_func.options(**opts) + with current_server(self): + output = remote_func.remote(*arglist, **kwargs) + ids = self.unify_and_track_outputs(output, task.client_id) + return ray_client_pb2.ClientTaskTicket(return_ids=ids) + + def _schedule_named_actor( + self, task: ray_client_pb2.ClientTask, context=None + ) -> ray_client_pb2.ClientTaskTicket: + assert len(task.payload_id) == 0 + # Convert empty string back to None. + actor = ray.get_actor(task.name, task.namespace or None) + bin_actor_id = actor._actor_id.binary() + if bin_actor_id not in self.actor_refs: + self.actor_refs[bin_actor_id] = actor + self.actor_owners[task.client_id].add(bin_actor_id) + self.named_actors.add(bin_actor_id) + return ray_client_pb2.ClientTaskTicket(return_ids=[actor._actor_id.binary()]) + + def lookup_or_register_func( + self, id: bytes, client_id: str, options: Optional[Dict] + ) -> ray.remote_function.RemoteFunction: + with disable_client_hook(): + if id not in self.function_refs: + funcref = self.object_refs[client_id][id] + func = ray.get(funcref) + if not inspect.isfunction(func): + raise Exception( + "Attempting to register function that isn't a function." + ) + if options is None or len(options) == 0: + self.function_refs[id] = ray.remote(func) + else: + self.function_refs[id] = ray.remote(**options)(func) + return self.function_refs[id] + + def lookup_or_register_actor( + self, id: bytes, client_id: str, options: Optional[Dict] + ): + with disable_client_hook(): + if id not in self.registered_actor_classes: + actor_class_ref = self.object_refs[client_id][id] + actor_class = ray.get(actor_class_ref) + if not inspect.isclass(actor_class): + raise Exception("Attempting to schedule actor that isn't a class.") + if options is None or len(options) == 0: + reg_class = ray.remote(actor_class) + else: + reg_class = ray.remote(**options)(actor_class) + self.registered_actor_classes[id] = reg_class + + return self.registered_actor_classes[id] + + def unify_and_track_outputs(self, output, client_id): + if output is None: + outputs = [] + elif isinstance(output, list): + outputs = output + else: + outputs = [output] + for out in outputs: + if out.binary() in self.object_refs[client_id]: + logger.warning(f"Already saw object_ref {out}") + self.object_refs[client_id][out.binary()] = out + return [out.binary() for out in outputs] + + +def return_exception_in_context(err, context): + if context is not None: + context.set_details(encode_exception(err)) + # Note: https://grpc.github.io/grpc/core/md_doc_statuscodes.html + # ABORTED used here since it should never be generated by the + # grpc lib -- this way we know the error was generated by ray logic + context.set_code(grpc.StatusCode.ABORTED) + + +def encode_exception(exception) -> str: + data = cloudpickle.dumps(exception) + return base64.standard_b64encode(data).decode() + + +def decode_options(options: ray_client_pb2.TaskOptions) -> Optional[Dict[str, Any]]: + if not options.pickled_options: + return None + opts = pickle.loads(options.pickled_options) + assert isinstance(opts, dict) + + return opts + + +def serve(host: str, port: int, ray_connect_handler=None): + def default_connect_handler( + job_config: JobConfig = None, **ray_init_kwargs: Dict[str, Any] + ): + with disable_client_hook(): + if not ray.is_initialized(): + return ray.init(job_config=job_config, **ray_init_kwargs) + + from ray._private.grpc_utils import create_grpc_server_with_interceptors + + ray_connect_handler = ray_connect_handler or default_connect_handler + server = create_grpc_server_with_interceptors( + max_workers=CLIENT_SERVER_MAX_THREADS, + thread_name_prefix="ray_client_server", + options=GRPC_OPTIONS, + asynchronous=False, + ) + task_servicer = RayletServicer(ray_connect_handler) + data_servicer = DataServicer(task_servicer) + logs_servicer = LogstreamServicer() + ray_client_pb2_grpc.add_RayletDriverServicer_to_server(task_servicer, server) + ray_client_pb2_grpc.add_RayletDataStreamerServicer_to_server(data_servicer, server) + ray_client_pb2_grpc.add_RayletLogStreamerServicer_to_server(logs_servicer, server) + if not is_localhost(host): + add_port_to_grpc_server(server, f"127.0.0.1:{port}") + add_port_to_grpc_server(server, f"{host}:{port}") + current_handle = ClientServerHandle( + task_servicer=task_servicer, + data_servicer=data_servicer, + logs_servicer=logs_servicer, + grpc_server=server, + ) + server.start() + return current_handle + + +def init_and_serve(host: str, port: int, *args, **kwargs): + with disable_client_hook(): + # Disable client mode inside the worker's environment + info = ray.init(*args, **kwargs) + + def ray_connect_handler(job_config=None, **ray_init_kwargs): + # Ray client will disconnect from ray when + # num_clients == 0. + if ray.is_initialized(): + return info + else: + return ray.init(job_config=job_config, *args, **kwargs) + + server_handle = serve(host, port, ray_connect_handler=ray_connect_handler) + return (server_handle, info) + + +def shutdown_with_server(server, _exiting_interpreter=False): + server.stop(1) + with disable_client_hook(): + ray.shutdown(_exiting_interpreter) + + +def create_ray_handler(address, redis_password, redis_username=None): + def ray_connect_handler(job_config: JobConfig = None, **ray_init_kwargs): + if address: + if redis_password: + ray.init( + address=address, + _redis_username=redis_username, + _redis_password=redis_password, + job_config=job_config, + **ray_init_kwargs, + ) + else: + ray.init(address=address, job_config=job_config, **ray_init_kwargs) + else: + ray.init(job_config=job_config, **ray_init_kwargs) + + return ray_connect_handler + + +def try_create_gcs_client(address: Optional[str]) -> Optional[GcsClient]: + """ + Try to create a gcs client based on the command line args or by + autodetecting a running Ray cluster. + """ + address = canonicalize_bootstrap_address_or_die(address) + return GcsClient(address=address) + + +def main(): + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Host IP to bind to" + ) + parser.add_argument("-p", "--port", type=int, default=10001, help="Port to bind to") + parser.add_argument( + "--mode", + type=str, + choices=["proxy", "legacy", "specific-server"], + default="proxy", + ) + parser.add_argument( + "--address", required=False, type=str, help="Address to use to connect to Ray" + ) + parser.add_argument( + "--redis-username", + required=False, + type=str, + help="username for connecting to Redis", + ) + parser.add_argument( + "--redis-password", + required=False, + type=str, + help="Password for connecting to Redis", + ) + parser.add_argument( + "--runtime-env-agent-address", + required=False, + type=str, + default=None, + help="The port to use for connecting to the runtime_env_agent.", + ) + args, _ = parser.parse_known_args() + setup_logger(ray_constants.LOGGER_LEVEL, ray_constants.LOGGER_FORMAT) + + ray_connect_handler = create_ray_handler( + args.address, args.redis_password, args.redis_username + ) + + hostport = build_address(args.host, args.port) + args_str = str(args) + if args.redis_password: + args_str = args_str.replace(args.redis_password, "****") + logger.info(f"Starting Ray Client server on {hostport}, args {args_str}") + if args.mode == "proxy": + server = serve_proxier( + args.host, + args.port, + args.address, + redis_username=args.redis_username, + redis_password=args.redis_password, + runtime_env_agent_address=args.runtime_env_agent_address, + ) + else: + server = serve(args.host, args.port, ray_connect_handler) + + try: + idle_checks_remaining = TIMEOUT_FOR_SPECIFIC_SERVER_S + while True: + health_report = { + "time": time.time(), + } + + try: + if not ray.experimental.internal_kv._internal_kv_initialized(): + gcs_client = try_create_gcs_client(args.address) + ray.experimental.internal_kv._initialize_internal_kv(gcs_client) + ray.experimental.internal_kv._internal_kv_put( + "ray_client_server", + json.dumps(health_report), + namespace=ray_constants.KV_NAMESPACE_HEALTHCHECK, + ) + except Exception as e: + logger.error( + f"[{args.mode}] Failed to put health check on {args.address}" + ) + logger.exception(e) + + time.sleep(1) + if args.mode == "specific-server": + if server.data_servicer.num_clients > 0: + idle_checks_remaining = TIMEOUT_FOR_SPECIFIC_SERVER_S + else: + idle_checks_remaining -= 1 + if idle_checks_remaining == 0: + raise KeyboardInterrupt() + if ( + idle_checks_remaining % 5 == 0 + and idle_checks_remaining != TIMEOUT_FOR_SPECIFIC_SERVER_S + ): + logger.info(f"{idle_checks_remaining} idle checks before shutdown.") + + except KeyboardInterrupt: + server.stop(0) + + +if __name__ == "__main__": + main() diff --git a/lib/python3.12/site-packages/ray/util/client/server/server_pickler.py b/lib/python3.12/site-packages/ray/util/client/server/server_pickler.py new file mode 100644 index 0000000000000000000000000000000000000000..5211a7991a867997b690f4050dd129e968415929 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/server/server_pickler.py @@ -0,0 +1,124 @@ +"""Implements the client side of the client/server pickling protocol. + +These picklers are aware of the server internals and can find the +references held for the client within the server. + +More discussion about the client/server pickling protocol can be found in: + + ray/util/client/client_pickler.py + +ServerPickler dumps ray objects from the server into the appropriate stubs. +ClientUnpickler loads stubs from the client and finds their associated handle +in the server instance. +""" +import io +from typing import TYPE_CHECKING, Any + +import ray +import ray.cloudpickle as cloudpickle +from ray._private.client_mode_hook import disable_client_hook +from ray.util.client.client_pickler import PickleStub +from ray.util.client.server.server_stubs import ( + ClientReferenceActor, + ClientReferenceFunction, +) + +if TYPE_CHECKING: + from ray.util.client.server.server import RayletServicer + +import pickle # noqa: F401 + + +class ServerPickler(cloudpickle.CloudPickler): + def __init__(self, client_id: str, server: "RayletServicer", *args, **kwargs): + super().__init__(*args, **kwargs) + self.client_id = client_id + self.server = server + + def persistent_id(self, obj): + if isinstance(obj, ray.ObjectRef): + obj_id = obj.binary() + if obj_id not in self.server.object_refs[self.client_id]: + # We're passing back a reference, probably inside a reference. + # Let's hold onto it. + self.server.object_refs[self.client_id][obj_id] = obj + return PickleStub( + type="Object", + client_id=self.client_id, + ref_id=obj_id, + name=None, + baseline_options=None, + ) + elif isinstance(obj, ray.actor.ActorHandle): + actor_id = obj._actor_id.binary() + if actor_id not in self.server.actor_refs: + # We're passing back a handle, probably inside a reference. + self.server.actor_refs[actor_id] = obj + if actor_id not in self.server.actor_owners[self.client_id]: + self.server.actor_owners[self.client_id].add(actor_id) + return PickleStub( + type="Actor", + client_id=self.client_id, + ref_id=obj._actor_id.binary(), + name=None, + baseline_options=None, + ) + return None + + +class ClientUnpickler(pickle.Unpickler): + def __init__(self, server, *args, **kwargs): + super().__init__(*args, **kwargs) + self.server = server + + def persistent_load(self, pid): + assert isinstance(pid, PickleStub) + if pid.type == "Ray": + return ray + elif pid.type == "Object": + return self.server.object_refs[pid.client_id][pid.ref_id] + elif pid.type == "Actor": + return self.server.actor_refs[pid.ref_id] + elif pid.type == "RemoteFuncSelfReference": + return ClientReferenceFunction(pid.client_id, pid.ref_id) + elif pid.type == "RemoteFunc": + return self.server.lookup_or_register_func( + pid.ref_id, pid.client_id, pid.baseline_options + ) + elif pid.type == "RemoteActorSelfReference": + return ClientReferenceActor(pid.client_id, pid.ref_id) + elif pid.type == "RemoteActor": + return self.server.lookup_or_register_actor( + pid.ref_id, pid.client_id, pid.baseline_options + ) + elif pid.type == "RemoteMethod": + actor = self.server.actor_refs[pid.ref_id] + return getattr(actor, pid.name) + else: + raise NotImplementedError("Uncovered client data type") + + +def dumps_from_server( + obj: Any, client_id: str, server_instance: "RayletServicer", protocol=None +) -> bytes: + with io.BytesIO() as file: + sp = ServerPickler(client_id, server_instance, file, protocol=protocol) + sp.dump(obj) + return file.getvalue() + + +def loads_from_client( + data: bytes, + server_instance: "RayletServicer", + *, + fix_imports=True, + encoding="ASCII", + errors="strict" +) -> Any: + with disable_client_hook(): + if isinstance(data, str): + raise TypeError("Can't load pickle from unicode string") + file = io.BytesIO(data) + return ClientUnpickler( + server_instance, file, fix_imports=fix_imports, encoding=encoding + ).load() diff --git a/lib/python3.12/site-packages/ray/util/client/server/server_stubs.py b/lib/python3.12/site-packages/ray/util/client/server/server_stubs.py new file mode 100644 index 0000000000000000000000000000000000000000..020ebf2aeb2cae5e80e624f31262eb4fc2b45467 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/server/server_stubs.py @@ -0,0 +1,66 @@ +from abc import ABC, abstractmethod +from contextlib import contextmanager + +_current_server = None + + +@contextmanager +def current_server(r): + global _current_server + remote = _current_server + _current_server = r + try: + yield + finally: + _current_server = remote + + +class ClientReferenceSentinel(ABC): + def __init__(self, client_id, id): + self.client_id = client_id + self.id = id + + def __reduce__(self): + remote_obj = self.get_remote_obj() + if remote_obj is None: + return (self.__class__, (self.client_id, self.id)) + return (identity, (remote_obj,)) + + @abstractmethod + def get_remote_obj(self): + pass + + def get_real_ref_from_server(self): + global _current_server + if _current_server is None: + return None + client_map = _current_server.client_side_ref_map.get(self.client_id, None) + if client_map is None: + return None + return client_map.get(self.id, None) + + +class ClientReferenceActor(ClientReferenceSentinel): + def get_remote_obj(self): + global _current_server + real_ref_id = self.get_real_ref_from_server() + if real_ref_id is None: + return None + return _current_server.lookup_or_register_actor( + real_ref_id, self.client_id, None + ) + + +class ClientReferenceFunction(ClientReferenceSentinel): + def get_remote_obj(self): + global _current_server + real_ref_id = self.get_real_ref_from_server() + if real_ref_id is None: + return None + return _current_server.lookup_or_register_func( + real_ref_id, self.client_id, None + ) + + +def identity(x): + return x diff --git a/lib/python3.12/site-packages/ray/util/client/worker.py b/lib/python3.12/site-packages/ray/util/client/worker.py new file mode 100644 index 0000000000000000000000000000000000000000..166ec6cb675b972b303a2eedf03647b146128e40 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client/worker.py @@ -0,0 +1,908 @@ +"""This file includes the Worker class which sits on the client side. +It implements the Ray API functions that are forwarded through grpc calls +to the server. +""" + +import base64 +import json +import logging +import os +import tempfile +import threading +import time +import uuid +import warnings +from collections import defaultdict +from concurrent.futures import Future +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union + +import grpc + +import ray.cloudpickle as cloudpickle +import ray.core.generated.ray_client_pb2 as ray_client_pb2 +import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc +from ray._private.ray_constants import ( + DEFAULT_CLIENT_RECONNECT_GRACE_PERIOD, + env_float, + env_integer, +) +from ray._private.runtime_env.py_modules import upload_py_modules_if_needed +from ray._private.runtime_env.working_dir import upload_working_dir_if_needed + +# Use cloudpickle's version of pickle for UnpicklingError +from ray.cloudpickle.compat import pickle +from ray.exceptions import GetTimeoutError +from ray.job_config import JobConfig +from ray.util.client.client_pickler import dumps_from_client, loads_from_server +from ray.util.client.common import ( + GRPC_OPTIONS, + GRPC_UNRECOVERABLE_ERRORS, + INT32_MAX, + OBJECT_TRANSFER_WARNING_SIZE, + ClientActorClass, + ClientActorHandle, + ClientActorRef, + ClientObjectRef, + ClientRemoteFunc, + ClientStub, +) +from ray.util.client.dataclient import DataClient +from ray.util.client.logsclient import LogstreamClient +from ray.util.debug import log_once + +if TYPE_CHECKING: + from ray.actor import ActorClass + from ray.remote_function import RemoteFunction + +logger = logging.getLogger(__name__) + +INITIAL_TIMEOUT_SEC = env_integer("RAY_CLIENT_INITIAL_CONNECTION_TIMEOUT_S", 5) +MAX_TIMEOUT_SEC = env_integer("RAY_CLIENT_MAX_CONNECTION_TIMEOUT_S", 30) +# The max amount of time an operation can run blocking in the server. This +# allows for Ctrl-C of the client to work without explicitly cancelling server +# operations. +MAX_BLOCKING_OPERATION_TIME_S: float = env_float( + "RAY_CLIENT_MAX_BLOCKING_OPERATION_TIME_S", 2.0 +) + +# If the total size (bytes) of all outbound messages to schedule tasks since +# the connection began exceeds this value, a warning should be raised +MESSAGE_SIZE_THRESHOLD = 10 * 2**20 # 10 MB + +# Links to the Ray Design Pattern doc to use in the task overhead warning +# message +DESIGN_PATTERN_FINE_GRAIN_TASKS_LINK = "https://docs.ray.io/en/latest/ray-core/patterns/too-fine-grained-tasks.html" # noqa E501 + +DESIGN_PATTERN_LARGE_OBJECTS_LINK = "https://docs.ray.io/en/latest/ray-core/patterns/closure-capture-large-objects.html" # noqa E501 + + +def backoff(timeout: int) -> int: + timeout = timeout + 5 + if timeout > MAX_TIMEOUT_SEC: + timeout = MAX_TIMEOUT_SEC + return timeout + + +class Worker: + def __init__( + self, + conn_str: str = "", + secure: bool = False, + metadata: List[Tuple[str, str]] = None, + connection_retries: int = 3, + _credentials: Optional[grpc.ChannelCredentials] = None, + ): + """Initializes the worker side grpc client. + + Args: + conn_str: The host:port connection string for the ray server. + secure: whether to use SSL secure channel or not. + metadata: additional metadata passed in the grpc request headers. + connection_retries: Number of times to attempt to reconnect to the + ray server if it doesn't respond immediately. Setting to 0 tries + at least once. For infinite retries, catch the ConnectionError + exception. + _credentials: gprc channel credentials. Default ones will be used + if None. + """ + self._client_id = make_client_id() + self.metadata = [("client_id", self._client_id)] + ( + metadata if metadata else [] + ) + self.channel = None + self.server = None + self._conn_state = grpc.ChannelConnectivity.IDLE + self._converted: Dict[str, ClientStub] = {} + self._secure = secure or os.environ.get("RAY_USE_TLS", "0").lower() in ( + "1", + "true", + ) + self._conn_str = conn_str + self._connection_retries = connection_retries + + if _credentials is not None: + self._credentials = _credentials + self._secure = True + else: + self._credentials = None + + self._reconnect_grace_period = DEFAULT_CLIENT_RECONNECT_GRACE_PERIOD + if "RAY_CLIENT_RECONNECT_GRACE_PERIOD" in os.environ: + # Use value in environment variable if available + self._reconnect_grace_period = int( + os.environ["RAY_CLIENT_RECONNECT_GRACE_PERIOD"] + ) + # Disable retries if grace period is set to 0 + self._reconnect_enabled = self._reconnect_grace_period != 0 + + # Set to True when the connection cannot be recovered and reconnect + # attempts should be stopped + self._in_shutdown = False + # Set to True after initial connection succeeds + self._has_connected = False + + self._connect_channel() + self._has_connected = True + + # Has Ray been initialized on the server? + self._serverside_ray_initialized = False + + # Initialize the streams to finish protocol negotiation. + self.data_client = DataClient(self, self._client_id, self.metadata) + self.reference_count: Dict[bytes, int] = defaultdict(int) + + self.log_client = LogstreamClient(self, self.metadata) + self.log_client.set_logstream_level(logging.INFO) + + self.closed = False + + # Track this value to raise a warning if a lot of data are transferred. + self.total_outbound_message_size_bytes = 0 + + # Used to create unique IDs for RPCs to the RayletServicer + self._req_id_lock = threading.Lock() + self._req_id = 0 + + def _connect_channel(self, reconnecting=False) -> None: + """ + Attempts to connect to the server specified by conn_str. If + reconnecting after an RPC error, cleans up the old channel and + continues to attempt to connect until the grace period is over. + """ + if self.channel is not None: + self.channel.unsubscribe(self._on_channel_state_change) + self.channel.close() + + from ray._private.grpc_utils import init_grpc_channel + + # Prepare credentials if secure connection is requested + credentials = None + if self._secure: + if self._credentials is not None: + credentials = self._credentials + elif os.environ.get("RAY_USE_TLS", "0").lower() in ("1", "true"): + # init_grpc_channel will handle this via load_certs_from_env() + credentials = None + else: + # Default SSL credentials (no specific certs) + credentials = grpc.ssl_channel_credentials() + + # Create channel with auth interceptors via helper + # This automatically adds auth interceptors when token auth is enabled + self.channel = init_grpc_channel( + self._conn_str, + options=GRPC_OPTIONS, + asynchronous=False, + credentials=credentials, + ) + + self.channel.subscribe(self._on_channel_state_change) + + # Retry the connection until the channel responds to something + # looking like a gRPC connection, though it may be a proxy. + start_time = time.time() + conn_attempts = 0 + timeout = INITIAL_TIMEOUT_SEC + service_ready = False + while conn_attempts < max(self._connection_retries, 1) or reconnecting: + conn_attempts += 1 + if self._in_shutdown: + # User manually closed the worker before connection finished + break + elapsed_time = time.time() - start_time + if reconnecting and elapsed_time > self._reconnect_grace_period: + self._in_shutdown = True + raise ConnectionError( + "Failed to reconnect within the reconnection grace period " + f"({self._reconnect_grace_period}s)" + ) + try: + # Let gRPC wait for us to see if the channel becomes ready. + # If it throws, we couldn't connect. + grpc.channel_ready_future(self.channel).result(timeout=timeout) + # The HTTP2 channel is ready. Wrap the channel with the + # RayletDriverStub, allowing for unary requests. + self.server = ray_client_pb2_grpc.RayletDriverStub(self.channel) + service_ready = bool(self.ping_server()) + if service_ready: + break + # Ray is not ready yet, wait a timeout + time.sleep(timeout) + except grpc.FutureTimeoutError: + logger.debug(f"Couldn't connect channel in {timeout} seconds, retrying") + # Note that channel_ready_future constitutes its own timeout, + # which is why we do not sleep here. + except grpc.RpcError as e: + logger.debug( + f"Ray client server unavailable, retrying in {timeout}s..." + ) + logger.debug(f"Received when checking init: {e.details()}") + # Ray is not ready yet, wait a timeout. + time.sleep(timeout) + # Fallthrough, backoff, and retry at the top of the loop + logger.debug( + f"Waiting for Ray to become ready on the server, retry in {timeout}s..." + ) + if not reconnecting: + # Don't increase backoff when trying to reconnect -- + # we already know the server exists, attempt to reconnect + # as soon as we can + timeout = backoff(timeout) + + # If we made it through the loop without service_ready + # it means we've used up our retries and + # should error back to the user. + if not service_ready: + self._in_shutdown = True + if log_once("ray_client_security_groups"): + warnings.warn( + "Ray Client connection timed out. Ensure that " + "the Ray Client port on the head node is reachable " + "from your local machine. See https://docs.ray.io/en" + "/latest/cluster/ray-client.html#step-2-check-ports for " + "more information." + ) + raise ConnectionError("ray client connection timeout") + + def _can_reconnect(self, e: grpc.RpcError) -> bool: + """ + Returns True if the RPC error can be recovered from and a retry is + appropriate, false otherwise. + """ + if not self._reconnect_enabled: + return False + if self._in_shutdown: + # Channel is being shutdown, don't try to reconnect + return False + if e.code() in GRPC_UNRECOVERABLE_ERRORS: + # Unrecoverable error -- These errors are specifically raised + # by the server's application logic + return False + if e.code() == grpc.StatusCode.INTERNAL: + details = e.details() + if details == "Exception serializing request!": + # The client failed tried to send a bad request (for example, + # passing "None" instead of a valid grpc message). Don't + # try to reconnect/retry. + return False + # All other errors can be treated as recoverable + return True + + def _call_stub(self, stub_name: str, *args, **kwargs) -> Any: + """ + Calls the stub specified by stub_name (Schedule, WaitObject, etc...). + If a recoverable error occurrs while calling the stub, attempts to + retry the RPC. + """ + while not self._in_shutdown: + try: + return getattr(self.server, stub_name)(*args, **kwargs) + except grpc.RpcError as e: + if self._can_reconnect(e): + time.sleep(0.5) + continue + raise + except ValueError: + # Trying to use the stub on a cancelled channel will raise + # ValueError. This should only happen when the data client + # is attempting to reset the connection -- sleep and try + # again. + time.sleep(0.5) + continue + raise ConnectionError("Client is shutting down.") + + def _get_object_iterator( + self, req: ray_client_pb2.GetRequest, *args, **kwargs + ) -> Any: + """ + Calls the stub for GetObject on the underlying server stub. If a + recoverable error occurs while streaming the response, attempts + to retry the get starting from the first chunk that hasn't been + received. + """ + last_seen_chunk = -1 + while not self._in_shutdown: + # If we disconnect partway through, restart the get request + # at the first chunk we haven't seen + req.start_chunk_id = last_seen_chunk + 1 + try: + for chunk in self.server.GetObject(req, *args, **kwargs): + if chunk.chunk_id <= last_seen_chunk: + # Ignore repeat chunks + logger.debug( + f"Received a repeated chunk {chunk.chunk_id} " + f"from request {req.req_id}." + ) + continue + if last_seen_chunk + 1 != chunk.chunk_id: + raise RuntimeError( + f"Received chunk {chunk.chunk_id} when we expected " + f"{self.last_seen_chunk + 1}" + ) + last_seen_chunk = chunk.chunk_id + yield chunk + if last_seen_chunk == chunk.total_chunks - 1: + # We've yielded the last chunk, exit early + return + return + except grpc.RpcError as e: + if self._can_reconnect(e): + time.sleep(0.5) + continue + raise + except ValueError: + # Trying to use the stub on a cancelled channel will raise + # ValueError. This should only happen when the data client + # is attempting to reset the connection -- sleep and try + # again. + time.sleep(0.5) + continue + raise ConnectionError("Client is shutting down.") + + def _add_ids_to_metadata(self, metadata: Any): + """ + Adds a unique req_id and the current thread's identifier to the + metadata. These values are useful for preventing mutating operations + from being replayed on the server side in the event that the client + must retry a requsest. + Args: + metadata: the gRPC metadata to append the IDs to + """ + if not self._reconnect_enabled: + # IDs not needed if the reconnects are disabled + return metadata + thread_id = str(threading.get_ident()) + with self._req_id_lock: + self._req_id += 1 + if self._req_id > INT32_MAX: + self._req_id = 1 + req_id = str(self._req_id) + return metadata + [("thread_id", thread_id), ("req_id", req_id)] + + def _on_channel_state_change(self, conn_state: grpc.ChannelConnectivity): + logger.debug(f"client gRPC channel state change: {conn_state}") + self._conn_state = conn_state + + def connection_info(self): + try: + data = self.data_client.ConnectionInfo() + except grpc.RpcError as e: + raise decode_exception(e) + return { + "num_clients": data.num_clients, + "python_version": data.python_version, + "ray_version": data.ray_version, + "ray_commit": data.ray_commit, + } + + def register_callback( + self, + ref: ClientObjectRef, + callback: Callable[[ray_client_pb2.DataResponse], None], + ) -> None: + req = ray_client_pb2.GetRequest(ids=[ref.id], asynchronous=True) + self.data_client.RegisterGetCallback(req, callback) + + def get(self, vals, *, timeout: Optional[float] = None) -> Any: + if isinstance(vals, list): + if not vals: + return [] + to_get = vals + elif isinstance(vals, ClientObjectRef): + to_get = [vals] + else: + raise Exception( + "Can't get something that's not a " + "list of IDs or just an ID: %s" % type(vals) + ) + + if timeout is None: + deadline = None + else: + deadline = time.monotonic() + timeout + + while True: + if deadline: + op_timeout = min( + MAX_BLOCKING_OPERATION_TIME_S, + max(deadline - time.monotonic(), 0.001), + ) + else: + op_timeout = MAX_BLOCKING_OPERATION_TIME_S + try: + res = self._get(to_get, op_timeout) + break + except GetTimeoutError: + if deadline and time.monotonic() > deadline: + raise + logger.debug("Internal retry for get {}".format(to_get)) + if len(to_get) != len(res): + raise Exception( + "Mismatched number of items in request ({}) and response ({})".format( + len(to_get), len(res) + ) + ) + if isinstance(vals, ClientObjectRef): + res = res[0] + return res + + def _get(self, ref: List[ClientObjectRef], timeout: float): + req = ray_client_pb2.GetRequest(ids=[r.id for r in ref], timeout=timeout) + data = bytearray() + try: + resp = self._get_object_iterator(req, metadata=self.metadata) + for chunk in resp: + if not chunk.valid: + try: + err = cloudpickle.loads(chunk.error) + except (pickle.UnpicklingError, TypeError): + logger.exception("Failed to deserialize {}".format(chunk.error)) + raise + raise err + if chunk.total_size > OBJECT_TRANSFER_WARNING_SIZE and log_once( + "client_object_transfer_size_warning" + ): + size_gb = chunk.total_size / 2**30 + warnings.warn( + "Ray Client is attempting to retrieve a " + f"{size_gb:.2f} GiB object over the network, which may " + "be slow. Consider serializing the object to a file " + "and using S3 or rsync instead.", + UserWarning, + stacklevel=5, + ) + data.extend(chunk.data) + except grpc.RpcError as e: + raise decode_exception(e) + return loads_from_server(data) + + def put( + self, + val, + *, + client_ref_id: bytes = None, + _owner: Optional[ClientActorHandle] = None, + ): + if isinstance(val, ClientObjectRef): + raise TypeError( + "Calling 'put' on an ObjectRef is not allowed " + "(similarly, returning an ObjectRef from a remote " + "function is not allowed). If you really want to " + "do this, you can wrap the ObjectRef in a list and " + "call 'put' on it (or return it)." + ) + data = dumps_from_client(val, self._client_id) + return self._put_pickled(data, client_ref_id, _owner) + + def _put_pickled( + self, data, client_ref_id: bytes, owner: Optional[ClientActorHandle] = None + ): + req = ray_client_pb2.PutRequest(data=data) + if client_ref_id is not None: + req.client_ref_id = client_ref_id + if owner is not None: + req.owner_id = owner.actor_ref.id + + resp = self.data_client.PutObject(req) + if not resp.valid: + try: + raise cloudpickle.loads(resp.error) + except (pickle.UnpicklingError, TypeError): + logger.exception("Failed to deserialize {}".format(resp.error)) + raise + return ClientObjectRef(resp.id) + + # TODO(ekl) respect MAX_BLOCKING_OPERATION_TIME_S for wait too + def wait( + self, + object_refs: List[ClientObjectRef], + *, + num_returns: int = 1, + timeout: float = None, + fetch_local: bool = True, + ) -> Tuple[List[ClientObjectRef], List[ClientObjectRef]]: + if not isinstance(object_refs, list): + raise TypeError( + f"wait() expected a list of ClientObjectRef, got {type(object_refs)}" + ) + for ref in object_refs: + if not isinstance(ref, ClientObjectRef): + raise TypeError( + "wait() expected a list of ClientObjectRef, " + f"got list containing {type(ref)}" + ) + data = { + "object_ids": [object_ref.id for object_ref in object_refs], + "num_returns": num_returns, + "timeout": timeout if (timeout is not None) else -1, + "client_id": self._client_id, + } + req = ray_client_pb2.WaitRequest(**data) + resp = self._call_stub("WaitObject", req, metadata=self.metadata) + if not resp.valid: + # TODO(ameer): improve error/exceptions messages. + raise Exception("Client Wait request failed. Reference invalid?") + client_ready_object_ids = [ + ClientObjectRef(ref) for ref in resp.ready_object_ids + ] + client_remaining_object_ids = [ + ClientObjectRef(ref) for ref in resp.remaining_object_ids + ] + + return (client_ready_object_ids, client_remaining_object_ids) + + def call_remote(self, instance, *args, **kwargs) -> List[Future]: + task = instance._prepare_client_task() + # data is serialized tuple of (args, kwargs) + task.data = dumps_from_client((args, kwargs), self._client_id) + num_returns = instance._num_returns() + if num_returns == "dynamic": + num_returns = -1 + if num_returns == "streaming": + raise RuntimeError( + 'Streaming actor methods (num_returns="streaming") ' + "are not currently supported when using Ray Client." + ) + + return self._call_schedule_for_task(task, num_returns) + + def _call_schedule_for_task( + self, task: ray_client_pb2.ClientTask, num_returns: Optional[int] + ) -> List[Future]: + logger.debug(f"Scheduling task {task.name} {task.type} {task.payload_id}") + task.client_id = self._client_id + if num_returns is None: + num_returns = 1 + + num_return_refs = num_returns + if num_return_refs == -1: + num_return_refs = 1 + id_futures = [Future() for _ in range(num_return_refs)] + + def populate_ids(resp: Union[ray_client_pb2.DataResponse, Exception]) -> None: + if isinstance(resp, Exception): + if isinstance(resp, grpc.RpcError): + resp = decode_exception(resp) + for future in id_futures: + future.set_exception(resp) + return + + ticket = resp.task_ticket + if not ticket.valid: + try: + ex = cloudpickle.loads(ticket.error) + except (pickle.UnpicklingError, TypeError) as e_new: + ex = e_new + for future in id_futures: + future.set_exception(ex) + return + + if len(ticket.return_ids) != num_return_refs: + exc = ValueError( + f"Expected {num_return_refs} returns but received " + f"{len(ticket.return_ids)}" + ) + for future, raw_id in zip(id_futures, ticket.return_ids): + future.set_exception(exc) + return + + for future, raw_id in zip(id_futures, ticket.return_ids): + future.set_result(raw_id) + + self.data_client.Schedule(task, populate_ids) + + self.total_outbound_message_size_bytes += task.ByteSize() + if ( + self.total_outbound_message_size_bytes > MESSAGE_SIZE_THRESHOLD + and log_once("client_communication_overhead_warning") + ): + warnings.warn( + "More than 10MB of messages have been created to schedule " + "tasks on the server. This can be slow on Ray Client due to " + "communication overhead over the network. If you're running " + "many fine-grained tasks, consider running them inside a " + 'single remote function. See the section on "Too ' + 'fine-grained tasks" in the Ray Design Patterns document for ' + f"more details: {DESIGN_PATTERN_FINE_GRAIN_TASKS_LINK}. If " + "your functions frequently use large objects, consider " + "storing the objects remotely with ray.put. An example of " + 'this is shown in the "Closure capture of large / ' + 'unserializable object" section of the Ray Design Patterns ' + "document, available here: " + f"{DESIGN_PATTERN_LARGE_OBJECTS_LINK}", + UserWarning, + ) + return id_futures + + def call_release(self, id: bytes) -> None: + if self.closed: + return + self.reference_count[id] -= 1 + if self.reference_count[id] == 0: + self._release_server(id) + del self.reference_count[id] + + def _release_server(self, id: bytes) -> None: + if self.data_client is not None: + logger.debug(f"Releasing {id.hex()}") + self.data_client.ReleaseObject(ray_client_pb2.ReleaseRequest(ids=[id])) + + def call_retain(self, id: bytes) -> None: + logger.debug(f"Retaining {id.hex()}") + self.reference_count[id] += 1 + + def close(self): + self._in_shutdown = True + self.closed = True + self.data_client.close() + self.log_client.close() + self.server = None + if self.channel: + self.channel.close() + self.channel = None + + def get_actor( + self, name: str, namespace: Optional[str] = None + ) -> ClientActorHandle: + task = ray_client_pb2.ClientTask() + task.type = ray_client_pb2.ClientTask.NAMED_ACTOR + task.name = name + task.namespace = namespace or "" + # Populate task.data with empty args and kwargs + task.data = dumps_from_client(([], {}), self._client_id) + futures = self._call_schedule_for_task(task, 1) + assert len(futures) == 1 + handle = ClientActorHandle(ClientActorRef(futures[0], weak_ref=True)) + # `actor_ref.is_nil()` waits until the underlying ID is resolved. + # This is needed because `get_actor` is often used to check the + # existence of an actor. + if handle.actor_ref.is_nil(): + raise ValueError(f"ActorID for {name} is empty") + return handle + + def terminate_actor(self, actor: ClientActorHandle, no_restart: bool) -> None: + if not isinstance(actor, ClientActorHandle): + raise ValueError( + "ray.kill() only supported for actors. Got: {}.".format(type(actor)) + ) + term_actor = ray_client_pb2.TerminateRequest.ActorTerminate() + term_actor.id = actor.actor_ref.id + term_actor.no_restart = no_restart + term = ray_client_pb2.TerminateRequest(actor=term_actor) + term.client_id = self._client_id + try: + self.data_client.Terminate(term) + except grpc.RpcError as e: + raise decode_exception(e) + + def terminate_task( + self, obj: ClientObjectRef, force: bool, recursive: bool + ) -> None: + if not isinstance(obj, ClientObjectRef): + raise TypeError( + "ray.cancel() only supported for non-actor object refs. " + f"Got: {type(obj)}." + ) + term_object = ray_client_pb2.TerminateRequest.TaskObjectTerminate() + term_object.id = obj.id + term_object.force = force + term_object.recursive = recursive + term = ray_client_pb2.TerminateRequest(task_object=term_object) + term.client_id = self._client_id + try: + self.data_client.Terminate(term) + except grpc.RpcError as e: + raise decode_exception(e) + + def get_cluster_info( + self, + req_type: ray_client_pb2.ClusterInfoType.TypeEnum, + timeout: Optional[float] = None, + ): + req = ray_client_pb2.ClusterInfoRequest() + req.type = req_type + resp = self.server.ClusterInfo(req, timeout=timeout, metadata=self.metadata) + if resp.WhichOneof("response_type") == "resource_table": + # translate from a proto map to a python dict + output_dict = dict(resp.resource_table.table) + return output_dict + elif resp.WhichOneof("response_type") == "runtime_context": + return resp.runtime_context + return json.loads(resp.json) + + def internal_kv_get(self, key: bytes, namespace: Optional[bytes]) -> bytes: + req = ray_client_pb2.KVGetRequest(key=key, namespace=namespace) + try: + resp = self._call_stub("KVGet", req, metadata=self.metadata) + except grpc.RpcError as e: + raise decode_exception(e) + if resp.HasField("value"): + return resp.value + # Value is None when the key does not exist in the KV. + return None + + def internal_kv_exists(self, key: bytes, namespace: Optional[bytes]) -> bool: + req = ray_client_pb2.KVExistsRequest(key=key, namespace=namespace) + try: + resp = self._call_stub("KVExists", req, metadata=self.metadata) + except grpc.RpcError as e: + raise decode_exception(e) + return resp.exists + + def internal_kv_put( + self, key: bytes, value: bytes, overwrite: bool, namespace: Optional[bytes] + ) -> bool: + req = ray_client_pb2.KVPutRequest( + key=key, value=value, overwrite=overwrite, namespace=namespace + ) + metadata = self._add_ids_to_metadata(self.metadata) + try: + resp = self._call_stub("KVPut", req, metadata=metadata) + except grpc.RpcError as e: + raise decode_exception(e) + return resp.already_exists + + def internal_kv_del( + self, key: bytes, del_by_prefix: bool, namespace: Optional[bytes] + ) -> int: + req = ray_client_pb2.KVDelRequest( + key=key, del_by_prefix=del_by_prefix, namespace=namespace + ) + metadata = self._add_ids_to_metadata(self.metadata) + try: + resp = self._call_stub("KVDel", req, metadata=metadata) + except grpc.RpcError as e: + raise decode_exception(e) + return resp.deleted_num + + def internal_kv_list( + self, prefix: bytes, namespace: Optional[bytes] + ) -> List[bytes]: + try: + req = ray_client_pb2.KVListRequest(prefix=prefix, namespace=namespace) + return self._call_stub("KVList", req, metadata=self.metadata).keys + except grpc.RpcError as e: + raise decode_exception(e) + + def pin_runtime_env_uri(self, uri: str, expiration_s: int) -> None: + req = ray_client_pb2.ClientPinRuntimeEnvURIRequest( + uri=uri, expiration_s=expiration_s + ) + self._call_stub("PinRuntimeEnvURI", req, metadata=self.metadata) + + def list_named_actors(self, all_namespaces: bool) -> List[Dict[str, str]]: + req = ray_client_pb2.ClientListNamedActorsRequest(all_namespaces=all_namespaces) + return json.loads(self.data_client.ListNamedActors(req).actors_json) + + def is_initialized(self) -> bool: + if not self.is_connected() or self.server is None: + return False + if not self._serverside_ray_initialized: + # We only check that Ray is initialized on the server once to + # avoid making an RPC every time this function is called. This is + # safe to do because Ray only 'un-initializes' on the server when + # the Client connection is torn down. + self._serverside_ray_initialized = self.get_cluster_info( + ray_client_pb2.ClusterInfoType.IS_INITIALIZED + ) + + return self._serverside_ray_initialized + + def ping_server(self, timeout=None) -> bool: + """Simple health check. + + Piggybacks the IS_INITIALIZED call to check if the server provides + an actual response. + """ + if self.server is not None: + logger.debug("Pinging server.") + result = self.get_cluster_info( + ray_client_pb2.ClusterInfoType.PING, timeout=timeout + ) + return result is not None + return False + + def is_connected(self) -> bool: + return not self._in_shutdown and self._has_connected + + def _server_init( + self, job_config: JobConfig, ray_init_kwargs: Optional[Dict[str, Any]] = None + ): + """Initialize the server""" + if ray_init_kwargs is None: + ray_init_kwargs = {} + try: + if job_config is None: + serialized_job_config = None + else: + with tempfile.TemporaryDirectory() as tmp_dir: + runtime_env = job_config.runtime_env or {} + runtime_env = upload_py_modules_if_needed( + runtime_env, tmp_dir, logger=logger + ) + runtime_env = upload_working_dir_if_needed( + runtime_env, tmp_dir, logger=logger + ) + # Remove excludes, it isn't relevant after the upload step. + runtime_env.pop("excludes", None) + job_config.set_runtime_env(runtime_env, validate=True) + + serialized_job_config = pickle.dumps(job_config) + + response = self.data_client.Init( + ray_client_pb2.InitRequest( + job_config=serialized_job_config, + ray_init_kwargs=json.dumps(ray_init_kwargs), + reconnect_grace_period=self._reconnect_grace_period, + ) + ) + if not response.ok: + raise ConnectionAbortedError( + f"Initialization failure from server:\n{response.msg}" + ) + + except grpc.RpcError as e: + raise decode_exception(e) + + def _convert_actor(self, actor: "ActorClass") -> str: + """Register a ClientActorClass for the ActorClass and return a UUID""" + key = uuid.uuid4().hex + cls = actor.__ray_metadata__.modified_class + self._converted[key] = ClientActorClass(cls, options=actor._default_options) + return key + + def _convert_function(self, func: "RemoteFunction") -> str: + """Register a ClientRemoteFunc for the ActorClass and return a UUID""" + key = uuid.uuid4().hex + self._converted[key] = ClientRemoteFunc( + func._function, options=func._default_options + ) + return key + + def _get_converted(self, key: str) -> "ClientStub": + """Given a UUID, return the converted object""" + return self._converted[key] + + def _converted_key_exists(self, key: str) -> bool: + """Check if a key UUID is present in the store of converted objects.""" + return key in self._converted + + def _dumps_from_client(self, val) -> bytes: + return dumps_from_client(val, self._client_id) + + +def make_client_id() -> str: + id = uuid.uuid4() + return id.hex + + +def decode_exception(e: grpc.RpcError) -> Exception: + if e.code() != grpc.StatusCode.ABORTED: + # The ABORTED status code is used by the server when an application + # error is serialized into the exception details. If the code + # isn't ABORTED, then return the original error since there's no + # serialized error to decode. + # See server.py::return_exception_in_context for details + return ConnectionError(f"GRPC connection failed: {e}") + data = base64.standard_b64decode(e.details()) + return loads_from_server(data) diff --git a/lib/python3.12/site-packages/ray/util/client_connect.py b/lib/python3.12/site-packages/ray/util/client_connect.py new file mode 100644 index 0000000000000000000000000000000000000000..8c64459a1436224e40c559ac1cd3957859a8e63f --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/client_connect.py @@ -0,0 +1,76 @@ +import logging +from typing import Any, Dict, List, Optional, Tuple + +from ray._private.client_mode_hook import ( + _explicitly_enable_client_mode, + _set_client_hook_status, +) +from ray._private.utils import get_ray_doc_version +from ray.job_config import JobConfig +from ray.util.annotations import Deprecated +from ray.util.client import ray + +logger = logging.getLogger(__name__) + + +@Deprecated( + message="Use ray.init(ray://:) " + "instead. See detailed usage at {}.".format( + f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-init" # noqa: E501 + ) +) +def connect( + conn_str: str, + secure: bool = False, + metadata: List[Tuple[str, str]] = None, + connection_retries: int = 3, + job_config: JobConfig = None, + namespace: str = None, + *, + ignore_version: bool = False, + _credentials: Optional["grpc.ChannelCredentials"] = None, # noqa: F821 + ray_init_kwargs: Optional[Dict[str, Any]] = None, +) -> Dict[str, Any]: + if ray.is_connected(): + ignore_reinit_error = ray_init_kwargs.get("ignore_reinit_error", False) + if ignore_reinit_error: + logger.info( + "Calling ray.init() again after it has already been called. " + "Reusing the existing Ray client connection." + ) + return ray.get_context().client_worker.connection_info() + raise RuntimeError( + "Ray Client is already connected. Maybe you called " + 'ray.init("ray://
") twice by accident?' + ) + + # Enable the same hooks that RAY_CLIENT_MODE does, as calling + # ray.init("ray://
") is specifically for using client mode. + _set_client_hook_status(True) + _explicitly_enable_client_mode() + + # TODO(barakmich): https://github.com/ray-project/ray/issues/13274 + # for supporting things like cert_path, ca_path, etc and creating + # the correct metadata + conn = ray.connect( + conn_str, + job_config=job_config, + secure=secure, + metadata=metadata, + connection_retries=connection_retries, + namespace=namespace, + ignore_version=ignore_version, + _credentials=_credentials, + ray_init_kwargs=ray_init_kwargs, + ) + return conn + + +@Deprecated( + message="Use ray.shutdown() instead. See detailed usage at {}.".format( + f"https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/package-ref.html#ray-shutdown" # noqa: E501 + ) +) +def disconnect(): + """Disconnects from server; is idempotent.""" + return ray.disconnect() diff --git a/lib/python3.12/site-packages/ray/util/collective/__init__.py b/lib/python3.12/site-packages/ray/util/collective/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..09423ad37c11c222d33a7e5e53c045927065a774 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/__init__.py @@ -0,0 +1,53 @@ +from ray.util.collective.collective import ( + allgather, + allgather_multigpu, + allreduce, + allreduce_multigpu, + barrier, + broadcast, + broadcast_multigpu, + create_collective_group, + destroy_collective_group, + get_collective_group_size, + get_group_handle, + get_rank, + gloo_available, + init_collective_group, + is_group_initialized, + nccl_available, + recv, + recv_multigpu, + reduce, + reduce_multigpu, + reducescatter, + reducescatter_multigpu, + send, + send_multigpu, +) + +__all__ = [ + "nccl_available", + "gloo_available", + "is_group_initialized", + "init_collective_group", + "destroy_collective_group", + "create_collective_group", + "get_rank", + "get_collective_group_size", + "allreduce", + "allreduce_multigpu", + "barrier", + "reduce", + "reduce_multigpu", + "broadcast", + "broadcast_multigpu", + "allgather", + "allgather_multigpu", + "reducescatter", + "reducescatter_multigpu", + "send", + "send_multigpu", + "recv", + "recv_multigpu", + "get_group_handle", +] diff --git a/lib/python3.12/site-packages/ray/util/collective/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83b70416f1d665bf076b2277245bdad111574500 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/__pycache__/collective.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/__pycache__/collective.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..308a38b448d8d840828eef558732e13cd5212ea4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/__pycache__/collective.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/__pycache__/const.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/__pycache__/const.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9724f9b1d5774ebf521f1189e8e982f91ac3fb3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/__pycache__/const.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/__pycache__/types.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/__pycache__/types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea78599c4fc9cd2d71add378d984e197c5f81ea5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/__pycache__/types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92b4076b48f7e148301e9d5f1b68b5fee431528e Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/collective.py b/lib/python3.12/site-packages/ray/util/collective/collective.py new file mode 100644 index 0000000000000000000000000000000000000000..5e1681015c252b63cf2db0dc407d1d972ae1ce9b --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/collective.py @@ -0,0 +1,897 @@ +"""APIs exposed under the namespace ray.util.collective.""" + +import logging +import os +import threading +import time +from typing import List + +import numpy as np + +import ray +import ray.experimental.internal_kv as _internal_kv +from . import types +from ray.experimental.collective.util import ( + get_address_and_port as _get_address_and_port, +) +from ray.util.collective.collective_group.torch_gloo_collective_group import ( + get_master_address_metadata_key as _get_master_addr_key, +) + +logger = logging.getLogger(__name__) + +try: + from ray.util.collective.collective_group.nccl_collective_group import NCCLGroup + + _NCCL_AVAILABLE = True + _LOG_NCCL_WARNING = False +except ImportError: + _NCCL_AVAILABLE = False + _LOG_NCCL_WARNING = True + + +try: + from ray.util.collective.collective_group.torch_gloo_collective_group import ( + TorchGLOOGroup, + ) + + _TORCH_DISTRIBUTED_AVAILABLE = True +except ImportError: + _TORCH_DISTRIBUTED_AVAILABLE = False + +try: + from ray.util.collective.collective_group.nixl_backend import NixlBackend + + _NIXL_AVAILABLE = True +except ImportError: + _NIXL_AVAILABLE = False + + +def nccl_available(): + global _LOG_NCCL_WARNING + if ray.get_gpu_ids() and _LOG_NCCL_WARNING: + logger.warning( + "NCCL seems unavailable. Please install Cupy " + "following the guide at: " + "https://docs.cupy.dev/en/stable/install.html." + ) + _LOG_NCCL_WARNING = False + return _NCCL_AVAILABLE + + +def gloo_available(): + # Since we use torch_gloo as the backend for Gloo, + # we can just return the availability of torch.distributed. + return _TORCH_DISTRIBUTED_AVAILABLE + + +def torch_distributed_available(): + return _TORCH_DISTRIBUTED_AVAILABLE + + +def nixl_available(): + return _NIXL_AVAILABLE + + +class GroupManager(object): + """Use this class to manage the collective groups we created so far. + + Each process will have an instance of `GroupManager`. Each process + could belong to multiple collective groups. The membership information + and other metadata are stored in the global `_group_mgr` object. + """ + + def __init__(self): + self._name_group_map = {} + + def create_collective_group( + self, backend, world_size, rank, group_name, gloo_timeout + ): + """The entry to create new collective groups in the manager. + + Put the registration and the group information into the manager + metadata as well. + """ + backend = types.Backend(backend) + if backend == types.Backend.MPI: + raise RuntimeError("Ray does not support MPI.") + elif backend == types.Backend.GLOO or backend == types.Backend.TORCH_GLOO: + # Rendezvous: ensure a MASTER_ADDR:MASTER_PORT is published in internal_kv. + metadata_key = _get_master_addr_key(group_name) + if rank == 0: + addr, port = _get_address_and_port() + _internal_kv._internal_kv_put(metadata_key, f"{addr}:{port}") + else: + # Wait until rank 0 publishes the metadata or timeout. + deadline_s = time.time() + ( + gloo_timeout / 1000.0 if gloo_timeout else 30.0 + ) + while True: + meta = _internal_kv._internal_kv_get(metadata_key) + if meta is not None: + break + if time.time() > deadline_s: + raise TimeoutError( + f"Timed out waiting for GLOO rendezvous metadata for group '{group_name}'." + ) + time.sleep(0.05) + + logger.debug( + "Creating torch.distributed GLOO group: '{}'...".format(group_name) + ) + g = TorchGLOOGroup(world_size, rank, group_name, gloo_timeout) + elif backend == types.Backend.NCCL: + _check_backend_availability(backend) + logger.debug("Creating NCCL group: '{}'...".format(group_name)) + g = NCCLGroup(world_size, rank, group_name) + elif backend == types.Backend.NIXL: + _check_backend_availability(backend) + logger.debug("Creating NIXL Backend: '{}'...".format(group_name)) + g = NixlBackend() + else: + raise RuntimeError(f"Unexpected backend: {backend}") + + self._name_group_map[group_name] = g + return self._name_group_map[group_name] + + def is_group_exist(self, group_name): + return group_name in self._name_group_map + + def get_group_by_name(self, group_name): + """Get the collective group handle by its name.""" + if not self.is_group_exist(group_name): + logger.warning("The group '{}' is not initialized.".format(group_name)) + return None + return self._name_group_map[group_name] + + def destroy_collective_group(self, group_name): + """Group destructor.""" + if not self.is_group_exist(group_name): + logger.warning("The group '{}' does not exist.".format(group_name)) + return + + # release the collective group resource + g = self._name_group_map[group_name] + # clean up the dicts + del self._name_group_map[group_name] + # Release the communicator resources + g.destroy_group() + + # Release the detached actors spawned by `create_collective_group()` + name = "info_" + group_name + try: + store = ray.get_actor(name) + ray.kill(store) + except ValueError: + pass + + +_group_mgr = GroupManager() +# This lock is used to make external calls to the _group_mgr thread-safe. +_group_mgr_lock = threading.Lock() + + +def is_group_initialized(group_name): + """Check if the group is initialized in this process by the group name.""" + global _group_mgr + global _group_mgr_lock + with _group_mgr_lock: + return _group_mgr.is_group_exist(group_name) + + +def init_collective_group( + world_size: int, + rank: int, + backend=types.Backend.NCCL, + group_name: str = "default", + gloo_timeout: int = 30000, +): + """Initialize a collective group inside an actor process. + + Args: + world_size: the total number of processes in the group. + rank: the rank of the current process. + backend: the CCL backend to use, NCCL or GLOO. + group_name: the name of the collective group. + + Returns: + None + """ + _check_inside_actor() + backend = types.Backend(backend) + _check_backend_availability(backend) + global _group_mgr + global _group_mgr_lock + + # TODO(Hao): implement a group auto-counter. + if not group_name: + raise ValueError("group_name '{}' needs to be a string.".format(group_name)) + + with _group_mgr_lock: + if _group_mgr.is_group_exist(group_name): + raise RuntimeError("Trying to initialize a group twice.") + + assert world_size > 0 + assert rank >= 0 + assert rank < world_size + _group_mgr.create_collective_group( + backend, world_size, rank, group_name, gloo_timeout + ) + + +def create_collective_group( + actors, + world_size: int, + ranks: List[int], + backend=types.Backend.NCCL, + group_name: str = "default", + gloo_timeout: int = 30000, +): + """Declare a list of actors as a collective group. + + Note: This function should be called in a driver process. + + Args: + actors: a list of actors to be set in a collective group. + world_size: the total number of processes in the group. + ranks (List[int]): the rank of each actor. + backend: the CCL backend to use, NCCL or GLOO. + group_name: the name of the collective group. + + Returns: + None + """ + backend = types.Backend(backend) + _check_backend_availability(backend) + + name = "info_" + group_name + try: + ray.get_actor(name) + raise RuntimeError("Trying to initialize a group twice.") + except ValueError: + pass + + if len(ranks) != len(actors): + raise RuntimeError( + "Each actor should correspond to one rank. Got '{}' " + "ranks but '{}' actors".format(len(ranks), len(actors)) + ) + + if set(ranks) != set(range(len(ranks))): + raise RuntimeError( + "Ranks must be a permutation from 0 to '{}'. Got '{}'.".format( + len(ranks), "".join([str(r) for r in ranks]) + ) + ) + + if world_size <= 0: + raise RuntimeError( + "World size must be greater than zero. Got '{}'.".format(world_size) + ) + if not all(ranks) >= 0: + raise RuntimeError("Ranks must be non-negative.") + if not all(ranks) < world_size: + raise RuntimeError("Ranks cannot be greater than world_size.") + + # avoid a circular dependency + from ray.util.collective.util import Info + + # store the information into a NamedActor that can be accessed later. + name = "info_" + group_name + actors_id = [a._ray_actor_id for a in actors] + # TODO (Dacheng): how do we recycle this name actor? + info = Info.options(name=name, lifetime="detached").remote() + ray.get([info.set_info.remote(actors_id, world_size, ranks, backend, gloo_timeout)]) + + +# TODO (we need a declarative destroy() API here.) +def destroy_collective_group(group_name: str = "default") -> None: + """Destroy a collective group given its group name.""" + _check_inside_actor() + global _group_mgr + global _group_mgr_lock + with _group_mgr_lock: + _group_mgr.destroy_collective_group(group_name) + + +def get_rank(group_name: str = "default") -> int: + """Return the rank of this process in the given group. + + Args: + group_name: the name of the group to query + + Returns: + the rank of this process in the named group, + -1 if the group does not exist or the process does + not belong to the group. + """ + _check_inside_actor() + + global _group_mgr + global _group_mgr_lock + with _group_mgr_lock: + if not _group_mgr.is_group_exist(group_name): + return -1 + g = _group_mgr.get_group_by_name(group_name) + return g.rank + + +def get_collective_group_size(group_name: str = "default") -> int: + """Return the size of the collective group with the given name. + + Args: + group_name: the name of the group to query + + Returns: + The world size of the collective group, -1 if the group does + not exist or the process does not belong to the group. + """ + _check_inside_actor() + global _group_mgr + global _group_mgr_lock + with _group_mgr_lock: + if not _group_mgr.is_group_exist(group_name): + return -1 + g = _group_mgr.get_group_by_name(group_name) + return g.world_size + + +def allreduce(tensor, group_name: str = "default", op=types.ReduceOp.SUM): + """Collective allreduce the tensor across the group. + + Args: + tensor: the tensor to be all-reduced on this process. + group_name: the collective group name to perform allreduce. + op: The reduce operation. + + Returns: + None + """ + _check_single_tensor_input(tensor) + g = get_group_handle(group_name) + opts = types.AllReduceOptions + opts.reduceOp = op + g.allreduce([tensor], opts) + + +def allreduce_multigpu( + tensor_list: list, group_name: str = "default", op=types.ReduceOp.SUM +): + """Collective allreduce a list of tensors across the group. + + Args: + tensor_list (List[tensor]): list of tensors to be allreduced, + each on a GPU. + group_name: the collective group name to perform allreduce. + + Returns: + None + """ + if not types.cupy_available(): + raise RuntimeError("Multigpu calls requires NCCL and Cupy.") + _check_tensor_list_input(tensor_list) + g = get_group_handle(group_name) + opts = types.AllReduceOptions + opts.reduceOp = op + g.allreduce(tensor_list, opts) + + +def barrier(group_name: str = "default"): + """Barrier all processes in the collective group. + + Args: + group_name: the name of the group to barrier. + + Returns: + None + """ + g = get_group_handle(group_name) + g.barrier() + + +def reduce( + tensor, dst_rank: int = 0, group_name: str = "default", op=types.ReduceOp.SUM +): + """Reduce the tensor across the group to the destination rank. + + Args: + tensor: the tensor to be reduced on this process. + dst_rank: the rank of the destination process. + group_name: the collective group name to perform reduce. + op: The reduce operation. + + Returns: + None + """ + _check_single_tensor_input(tensor) + g = get_group_handle(group_name) + + # check dst rank + _check_rank_valid(g, dst_rank) + opts = types.ReduceOptions() + opts.reduceOp = op + opts.root_rank = dst_rank + opts.root_tensor = 0 + g.reduce([tensor], opts) + + +def reduce_multigpu( + tensor_list: list, + dst_rank: int = 0, + dst_tensor: int = 0, + group_name: str = "default", + op=types.ReduceOp.SUM, +): + """Reduce the tensor across the group to the destination rank + and destination tensor. + + Args: + tensor_list: the list of tensors to be reduced on this process; + each tensor located on a GPU. + dst_rank: the rank of the destination process. + dst_tensor: the index of GPU at the destination. + group_name: the collective group name to perform reduce. + op: The reduce operation. + + Returns: + None + """ + if not types.cupy_available(): + raise RuntimeError("Multigpu calls requires NCCL and Cupy.") + _check_tensor_list_input(tensor_list) + g = get_group_handle(group_name) + + # check dst rank + _check_rank_valid(g, dst_rank) + _check_root_tensor_valid(len(tensor_list), dst_tensor) + opts = types.ReduceOptions() + opts.reduceOp = op + opts.root_rank = dst_rank + opts.root_tensor = dst_tensor + g.reduce(tensor_list, opts) + + +def broadcast(tensor, src_rank: int = 0, group_name: str = "default"): + """Broadcast the tensor from a source process to all others. + + Args: + tensor: the tensor to be broadcasted (src) or received (destination). + src_rank: the rank of the source process. + group_name: the collective group name to perform broadcast. + + Returns: + None + """ + _check_single_tensor_input(tensor) + g = get_group_handle(group_name) + + # check src rank + _check_rank_valid(g, src_rank) + opts = types.BroadcastOptions() + opts.root_rank = src_rank + opts.root_tensor = 0 + g.broadcast([tensor], opts) + + +def broadcast_multigpu( + tensor_list, src_rank: int = 0, src_tensor: int = 0, group_name: str = "default" +): + """Broadcast the tensor from a source GPU to all other GPUs. + + Args: + tensor_list: the tensors to broadcast (src) or receive (dst). + src_rank: the rank of the source process. + src_tensor: the index of the source GPU on the source process. + group_name: the collective group name to perform broadcast. + + Returns: + None + """ + if not types.cupy_available(): + raise RuntimeError("Multigpu calls requires NCCL and Cupy.") + _check_tensor_list_input(tensor_list) + g = get_group_handle(group_name) + + # check src rank + _check_rank_valid(g, src_rank) + _check_root_tensor_valid(len(tensor_list), src_tensor) + opts = types.BroadcastOptions() + opts.root_rank = src_rank + opts.root_tensor = src_tensor + g.broadcast(tensor_list, opts) + + +def allgather(tensor_list: list, tensor, group_name: str = "default"): + """Allgather tensors from each process of the group into a list. + + Args: + tensor_list: the results, stored as a list of tensors. + tensor: the tensor (to be gathered) in the current process + group_name: the name of the collective group. + + Returns: + None + """ + _check_single_tensor_input(tensor) + _check_tensor_list_input(tensor_list) + g = get_group_handle(group_name) + if len(tensor_list) != g.world_size: + # Typically CLL lib requires len(tensor_list) >= world_size; + # Here we make it more strict: len(tensor_list) == world_size. + raise RuntimeError( + "The length of the tensor list operands to allgather " + "must be equal to world_size." + ) + opts = types.AllGatherOptions() + g.allgather([tensor_list], [tensor], opts) + + +def allgather_multigpu( + output_tensor_lists: list, input_tensor_list: list, group_name: str = "default" +): + """Allgather tensors from each gpus of the group into lists. + + Args: + output_tensor_lists (List[List[tensor]]): gathered results, with shape + must be num_gpus * world_size * shape(tensor). + input_tensor_list: (List[tensor]): a list of tensors, with shape + num_gpus * shape(tensor). + group_name: the name of the collective group. + + Returns: + None + """ + if not types.cupy_available(): + raise RuntimeError("Multigpu calls requires NCCL and Cupy.") + _check_tensor_lists_input(output_tensor_lists) + _check_tensor_list_input(input_tensor_list) + g = get_group_handle(group_name) + opts = types.AllGatherOptions() + g.allgather(output_tensor_lists, input_tensor_list, opts) + + +def reducescatter( + tensor, tensor_list: list, group_name: str = "default", op=types.ReduceOp.SUM +): + """Reducescatter a list of tensors across the group. + + Reduce the list of the tensors across each process in the group, then + scatter the reduced list of tensors -- one tensor for each process. + + Args: + tensor: the resulted tensor on this process. + tensor_list: The list of tensors to be reduced and scattered. + group_name: the name of the collective group. + op: The reduce operation. + + Returns: + None + """ + _check_single_tensor_input(tensor) + _check_tensor_list_input(tensor_list) + g = get_group_handle(group_name) + opts = types.ReduceScatterOptions() + opts.reduceOp = op + if len(tensor_list) != g.world_size: + raise RuntimeError( + "The length of the tensor list operands to reducescatter " + "must not be equal to world_size." + ) + g.reducescatter([tensor], [tensor_list], opts) + + +def reducescatter_multigpu( + output_tensor_list, + input_tensor_lists, + group_name: str = "default", + op=types.ReduceOp.SUM, +): + """Reducescatter a list of tensors across all GPUs. + + Args: + output_tensor_list: the resulted list of tensors, with + shape: num_gpus * shape(tensor). + input_tensor_lists: the original tensors, with shape: + num_gpus * world_size * shape(tensor). + group_name: the name of the collective group. + op: The reduce operation. + + Returns: + None. + """ + if not types.cupy_available(): + raise RuntimeError("Multigpu calls requires NCCL and Cupy.") + _check_tensor_lists_input(input_tensor_lists) + _check_tensor_list_input(output_tensor_list) + g = get_group_handle(group_name) + opts = types.ReduceScatterOptions() + opts.reduceOp = op + g.reducescatter(output_tensor_list, input_tensor_lists, opts) + + +def send(tensor, dst_rank: int, group_name: str = "default"): + """Send a tensor to a remote process synchronously. + + Args: + tensor: the tensor to send. + dst_rank: the rank of the destination process. + group_name: the name of the collective group. + + Returns: + None + """ + _check_single_tensor_input(tensor) + g = get_group_handle(group_name) + _check_rank_valid(g, dst_rank) + if dst_rank == g.rank: + raise RuntimeError("The destination rank '{}' is self.".format(dst_rank)) + opts = types.SendOptions() + opts.dst_rank = dst_rank + g.send([tensor], opts) + + +def send_multigpu( + tensor, + dst_rank: int, + dst_gpu_index: int, + group_name: str = "default", + n_elements: int = 0, +): + """Send a tensor to a remote GPU synchronously. + + The function assumes each process owns >1 GPUs, and the sender + process and receiver process has equal number of GPUs. + + Args: + tensor: the tensor to send, located on a GPU. + dst_rank: the rank of the destination process. + dst_gpu_index: the destination gpu index. + group_name: the name of the collective group. + n_elements: if specified, send the next n elements + from the starting address of tensor. + + Returns: + None + """ + if not types.cupy_available(): + raise RuntimeError("send_multigpu call requires NCCL.") + _check_single_tensor_input(tensor) + g = get_group_handle(group_name) + _check_rank_valid(g, dst_rank) + if dst_rank == g.rank: + raise RuntimeError( + "The dst_rank '{}' is self. Considering " + "doing GPU to GPU memcpy instead?".format(dst_rank) + ) + if n_elements < 0: + raise RuntimeError("The n_elements '{}' should >= 0.".format(n_elements)) + opts = types.SendOptions() + opts.dst_rank = dst_rank + opts.dst_gpu_index = dst_gpu_index + opts.n_elements = n_elements + g.send([tensor], opts) + + +def recv(tensor, src_rank: int, group_name: str = "default"): + """Receive a tensor from a remote process synchronously. + + Args: + tensor: the received tensor. + src_rank: the rank of the source process. + group_name: the name of the collective group. + + Returns: + None + """ + _check_single_tensor_input(tensor) + g = get_group_handle(group_name) + _check_rank_valid(g, src_rank) + if src_rank == g.rank: + raise RuntimeError("The destination rank '{}' is self.".format(src_rank)) + opts = types.RecvOptions() + opts.src_rank = src_rank + g.recv([tensor], opts) + + +def recv_multigpu( + tensor, + src_rank: int, + src_gpu_index: int, + group_name: str = "default", + n_elements: int = 0, +): + """Receive a tensor from a remote GPU synchronously. + + The function asssume each process owns >1 GPUs, and the sender + process and receiver process has equal nubmer of GPUs. + + Args: + tensor: The received tensor, located on a GPU. + src_rank: The rank of the source process. + src_gpu_index: The index of the source GPU on the src process. + group_name: The name of the collective group. + + Returns: + None + """ + if not types.cupy_available(): + raise RuntimeError("recv_multigpu call requires NCCL.") + _check_single_tensor_input(tensor) + g = get_group_handle(group_name) + _check_rank_valid(g, src_rank) + if src_rank == g.rank: + raise RuntimeError( + "The dst_rank '{}' is self. Considering " + "doing GPU to GPU memcpy instead?".format(src_rank) + ) + if n_elements < 0: + raise RuntimeError("The n_elements '{}' should be >= 0.".format(n_elements)) + opts = types.RecvOptions() + opts.src_rank = src_rank + opts.src_gpu_index = src_gpu_index + opts.n_elements = n_elements + g.recv([tensor], opts) + + +def synchronize(gpu_id: int): + """Synchronize the current process to a give device. + + Args: + gpu_id: the GPU device id to synchronize. + + Returns: + None + """ + if not types.cupy_available(): + raise RuntimeError("synchronize call requires CUDA and NCCL.") + import cupy as cp + + cp.cuda.Device(gpu_id).synchronize() + + +def get_group_handle(group_name: str = "default"): + """Check if the group is initialized and return the group handle. + + Args: + group_name: the name of the collective group. + + Returns: + The collective group handle. + """ + if group_name != types.NIXL_GROUP_NAME: + _check_inside_actor() + global _group_mgr + global _group_mgr_lock + with _group_mgr_lock: + if not _group_mgr.is_group_exist(group_name): + # try loading from remote info store + try: + if group_name == types.NIXL_GROUP_NAME: + _group_mgr.create_collective_group( + types.Backend.NIXL, None, None, group_name, None + ) + else: + # if the information is stored in an Info object, + # get and create the group. + name = "info_" + group_name + mgr = ray.get_actor(name=name) + ids, world_size, rank, backend, gloo_timeout = ray.get( + mgr.get_info.remote() + ) + worker = ray._private.worker.global_worker + id_ = worker.core_worker.get_actor_id() + r = rank[ids.index(id_)] + _group_mgr.create_collective_group( + backend, world_size, r, group_name, gloo_timeout + ) + except ValueError as exc: + # check if this group is initialized using options() + if ( + "collective_group_name" in os.environ + and os.environ["collective_group_name"] == group_name + ): + rank = int(os.environ["collective_rank"]) + world_size = int(os.environ["collective_world_size"]) + backend = os.environ["collective_backend"] + gloo_timeout = os.getenv("collective_gloo_timeout", 30000) + _group_mgr.create_collective_group( + backend, world_size, rank, group_name, gloo_timeout + ) + else: + raise RuntimeError( + "The collective group '{}' is not " + "initialized in the process.".format(group_name) + ) from exc + g = _group_mgr.get_group_by_name(group_name) + return g + + +def _check_single_tensor_input(tensor): + """Check if the tensor is with a supported type.""" + if isinstance(tensor, np.ndarray): + return + if types.cupy_available(): + if isinstance(tensor, types.cp.ndarray): + return + if types.torch_available(): + if isinstance(tensor, types.th.Tensor): + return + raise RuntimeError( + "Unrecognized tensor type '{}'. Supported types are: " + "np.ndarray, torch.Tensor, cupy.ndarray.".format(type(tensor)) + ) + + +def _check_backend_availability(backend: types.Backend): + """Check whether the backend is available.""" + if backend == types.Backend.GLOO: + # Now we have deprecated pygloo, and use torch_gloo in all cases. + if not torch_distributed_available(): + raise RuntimeError("torch.distributed is not available.") + elif backend == types.Backend.NCCL: + if not nccl_available(): + raise RuntimeError("NCCL is not available.") + elif backend == types.Backend.TORCH_GLOO: + if not torch_distributed_available(): + raise RuntimeError("torch.distributed is not available.") + elif backend == types.Backend.NIXL: + if not nixl_available(): + raise RuntimeError("NIXL is not available.") + + +def _check_inside_actor(): + """Check if currently it is inside a Ray actor/task.""" + worker = ray._private.worker.global_worker + if worker.mode == ray.WORKER_MODE: + return + else: + raise RuntimeError( + "The collective APIs shall be only used inside a Ray actor or task." + ) + + +def _check_rank_valid(g, rank: int): + """Check the rank: 0 <= rank < world_size.""" + if rank < 0: + raise ValueError("rank '{}' is negative.".format(rank)) + if rank >= g.world_size: + raise ValueError( + "rank '{}' must be less than world size '{}'".format(rank, g.world_size) + ) + + +def _check_tensor_list_input(tensor_list): + """Check if the input is a list of supported tensor types.""" + if not isinstance(tensor_list, list): + raise RuntimeError( + "The input must be a list of tensors. " + "Got '{}'.".format(type(tensor_list)) + ) + if not tensor_list: + raise RuntimeError("Got an empty list of tensors.") + for t in tensor_list: + _check_single_tensor_input(t) + + +def _check_tensor_lists_input(tensor_lists): + """Check if the input is a list of lists of supported tensor types.""" + if not isinstance(tensor_lists, list): + raise RuntimeError( + "The input must be a list of lists of tensors. " + "Got '{}'.".format(type(tensor_lists)) + ) + if not tensor_lists: + raise RuntimeError(f"Did not receive tensors. Got: {tensor_lists}") + for t in tensor_lists: + _check_tensor_list_input(t) + + +def _check_root_tensor_valid(length, root_tensor): + """Check the root_tensor device is 0 <= root_tensor < length""" + if root_tensor < 0: + raise ValueError("root_tensor '{}' is negative.".format(root_tensor)) + if root_tensor >= length: + raise ValueError( + "root_tensor '{}' is greater than the number of GPUs: " + "'{}'".format(root_tensor, length) + ) diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/__init__.py b/lib/python3.12/site-packages/ray/util/collective/collective_group/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..63fc8a4ebdb101f7524ae547d9dbb1d1f1066f06 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/base_collective_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/base_collective_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f6fe84feff21ea8e2b74169869eb7e08366ad270 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/base_collective_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/cuda_stream.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/cuda_stream.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3806dc175237c853bf2e57a081a3750cdc0afc1f Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/cuda_stream.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nccl_collective_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nccl_collective_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e5190826f18b19cab7ad9fef968665e035ecc289 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nccl_collective_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nccl_util.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nccl_util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0284a666c53acf8ee4d7846756716e5d3e2ca673 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nccl_util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nixl_backend.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nixl_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc9df7bfddcbf187783a52db44861b77e49a0863 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/nixl_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/torch_gloo_collective_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/torch_gloo_collective_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..79b78694723b7bdb1afb294a97c0d504c3935ea9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/collective/collective_group/__pycache__/torch_gloo_collective_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/base_collective_group.py b/lib/python3.12/site-packages/ray/util/collective/collective_group/base_collective_group.py new file mode 100644 index 0000000000000000000000000000000000000000..eff07fb16c671f871d433bced5906d1411850762 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/collective_group/base_collective_group.py @@ -0,0 +1,85 @@ +"""Abstract class for collective groups.""" +from abc import ABCMeta, abstractmethod + +from ray.util.collective.types import ( + AllGatherOptions, + AllReduceOptions, + BarrierOptions, + BroadcastOptions, + RecvOptions, + ReduceOptions, + ReduceScatterOptions, + SendOptions, +) + + +class BaseGroup(metaclass=ABCMeta): + def __init__(self, world_size, rank, group_name): + """Init the process group with basic information. + + Args: + world_size: The total number of processes in the group. + rank: The rank of the current process. + group_name: The group name. + """ + self._world_size = world_size + self._rank = rank + self._group_name = group_name + + @property + def rank(self): + """Return the rank of the current process.""" + return self._rank + + @property + def world_size(self): + """Return the number of processes in this group.""" + return self._world_size + + @property + def group_name(self): + """Return the group name of this group.""" + return self._group_name + + def destroy_group(self): + """GC the communicators.""" + pass + + @classmethod + def backend(cls): + """The backend of this collective group.""" + raise NotImplementedError() + + @abstractmethod + def allreduce(self, tensor, allreduce_options=AllReduceOptions()): + raise NotImplementedError() + + @abstractmethod + def barrier(self, barrier_options=BarrierOptions()): + raise NotImplementedError() + + @abstractmethod + def reduce(self, tensor, reduce_options=ReduceOptions()): + raise NotImplementedError() + + @abstractmethod + def allgather(self, tensor_list, tensor, allgather_options=AllGatherOptions()): + raise NotImplementedError() + + @abstractmethod + def broadcast(self, tensor, broadcast_options=BroadcastOptions()): + raise NotImplementedError() + + @abstractmethod + def reducescatter( + self, tensor, tensor_list, reducescatter_options=ReduceScatterOptions() + ): + raise NotImplementedError() + + @abstractmethod + def send(self, tensor, send_options: SendOptions): + raise NotImplementedError() + + @abstractmethod + def recv(self, tensor, recv_options: RecvOptions): + raise NotImplementedError() diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/cuda_stream.py b/lib/python3.12/site-packages/ray/util/collective/collective_group/cuda_stream.py new file mode 100644 index 0000000000000000000000000000000000000000..dbccb00c1a1742f83f97ad8af788c500af603dca --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/collective_group/cuda_stream.py @@ -0,0 +1,95 @@ +import logging +import threading + +import cupy + +from ray.util.collective.collective_group import nccl_util +from ray.util.collective.const import ENV + +NCCL_STREAM_POOL_SIZE = 32 +MAX_GPU_PER_ACTOR = 16 + +logger = logging.getLogger(__name__) + + +class StreamPool: + """The class that represents a stream pool associated with a GPU. + + When multistream is enabled, we will allocate a pool of streams for each + GPU, and get available stream from this pool when a collective kernel is + initialized. This enables overlapping computation/communication kernels + using multiple CUDA streams, given that the streams a appropriately + synchronized. The class is thread-safe. + + + Args: + device_idx: the absolute index of the device for this pool. + """ + + def __init__(self, device_idx): + self.device_idx = device_idx + + self._initialized = False + self._initialized_lock = threading.Lock() + + self._pool = [None] * NCCL_STREAM_POOL_SIZE + self._counter = 0 + self._pool_lock = threading.Lock() + + def get_stream(self): + """Get an available stream from the pool. + + The function locks the stream pool and releases the lock before + returning. + + Returns: + stream (cupy.cuda.Stream): the returned stream from pool. + """ + + # check the flag + self._initialized_lock.acquire() + if not self._initialized: + self._init_once() + self._initialized_lock.release() + + # Get the stream from the pool. + self._pool_lock.acquire() + stream = self._pool[self._counter] + self._counter = (self._counter + 1) % NCCL_STREAM_POOL_SIZE + self._pool_lock.release() + return stream + + def _init_once(self): + """Initialize the stream pool only for once.""" + with nccl_util.Device(self.device_idx): + for i in range(NCCL_STREAM_POOL_SIZE): + # this is the only place where self._pool will be written. + if ENV.NCCL_USE_MULTISTREAM.val: + logger.debug("NCCL multistream enabled.") + self._pool[i] = cupy.cuda.Stream(null=False, non_blocking=False) + else: + logger.debug("NCCL multistream disabled.") + self._pool[i] = cupy.cuda.Stream.null + self._init_flag = True + + +# This is a map from GPU index to its stream pool. +# It is supposed to be READ-ONLY out of this file +_device_stream_pool_map = dict() + + +def _init_stream_pool(): + global _device_stream_pool_map + for i in range(MAX_GPU_PER_ACTOR): + _device_stream_pool_map[i] = StreamPool(i) + + +def get_stream_pool(device_idx): + """Get the CUDA stream pool of a GPU device.""" + # In case there will be multiple threads writing to the pool. + lock = threading.Lock() + lock.acquire() + if not _device_stream_pool_map: + _init_stream_pool() + lock.release() + return _device_stream_pool_map[device_idx] diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/nccl_collective_group.py b/lib/python3.12/site-packages/ray/util/collective/collective_group/nccl_collective_group.py new file mode 100644 index 0000000000000000000000000000000000000000..07e3da29686af5ff09c7d1aaf3b344e8f33cacda --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/collective_group/nccl_collective_group.py @@ -0,0 +1,836 @@ +import datetime +import logging +import time + +import cupy +import torch + +import ray +from ray.util.collective.collective_group import nccl_util +from ray.util.collective.collective_group.base_collective_group import BaseGroup +from ray.util.collective.collective_group.cuda_stream import get_stream_pool +from ray.util.collective.const import ENV, get_store_name +from ray.util.collective.types import ( + AllGatherOptions, + AllReduceOptions, + Backend, + BarrierOptions, + BroadcastOptions, + RecvOptions, + ReduceOptions, + ReduceScatterOptions, + SendOptions, + torch_available, +) + +logger = logging.getLogger(__name__) + + +class Rendezvous: + """A rendezvous class for different actor/task processes to meet. + + To initialize an NCCL collective communication group, different + actors/tasks spawned in Ray in a collective group needs to meet + each other to synchronize the NCCLUniqueID. This class guarantees + they meet via the NCCLUniqueIDStore, initialized on the rank=0 + process. + + Args: + store_key: the unique store key, usually as a concatanation + of group_name and communicator key. See `get_nccl_communicator` + for more details. + """ + + def __init__(self, store_key): + if not store_key: + raise ValueError( + "Invalid store_key. The store_key is a concatenation of " + "'group_name' and the 'communicator_key'. See the " + "docstring of `get_nccl_communicator` for details." + ) + self._store_key = store_key + self._store_name = None + self._store = None + + def meet(self, timeout_s=180): + """Meet at the named actor store. + + Args: + timeout_s: timeout in seconds. + + Return: + None + """ + if timeout_s <= 0: + raise ValueError( + "The 'timeout' argument must be positive. " + "Got '{}'.".format(timeout_s) + ) + self._store_name = get_store_name(self._store_key) + timeout_delta = datetime.timedelta(seconds=timeout_s) + elapsed = datetime.timedelta(seconds=0) + start_time = datetime.datetime.now() + while elapsed < timeout_delta: + try: + logger.debug( + "Trying to meet at the store '{}'".format(self._store_name) + ) + self._store = ray.get_actor(self._store_name) + except ValueError: + logger.debug( + "Failed to meet at the store '{}'." + "Trying again...".format(self._store_name) + ) + time.sleep(1) + elapsed = datetime.datetime.now() - start_time + continue + logger.debug("Successful rendezvous!") + break + if not self._store: + raise RuntimeError( + "Unable to meet other processes " + "at the rendezvous store. If you are using " + "P2P communication, please check if tensors " + "are put in the correct GPU. " + ) + + @property + def store(self): + return self._store + + def get_nccl_id(self, timeout_s=180): + """Get the NCCLUniqueID from the store through Ray. + + Args: + timeout_s: timeout in seconds. + + Return: + uid: the NCCLUniqueID if successful. + """ + if not self._store: + raise ValueError("Rendezvous store is not setup.") + try: + uid = ray.get(self._store.wait_and_get_id.remote(), timeout=timeout_s) + except ray.exceptions.GetTimeoutError: + raise RuntimeError( + f"Unable to get the NCCLUniqueID from the store within {timeout_s} seconds." + ) from None + return uid + + +class NCCLGroup(BaseGroup): + def __init__(self, world_size, rank, group_name): + """Init an NCCL collective group.""" + super(NCCLGroup, self).__init__(world_size, rank, group_name) + + # communicator and stream cache. + # TODO (Hao): we need a lock here... + self._dev_comm_map = {} + self._dev_streams_map = {} + + # record the used GPU IDs. + self._used_gpu_indices = set() + + # TODO(Fu): might need an event map + self._dev_event_map = {} + + if nccl_util.get_nccl_build_version() < 2000: + raise RuntimeError("NCCL in Ray requires NCCL >= 2.0.") + if nccl_util.get_nccl_runtime_version() < 2704: + logger.warning("NCCL send/recv calls requires NCCL>=2.7.4") + + def destroy_group(self): + """Destroy the group and release NCCL communicators.""" + if len(self._dev_comm_map.keys()) > 0: + # TODO(Hao): check this barrier call + # self.barrier() + + # Destroy the communicators and streams. + for comm_key, comms in self._dev_comm_map.items(): + for c in comms: + c.destroy() + self._dev_comm_map[comm_key] = None + + if self.rank == 0: + for comm_key in self._dev_comm_map: + assert not self._dev_comm_map[comm_key] + group_key = self._generate_group_key(comm_key) + self._destroy_store(group_key) + self._barrier_tensor = None + self._dev_comm_map = None + self._dev_streams_map = None + super(NCCLGroup, self).destroy_group() + + @classmethod + def backend(cls): + return Backend.NCCL + + def allreduce(self, tensors, allreduce_options=AllReduceOptions()): + """AllReduce tensors across the collective group following options. + + Args: + tensors: the list of tensors to be reduced. Each tensor must + reside on one GPU of the current process. + allreduce_options: allreduce options. + + Returns: + None + """ + + def collective_fn(input_tensor, output_tensor, comm, stream): + comm.allReduce( + nccl_util.get_tensor_ptr(input_tensor), + nccl_util.get_tensor_ptr(output_tensor), + nccl_util.get_tensor_n_elements(input_tensor), + nccl_util.get_nccl_tensor_dtype(input_tensor), + nccl_util.get_nccl_reduce_op(allreduce_options.reduceOp), + stream.ptr, + ) + + self._collective(tensors, tensors, collective_fn) + + def barrier(self, barrier_options=BarrierOptions()): + """Blocks until all processes reach this barrier. + + Args: + barrier_options: barrier options. + + Returns: + None + """ + # Get the device list. + if self._used_gpu_indices: + devices = list(self._used_gpu_indices) + else: + devices = list(range(nccl_util.get_num_gpus())) + barrier_tensors = [None] * len(devices) + for i, d in enumerate(devices): + with nccl_util.Device(d): + barrier_tensors[i] = cupy.array([1]) + self.allreduce(barrier_tensors) + + def reduce(self, tensors, reduce_options=ReduceOptions()): + """Reduce tensors to a destination gpu following options. + + Args: + tensors: the list of tensors to be reduced, each tensor + must reside on one gpu of the current process. + reduce_options: reduce options. + + Returns: + None + """ + root_rank = len(tensors) * reduce_options.root_rank + reduce_options.root_tensor + + def collective_fn(input_tensor, output_tensor, comm, stream): + comm.reduce( + nccl_util.get_tensor_ptr(input_tensor), + nccl_util.get_tensor_ptr(output_tensor), + nccl_util.get_tensor_n_elements(input_tensor), + nccl_util.get_nccl_tensor_dtype(input_tensor), + nccl_util.get_nccl_reduce_op(reduce_options.reduceOp), + root_rank, + stream.ptr, + ) + + self._collective(tensors, tensors, collective_fn) + + def broadcast(self, tensors, broadcast_options=BroadcastOptions()): + """Broadcast tensors to all other gpus following options. + + Args: + tensors: tensors to be broadcast or received. + broadcast_options: broadcast options. + + Returns: + None + """ + root_rank = ( + len(tensors) * broadcast_options.root_rank + broadcast_options.root_tensor + ) + + def collective_fn(input_tensor, output_tensor, comm, stream): + comm.broadcast( + nccl_util.get_tensor_ptr(input_tensor), + nccl_util.get_tensor_ptr(output_tensor), + nccl_util.get_tensor_n_elements(input_tensor), + nccl_util.get_nccl_tensor_dtype(input_tensor), + root_rank, + stream.ptr, + ) + + self._collective(tensors, tensors, collective_fn) + + def allgather(self, tensor_lists, tensors, allgather_options=AllGatherOptions()): + """Allgather tensors across gpus into a list of tensors. + + Args: + tensor_lists (List[List[Tensor]]): allgathered tensors. + tensors: the list of tensors to allgather across the group. + Each tensor must lolcate on a GPU of the process. + allgather_options: allgather options. + + Returns: + None + """ + + def collective_fn(input_tensor, output_tensor, comm, stream): + comm.allGather( + nccl_util.get_tensor_ptr(input_tensor), + nccl_util.get_tensor_ptr(output_tensor), + nccl_util.get_tensor_n_elements(input_tensor), + nccl_util.get_nccl_tensor_dtype(input_tensor), + stream.ptr, + ) + + _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists) + output_flattened = [ + _flatten_for_scatter_gather(tensor_list, copy=False) + for tensor_list in tensor_lists + ] + + def postprocess_fn(stream): + # TODO(Hao): designate a copy stream. + for i, tensor_list in enumerate(tensor_lists): + for j, tensor in enumerate(tensor_list): + nccl_util.copy_tensor(tensor, output_flattened[i][j]) + + self._collective( + tensors, output_flattened, collective_fn, postprocess_fn=postprocess_fn + ) + + def reducescatter( + self, tensors, tensor_lists, reducescatter_options=ReduceScatterOptions() + ): + """Reduce then scatter a list of tensors across the group. + + Args: + tensors: the output tensors (could be unspecified), each + located on a GPU of the current process. + tensor_lists (List[List]): the list of tensors to be reduced then + scattered. + reducescatter_options: reduce-scatter options. + + Returns: + None + """ + + def collective_fn(input_tensor, output_tensor, comm, stream): + comm.reduceScatter( + nccl_util.get_tensor_ptr(input_tensor), + nccl_util.get_tensor_ptr(output_tensor), + nccl_util.get_tensor_n_elements(output_tensor), + nccl_util.get_nccl_tensor_dtype(output_tensor), + nccl_util.get_nccl_reduce_op(reducescatter_options.reduceOp), + stream.ptr, + ) + + _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists) + input_flattened = [ + _flatten_for_scatter_gather(tensor_list, copy=False) + for tensor_list in tensor_lists + ] + + def preprocess_fn(stream): + for i, tensor_list in enumerate(tensor_lists): + for j, tensor in enumerate(tensor_list): + nccl_util.copy_tensor(input_flattened[i][j], tensor) + + self._collective( + input_flattened, tensors, collective_fn, preprocess_fn=preprocess_fn + ) + + def send(self, tensors, send_options=SendOptions()): + """Send a tensor to a destination gpu in the group. + + Args: + tensors: the tensor to send. + send_options: send options. + + Returns: + None + """ + + def p2p_fn(tensor, comm, stream, peer): + comm.send( + nccl_util.get_tensor_ptr(tensor), + send_options.n_elements + if send_options.n_elements > 0 + else nccl_util.get_tensor_n_elements(tensor), + nccl_util.get_nccl_tensor_dtype(tensor), + peer, + stream.ptr, + ) + + self._point2point( + tensors, p2p_fn, send_options.dst_rank, send_options.dst_gpu_index + ) + + def recv(self, tensors, recv_options=RecvOptions()): + """Receive a tensor from a source gpu in the group. + + Args: + tensors: the received tensor. + recv_options: Receive options. + + Returns: + None + """ + + def p2p_fn(tensor, comm, stream, peer): + comm.recv( + nccl_util.get_tensor_ptr(tensor), + recv_options.n_elements + if recv_options.n_elements > 0 + else nccl_util.get_tensor_n_elements(tensor), + nccl_util.get_nccl_tensor_dtype(tensor), + peer, + stream.ptr, + ) + + self._point2point( + tensors, p2p_fn, recv_options.src_rank, recv_options.src_gpu_index + ) + + def _get_nccl_collective_communicator(self, comm_key, device_list): + """Create or retrieve an NCCL communicator from cache. + + If the communicator is found in cache, return the communicator. If not, + a communicator and a stream will be created and put in cache. + TODO(Hao): this function is not thread-safe now. + + Args: + comm_key: the key to query the communicator cache. + device_list: a list of GPU devices of the current process + that participates into the collective. + + Returns: + communicator: the NCCL communicator corresponded to the devices. + """ + if not comm_key: + raise RuntimeError("Got empty communicator key.") + for d in device_list: + self._used_gpu_indices.add(d) + + # TODO(Hao): lock the _dev_comm_map here. + if comm_key in self._dev_comm_map: + return self._dev_comm_map[comm_key] + + group_key = self._generate_group_key(comm_key) + if self.rank == 0: + nccl_uid = self._generate_nccl_uid(group_key) + else: + rendezvous = Rendezvous(group_key) + rendezvous.meet() + nccl_uid = rendezvous.get_nccl_id() + + # Now create the communicators + actual_world_size = len(device_list) * self.world_size + comms = [None] * len(device_list) + streams = [None] * len(device_list) + events = [None] * len(device_list) + nccl_util.groupStart() + for i, device in enumerate(device_list): + actual_rank = self.rank * len(device_list) + i + with nccl_util.Device(device): + comms[i] = nccl_util.create_nccl_communicator( + actual_world_size, nccl_uid, actual_rank + ) + # request a stream from the pool + # note the device_idx is absolute index. + streams[i] = get_stream_pool(device).get_stream() + # TODO(Fu): double check the parameters + events[i] = cupy.cuda.Event() + nccl_util.groupEnd() + # TODO(Fu): lock + self._dev_comm_map[comm_key] = comms + self._dev_streams_map[comm_key] = streams + self._dev_event_map[comm_key] = events + return comms + + @staticmethod + def _sync_streams(device_list, events, streams): + """Let NCCL streams wait for current streams for every device.""" + # TODO(Fu): recordStream besides calling this function? + if ENV.NCCL_USE_MULTISTREAM.val: + for i, device in enumerate(device_list): + with nccl_util.Device(device): + events[i].record(cupy.cuda.get_current_stream()) + streams[i].wait_event(events[i]) + + def _get_nccl_p2p_communicator(self, comm_key, my_gpu_idx, peer_rank, peer_gpu_idx): + """Create or retrieve an NCCL communicator for p2p tasks. + + Note(Hao): this function is not thread-safe now. + + Args: + comm_key: communicator key. + my_gpu_idx: the gpu index on the current process. + peer_rank: the rank of the destination process. + peer_gpu_idx: the gpu index on the peer process. + Returns: + communicator + """ + if not comm_key: + raise RuntimeError("Got empty communicator key.") + + # TODO(Hao): lock the _dev_comm_map here. + if comm_key in self._dev_comm_map: + return self._dev_comm_map[comm_key] + + # Note (Hao): This is a bit complex so I decide to take a note here. + # Here we need to consider three cases: + # Case 1: src_rank != dst_rank, hence the send and recv happen on + # different process (actors/tasks); each process makes independent + # collective calls and manages corresponding communicators. + # Case 2: src_rank == dst_rank, src_gpu_idx == dst_gpu_idx; for + # this case, we simply throw a RuntimeError; + # Case 3: src_rank == dst_rank, src_gpu_idx != dst_gpu_idx, which + # means the send and recv will be called on the same process. We + # DO NOT support this case for now. We need to properly scope: + # (1) communicators creation, and + # (2) send/recv calls + # using groupStart(（ and groupEnd() calls to avoid deadlocks. + if self.rank < peer_rank: + my_p2p_rank = 0 + elif self.rank > peer_rank: + my_p2p_rank = 1 + else: + raise RuntimeError( + "Send and recv happens on the same process! " + "ray.util.collective does not support this case as of now. " + "Alternatively, consider doing GPU to GPU memcpy?" + ) + + group_key = self._generate_group_key(comm_key) + if my_p2p_rank == 0: + nccl_uid = self._generate_nccl_uid(group_key) + else: + rendezvous = Rendezvous(group_key) + rendezvous.meet() + nccl_uid = rendezvous.get_nccl_id() + + # create the p2p communicators + with nccl_util.Device(my_gpu_idx): + comm = nccl_util.create_nccl_communicator(2, nccl_uid, my_p2p_rank) + stream = get_stream_pool(my_gpu_idx).get_stream() + event = cupy.cuda.Event() + + # TODO(Fu): lock and might need to add event + self._dev_comm_map[comm_key] = [comm] + self._dev_streams_map[comm_key] = [stream] + self._dev_event_map[comm_key] = [event] + return [comm] + + def _generate_group_key(self, comm_key): + """Generate a unique key used to initialize the KV store. + + The group key is a concatenation of the communicator key and + the group name, following: [comm_key]@[group_name]. + """ + return comm_key + "@" + self.group_name + + @staticmethod + def _destroy_store(group_key): + """Destroy the KV store (Ray named actor). + + Args: + group_key: the unique key to retrieve the KV store. + + Returns: + None + """ + store_name = get_store_name(group_key) + store = ray.get_actor(store_name) + # ray.get([store.__ray_terminate__.remote()]) + ray.kill(store) + + def _generate_nccl_uid(self, key): + """Generate an NCCL unique ID for initializing communicators. + + The method will also create a KV store using Ray named actor and store + the NCCLUniqueID in the store. The store needs to be garbage collected + when destroying the collective group. + + Args: + key: the key of the . + + Returns: + NCCLUniqueID (str): NCCL unique ID. + """ + group_uid = nccl_util.get_nccl_unique_id() + store_name = get_store_name(key) + # Avoid a potential circular dependency in ray/actor.py + from ray.util.collective.util import NCCLUniqueIDStore + + store = NCCLUniqueIDStore.options(name=store_name, lifetime="detached").remote( + store_name + ) + ray.get([store.set_id.remote(group_uid)]) + return group_uid + + def _collective( + self, + input_tensors, + output_tensors, + collective_fn, + preprocess_fn=None, + postprocess_fn=None, + ): + """A method to encapsulate all collective calls. + + Args: + input_tensors: the list of the input tensors. + output_tensors: the list of the output tensors. + collective_fn: the collective function call. + preprocess_fn: preprocess procedures before collective calls. + postprocess_fn: postprocess procedures after collective calls. + + Returns: + None + """ + _check_gpu_tensors(input_tensors) + _check_gpu_tensors(output_tensors) + + devices = nccl_util.get_tensor_device_list(input_tensors) + key = _get_comm_key_from_devices(devices) + comms = self._get_nccl_collective_communicator(key, devices) + streams = self._dev_streams_map[key] + events = self._dev_event_map[key] + + # TODO(Hao): sync streams and events + self._sync_streams(devices, events, streams) + + # Make the collective call + if preprocess_fn: + preprocess_fn(streams) + + nccl_util.groupStart() + # TODO(Fu): how to recordStreams as there are no library functions + # We also need to make sure input tensors are not freed before their + # usages on ncclStreams finish. This can be achieved by calling + # c10::cuda::CUDACachingAllocator::recordStream, which remembers the + # usage stream (ncclStream), creates an event on the usage stream + # when GC attempts to free the input tensor, and delays GC until that + # event is done. + for i, tensor in enumerate(input_tensors): + collective_fn(tensor, output_tensors[i], comms[i], streams[i]) + nccl_util.groupEnd() + if postprocess_fn: + postprocess_fn(streams) + + def _point2point(self, tensors, p2p_fn, peer_rank: int, peer_gpu_idx: int): + """A method to encapsulate all peer-to-peer calls (i.e., send/recv). + + Args: + tensors: the tensor to send or receive. + p2p_fn: the p2p function call. + peer_rank: the rank of the peer process. + peer_gpu_idx: the index of the gpu on the peer process. + + Returns: + None + """ + # check send/recv availability. + if nccl_util.get_nccl_runtime_version() < 2704: + raise RuntimeError( + "P2p send/recv requires NCCL >= 2.7.4. " + "Got '{}'.".format(nccl_util.get_nccl_runtime_version()) + ) + _check_gpu_tensors(tensors) + + # we currently only support single device to single device send/recv. + assert len(tensors) == 1 + my_gpu_idx = nccl_util.get_tensor_device(tensors[0]) + comm_key = _get_comm_key_send_recv( + self.rank, my_gpu_idx, peer_rank, peer_gpu_idx + ) + comms = self._get_nccl_p2p_communicator( + comm_key, my_gpu_idx, peer_rank, peer_gpu_idx + ) + streams = self._dev_streams_map[comm_key] + events = self._dev_event_map[comm_key] + + # TODO(Hao): sync streams and events + self._sync_streams([my_gpu_idx], events, streams) + + # We have made sure that self.rank != peer_rank during API check. + peer_p2p_rank = 0 if self.rank > peer_rank else 1 + for i, tensor in enumerate(tensors): + p2p_fn(tensor, comms[i], streams[i], peer_p2p_rank) + # Record the stream to avoid tensor being freed before the send/recv is completed. + torch_stream = torch.cuda.ExternalStream(streams[i].ptr) + tensor.record_stream(torch_stream) + + +def _flatten_for_scatter_gather(tensor_list, copy=False): + """Flatten the tensor for gather/scatter operations. + + Args: + tensor_list: the list of tensors to be scattered/gathered. + copy: whether the copy the tensors in tensor_list into the buffer. + + Returns: + The flattened tensor buffer. + """ + if not tensor_list: + raise RuntimeError("Received an empty list.") + t = tensor_list[0] + buffer_shape = [len(tensor_list)] + nccl_util.get_tensor_shape(t) + + # TODO(wuxibin): cupy doesn't support bfloat16 for now, + # once it is supported, we can eliminate this if statement. + # + # Allocate using the same backend as the tensors in `tensor_list`. + # Use torch only when the tensors are torch.Tensor; otherwise fall back to CuPy. + use_torch = False + if torch_available(): + try: + import torch + + use_torch = isinstance(t, torch.Tensor) + except ImportError: + use_torch = False + + if use_torch: + buffer = torch.empty(tuple(buffer_shape), dtype=t.dtype, device=t.device) + else: + # note we need a cupy dtype here. + dtype = nccl_util.get_cupy_tensor_dtype(t) + device = nccl_util.get_tensor_device(t) + with nccl_util.Device(device): + buffer = cupy.empty(buffer_shape, dtype=dtype) + + if copy: + for i, tensor in enumerate(tensor_list): + nccl_util.copy_tensor(buffer[i], tensor) + return buffer + + +def _check_inputs_compatibility_for_scatter_gather(tensors, tensor_lists): + """Check the compatibility between tensor input and tensor list input.""" + if not tensors or not isinstance(tensors, list): + raise RuntimeError("The first argument 'tensors' expects a list of tensors.") + if not tensor_lists or not isinstance(tensor_lists, list): + raise RuntimeError( + "The second argument 'tensor_lists' expects a list of tensor list." + ) + dtype = nccl_util.get_nccl_tensor_dtype(tensors[0]) + shape = nccl_util.get_tensor_shape(tensors[0]) + for i, tensor_list in enumerate(tensor_lists): + # check all tensor in `tensors` match. + dt = nccl_util.get_nccl_tensor_dtype(tensors[i]) + if dt != dtype: + raise RuntimeError( + "All tensor operands to scatter/gather must " + "have the same dtype. Got '{}' and '{}'.".format(dt, dtype) + ) + # Note: typically CCL libraries only requires they have the same + # number of elements; Here we make it more strict -- we require + # exact shape match. + s = nccl_util.get_tensor_shape(tensors[i]) + if s != shape: + raise RuntimeError( + "All tensor operands to scatter/gather must " + "have the same shape. Got '{}' and '{}'.".format(s, shape) + ) + # check all tensors in `tensor_lists` match. + for t in tensor_lists[i]: + # check dtype + dt = nccl_util.get_nccl_tensor_dtype(t) + if dt != dtype: + raise RuntimeError( + "All tensor operands to scatter/gather must " + "have the same dtype. Got '{}' and '{}'.".format(dt, dtype) + ) + s = nccl_util.get_tensor_shape(t) + if s != shape: + raise RuntimeError( + "All tensor operands to scatter/gather must " + "have the same shape. Got '{}' and '{}'.".format(s, shape) + ) + + +def _check_gpu_tensors(tensors): + """Check all tensors are distributed on different GPUs.""" + if not tensors or not isinstance(tensors, list): + raise RuntimeError("'tensors' must be a nonempty list.") + if len(tensors) > nccl_util.get_num_gpus(): + raise RuntimeError( + "Tensor list cannot be larger than the number" + "of available GPUs. Got {} > {}.".format( + len(tensors), nccl_util.get_num_gpus() + ) + ) + t0 = tensors[0] + dt = nccl_util.get_nccl_tensor_dtype(t0) + s = nccl_util.get_tensor_shape(t0) + d = nccl_util.get_tensor_device(t0) + for i, t in enumerate(tensors): + if i == 0: + continue + # We need to check the following: + # (1) tensor is cuda (already checked during API) + # (2) tensor dtype + # (3) tensor shape match + # (4) each tensor is on a different GPU + dtype = nccl_util.get_nccl_tensor_dtype(t) + if dt != dtype: + raise RuntimeError( + "Tensors must have identical dtype. Got: '{}'.".format(dtype) + ) + shape = nccl_util.get_tensor_shape(t) + if s != shape: + raise RuntimeError( + "Tensor must have identical shape. Got: '{}'.".format(shape) + ) + device = nccl_util.get_tensor_device(t) + if device == d: + raise RuntimeError("Tensor must be on distinct GPUs.") + + +def _get_comm_key_from_devices(devices): + """Return a key from a list of devices for collective calls. + + For example, if the tensors are on gpus 0, 1, 2, 3, + then the key would be "0,1,2,3". + + Args: + devices: a list of GPU device indices + + Returns: + str: a string represents the key to query the communicator cache. + + """ + return ",".join([str(d) for d in devices]) + + +def _get_comm_key_send_recv(my_rank, my_gpu_idx, peer_rank, peer_gpu_idx): + """Return a key given source and destination ranks for p2p tasks. + + The p2p key is in the following form: + [min_rank]_[gpu_index]:[max_rank]_[gpu_index]. + + Args: + my_rank: the rank of the source process. + my_gpu_idx: the source gpu index on the process. + peer_rank: the rank of the destination process. + peer_gpu_idx: the destination gpu index on the process. + + Returns: + comm_key: a string key to query the communication cache. + """ + if my_rank < peer_rank: + lower_key = str(my_rank) + "_" + str(my_gpu_idx) + higher_key = str(peer_rank) + "_" + str(peer_gpu_idx) + elif my_rank > peer_rank: + lower_key = str(peer_rank) + "_" + str(peer_gpu_idx) + higher_key = str(my_rank) + "_" + str(my_gpu_idx) + else: + raise RuntimeError( + "Send and recv happens on the same process. ray.util.collective " + "does not support this case as of now. Alternatively, consider " + "doing GPU to GPU memcpy?" + ) + comm_key = lower_key + ":" + higher_key + return comm_key diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/nccl_util.py b/lib/python3.12/site-packages/ray/util/collective/collective_group/nccl_util.py new file mode 100644 index 0000000000000000000000000000000000000000..7f68d843020875d58932722fbe73b4baabaec080 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/collective_group/nccl_util.py @@ -0,0 +1,297 @@ +"""Code to wrap some NCCL API calls.""" +import numpy + +try: + import cupy + from cupy.cuda import ( + Device, # noqa: F401 + nccl, + ) + from cupy.cuda.nccl import ( + NcclCommunicator, + get_build_version, + get_version, + groupEnd, # noqa: F401 + groupStart, # noqa: F401 + ) +except ImportError: + raise ImportError("NCCL in Ray requires Cupy being available!") + +from ray.util.collective.types import ReduceOp, torch_available + +NCCL_REDUCE_OP_MAP = { + ReduceOp.SUM: nccl.NCCL_SUM, + ReduceOp.PRODUCT: nccl.NCCL_PROD, + ReduceOp.MIN: nccl.NCCL_MIN, + ReduceOp.MAX: nccl.NCCL_MAX, +} + +# cupy types are the same with numpy types +NUMPY_NCCL_DTYPE_MAP = { + # INT types + numpy.int_: nccl.NCCL_INT64, + numpy.uint8: nccl.NCCL_UINT8, + numpy.uint32: nccl.NCCL_UINT32, + numpy.uint64: nccl.NCCL_UINT64, + numpy.int8: nccl.NCCL_INT8, + numpy.int32: nccl.NCCL_INT32, + numpy.int64: nccl.NCCL_INT64, + # FLOAT types + numpy.half: nccl.NCCL_HALF, + numpy.float16: nccl.NCCL_FLOAT16, + numpy.float32: nccl.NCCL_FLOAT32, + numpy.float64: nccl.NCCL_FLOAT64, + numpy.double: nccl.NCCL_DOUBLE, +} + +if torch_available(): + import torch + import torch.utils.dlpack + + TORCH_NCCL_DTYPE_MAP = { + torch.bool: nccl.NCCL_INT8, + # INT types + torch.int: nccl.NCCL_INT, + torch.uint8: nccl.NCCL_UINT8, + torch.int8: nccl.NCCL_INT8, + torch.int32: nccl.NCCL_INT32, + torch.int64: nccl.NCCL_INT64, + torch.long: nccl.NCCL_INT64, + # FLOAT types + torch.half: nccl.NCCL_HALF, + torch.float: nccl.NCCL_FLOAT, + torch.float16: nccl.NCCL_FLOAT16, + torch.float32: nccl.NCCL_FLOAT32, + torch.float64: nccl.NCCL_FLOAT64, + torch.double: nccl.NCCL_DOUBLE, + } + + # Older versions of cupy don't support bfloat16. + if hasattr(nccl, "NCCL_BFLOAT16"): + TORCH_NCCL_DTYPE_MAP[torch.bfloat16] = nccl.NCCL_BFLOAT16 + + TORCH_NUMPY_DTYPE_MAP = { + # INT types + torch.int: numpy.int32, + torch.uint8: numpy.uint8, + torch.int8: numpy.int8, + torch.int32: numpy.int32, + torch.int64: numpy.int64, + torch.long: numpy.int64, + # FLOAT types + torch.half: numpy.half, + torch.float: numpy.float32, + torch.float16: numpy.float16, + torch.float32: numpy.float32, + torch.float64: numpy.float64, + } + + +def get_num_gpus(): + """Returns the number of compute-capable GPUs.""" + return cupy.cuda.runtime.getDeviceCount() + + +def get_nccl_build_version(): + return get_build_version() + + +def get_nccl_runtime_version(): + return get_version() + + +def get_nccl_unique_id(): + return nccl.get_unique_id() + + +def create_nccl_communicator(world_size, nccl_unique_id, rank): + """Create an NCCL communicator using NCCL APIs. + + Args: + world_size: the number of processes of this communicator group. + nccl_unique_id: the NCCLUniqueID for this group. + rank: the rank of this process. + Returns: + comm (nccl.ncclComm_t): an NCCL communicator. + """ + comm = NcclCommunicator(world_size, nccl_unique_id, rank) + return comm + + +def get_nccl_reduce_op(reduce_op): + """Map the reduce op to NCCL reduce op type. + + Args: + reduce_op: ReduceOp Enum (SUM/PRODUCT/MIN/MAX). + Returns: + (nccl.ncclRedOp_t): the mapped NCCL reduce op. + """ + if reduce_op not in NCCL_REDUCE_OP_MAP: + raise RuntimeError("NCCL does not support reduce op: '{}'.".format(reduce_op)) + return NCCL_REDUCE_OP_MAP[reduce_op] + + +def get_nccl_tensor_dtype(tensor): + """Return the corresponded NCCL dtype given a tensor.""" + if isinstance(tensor, cupy.ndarray): + return NUMPY_NCCL_DTYPE_MAP[tensor.dtype.type] + if torch_available(): + if isinstance(tensor, torch.Tensor): + return TORCH_NCCL_DTYPE_MAP[tensor.dtype] + raise ValueError( + "Unsupported tensor type. Got: {}. Supported " + "GPU tensor types are: torch.Tensor, " + "cupy.ndarray.".format(type(tensor)) + ) + + +def get_cupy_tensor_dtype(tensor): + """Return the corresponded Cupy dtype given a tensor.""" + if isinstance(tensor, cupy.ndarray): + return tensor.dtype.type + if torch_available(): + if isinstance(tensor, torch.Tensor): + return TORCH_NUMPY_DTYPE_MAP[tensor.dtype] + raise ValueError( + "Unsupported tensor type. Got: {}. Supported " + "GPU tensor types are: torch.Tensor, " + "cupy.ndarray.".format(type(tensor)) + ) + + +def get_tensor_ptr(tensor): + """Return the pointer to the underlying memory storage of a tensor.""" + if isinstance(tensor, cupy.ndarray): + return tensor.data.ptr + if isinstance(tensor, numpy.ndarray): + return tensor.data + if torch_available(): + if isinstance(tensor, torch.Tensor): + if not tensor.is_cuda: + raise RuntimeError( + "Torch tensor must be on GPU when using NCCL collectives." + ) + return tensor.data_ptr() + raise ValueError( + "Unsupported tensor type. Got: {}. Supported " + "GPU tensor types are: torch.Tensor, " + "cupy.ndarray.".format(type(tensor)) + ) + + +def get_tensor_n_elements(tensor): + """Return the number of elements in a tensor.""" + if isinstance(tensor, cupy.ndarray) or isinstance(tensor, numpy.ndarray): + return tensor.size + if torch_available(): + if isinstance(tensor, torch.Tensor): + return torch.numel(tensor) + raise ValueError( + "Unsupported tensor type. Got: {}. Supported " + "GPU tensor types are: torch.Tensor, " + "cupy.ndarray.".format(type(tensor)) + ) + + +def get_tensor_shape(tensor): + """Return the shape of the tensor as a list.""" + if isinstance(tensor, cupy.ndarray): + return list(tensor.shape) + if torch_available(): + if isinstance(tensor, torch.Tensor): + return list(tensor.size()) + raise ValueError( + "Unsupported tensor type. Got: {}. Supported " + "GPU tensor types are: torch.Tensor, " + "cupy.ndarray.".format(type(tensor)) + ) + + +def get_tensor_strides(tensor): + """Return the strides of the tensor as a list.""" + if isinstance(tensor, cupy.ndarray): + return [int(stride / tensor.dtype.itemsize) for stride in tensor.strides] + if torch_available(): + if isinstance(tensor, torch.Tensor): + return list(tensor.stride()) + raise ValueError( + "Unsupported tensor type. Got: {}. Supported " + "GPU tensor types are: torch.Tensor, " + "cupy.ndarray.".format(type(tensor)) + ) + + +def get_tensor_device(tensor): + """Return the GPU index of a tensor.""" + if isinstance(tensor, cupy.ndarray): + try: + device = tensor.device.id + except AttributeError as exec: + raise RuntimeError("The tensor is not on a valid GPU.") from exec + elif torch_available() and isinstance(tensor, torch.Tensor): + device = tensor.device.index + if not isinstance(device, int): + raise RuntimeError("The tensor is not on a valid GPU.") + else: + raise ValueError("Unsupported tensor type. Got: {}.".format(type(tensor))) + return device + + +def copy_tensor(dst_tensor, src_tensor): + """Copy the content from src_tensor to dst_tensor. + + Args: + dst_tensor: the tensor to copy from. + src_tensor: the tensor to copy to. + + Returns: + None + """ + copied = True + if isinstance(dst_tensor, cupy.ndarray) and isinstance(src_tensor, cupy.ndarray): + cupy.copyto(dst_tensor, src_tensor) + elif torch_available(): + if isinstance(dst_tensor, torch.Tensor) and isinstance( + src_tensor, torch.Tensor + ): + dst_tensor.copy_(src_tensor) + elif isinstance(dst_tensor, torch.Tensor) and isinstance( + src_tensor, cupy.ndarray + ): + t = torch.utils.dlpack.from_dlpack(src_tensor.toDlpack()) + dst_tensor.copy_(t) + elif isinstance(dst_tensor, cupy.ndarray) and isinstance( + src_tensor, torch.Tensor + ): + t = cupy.fromDlpack(torch.utils.dlpack.to_dlpack(src_tensor)) + cupy.copyto(dst_tensor, t) + else: + copied = False + else: + copied = False + if not copied: + raise ValueError( + "Unsupported tensor type. Got: {} and {}. Supported " + "GPU tensor types are: torch.Tensor, cupy.ndarray.".format( + type(dst_tensor), type(src_tensor) + ) + ) + + +def get_tensor_device_list(tensors): + """Returns the gpu devices of the list of input tensors. + + Args: + tensors: a list of tensors, each locates on a GPU. + + Returns: + list: the list of GPU devices. + + """ + if not isinstance(tensors, list): + raise RuntimeError( + "Expect a list of tensors each locates on a GPU device. " + "Got: '{}'.".format(type(tensors)) + ) + devices = [get_tensor_device(t) for t in tensors] + return devices diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/nixl_backend.py b/lib/python3.12/site-packages/ray/util/collective/collective_group/nixl_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..beff753b055a772b83ddca4bd0106a2f8a30bc40 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/collective_group/nixl_backend.py @@ -0,0 +1,147 @@ +import threading +import time +from typing import TYPE_CHECKING, Any, List, Tuple + +from nixl._api import nixl_agent, nixl_agent_config + +import ray +from ray.util.collective.types import Backend + +if TYPE_CHECKING: + import torch + + +class NixlBackend: + """Backend implementation for NIXL tensor transport. + + This class provides functionality for transferring tensors using NIXL. It handles + initialization of the NIXL agent, receiving tensors, and managing NIXL metadata. + """ + + def __init__(self): + """Initialize the NIXL backend. + + Creates a NIXL agent with UCX backend. + """ + agent_config = nixl_agent_config(backends=["UCX"]) + ctx = ray.get_runtime_context() + actor_id = ctx.get_actor_id() + if actor_id is None: + # If the actor id is None, it means the current process is a driver. + import uuid + + actor_id = f"RAY-DRIVER-{uuid.uuid4()}" + self._nixl_agent = nixl_agent(actor_id, agent_config) + self._aborted_transfer_obj_ids = set() + self._aborted_transfer_obj_ids_lock = threading.Lock() + + @classmethod + def backend(cls): + """Get the backend type. + + Returns: + Backend.NIXL: The backend type enum value for NIXL. + """ + return Backend.NIXL + + def recv( + self, + tensors: List["torch.Tensor"], + obj_id: str, + nixl_serialized_descs: bytes, + remote_nixl_agent_meta: bytes, + ): + """Receive tensors from a remote NIXL agent. + + Args: + tensors: List of tensors to receive into. + obj_id: The object ID for related GPU object. + nixl_serialized_descs: Serialized NIXL descriptors for the remote tensors. + remote_nixl_agent_meta: Metadata about the remote NIXL agent. + + Raises: + RuntimeError: If the NIXL transfer enters an error state. + """ + with self._aborted_transfer_obj_ids_lock: + if obj_id in self._aborted_transfer_obj_ids: + self._aborted_transfer_obj_ids.remove(obj_id) + raise RuntimeError(f"NIXL transfer aborted for object id: {obj_id}") + + local_descs = None + remote_name = None + xfer_handle = None + try: + nixl_agent = self._nixl_agent + remote_descs = nixl_agent.deserialize_descs(nixl_serialized_descs) + local_descs = nixl_agent.register_memory(tensors) + remote_name = nixl_agent.add_remote_agent(remote_nixl_agent_meta) + + xfer_handle = nixl_agent.initialize_xfer( + # "UUID" here is just a placeholder, can be any bytes, but without it, + # nixl will fail to transfer multiple times. + "READ", + local_descs.trim(), + remote_descs, + remote_name, + "UUID", + ) + + state = nixl_agent.transfer(xfer_handle) + if state == "ERR": + raise RuntimeError("NIXL transfer got to Error state.") + # Since current nixl does not provide a better way, we need to check the state of + # the transfer continuously. + while True: + state = nixl_agent.check_xfer_state(xfer_handle) + if state == "ERR": + raise RuntimeError("NIXL transfer got to Error state.") + if state == "PROC": + with self._aborted_transfer_obj_ids_lock: + if obj_id in self._aborted_transfer_obj_ids: + self._aborted_transfer_obj_ids.remove(obj_id) + raise RuntimeError( + f"NIXL transfer aborted for object id: {obj_id}" + ) + time.sleep(0.001) # Avoid busy waiting + elif state == "DONE": + break + finally: + # We could raise errors or NIXL could raise errors like NIXL_ERR_REMOTE_DISCONNECT, + # so doing best effort cleanup. + with self._aborted_transfer_obj_ids_lock: + self._aborted_transfer_obj_ids.discard(obj_id) + if xfer_handle: + nixl_agent.release_xfer_handle(xfer_handle) + if remote_name: + nixl_agent.remove_remote_agent(remote_name) + if local_descs: + nixl_agent.deregister_memory(local_descs) + + def get_nixl_metadata( + self, tensors: List["torch.Tensor"] + ) -> Tuple[Any, bytes, bytes]: + """Get NIXL metadata for a set of tensors. + + Args: + tensors: List of tensors to get metadata for. + + Returns: + tuple: A tuple containing: + - Serialized NIXL descriptors for the tensors + - Metadata about this NIXL agent + """ + nixl_agent = self._nixl_agent + reg_descs = nixl_agent.register_memory(tensors) + xfer_descs = reg_descs.trim() + return ( + reg_descs, + nixl_agent.get_serialized_descs(xfer_descs), + nixl_agent.get_agent_metadata(), + ) + + def deregister_memory(self, descs: Any): + self._nixl_agent.deregister_memory(descs) + + def abort(self, obj_id: str): + with self._aborted_transfer_obj_ids_lock: + self._aborted_transfer_obj_ids.add(obj_id) diff --git a/lib/python3.12/site-packages/ray/util/collective/collective_group/torch_gloo_collective_group.py b/lib/python3.12/site-packages/ray/util/collective/collective_group/torch_gloo_collective_group.py new file mode 100644 index 0000000000000000000000000000000000000000..d2314c5ea54a72d769e2ea0b70782baca9df6cb1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/collective_group/torch_gloo_collective_group.py @@ -0,0 +1,229 @@ +import os +from typing import TYPE_CHECKING, List, Optional + +import numpy as np +import torch +import torch.distributed as dist + +import ray.experimental.internal_kv as internal_kv +from ray.util.collective.collective_group.base_collective_group import BaseGroup +from ray.util.collective.types import ( + AllGatherOptions, + AllReduceOptions, + Backend, + BarrierOptions, + BroadcastOptions, + RecvOptions, + ReduceOp, + ReduceOptions, + ReduceScatterOptions, + SendOptions, +) + +if TYPE_CHECKING: + import torch + + +TORCH_REDUCE_OP_MAP = { + ReduceOp.SUM: dist.ReduceOp.SUM, + ReduceOp.PRODUCT: dist.ReduceOp.PRODUCT, + ReduceOp.MIN: dist.ReduceOp.MIN, + ReduceOp.MAX: dist.ReduceOp.MAX, +} + + +def get_master_address_metadata_key(group_name: str): + return f"collective_group_master_address_{group_name}" + + +class TorchGLOOGroup(BaseGroup): + def __init__( + self, + world_size: int, + rank: int, + group_name: str, + gloo_timeout: Optional[int] = None, + ): + # Initialize the default process group only once per process. + if not dist.is_initialized(): + metadata_key = get_master_address_metadata_key(group_name) + try: + metadata = internal_kv._internal_kv_get(metadata_key) + except ValueError: + raise RuntimeError( + f"TorchGLOOGroup expected metadata in internal_kv with name `{metadata_key}`. " + "TorchGLOOGroup should not be instantiated directly. " + "Use ray.experimental.collective.create_collective_group to create the group." + ) + if metadata is None: + raise RuntimeError( + f"Missing rendezvous metadata for group `{group_name}` under key `{metadata_key}`." + ) + metadata = metadata.decode() + master_addr, master_port = metadata.split(":") + os.environ["MASTER_ADDR"] = master_addr + os.environ["MASTER_PORT"] = master_port + + dist.init_process_group( + backend="gloo", init_method="env://", world_size=world_size, rank=rank + ) + + super().__init__(world_size, rank, group_name) + + # Create a subgroup for this logical group. For the default group, use WORLD. + self._is_default_group = group_name == "default" + if self._is_default_group: + self._pg = dist.group.WORLD + else: + # All ranks participate in this subgroup with global ranks [0..world_size-1]. + ranks = list(range(world_size)) + self._pg = dist.new_group(ranks=ranks, backend="gloo") + + # Compatibility shim for legacy tests expecting a pygloo context with getTimeout(). + # Store the rendezvous timeout in milliseconds, defaulting to 30000 if unspecified. + class _GlooCompatContext: + def __init__(self, timeout_ms: int): + self._timeout_ms = timeout_ms + + def getTimeout(self) -> int: + return self._timeout_ms + + self._gloo_context = _GlooCompatContext( + gloo_timeout if gloo_timeout is not None else 30000 + ) + + def destroy_group(self): + """GC the communicators.""" + # Destroy only the subgroup for non-default groups. Allow default to be torn down explicitly. + if self._is_default_group: + # Destroy default process group to allow re-init in tests that recreate the same group. + dist.destroy_process_group() + else: + # Destroy just this subgroup. + if self._pg is not None: + dist.destroy_process_group(self._pg) + + @classmethod + def backend(cls): + """The backend of this collective group.""" + return Backend.TORCH_GLOO + + def _check_tensor_input(self, tensor: List["torch.Tensor"]) -> "torch.Tensor": + """ray.util.collective wraps tensor arguments in a list. + Accept a single torch.Tensor or numpy.ndarray and unwrap/convert it. + """ + assert isinstance(tensor, list) and len(tensor) == 1 + t = tensor[0] + if isinstance(t, torch.Tensor): + return t + if isinstance(t, np.ndarray): + return torch.from_numpy(t) + raise ValueError( + f"torch_gloo group only accepts torch.Tensor or numpy.ndarray, received {type(t)}" + ) + + def _check_tensor_list_input( + self, tensor_list: List[List["torch.Tensor"]] + ) -> List["torch.Tensor"]: + """ray.util.collective wraps tensor arguments in a list. + Accept a single list containing torch.Tensors or numpy.ndarrays and + unwrap/convert items as needed. + """ + assert isinstance(tensor_list, list) and len(tensor_list) == 1 + tensor_list = tensor_list[0] + converted_tensor_list = [] + for tensor in tensor_list: + if isinstance(tensor, np.ndarray): + tensor = torch.from_numpy(tensor) + converted_tensor_list.append(tensor) + elif isinstance(tensor, torch.Tensor): + converted_tensor_list.append(tensor) + else: + raise ValueError( + f"torch_gloo group only accepts torch.Tensor or numpy.ndarray types, received tensor list with value {tensor}" + ) + return converted_tensor_list + + def allreduce( + self, + tensor: List["torch.Tensor"], + allreduce_options: Optional[AllReduceOptions] = None, + ) -> None: + if allreduce_options is None: + allreduce_options = AllReduceOptions() + tensor = self._check_tensor_input(tensor) + torch_reduce_op = TORCH_REDUCE_OP_MAP[allreduce_options.reduceOp] + dist.all_reduce(tensor, op=torch_reduce_op, group=self._pg) + + def barrier(self, barrier_options=BarrierOptions()) -> None: + dist.barrier(group=self._pg) + + def reduce( + self, + tensor: List["torch.Tensor"], + reduce_options: Optional[ReduceOptions] = None, + ) -> None: + if reduce_options is None: + reduce_options = ReduceOptions() + t = self._check_tensor_input(tensor) + torch_reduce_op = TORCH_REDUCE_OP_MAP[reduce_options.reduceOp] + # Avoid mutating non-root ranks' user tensors to match util.collective semantics. + if self._rank == reduce_options.root_rank: + dist.reduce( + t, dst=reduce_options.root_rank, op=torch_reduce_op, group=self._pg + ) + else: + tmp = t.detach().clone() + dist.reduce( + tmp, dst=reduce_options.root_rank, op=torch_reduce_op, group=self._pg + ) + + def allgather( + self, + tensor_list: List[List["torch.Tensor"]], + tensor: List["torch.Tensor"], + allgather_options: Optional[AllGatherOptions] = None, + ) -> None: + if allgather_options is None: + allgather_options = AllGatherOptions() + tensor_list = self._check_tensor_list_input(tensor_list) + tensor = self._check_tensor_input(tensor) + dist.all_gather(tensor_list, tensor, group=self._pg) + + def broadcast( + self, tensor: List["torch.Tensor"], broadcast_options=BroadcastOptions() + ) -> None: + tensor = self._check_tensor_input(tensor) + dist.broadcast(tensor, src=broadcast_options.root_rank, group=self._pg) + + def reducescatter( + self, + output_tensor: List["torch.Tensor"], + tensor_list: List[List["torch.Tensor"]], + reducescatter_options: Optional[ReduceScatterOptions] = None, + ) -> None: + if reducescatter_options is None: + reducescatter_options = ReduceScatterOptions() + tensor_list = self._check_tensor_list_input(tensor_list) + output_tensor = self._check_tensor_input(output_tensor) + if output_tensor.shape != tensor_list[self._rank].shape: + raise ValueError( + "Output tensor has wrong shape {output_tensor.shape}, expected {tensor_list[self._rank].shape}" + ) + torch_reduce_op = TORCH_REDUCE_OP_MAP[reducescatter_options.reduceOp] + + # torch.distributed gloo doesn't support reducescatter. Implement a + # simple version using allreduce. + for tensor in tensor_list: + dist.all_reduce(tensor, op=torch_reduce_op, group=self._pg) + + if output_tensor.data_ptr() != tensor_list[self._rank].data_ptr(): + output_tensor.copy_(tensor_list[self._rank]) + + def send(self, tensor: List["torch.Tensor"], send_options: SendOptions) -> None: + tensor = self._check_tensor_input(tensor) + dist.send(tensor, dst=send_options.dst_rank) + + def recv(self, tensor: List["torch.Tensor"], recv_options: RecvOptions) -> None: + tensor = self._check_tensor_input(tensor) + dist.recv(tensor, src=recv_options.src_rank) diff --git a/lib/python3.12/site-packages/ray/util/collective/const.py b/lib/python3.12/site-packages/ray/util/collective/const.py new file mode 100644 index 0000000000000000000000000000000000000000..35a11d23abbf4f3708acc3c76cb4f9dde93e10cc --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/const.py @@ -0,0 +1,34 @@ +""" +Constants. + +Contains constants used to setup collective groups. +""" +import hashlib +import os +from enum import Enum, auto + + +def get_store_name(group_name): + """Generate the unique name for the NCCLUniqueID store (named actor). + + Args: + group_name: unique user name for the store. + Return: + str: SHA1-hexlified name for the store. + """ + if not group_name: + raise ValueError("group_name is None.") + hexlified_name = hashlib.sha1(group_name.encode()).hexdigest() + return hexlified_name + + +class ENV(Enum): + """ray.util.collective environment variables.""" + + NCCL_USE_MULTISTREAM = auto(), lambda v: (v or "True") == "True" + + @property + def val(self): + """Return the output of the lambda against the system's env value.""" + _, default_fn = self.value + return default_fn(os.getenv(self.name)) diff --git a/lib/python3.12/site-packages/ray/util/collective/types.py b/lib/python3.12/site-packages/ray/util/collective/types.py new file mode 100644 index 0000000000000000000000000000000000000000..c0f395d6e5d770ce3b5ebe56a1f2b6bb3ebf8de5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/types.py @@ -0,0 +1,195 @@ +"""Types conversion between different backends.""" + +from dataclasses import dataclass +from datetime import timedelta +from enum import Enum +from typing import TYPE_CHECKING, Any, List, Optional, Tuple + +from numpy import int32 + +_NUMPY_AVAILABLE = True +_TORCH_AVAILABLE = True +_CUPY_AVAILABLE = True + +if TYPE_CHECKING: + import torch + +try: + import torch as th # noqa: F401 +except ImportError: + _TORCH_AVAILABLE = False + +try: + import cupy as cp # noqa: F401 +except ImportError: + _CUPY_AVAILABLE = False + + +def cupy_available(): + return _CUPY_AVAILABLE + + +def torch_available(): + return _TORCH_AVAILABLE + + +class Backend(object): + """A class to represent different backends.""" + + NCCL = "nccl" + MPI = "mpi" + # `pygloo` is deprecated. Use gloo through torch.distributed for both + # `GLOO` and `TORCH_GLOO`. + GLOO = "gloo" + # Use gloo through torch.distributed. + TORCH_GLOO = "torch_gloo" + NIXL = "nixl" + UNRECOGNIZED = "unrecognized" + + def __new__(cls, name: str): + backend = getattr(Backend, name.upper(), Backend.UNRECOGNIZED) + if backend == Backend.UNRECOGNIZED: + raise ValueError( + "Unrecognized backend: '{}'. Only NCCL is supported".format(name) + ) + if backend == Backend.MPI: + raise RuntimeError("Ray does not support MPI backend.") + return backend + + +@dataclass +class TensorTransportMetadata: + """Metadata for tensors stored in the GPU object store. + + Args: + tensor_meta: A list of tuples, each containing the shape and dtype of a tensor. + tensor_device: The device of the tensor. Currently, we require all tensors in the + list have the same device type. + """ + + tensor_meta: List[Tuple["torch.Size", "torch.dtype"]] + tensor_device: Optional["torch.device"] = None + + +@dataclass +class NixlTransportMetadata(TensorTransportMetadata): + """Metadata for tensors stored in the GPU object store for NIXL transport. + + Args: + nixl_serialized_descs: Serialized tensor descriptors for NIXL transport. + nixl_agent_meta: The additional metadata of the remote NIXL agent. + """ + + nixl_reg_descs: Optional[Any] = None + nixl_serialized_descs: Optional[bytes] = None + nixl_agent_meta: Optional[bytes] = None + + __eq__ = object.__eq__ + __hash__ = object.__hash__ + + +@dataclass +class CollectiveTransportMetadata(TensorTransportMetadata): + """Metadata for tensors stored in the GPU object store for collective transport.""" + + +@dataclass +class CommunicatorMetadata: + """Metadata for the communicator. + + Args: + communicator_name: The name of the communicator. + """ + + communicator_name: str = "" + + +@dataclass +class CollectiveCommunicatorMetadata(CommunicatorMetadata): + """Metadata for the collective communicator (e.g. NCCL, GLOO). + + Args: + src_rank: The rank of the source actor. + dst_rank: The rank of the destination actor. + """ + + src_rank: Optional[int32] = None + dst_rank: Optional[int32] = None + + +@dataclass +class NixlCommunicatorMetadata(CommunicatorMetadata): + """Metadata for the NIXL communicator.""" + + +class ReduceOp(Enum): + SUM = 0 + PRODUCT = 1 + MIN = 2 + MAX = 3 + + +unset_timeout_ms = timedelta(milliseconds=-1) + +# This is used to identify the collective group for NIXL. +NIXL_GROUP_NAME = "ray_internal_nixl_group" + + +@dataclass +class AllReduceOptions: + reduceOp = ReduceOp.SUM + timeout_ms = unset_timeout_ms + + +@dataclass +class BarrierOptions: + timeout_ms = unset_timeout_ms + + +@dataclass +class ReduceOptions: + reduceOp = ReduceOp.SUM + root_rank = 0 + root_tensor = 0 # index for multi-gpu reduce operations + timeout_ms = unset_timeout_ms + + +@dataclass +class AllGatherOptions: + timeout_ms = unset_timeout_ms + + +# +# @dataclass +# class GatherOptions: +# root_rank = 0 +# timeout = unset_timeout + + +@dataclass +class BroadcastOptions: + root_rank = 0 + root_tensor = 0 + timeout_ms = unset_timeout_ms + + +@dataclass +class ReduceScatterOptions: + reduceOp = ReduceOp.SUM + timeout_ms = unset_timeout_ms + + +@dataclass +class SendOptions: + dst_rank = 0 + dst_gpu_index = 0 + n_elements = 0 + timeout_ms = unset_timeout_ms + + +@dataclass +class RecvOptions: + src_rank = 0 + src_gpu_index = 0 + n_elements = 0 + unset_timeout_ms = unset_timeout_ms diff --git a/lib/python3.12/site-packages/ray/util/collective/util.py b/lib/python3.12/site-packages/ray/util/collective/util.py new file mode 100644 index 0000000000000000000000000000000000000000..02221995fd606fc6cd7aed944fa38dee21c7d4ed --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/collective/util.py @@ -0,0 +1,85 @@ +"""Some utility class for Collectives.""" +import asyncio +import logging + +import ray + +logger = logging.getLogger(__name__) + + +@ray.remote +class NCCLUniqueIDStore: + """NCCLUniqueID Store as a named actor class. + + Args: + name: the unique name for this named actor. + + Attributes: + name: the unique name for this named actor. + nccl_id: the NCCLUniqueID held in this store. + """ + + def __init__(self, name): + self.name = name + self.nccl_id = None + self.event = asyncio.Event() + + async def set_id(self, uid): + """ + Initialize the NCCL unique ID for this store. + + Args: + uid: the unique ID generated via the NCCL generate_communicator_id API. + + Returns: + The NCCL unique ID set. + """ + self.nccl_id = uid + self.event.set() + return uid + + async def wait_and_get_id(self): + """Wait for the NCCL unique ID to be set and return it.""" + await self.event.wait() + return self.nccl_id + + def get_id(self): + """Get the NCCL unique ID held in this store.""" + if not self.nccl_id: + logger.warning( + "The NCCL ID has not been set yet for store {}.".format(self.name) + ) + return self.nccl_id + + +@ray.remote +class Info: + """Store the group information created via `create_collective_group`. + + Note: Should be used as a NamedActor. + """ + + def __init__(self): + self.ids = None + self.world_size = -1 + self.rank = -1 + self.backend = None + self.gloo_timeout = 30000 + + def set_info(self, ids, world_size, rank, backend, gloo_timeout): + """Store collective information.""" + self.ids = ids + self.world_size = world_size + self.rank = rank + self.backend = backend + self.gloo_timeout = gloo_timeout + + def get_info(self): + """Get previously stored collective information.""" + return ( + self.ids, + self.world_size, + self.rank, + self.backend, + self.gloo_timeout, + ) diff --git a/lib/python3.12/site-packages/ray/util/common.py b/lib/python3.12/site-packages/ray/util/common.py new file mode 100644 index 0000000000000000000000000000000000000000..081ee03ef25e9678d978e9192040ad19b54b2187 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/common.py @@ -0,0 +1 @@ +INT32_MAX = (2**31) - 1 diff --git a/lib/python3.12/site-packages/ray/util/debug.py b/lib/python3.12/site-packages/ray/util/debug.py new file mode 100644 index 0000000000000000000000000000000000000000..29d9a3e1d4978c3aea86aa5330dec66a14c0309f --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/debug.py @@ -0,0 +1,275 @@ +import gc +import os +import re +import time +import tracemalloc +from collections import defaultdict, namedtuple +from typing import Callable, List, Optional + +from ray.util.annotations import DeveloperAPI + +_logged = set() +_disabled = False +_periodic_log = False +_last_logged = 0.0 + + +@DeveloperAPI +def log_once(key): + """Returns True if this is the "first" call for a given key. + + Various logging settings can adjust the definition of "first". + + Example: + + .. testcode:: + + import logging + from ray.util.debug import log_once + + logger = logging.getLogger(__name__) + if log_once("some_key"): + logger.info("Some verbose logging statement") + """ + + global _last_logged + + if _disabled: + return False + elif key not in _logged: + _logged.add(key) + _last_logged = time.time() + return True + elif _periodic_log and time.time() - _last_logged > 60.0: + _logged.clear() + _last_logged = time.time() + return False + else: + return False + + +@DeveloperAPI +def disable_log_once_globally(): + """Make log_once() return False in this process.""" + + global _disabled + _disabled = True + + +@DeveloperAPI +def enable_periodic_logging(): + """Make log_once() periodically return True in this process.""" + + global _periodic_log + _periodic_log = True + + +@DeveloperAPI +def reset_log_once(key: Optional[str] = None): + """Resets log_once for the provided key. + + If you don't provide a key, resets log_once for all keys. + """ + if key is None: + _logged.clear() + else: + _logged.discard(key) + + +# A suspicious memory-allocating stack-trace that we should re-test +# to make sure it's not a false positive. +Suspect = DeveloperAPI( + namedtuple( + "Suspect", + [ + # The stack trace of the allocation, going back n frames, depending + # on the tracemalloc.start(n) call. + "traceback", + # The amount of memory taken by this particular stack trace + # over the course of the experiment. + "memory_increase", + # The slope of the scipy linear regression (x=iteration; y=memory size). + "slope", + # The rvalue of the scipy linear regression. + "rvalue", + # The memory size history (list of all memory sizes over all iterations). + "hist", + ], + ) +) + + +def _test_some_code_for_memory_leaks( + desc: str, + init: Optional[Callable[[], None]], + code: Callable[[], None], + repeats: int, + max_num_trials: int = 1, +) -> List[Suspect]: + """Runs given code (and init code) n times and checks for memory leaks. + + Args: + desc: A descriptor of the test. + init: Optional code to be executed initially. + code: The actual code to be checked for producing memory leaks. + repeats: How many times to repeatedly execute `code`. + max_num_trials: The maximum number of trials to run. A new trial is only + run, if the previous one produced a memory leak. For all non-1st trials, + `repeats` calculates as: actual_repeats = `repeats` * (trial + 1), where + the first trial is 0. + + Returns: + A list of Suspect objects, describing possible memory leaks. If list + is empty, no leaks have been found. + """ + + def _i_print(i): + if (i + 1) % 10 == 0: + print(".", end="" if (i + 1) % 100 else f" {i + 1}\n", flush=True) + + # Do n trials to make sure a found leak is really one. + suspicious = set() + suspicious_stats = [] + for trial in range(max_num_trials): + # Store up to n frames of each call stack. + tracemalloc.start(20) + + table = defaultdict(list) + + # Repeat running code for n times. + # Increase repeat value with each trial to make sure stats are more + # solid each time (avoiding false positives). + actual_repeats = repeats * (trial + 1) + + print(f"{desc} {actual_repeats} times.") + + # Initialize if necessary. + if init is not None: + init() + # Run `code` n times, each time taking a memory snapshot. + for i in range(actual_repeats): + _i_print(i) + # Manually trigger garbage collection before and after code runs in order to + # make tracemalloc snapshots as accurate as possible. + gc.collect() + code() + gc.collect() + _take_snapshot(table, suspicious) + print("\n") + + # Check, which traces have moved up in their memory consumption + # constantly over time. + suspicious.clear() + suspicious_stats.clear() + # Suspicious memory allocation found? + suspects = _find_memory_leaks_in_table(table) + for suspect in sorted(suspects, key=lambda s: s.memory_increase, reverse=True): + # Only print out the biggest offender: + if len(suspicious) == 0: + _pprint_suspect(suspect) + print("-> added to retry list") + suspicious.add(suspect.traceback) + suspicious_stats.append(suspect) + + tracemalloc.stop() + + # Some suspicious memory allocations found. + if len(suspicious) > 0: + print(f"{len(suspicious)} suspects found. Top-ten:") + for i, s in enumerate(suspicious_stats): + if i > 10: + break + print( + f"{i}) line={s.traceback[-1]} mem-increase={s.memory_increase}B " + f"slope={s.slope}B/detection rval={s.rvalue}" + ) + # Nothing suspicious found -> Exit trial loop and return. + else: + print("No remaining suspects found -> returning") + break + + # Print out final top offender. + if len(suspicious_stats) > 0: + _pprint_suspect(suspicious_stats[0]) + + return suspicious_stats + + +def _take_snapshot(table, suspicious=None): + # Take a memory snapshot. + snapshot = tracemalloc.take_snapshot() + # Group all memory allocations by their stacktrace (going n frames + # deep as defined above in tracemalloc.start(n)). + # Then sort groups by size, then count, then trace. + top_stats = snapshot.statistics("traceback") + + # For the first m largest increases, keep only, if a) first trial or b) those + # that are already in the `suspicious` set. + for stat in top_stats[:100]: + if not suspicious or stat.traceback in suspicious: + table[stat.traceback].append(stat.size) + + +def _find_memory_leaks_in_table(table): + import numpy as np + import scipy.stats + + suspects = [] + + for traceback, hist in table.items(): + # Do a quick mem increase check. + memory_increase = hist[-1] - hist[0] + + # Only if memory increased, do we check further. + if memory_increase <= 0.0: + continue + + # Ignore this very module here (we are collecting lots of data + # so an increase is expected). + top_stack = str(traceback[-1]) + drive_separator = "\\\\" if os.name == "nt" else "/" + if any( + s in top_stack + for s in [ + "tracemalloc", + "pycharm", + "thirdparty_files/psutil", + re.sub("\\.", drive_separator, __name__) + ".py", + ] + ): + continue + + # Do a linear regression to get the slope and R-value. + line = scipy.stats.linregress(x=np.arange(len(hist)), y=np.array(hist)) + + # - If weak positive slope and some confidence and + # increase > n bytes -> error. + # - If stronger positive slope -> error. + if memory_increase > 1000 and ( + (line.slope > 60.0 and line.rvalue > 0.875) + or (line.slope > 20.0 and line.rvalue > 0.9) + or (line.slope > 10.0 and line.rvalue > 0.95) + ): + suspects.append( + Suspect( + traceback=traceback, + memory_increase=memory_increase, + slope=line.slope, + rvalue=line.rvalue, + hist=hist, + ) + ) + + return suspects + + +def _pprint_suspect(suspect): + print( + "Most suspicious memory allocation in traceback " + "(only printing out this one, but all (less suspicious)" + " suspects will be investigated as well):" + ) + print("\n".join(suspect.traceback.format())) + print(f"Increase total={suspect.memory_increase}B") + print(f"Slope={suspect.slope} B/detection") + print(f"Rval={suspect.rvalue}") diff --git a/lib/python3.12/site-packages/ray/util/debugpy.py b/lib/python3.12/site-packages/ray/util/debugpy.py new file mode 100644 index 0000000000000000000000000000000000000000..1f5a0157f2b6eed7fa8cd510fd3ac9eae6c7d1cc --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/debugpy.py @@ -0,0 +1,137 @@ +import importlib +import logging +import os +import sys +import threading + +import ray +from ray._common.network_utils import build_address +from ray.util.annotations import DeveloperAPI + +log = logging.getLogger(__name__) + +POST_MORTEM_ERROR_UUID = "post_mortem_error_uuid" + + +def _try_import_debugpy(): + try: + debugpy = importlib.import_module("debugpy") + if not hasattr(debugpy, "__version__") or debugpy.__version__ < "1.8.0": + raise ImportError() + return debugpy + except (ModuleNotFoundError, ImportError): + log.error( + "Module 'debugpy>=1.8.0' cannot be loaded. " + "Ray Debugpy Debugger will not work without 'debugpy>=1.8.0' installed. " + "Install this module using 'pip install debugpy==1.8.0' " + ) + return None + + +# A lock to ensure that only one thread can open the debugger port. +debugger_port_lock = threading.Lock() + + +def _override_breakpoint_hooks(): + """ + This method overrides the breakpoint() function to set_trace() + so that other threads can reuse the same setup logic. + This is based on: https://github.com/microsoft/debugpy/blob/ef9a67fe150179ee4df9997f9273723c26687fab/src/debugpy/_vendored/pydevd/pydev_sitecustomize/sitecustomize.py#L87 # noqa: E501 + """ + sys.__breakpointhook__ = set_trace + sys.breakpointhook = set_trace + import builtins as __builtin__ + + __builtin__.breakpoint = set_trace + + +def _ensure_debugger_port_open_thread_safe(): + """ + This is a thread safe method that ensure that the debugger port + is open, and if not, open it. + """ + + # The lock is acquired before checking the debugger port so only + # one thread can open the debugger port. + with debugger_port_lock: + debugpy = _try_import_debugpy() + if not debugpy: + return + + debugger_port = ray._private.worker.global_worker.debugger_port + if not debugger_port: + (host, port) = debugpy.listen( + (ray._private.worker.global_worker.node_ip_address, 0) + ) + ray._private.worker.global_worker.set_debugger_port(port) + log.info(f"Ray debugger is listening on {build_address(host, port)}") + else: + log.info(f"Ray debugger is already open on {debugger_port}") + + +@DeveloperAPI +def set_trace(breakpoint_uuid=None): + """Interrupt the flow of the program and drop into the Ray debugger. + Can be used within a Ray task or actor. + """ + debugpy = _try_import_debugpy() + if not debugpy: + return + + _ensure_debugger_port_open_thread_safe() + + # debugpy overrides the breakpoint() function, so we need to set it back + # so other threads can reuse it. + _override_breakpoint_hooks() + + with ray._private.worker.global_worker.worker_paused_by_debugger(): + msg = ( + "Waiting for debugger to attach (see " + "https://docs.ray.io/en/latest/ray-observability/" + "ray-distributed-debugger.html)..." + ) + log.info(msg) + debugpy.wait_for_client() + + log.info("Debugger client is connected") + if breakpoint_uuid == POST_MORTEM_ERROR_UUID: + _debugpy_excepthook() + else: + _debugpy_breakpoint() + + +def _debugpy_breakpoint(): + """ + Drop the user into the debugger on a breakpoint. + """ + import pydevd + + pydevd.settrace(stop_at_frame=sys._getframe().f_back) + + +def _debugpy_excepthook(): + """ + Drop the user into the debugger on an unhandled exception. + """ + import threading + + import pydevd + + py_db = pydevd.get_global_debugger() + thread = threading.current_thread() + additional_info = py_db.set_additional_thread_info(thread) + additional_info.is_tracing += 1 + try: + error = sys.exc_info() + py_db.stop_on_unhandled_exception(py_db, thread, additional_info, error) + sys.excepthook(error[0], error[1], error[2]) + finally: + additional_info.is_tracing -= 1 + + +def _is_ray_debugger_post_mortem_enabled(): + return os.environ.get("RAY_DEBUG_POST_MORTEM", "0") == "1" + + +def _post_mortem(): + return set_trace(POST_MORTEM_ERROR_UUID) diff --git a/lib/python3.12/site-packages/ray/util/horovod/__init__.py b/lib/python3.12/site-packages/ray/util/horovod/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a08001abc7961a19474b54b31bf89be0757df6cd --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/horovod/__init__.py @@ -0,0 +1,4 @@ +raise DeprecationWarning( + "ray.util.horovod has been removed as of Ray 2.0. Instead, use the `horovod` " + "library directly or the `HorovodTrainer` in Ray Train." +) diff --git a/lib/python3.12/site-packages/ray/util/horovod/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/horovod/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2b574d21636ef3fdcf268ac3c4f7822519279154 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/horovod/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/iter_metrics.py b/lib/python3.12/site-packages/ray/util/iter_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..eb06a97c3ace41d230e1d16b329126272d71a3ff --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/iter_metrics.py @@ -0,0 +1,69 @@ +import collections +from typing import List + +from ray.util.annotations import Deprecated +from ray.util.timer import _Timer + + +@Deprecated +class MetricsContext: + """Metrics context object for a local iterator. + + This object is accessible by all operators of a local iterator. It can be + used to store and retrieve global execution metrics for the iterator. + It can be accessed by calling LocalIterator.get_metrics(), which is only + allowable inside iterator functions. + + Attributes: + counters: dict storing increasing metrics. + timers: dict storing latency timers. + info: dict storing misc metric values. + current_actor: reference to the actor handle that + produced the current iterator output. This is automatically set + for gather_async(). + """ + + def __init__(self): + self.counters = collections.defaultdict(int) + self.timers = collections.defaultdict(_Timer) + self.info = {} + self.current_actor = None + + def save(self): + """Return a serializable copy of this context.""" + return { + "counters": dict(self.counters), + "info": dict(self.info), + "timers": None, # TODO(ekl) consider persisting timers too + } + + def restore(self, values): + """Restores state given the output of save().""" + self.counters.clear() + self.counters.update(values["counters"]) + self.timers.clear() + self.info = values["info"] + + +@Deprecated +class SharedMetrics: + """Holds an indirect reference to a (shared) metrics context. + + This is used by LocalIterator.union() to point the metrics contexts of + entirely separate iterator chains to the same underlying context.""" + + def __init__( + self, metrics: MetricsContext = None, parents: List["SharedMetrics"] = None + ): + self.metrics = metrics or MetricsContext() + self.parents = parents or [] + self.set(self.metrics) + + def set(self, metrics): + """Recursively set self and parents to point to the same metrics.""" + self.metrics = metrics + for parent in self.parents: + parent.set(metrics) + + def get(self): + return self.metrics diff --git a/lib/python3.12/site-packages/ray/util/lightgbm/__init__.py b/lib/python3.12/site-packages/ray/util/lightgbm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9a46aefde633525e93463c500b0a229d05391449 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/lightgbm/__init__.py @@ -0,0 +1,4 @@ +raise DeprecationWarning( + "ray.util.lightgbm has been removed as of Ray 2.0. Instead, use the `lightgbm-ray` " + "library directly or the `LightGBMTrainer` in Ray Train." +) diff --git a/lib/python3.12/site-packages/ray/util/lightgbm/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/lightgbm/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e688af3cde02f730ea3ca31bf96b0fd46fdcbfe Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/lightgbm/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/multiprocessing/__init__.py b/lib/python3.12/site-packages/ray/util/multiprocessing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..75c07d911814581664e77fafefb8f21eee2f2945 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/multiprocessing/__init__.py @@ -0,0 +1,5 @@ +from multiprocessing import JoinableQueue, TimeoutError + +from .pool import Pool + +__all__ = ["Pool", "TimeoutError", "JoinableQueue"] diff --git a/lib/python3.12/site-packages/ray/util/multiprocessing/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/multiprocessing/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..83b414d38606f83a144fe49e20a4d7d158f7396e Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/multiprocessing/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/multiprocessing/__pycache__/pool.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/multiprocessing/__pycache__/pool.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..702781f57d39e0d7f0adb0d1b1877ea3a24e7224 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/multiprocessing/__pycache__/pool.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/multiprocessing/pool.py b/lib/python3.12/site-packages/ray/util/multiprocessing/pool.py new file mode 100644 index 0000000000000000000000000000000000000000..980626ca6eba0cbb2ad72c67de6f5655b1f1f441 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/multiprocessing/pool.py @@ -0,0 +1,1008 @@ +import collections +import copy +import gc +import itertools +import logging +import os +import queue +import sys +import threading +import time +from multiprocessing import TimeoutError +from typing import Any, Callable, Dict, Hashable, Iterable, List, Optional, Tuple + +import ray +from ray._common.usage import usage_lib +from ray.util import log_once + +try: + from joblib._parallel_backends import SafeFunction + from joblib.parallel import BatchedCalls, parallel_backend +except ImportError: + BatchedCalls = None + parallel_backend = None + SafeFunction = None + + +logger = logging.getLogger(__name__) + +RAY_ADDRESS_ENV = "RAY_ADDRESS" + + +def _put_in_dict_registry( + obj: Any, registry_hashable: Dict[Hashable, ray.ObjectRef] +) -> ray.ObjectRef: + if obj not in registry_hashable: + ret = ray.put(obj) + registry_hashable[obj] = ret + else: + ret = registry_hashable[obj] + return ret + + +def _put_in_list_registry( + obj: Any, registry: List[Tuple[Any, ray.ObjectRef]] +) -> ray.ObjectRef: + try: + ret = next((ref for o, ref in registry if o is obj)) + except StopIteration: + ret = ray.put(obj) + registry.append((obj, ret)) + return ret + + +def ray_put_if_needed( + obj: Any, + registry: Optional[List[Tuple[Any, ray.ObjectRef]]] = None, + registry_hashable: Optional[Dict[Hashable, ray.ObjectRef]] = None, +) -> ray.ObjectRef: + """ray.put obj in object store if it's not an ObjRef and bigger than 100 bytes, + with support for list and dict registries""" + if isinstance(obj, ray.ObjectRef) or sys.getsizeof(obj) < 100: + return obj + ret = obj + if registry_hashable is not None: + try: + ret = _put_in_dict_registry(obj, registry_hashable) + except TypeError: + if registry is not None: + ret = _put_in_list_registry(obj, registry) + elif registry is not None: + ret = _put_in_list_registry(obj, registry) + return ret + + +def ray_get_if_needed(obj: Any) -> Any: + """If obj is an ObjectRef, do ray.get, otherwise return obj""" + if isinstance(obj, ray.ObjectRef): + return ray.get(obj) + return obj + + +if BatchedCalls is not None: + + class RayBatchedCalls(BatchedCalls): + """Joblib's BatchedCalls with basic Ray object store management + + This functionality is provided through the put_items_in_object_store, + which uses external registries (list and dict) containing objects + and their ObjectRefs.""" + + def put_items_in_object_store( + self, + registry: Optional[List[Tuple[Any, ray.ObjectRef]]] = None, + registry_hashable: Optional[Dict[Hashable, ray.ObjectRef]] = None, + ): + """Puts all applicable (kw)args in self.items in object store + + Takes two registries - list for unhashable objects and dict + for hashable objects. The registries are a part of a Pool object. + The method iterates through all entries in items list (usually, + there will be only one, but the number depends on joblib Parallel + settings) and puts all of the args and kwargs into the object + store, updating the registries. + If an arg or kwarg is already in a registry, it will not be + put again, and instead, the cached object ref will be used.""" + new_items = [] + for func, args, kwargs in self.items: + args = [ + ray_put_if_needed(arg, registry, registry_hashable) for arg in args + ] + kwargs = { + k: ray_put_if_needed(v, registry, registry_hashable) + for k, v in kwargs.items() + } + new_items.append((func, args, kwargs)) + self.items = new_items + + def __call__(self): + # Exactly the same as in BatchedCalls, with the + # difference being that it gets args and kwargs from + # object store (which have been put in there by + # put_items_in_object_store) + + # Set the default nested backend to self._backend but do + # not set the change the default number of processes to -1 + with parallel_backend(self._backend, n_jobs=self._n_jobs): + return [ + func( + *[ray_get_if_needed(arg) for arg in args], + **{k: ray_get_if_needed(v) for k, v in kwargs.items()}, + ) + for func, args, kwargs in self.items + ] + + def __reduce__(self): + # Exactly the same as in BatchedCalls, with the + # difference being that it returns RayBatchedCalls + # instead + if self._reducer_callback is not None: + self._reducer_callback() + # no need pickle the callback. + return ( + RayBatchedCalls, + (self.items, (self._backend, self._n_jobs), None, self._pickle_cache), + ) + +else: + RayBatchedCalls = None + + +# Helper function to divide a by b and round the result up. +def div_round_up(a, b): + return -(-a // b) + + +class PoolTaskError(Exception): + def __init__(self, underlying): + self.underlying = underlying + + +class ResultThread(threading.Thread): + """Thread that collects results from distributed actors. + + It winds down when either: + - A pre-specified number of objects has been processed + - When the END_SENTINEL (submitted through self.add_object_ref()) + has been received and all objects received before that have been + processed. + + Initialize the thread with total_object_refs = float('inf') to wait for the + END_SENTINEL. + + Args: + object_refs (List[RayActorObjectRefs]): ObjectRefs to Ray Actor calls. + Thread tracks whether they are ready. More ObjectRefs may be added + with add_object_ref (or _add_object_ref internally) until the object + count reaches total_object_refs. + single_result: Should be True if the thread is managing function + with a single result (like apply_async). False if the thread is managing + a function with a List of results. + callback: called only once at the end of the thread + if no results were errors. If single_result=True, and result is + not an error, callback is invoked with the result as the only + argument. If single_result=False, callback is invoked with + a list of all the results as the only argument. + error_callback: called only once on the first result + that errors. Should take an Exception as the only argument. + If no result errors, this callback is not called. + total_object_refs: Number of ObjectRefs that this thread + expects to be ready. May be more than len(object_refs) since + more ObjectRefs can be submitted after the thread starts. + If None, defaults to len(object_refs). If float("inf"), thread runs + until END_SENTINEL (submitted through self.add_object_ref()) + has been received and all objects received before that have + been processed. + """ + + END_SENTINEL = None + + def __init__( + self, + object_refs: list, + single_result: bool = False, + callback: callable = None, + error_callback: callable = None, + total_object_refs: Optional[int] = None, + ): + threading.Thread.__init__(self, daemon=True) + self._got_error = False + self._object_refs = [] + self._num_ready = 0 + self._results = [] + self._ready_index_queue = queue.Queue() + self._single_result = single_result + self._callback = callback + self._error_callback = error_callback + self._total_object_refs = total_object_refs or len(object_refs) + self._indices = {} + # Thread-safe queue used to add ObjectRefs to fetch after creating + # this thread (used to lazily submit for imap and imap_unordered). + self._new_object_refs = queue.Queue() + for object_ref in object_refs: + self._add_object_ref(object_ref) + + def _add_object_ref(self, object_ref): + self._indices[object_ref] = len(self._object_refs) + self._object_refs.append(object_ref) + self._results.append(None) + + def add_object_ref(self, object_ref): + self._new_object_refs.put(object_ref) + + def run(self): + unready = copy.copy(self._object_refs) + aggregated_batch_results = [] + + # Run for a specific number of objects if self._total_object_refs is finite. + # Otherwise, process all objects received prior to the stop signal, given by + # self.add_object(END_SENTINEL). + while self._num_ready < self._total_object_refs: + # Get as many new IDs from the queue as possible without blocking, + # unless we have no IDs to wait on, in which case we block. + ready_id = None + while ready_id is None: + try: + block = len(unready) == 0 + new_object_ref = self._new_object_refs.get(block=block) + if new_object_ref is self.END_SENTINEL: + # Receiving the END_SENTINEL object is the signal to stop. + # Store the total number of objects. + self._total_object_refs = len(self._object_refs) + else: + self._add_object_ref(new_object_ref) + unready.append(new_object_ref) + except queue.Empty: + # queue.Empty means no result was retrieved if block=False. + pass + + # Check if any of the available IDs are done. The timeout is required + # here to periodically check for new IDs from self._new_object_refs. + # NOTE(edoakes): the choice of a 100ms timeout here is arbitrary. Too + # low of a timeout would cause higher overhead from busy spinning and + # too high would cause higher tail latency to fetch the first result in + # some cases. + ready, unready = ray.wait(unready, num_returns=1, timeout=0.1) + if len(ready) > 0: + ready_id = ready[0] + + try: + batch = ray.get(ready_id) + except ray.exceptions.RayError as e: + batch = [e] + + # The exception callback is called only once on the first result + # that errors. If no result errors, it is never called. + if not self._got_error: + for result in batch: + if isinstance(result, Exception): + self._got_error = True + if self._error_callback is not None: + self._error_callback(result) + break + else: + aggregated_batch_results.append(result) + + self._num_ready += 1 + self._results[self._indices[ready_id]] = batch + self._ready_index_queue.put(self._indices[ready_id]) + + # The regular callback is called only once on the entire List of + # results as long as none of the results were errors. If any results + # were errors, the regular callback is never called; instead, the + # exception callback is called on the first erroring result. + # + # This callback is called outside the while loop to ensure that it's + # called on the entire list of results– not just a single batch. + if not self._got_error and self._callback is not None: + if not self._single_result: + self._callback(aggregated_batch_results) + else: + # On a thread handling a function with a single result + # (e.g. apply_async), we call the callback on just that result + # instead of on a list encaspulating that result + self._callback(aggregated_batch_results[0]) + + def got_error(self): + # Should only be called after the thread finishes. + return self._got_error + + def result(self, index): + # Should only be called on results that are ready. + return self._results[index] + + def results(self): + # Should only be called after the thread finishes. + return self._results + + def next_ready_index(self, timeout=None): + try: + return self._ready_index_queue.get(timeout=timeout) + except queue.Empty: + # queue.Queue signals a timeout by raising queue.Empty. + raise TimeoutError + + +class AsyncResult: + """An asynchronous interface to task results. + + This should not be constructed directly. + """ + + def __init__( + self, chunk_object_refs, callback=None, error_callback=None, single_result=False + ): + self._single_result = single_result + self._result_thread = ResultThread( + chunk_object_refs, single_result, callback, error_callback + ) + self._result_thread.start() + + def wait(self, timeout=None): + """ + Returns once the result is ready or the timeout expires (does not + raise TimeoutError). + + Args: + timeout: timeout in milliseconds. + """ + + self._result_thread.join(timeout) + + def get(self, timeout=None): + self.wait(timeout) + if self._result_thread.is_alive(): + raise TimeoutError + + results = [] + for batch in self._result_thread.results(): + for result in batch: + if isinstance(result, PoolTaskError): + raise result.underlying + elif isinstance(result, Exception): + raise result + results.extend(batch) + + if self._single_result: + return results[0] + + return results + + def ready(self): + """ + Returns true if the result is ready, else false if the tasks are still + running. + """ + + return not self._result_thread.is_alive() + + def successful(self): + """ + Returns true if none of the submitted tasks errored, else false. Should + only be called once the result is ready (can be checked using `ready`). + """ + + if not self.ready(): + raise ValueError(f"{self!r} not ready") + return not self._result_thread.got_error() + + +class IMapIterator: + """Base class for OrderedIMapIterator and UnorderedIMapIterator.""" + + def __init__(self, pool, func, iterable, chunksize=None): + self._pool = pool + self._func = func + self._next_chunk_index = 0 + self._finished_iterating = False + # List of bools indicating if the given chunk is ready or not for all + # submitted chunks. Ordering mirrors that in the in the ResultThread. + self._submitted_chunks = [] + self._ready_objects = collections.deque() + self._iterator = iter(iterable) + if isinstance(iterable, collections.abc.Iterator): + # Got iterator (which has no len() function). + # Make default chunksize 1 instead of using _calculate_chunksize(). + # Indicate unknown queue length, requiring explicit stopping. + self._chunksize = chunksize or 1 + result_list_size = float("inf") + else: + self._chunksize = chunksize or pool._calculate_chunksize(iterable) + result_list_size = div_round_up(len(iterable), chunksize) + + self._result_thread = ResultThread([], total_object_refs=result_list_size) + self._result_thread.start() + + for _ in range(len(self._pool._actor_pool)): + self._submit_next_chunk() + + def _submit_next_chunk(self): + # The full iterable has already been submitted, so no-op. + if self._finished_iterating: + return + + actor_index = len(self._submitted_chunks) % len(self._pool._actor_pool) + chunk_iterator = itertools.islice(self._iterator, self._chunksize) + + # Check whether we have run out of samples. + # This consumes the original iterator, so we convert to a list and back + chunk_list = list(chunk_iterator) + if len(chunk_list) < self._chunksize: + # Reached end of self._iterator + self._finished_iterating = True + if len(chunk_list) == 0: + # Nothing to do, return. + return + chunk_iterator = iter(chunk_list) + + new_chunk_id = self._pool._submit_chunk( + self._func, chunk_iterator, self._chunksize, actor_index + ) + self._submitted_chunks.append(False) + # Wait for the result + self._result_thread.add_object_ref(new_chunk_id) + # If we submitted the final chunk, notify the result thread + if self._finished_iterating: + self._result_thread.add_object_ref(ResultThread.END_SENTINEL) + + def __iter__(self): + return self + + def __next__(self): + return self.next() + + def next(self): + # Should be implemented by subclasses. + raise NotImplementedError + + +class OrderedIMapIterator(IMapIterator): + """Iterator to the results of tasks submitted using `imap`. + + The results are returned in the same order that they were submitted, even + if they don't finish in that order. Only one batch of tasks per actor + process is submitted at a time - the rest are submitted as results come in. + + Should not be constructed directly. + """ + + def next(self, timeout=None): + if len(self._ready_objects) == 0: + if self._finished_iterating and ( + self._next_chunk_index == len(self._submitted_chunks) + ): + # Finish when all chunks have been dispatched and processed + # Notify the calling process that the work is done. + raise StopIteration + + # This loop will break when the next index in order is ready or + # self._result_thread.next_ready_index() raises a timeout. + index = -1 + while index != self._next_chunk_index: + start = time.time() + index = self._result_thread.next_ready_index(timeout=timeout) + self._submit_next_chunk() + self._submitted_chunks[index] = True + if timeout is not None: + timeout = max(0, timeout - (time.time() - start)) + + while ( + self._next_chunk_index < len(self._submitted_chunks) + and self._submitted_chunks[self._next_chunk_index] + ): + for result in self._result_thread.result(self._next_chunk_index): + self._ready_objects.append(result) + self._next_chunk_index += 1 + + return self._ready_objects.popleft() + + +class UnorderedIMapIterator(IMapIterator): + """Iterator to the results of tasks submitted using `imap`. + + The results are returned in the order that they finish. Only one batch of + tasks per actor process is submitted at a time - the rest are submitted as + results come in. + + Should not be constructed directly. + """ + + def next(self, timeout=None): + if len(self._ready_objects) == 0: + if self._finished_iterating and ( + self._next_chunk_index == len(self._submitted_chunks) + ): + # Finish when all chunks have been dispatched and processed + # Notify the calling process that the work is done. + raise StopIteration + + index = self._result_thread.next_ready_index(timeout=timeout) + self._submit_next_chunk() + + for result in self._result_thread.result(index): + self._ready_objects.append(result) + self._next_chunk_index += 1 + + return self._ready_objects.popleft() + + +@ray.remote(num_cpus=0) +class PoolActor: + """Actor used to process tasks submitted to a Pool.""" + + def __init__(self, initializer=None, initargs=None): + if initializer: + initargs = initargs or () + initializer(*initargs) + + def ping(self): + # Used to wait for this actor to be initialized. + pass + + def run_batch(self, func, batch): + results = [] + for args, kwargs in batch: + args = args or () + kwargs = kwargs or {} + try: + results.append(func(*args, **kwargs)) + except Exception as e: + results.append(PoolTaskError(e)) + return results + + +# https://docs.python.org/3/library/multiprocessing.html#module-multiprocessing.pool +class Pool: + """A pool of actor processes that is used to process tasks in parallel. + + Args: + processes: number of actor processes to start in the pool. Defaults to + the number of cores in the Ray cluster if one is already running, + otherwise the number of cores on this machine. + initializer: function to be run in each actor when it starts up. + initargs: iterable of arguments to the initializer function. + maxtasksperchild: maximum number of tasks to run in each actor process. + After a process has executed this many tasks, it will be killed and + replaced with a new one. + ray_address: address of the Ray cluster to run on. If None, a new local + Ray cluster will be started on this machine. Otherwise, this will + be passed to `ray.init()` to connect to a running cluster. This may + also be specified using the `RAY_ADDRESS` environment variable. + ray_remote_args: arguments used to configure the Ray Actors making up + the pool. See :func:`ray.remote` for details. + """ + + def __init__( + self, + processes: Optional[int] = None, + initializer: Optional[Callable] = None, + initargs: Optional[Iterable] = None, + maxtasksperchild: Optional[int] = None, + context: Any = None, + ray_address: Optional[str] = None, + ray_remote_args: Optional[Dict[str, Any]] = None, + ): + usage_lib.record_library_usage("util.multiprocessing.Pool") + + self._closed = False + self._initializer = initializer + self._initargs = initargs + self._maxtasksperchild = maxtasksperchild or -1 + self._actor_deletion_ids = [] + self._registry: List[Tuple[Any, ray.ObjectRef]] = [] + self._registry_hashable: Dict[Hashable, ray.ObjectRef] = {} + self._current_index = 0 + self._ray_remote_args = ray_remote_args or {} + self._pool_actor = None + + if context and log_once("context_argument_warning"): + logger.warning( + "The 'context' argument is not supported using " + "ray. Please refer to the documentation for how " + "to control ray initialization." + ) + + processes = self._init_ray(processes, ray_address) + self._start_actor_pool(processes) + + def _init_ray(self, processes=None, ray_address=None): + # Initialize ray. If ray is already initialized, we do nothing. + # Else, the priority is: + # ray_address argument > RAY_ADDRESS > start new local cluster. + if not ray.is_initialized(): + # Cluster mode. + if ray_address is None and ( + RAY_ADDRESS_ENV in os.environ + or ray._private.utils.read_ray_address() is not None + ): + init_kwargs = {} + if os.environ.get(RAY_ADDRESS_ENV) == "local": + init_kwargs["num_cpus"] = processes + ray.init(**init_kwargs) + elif ray_address is not None: + init_kwargs = {} + if ray_address == "local": + init_kwargs["num_cpus"] = processes + ray.init(address=ray_address, **init_kwargs) + # Local mode. + else: + ray.init(num_cpus=processes) + + ray_cpus = int(ray._private.state.cluster_resources()["CPU"]) + if processes is None: + processes = ray_cpus + if processes <= 0: + raise ValueError("Processes in the pool must be >0.") + if ray_cpus < processes: + raise ValueError( + "Tried to start a pool with {} processes on an " + "existing ray cluster, but there are only {} " + "CPUs in the ray cluster.".format(processes, ray_cpus) + ) + + return processes + + def _start_actor_pool(self, processes): + self._pool_actor = None + self._actor_pool = [self._new_actor_entry() for _ in range(processes)] + ray.get([actor.ping.remote() for actor, _ in self._actor_pool]) + + def _wait_for_stopping_actors(self, timeout=None): + if len(self._actor_deletion_ids) == 0: + return + if timeout is not None: + timeout = float(timeout) + + _, deleting = ray.wait( + self._actor_deletion_ids, + num_returns=len(self._actor_deletion_ids), + timeout=timeout, + ) + self._actor_deletion_ids = deleting + + def _stop_actor(self, actor): + # Check and clean up any outstanding IDs corresponding to deletions. + self._wait_for_stopping_actors(timeout=0.0) + # The deletion task will block until the actor has finished executing + # all pending tasks. + self._actor_deletion_ids.append(actor.__ray_terminate__.remote()) + + def _new_actor_entry(self): + # NOTE(edoakes): The initializer function can't currently be used to + # modify the global namespace (e.g., import packages or set globals) + # due to a limitation in cloudpickle. + # Cache the PoolActor with options + if not self._pool_actor: + self._pool_actor = PoolActor.options(**self._ray_remote_args) + return (self._pool_actor.remote(self._initializer, self._initargs), 0) + + def _next_actor_index(self): + if self._current_index == len(self._actor_pool) - 1: + self._current_index = 0 + else: + self._current_index += 1 + return self._current_index + + # Batch should be a list of tuples: (args, kwargs). + def _run_batch(self, actor_index, func, batch): + actor, count = self._actor_pool[actor_index] + object_ref = actor.run_batch.remote(func, batch) + count += 1 + assert self._maxtasksperchild == -1 or count <= self._maxtasksperchild + if count == self._maxtasksperchild: + self._stop_actor(actor) + actor, count = self._new_actor_entry() + self._actor_pool[actor_index] = (actor, count) + return object_ref + + def apply( + self, + func: Callable, + args: Optional[Tuple] = None, + kwargs: Optional[Dict] = None, + ): + """Run the given function on a random actor process and return the + result synchronously. + + Args: + func: function to run. + args: optional arguments to the function. + kwargs: optional keyword arguments to the function. + + Returns: + The result. + """ + + return self.apply_async(func, args, kwargs).get() + + def apply_async( + self, + func: Callable, + args: Optional[Tuple] = None, + kwargs: Optional[Dict] = None, + callback: Callable[[Any], None] = None, + error_callback: Callable[[Exception], None] = None, + ): + """Run the given function on a random actor process and return an + asynchronous interface to the result. + + Args: + func: function to run. + args: optional arguments to the function. + kwargs: optional keyword arguments to the function. + callback: callback to be executed on the result once it is finished + only if it succeeds. + error_callback: callback to be executed the result once it is + finished only if the task errors. The exception raised by the + task will be passed as the only argument to the callback. + + Returns: + AsyncResult containing the result. + """ + + self._check_running() + func = self._convert_to_ray_batched_calls_if_needed(func) + object_ref = self._run_batch(self._next_actor_index(), func, [(args, kwargs)]) + return AsyncResult([object_ref], callback, error_callback, single_result=True) + + def _convert_to_ray_batched_calls_if_needed(self, func: Callable) -> Callable: + """Convert joblib's BatchedCalls to RayBatchedCalls for ObjectRef caching. + + This converts joblib's BatchedCalls callable, which is a collection of + functions with their args and kwargs to be ran sequentially in an + Actor, to a RayBatchedCalls callable, which provides identical + functionality in addition to a method which ensures that common + args and kwargs are put into the object store just once, saving time + and memory. That method is then ran. + + If func is not a BatchedCalls instance, it is returned without changes. + + The ObjectRefs are cached inside two registries (_registry and + _registry_hashable), which are common for the entire Pool and are + cleaned on close.""" + if RayBatchedCalls is None: + return func + orginal_func = func + # SafeFunction is a Python 2 leftover and can be + # safely removed. + if isinstance(func, SafeFunction): + func = func.func + if isinstance(func, BatchedCalls): + func = RayBatchedCalls( + func.items, + (func._backend, func._n_jobs), + func._reducer_callback, + func._pickle_cache, + ) + # go through all the items and replace args and kwargs with + # ObjectRefs, caching them in registries + func.put_items_in_object_store(self._registry, self._registry_hashable) + else: + func = orginal_func + return func + + def _calculate_chunksize(self, iterable): + chunksize, extra = divmod(len(iterable), len(self._actor_pool) * 4) + if extra: + chunksize += 1 + return chunksize + + def _submit_chunk(self, func, iterator, chunksize, actor_index, unpack_args=False): + chunk = [] + while len(chunk) < chunksize: + try: + args = next(iterator) + if not unpack_args: + args = (args,) + chunk.append((args, {})) + except StopIteration: + break + + # Nothing to submit. The caller should prevent this. + assert len(chunk) > 0 + + return self._run_batch(actor_index, func, chunk) + + def _chunk_and_run(self, func, iterable, chunksize=None, unpack_args=False): + if not hasattr(iterable, "__len__"): + iterable = list(iterable) + + if chunksize is None: + chunksize = self._calculate_chunksize(iterable) + + iterator = iter(iterable) + chunk_object_refs = [] + while len(chunk_object_refs) * chunksize < len(iterable): + actor_index = len(chunk_object_refs) % len(self._actor_pool) + chunk_object_refs.append( + self._submit_chunk( + func, iterator, chunksize, actor_index, unpack_args=unpack_args + ) + ) + + return chunk_object_refs + + def _map_async( + self, + func, + iterable, + chunksize=None, + unpack_args=False, + callback=None, + error_callback=None, + ): + self._check_running() + object_refs = self._chunk_and_run( + func, iterable, chunksize=chunksize, unpack_args=unpack_args + ) + return AsyncResult(object_refs, callback, error_callback) + + def map(self, func: Callable, iterable: Iterable, chunksize: Optional[int] = None): + """Run the given function on each element in the iterable round-robin + on the actor processes and return the results synchronously. + + Args: + func: function to run. + iterable: iterable of objects to be passed as the sole argument to + func. + chunksize: number of tasks to submit as a batch to each actor + process. If unspecified, a suitable chunksize will be chosen. + + Returns: + A list of results. + """ + + return self._map_async( + func, iterable, chunksize=chunksize, unpack_args=False + ).get() + + def map_async( + self, + func: Callable, + iterable: Iterable, + chunksize: Optional[int] = None, + callback: Callable[[List], None] = None, + error_callback: Callable[[Exception], None] = None, + ): + """Run the given function on each element in the iterable round-robin + on the actor processes and return an asynchronous interface to the + results. + + Args: + func: function to run. + iterable: iterable of objects to be passed as the only argument to + func. + chunksize: number of tasks to submit as a batch to each actor + process. If unspecified, a suitable chunksize will be chosen. + callback: Will only be called if none of the results were errors, + and will only be called once after all results are finished. + A Python List of all the finished results will be passed as the + only argument to the callback. + error_callback: callback executed on the first errored result. + The Exception raised by the task will be passed as the only + argument to the callback. + + Returns: + AsyncResult + """ + return self._map_async( + func, + iterable, + chunksize=chunksize, + unpack_args=False, + callback=callback, + error_callback=error_callback, + ) + + def starmap(self, func, iterable, chunksize=None): + """Same as `map`, but unpacks each element of the iterable as the + arguments to func like: [func(*args) for args in iterable]. + """ + + return self._map_async( + func, iterable, chunksize=chunksize, unpack_args=True + ).get() + + def starmap_async( + self, + func: Callable, + iterable: Iterable, + callback: Callable[[List], None] = None, + error_callback: Callable[[Exception], None] = None, + ): + """Same as `map_async`, but unpacks each element of the iterable as the + arguments to func like: [func(*args) for args in iterable]. + """ + + return self._map_async( + func, + iterable, + unpack_args=True, + callback=callback, + error_callback=error_callback, + ) + + def imap(self, func: Callable, iterable: Iterable, chunksize: Optional[int] = 1): + """Same as `map`, but only submits one batch of tasks to each actor + process at a time. + + This can be useful if the iterable of arguments is very large or each + task's arguments consumes a large amount of resources. + + The results are returned in the order corresponding to their arguments + in the iterable. + + Returns: + OrderedIMapIterator + """ + + self._check_running() + return OrderedIMapIterator(self, func, iterable, chunksize=chunksize) + + def imap_unordered( + self, func: Callable, iterable: Iterable, chunksize: Optional[int] = 1 + ): + """Same as `map`, but only submits one batch of tasks to each actor + process at a time. + + This can be useful if the iterable of arguments is very large or each + task's arguments consumes a large amount of resources. + + The results are returned in the order that they finish. + + Returns: + UnorderedIMapIterator + """ + + self._check_running() + return UnorderedIMapIterator(self, func, iterable, chunksize=chunksize) + + def _check_running(self): + if self._closed: + raise ValueError("Pool not running") + + def __enter__(self): + self._check_running() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.terminate() + + def close(self): + """Close the pool. + + Prevents any more tasks from being submitted on the pool but allows + outstanding work to finish. + """ + + self._registry.clear() + self._registry_hashable.clear() + for actor, _ in self._actor_pool: + self._stop_actor(actor) + self._closed = True + gc.collect() + + def terminate(self): + """Close the pool. + + Prevents any more tasks from being submitted on the pool and stops + outstanding work. + """ + + if not self._closed: + self.close() + for actor, _ in self._actor_pool: + ray.kill(actor) + + def join(self): + """Wait for the actors in a closed pool to exit. + + If the pool was closed using `close`, this will return once all + outstanding work is completed. + + If the pool was closed using `terminate`, this will return quickly. + """ + + if not self._closed: + raise ValueError("Pool is still running") + self._wait_for_stopping_actors() diff --git a/lib/python3.12/site-packages/ray/util/placement_group.py b/lib/python3.12/site-packages/ray/util/placement_group.py new file mode 100644 index 0000000000000000000000000000000000000000..d2e29b81c536374a0519ce75e779d1df2fed25ee --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/placement_group.py @@ -0,0 +1,583 @@ +import warnings +from typing import Dict, List, Optional, Union + +import ray +from ray._common.utils import PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME, hex_to_binary +from ray._private.auto_init_hook import auto_init_ray +from ray._private.client_mode_hook import client_mode_should_convert, client_mode_wrap +from ray._private.label_utils import validate_label_selector +from ray._private.utils import get_ray_doc_version +from ray._raylet import PlacementGroupID +from ray.util.annotations import DeveloperAPI, PublicAPI +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +bundle_reservation_check = None + +VALID_PLACEMENT_GROUP_STRATEGIES = { + "PACK", + "SPREAD", + "STRICT_PACK", + "STRICT_SPREAD", +} + + +# We need to import this method to use for ready API. +# But ray.remote is only available in runtime, and +# if we define this method inside ready method, this function is +# exported whenever ready is called, which can impact performance, +# https://github.com/ray-project/ray/issues/6240. +def _export_bundle_reservation_check_method_if_needed(): + global bundle_reservation_check + if bundle_reservation_check: + return + + @ray.remote(num_cpus=0) + def bundle_reservation_check_func(placement_group): + return placement_group + + bundle_reservation_check = bundle_reservation_check_func + + +@PublicAPI +class PlacementGroup: + """A handle to a placement group.""" + + @staticmethod + def empty() -> "PlacementGroup": + return PlacementGroup(PlacementGroupID.nil()) + + def __init__( + self, + id: "ray._raylet.PlacementGroupID", + bundle_cache: Optional[List[Dict]] = None, + ): + self.id = id + self.bundle_cache = bundle_cache + + @property + def is_empty(self): + return self.id.is_nil() + + def ready(self) -> "ray._raylet.ObjectRef": + """Returns an ObjectRef to check ready status. + + This API runs a small dummy task to wait for placement group creation. + It is compatible to ray.get and ray.wait. + + Example: + .. testcode:: + + import ray + + pg = ray.util.placement_group([{"CPU": 1}]) + ray.get(pg.ready()) + + pg = ray.util.placement_group([{"CPU": 1}]) + ray.wait([pg.ready()]) + + """ + self._fill_bundle_cache_if_needed() + + _export_bundle_reservation_check_method_if_needed() + + assert len(self.bundle_cache) != 0, ( + "ready() cannot be called on placement group object with a " + "bundle length == 0, current bundle length: " + f"{len(self.bundle_cache)}" + ) + + return bundle_reservation_check.options( + scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=self), + ).remote(self) + + def wait(self, timeout_seconds: Union[float, int] = 30) -> bool: + """Wait for the placement group to be ready within the specified time. + Args: + timeout_seconds(float|int): Timeout in seconds. + Return: + True if the placement group is created. False otherwise. + """ + return _call_placement_group_ready(self.id, timeout_seconds) + + @property + def bundle_specs(self) -> List[Dict]: + """List[Dict]: Return bundles belonging to this placement group.""" + self._fill_bundle_cache_if_needed() + return self.bundle_cache + + @property + def bundle_count(self) -> int: + self._fill_bundle_cache_if_needed() + return len(self.bundle_cache) + + def _fill_bundle_cache_if_needed(self) -> None: + if not self.bundle_cache: + self.bundle_cache = _get_bundle_cache(self.id) + + def __eq__(self, other): + if not isinstance(other, PlacementGroup): + return False + return self.id == other.id + + def __hash__(self): + return hash(self.id) + + +@client_mode_wrap +def _call_placement_group_ready(pg_id: PlacementGroupID, timeout_seconds: int) -> bool: + worker = ray._private.worker.global_worker + worker.check_connected() + + return worker.core_worker.wait_placement_group_ready(pg_id, timeout_seconds) + + +@client_mode_wrap +def _get_bundle_cache(pg_id: PlacementGroupID) -> List[Dict]: + worker = ray._private.worker.global_worker + worker.check_connected() + + return list( + ray._private.state.state.placement_group_table(pg_id)["bundles"].values() + ) + + +@PublicAPI +@client_mode_wrap +def placement_group( + bundles: List[Dict[str, float]], + strategy: str = "PACK", + name: str = "", + lifetime: Optional[str] = None, + _soft_target_node_id: Optional[str] = None, + bundle_label_selector: List[Dict[str, str]] = None, +) -> PlacementGroup: + """Asynchronously creates a PlacementGroup. + + Args: + bundles: A list of bundles which + represent the resources requirements. + strategy: The strategy to create the placement group. + + - "PACK": Packs Bundles into as few nodes as possible. + - "SPREAD": Places Bundles across distinct nodes as even as possible. + - "STRICT_PACK": Packs Bundles into one node. The group is + not allowed to span multiple nodes. + - "STRICT_SPREAD": Packs Bundles across distinct nodes. + + name: The name of the placement group. + lifetime: Either `None`, which defaults to the placement group + will fate share with its creator and will be deleted once its + creator is dead, or "detached", which means the placement group + will live as a global object independent of the creator. + _soft_target_node_id: (Private, Experimental) Soft hint where bundles of + this placement group should be placed. + The target node is specified by it's hex ID. + If the target node has no available resources or died, + bundles can be placed elsewhere. + This currently only works with STRICT_PACK pg. + bundle_label_selector: A list of label selectors to apply to a + placement group on a per-bundle level. + + Raises: + ValueError: if bundle type is not a list. + ValueError: if empty bundle or empty resource bundles are given. + ValueError: if the wrong lifetime arguments are given. + + Return: + PlacementGroup: Placement group object. + """ + worker = ray._private.worker.global_worker + worker.check_connected() + + validate_placement_group( + bundles=bundles, + strategy=strategy, + lifetime=lifetime, + _soft_target_node_id=_soft_target_node_id, + bundle_label_selector=bundle_label_selector, + ) + + if bundle_label_selector is None: + bundle_label_selector = [] + + if lifetime == "detached": + detached = True + else: + detached = False + + placement_group_id = worker.core_worker.create_placement_group( + name, + bundles, + strategy, + detached, + _soft_target_node_id, + bundle_label_selector, + ) + + return PlacementGroup(placement_group_id) + + +@PublicAPI +@client_mode_wrap +def remove_placement_group(placement_group: PlacementGroup) -> None: + """Asynchronously remove placement group. + + Args: + placement_group: The placement group to delete. + """ + assert placement_group is not None + worker = ray._private.worker.global_worker + worker.check_connected() + + worker.core_worker.remove_placement_group(placement_group.id) + + +@PublicAPI +@client_mode_wrap +def get_placement_group(placement_group_name: str) -> PlacementGroup: + """Get a placement group object with a global name. + + Returns: + None if can't find a placement group with the given name. + The placement group object otherwise. + """ + if not placement_group_name: + raise ValueError("Please supply a non-empty value to get_placement_group") + worker = ray._private.worker.global_worker + worker.check_connected() + placement_group_info = ray._private.state.state.get_placement_group_by_name( + placement_group_name, worker.namespace + ) + if placement_group_info is None: + raise ValueError( + f"Failed to look up placement group with name: {placement_group_name}" + ) + else: + return PlacementGroup( + PlacementGroupID(hex_to_binary(placement_group_info["placement_group_id"])) + ) + + +@DeveloperAPI +@client_mode_wrap +def placement_group_table(placement_group: PlacementGroup = None) -> dict: + """Get the state of the placement group from GCS. + + Args: + placement_group: placement group to see + states. + """ + worker = ray._private.worker.global_worker + worker.check_connected() + placement_group_id = placement_group.id if (placement_group is not None) else None + return ray._private.state.state.placement_group_table(placement_group_id) + + +@PublicAPI +def get_current_placement_group() -> Optional[PlacementGroup]: + """Get the current placement group which a task or actor is using. + + It returns None if there's no current placement group for the worker. + For example, if you call this method in your driver, it returns None + (because drivers never belong to any placement group). + + Examples: + .. testcode:: + + import ray + from ray.util.placement_group import get_current_placement_group + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + + @ray.remote + def f(): + # This returns the placement group the task f belongs to. + # It means this pg is identical to the pg created below. + return get_current_placement_group() + + pg = ray.util.placement_group([{"CPU": 2}]) + assert ray.get(f.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg)).remote()) == pg + + # Driver doesn't belong to any placement group, + # so it returns None. + assert get_current_placement_group() is None + + Return: + PlacementGroup: Placement group object. + None if the current task or actor wasn't + created with any placement group. + """ + auto_init_ray() + if client_mode_should_convert(): + # Client mode is only a driver. + return None + worker = ray._private.worker.global_worker + worker.check_connected() + pg_id = worker.placement_group_id + if pg_id.is_nil(): + return None + return PlacementGroup(pg_id) + + +def check_placement_group_index( + placement_group: PlacementGroup, bundle_index: int +) -> None: + assert placement_group is not None + if placement_group.id.is_nil(): + if bundle_index != -1: + raise ValueError( + "If placement group is not set, " + "the value of bundle index must be -1." + ) + elif bundle_index >= placement_group.bundle_count or bundle_index < -1: + raise ValueError( + f"placement group bundle index {bundle_index} " + f"is invalid. Valid placement group indexes: " + f"0-{placement_group.bundle_count}" + ) + + +def validate_placement_group( + bundles: List[Dict[str, float]], + strategy: str = "PACK", + lifetime: Optional[str] = None, + _soft_target_node_id: Optional[str] = None, + bundle_label_selector: List[Dict[str, str]] = None, +) -> bool: + """Validates inputs for placement_group. + + Raises ValueError if inputs are invalid. + """ + if _soft_target_node_id and strategy != "STRICT_PACK": + raise ValueError( + "_soft_target_node_id currently only works " + f"with STRICT_PACK but got {strategy}" + ) + + if _soft_target_node_id and ray.NodeID.from_hex(_soft_target_node_id).is_nil(): + raise ValueError( + f"Invalid hex ID of _soft_target_node_id, got {_soft_target_node_id}" + ) + + _validate_bundles(bundles) + + if bundle_label_selector is not None: + if len(bundles) != len(bundle_label_selector): + raise ValueError( + f"Invalid bundle label selector {bundle_label_selector}. " + f"The length of `bundle_label_selector` should equal the length of `bundles`." + ) + _validate_bundle_label_selector(bundle_label_selector) + + if strategy not in VALID_PLACEMENT_GROUP_STRATEGIES: + raise ValueError( + f"Invalid placement group strategy {strategy}. " + f"Supported strategies are: {VALID_PLACEMENT_GROUP_STRATEGIES}." + ) + + if lifetime not in [None, "detached"]: + raise ValueError( + "Placement group `lifetime` argument must be either `None` or " + f"'detached'. Got {lifetime}." + ) + + +def _validate_bundles(bundles: List[Dict[str, float]]): + """Validates each bundle and raises a ValueError if any bundle is invalid.""" + + if not isinstance(bundles, list): + raise ValueError( + "Placement group bundles must be a list, " f"got {type(bundles)}." + ) + + if len(bundles) == 0: + raise ValueError( + "Bundles must be a non-empty list of resource " + 'dictionaries. For example: `[{"CPU": 1.0}, {"GPU": 1.0}]`. ' + "Got empty list instead." + ) + + for bundle in bundles: + if ( + not isinstance(bundle, dict) + or not all(isinstance(k, str) for k in bundle.keys()) + or not all(isinstance(v, (int, float)) for v in bundle.values()) + ): + raise ValueError( + "Bundles must be a non-empty list of " + "resource dictionaries. For example: " + '`[{"CPU": 1.0}, {"GPU": 1.0}]`.' + ) + + if len(bundle) == 0 or all( + resource_value == 0 for resource_value in bundle.values() + ): + raise ValueError( + "Bundles cannot be an empty dictionary or " + f"resources with only 0 values. Bundles: {bundles}" + ) + + if "object_store_memory" in bundle.keys(): + warnings.warn( + "Setting 'object_store_memory' for" + " bundles is deprecated since it doesn't actually" + " reserve the required object store memory." + f" Use object spilling that's enabled by default (https://docs.ray.io/en/{get_ray_doc_version()}/ray-core/objects/object-spilling.html) " # noqa: E501 + "instead to bypass the object store memory size limitation.", + DeprecationWarning, + stacklevel=1, + ) + + +def _validate_bundle_label_selector(bundle_label_selector: List[Dict[str, str]]): + """Validates each label selector and raises a ValueError if any label selector is invalid.""" + + if not isinstance(bundle_label_selector, list): + raise ValueError( + "Placement group bundle_label_selector must be a list, " + f"got {type(bundle_label_selector)}." + ) + + if len(bundle_label_selector) == 0: + # No label selectors provided, no-op. + return + + for label_selector in bundle_label_selector: + if ( + not isinstance(label_selector, dict) + or not all(isinstance(k, str) for k in label_selector.keys()) + or not all(isinstance(v, str) for v in label_selector.values()) + ): + raise ValueError( + "Bundle label selector must be a list of string dictionary" + " label selectors. For example: " + '`[{ray.io/market_type": "spot"}, {"ray.io/accelerator-type": "A100"}]`.' + ) + # Call helper function to validate label selector key-value syntax. + error_message = validate_label_selector(label_selector) + if error_message: + raise ValueError( + f"Invalid label selector provided in bundle_label_selector list." + f" Detailed error: '{error_message}'" + ) + + +def _valid_resource_shape(resources, bundle_specs): + """ + If the resource shape cannot fit into every + bundle spec, return False + """ + for bundle in bundle_specs: + fit_in_bundle = True + for resource, requested_val in resources.items(): + # Skip "bundle" resource as it is automatically added + # to all nodes with bundles by the placement group. + if resource == PLACEMENT_GROUP_BUNDLE_RESOURCE_NAME: + continue + if bundle.get(resource, 0) < requested_val: + fit_in_bundle = False + break + if fit_in_bundle: + # If resource request fits in any bundle, it is valid. + return True + return False + + +def _validate_resource_shape( + placement_group, resources, placement_resources, task_or_actor_repr +): + bundles = placement_group.bundle_specs + resources_valid = _valid_resource_shape(resources, bundles) + placement_resources_valid = _valid_resource_shape(placement_resources, bundles) + + if not resources_valid: + raise ValueError( + f"Cannot schedule {task_or_actor_repr} with " + "the placement group because the resource request " + f"{resources} cannot fit into any bundles for " + f"the placement group, {bundles}." + ) + if not placement_resources_valid: + # Happens for the default actor case. + # placement_resources is not an exposed concept to users, + # so we should write more specialized error messages. + raise ValueError( + f"Cannot schedule {task_or_actor_repr} with " + "the placement group because the actor requires " + f"{placement_resources.get('CPU', 0)} CPU for " + "creation, but it cannot " + f"fit into any bundles for the placement group, " + f"{bundles}. Consider " + "creating a placement group with CPU resources." + ) + + +def _configure_placement_group_based_on_context( + placement_group_capture_child_tasks: bool, + bundle_index: int, + resources: Dict, + placement_resources: Dict, + task_or_actor_repr: str, + placement_group: Union[PlacementGroup, str, None] = "default", +) -> PlacementGroup: + """Configure the placement group based on the given context. + + Based on the given context, this API returns the placement group instance + for task/actor scheduling. + + Params: + placement_group_capture_child_tasks: Whether or not the + placement group needs to be captured from the global + context. + bundle_index: The bundle index for tasks/actor scheduling. + resources: The scheduling resources. + placement_resources: The scheduling placement resources for + actors. + task_or_actor_repr: The repr of task or actor + function/class descriptor. + placement_group: The placement group instance. + - "default": Default placement group argument. Currently, + the default behavior is to capture the parent task' + placement group if placement_group_capture_child_tasks + is set. + - None: means placement group is explicitly not configured. + - Placement group instance: In this case, do nothing. + + Returns: + Placement group instance based on the given context. + + Raises: + ValueError: If the bundle index is invalid for the placement group + or the requested resources shape doesn't fit to any + bundles. + """ + # Validate inputs. + assert placement_group_capture_child_tasks is not None + assert resources is not None + + # Validate and get the PlacementGroup instance. + # Placement group could be None, default, or placement group. + # Default behavior is "do not capture child tasks". + if placement_group != "default": + if not placement_group: + placement_group = PlacementGroup.empty() + elif placement_group == "default": + if placement_group_capture_child_tasks: + placement_group = get_current_placement_group() + else: + placement_group = PlacementGroup.empty() + + if not placement_group: + placement_group = PlacementGroup.empty() + assert isinstance(placement_group, PlacementGroup) + + # Validate the index. + check_placement_group_index(placement_group, bundle_index) + + # Validate the shape. + if not placement_group.is_empty: + _validate_resource_shape( + placement_group, resources, placement_resources, task_or_actor_repr + ) + return placement_group diff --git a/lib/python3.12/site-packages/ray/util/queue.py b/lib/python3.12/site-packages/ray/util/queue.py new file mode 100644 index 0000000000000000000000000000000000000000..b18075c801ac71800a2cb3dffa6ad187935c1a0a --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/queue.py @@ -0,0 +1,306 @@ +import asyncio +import queue +from collections.abc import Iterable +from typing import Any, Dict, List, Optional + +import ray +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="beta") +class Empty(queue.Empty): + pass + + +@PublicAPI(stability="beta") +class Full(queue.Full): + pass + + +@PublicAPI(stability="beta") +class Queue: + """A first-in, first-out queue implementation on Ray. + + The behavior and use cases are similar to those of the asyncio.Queue class. + + Features both sync and async put and get methods. Provides the option to + block until space is available when calling put on a full queue, + or to block until items are available when calling get on an empty queue. + + Optionally supports batched put and get operations to minimize + serialization overhead. + + Args: + maxsize (optional, int): maximum size of the queue. If zero, size is + unbounded. + actor_options (optional, Dict): Dictionary of options to pass into + the QueueActor during creation. These are directly passed into + QueueActor.options(...). This could be useful if you + need to pass in custom resource requirements, for example. + + Examples: + .. testcode:: + + from ray.util.queue import Queue + q = Queue() + items = list(range(10)) + for item in items: + q.put(item) + for item in items: + assert item == q.get() + # Create Queue with the underlying actor reserving 1 CPU. + q = Queue(actor_options={"num_cpus": 1}) + """ + + def __init__(self, maxsize: int = 0, actor_options: Optional[Dict] = None) -> None: + from ray._common.usage.usage_lib import record_library_usage + + record_library_usage("util.Queue") + + actor_options = actor_options or {} + self.maxsize = maxsize + self.actor = ( + ray.remote(_QueueActor).options(**actor_options).remote(self.maxsize) + ) + + def __len__(self) -> int: + return self.size() + + def size(self) -> int: + """The size of the queue.""" + return ray.get(self.actor.qsize.remote()) + + def qsize(self) -> int: + """The size of the queue.""" + return self.size() + + def empty(self) -> bool: + """Whether the queue is empty.""" + return ray.get(self.actor.empty.remote()) + + def full(self) -> bool: + """Whether the queue is full.""" + return ray.get(self.actor.full.remote()) + + def put( + self, item: Any, block: bool = True, timeout: Optional[float] = None + ) -> None: + """Adds an item to the queue. + + If block is True and the queue is full, blocks until the queue is no + longer full or until timeout. + + There is no guarantee of order if multiple producers put to the same + full queue. + + Raises: + Full: if the queue is full and blocking is False. + Full: if the queue is full, blocking is True, and it timed out. + ValueError: if timeout is negative. + """ + if not block: + try: + ray.get(self.actor.put_nowait.remote(item)) + except asyncio.QueueFull: + raise Full + else: + if timeout is not None and timeout < 0: + raise ValueError("'timeout' must be a non-negative number") + else: + ray.get(self.actor.put.remote(item, timeout)) + + async def put_async( + self, item: Any, block: bool = True, timeout: Optional[float] = None + ) -> None: + """Adds an item to the queue. + + If block is True and the queue is full, + blocks until the queue is no longer full or until timeout. + + There is no guarantee of order if multiple producers put to the same + full queue. + + Raises: + Full: if the queue is full and blocking is False. + Full: if the queue is full, blocking is True, and it timed out. + ValueError: if timeout is negative. + """ + if not block: + try: + await self.actor.put_nowait.remote(item) + except asyncio.QueueFull: + raise Full + else: + if timeout is not None and timeout < 0: + raise ValueError("'timeout' must be a non-negative number") + else: + await self.actor.put.remote(item, timeout) + + def get(self, block: bool = True, timeout: Optional[float] = None) -> Any: + """Gets an item from the queue. + + If block is True and the queue is empty, blocks until the queue is no + longer empty or until timeout. + + There is no guarantee of order if multiple consumers get from the + same empty queue. + + Returns: + The next item in the queue. + + Raises: + Empty: if the queue is empty and blocking is False. + Empty: if the queue is empty, blocking is True, and it timed out. + ValueError: if timeout is negative. + """ + if not block: + try: + return ray.get(self.actor.get_nowait.remote()) + except asyncio.QueueEmpty: + raise Empty + else: + if timeout is not None and timeout < 0: + raise ValueError("'timeout' must be a non-negative number") + else: + return ray.get(self.actor.get.remote(timeout)) + + async def get_async( + self, block: bool = True, timeout: Optional[float] = None + ) -> Any: + """Gets an item from the queue. + + There is no guarantee of order if multiple consumers get from the + same empty queue. + + Returns: + The next item in the queue. + Raises: + Empty: if the queue is empty and blocking is False. + Empty: if the queue is empty, blocking is True, and it timed out. + ValueError: if timeout is negative. + """ + if not block: + try: + return await self.actor.get_nowait.remote() + except asyncio.QueueEmpty: + raise Empty + else: + if timeout is not None and timeout < 0: + raise ValueError("'timeout' must be a non-negative number") + else: + return await self.actor.get.remote(timeout) + + def put_nowait(self, item: Any) -> None: + """Equivalent to put(item, block=False). + + Raises: + Full: if the queue is full. + """ + return self.put(item, block=False) + + def put_nowait_batch(self, items: Iterable) -> None: + """Takes in a list of items and puts them into the queue in order. + + Raises: + Full: if the items will not fit in the queue + """ + if not isinstance(items, Iterable): + raise TypeError("Argument 'items' must be an Iterable") + + ray.get(self.actor.put_nowait_batch.remote(items)) + + def get_nowait(self) -> Any: + """Equivalent to get(block=False). + + Raises: + Empty: if the queue is empty. + """ + return self.get(block=False) + + def get_nowait_batch(self, num_items: int) -> List[Any]: + """Gets items from the queue and returns them in a + list in order. + + Raises: + Empty: if the queue does not contain the desired number of items + """ + if not isinstance(num_items, int): + raise TypeError("Argument 'num_items' must be an int") + if num_items < 0: + raise ValueError("'num_items' must be nonnegative") + + return ray.get(self.actor.get_nowait_batch.remote(num_items)) + + def shutdown(self, force: bool = False, grace_period_s: int = 5) -> None: + """Terminates the underlying QueueActor. + + All of the resources reserved by the queue will be released. + + Args: + force: If True, forcefully kill the actor, causing an + immediate failure. If False, graceful + actor termination will be attempted first, before falling back + to a forceful kill. + grace_period_s: If force is False, how long in seconds to + wait for graceful termination before falling back to + forceful kill. + """ + if self.actor: + if force: + ray.kill(self.actor, no_restart=True) + else: + done_ref = self.actor.__ray_terminate__.remote() + done, not_done = ray.wait([done_ref], timeout=grace_period_s) + if not_done: + ray.kill(self.actor, no_restart=True) + self.actor = None + + +class _QueueActor: + def __init__(self, maxsize): + self.maxsize = maxsize + self.queue = asyncio.Queue(self.maxsize) + + def qsize(self): + return self.queue.qsize() + + def empty(self): + return self.queue.empty() + + def full(self): + return self.queue.full() + + async def put(self, item, timeout=None): + try: + await asyncio.wait_for(self.queue.put(item), timeout) + except asyncio.TimeoutError: + raise Full + + async def get(self, timeout=None): + try: + return await asyncio.wait_for(self.queue.get(), timeout) + except asyncio.TimeoutError: + raise Empty + + def put_nowait(self, item): + self.queue.put_nowait(item) + + def put_nowait_batch(self, items): + # If maxsize is 0, queue is unbounded, so no need to check size. + if self.maxsize > 0 and len(items) + self.qsize() > self.maxsize: + raise Full( + f"Cannot add {len(items)} items to queue of size " + f"{self.qsize()} and maxsize {self.maxsize}." + ) + for item in items: + self.queue.put_nowait(item) + + def get_nowait(self): + return self.queue.get_nowait() + + def get_nowait_batch(self, num_items): + if num_items > self.qsize(): + raise Empty( + f"Cannot get {num_items} items from queue of size " f"{self.qsize()}." + ) + return [self.queue.get_nowait() for _ in range(num_items)] diff --git a/lib/python3.12/site-packages/ray/util/scheduling_strategies.py b/lib/python3.12/site-packages/ray/util/scheduling_strategies.py new file mode 100644 index 0000000000000000000000000000000000000000..68cb6a9d4700d834b66c6b944cdb3f7d4dbde1a1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/scheduling_strategies.py @@ -0,0 +1,226 @@ +from typing import TYPE_CHECKING, Dict, Optional, Union + +from ray._raylet import NodeID +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.util.placement_group import PlacementGroup + +# "DEFAULT": The default hybrid scheduling strategy +# based on config scheduler_spread_threshold. +# This disables any potential placement group capture. + +# "SPREAD": Spread scheduling on a best effort basis. + + +@PublicAPI +class PlacementGroupSchedulingStrategy: + """Placement group based scheduling strategy. + + Attributes: + placement_group: the placement group this actor belongs to, + or None if it doesn't belong to any group. + placement_group_bundle_index: the index of the bundle + if the actor belongs to a placement group, which may be -1 to + specify any available bundle. + placement_group_capture_child_tasks: Whether or not children tasks + of this actor should implicitly use the same placement group + as its parent. It is False by default. + """ + + def __init__( + self, + placement_group: "PlacementGroup", + placement_group_bundle_index: int = -1, + placement_group_capture_child_tasks: Optional[bool] = None, + ): + self.placement_group = placement_group + self.placement_group_bundle_index = placement_group_bundle_index + self.placement_group_capture_child_tasks = placement_group_capture_child_tasks + + +@PublicAPI +class NodeAffinitySchedulingStrategy: + """Static scheduling strategy used to run a task or actor on a particular node. + + Attributes: + node_id: the hex id of the node where the task or actor should run. + soft: whether the scheduler should run the task or actor somewhere else + if the target node doesn't exist (e.g. the node dies) or is infeasible + during scheduling. + If the node exists and is feasible, the task or actor + will only be scheduled there. + This means if the node doesn't have the available resources, + the task or actor will wait indefinitely until resources become available. + If the node doesn't exist or is infeasible, the task or actor + will fail if soft is False + or be scheduled somewhere else if soft is True. + """ + + def __init__( + self, + node_id: str, + soft: bool, + _spill_on_unavailable: bool = False, + _fail_on_unavailable: bool = False, + ): + # This will be removed once we standardize on node id being hex string. + if not isinstance(node_id, str): + node_id = node_id.hex() + + self.node_id = node_id + self.soft = soft + self._spill_on_unavailable = _spill_on_unavailable + self._fail_on_unavailable = _fail_on_unavailable + + self._validate_attributes() + + def _validate_attributes(self): + invalid_node_id_error = ValueError( + f"Invalid node_id '{self.node_id}'. Node ID must be a valid " + "hex string. To get a list of all nodes and their IDs in your cluster, " + "use ray.nodes(). See https://docs.ray.io/en/latest/ray-core/miscellaneous.html#node-information for more details." + ) + try: + node_id = NodeID.from_hex(self.node_id) + except Exception as e: + raise invalid_node_id_error from e + + if node_id.is_nil(): + raise invalid_node_id_error + + if self._spill_on_unavailable and not self.soft: + raise ValueError( + "_spill_on_unavailable cannot be set when soft is " + "False. Please set soft to True to use _spill_on_unavailable." + ) + if self._fail_on_unavailable and self.soft: + raise ValueError( + "_fail_on_unavailable cannot be set when soft is " + "True. Please set soft to False to use _fail_on_unavailable." + ) + + +def _validate_label_match_operator_values(values, operator): + if not values: + raise ValueError( + f"The variadic parameter of the {operator} operator" + f' must be a non-empty tuple: e.g. {operator}("value1", "value2").' + ) + + index = 0 + for value in values: + if not isinstance(value, str): + raise ValueError( + f"Type of value in position {index} for the {operator} operator " + f'must be str (e.g. {operator}("value1", "value2")) ' + f"but got {str(value)} of type {type(value)}." + ) + index = index + 1 + + +@PublicAPI(stability="alpha") +class In: + def __init__(self, *values): + _validate_label_match_operator_values(values, "In") + self.values = list(values) + + +@PublicAPI(stability="alpha") +class NotIn: + def __init__(self, *values): + _validate_label_match_operator_values(values, "NotIn") + self.values = list(values) + + +@PublicAPI(stability="alpha") +class Exists: + def __init__(self): + pass + + +@PublicAPI(stability="alpha") +class DoesNotExist: + def __init__(self): + pass + + +class _LabelMatchExpression: + """An expression used to select node by node's labels + Attributes: + key: the key of label + operator: In、NotIn、Exists、DoesNotExist + """ + + def __init__(self, key: str, operator: Union[In, NotIn, Exists, DoesNotExist]): + self.key = key + self.operator = operator + + +LabelMatchExpressionsT = Dict[str, Union[In, NotIn, Exists, DoesNotExist]] + + +@PublicAPI(stability="alpha") +class NodeLabelSchedulingStrategy: + """ + Label based node affinity scheduling strategy + + scheduling_strategy=NodeLabelSchedulingStrategy({ + "region": In("us"), + "gpu_type": Exists(), + }) + """ + + def __init__( + self, hard: LabelMatchExpressionsT, *, soft: LabelMatchExpressionsT = None + ): + self.hard = _convert_map_to_expressions(hard, "hard") + self.soft = _convert_map_to_expressions(soft, "soft") + self._check_usage() + + def _check_usage(self): + if not (self.hard or self.soft): + raise ValueError( + "The `hard` and `soft` parameter " + "of NodeLabelSchedulingStrategy cannot both be empty." + ) + + +def _convert_map_to_expressions(map_expressions: LabelMatchExpressionsT, param: str): + expressions = [] + if map_expressions is None: + return expressions + + if not isinstance(map_expressions, Dict): + raise ValueError( + f'The {param} parameter must be a map (e.g. {{"key1": In("value1")}}) ' + f"but got type {type(map_expressions)}." + ) + + for key, value in map_expressions.items(): + if not isinstance(key, str): + raise ValueError( + f"The map key of the {param} parameter must " + f'be of type str (e.g. {{"key1": In("value1")}}) ' + f"but got {str(key)} of type {type(key)}." + ) + + if not isinstance(value, (In, NotIn, Exists, DoesNotExist)): + raise ValueError( + f"The map value for key {key} of the {param} parameter " + f"must be one of the `In`, `NotIn`, `Exists` or `DoesNotExist` " + f'operator (e.g. {{"key1": In("value1")}}) ' + f"but got {str(value)} of type {type(value)}." + ) + + expressions.append(_LabelMatchExpression(key, value)) + return expressions + + +SchedulingStrategyT = Union[ + None, + str, # Literal["DEFAULT", "SPREAD"] + PlacementGroupSchedulingStrategy, + NodeAffinitySchedulingStrategy, + NodeLabelSchedulingStrategy, +] diff --git a/lib/python3.12/site-packages/ray/util/serialization.py b/lib/python3.12/site-packages/ray/util/serialization.py new file mode 100644 index 0000000000000000000000000000000000000000..106e06c8681bc959db3f9ae7e0c64e5c0187bb38 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/serialization.py @@ -0,0 +1,55 @@ +import ray +import ray.cloudpickle as pickle +from ray.util.annotations import DeveloperAPI, PublicAPI + + +@PublicAPI +def register_serializer(cls: type, *, serializer: callable, deserializer: callable): + """Use the given serializer to serialize instances of type ``cls``, + and use the deserializer to deserialize the serialized object. + + Args: + cls: A Python class/type. + serializer: A function that converts an instances of + type ``cls`` into a serializable object (e.g. python dict + of basic objects). + deserializer: A function that constructs the + instance of type ``cls`` from the serialized object. + This function itself must be serializable. + """ + context = ray._private.worker.global_worker.get_serialization_context() + context._register_cloudpickle_serializer(cls, serializer, deserializer) + + +@PublicAPI +def deregister_serializer(cls: type): + """Deregister the serializer associated with the type ``cls``. + There is no effect if the serializer is unavailable. + + Args: + cls: A Python class/type. + """ + context = ray._private.worker.global_worker.get_serialization_context() + context._unregister_cloudpickle_reducer(cls) + + +@DeveloperAPI +class StandaloneSerializationContext: + # NOTE(simon): Used for registering custom serializers. We cannot directly + # use the SerializationContext because it requires Ray workers. Please + # make sure to keep the API consistent. + + def _register_cloudpickle_reducer(self, cls, reducer): + pickle.CloudPickler.dispatch[cls] = reducer + + def _unregister_cloudpickle_reducer(self, cls): + pickle.CloudPickler.dispatch.pop(cls, None) + + def _register_cloudpickle_serializer( + self, cls, custom_serializer, custom_deserializer + ): + def _CloudPicklerReducer(obj): + return custom_deserializer, (custom_serializer(obj),) + + # construct a reducer + pickle.CloudPickler.dispatch[cls] = _CloudPicklerReducer diff --git a/lib/python3.12/site-packages/ray/util/serialization_addons.py b/lib/python3.12/site-packages/ray/util/serialization_addons.py new file mode 100644 index 0000000000000000000000000000000000000000..497a4269b0df731d78cde51d07910f1c0b48af69 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/serialization_addons.py @@ -0,0 +1,39 @@ +""" +This module is intended for implementing internal serializers for some +site packages. +""" + +import sys + +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +def register_starlette_serializer(serialization_context): + try: + import starlette.datastructures + except ImportError: + return + + # Starlette's app.state object is not serializable + # because it overrides __getattr__ + serialization_context._register_cloudpickle_serializer( + starlette.datastructures.State, + custom_serializer=lambda s: s._state, + custom_deserializer=lambda s: starlette.datastructures.State(s), + ) + + +@DeveloperAPI +def apply(serialization_context): + from ray._common.pydantic_compat import register_pydantic_serializers + + register_pydantic_serializers(serialization_context) + register_starlette_serializer(serialization_context) + + if sys.platform != "win32": + from ray._private.arrow_serialization import ( + _register_custom_datasets_serializers, + ) + + _register_custom_datasets_serializers(serialization_context) diff --git a/lib/python3.12/site-packages/ray/util/spark/__init__.py b/lib/python3.12/site-packages/ray/util/spark/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..69d68172eb19551e52a1754cdeac9a5166fb2476 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/spark/__init__.py @@ -0,0 +1,13 @@ +from ray.util.spark.cluster_init import ( + MAX_NUM_WORKER_NODES, + setup_global_ray_cluster, + setup_ray_cluster, + shutdown_ray_cluster, +) + +__all__ = [ + "setup_ray_cluster", + "shutdown_ray_cluster", + "MAX_NUM_WORKER_NODES", + "setup_global_ray_cluster", +] diff --git a/lib/python3.12/site-packages/ray/util/spark/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/spark/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bbadfce0b316b83e40578607968e0ace3ed88a16 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/spark/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/spark/__pycache__/cluster_init.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/spark/__pycache__/cluster_init.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..579625698dc03fb047def426082569f0c45b54ae Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/spark/__pycache__/cluster_init.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/spark/__pycache__/databricks_hook.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/spark/__pycache__/databricks_hook.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..408753a1a351e72b41b8977bb867e0097a66e09d Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/spark/__pycache__/databricks_hook.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/spark/__pycache__/start_hook_base.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/spark/__pycache__/start_hook_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0b10027fe2dc97b2e748a307513a67ceab611f98 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/spark/__pycache__/start_hook_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/spark/__pycache__/start_ray_node.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/spark/__pycache__/start_ray_node.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cbfe6abf7fdbd01a0ffceeddf194c97ca3e55a37 Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/spark/__pycache__/start_ray_node.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/spark/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/util/spark/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..951cbcfb74b37661e0b4415d33dba636b336391e Binary files /dev/null and b/lib/python3.12/site-packages/ray/util/spark/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/util/spark/cluster_init.py b/lib/python3.12/site-packages/ray/util/spark/cluster_init.py new file mode 100644 index 0000000000000000000000000000000000000000..9c5c696fa4592aa0f4496890d8cb5866dc14e59d --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/spark/cluster_init.py @@ -0,0 +1,1922 @@ +import copy +import json +import logging +import os +import signal +import socket +import sys +import threading +import time +import uuid +import warnings +from threading import Event +from typing import Dict, Optional, Tuple, Type + +import requests +import yaml +from packaging.version import Version + +import ray +import ray._private.services +from .databricks_hook import DefaultDatabricksRayOnSparkStartHook +from .start_hook_base import RayOnSparkStartHook +from .utils import ( + _get_cpu_cores, + _get_local_ray_node_slots, + _get_num_physical_gpus, + _wait_service_up, + calc_mem_ray_head_node, + exec_cmd, + gen_cmd_exec_failure_msg, + get_avail_mem_per_ray_worker_node, + get_configured_spark_executor_memory_bytes, + get_max_num_concurrent_tasks, + get_random_unused_port, + get_spark_application_driver_host, + get_spark_session, + get_spark_task_assigned_physical_gpus, + is_in_databricks_runtime, + is_port_in_use, +) +from ray._common.network_utils import build_address, parse_address +from ray._common.utils import load_class +from ray.autoscaler._private.spark.node_provider import HEAD_NODE_ID +from ray.util.annotations import DeveloperAPI, PublicAPI + +_logger = logging.getLogger("ray.util.spark") +_logger.setLevel(logging.INFO) + + +RAY_ON_SPARK_START_HOOK = "RAY_ON_SPARK_START_HOOK" + +MAX_NUM_WORKER_NODES = -1 + +RAY_ON_SPARK_COLLECT_LOG_TO_PATH = "RAY_ON_SPARK_COLLECT_LOG_TO_PATH" +RAY_ON_SPARK_START_RAY_PARENT_PID = "RAY_ON_SPARK_START_RAY_PARENT_PID" + + +def _check_system_environment(): + if os.name != "posix": + raise RuntimeError("Ray on spark only supports running on POSIX system.") + + spark_dependency_error = "ray.util.spark module requires pyspark >= 3.3" + try: + import pyspark + + if Version(pyspark.__version__).release < (3, 3, 0): + raise RuntimeError(spark_dependency_error) + except ImportError: + raise RuntimeError(spark_dependency_error) + + +class RayClusterOnSpark: + """ + This class is the type of instance returned by the `_setup_ray_cluster` interface. + Its main functionality is to: + Connect to, disconnect from, and shutdown the Ray cluster running on Apache Spark. + Serve as a Python context manager for the `RayClusterOnSpark` instance. + + Args + address: The url for the ray head node (defined as the hostname and unused + port on Spark driver node) + head_proc: Ray head process + spark_job_group_id: The Spark job id for a submitted ray job + num_workers_node: The number of workers in the ray cluster. + """ + + def __init__( + self, + address, + head_proc, + min_worker_nodes, + max_worker_nodes, + temp_dir, + cluster_unique_id, + start_hook, + ray_dashboard_port, + spark_job_server, + global_cluster_lock_fd, + ray_client_server_port, + ): + self.address = address + self.head_proc = head_proc + self.min_worker_nodes = min_worker_nodes + self.max_worker_nodes = max_worker_nodes + self.temp_dir = temp_dir + self.cluster_unique_id = cluster_unique_id + self.start_hook = start_hook + self.ray_dashboard_port = ray_dashboard_port + self.spark_job_server = spark_job_server + self.global_cluster_lock_fd = global_cluster_lock_fd + self.ray_client_server_port = ray_client_server_port + + self.is_shutdown = False + self.spark_job_is_canceled = False + self.background_job_exception = None + + # Ray client context returns by `ray.init` + self.ray_ctx = None + + def wait_until_ready(self): + import ray + + if self.is_shutdown: + raise RuntimeError( + "The ray cluster has been shut down or it failed to start." + ) + + try: + ray.init(address=self.address) + + if self.ray_dashboard_port is not None and _wait_service_up( + parse_address(self.address)[0], + self.ray_dashboard_port, + _RAY_DASHBOARD_STARTUP_TIMEOUT, + ): + self.start_hook.on_ray_dashboard_created(self.ray_dashboard_port) + else: + try: + __import__("ray.dashboard.optional_deps") + except ModuleNotFoundError as e: + _logger.warning( + "Dependencies to launch the optional dashboard API " + "server cannot be found. They can be installed with " + f"pip install ray[default], root cause: ({repr(e)})" + ) + + last_alive_worker_count = 0 + last_progress_move_time = time.time() + while True: + time.sleep(_RAY_CLUSTER_STARTUP_PROGRESS_CHECKING_INTERVAL) + + # Inside the waiting ready loop, + # checking `self.background_job_exception`, if it is not None, + # it means the background spark job has failed, + # in this case, raise error directly. + if self.background_job_exception is not None: + raise RuntimeError( + "Ray workers failed to start." + ) from self.background_job_exception + + cur_alive_worker_count = ( + len([node for node in ray.nodes() if node["Alive"]]) - 1 + ) # Minus 1 means excluding the head node. + + if cur_alive_worker_count >= self.min_worker_nodes: + _logger.info( + f"Started {cur_alive_worker_count} Ray worker nodes, " + f"meet the minimum number of Ray worker nodes required." + ) + return + + if cur_alive_worker_count > last_alive_worker_count: + last_alive_worker_count = cur_alive_worker_count + last_progress_move_time = time.time() + _logger.info( + "Ray worker nodes are starting. Progress: " + f"({cur_alive_worker_count} / {self.max_worker_nodes})" + ) + else: + if ( + time.time() - last_progress_move_time + > _RAY_CONNECT_CLUSTER_POLL_PROGRESS_TIMEOUT + ): + if cur_alive_worker_count == 0: + ( + job_server_host, + job_server_port, + ) = self.spark_job_server.server_address[:2] + response = requests.post( + url=( + f"http://{build_address(job_server_host, job_server_port)}" + "/query_last_worker_err" + ), + json={"spark_job_group_id": None}, + ) + response.raise_for_status() + + decoded_resp = response.content.decode("utf-8") + json_res = json.loads(decoded_resp) + last_worker_err = json_res["last_worker_err"] + + if last_worker_err: + raise RuntimeError( + "Starting Ray worker node failed, error:\n" + f"{last_worker_err}" + ) + else: + raise RuntimeError( + "Current spark cluster has no resources to launch " + "Ray worker nodes." + ) + _logger.warning( + "Timeout in waiting for minimal ray workers to start. " + "Started / Total requested: " + f"({cur_alive_worker_count} / {self.min_worker_nodes}). " + "Current spark cluster does not have sufficient resources " + "to launch requested minimal number of Ray worker nodes." + ) + return + finally: + ray.shutdown() + + def connect(self): + if ray.is_initialized(): + raise RuntimeError("Already connected to Ray cluster.") + self.ray_ctx = ray.init(address=self.address) + + def disconnect(self): + ray.shutdown() + self.ray_ctx = None + + def shutdown(self): + """ + Shutdown the ray cluster created by the `setup_ray_cluster` API. + """ + import fcntl + + if not self.is_shutdown: + try: + self.disconnect() + except Exception: + pass + os.environ.pop("RAY_ADDRESS", None) + + if self.global_cluster_lock_fd is not None: + # release global mode cluster lock. + fcntl.flock(self.global_cluster_lock_fd, fcntl.LOCK_UN) + + self.spark_job_server.shutdown() + try: + self.head_proc.terminate() + except Exception as e: + # swallow exception. + _logger.warning( + "An Error occurred during shutdown of ray head node: " f"{repr(e)}" + ) + self.is_shutdown = True + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.shutdown() + + +def _convert_ray_node_option(key, value): + converted_key = f"--{key.replace('_', '-')}" + if key in ["system_config", "resources", "labels"]: + return f"{converted_key}={json.dumps(value)}" + if value is None: + return converted_key + return f"{converted_key}={str(value)}" + + +def _convert_ray_node_options(options): + return [_convert_ray_node_option(k, v) for k, v in options.items()] + + +_RAY_HEAD_STARTUP_TIMEOUT = 20 +_RAY_DASHBOARD_STARTUP_TIMEOUT = 60 +_BACKGROUND_JOB_STARTUP_WAIT = int( + os.environ.get("RAY_ON_SPARK_BACKGROUND_JOB_STARTUP_WAIT", "30") +) +_RAY_CLUSTER_STARTUP_PROGRESS_CHECKING_INTERVAL = 3 +_RAY_WORKER_NODE_STARTUP_INTERVAL = int( + os.environ.get("RAY_ON_SPARK_RAY_WORKER_NODE_STARTUP_INTERVAL", "10") +) +_RAY_CONNECT_CLUSTER_POLL_PROGRESS_TIMEOUT = 120 + + +def _preallocate_ray_worker_port_range(): + """ + If we start multiple ray workers on a machine concurrently, some ray worker + processes might fail due to ray port conflicts, this is because race condition + on getting free port and opening the free port. + To address the issue, this function use an exclusive file lock to delay the + worker processes to ensure that port acquisition does not create a resource + contention issue due to a race condition. + + After acquiring lock, it will allocate port range for worker ports + (for ray node config --min-worker-port and --max-worker-port). + Because on a spark cluster, multiple ray cluster might be created, so on one spark + worker machine, there might be multiple ray worker nodes running, these worker + nodes might belong to different ray cluster, and we must ensure these ray nodes on + the same machine using non-overlapping worker port range, to achieve this, in this + function, it creates a file `/tmp/ray_on_spark_worker_port_allocation.txt` file, + the file format is composed of multiple lines, each line contains 2 number: `pid` + and `port_range_slot_index`, each port range slot allocates 1000 ports, and + corresponding port range is: + - range_begin (inclusive): 20000 + port_range_slot_index * 1000 + - range_end (exclusive): range_begin + 1000 + In this function, it first scans `/tmp/ray_on_spark_worker_port_allocation.txt` + file, removing lines that containing dead process pid, then find the first unused + port_range_slot_index, then regenerate this file, and return the allocated port + range. + + Returns: Allocated port range for current worker ports + """ + import fcntl + + import psutil + + def acquire_lock(file_path): + mode = os.O_RDWR | os.O_CREAT | os.O_TRUNC + try: + fd = os.open(file_path, mode) + # The lock file must be readable / writable to all users. + os.chmod(file_path, 0o0777) + # Allow for retrying getting a file lock a maximum number of seconds + max_lock_iter = 600 + for _ in range(max_lock_iter): + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + # Lock is used by other processes, continue loop to wait for lock + # available + pass + else: + # Acquire lock successfully. + return fd + time.sleep(10) + raise TimeoutError(f"Acquiring lock on file {file_path} timeout.") + except Exception: + os.close(fd) + + lock_file_path = "/tmp/ray_on_spark_worker_startup_barrier_lock.lock" + try: + lock_fd = acquire_lock(lock_file_path) + except TimeoutError: + # If timeout happens, the file lock might be hold by another process and that + # process does not release the lock in time by some unexpected reason. + # In this case, remove the existing lock file and create the file again, and + # then acquire file lock on the new file. + try: + os.remove(lock_file_path) + except Exception: + pass + lock_fd = acquire_lock(lock_file_path) + + def release_lock(): + fcntl.flock(lock_fd, fcntl.LOCK_UN) + os.close(lock_fd) + + try: + port_alloc_file = "/tmp/ray_on_spark_worker_port_allocation.txt" + + # NB: reading / writing `port_alloc_file` is protected by exclusive lock + # on file `lock_file_path` + if os.path.exists(port_alloc_file): + with open(port_alloc_file, mode="r") as fp: + port_alloc_data = fp.read() + port_alloc_table = [ + line.split(" ") for line in port_alloc_data.strip().split("\n") + ] + port_alloc_table = [ + (int(pid_str), int(slot_index_str)) + for pid_str, slot_index_str in port_alloc_table + ] + else: + port_alloc_table = [] + with open(port_alloc_file, mode="w"): + pass + # The port range allocation file must be readable / writable to all users. + os.chmod(port_alloc_file, 0o0777) + + port_alloc_map = { + pid: slot_index + for pid, slot_index in port_alloc_table + if psutil.pid_exists(pid) # remove slot used by dead process + } + + allocated_slot_set = set(port_alloc_map.values()) + + if len(allocated_slot_set) == 0: + new_slot_index = 0 + else: + new_slot_index = max(allocated_slot_set) + 1 + for index in range(new_slot_index): + if index not in allocated_slot_set: + new_slot_index = index + break + + port_alloc_map[os.getpid()] = new_slot_index + + with open(port_alloc_file, mode="w") as fp: + for pid, slot_index in port_alloc_map.items(): + fp.write(f"{pid} {slot_index}\n") + + worker_port_range_begin = 20000 + new_slot_index * 1000 + worker_port_range_end = worker_port_range_begin + 1000 + + if worker_port_range_end > 65536: + raise RuntimeError( + "Too many ray worker nodes are running on this machine, cannot " + "allocate worker port range for new ray worker node." + ) + except Exception: + release_lock() + raise + + def hold_lock(): + time.sleep(_RAY_WORKER_NODE_STARTUP_INTERVAL) + release_lock() + + threading.Thread(target=hold_lock, args=()).start() + + return worker_port_range_begin, worker_port_range_end + + +def _append_default_spilling_dir_config(head_node_options, object_spilling_dir): + if "system_config" not in head_node_options: + head_node_options["system_config"] = {} + sys_conf = head_node_options["system_config"] + if "object_spilling_config" not in sys_conf: + sys_conf["object_spilling_config"] = json.dumps( + { + "type": "filesystem", + "params": { + "directory_path": object_spilling_dir, + }, + } + ) + return head_node_options + + +def _append_resources_config(node_options, resources): + if "resources" not in node_options: + node_options["resources"] = {} + + node_options["resources"].update(resources) + return node_options + + +def _get_default_ray_tmp_dir(): + return os.path.join(os.environ.get("RAY_TMPDIR", "/tmp"), "ray") + + +def _create_hook_entry(is_global): + if RAY_ON_SPARK_START_HOOK in os.environ: + return load_class(os.environ[RAY_ON_SPARK_START_HOOK])() + elif is_in_databricks_runtime(): + return DefaultDatabricksRayOnSparkStartHook(is_global) + else: + return RayOnSparkStartHook(is_global) + + +def _setup_ray_cluster( + *, + max_worker_nodes: int, + min_worker_nodes: int, + num_cpus_worker_node: int, + num_cpus_head_node: int, + num_gpus_worker_node: int, + num_gpus_head_node: int, + using_stage_scheduling: bool, + heap_memory_worker_node: int, + heap_memory_head_node: int, + object_store_memory_worker_node: int, + object_store_memory_head_node: int, + head_node_options: Dict, + worker_node_options: Dict, + ray_temp_root_dir: str, + collect_log_to_path: str, + autoscale_upscaling_speed: float, + autoscale_idle_timeout_minutes: float, + is_global: bool, +) -> Type[RayClusterOnSpark]: + """ + The public API `ray.util.spark.setup_ray_cluster` does some argument + validation and then pass validated arguments to this interface. + and it returns a `RayClusterOnSpark` instance. + + The returned instance can be used to connect to, disconnect from and shutdown the + ray cluster. This instance can also be used as a context manager (used by + encapsulating operations within `with _setup_ray_cluster(...):`). Upon entering the + managed scope, the ray cluster is initiated and connected to. When exiting the + scope, the ray cluster is disconnected and shut down. + + Note: This function interface is stable and can be used for + instrumentation logging patching. + """ + import fcntl + + start_hook = _create_hook_entry(is_global) + spark = get_spark_session() + + ray_head_ip = socket.gethostbyname(get_spark_application_driver_host(spark)) + ray_head_port = get_random_unused_port(ray_head_ip, min_port=9000, max_port=10000) + port_exclude_list = [ray_head_port] + + # Make a copy for head_node_options to avoid changing original dict in user code. + head_node_options = head_node_options.copy() + include_dashboard = head_node_options.pop("include_dashboard", None) + ray_dashboard_port = head_node_options.pop("dashboard_port", None) + + if is_global: + ray_client_server_port = 10001 + else: + ray_client_server_port = get_random_unused_port( + ray_head_ip, + min_port=9000, + max_port=10000, + exclude_list=port_exclude_list, + ) + + port_exclude_list.append(ray_client_server_port) + + spark_job_server_port = get_random_unused_port( + ray_head_ip, + min_port=9000, + max_port=10000, + exclude_list=port_exclude_list, + ) + port_exclude_list.append(spark_job_server_port) + + if include_dashboard is None or include_dashboard is True: + if ray_dashboard_port is None: + ray_dashboard_port = get_random_unused_port( + ray_head_ip, + min_port=9000, + max_port=10000, + exclude_list=port_exclude_list, + ) + port_exclude_list.append(ray_dashboard_port) + ray_dashboard_agent_port = get_random_unused_port( + ray_head_ip, + min_port=9000, + max_port=10000, + exclude_list=port_exclude_list, + ) + port_exclude_list.append(ray_dashboard_agent_port) + + dashboard_options = [ + "--dashboard-host=0.0.0.0", + f"--dashboard-port={ray_dashboard_port}", + f"--dashboard-agent-listen-port={ray_dashboard_agent_port}", + ] + # If include_dashboard is None, we don't set `--include-dashboard` option, + # in this case Ray will decide whether dashboard can be started + # (e.g. checking any missing dependencies). + if include_dashboard is True: + dashboard_options += ["--include-dashboard=true"] + else: + dashboard_options = [ + "--include-dashboard=false", + ] + + _logger.info( + f"Ray head hostname: {ray_head_ip}, port: {ray_head_port}, " + f"ray client server port: {ray_client_server_port}." + ) + + cluster_unique_id = uuid.uuid4().hex[:8] + + if is_global: + # global mode enabled + # for global mode, Ray always uses default temp dir + # so that local Ray client can discover it without specifying + # head node address. + if ray_temp_root_dir is not None: + raise ValueError( + "Ray on spark global mode cluster does not allow you to set " + "'ray_temp_root_dir' argument." + ) + + # We only allow user to launch one active Ray on spark global cluster + # at a time. So acquiring a global file lock before setting up a new + # Ray on spark global cluster. + global_cluster_lock_fd = os.open( + "/tmp/ray_on_spark_global_cluster.lock", os.O_RDWR | os.O_CREAT | os.O_TRUNC + ) + + try: + # acquiring exclusive lock to ensure copy logs and removing dir safely. + fcntl.flock(global_cluster_lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError: + # acquiring global lock failed. + raise ValueError( + "Acquiring global lock failed for setting up new global mode Ray on " + "spark cluster. If there is an active global mode Ray on spark " + "cluster, please shut down it before you create a new one." + ) + + ray_temp_dir = None + ray_default_tmp_dir = _get_default_ray_tmp_dir() + os.makedirs(ray_default_tmp_dir, exist_ok=True) + object_spilling_dir = os.path.join(ray_default_tmp_dir, "spill") + else: + global_cluster_lock_fd = None + if ray_temp_root_dir is None: + ray_temp_root_dir = start_hook.get_default_temp_root_dir() + ray_temp_dir = os.path.join( + ray_temp_root_dir, f"ray-{ray_head_port}-{cluster_unique_id}" + ) + os.makedirs(ray_temp_dir, exist_ok=True) + object_spilling_dir = os.path.join(ray_temp_dir, "spill") + + os.makedirs(object_spilling_dir, exist_ok=True) + + head_node_options = _append_default_spilling_dir_config( + head_node_options, object_spilling_dir + ) + + from ray.autoscaler._private.spark.spark_job_server import ( + _start_spark_job_server, + ) + + ray_node_custom_env = start_hook.custom_environment_variables() + spark_job_server = _start_spark_job_server( + ray_head_ip, spark_job_server_port, spark, ray_node_custom_env + ) + autoscaling_cluster = AutoscalingCluster( + head_resources={ + "CPU": num_cpus_head_node, + "GPU": num_gpus_head_node, + "memory": heap_memory_head_node, + "object_store_memory": object_store_memory_head_node, + }, + worker_node_types={ + "ray.worker": { + "resources": { + "CPU": num_cpus_worker_node, + "GPU": num_gpus_worker_node, + "memory": heap_memory_worker_node, + "object_store_memory": object_store_memory_worker_node, + }, + "node_config": {}, + "min_workers": min_worker_nodes, + "max_workers": max_worker_nodes, + }, + }, + extra_provider_config={ + "ray_head_ip": ray_head_ip, + "ray_head_port": ray_head_port, + "cluster_unique_id": cluster_unique_id, + "using_stage_scheduling": using_stage_scheduling, + "ray_temp_dir": ray_temp_dir, + "worker_node_options": worker_node_options, + "collect_log_to_path": collect_log_to_path, + "spark_job_server_port": spark_job_server_port, + }, + upscaling_speed=autoscale_upscaling_speed, + idle_timeout_minutes=autoscale_idle_timeout_minutes, + ) + ray_head_proc, tail_output_deque = autoscaling_cluster.start( + ray_head_ip, + ray_head_port, + ray_client_server_port, + ray_temp_dir, + dashboard_options, + head_node_options, + collect_log_to_path, + ray_node_custom_env, + ) + ray_head_node_cmd = autoscaling_cluster.ray_head_node_cmd + + # wait ray head node spin up. + time.sleep(_RAY_HEAD_STARTUP_TIMEOUT) + + if not is_port_in_use(ray_head_ip, ray_head_port): + if ray_head_proc.poll() is None: + # Ray head GCS service is down. Kill ray head node. + ray_head_proc.terminate() + # wait killing complete. + time.sleep(0.5) + + cmd_exec_failure_msg = gen_cmd_exec_failure_msg( + ray_head_node_cmd, ray_head_proc.returncode, tail_output_deque + ) + raise RuntimeError("Start Ray head node failed!\n" + cmd_exec_failure_msg) + + _logger.info("Ray head node started.") + + cluster_address = build_address(ray_head_ip, ray_head_port) + # Set RAY_ADDRESS environment variable to the cluster address. + os.environ["RAY_ADDRESS"] = cluster_address + + ray_cluster_handler = RayClusterOnSpark( + address=cluster_address, + head_proc=ray_head_proc, + min_worker_nodes=min_worker_nodes, + max_worker_nodes=max_worker_nodes, + temp_dir=ray_temp_dir, + cluster_unique_id=cluster_unique_id, + start_hook=start_hook, + ray_dashboard_port=ray_dashboard_port, + spark_job_server=spark_job_server, + global_cluster_lock_fd=global_cluster_lock_fd, + ray_client_server_port=ray_client_server_port, + ) + + start_hook.on_cluster_created(ray_cluster_handler) + + return ray_cluster_handler + + +_active_ray_cluster = None +_active_ray_cluster_rwlock = threading.RLock() + + +def _create_resource_profile(num_cpus_per_node, num_gpus_per_node): + from pyspark.resource.profile import ResourceProfileBuilder + from pyspark.resource.requests import TaskResourceRequests + + task_res_req = TaskResourceRequests().cpus(num_cpus_per_node) + if num_gpus_per_node > 0: + task_res_req = task_res_req.resource("gpu", num_gpus_per_node) + return ResourceProfileBuilder().require(task_res_req).build + + +# A dict storing blocked key to replacement argument you should use. +_head_node_option_block_keys = { + "temp_dir": "ray_temp_root_dir", + "block": None, + "head": None, + "node_ip_address": None, + "port": None, + "num_cpus": None, + "num_gpus": None, + "dashboard_host": None, + "dashboard_agent_listen_port": None, +} + +_worker_node_option_block_keys = { + "temp_dir": "ray_temp_root_dir", + "block": None, + "head": None, + "address": None, + "num_cpus": "num_cpus_worker_node", + "num_gpus": "num_gpus_worker_node", + "memory": None, + "object_store_memory": "object_store_memory_worker_node", + "dashboard_agent_listen_port": None, + "min_worker_port": None, + "max_worker_port": None, +} + + +def _verify_node_options(node_options, block_keys, node_type): + for key in node_options: + if key.startswith("--") or "-" in key: + raise ValueError( + "For a ray node option like '--foo-bar', you should convert it to " + "following format 'foo_bar' in 'head_node_options' / " + "'worker_node_options' arguments." + ) + + if key in block_keys: + common_err_msg = ( + f"Setting the option '{key}' for {node_type} nodes is not allowed." + ) + replacement_arg = block_keys[key] + if replacement_arg: + raise ValueError( + f"{common_err_msg} You should set the '{replacement_arg}' option " + "instead." + ) + else: + raise ValueError( + f"{common_err_msg} This option is controlled by Ray on Spark." + ) + + +def _setup_ray_cluster_internal( + max_worker_nodes: int, + min_worker_nodes: Optional[int], + num_cpus_worker_node: Optional[int], + num_cpus_head_node: Optional[int], + num_gpus_worker_node: Optional[int], + num_gpus_head_node: Optional[int], + heap_memory_worker_node: Optional[int], + heap_memory_head_node: Optional[int], + object_store_memory_worker_node: Optional[int], + object_store_memory_head_node: Optional[int], + head_node_options: Optional[Dict], + worker_node_options: Optional[Dict], + ray_temp_root_dir: Optional[str], + strict_mode: bool, + collect_log_to_path: Optional[str], + autoscale_upscaling_speed: Optional[float], + autoscale_idle_timeout_minutes: Optional[float], + is_global: bool, + **kwargs, +) -> Tuple[str, str]: + global _active_ray_cluster + + _check_system_environment() + _install_sigterm_signal() + + head_node_options = head_node_options or {} + worker_node_options = worker_node_options or {} + + _verify_node_options( + head_node_options, + _head_node_option_block_keys, + "Ray head node on spark", + ) + _verify_node_options( + worker_node_options, + _worker_node_option_block_keys, + "Ray worker node on spark", + ) + + if _active_ray_cluster is not None: + raise RuntimeError( + "Current active ray cluster on spark haven't shut down. Please call " + "`ray.util.spark.shutdown_ray_cluster()` before initiating a new Ray " + "cluster on spark." + ) + + if ray.is_initialized(): + raise RuntimeError( + "Current python process already initialized Ray, Please shut down it " + "by `ray.shutdown()` before initiating a Ray cluster on spark." + ) + + spark = get_spark_session() + + spark_master = spark.sparkContext.master + + is_spark_local_mode = spark_master == "local" or spark_master.startswith("local[") + + if not ( + spark_master.startswith("spark://") + or spark_master.startswith("local-cluster[") + or spark_master == "yarn" + or is_spark_local_mode + ): + raise RuntimeError( + "Ray on Spark only supports spark cluster in standalone mode, " + "local-cluster mode, spark on yarn mode or spark local mode." + ) + + if is_spark_local_mode: + support_stage_scheduling = False + elif ( + is_in_databricks_runtime() + and Version(os.environ["DATABRICKS_RUNTIME_VERSION"]).major >= 12 + ): + support_stage_scheduling = True + else: + import pyspark + + if Version(pyspark.__version__).release >= (3, 4, 0): + support_stage_scheduling = True + else: + support_stage_scheduling = False + + if "num_cpus_per_node" in kwargs: + if num_cpus_worker_node is not None: + raise ValueError( + "'num_cpus_per_node' and 'num_cpus_worker_node' arguments are " + "equivalent. Only set 'num_cpus_worker_node'." + ) + num_cpus_worker_node = kwargs["num_cpus_per_node"] + warnings.warn( + "'num_cpus_per_node' argument is deprecated, please use " + "'num_cpus_worker_node' argument instead.", + DeprecationWarning, + ) + + if "num_gpus_per_node" in kwargs: + if num_gpus_worker_node is not None: + raise ValueError( + "'num_gpus_per_node' and 'num_gpus_worker_node' arguments are " + "equivalent. Only set 'num_gpus_worker_node'." + ) + num_gpus_worker_node = kwargs["num_gpus_per_node"] + warnings.warn( + "'num_gpus_per_node' argument is deprecated, please use " + "'num_gpus_worker_node' argument instead.", + DeprecationWarning, + ) + + if "object_store_memory_per_node" in kwargs: + if object_store_memory_worker_node is not None: + raise ValueError( + "'object_store_memory_per_node' and 'object_store_memory_worker_node' " + "arguments are equivalent. Only set " + "'object_store_memory_worker_node'." + ) + object_store_memory_worker_node = kwargs["object_store_memory_per_node"] + warnings.warn( + "'object_store_memory_per_node' argument is deprecated, please use " + "'object_store_memory_worker_node' argument instead.", + DeprecationWarning, + ) + + # Environment configurations within the Spark Session that dictate how many cpus + # and gpus to use for each submitted spark task. + num_spark_task_cpus = int(spark.sparkContext.getConf().get("spark.task.cpus", "1")) + + if num_cpus_worker_node is not None and num_cpus_worker_node <= 0: + raise ValueError("Argument `num_cpus_worker_node` value must be > 0.") + + # note: spark.task.resource.gpu.amount config might be fractional value like 0.5 + default_num_spark_task_gpus = float( + spark.sparkContext.getConf().get("spark.task.resource.gpu.amount", "0") + ) + rounded_num_spark_task_gpus = int(default_num_spark_task_gpus) + if default_num_spark_task_gpus > 0: + warn_msg = ( + "You configured 'spark.task.resource.gpu.amount' to " + f"{default_num_spark_task_gpus}," + "we recommend setting this value to 0 so that Spark jobs do not " + "reserve GPU resources, preventing Ray-on-Spark workloads from having the " + "maximum number of GPUs available." + ) + + if is_in_databricks_runtime(): + from ray.util.spark.databricks_hook import ( + get_databricks_display_html_function, + ) + + get_databricks_display_html_function()( + f"{warn_msg}" + ) + else: + _logger.warning(warn_msg) + + if num_gpus_worker_node is not None and num_gpus_worker_node < 0: + raise ValueError("Argument `num_gpus_worker_node` value must be >= 0.") + + def _get_spark_worker_resources(_): + from ray.util.spark.utils import ( + _get_cpu_cores, + _get_num_physical_gpus, + _get_spark_worker_total_physical_memory, + ) + + num_cpus_spark_worker = _get_cpu_cores() + num_gpus_spark_worker = _get_num_physical_gpus() + total_mem_bytes = _get_spark_worker_total_physical_memory() + + return ( + num_cpus_spark_worker, + num_gpus_spark_worker, + total_mem_bytes, + ) + + (num_cpus_spark_worker, num_gpus_spark_worker, spark_worker_mem_bytes,) = ( + spark.sparkContext.parallelize([1], 1) + .map(_get_spark_worker_resources) + .collect()[0] + ) + + if num_cpus_worker_node is not None and num_gpus_worker_node is not None: + if support_stage_scheduling: + using_stage_scheduling = True + res_profile = _create_resource_profile( + num_cpus_worker_node, num_gpus_worker_node + ) + else: + raise ValueError( + "Current spark version does not support stage scheduling, so that " + "you cannot set the argument `num_cpus_worker_node` and " + "`num_gpus_worker_node` values. Without setting the 2 arguments, " + "per-Ray worker node will be assigned with number of " + f"'spark.task.cpus' (equals to {num_spark_task_cpus}) cpu cores " + "and rounded down number of 'spark.task.resource.gpu.amount' " + f"(equals to {rounded_num_spark_task_gpus}) GPUs. To enable spark " + f"stage scheduling, you need to upgrade spark to 3.4 version or use " + "Databricks Runtime 12.x, and you cannot use spark local mode." + ) + elif num_cpus_worker_node is None and num_gpus_worker_node is None: + if support_stage_scheduling: + # Make one Ray worker node using maximum CPU / GPU resources + # of the whole spark worker node, this is the optimal + # configuration. + num_cpus_worker_node = num_cpus_spark_worker + num_gpus_worker_node = num_gpus_spark_worker + using_stage_scheduling = True + res_profile = _create_resource_profile( + num_cpus_worker_node, num_gpus_worker_node + ) + else: + using_stage_scheduling = False + res_profile = None + + num_cpus_worker_node = num_spark_task_cpus + num_gpus_worker_node = rounded_num_spark_task_gpus + else: + raise ValueError( + "'num_cpus_worker_node' and 'num_gpus_worker_node' arguments must be" + "set together or unset together." + ) + + ( + ray_worker_node_heap_mem_bytes, + ray_worker_node_object_store_mem_bytes, + ) = get_avail_mem_per_ray_worker_node( + spark, + heap_memory_worker_node, + object_store_memory_worker_node, + num_cpus_worker_node, + num_gpus_worker_node, + ) + + spark_worker_ray_node_slots = _get_local_ray_node_slots( + num_cpus_spark_worker, + num_gpus_spark_worker, + num_cpus_worker_node, + num_gpus_worker_node, + ) + + spark_executor_memory_bytes = get_configured_spark_executor_memory_bytes(spark) + spark_worker_required_memory_bytes = ( + spark_executor_memory_bytes + + spark_worker_ray_node_slots + * (ray_worker_node_heap_mem_bytes + ray_worker_node_object_store_mem_bytes) + ) + if spark_worker_required_memory_bytes > 0.8 * spark_worker_mem_bytes: + warn_msg = ( + "In each spark worker node, we recommend making the sum of " + "'spark_executor_memory + num_Ray_worker_nodes_per_spark_worker * " + "(memory_worker_node + object_store_memory_worker_node)' to be less than " + "'spark_worker_physical_memory * 0.8', otherwise it might lead to " + "spark worker physical memory exhaustion and Ray task OOM errors." + ) + + if is_in_databricks_runtime(): + from ray.util.spark.databricks_hook import ( + get_databricks_display_html_function, + ) + + get_databricks_display_html_function()( + f"{warn_msg}
" + ) + else: + _logger.warning(warn_msg) + + if "num_worker_nodes" in kwargs: + raise ValueError( + "'num_worker_nodes' argument is removed, please set " + "'max_worker_nodes' and 'min_worker_nodes' argument instead." + ) + + if max_worker_nodes == MAX_NUM_WORKER_NODES: + if min_worker_nodes is not None: + raise ValueError( + "If you set 'max_worker_nodes' to 'MAX_NUM_WORKER_NODES', autoscaling " + "is not supported, so that you cannot set 'min_worker_nodes' argument " + "and 'min_worker_nodes' is automatically set to be equal to " + "'max_worker_nodes'." + ) + + # max_worker_nodes=MAX_NUM_WORKER_NODES represents using all available + # spark task slots + max_worker_nodes = get_max_num_concurrent_tasks(spark.sparkContext, res_profile) + min_worker_nodes = max_worker_nodes + elif max_worker_nodes <= 0: + raise ValueError( + "The value of 'max_worker_nodes' argument must be either a positive " + "integer or 'ray.util.spark.MAX_NUM_WORKER_NODES'." + ) + + if "autoscale" in kwargs: + raise ValueError( + "'autoscale' argument is removed. You can set 'min_worker_nodes' argument " + "to be less than 'max_worker_nodes' to make autoscaling enabled." + ) + + if min_worker_nodes is None: + min_worker_nodes = max_worker_nodes + elif not (0 <= min_worker_nodes <= max_worker_nodes): + raise ValueError( + "The value of 'max_worker_nodes' argument must be an integer >= 0 " + "and <= 'max_worker_nodes'" + ) + + insufficient_resources = [] + + if num_cpus_worker_node < 4: + insufficient_resources.append( + "The provided CPU resources for each ray worker are inadequate to start " + "a ray cluster. Based on the total cpu resources available and the " + "configured task sizing, each ray worker node would start with " + f"{num_cpus_worker_node} CPU cores. This is less than the recommended " + "value of `4` CPUs per worker. On spark version >= 3.4 or Databricks " + "Runtime 12.x, you can set the argument `num_cpus_worker_node` to " + "a value >= 4 to address it, otherwise you need to increase the spark " + "application configuration 'spark.task.cpus' to a minimum of `4` to " + "address it." + ) + + if ray_worker_node_heap_mem_bytes < 10 * 1024 * 1024 * 1024: + insufficient_resources.append( + "The provided memory resources for each ray worker node are inadequate. " + "Based on the total memory available on the spark cluster and the " + "configured task sizing, each ray worker would start with " + f"{ray_worker_node_heap_mem_bytes} bytes heap memory. This is less than " + "the recommended value of 10GB. The ray worker node heap memory size is " + "calculated by " + "(SPARK_WORKER_PHYSICAL_MEMORY / num_local_spark_task_slots * 0.8) - " + "object_store_memory_worker_node. To increase the heap space available, " + "increase the memory in the spark cluster by using instance types with " + "larger memory, or increase number of CPU/GPU per Ray worker node " + "(so it leads to less Ray worker node slots per spark worker node), " + "or apply a lower `object_store_memory_worker_node`." + ) + if insufficient_resources: + if strict_mode: + raise ValueError( + "You are creating ray cluster on spark with strict mode (it can be " + "disabled by setting argument 'strict_mode=False' when calling API " + "'setup_ray_cluster'), strict mode requires the spark cluster config " + "satisfying following criterion: " + "\n".join(insufficient_resources) + ) + else: + _logger.warning("\n".join(insufficient_resources)) + + if num_cpus_head_node is None: + if is_global: + num_cpus_head_node = _get_cpu_cores() + else: + num_cpus_head_node = 0 + else: + if num_cpus_head_node < 0: + raise ValueError( + "Argument `num_cpus_head_node` value must be >= 0. " + f"Current value is {num_cpus_head_node}." + ) + + if num_gpus_head_node is None: + if is_global: + try: + num_gpus_head_node = _get_num_physical_gpus() + except Exception: + num_gpus_head_node = 0 + else: + num_gpus_head_node = 0 + else: + if num_gpus_head_node < 0: + raise ValueError( + "Argument `num_gpus_head_node` value must be >= 0." + f"Current value is {num_gpus_head_node}." + ) + + if ( + num_cpus_head_node == 0 + and num_gpus_head_node == 0 + and object_store_memory_head_node is None + ): + # Because tasks that require CPU or GPU resources are not scheduled to Ray + # head node, and user does not set `object_store_memory_head_node` explicitly, + # limit the heap memory and object store memory allocation to the + # head node, in order to save spark driver memory. + heap_memory_head_node = 1024 * 1024 * 1024 + object_store_memory_head_node = 1024 * 1024 * 1024 + else: + heap_memory_head_node, object_store_memory_head_node = calc_mem_ray_head_node( + heap_memory_head_node, object_store_memory_head_node + ) + + with _active_ray_cluster_rwlock: + cluster = _setup_ray_cluster( + max_worker_nodes=max_worker_nodes, + min_worker_nodes=min_worker_nodes, + num_cpus_worker_node=num_cpus_worker_node, + num_cpus_head_node=num_cpus_head_node, + num_gpus_worker_node=num_gpus_worker_node, + num_gpus_head_node=num_gpus_head_node, + using_stage_scheduling=using_stage_scheduling, + heap_memory_worker_node=ray_worker_node_heap_mem_bytes, + heap_memory_head_node=heap_memory_head_node, + object_store_memory_worker_node=ray_worker_node_object_store_mem_bytes, + object_store_memory_head_node=object_store_memory_head_node, + head_node_options=head_node_options, + worker_node_options=worker_node_options, + ray_temp_root_dir=ray_temp_root_dir, + collect_log_to_path=collect_log_to_path, + autoscale_upscaling_speed=autoscale_upscaling_speed, + autoscale_idle_timeout_minutes=autoscale_idle_timeout_minutes, + is_global=is_global, + ) + # set global _active_ray_cluster to be the + # started cluster. + _active_ray_cluster = cluster + + try: + cluster.wait_until_ready() # NB: this line might raise error. + except Exception as e: + try: + shutdown_ray_cluster() + except Exception: + pass + raise RuntimeError("Launch Ray-on-Spark cluster failed") from e + + head_ip = parse_address(cluster.address)[0] + remote_connection_address = ( + f"ray://{build_address(head_ip, cluster.ray_client_server_port)}" + ) + return cluster.address, remote_connection_address + + +@PublicAPI +def setup_ray_cluster( + *, + max_worker_nodes: int, + min_worker_nodes: Optional[int] = None, + num_cpus_worker_node: Optional[int] = None, + num_cpus_head_node: Optional[int] = None, + num_gpus_worker_node: Optional[int] = None, + num_gpus_head_node: Optional[int] = None, + memory_worker_node: Optional[int] = None, + memory_head_node: Optional[int] = None, + object_store_memory_worker_node: Optional[int] = None, + object_store_memory_head_node: Optional[int] = None, + head_node_options: Optional[Dict] = None, + worker_node_options: Optional[Dict] = None, + ray_temp_root_dir: Optional[str] = None, + strict_mode: bool = False, + collect_log_to_path: Optional[str] = None, + autoscale_upscaling_speed: Optional[float] = 1.0, + autoscale_idle_timeout_minutes: Optional[float] = 1.0, + **kwargs, +) -> Tuple[str, str]: + """ + Set up a ray cluster on the spark cluster by starting a ray head node in the + spark application's driver side node. + After creating the head node, a background spark job is created that + generates an instance of `RayClusterOnSpark` that contains configuration for the + ray cluster that will run on the Spark cluster's worker nodes. + After a ray cluster is set up, "RAY_ADDRESS" environment variable is set to + the cluster address, so you can call `ray.init()` without specifying ray cluster + address to connect to the cluster. To shut down the cluster you can call + `ray.util.spark.shutdown_ray_cluster()`. + Note: If the active ray cluster haven't shut down, you cannot create a new ray + cluster. + + Args: + max_worker_nodes: This argument represents maximum ray worker nodes to start + for the ray cluster. you can + specify the `max_worker_nodes` as `ray.util.spark.MAX_NUM_WORKER_NODES` + represents a ray cluster + configuration that will use all available resources configured for the + spark application. + To create a spark application that is intended to exclusively run a + shared ray cluster in non-scaling, it is recommended to set this argument + to `ray.util.spark.MAX_NUM_WORKER_NODES`. + min_worker_nodes: Minimal number of worker nodes (default `None`), + if "max_worker_nodes" value is equal to "min_worker_nodes" argument, + or "min_worker_nodes" argument value is None, then autoscaling is disabled + and Ray cluster is launched with fixed number "max_worker_nodes" of + Ray worker nodes, otherwise autoscaling is enabled. + num_cpus_worker_node: Number of cpus available to per-ray worker node, if not + provided, if spark stage scheduling is supported, 'num_cpus_head_node' + value equals to number of cpu cores per spark worker node, otherwise + it uses spark application configuration 'spark.task.cpus' instead. + **Limitation** Only spark version >= 3.4 or Databricks Runtime 12.x + supports setting this argument. + num_cpus_head_node: Number of cpus available to Ray head node, if not provide, + if it is global mode Ray cluster, use number of cpu cores in spark driver + node, otherwise use 0 instead. + use 0 instead. Number 0 means tasks requiring CPU resources are not + scheduled to Ray head node. + num_gpus_worker_node: Number of gpus available to per-ray worker node, if not + provided, if spark stage scheduling is supported, 'num_gpus_worker_node' + value equals to number of GPUs per spark worker node, otherwise + it uses rounded down value of spark application configuration + 'spark.task.resource.gpu.amount' instead. + This argument is only available on spark cluster that is configured with + 'gpu' resources. + **Limitation** Only spark version >= 3.4 or Databricks Runtime 12.x + supports setting this argument. + num_gpus_head_node: Number of gpus available to Ray head node, if not provide, + if it is global mode Ray cluster, use number of GPUs in spark driver node, + otherwise use 0 instead. + This argument is only available on spark cluster which spark driver node + has GPUs. + memory_worker_node: Optional[int]: + Heap memory configured for Ray worker node. This is basically setting + `--memory` option when starting Ray node by `ray start` command. + memory_head_node: Optional[int]: + Heap memory configured for Ray head node. This is basically setting + `--memory` option when starting Ray node by `ray start` command. + object_store_memory_worker_node: Object store memory available to per-ray worker + node, but it is capped by + "dev_shm_available_size * 0.8 / num_tasks_per_spark_worker". + The default value equals to + "0.3 * spark_worker_physical_memory * 0.8 / num_tasks_per_spark_worker". + object_store_memory_head_node: Object store memory available to Ray head + node, but it is capped by "dev_shm_available_size * 0.8". + The default value equals to + "0.3 * spark_driver_physical_memory * 0.8". + head_node_options: A dict representing Ray head node extra options, these + options will be passed to `ray start` script. Note you need to convert + `ray start` options key from `--foo-bar` format to `foo_bar` format. + For flag options (e.g. '--disable-usage-stats'), you should set the value + to None in the option dict, like `{"disable_usage_stats": None}`. + Note: Short name options (e.g. '-v') are not supported. + worker_node_options: A dict representing Ray worker node extra options, + these options will be passed to `ray start` script. Note you need to + convert `ray start` options key from `--foo-bar` format to `foo_bar` + format. + For flag options (e.g. '--disable-usage-stats'), you should set the value + to None in the option dict, like `{"disable_usage_stats": None}`. + Note: Short name options (e.g. '-v') are not supported. + ray_temp_root_dir: A local disk path to store the ray temporary data. The + created cluster will create a subdirectory + "ray-{head_port}-{random_suffix}" beneath this path. + strict_mode: Boolean flag to fast-fail initialization of the ray cluster if + the available spark cluster does not have sufficient resources to fulfill + the resource allocation for memory, cpu and gpu. When set to true, if the + requested resources are not available for recommended minimum recommended + functionality, an exception will be raised that details the inadequate + spark cluster configuration settings. If overridden as `False`, + a warning is raised. + collect_log_to_path: If specified, after ray head / worker nodes terminated, + collect their logs to the specified path. On Databricks Runtime, we + recommend you to specify a local path starts with '/dbfs/', because the + path mounts with a centralized storage device and stored data is persisted + after Databricks spark cluster terminated. + autoscale_upscaling_speed: If autoscale enabled, it represents the number of + nodes allowed to be pending as a multiple of the current number of nodes. + The higher the value, the more aggressive upscaling will be. For example, + if this is set to 1.0, the cluster can grow in size by at most 100% at any + time, so if the cluster currently has 20 nodes, at most 20 pending launches + are allowed. The minimum number of pending launches is 5 regardless of + this setting. + Default value is 1.0, minimum value is 1.0 + autoscale_idle_timeout_minutes: If autoscale enabled, it represents the number + of minutes that need to pass before an idle worker node is removed by the + autoscaler. The smaller the value, the more aggressive downscaling will be. + Worker nodes are considered idle when they hold no active tasks, actors, + or referenced objects (either in-memory or spilled to disk). This parameter + does not affect the head node. + Default value is 1.0, minimum value is 0 + Returns: + returns a tuple of (address, remote_connection_address) + "address" is in format of ":" + "remote_connection_address" is in format of + "ray://:", + if your client runs on a machine that also hosts a Ray cluster node locally, + you can connect to the Ray cluster via ``ray.init(address)``, + otherwise you can connect to the Ray cluster via + ``ray.init(remote_connection_address)``. + """ + + return _setup_ray_cluster_internal( + max_worker_nodes=max_worker_nodes, + min_worker_nodes=min_worker_nodes, + num_cpus_worker_node=num_cpus_worker_node, + num_cpus_head_node=num_cpus_head_node, + num_gpus_worker_node=num_gpus_worker_node, + num_gpus_head_node=num_gpus_head_node, + heap_memory_worker_node=memory_worker_node, + heap_memory_head_node=memory_head_node, + object_store_memory_worker_node=object_store_memory_worker_node, + object_store_memory_head_node=object_store_memory_head_node, + head_node_options=head_node_options, + worker_node_options=worker_node_options, + ray_temp_root_dir=ray_temp_root_dir, + strict_mode=strict_mode, + collect_log_to_path=collect_log_to_path, + autoscale_upscaling_speed=autoscale_upscaling_speed, + autoscale_idle_timeout_minutes=autoscale_idle_timeout_minutes, + is_global=False, + **kwargs, + ) + + +@PublicAPI +def setup_global_ray_cluster( + *, + max_worker_nodes: int, + is_blocking: bool = True, + min_worker_nodes: Optional[int] = None, + num_cpus_worker_node: Optional[int] = None, + num_cpus_head_node: Optional[int] = None, + num_gpus_worker_node: Optional[int] = None, + num_gpus_head_node: Optional[int] = None, + memory_worker_node: Optional[int] = None, + memory_head_node: Optional[int] = None, + object_store_memory_worker_node: Optional[int] = None, + object_store_memory_head_node: Optional[int] = None, + head_node_options: Optional[Dict] = None, + worker_node_options: Optional[Dict] = None, + strict_mode: bool = False, + collect_log_to_path: Optional[str] = None, + autoscale_upscaling_speed: Optional[float] = 1.0, + autoscale_idle_timeout_minutes: Optional[float] = 1.0, +): + """ + Set up a global mode cluster. + The global Ray on spark cluster means: + - You can only create one active global Ray on spark cluster at a time. + On databricks cluster, the global Ray cluster can be used by all users, + - as contrast, non-global Ray cluster can only be used by current notebook + user. + - It is up persistently without automatic shutdown. + - On databricks notebook, you can connect to the global cluster by calling + ``ray.init()`` without specifying its address, it will discover the + global cluster automatically if it is up. + + For global mode, the ``ray_temp_root_dir`` argument is not supported. + Global model Ray cluster always use the default Ray temporary directory + path. + + All arguments are the same with ``setup_ray_cluster`` API except that: + - the ``ray_temp_root_dir`` argument is not supported. + Global model Ray cluster always use the default Ray temporary directory + path. + - A new argument "is_blocking" (default ``True``) is added. + If "is_blocking" is True, + then keep the call blocking until it is interrupted. + once the call is interrupted, the global Ray on spark cluster is shut down and + `setup_global_ray_cluster` call terminates. + If "is_blocking" is False, + once Ray cluster setup completes, return immediately. + """ + + cluster_address = _setup_ray_cluster_internal( + max_worker_nodes=max_worker_nodes, + min_worker_nodes=min_worker_nodes, + num_cpus_worker_node=num_cpus_worker_node, + num_cpus_head_node=num_cpus_head_node, + num_gpus_worker_node=num_gpus_worker_node, + num_gpus_head_node=num_gpus_head_node, + heap_memory_worker_node=memory_worker_node, + heap_memory_head_node=memory_head_node, + object_store_memory_worker_node=object_store_memory_worker_node, + object_store_memory_head_node=object_store_memory_head_node, + head_node_options=head_node_options, + worker_node_options=worker_node_options, + ray_temp_root_dir=None, + strict_mode=strict_mode, + collect_log_to_path=collect_log_to_path, + autoscale_upscaling_speed=autoscale_upscaling_speed, + autoscale_idle_timeout_minutes=autoscale_idle_timeout_minutes, + is_global=True, + ) + + if not is_blocking: + return cluster_address + + global _global_ray_cluster_cancel_event + try: + _global_ray_cluster_cancel_event = Event() + # serve forever until user cancel the command. + _global_ray_cluster_cancel_event.wait() + finally: + _global_ray_cluster_cancel_event = None + # once the program is interrupted, + # or the corresponding databricks notebook command is interrupted + # shut down the Ray cluster. + shutdown_ray_cluster() + + +def _start_ray_worker_nodes( + *, + spark_job_server, + spark_job_group_id, + spark_job_group_desc, + num_worker_nodes, + using_stage_scheduling, + ray_head_ip, + ray_head_port, + ray_temp_dir, + num_cpus_per_node, + num_gpus_per_node, + heap_memory_per_node, + object_store_memory_per_node, + worker_node_options, + collect_log_to_path, + node_id, +): + # NB: + # In order to start ray worker nodes on spark cluster worker machines, + # We launch a background spark job: + # 1. Each spark task launches one ray worker node. This design ensures all ray + # worker nodes have the same shape (same cpus / gpus / memory configuration). + # If ray worker nodes have a non-uniform shape, the Ray cluster setup will + # be non-deterministic and could create issues with node sizing. + # 2. A ray worker node is started via the `ray start` CLI. In each spark task, + # a child process is started and will execute a `ray start ...` command in + # blocking mode. + # 3. Each task will acquire a file lock for 10s to ensure that the ray worker + # init will acquire a port connection to the ray head node that does not + # contend with other worker processes on the same Spark worker node. + # 4. When the ray cluster is shutdown, killing ray worker nodes is implemented by + # `sparkContext.cancelJobGroup` to cancel the background spark job, sending a + # SIGKILL signal to all spark tasks. Once the spark tasks are killed, + # `ray_start_node` process detects parent died event then it kills ray + # worker node. + spark = spark_job_server.spark + spark_job_server_port = spark_job_server.server_address[1] + ray_node_custom_env = spark_job_server.ray_node_custom_env + + def ray_cluster_job_mapper(_): + from pyspark.taskcontext import TaskContext + + _worker_logger = logging.getLogger("ray.util.spark.worker") + + context = TaskContext.get() + + ( + worker_port_range_begin, + worker_port_range_end, + ) = _preallocate_ray_worker_port_range() + + # 10001 is used as ray client server port of global mode ray cluster. + ray_worker_node_dashboard_agent_port = get_random_unused_port( + ray_head_ip, min_port=10002, max_port=20000 + ) + ray_worker_node_cmd = [ + sys.executable, + "-m", + "ray.util.spark.start_ray_node", + f"--num-cpus={num_cpus_per_node}", + "--block", + f"--address={build_address(ray_head_ip, ray_head_port)}", + f"--memory={heap_memory_per_node}", + f"--object-store-memory={object_store_memory_per_node}", + f"--min-worker-port={worker_port_range_begin}", + f"--max-worker-port={worker_port_range_end - 1}", + f"--dashboard-agent-listen-port={ray_worker_node_dashboard_agent_port}", + *_convert_ray_node_options(worker_node_options), + ] + if ray_temp_dir is not None: + ray_worker_node_cmd.append(f"--temp-dir={ray_temp_dir}") + + ray_worker_node_extra_envs = { + RAY_ON_SPARK_COLLECT_LOG_TO_PATH: collect_log_to_path or "", + RAY_ON_SPARK_START_RAY_PARENT_PID: str(os.getpid()), + "RAY_ENABLE_WINDOWS_OR_OSX_CLUSTER": "1", + **ray_node_custom_env, + } + + if num_gpus_per_node > 0: + task_resources = context.resources() + + if "gpu" not in task_resources: + raise RuntimeError( + "Couldn't get the gpu id, Please check the GPU resource " + "configuration" + ) + gpu_addr_list = [ + int(addr.strip()) for addr in task_resources["gpu"].addresses + ] + + available_physical_gpus = get_spark_task_assigned_physical_gpus( + gpu_addr_list + ) + ray_worker_node_cmd.append( + f"--num-gpus={len(available_physical_gpus)}", + ) + ray_worker_node_extra_envs["CUDA_VISIBLE_DEVICES"] = ",".join( + [str(gpu_id) for gpu_id in available_physical_gpus] + ) + + _worker_logger.info( + f"Start Ray worker, command: {' '.join(ray_worker_node_cmd)}" + ) + + try: + is_task_reschedule_failure = False + # Check node id availability + response = requests.post( + url=( + f"http://{build_address(ray_head_ip, spark_job_server_port)}" + "/check_node_id_availability" + ), + json={ + "node_id": node_id, + "spark_job_group_id": spark_job_group_id, + }, + ) + if not response.json()["available"]: + # The case happens when a Ray node is down unexpected + # caused by spark worker node down and spark tries to + # reschedule the spark task, so it triggers node + # creation with duplicated node id. + # in this case, finish the spark task immediately + # so spark won't try to reschedule this task + # and Ray autoscaler will trigger a new node creation + # with new node id, and a new spark job will be created + # for holding it. + is_task_reschedule_failure = True + raise RuntimeError( + "Starting Ray worker node twice with the same node id " + "is not allowed." + ) + + # Notify job server the task has been launched. + requests.post( + url=( + f"http://{build_address(ray_head_ip, spark_job_server_port)}" + "/notify_task_launched" + ), + json={ + "spark_job_group_id": spark_job_group_id, + }, + ) + + # Note: + # When a pyspark job cancelled, the UDF python worker process are killed by + # signal "SIGKILL", then `start_ray_node` process will detect the parent + # died event (see `ray.util.spark.start_ray_node.check_parent_alive`) and + # then kill ray worker node process and execute cleanup routine. + exec_cmd( + ray_worker_node_cmd, + synchronous=True, + extra_env=ray_worker_node_extra_envs, + ) + except Exception as e: + # In the following 2 cases, exception is raised: + # (1) + # Starting Ray worker node fails, the `e` will contain detail + # subprocess stdout/stderr output. + # (2) + # In autoscaling mode, when Ray worker node is down, autoscaler will + # try to start new Ray worker node if necessary, + # and it creates a new spark job to launch Ray worker node process, + # note the old spark job will reschedule the failed spark task + # and raise error of "Starting Ray worker node twice with the same + # node id is not allowed". + # + # For either case (1) or case (2), + # to avoid Spark triggers more spark task retries, we swallow + # exception here to make spark the task exit normally. + err_msg = f"Ray worker node process exit, reason: {e}." + _logger.warning(err_msg) + + yield err_msg, is_task_reschedule_failure + + spark.sparkContext.setJobGroup( + spark_job_group_id, + spark_job_group_desc, + ) + + # Starting a normal spark job (not barrier spark job) to run ray worker + # nodes, the design purpose is: + # 1. Using normal spark job, spark tasks can automatically retry + # individually, we don't need to write additional retry logic, But, in + # barrier mode, if one spark task fails, it will cause all other spark + # tasks killed. + # 2. Using normal spark job, we can support failover when a spark worker + # physical machine crashes. (spark will try to re-schedule the spark task + # to other spark worker nodes) + # 3. Using barrier mode job, if the cluster resources does not satisfy + # "idle spark task slots >= argument num_spark_task", then the barrier + # job gets stuck and waits until enough idle task slots available, this + # behavior is not user-friendly, on a shared spark cluster, user is hard + # to estimate how many idle tasks available at a time, But, if using normal + # spark job, it can launch job with less spark tasks (i.e. user will see a + # ray cluster setup with less worker number initially), and when more task + # slots become available, it continues to launch tasks on new available + # slots, and user can see the ray cluster worker number increases when more + # slots available. + job_rdd = spark.sparkContext.parallelize( + list(range(num_worker_nodes)), num_worker_nodes + ) + + if using_stage_scheduling: + resource_profile = _create_resource_profile( + num_cpus_per_node, + num_gpus_per_node, + ) + job_rdd = job_rdd.withResources(resource_profile) + + hook_entry = _create_hook_entry(is_global=(ray_temp_dir is None)) + hook_entry.on_spark_job_created(spark_job_group_id) + + err_msg, is_task_reschedule_failure = job_rdd.mapPartitions( + ray_cluster_job_mapper + ).collect()[0] + if not is_task_reschedule_failure: + spark_job_server.last_worker_error = err_msg + return err_msg + + return None + + +@PublicAPI +def shutdown_ray_cluster() -> None: + """ + Shut down the active ray cluster. + """ + global _active_ray_cluster + + with _active_ray_cluster_rwlock: + if _active_ray_cluster is None: + raise RuntimeError("No active ray cluster to shut down.") + + _active_ray_cluster.shutdown() + _active_ray_cluster = None + + +_global_ray_cluster_cancel_event = None + + +@DeveloperAPI +class AutoscalingCluster: + """Create a ray on spark autoscaling cluster.""" + + def __init__( + self, + head_resources: dict, + worker_node_types: dict, + extra_provider_config: dict, + upscaling_speed: float, + idle_timeout_minutes: float, + ): + """Create the cluster. + + Args: + head_resources: resources of the head node, including CPU. + worker_node_types: autoscaler node types config for worker nodes. + """ + self._head_resources = head_resources.copy() + self._head_resources["NODE_ID_AS_RESOURCE"] = HEAD_NODE_ID + self._config = self._generate_config( + head_resources, + worker_node_types, + extra_provider_config, + upscaling_speed, + idle_timeout_minutes, + ) + + def _generate_config( + self, + head_resources, + worker_node_types, + extra_provider_config, + upscaling_speed, + idle_timeout_minutes, + ): + base_config = yaml.safe_load( + open( + os.path.join( + os.path.dirname(ray.__file__), + "autoscaler/spark/defaults.yaml", + ) + ) + ) + custom_config = copy.deepcopy(base_config) + custom_config["available_node_types"] = worker_node_types + custom_config["available_node_types"]["ray.head.default"] = { + "resources": head_resources, + "node_config": {}, + "max_workers": 0, + } + + custom_config["max_workers"] = sum( + v["max_workers"] for _, v in worker_node_types.items() + ) + + custom_config["provider"].update(extra_provider_config) + + custom_config["upscaling_speed"] = upscaling_speed + custom_config["idle_timeout_minutes"] = idle_timeout_minutes + + return custom_config + + def start( + self, + ray_head_ip, + ray_head_port, + ray_client_server_port, + ray_temp_dir, + dashboard_options, + head_node_options, + collect_log_to_path, + ray_node_custom_env, + ): + """Start the cluster. + + After this call returns, you can connect to the cluster with + ray.init("auto"). + """ + from ray.util.spark.cluster_init import ( + RAY_ON_SPARK_COLLECT_LOG_TO_PATH, + _append_resources_config, + _convert_ray_node_options, + ) + + if ray_temp_dir is not None: + autoscale_config = os.path.join(ray_temp_dir, "autoscaling_config.json") + else: + autoscale_config = os.path.join( + _get_default_ray_tmp_dir(), "autoscaling_config.json" + ) + with open(autoscale_config, "w") as f: + f.write(json.dumps(self._config)) + + ( + worker_port_range_begin, + worker_port_range_end, + ) = _preallocate_ray_worker_port_range() + + ray_head_node_cmd = [ + sys.executable, + "-m", + "ray.util.spark.start_ray_node", + "--block", + "--head", + f"--node-ip-address={ray_head_ip}", + f"--port={ray_head_port}", + f"--ray-client-server-port={ray_client_server_port}", + f"--autoscaling-config={autoscale_config}", + f"--min-worker-port={worker_port_range_begin}", + f"--max-worker-port={worker_port_range_end - 1}", + *dashboard_options, + ] + + if ray_temp_dir is not None: + ray_head_node_cmd.append(f"--temp-dir={ray_temp_dir}") + + if "CPU" in self._head_resources: + ray_head_node_cmd.append( + "--num-cpus={}".format(self._head_resources.pop("CPU")) + ) + if "GPU" in self._head_resources: + ray_head_node_cmd.append( + "--num-gpus={}".format(self._head_resources.pop("GPU")) + ) + if "memory" in self._head_resources: + ray_head_node_cmd.append( + "--memory={}".format(self._head_resources.pop("memory")) + ) + if "object_store_memory" in self._head_resources: + ray_head_node_cmd.append( + "--object-store-memory={}".format( + self._head_resources.pop("object_store_memory") + ) + ) + + head_node_options = _append_resources_config( + head_node_options, self._head_resources + ) + ray_head_node_cmd.extend(_convert_ray_node_options(head_node_options)) + + extra_env = { + "AUTOSCALER_UPDATE_INTERVAL_S": "1", + RAY_ON_SPARK_COLLECT_LOG_TO_PATH: collect_log_to_path or "", + RAY_ON_SPARK_START_RAY_PARENT_PID: str(os.getpid()), + **ray_node_custom_env, + } + + self.ray_head_node_cmd = ray_head_node_cmd + + return _start_ray_head_node( + ray_head_node_cmd, synchronous=False, extra_env=extra_env + ) + + +def _start_ray_head_node(ray_head_node_cmd, synchronous, extra_env): + def preexec_function(): + # Make `start_ray_node` script and Ray node process run + # in a separate group, + # otherwise Ray node will be in the same group of parent process, + # if parent process is a Jupyter notebook kernel, when user + # clicks interrupt cell button, SIGINT signal is sent, then Ray node will + # receive SIGINT signal, and it causes Ray node process dies. + # `start_ray_node` script should also run in a separate group + # because on Databricks Runtime, because if Databricks notebook + # is detached, if the children processes don't exit within 1s, + # they will receive SIGKILL, this behavior makes start_ray_node + # doesn't have enough time to complete cleanup work like removing + # temp directory and collecting logs. + os.setpgrp() + + return exec_cmd( + ray_head_node_cmd, + synchronous=synchronous, + extra_env=extra_env, + preexec_fn=preexec_function, + ) + + +_sigterm_signal_installed = False + + +def _install_sigterm_signal(): + global _sigterm_signal_installed + + if _sigterm_signal_installed: + return + + try: + _origin_sigterm_handler = signal.getsignal(signal.SIGTERM) + + def _sigterm_handler(signum, frame): + try: + shutdown_ray_cluster() + except Exception: + # swallow exception to continue executing the following code in the + # handler + pass + signal.signal( + signal.SIGTERM, _origin_sigterm_handler + ) # Reset to original signal + os.kill( + os.getpid(), signal.SIGTERM + ) # Re-raise the signal to trigger original behavior + + signal.signal(signal.SIGTERM, _sigterm_handler) + _sigterm_signal_installed = True + except Exception: + _logger.warning("Install Ray-on-Spark SIGTERM handler failed.") diff --git a/lib/python3.12/site-packages/ray/util/spark/databricks_hook.py b/lib/python3.12/site-packages/ray/util/spark/databricks_hook.py new file mode 100644 index 0000000000000000000000000000000000000000..491f35c1419f2440e5d79eb4abfb88b7d70c14ce --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/spark/databricks_hook.py @@ -0,0 +1,221 @@ +import logging +import os +import threading +import time + +from .start_hook_base import RayOnSparkStartHook +from .utils import get_spark_session + +_logger = logging.getLogger(__name__) + +DATABRICKS_HOST = "DATABRICKS_HOST" +DATABRICKS_TOKEN = "DATABRICKS_TOKEN" +DATABRICKS_CLIENT_ID = "DATABRICKS_CLIENT_ID" +DATABRICKS_CLIENT_SECRET = "DATABRICKS_CLIENT_SECRET" + + +def verify_databricks_auth_env(): + return (DATABRICKS_HOST in os.environ and DATABRICKS_TOKEN in os.environ) or ( + DATABRICKS_HOST in os.environ + and DATABRICKS_CLIENT_ID in os.environ + and DATABRICKS_CLIENT_SECRET in os.environ + ) + + +def get_databricks_function(func_name): + import IPython + + ip_shell = IPython.get_ipython() + if ip_shell is None: + raise RuntimeError("No IPython environment.") + return ip_shell.ns_table["user_global"][func_name] + + +def get_databricks_display_html_function(): + return get_databricks_function("displayHTML") + + +def get_db_entry_point(): + """ + Return databricks entry_point instance, it is for calling some + internal API in databricks runtime + """ + from dbruntime import UserNamespaceInitializer + + user_namespace_initializer = UserNamespaceInitializer.getOrCreate() + return user_namespace_initializer.get_spark_entry_point() + + +def display_databricks_driver_proxy_url(spark_context, port, title): + """ + This helper function create a proxy URL for databricks driver webapp forwarding. + In databricks runtime, user does not have permission to directly access web + service binding on driver machine port, but user can visit it by a proxy URL with + following format: "/driver-proxy/o/{orgId}/{clusterId}/{port}/". + """ + driverLocal = spark_context._jvm.com.databricks.backend.daemon.driver.DriverLocal + commandContextTags = driverLocal.commandContext().get().toStringMap().apply("tags") + orgId = commandContextTags.apply("orgId") + clusterId = commandContextTags.apply("clusterId") + + proxy_link = f"/driver-proxy/o/{orgId}/{clusterId}/{port}/" + proxy_url = f"https://dbc-dp-{orgId}.cloud.databricks.com{proxy_link}" + + print("To monitor and debug Ray from Databricks, view the dashboard at ") + print(f" {proxy_url}") + + get_databricks_display_html_function()( + f""" + + """ + ) + + +DATABRICKS_AUTO_SHUTDOWN_POLL_INTERVAL_SECONDS = 3 +DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES = ( + "DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES" +) + + +_DATABRICKS_DEFAULT_TMP_ROOT_DIR = "/local_disk0/tmp" + + +class DefaultDatabricksRayOnSparkStartHook(RayOnSparkStartHook): + def get_default_temp_root_dir(self): + return _DATABRICKS_DEFAULT_TMP_ROOT_DIR + + def on_ray_dashboard_created(self, port): + display_databricks_driver_proxy_url( + get_spark_session().sparkContext, port, "Ray Cluster Dashboard" + ) + + def on_cluster_created(self, ray_cluster_handler): + db_api_entry = get_db_entry_point() + + if self.is_global: + # Disable auto shutdown if + # 1) autoscaling enabled + # because in autoscaling mode, background spark job will be killed + # automatically when ray cluster is idle. + # 2) global mode cluster + # Because global mode cluster is designed to keep running until + # user request to shut down it, and global mode cluster is shared + # by other users, the code here cannot track usage from other users + # so that we don't know whether it is safe to shut down the global + # cluster automatically. + auto_shutdown_minutes = 0 + else: + auto_shutdown_minutes = float( + os.environ.get(DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES, "30") + ) + if auto_shutdown_minutes == 0: + _logger.info( + "The Ray cluster will keep running until you manually detach the " + "Databricks notebook or call " + "`ray.util.spark.shutdown_ray_cluster()`." + ) + return + if auto_shutdown_minutes < 0: + raise ValueError( + "You must set " + f"'{DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES}' " + "to a value >= 0." + ) + + try: + db_api_entry.getIdleTimeMillisSinceLastNotebookExecution() + except Exception: + _logger.warning( + "Failed to retrieve idle time since last notebook execution, " + "so that we cannot automatically shut down Ray cluster when " + "Databricks notebook is inactive for the specified minutes. " + "You need to manually detach Databricks notebook " + "or call `ray.util.spark.shutdown_ray_cluster()` to shut down " + "Ray cluster on spark." + ) + return + + _logger.info( + "The Ray cluster will be shut down automatically if you don't run " + "commands on the Databricks notebook for " + f"{auto_shutdown_minutes} minutes. You can change the " + "auto-shutdown minutes by setting " + f"'{DATABRICKS_RAY_ON_SPARK_AUTOSHUTDOWN_MINUTES}' environment " + "variable, setting it to 0 means that the Ray cluster keeps running " + "until you manually call `ray.util.spark.shutdown_ray_cluster()` or " + "detach Databricks notebook." + ) + + def auto_shutdown_watcher(): + auto_shutdown_millis = auto_shutdown_minutes * 60 * 1000 + while True: + if ray_cluster_handler.is_shutdown: + # The cluster is shut down. The watcher thread exits. + return + + idle_time = db_api_entry.getIdleTimeMillisSinceLastNotebookExecution() + + if idle_time > auto_shutdown_millis: + from ray.util.spark import cluster_init + + with cluster_init._active_ray_cluster_rwlock: + if ray_cluster_handler is cluster_init._active_ray_cluster: + cluster_init.shutdown_ray_cluster() + return + + time.sleep(DATABRICKS_AUTO_SHUTDOWN_POLL_INTERVAL_SECONDS) + + threading.Thread(target=auto_shutdown_watcher, daemon=True).start() + + def on_spark_job_created(self, job_group_id): + db_api_entry = get_db_entry_point() + db_api_entry.registerBackgroundSparkJobGroup(job_group_id) + + def custom_environment_variables(self): + conf = { + **super().custom_environment_variables(), + # Hardcode `GLOO_SOCKET_IFNAME` to `eth0` for Databricks runtime. + # Torch on DBR does not reliably detect the correct interface to use, + # and ends up selecting the loopback interface, breaking cross-node + # commnication. + "GLOO_SOCKET_IFNAME": "eth0", + # 'DISABLE_MLFLOW_INTEGRATION' is the environmental variable to disable + # huggingface transformers MLflow integration, + # it doesn't work well in Databricks runtime, + # So disable it by default. + "DISABLE_MLFLOW_INTEGRATION": "TRUE", + } + + if verify_databricks_auth_env(): + conf[DATABRICKS_HOST] = os.environ[DATABRICKS_HOST] + if DATABRICKS_TOKEN in os.environ: + # PAT auth + conf[DATABRICKS_TOKEN] = os.environ[DATABRICKS_TOKEN] + else: + # OAuth + conf[DATABRICKS_CLIENT_ID] = os.environ[DATABRICKS_CLIENT_ID] + conf[DATABRICKS_CLIENT_SECRET] = os.environ[DATABRICKS_CLIENT_SECRET] + else: + warn_msg = ( + "MLflow support is not correctly configured within Ray tasks." + "To enable MLflow integration, " + "you need to set environmental variables DATABRICKS_HOST + " + "DATABRICKS_TOKEN, or set environmental variables " + "DATABRICKS_HOST + DATABRICKS_CLIENT_ID + DATABRICKS_CLIENT_SECRET " + "before calling `ray.util.spark.setup_ray_cluster`, these variables " + "are used to set up authentication with Databricks MLflow " + "service. For details, you can refer to Databricks documentation at " + "" + "Databricks PAT auth or " + "Databricks OAuth." + ) + get_databricks_display_html_function()( + f"{warn_msg}
" + ) + + return conf diff --git a/lib/python3.12/site-packages/ray/util/spark/start_hook_base.py b/lib/python3.12/site-packages/ray/util/spark/start_hook_base.py new file mode 100644 index 0000000000000000000000000000000000000000..d51dbec3a02bb5616b3b77646793dbcbf333c385 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/spark/start_hook_base.py @@ -0,0 +1,18 @@ +class RayOnSparkStartHook: + def __init__(self, is_global): + self.is_global = is_global + + def get_default_temp_root_dir(self): + return "/tmp" + + def on_ray_dashboard_created(self, port): + pass + + def on_cluster_created(self, ray_cluster_handler): + pass + + def on_spark_job_created(self, job_group_id): + pass + + def custom_environment_variables(self): + return {} diff --git a/lib/python3.12/site-packages/ray/util/spark/start_ray_node.py b/lib/python3.12/site-packages/ray/util/spark/start_ray_node.py new file mode 100644 index 0000000000000000000000000000000000000000..e03c99d74e55abed50c0033ba21397a5a7f509f1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/spark/start_ray_node.py @@ -0,0 +1,211 @@ +import fcntl +import logging +import os.path +import shutil +import signal +import socket +import subprocess +import sys +import threading +import time + +from ray._private.ray_process_reaper import SIGTERM_GRACE_PERIOD_SECONDS +from ray.util.spark.cluster_init import ( + RAY_ON_SPARK_COLLECT_LOG_TO_PATH, + RAY_ON_SPARK_START_RAY_PARENT_PID, +) + +# Spark on ray implementation does not directly invoke `ray start ...` script to create +# ray node subprocess, instead, it creates a subprocess to run this +# `ray.util.spark.start_ray_node` module, and in this module it invokes `ray start ...` +# script to start ray node, the purpose of `start_ray_node` module is to set up a +# exit handler for cleaning ray temp directory when ray node exits. +# When spark driver python process dies, or spark python worker dies, because +# `start_ray_node` starts a daemon thread of `check_parent_alive`, it will detect +# parent process died event and then trigger cleanup work. + + +_logger = logging.getLogger(__name__) + + +if __name__ == "__main__": + arg_list = sys.argv[1:] + + collect_log_to_path = os.environ[RAY_ON_SPARK_COLLECT_LOG_TO_PATH] + + temp_dir_arg_prefix = "--temp-dir=" + temp_dir = None + + for arg in arg_list: + if arg.startswith(temp_dir_arg_prefix): + temp_dir = arg[len(temp_dir_arg_prefix) :] + + if temp_dir is not None: + temp_dir = os.path.normpath(temp_dir) + else: + # This case is for global mode Ray on spark cluster + from ray.util.spark.cluster_init import _get_default_ray_tmp_dir + + temp_dir = _get_default_ray_tmp_dir() + + # Multiple Ray nodes might be launched in the same machine, + # so set `exist_ok` to True + os.makedirs(temp_dir, exist_ok=True) + + ray_cli_cmd = "ray" + lock_file = temp_dir + ".lock" + + lock_fd = os.open(lock_file, os.O_RDWR | os.O_CREAT | os.O_TRUNC) + + # Mutilple ray nodes might start on the same machine, and they are using the + # same temp directory, adding a shared lock representing current ray node is + # using the temp directory. + fcntl.flock(lock_fd, fcntl.LOCK_SH) + + process = subprocess.Popen( + # 'ray start ...' command uses python that is set by + # Shebang #! ..., the Shebang line is hardcoded in ray script, + # it can't be changed to other python executable path. + # to enforce using current python executable, + # turn the subprocess command to + # '`sys.executable` `which ray` start ...' + [sys.executable, shutil.which(ray_cli_cmd), "start", *arg_list], + text=True, + ) + + exit_handler_executed = False + sigterm_handler_executed = False + ON_EXIT_HANDLER_WAIT_TIME = 3 + + def on_exit_handler(): + global exit_handler_executed + + if exit_handler_executed: + # wait for exit_handler execution completed in other threads. + time.sleep(ON_EXIT_HANDLER_WAIT_TIME) + return + + exit_handler_executed = True + + try: + # Wait for a while to ensure the children processes of the ray node all + # exited. + time.sleep(SIGTERM_GRACE_PERIOD_SECONDS + 0.5) + + if process.poll() is None: + # "ray start ..." command process is still alive. Force to kill it. + process.kill() + + # Release the shared lock, representing current ray node does not use the + # temp dir. + fcntl.flock(lock_fd, fcntl.LOCK_UN) + + try: + # acquiring exclusive lock to ensure copy logs and removing dir safely. + fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + lock_acquired = True + except BlockingIOError: + # The file has active shared lock or exclusive lock, representing there + # are other ray nodes running, or other node running cleanup temp-dir + # routine. skip cleaning temp-dir, and skip copy logs to destination + # directory as well. + lock_acquired = False + + if lock_acquired: + # This is the final terminated ray node on current spark worker, + # start copy logs (including all local ray nodes logs) to destination. + if collect_log_to_path: + try: + log_dir_prefix = os.path.basename(temp_dir) + if log_dir_prefix == "ray": + # global mode cluster case, append a timestamp to it to + # avoid name conflict with last Ray global cluster log dir. + log_dir_prefix = ( + log_dir_prefix + f"-global-{int(time.time())}" + ) + base_dir = os.path.join( + collect_log_to_path, log_dir_prefix + "-logs" + ) + # Note: multiple Ray node launcher process might + # execute this line code, so we set exist_ok=True here. + os.makedirs(base_dir, exist_ok=True) + copy_log_dest_path = os.path.join( + base_dir, + socket.gethostname(), + ) + ray_session_dir = os.readlink( + os.path.join(temp_dir, "session_latest") + ) + shutil.copytree( + os.path.join(ray_session_dir, "logs"), + copy_log_dest_path, + ) + except Exception as e: + _logger.warning( + "Collect logs to destination directory failed, " + f"error: {repr(e)}." + ) + + # Start cleaning the temp-dir, + shutil.rmtree(temp_dir, ignore_errors=True) + except Exception: + # swallow any exception. + pass + finally: + fcntl.flock(lock_fd, fcntl.LOCK_UN) + os.close(lock_fd) + + def check_parent_alive() -> None: + orig_parent_pid = int(os.environ[RAY_ON_SPARK_START_RAY_PARENT_PID]) + while True: + time.sleep(0.5) + if os.getppid() != orig_parent_pid: + # Note raising SIGTERM signal in a background thread + # doesn't work + sigterm_handler() + break + + threading.Thread(target=check_parent_alive, daemon=True).start() + + try: + + def sighup_handler(*args): + pass + + # When spark application is terminated, this process will receive + # SIGHUP (comes from pyspark application termination). + # Ignore the SIGHUP signal, because in this case, + # `check_parent_alive` will capture parent process died event + # and execute killing node and cleanup routine + # but if we enable default SIGHUP handler, it will kill + # the process immediately and it causes `check_parent_alive` + # have no time to exeucte cleanup routine. + signal.signal(signal.SIGHUP, sighup_handler) + + def sigterm_handler(*args): + global sigterm_handler_executed + if not sigterm_handler_executed: + sigterm_handler_executed = True + process.terminate() + on_exit_handler() + else: + # wait for exit_handler execution completed in other threads. + time.sleep(ON_EXIT_HANDLER_WAIT_TIME) + # Sigterm exit code is 143. + os._exit(143) + + signal.signal(signal.SIGTERM, sigterm_handler) + while True: + try: + ret_code = process.wait() + break + except KeyboardInterrupt: + # Jupyter notebook interrupt button triggers SIGINT signal and + # `start_ray_node` (subprocess) will receive SIGINT signal and it + # causes KeyboardInterrupt exception being raised. + pass + on_exit_handler() + sys.exit(ret_code) + except Exception: + on_exit_handler() + raise diff --git a/lib/python3.12/site-packages/ray/util/spark/utils.py b/lib/python3.12/site-packages/ray/util/spark/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..2d47c525d9316a9d45840a54411a68da074a6f95 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/spark/utils.py @@ -0,0 +1,528 @@ +import collections +import logging +import os +import random +import shutil +import subprocess +import sys +import threading +import time + +from ray._common.network_utils import is_ipv6 + +_logger = logging.getLogger("ray.util.spark.utils") + + +def is_in_databricks_runtime(): + return "DATABRICKS_RUNTIME_VERSION" in os.environ + + +def gen_cmd_exec_failure_msg(cmd, return_code, tail_output_deque): + cmd_str = " ".join(cmd) + tail_output = "".join(tail_output_deque) + return ( + f"Command {cmd_str} failed with return code {return_code}, tail output are " + f"included below.\n{tail_output}\n" + ) + + +def get_configured_spark_executor_memory_bytes(spark): + value_str = spark.conf.get("spark.executor.memory", "1g").lower() + value_num = int(value_str[:-1]) + value_unit = value_str[-1] + unit_map = { + "k": 1024, + "m": 1024 * 1024, + "g": 1024 * 1024 * 1024, + "t": 1024 * 1024 * 1024 * 1024, + } + return value_num * unit_map[value_unit] + + +def exec_cmd( + cmd, + *, + extra_env=None, + synchronous=True, + **kwargs, +): + """ + A convenience wrapper of `subprocess.Popen` for running a command from a Python + script. + If `synchronous` is True, wait until the process terminated and if subprocess + return code is not 0, raise error containing last 100 lines output. + If `synchronous` is False, return an `Popen` instance and a deque instance holding + tail outputs. + The subprocess stdout / stderr output will be streamly redirected to current + process stdout. + """ + illegal_kwargs = set(kwargs.keys()).intersection({"text", "stdout", "stderr"}) + if illegal_kwargs: + raise ValueError(f"`kwargs` cannot contain {list(illegal_kwargs)}") + + env = kwargs.pop("env", None) + if extra_env is not None and env is not None: + raise ValueError("`extra_env` and `env` cannot be used at the same time") + + env = env if extra_env is None else {**os.environ, **extra_env} + + process = subprocess.Popen( + cmd, + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + **kwargs, + ) + + tail_output_deque = collections.deque(maxlen=100) + + def redirect_log_thread_fn(): + for line in process.stdout: + # collect tail logs by `tail_output_deque` + tail_output_deque.append(line) + + # redirect to stdout. + sys.stdout.write(line) + + threading.Thread(target=redirect_log_thread_fn, args=()).start() + + if not synchronous: + return process, tail_output_deque + + return_code = process.wait() + if return_code != 0: + raise RuntimeError( + gen_cmd_exec_failure_msg(cmd, return_code, tail_output_deque) + ) + + +def is_port_in_use(host, port): + import socket + from contextlib import closing + + with closing( + socket.socket( + socket.AF_INET6 if is_ipv6(host) else socket.AF_INET, socket.SOCK_STREAM + ) + ) as sock: + return sock.connect_ex((host, port)) == 0 + + +def _wait_service_up(host, port, timeout): + beg_time = time.time() + + while time.time() - beg_time < timeout: + if is_port_in_use(host, port): + return True + time.sleep(1) + + return False + + +def get_random_unused_port( + host, min_port=1024, max_port=65535, max_retries=100, exclude_list=None +): + """ + Get random unused port. + """ + # Use true random generator + rng = random.SystemRandom() + + exclude_list = exclude_list or [] + for _ in range(max_retries): + port = rng.randint(min_port, max_port) + if port in exclude_list: + continue + if not is_port_in_use(host, port): + return port + raise RuntimeError( + f"Get available port between range {min_port} and {max_port} failed." + ) + + +def get_spark_session(): + from pyspark.sql import SparkSession + + spark_session = SparkSession.getActiveSession() + if spark_session is None: + raise RuntimeError( + "Spark session haven't been initiated yet. Please use " + "`SparkSession.builder` to create a spark session and connect to a spark " + "cluster." + ) + return spark_session + + +def get_spark_application_driver_host(spark): + return spark.conf.get("spark.driver.host") + + +def get_max_num_concurrent_tasks(spark_context, resource_profile): + """Gets the current max number of concurrent tasks.""" + # pylint: disable=protected-access= + ssc = spark_context._jsc.sc() + if resource_profile is not None: + + def dummpy_mapper(_): + pass + + # Runs a dummy spark job to register the `res_profile` + spark_context.parallelize([1], 1).withResources(resource_profile).map( + dummpy_mapper + ).collect() + + return ssc.maxNumConcurrentTasks(resource_profile._java_resource_profile) + else: + return ssc.maxNumConcurrentTasks( + ssc.resourceProfileManager().defaultResourceProfile() + ) + + +def _get_spark_worker_total_physical_memory(): + import psutil + + if RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES in os.environ: + return int(os.environ[RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES]) + return psutil.virtual_memory().total + + +def _get_spark_worker_total_shared_memory(): + import shutil + + if RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES in os.environ: + return int(os.environ[RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES]) + + return shutil.disk_usage("/dev/shm").total + + +# The maximum proportion for Ray worker node object store memory size +_RAY_ON_SPARK_MAX_OBJECT_STORE_MEMORY_PROPORTION = 0.8 + +# The buffer offset for calculating Ray node memory. +_RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET = 0.8 + + +def calc_mem_ray_head_node(configured_heap_memory_bytes, configured_object_store_bytes): + import shutil + + import psutil + + if RAY_ON_SPARK_DRIVER_PHYSICAL_MEMORY_BYTES in os.environ: + available_physical_mem = int( + os.environ[RAY_ON_SPARK_DRIVER_PHYSICAL_MEMORY_BYTES] + ) + else: + available_physical_mem = psutil.virtual_memory().total + + available_physical_mem = ( + available_physical_mem * _RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET + ) + + if RAY_ON_SPARK_DRIVER_SHARED_MEMORY_BYTES in os.environ: + available_shared_mem = int(os.environ[RAY_ON_SPARK_DRIVER_SHARED_MEMORY_BYTES]) + else: + available_shared_mem = shutil.disk_usage("/dev/shm").total + + available_shared_mem = ( + available_shared_mem * _RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET + ) + + heap_mem_bytes, object_store_bytes, warning_msg = _calc_mem_per_ray_node( + available_physical_mem, + available_shared_mem, + configured_heap_memory_bytes, + configured_object_store_bytes, + ) + + if warning_msg is not None: + _logger.warning(warning_msg) + + return heap_mem_bytes, object_store_bytes + + +def _calc_mem_per_ray_worker_node( + num_task_slots, + physical_mem_bytes, + shared_mem_bytes, + configured_heap_memory_bytes, + configured_object_store_bytes, +): + available_physical_mem_per_node = int( + physical_mem_bytes / num_task_slots * _RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET + ) + available_shared_mem_per_node = int( + shared_mem_bytes / num_task_slots * _RAY_ON_SPARK_NODE_MEMORY_BUFFER_OFFSET + ) + return _calc_mem_per_ray_node( + available_physical_mem_per_node, + available_shared_mem_per_node, + configured_heap_memory_bytes, + configured_object_store_bytes, + ) + + +def _calc_mem_per_ray_node( + available_physical_mem_per_node, + available_shared_mem_per_node, + configured_heap_memory_bytes, + configured_object_store_bytes, +): + from ray._private.ray_constants import ( + DEFAULT_OBJECT_STORE_MEMORY_PROPORTION, + OBJECT_STORE_MINIMUM_MEMORY_BYTES, + ) + + warning_msg = None + + object_store_bytes = configured_object_store_bytes or ( + available_physical_mem_per_node * DEFAULT_OBJECT_STORE_MEMORY_PROPORTION + ) + + # If allow Ray using slow storage oas object store, + # we don't need to cap object store size by /dev/shm capacity + if not os.environ.get("RAY_OBJECT_STORE_ALLOW_SLOW_STORAGE"): + if object_store_bytes > available_shared_mem_per_node: + object_store_bytes = available_shared_mem_per_node + + object_store_bytes_upper_bound = ( + available_physical_mem_per_node + * _RAY_ON_SPARK_MAX_OBJECT_STORE_MEMORY_PROPORTION + ) + + if object_store_bytes > object_store_bytes_upper_bound: + object_store_bytes = object_store_bytes_upper_bound + warning_msg = ( + "Your configured `object_store_memory_per_node` value " + "is too high and it is capped by 80% of per-Ray node " + "allocated memory." + ) + + if object_store_bytes < OBJECT_STORE_MINIMUM_MEMORY_BYTES: + if object_store_bytes == available_shared_mem_per_node: + warning_msg = ( + "Your operating system is configured with too small /dev/shm " + "size, so `object_store_memory_worker_node` value is configured " + f"to minimal size ({OBJECT_STORE_MINIMUM_MEMORY_BYTES} bytes)," + f"Please increase system /dev/shm size." + ) + else: + warning_msg = ( + "You configured too small Ray node object store memory size, " + "so `object_store_memory_worker_node` value is configured " + f"to minimal size ({OBJECT_STORE_MINIMUM_MEMORY_BYTES} bytes)," + "Please increase 'object_store_memory_worker_node' argument value." + ) + + object_store_bytes = OBJECT_STORE_MINIMUM_MEMORY_BYTES + + object_store_bytes = int(object_store_bytes) + + if configured_heap_memory_bytes is None: + heap_mem_bytes = int(available_physical_mem_per_node - object_store_bytes) + else: + heap_mem_bytes = int(configured_heap_memory_bytes) + + return heap_mem_bytes, object_store_bytes, warning_msg + + +# User can manually set these environment variables +# if ray on spark code accessing corresponding information failed. +# Note these environment variables must be set in spark executor side, +# you should set them via setting spark config of +# `spark.executorEnv.[EnvironmentVariableName]` +RAY_ON_SPARK_WORKER_CPU_CORES = "RAY_ON_SPARK_WORKER_CPU_CORES" +RAY_ON_SPARK_WORKER_GPU_NUM = "RAY_ON_SPARK_WORKER_GPU_NUM" +RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES = "RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES" +RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES = "RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES" + +# User can manually set these environment variables on spark driver node +# if ray on spark code accessing corresponding information failed. +RAY_ON_SPARK_DRIVER_PHYSICAL_MEMORY_BYTES = "RAY_ON_SPARK_DRIVER_PHYSICAL_MEMORY_BYTES" +RAY_ON_SPARK_DRIVER_SHARED_MEMORY_BYTES = "RAY_ON_SPARK_DRIVER_SHARED_MEMORY_BYTES" + + +def _get_cpu_cores(): + import multiprocessing + + if RAY_ON_SPARK_WORKER_CPU_CORES in os.environ: + # In some cases, spark standalone cluster might configure virtual cpu cores + # for spark worker that different with number of physical cpu cores, + # but we cannot easily get the virtual cpu cores configured for spark + # worker, as a workaround, we provide an environmental variable config + # `RAY_ON_SPARK_WORKER_CPU_CORES` for user. + return int(os.environ[RAY_ON_SPARK_WORKER_CPU_CORES]) + + return multiprocessing.cpu_count() + + +def _get_num_physical_gpus(): + if RAY_ON_SPARK_WORKER_GPU_NUM in os.environ: + # In some cases, spark standalone cluster might configure part of physical + # GPUs for spark worker, + # but we cannot easily get related configuration, + # as a workaround, we provide an environmental variable config + # `RAY_ON_SPARK_WORKER_CPU_CORES` for user. + return int(os.environ[RAY_ON_SPARK_WORKER_GPU_NUM]) + + if shutil.which("nvidia-smi") is None: + # GPU driver is not installed. + return 0 + try: + completed_proc = subprocess.run( + "nvidia-smi --query-gpu=name --format=csv,noheader", + shell=True, + check=True, + text=True, + capture_output=True, + ) + return len(completed_proc.stdout.strip().split("\n")) + except Exception as e: + _logger.info( + "'nvidia-smi --query-gpu=name --format=csv,noheader' command execution " + f"failed, error: {repr(e)}" + ) + return 0 + + +def _get_local_ray_node_slots( + num_cpus, + num_gpus, + num_cpus_per_node, + num_gpus_per_node, +): + if num_cpus_per_node > num_cpus: + raise ValueError( + "cpu number per Ray worker node should be <= spark worker node CPU cores, " + f"you set cpu number per Ray worker node to {num_cpus_per_node} but " + f"spark worker node CPU core number is {num_cpus}." + ) + num_ray_node_slots = num_cpus // num_cpus_per_node + + if num_gpus_per_node > 0: + if num_gpus_per_node > num_gpus: + raise ValueError( + "gpu number per Ray worker node should be <= spark worker node " + "GPU number, you set GPU devices number per Ray worker node to " + f"{num_gpus_per_node} but spark worker node GPU devices number " + f"is {num_gpus}." + ) + if num_ray_node_slots > num_gpus // num_gpus_per_node: + num_ray_node_slots = num_gpus // num_gpus_per_node + + return num_ray_node_slots + + +def _get_avail_mem_per_ray_worker_node( + num_cpus_per_node, + num_gpus_per_node, + heap_memory_per_node, + object_store_memory_per_node, +): + """ + Returns tuple of ( + ray_worker_node_heap_mem_bytes, + ray_worker_node_object_store_bytes, + error_message, # always None + warning_message, + ) + """ + num_cpus = _get_cpu_cores() + if num_gpus_per_node > 0: + num_gpus = _get_num_physical_gpus() + else: + num_gpus = 0 + + num_ray_node_slots = _get_local_ray_node_slots( + num_cpus, num_gpus, num_cpus_per_node, num_gpus_per_node + ) + + physical_mem_bytes = _get_spark_worker_total_physical_memory() + shared_mem_bytes = _get_spark_worker_total_shared_memory() + + ( + ray_worker_node_heap_mem_bytes, + ray_worker_node_object_store_bytes, + warning_msg, + ) = _calc_mem_per_ray_worker_node( + num_ray_node_slots, + physical_mem_bytes, + shared_mem_bytes, + heap_memory_per_node, + object_store_memory_per_node, + ) + return ( + ray_worker_node_heap_mem_bytes, + ray_worker_node_object_store_bytes, + None, + warning_msg, + ) + + +def get_avail_mem_per_ray_worker_node( + spark, + heap_memory_per_node, + object_store_memory_per_node, + num_cpus_per_node, + num_gpus_per_node, +): + """ + Return the available heap memory and object store memory for each ray worker, + and error / warning message if it has. + Return value is a tuple of + (ray_worker_node_heap_mem_bytes, ray_worker_node_object_store_bytes, + error_message, warning_message) + NB: We have one ray node per spark task. + """ + + def mapper(_): + try: + return _get_avail_mem_per_ray_worker_node( + num_cpus_per_node, + num_gpus_per_node, + heap_memory_per_node, + object_store_memory_per_node, + ) + except Exception as e: + import traceback + + trace_msg = "\n".join(traceback.format_tb(e.__traceback__)) + return -1, -1, repr(e) + trace_msg, None + + # Running memory inference routine on spark executor side since the spark worker + # nodes may have a different machine configuration compared to the spark driver + # node. + ( + inferred_ray_worker_node_heap_mem_bytes, + inferred_ray_worker_node_object_store_bytes, + err, + warning_msg, + ) = ( + spark.sparkContext.parallelize([1], 1).map(mapper).collect()[0] + ) + + if err is not None: + raise RuntimeError( + f"Inferring ray worker node available memory failed, error: {err}. " + "You can bypass this error by setting following spark configs: " + "spark.executorEnv.RAY_ON_SPARK_WORKER_CPU_CORES, " + "spark.executorEnv.RAY_ON_SPARK_WORKER_GPU_NUM, " + "spark.executorEnv.RAY_ON_SPARK_WORKER_PHYSICAL_MEMORY_BYTES, " + "spark.executorEnv.RAY_ON_SPARK_WORKER_SHARED_MEMORY_BYTES." + ) + if warning_msg is not None: + _logger.warning(warning_msg) + return ( + inferred_ray_worker_node_heap_mem_bytes, + inferred_ray_worker_node_object_store_bytes, + ) + + +def get_spark_task_assigned_physical_gpus(gpu_addr_list): + if "CUDA_VISIBLE_DEVICES" in os.environ: + visible_cuda_dev_list = [ + int(dev.strip()) for dev in os.environ["CUDA_VISIBLE_DEVICES"].split(",") + ] + return [visible_cuda_dev_list[addr] for addr in gpu_addr_list] + else: + return gpu_addr_list diff --git a/lib/python3.12/site-packages/ray/util/timer.py b/lib/python3.12/site-packages/ray/util/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..2f36aef155ea4dfaf4c693642b4bcf4e1054d046 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/timer.py @@ -0,0 +1,65 @@ +import time + + +class _Timer: + """A running stat for conveniently logging the duration of a code block. + + Example: + wait_timer = TimerStat() + with wait_timer: + ray.wait(...) + + Note that this class is *not* thread-safe. + """ + + def __init__(self, window_size=10): + self._window_size = window_size + self._samples = [] + self._units_processed = [] + self._start_time = None + self._total_time = 0.0 + self.count = 0 + + def __enter__(self): + assert self._start_time is None, "concurrent updates not supported" + self._start_time = time.time() + + def __exit__(self, exc_type, exc_value, tb): + assert self._start_time is not None + time_delta = time.time() - self._start_time + self.push(time_delta) + self._start_time = None + + def push(self, time_delta): + self._samples.append(time_delta) + if len(self._samples) > self._window_size: + self._samples.pop(0) + self.count += 1 + self._total_time += time_delta + + def push_units_processed(self, n): + self._units_processed.append(n) + if len(self._units_processed) > self._window_size: + self._units_processed.pop(0) + + def has_units_processed(self): + return len(self._units_processed) > 0 + + @property + def mean(self): + if len(self._samples) == 0: + return 0.0 + return float(sum(self._samples)) / len(self._samples) + + @property + def mean_units_processed(self): + if len(self._units_processed) == 0: + return 0.0 + return float(sum(self._units_processed)) / len(self._units_processed) + + @property + def mean_throughput(self): + time_total = float(sum(self._samples)) + if not time_total: + return 0.0 + return float(sum(self._units_processed)) / time_total diff --git a/lib/python3.12/site-packages/ray/util/tpu.py b/lib/python3.12/site-packages/ray/util/tpu.py new file mode 100644 index 0000000000000000000000000000000000000000..49cd8a2581e214d31e1a19c4b369af1e287698e4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/util/tpu.py @@ -0,0 +1,253 @@ +from typing import Optional + +import ray +from ray._private.accelerators import TPUAcceleratorManager +from ray._private.accelerators.tpu import ( + VALID_TPU_TYPES, + get_chips_per_host, + reserve_tpu_slice, +) +from ray._private.client_mode_hook import client_mode_wrap +from ray.util.annotations import PublicAPI +from ray.util.placement_group import PlacementGroup, placement_group + + +@PublicAPI(stability="alpha") +def get_current_pod_name() -> Optional[str]: + """ + Return the name of the TPU pod that the worker is a part of. + + Returns: + The name of the TPU pod. Returns None if not part of a TPU pod. + """ + tpu_name = TPUAcceleratorManager.get_current_node_tpu_name() + if tpu_name == "": + tpu_name = None + return tpu_name + + +@PublicAPI(stability="alpha") +def get_current_pod_worker_count() -> Optional[int]: + """ + Count the number of workers associated with the TPU pod that the worker belongs to. + + Returns: + The total number of workers in the TPU pod. Returns None if the worker is not + part of a TPU pod. + """ + return TPUAcceleratorManager.get_num_workers_in_current_tpu_pod() + + +@PublicAPI(stability="alpha") +def get_num_tpu_chips_on_node() -> int: + """ + Return the number of TPU chips on the node. + Returns: + The total number of chips on the TPU node. Returns 0 if none are found. + """ + return TPUAcceleratorManager.get_current_node_num_accelerators() + + +@PublicAPI(stability="alpha") +class SlicePlacementGroup: + """ + A handle to a placement group reservation for a TPU slice. + + The following definitions are added for clarity: + + - Accelerator type: A string describing the accelerator type and version (e.g. TPU-V2, TPU-V6E). + - Accelerator version: The accelerator generation only (e.g. v6e, v5p, v5litepod). + - Pod type: The TPU accelerator version and the number of chips in a topology. (e.g. v6e-128, v5p-8). + - Accelerator topology: The physical topology representing the structure (e.g. 2x2x2, 16x16). + + Args: + topology: The TPU topology string (e.g. "2x2x2"). + accelerator_version: The TPU accelerator generation (e.g. "v6e", "v5p", "v4"). + strategy: PlacementGroup parameter. The strategy to create the placement group. Currently default to "SPREAD" + + - "PACK": Packs Bundles into as few nodes as possible. + - "SPREAD": Places Bundles across distinct nodes as even as possible. + - "STRICT_PACK": Packs Bundles into one node. The group is + not allowed to span multiple nodes. + - "STRICT_SPREAD": Packs Bundles across distinct nodes. + + lifetime: PlacementGroup parameter. Either `None`, which defaults to the placement group + will fate share with its creator and will be deleted once its + creator is dead, or "detached", which means the placement group + will live as a global object independent of the creator. + + num_slices: Number of TPU slices in the SlicePlacementGroup. Defaults to 1 when unspecified. + + Examples: + + .. testcode:: python + :skipif: True + + import ray + from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + from ray.util.tpu import SlicePlacementGroup + + slice_handle = SlicePlacementGroup(topology="4x4", accelerator_version="v6e") + slice_pg = slice_handle.placement_group + ray.get(slice_pg.ready(), timeout=10) + + @ray.remote(num_cpus=0, resources={'TPU': 4}) + def spmd_task(world, rank): + print(f"Current TPU is rank {rank} of {world}") + + tasks = [ + spmd_task.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=slice_pg, + ) + ).remote(world=4, rank=i) + for i in range(slice_handle.num_workers) + ] + + """ + + def __init__( + self, + topology: str, + accelerator_version: str, + # below are args related to PG + strategy: str = "SPREAD", + name: str = "", + lifetime: Optional[str] = None, + # default + num_slices=1, + ): + self._topology = topology.strip().lower() + self._accelerator_version = accelerator_version.strip().lower() + self._num_slices = num_slices + self._validate_tpu_config() + + # Reserve a TPU slice of the provided accelerator version and topology. + self._placement_group = self._reserve_slice( + strategy, + name, + lifetime, + ) + + def _accelerator_version_check(self, accelerator_version: str): + if accelerator_version not in VALID_TPU_TYPES: + raise ValueError( + f"Invalid accelerator version: {accelerator_version}. Must be one of: {VALID_TPU_TYPES}" + ) + + def _validate_tpu_config(self): + # Should validate topology and generation values, calculate and + # set self._num_workers, and self._chips_per_host, and return a + # ValueError if invalid. + self._accelerator_version_check(self.accelerator_version) + if not TPUAcceleratorManager.is_valid_tpu_accelerator_topology( + tpu_accelerator_version=self.accelerator_version, + tpu_topology=self._topology, + ): + raise ValueError( + f"Invalid accelerator topology: '{self._topology}' for " + f"accelerator version: '{self.accelerator_version}'" + ) + + total_chips = 1 + for value in self._topology.strip().lower().split("x"): + total_chips *= int(value) + + self._chips_per_host = get_chips_per_host( + self._topology, self.accelerator_version + ) + self._num_workers_per_slice = total_chips // self._chips_per_host + self._num_workers = self._num_workers_per_slice * self._num_slices + + def _reserve_slice( + self, + strategy: str = "SPREAD", + name: str = "", + lifetime: Optional[str] = None, + ) -> PlacementGroup: + """Performs the two-step scheduling to reserve a TPU slice.""" + bundle_label_selector = [] + bundles = [] + + # Construct accelerator format for reserve_tpu_slice. e.g. From "v6e" to "TPU-V6E", "v5p" to "TPU-V5P". + accelerator_type = "TPU-" + self.accelerator_version.upper() + for _ in range(self.num_slices): + # Reserving a slice is done through constructing num_workers bundles, each with a label selector for + # the unique name of an available TPU slice. + slice_name = reserve_tpu_slice(self._topology, accelerator_type) + bundle_label_selector += [ + {ray._raylet.RAY_NODE_TPU_SLICE_NAME_KEY: slice_name} + ] * self._num_workers_per_slice + bundles += [{"TPU": self._chips_per_host}] * self._num_workers_per_slice + + pg = placement_group( + bundles=bundles, + strategy=strategy, + name=name, + lifetime=lifetime, + bundle_label_selector=bundle_label_selector, + ) + + return pg + + @property + def placement_group(self) -> PlacementGroup: + """The underlying PlacementGroup object.""" + return self._placement_group + + @property + def chips_per_host(self) -> int: + """The number of chips per host for this TPU slice.""" + # This is the same value as resources per worker for TPU. + return self._chips_per_host + + @property + def num_workers(self) -> int: + """The total number of hosts in the SlicePlacementGroup.""" + return self._num_workers + + @property + def topology(self) -> str: + """The physical topology of the TPU slice.""" + return self._topology + + @property + def accelerator_version(self) -> str: + """The TPU accelerator type of the slice.""" + return self._accelerator_version + + @property + def num_slices(self) -> int: + """The number of TPU slices this SlicePlacementGroup spans.""" + return self._num_slices + + +@PublicAPI(stability="alpha") +@client_mode_wrap +def slice_placement_group( + topology: str, + accelerator_version: str, + num_slices: int = 1, + **kwargs, +) -> SlicePlacementGroup: + """Asynchronously creates a PlacementGroup for a TPU slice. + + A slice placement group reserves num_slices TPU slice(s) and creates a placement + group for scheduling tasks. + + Args: + topology: The desired TPU pod topology (e.g. "4x4", "2x8"). + accelerator_version: The TPU accelerator generation, (e.g. "V4", "V5P", "V6E"). + num_slices: The number of tpu slices within the placement group + **kwargs: Additional arguments for the placement group, such as 'name', 'lifetime', or 'strategy'. + + Returns: + The handle for the created SlicePlacementGroup. + """ + + return SlicePlacementGroup( + topology=topology, + accelerator_version=accelerator_version, + num_slices=num_slices, + **kwargs, + ) diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..001b487ce79a74b86376f18d63e1c12496fbeb87 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_contextlib.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_contextlib.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2344ab9d94c363d497b0a7e5304e81b02b86e7e3 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_contextlib.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_datasets.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_datasets.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f75e154a1c306633067e2ea75a43f7601469112c Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_datasets.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_nestedkey.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_nestedkey.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d43197976b6db6ce396cf4a803f86c697c0566a Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_nestedkey.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_pytree.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_pytree.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cc528c053a809c1bc475aa2000548f1dd45dd61a Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_pytree.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_reductions.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_reductions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c82c379719baa9bb5a86f5de64a5f5b69cfd80e7 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_reductions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_tensorcollection.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_tensorcollection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1595d2b6b79aa4064990af2f476ed4d00011feaf Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_tensorcollection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_torch_func.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_torch_func.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a9398e1e4de778dbd1c1d36f9a36b6609dacbb15 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_torch_func.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_unbatched.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_unbatched.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1e7813849496826a216e419b34c411492d5a7c3e Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_unbatched.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/_version.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/_version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd3a8e94bcf081fb31b012893e08b8ff7838aa88 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/_version.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/functional.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/functional.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4e7dce214537fb9e1fc44f3bfd7eeb8e123b078e Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/functional.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/memmap.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/memmap.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..275f215de4d09705ab71eedd1752e9d22e09e1f5 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/memmap.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/persistent.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/persistent.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a710949584c26d86263a8717ec391a3a8f7f727f Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/persistent.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/__pycache__/tensordict.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/__pycache__/tensordict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11bea5dee60a08521d83081c5f96353bc66fb509 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/__pycache__/tensordict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9f6249af81bcdf9cbddbabf3cfe50b32d02d2bc2 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee133e6fcf2c830202544b9b428b7f48df4138c5 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/cudagraphs.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/cudagraphs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dec4584253a9b891c8c8c035a9fbd59afbad712d Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/cudagraphs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/ensemble.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/ensemble.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5a58ca48877230eafecc24bac53a8ac4c9b93cef Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/ensemble.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/functional_modules.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/functional_modules.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d45ef8c7fe057e1723bd69f182300bc6699113e Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/functional_modules.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/params.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/params.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f154f92be6e6cea20aec65ece6699e0b441eac8c Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/params.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/probabilistic.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/probabilistic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..58cf781bd90d77d9076304786bd346686476c82f Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/probabilistic.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/sequence.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/sequence.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6b8ad46310daefabf0221e4b34f60c779a3ee900 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/sequence.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/tensorclass_module.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/tensorclass_module.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..65f3b796ec61726f5be665d1e5f607f139bc986c Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/tensorclass_module.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0601039f72d0c29f38718b83b5eaa82c17965d3e Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/common.py b/lib/python3.12/site-packages/tensordict/nn/common.py new file mode 100644 index 0000000000000000000000000000000000000000..e1255bdba173b8eccdb33ea077eedb0535be1ee0 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/common.py @@ -0,0 +1,1494 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import functools +import inspect +import warnings +from functools import wraps +from textwrap import indent +from typing import ( + Any, + Callable, + Dict, + Iterable, + List, + MutableSequence, + Optional, + Sequence, + Tuple, + Union, +) + +import torch +from cloudpickle import dumps as cloudpickle_dumps, loads as cloudpickle_loads +from tensordict._td import TensorDict + +from tensordict.base import is_tensor_collection, NO_DEFAULT, TensorDictBase +from tensordict.functional import make_tensordict +from tensordict.nn.utils import _dispatch_td_nn_modules, _set_skip_existing_None +from tensordict.utils import ( + _unravel_key_to_tuple, + _zip_strict, + NestedKey, + unravel_key_list, +) +from torch import nn, Tensor + +try: + from torch.compiler import is_compiling +except ImportError: # torch 2.0 + from torch._dynamo import is_compiling + +try: + from functorch import FunctionalModule, FunctionalModuleWithBuffers + + _has_functorch = True +except ImportError: + _has_functorch = False + + class FunctionalModule: + pass + + class FunctionalModuleWithBuffers: + pass + + +__all__ = [ + "TensorDictModule", + "TensorDictModuleWrapper", +] + + +class dispatch: + """Allows for a function expecting a TensorDict to be called using kwargs. + + :func:`dispatch` must be used within modules that have an ``in_keys`` (or + another source of keys indicated by the ``source`` keyword argument) and + ``out_keys`` (or another ``dest`` key list) attributes indicating what keys + to be read and written from the tensordict. The wrapped function should + also have a ``tensordict`` leading argument. + + The resulting function will return a single tensor (if there is a single + element in out_keys), otherwise it will return a tuple sorted as the ``out_keys`` + of the module. + + :func:`dispatch` can be used either as a method or as a class when extra arguments + need to be passed. + + Args: + separator (str, optional): separator that combines sub-keys together + for ``in_keys`` that are tuples of strings. + Defaults to ``"_"``. + source (str or list of keys, optional): if a string is provided, + it points to the module attribute that contains the + list of input keys to be used. If a list is provided instead, it + will contain the keys used as input to the module. + Defaults to ``"in_keys"`` which is the attribute name of + :class:`~.TensorDictModule` list of input keys. + dest (str or list of keys, optional): if a string is provided, + it points to the module attribute that contains the + list of output keys to be used. If a list is provided instead, it + will contain the keys used as output to the module. + Defaults to ``"out_keys"`` which is the attribute name of + :class:`~.TensorDictModule` list of output keys. + auto_batch_size (bool, optional): if ``True``, the batch-size of the + input tensordict is determined automatically as the maximum number + of common dimensions across all the input tensors. + Defaults to ``True``. + + Examples: + >>> class MyModule(nn.Module): + ... in_keys = ["a"] + ... out_keys = ["b"] + ... + ... @dispatch + ... def forward(self, tensordict): + ... tensordict['b'] = tensordict['a'] + 1 + ... return tensordict + ... + >>> module = MyModule() + >>> b = module(a=torch.zeros(1, 2)) + >>> assert (b == 1).all() + >>> # equivalently + >>> class MyModule(nn.Module): + ... keys_in = ["a"] + ... keys_out = ["b"] + ... + ... @dispatch(source="keys_in", dest="keys_out") + ... def forward(self, tensordict): + ... tensordict['b'] = tensordict['a'] + 1 + ... return tensordict + ... + >>> module = MyModule() + >>> b = module(a=torch.zeros(1, 2)) + >>> assert (b == 1).all() + >>> # or this + >>> class MyModule(nn.Module): + ... @dispatch(source=["a"], dest=["b"]) + ... def forward(self, tensordict): + ... tensordict['b'] = tensordict['a'] + 1 + ... return tensordict + ... + >>> module = MyModule() + >>> b = module(a=torch.zeros(1, 2)) + >>> assert (b == 1).all() + + :func:`dispatch_kwargs` will also work with nested keys with the default + ``"_"`` separator. + + Examples: + >>> class MyModuleNest(nn.Module): + ... in_keys = [("a", "c")] + ... out_keys = ["b"] + ... + ... @dispatch + ... def forward(self, tensordict): + ... tensordict['b'] = tensordict['a', 'c'] + 1 + ... return tensordict + ... + >>> module = MyModuleNest() + >>> b, = module(a_c=torch.zeros(1, 2)) + >>> assert (b == 1).all() + + If another separator is wanted, it can be indicated with the ``separator`` + argument in the constructor: + + Examples: + >>> class MyModuleNest(nn.Module): + ... in_keys = [("a", "c")] + ... out_keys = ["b"] + ... + ... @dispatch(separator="sep") + ... def forward(self, tensordict): + ... tensordict['b'] = tensordict['a', 'c'] + 1 + ... return tensordict + ... + >>> module = MyModuleNest() + >>> b, = module(asepc=torch.zeros(1, 2)) + >>> assert (b == 1).all() + + + Since the input keys is a sorted sequence of strings, + :func:`dispatch` can also be used with unnamed arguments where the order + must match the order of the input keys. + + .. note:: + If the first argument is a :class:`~.TensorDictBase` instance, it is + assumed that dispatch is __not__ being used and that this tensordict + contains all the necessary information to be run through the module. + In other words, one cannot decompose a tensordict with the first key + of the module inputs pointing to a tensordict instance. + In general, it is preferred to use :func:`dispatch` with tensordict + leaves only. + + Examples: + >>> class MyModuleNest(nn.Module): + ... in_keys = [("a", "c"), "d"] + ... out_keys = ["b"] + ... + ... @dispatch + ... def forward(self, tensordict): + ... tensordict['b'] = tensordict['a', 'c'] + tensordict["d"] + ... return tensordict + ... + >>> module = MyModuleNest() + >>> b, = module(torch.zeros(1, 2), d=torch.ones(1, 2)) # works + >>> assert (b == 1).all() + >>> b, = module(torch.zeros(1, 2), torch.ones(1, 2)) # works + >>> assert (b == 1).all() + >>> try: + ... b, = module(torch.zeros(1, 2), a_c=torch.ones(1, 2)) # fails + ... except: + ... print("oopsy!") + ... + + """ + + DEFAULT_SEPARATOR = "_" + DEFAULT_SOURCE = "in_keys" + DEFAULT_DEST = "out_keys" + + def __new__( + cls, + separator=DEFAULT_SEPARATOR, + source=DEFAULT_SOURCE, + dest=DEFAULT_DEST, + auto_batch_size: bool = True, + ): + if callable(separator): + func = separator + separator = dispatch.DEFAULT_SEPARATOR + self = super().__new__(cls) + self.__init__(separator, source, dest) + return self.__call__(func) + return super().__new__(cls) + + def __init__( + self, + separator=DEFAULT_SEPARATOR, + source=DEFAULT_SOURCE, + dest=DEFAULT_DEST, + auto_batch_size: bool = True, + ): + self.separator = separator + self.source = source + self.dest = dest + self.auto_batch_size = auto_batch_size + + def __call__(self, func: Callable) -> Callable: + + is_method = inspect.ismethod(func) or ( + inspect.isfunction(func) + and func.__code__.co_argcount > 0 + and func.__code__.co_varnames[0] == "self" + ) + # sanity check + for i, key in enumerate(inspect.signature(func).parameters): + if (is_method or (key == "self")) and (i == 0): + is_method = True + # skip self + continue + if key != "tensordict": + raise RuntimeError( + "the first argument of the wrapped function must be " + f"named 'tensordict'. Got {key} instead." + ) + break + # if the env variable was used, we can skip the wrapper altogether + if not _dispatch_td_nn_modules(): + return func + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + if is_method: + _self = args[0] + args = args[1:] + else: + _self = None + if not _dispatch_td_nn_modules(): + return func(_self, *args, **kwargs) + + source = self.source + if isinstance(source, str): + if _self is None: + raise RuntimeError( + "The in keys must be passed to dispatch when func is not a method but a function." + ) + source = getattr(_self, source) + tensordict = None + if len(args): + if is_tensor_collection(args[0]): + tensordict, args = args[0], args[1:] + if tensordict is None: + tensordict_values = {} + dest = self.dest + if isinstance(dest, str): + if _self is None: + raise RuntimeError( + "The in keys must be passed to dispatch when func is not a method but a function." + ) + dest = getattr(_self, dest) + for key in source: + expected_key = self.separator.join(_unravel_key_to_tuple(key)) + if len(args): + tensordict_values[key] = args[0] + args = args[1:] + if expected_key in kwargs: + raise RuntimeError( + "Duplicated argument in args and kwargs." + ) + elif expected_key in kwargs: + try: + tensordict_values[key] = kwargs.pop(expected_key) + except KeyError: + raise KeyError( + f"The key {expected_key} wasn't found in the keyword arguments " + f"but is expected to execute that function." + ) + batch_size = torch.Size([]) if not self.auto_batch_size else None + tensordict = make_tensordict( + tensordict_values, + batch_size=batch_size, + auto_batch_size=self.auto_batch_size, + ) + if _self is not None: + out = func(_self, tensordict, *args, **kwargs) + else: + out = func(tensordict, *args, **kwargs) + + # This makes dispatch responsible of handling partial outputs (such as selected through select_out_keys) + out = tuple(out[key] for key in dest) + return out[0] if len(out) == 1 else out + + if _self is not None: + return func(_self, tensordict, *args, **kwargs) + return func(tensordict, *args, **kwargs) + + return self._update_func_signature(func, wrapper) + + def _update_func_signature(self, func, wrapper): + # Create a new signature with the desired parameters + # Get the original function's signature + orig_signature = inspect.signature(func) + + # params = [inspect.Parameter(name='', kind=inspect.Parameter.VAR_POSITIONAL)] + params = [] + i = -1 + for i, param in enumerate(orig_signature.parameters.values()): + if param.kind in ( + inspect.Parameter.VAR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ): + i = i - 1 + break + if param.default is inspect._empty: + params.append( + inspect.Parameter( + name=param.name, + kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, + default=None, + ) + ) + else: + params.append(param) + + # Add the **kwargs parameter + + # for key in self.get_source(func, self_func): + if i >= 0: + params.extend(list(orig_signature.parameters.values())[i + 1 :]) + elif i == -1: + params.extend(list(orig_signature.parameters.values())) + + # Update the wrapper's signature + wrapper.__signature__ = inspect.Signature(params) + + return wrapper + + def get_source(self, func, self_func): + source = self.source + if isinstance(source, str): + return getattr(self_func, source) + return source + + +class _OutKeysSelect: + module: nn.Module | None = None + + def __init__(self, out_keys): + self.out_keys = list(out_keys) + self._initialized = None + self._is_dispatched = None + + def _init(self, module): + if self._initialized: + return + self._initialized = True + self.module = module + if not all(key in module.out_keys for key in self.out_keys): + raise RuntimeError("Some keys are not part of the module out_keys.") + module.out_keys = self.out_keys + + def __call__( # noqa: F811 + self, + module: TensorDictModuleBase, + tensordict_in: TensorDictBase, + kwargs: Dict, + tensordict_out: TensorDictBase, + ): + if not self._initialized: + raise RuntimeError( + "_OutKeysSelect must be initialized before being called." + ) + # detect dispatch calls + in_keys = module.in_keys + if not tensordict_in and kwargs.get("tensordict") is not None: + tensordict_in = kwargs.pop("tensordict") + is_dispatched = self._detect_dispatch(tensordict_in, kwargs, in_keys) + out_keys = self.out_keys + # if dispatch filtered the out keys as they should we're happy + if is_dispatched: + if (not isinstance(tensordict_out, tuple) and len(out_keys) == 1) or ( + isinstance(tensordict_out, tuple) + and len(out_keys) == len(tensordict_out) + ): + return tensordict_out + if is_dispatched: + # it might be the case that dispatch was not aware of what the out-keys were. + if isinstance(tensordict_out, tuple): + out = tuple( + item + for i, item in enumerate(tensordict_out) + if module._out_keys[i] in module.out_keys + ) + if len(out) == 1: + return out[0] + return out + elif module._out_keys[0] in module.out_keys and len(module._out_keys) == 1: + return tensordict_out + elif ( + module._out_keys[0] not in module.out_keys + and len(module._out_keys) == 1 + ): + return () + else: + raise RuntimeError( + f"Selecting out-keys failed. Original out_keys: {module._out_keys}, selected: {module.out_keys}." + ) + return tensordict_out.select( + *in_keys, *out_keys, inplace=True, strict=tensordict_out is tensordict_in + ) + + def _detect_dispatch(self, tensordict_in, kwargs, in_keys): # noqa: F811 + if isinstance(tensordict_in, TensorDictBase) and all( + key in tensordict_in.keys(include_nested=True) for key in in_keys + ): + return False + elif isinstance(tensordict_in, tuple): + if len(tensordict_in) or len(kwargs): + if len(tensordict_in) and isinstance(tensordict_in[0], TensorDictBase): + return self._detect_dispatch(tensordict_in[0], kwargs, in_keys) + elif ( + not len(tensordict_in) + and len(kwargs) + and isinstance(kwargs.get("tensordict"), TensorDictBase) + ): + return self._detect_dispatch(kwargs["tensordict"], in_keys) + return True + return not len(in_keys) + # not a TDBase: must be True + return True + + def remove(self): + # reset ground truth + if self.module is None: + return + if self.module._out_keys is not None: + self.module.out_keys = self.module._out_keys + + def __del__(self): + self.remove() + + +class TensorDictModuleBase(nn.Module): + """Base class to TensorDict modules. + + TensorDictModule subclasses are characterized by ``in_keys`` and ``out_keys`` + key-lists that indicate what input entries are to be read and what output + entries should be expected to be written. + + The forward method input/output signature should always follow the + convention: + + >>> tensordict_out = module.forward(tensordict_in) + + Unlike :class:`~tensordict.nn.TensorDictModule`, `TensorDictModuleBase` is typically used via subclassing: + you can wrap any python function in a `TensorDictModuleBase` subclass, as long as the subclass forward reads and + writes tensordict (or related types) instances. + + The `in_keys` and `out_keys` should be properly specified. For example, `out_keys` can be dynamically reduced using + :meth:`~tensordict.nn.TensorDictBase.select_out_keys`. + + Examples: + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictModuleBase + >>> class Mod(TensorDictModuleBase): + ... in_keys = ["a"] # can also be specified during __init__ + ... out_keys = ["b", "c"] + ... def forward(self, tensordict): + ... b = tensordict["a"].clone() + ... c = b + 1 + ... return tensordict.replace({"b": b, "c": c}) + >>> mod = Mod() + >>> td = mod(TensorDict(a=0)) + >>> td["b"] + tensor(0) + >>> td["c"] + tensor(1) + >>> mod.select_out_keys("c") + >>> td = mod(TensorDict(a=0)) + >>> td["c"] + tensor(1) + >>> assert "b" not in td + + """ + + def __new__(cls, *args, **kwargs): + # check the out_keys and in_keys in the dict + if "in_keys" in cls.__dict__ and not isinstance( + cls.__dict__.get("in_keys"), property + ): + in_keys = cls.__dict__.get("in_keys") + # now let's remove it + delattr(cls, "in_keys") + cls._in_keys = unravel_key_list(in_keys) + cls.in_keys = TensorDictModuleBase.in_keys + if "out_keys" in cls.__dict__ and not isinstance( + cls.__dict__.get("out_keys"), property + ): + out_keys = cls.__dict__.get("out_keys") + # now let's remove it + delattr(cls, "out_keys") + out_keys = unravel_key_list(out_keys) + cls._out_keys = out_keys + cls._out_keys_apparent = out_keys + cls.out_keys = TensorDictModuleBase.out_keys + out = super().__new__(cls) + return out + + @staticmethod + def is_tdmodule_compatible(module): + """Checks if a module is compatible with TensorDictModule API.""" + return hasattr(module, "in_keys") and hasattr(module, "out_keys") + + @property + def in_keys(self): + return self._in_keys + + @in_keys.setter + def in_keys(self, value: List[Union[str, Tuple[str]]]): + self._in_keys = unravel_key_list(value) + + @property + def out_keys(self): + return self._out_keys_apparent + + @property + def out_keys_source(self): + return self._out_keys + + @out_keys.setter + def out_keys(self, value: List[Union[str, Tuple[str]]]): + # the first time out_keys are set, they are marked as ground truth + value = unravel_key_list(list(value)) + if not hasattr(self, "_out_keys"): + self._out_keys = value + self._out_keys_apparent = value + + def select_out_keys(self, *out_keys) -> TensorDictModuleBase: # noqa: F811 + """Selects the keys that will be found in the output tensordict. + + This is useful whenever one wants to get rid of intermediate keys in a + complicated graph, or when the presence of these keys may trigger unexpected + behaviours. + + The original ``out_keys`` can still be accessed via ``module.out_keys_source``. + + Args: + *out_keys (a sequence of strings or tuples of strings): the + out_keys that should be found in the output tensordict. + + Returns: the same module, modified in-place with updated ``out_keys``. + + The simplest usage is with :class:`~.TensorDictModule`: + + Examples: + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictModule, TensorDictSequential + >>> import torch + >>> mod = TensorDictModule(lambda x, y: (x+2, y+2), in_keys=["a", "b"], out_keys=["c", "d"]) + >>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []) + >>> mod(td) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> mod.select_out_keys("d") + >>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []) + >>> mod(td) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + This feature will also work with dispatched arguments: + Examples: + >>> mod(torch.zeros(()), torch.ones(())) + tensor(2.) + + This change will occur in-place (ie the same module will be returned + with an updated list of out_keys). It can be reverted using the + :meth:`TensorDictModuleBase.reset_out_keys` method. + + Examples: + >>> mod.reset_out_keys() + >>> mod(TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, [])) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + This will work with other classes too, such as Sequential: + Examples: + >>> from tensordict.nn import TensorDictSequential + >>> seq = TensorDictSequential( + ... TensorDictModule(lambda x: x+1, in_keys=["x"], out_keys=["y"]), + ... TensorDictModule(lambda x: x+1, in_keys=["y"], out_keys=["z"]), + ... ) + >>> td = TensorDict({"x": torch.zeros(())}, []) + >>> seq(td) + TensorDict( + fields={ + x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + y: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + z: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> seq.select_out_keys("z") + >>> td = TensorDict({"x": torch.zeros(())}, []) + >>> seq(td) + TensorDict( + fields={ + x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + z: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + out_keys = unravel_key_list(list(out_keys)) + if len(out_keys) == 1: + if out_keys[0] not in self.out_keys: + err_msg = f"Can't select non existent key: {out_keys[0]}. " + if ( + out_keys[0] + and isinstance(out_keys[0], (tuple, list)) + and out_keys[0][0] in self.out_keys + ): + err_msg += f"Are you passing the keys in a list? Try unpacking as: `{', '.join(out_keys[0])}`" + raise ValueError(err_msg) + self.register_forward_hook(_OutKeysSelect(out_keys), with_kwargs=True) + for hook in self._forward_hooks.values(): + if isinstance(hook, _OutKeysSelect): + hook._init(self) + return self + + def reset_out_keys(self): + """Resets the ``out_keys`` attribute to its orignal value. + + Returns: the same module, with its original ``out_keys`` values. + + Examples: + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictModule, TensorDictSequential + >>> import torch + >>> mod = TensorDictModule(lambda x, y: (x+2, y+2), in_keys=["a", "b"], out_keys=["c", "d"]) + >>> mod.select_out_keys("d") + >>> td = TensorDict({"a": torch.zeros(()), "b": torch.ones(())}, []) + >>> mod(td) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> mod.reset_out_keys() + >>> mod(td) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + """ + for i, hook in list(self._forward_hooks.items()): + if isinstance(hook, _OutKeysSelect): + hook.remove() + del self._forward_hooks[i] + return self + + def reset_parameters_recursive( + self, parameters: Optional[TensorDictBase] = None + ) -> Optional[TensorDictBase]: + """Recursively reset the parameters of the module and its children. + + Args: + parameters (TensorDict of parameters, optional): If set to None, the module will reset using self.parameters(). + Otherwise, we will reset the parameters in the tensordict in-place. This is + useful for functional modules where the parameters are not stored in the module itself. + + Returns: + A tensordict of the new parameters, only if parameters was not None. + + Examples: + >>> from tensordict.nn import TensorDictModule + >>> from torch import nn + >>> net = nn.Sequential(nn.Linear(2,3), nn.ReLU()) + >>> old_param = net[0].weight.clone() + >>> module = TensorDictModule(net, in_keys=['bork'], out_keys=['dork']) + >>> module.reset_parameters() + >>> (old_param == net[0].weight).any() + tensor(False) + + This method also supports functional parameter sampling: + + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictModule + >>> from torch import nn + >>> net = nn.Sequential(nn.Linear(2,3), nn.ReLU()) + >>> module = TensorDictModule(net, in_keys=['bork'], out_keys=['dork']) + >>> params = TensorDict.from_module(module) + >>> old_params = params.clone(recurse=True) + >>> module.reset_parameters(params) + >>> (old_params == params).any() + False + """ + if parameters is None: + any_reset = self._reset_parameters(self) + if not any_reset: + warnings.warn( + "reset_parameters_recursive was called without the parameters argument and did not find any parameters to reset" + ) + return + elif parameters.ndim: + raise RuntimeError( + "reset_parameters_recursive does not support batched TensorDicts, ensure `batch_size` is empty and the parameters shape match their original shape." + ) + + sanitized_parameters = parameters.apply( + lambda x: x.detach().requires_grad_(), inplace=False + ) + + with sanitized_parameters.to_module(self): + self._reset_parameters(self) + return sanitized_parameters + + def _reset_parameters(self, module: nn.Module) -> bool: + any_reset = False + for child in module.children(): + if isinstance(child, nn.Module): + any_reset |= self._reset_parameters(child) + + if hasattr(child, "reset_parameters"): + child.reset_parameters() + any_reset |= True + return any_reset + + @property + def __name__(self): + # This is necessary to make compiled vmap over TDModule happy + return self.__class__.__name__ + + def __repr__(self): + return f"{self.__class__.__name__}()" + + +class TensorDictModule(TensorDictModuleBase): + """A TensorDictModule, is a python wrapper around a :obj:`nn.Module` that reads and writes to a TensorDict. + + Args: + module (Callable[[Any], Any]): a callable, typically a :class:`torch.nn.Module`, + used to map the input to the output parameter space. Its forward method + can return a single tensor, a tuple of tensors or even a dictionary. + In the latter case, the output keys of the :class:`TensorDictModule` + will be used to populate the output tensordict (ie. the keys present + in ``out_keys`` should be present in the dictionary returned by the + ``module`` forward method). + in_keys (iterable of NestedKeys, Dict[NestedStr, str]): keys to be read + from input tensordict and passed to the module. If it + contains more than one element, the values will be passed in the + order given by the in_keys iterable. + If ``in_keys`` is a dictionary, its keys must correspond to the key + to be read in the tensordict and its values must match the name of + the keyword argument in the function signature. If `out_to_in_map` is ``True``, + the mapping gets inverted so that the keys correspond to the keyword + arguments in the function signature. + out_keys (iterable of str): keys to be written to the input tensordict. The length of out_keys must match the + number of tensors returned by the embedded module. Using "_" as a key avoid writing tensor to output. + + Keyword Args: + out_to_in_map (bool, optional): if ``True`` (default), `in_keys` is read as if the keys are the arguments keys of + the :meth:`~.forward` method and the values are the keys in the input :class:`~tensordict.TensorDict`. If + ``False``, keys are considered to be the input keys and values the method's arguments keys. + inplace (bool or string, optional): if ``True`` (default), the output of the module are written in the tensordict + provided to the :meth:`~.forward` method. If ``False``, a new :class:`~tensordict.TensorDict` with and empty + batch-size and no device is created. if ``"empty"``, :meth:`~tensordict.TensorDict.empty` will be used to + create the output tensordict. + + .. note:: + If ``inplace=False`` and the tensordict passed to the module is another + :class:`~tensordict.TensorDictBase` subclass than :class:`~tensordict.TensorDict`, the output will still + be a :class:`~tensordict.TensorDict` instance. Its batch-size will be empty, and it will have no device. + Set to ``"empty"`` to get the same :class:`~tensordict.TensorDictBase` subtype, an identical batch-size + and device. Use ``tensordict_out`` at runtime (see below) to have a more fine-grained control over the + output. + + .. note:: + If ``inplace=False`` and a `tensordict_out` is passed to the :meth:`~.forward` method, + the ``tensordict_out`` will prevail. This is the way one can get a tensordict_out taensordict passed to the module is another + :class:`~tensordict.TensorDictBase` subclass than :class:`~tensordict.TensorDict`, the output will still + be a :class:`~tensordict.TensorDict` instance. + + method (str, optional): the method to be called in the module, if any. Defaults to `__call__`. + method_kwargs (Dict[str, Any], optional): additional keyword arguments to be passed to the module's method being called. + strict (bool, optional): if ``True``, the module will raise an exception if any of the inputs is missing from + the input tensordict. Otherwise, a `None` value will be used as placeholder. Defaults to ``False``. + get_kwargs (dict[str, Any], optional): additional keyword arguments to be passed to the :meth:`~tensordict.TensorDictBase.get` + method. This is particularily useful when dealing with ragged tensors (see :meth:`~tensordict.LazyStackedTensorDict.get`). + Defaults to ``{}``. + + Embedding a neural network in a TensorDictModule only requires to specify the input + and output keys. TensorDictModule support functional and regular :obj:`nn.Module` + objects. In the functional case, the 'params' (and 'buffers') keyword argument must + be specified: + + Examples: + >>> from tensordict import TensorDict + >>> # one can wrap regular nn.Module + >>> module = TensorDictModule(nn.Transformer(128), in_keys=["input", "tgt"], out_keys=["out"]) + >>> input = torch.ones(2, 3, 128) + >>> tgt = torch.zeros(2, 3, 128) + >>> data = TensorDict({"input": input, "tgt": tgt}, batch_size=[2, 3]) + >>> data = module(data) + >>> print(data) + TensorDict( + fields={ + input: Tensor(shape=torch.Size([2, 3, 128]), device=cpu, dtype=torch.float32, is_shared=False), + out: Tensor(shape=torch.Size([2, 3, 128]), device=cpu, dtype=torch.float32, is_shared=False), + tgt: Tensor(shape=torch.Size([2, 3, 128]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([2, 3]), + device=None, + is_shared=False) + + We can also pass directly the tensors + + Examples: + >>> out = module(input, tgt) + >>> assert out.shape == input.shape + >>> # we can also wrap regular functions + >>> module = TensorDictModule(lambda x: (x-1, x+1), in_keys=[("input", "x")], out_keys=[("output", "x-1"), ("output", "x+1")]) + >>> module(TensorDict({("input", "x"): torch.zeros(())}, batch_size=[])) + TensorDict( + fields={ + input: TensorDict( + fields={ + x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False), + output: TensorDict( + fields={ + x+1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + x-1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + We can use TensorDictModule to populate a tensordict: + + Examples: + >>> module = TensorDictModule(lambda: torch.randn(3), in_keys=[], out_keys=["x"]) + >>> print(module(TensorDict({}, batch_size=[]))) + TensorDict( + fields={ + x: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + Another feature is passing a dictionary as input keys, to control the + dispatching of values to specific keyword arguments. + + Examples: + >>> module = TensorDictModule(lambda x, *, y: x+y, + ... in_keys={'1': 'x', '2': 'y'}, out_keys=['z'], out_to_in_map=False + ... ) + >>> td = module(TensorDict({'1': torch.ones(()), '2': torch.ones(())*2}, [])) + >>> td['z'] + tensor(3.) + + If `out_to_in_map` is set to ``True``, then the `in_keys` mapping is reversed. This way, + one can use the same input key for different keyword arguments. + + Examples: + >>> module = TensorDictModule(lambda x, *, y, z: x+y+z, + ... in_keys={'x': '1', 'y': '2', z: '2'}, out_keys=['t'], out_to_in_map=True + ... ) + >>> td = module(TensorDict({'1': torch.ones(()), '2': torch.ones(())*2}, [])) + >>> td['t'] + tensor(5.) + + We can specify the method to be called within a module. Compared to using a lambda function or similar around the + module's method, this has the advantage that the module attributes (params, buffers, submodules) will be exposed. + + Examples: + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod + >>> from torch import nn + >>> import torch + >>> + >>> class MyNet(nn.Module): + ... def my_func(self, tensor: torch.Tensor, *, an_integer: int): + ... return tensor + an_integer + ... + >>> s = Seq( + ... { + ... "a": lambda td: td+1, + ... "b": lambda td: td * 2, + ... "c": Mod(MyNet(), in_keys=["a"], out_keys=["b"], method="my_func", method_kwargs={"an_integer": 2}), + ... } + ... ) + >>> td = s(TensorDict(a=0)) + >>> print(td) + >>> + >>> assert td["b"] == 4 + + Functional calls to a tensordict module is easy: + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictModule + >>> td = TensorDict({"input": torch.randn(3, 4), "hidden": torch.randn(3, 8)}, [3,]) + >>> module = torch.nn.GRUCell(4, 8) + >>> td_module = TensorDictModule( + ... module=module, in_keys=["input", "hidden"], out_keys=["output"] + ... ) + >>> params = TensorDict.from_module(td_module) + >>> # functional API + >>> with params.to_module(td_module): + ... td_functional = td_module(td.clone()) + >>> print(td_functional) + TensorDict( + fields={ + hidden: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + input: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + output: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + + In the stateful case: + >>> module = torch.nn.GRUCell(4, 8) + >>> td_module = TensorDictModule( + ... module=module, in_keys=["input", "hidden"], out_keys=["output"] + ... ) + >>> td_stateful = td_module(td.clone()) + >>> print(td_stateful) + TensorDict( + fields={ + hidden: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + input: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + output: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + + """ + + _IN_KEY_ERR = "in_keys must be of type list, str or tuples of str, or dict." + _OUT_KEY_ERR = "out_keys must be of type list, str or tuples of str." + + def __init__( + self, + module: Callable, + in_keys: NestedKey | List[NestedKey] | Dict[NestedKey:str], + out_keys: NestedKey | List[NestedKey], + *, + out_to_in_map: bool | None = None, + inplace: bool | str = True, + method: str | None = None, + method_kwargs: dict | None = None, + strict: bool = False, + get_kwargs: dict | None = None, + ) -> None: + super().__init__() + + if out_to_in_map is not None and not isinstance(in_keys, dict): + warnings.warn( + "out_to_in_map is not None but is only used when in_key is a dictionary." + ) + + if isinstance(in_keys, dict): + if out_to_in_map is None: + out_to_in_map = True + + # write the kwargs and create a list instead + _in_keys = [] + self._kwargs = [] + for key, value in in_keys.items(): + if out_to_in_map: # arg: td_key + self._kwargs.append(key) + _in_keys.append(value) + else: # td_key: arg + self._kwargs.append(value) + _in_keys.append(key) + in_keys = _in_keys + else: + if isinstance(in_keys, (str, tuple)): + in_keys = [in_keys] + elif not isinstance(in_keys, MutableSequence): + raise ValueError(self._IN_KEY_ERR) + self._kwargs = None + + if isinstance(out_keys, (str, tuple)): + out_keys = [out_keys] + elif not isinstance(out_keys, MutableSequence): + raise ValueError(self._OUT_KEY_ERR) + try: + in_keys = unravel_key_list(list(in_keys)) + except Exception: + raise ValueError(self._IN_KEY_ERR) + try: + out_keys = unravel_key_list(list(out_keys)) + except Exception: + raise ValueError(self._OUT_KEY_ERR) + + if type(module) is type or (method is None and not callable(module)): + raise ValueError( + f"Module {module} if type {type(module)} is not callable. " + f"Typical accepted types are nn.Module or TensorDictModule. " + f"If you need to call a specific method from your module, pass the " + f"`method` keyword argument to the TensorDictModule constructor." + ) + self.out_keys = out_keys + self.in_keys = in_keys + + self.strict = strict + + if "_" in self.in_keys: + warnings.warn( + 'key "_" is for ignoring output, it should not be used in input keys', + stacklevel=2, + ) + + self.module = module + if inplace not in (None, True, False, "empty"): + raise ValueError( + f"The only accepted valued for inplace is `None`, `True`, `False`, or `'empty'`. Got inplace={inplace} " + "instead." + ) + self.inplace = inplace + self.method = method + self.method_kwargs = method_kwargs if method_kwargs is not None else {} + self._get_kwargs = get_kwargs if get_kwargs is not None else {} + + @property + def is_functional(self) -> bool: + return _has_functorch and isinstance( + self.module, + (FunctionalModule, FunctionalModuleWithBuffers), + ) + + def _write_to_tensordict( + self, + tensordict: TensorDictBase, + tensors: list[Tensor], + tensordict_out: TensorDictBase | None = None, + out_keys: Iterable[NestedKey] | None = None, + ) -> TensorDictBase: + if out_keys is None: + out_keys = self.out_keys_source + if tensordict_out is None: + if self.inplace is not True: + if self.inplace == "empty": + tensordict_out = tensordict.empty() + else: + tensordict_out = TensorDict() + else: + tensordict_out = tensordict + if len(tensors) > len(out_keys): + raise RuntimeError( + f"There are more tensors ({len(tensors)=}) than out_keys ({out_keys=})." + ) + elif len(out_keys) > len(tensors): + raise RuntimeError("There are more out_keys than tensors.") + for _out_key, _tensor in zip(out_keys, tensors): + if _out_key != "_": + tensordict_out.set(_out_key, TensorDict.from_any(_tensor)) + return tensordict_out + + def _call_module( + self, tensors: Sequence[Tensor], **kwargs: Any + ) -> Tensor | Sequence[Tensor]: + kwargs.update(self.method_kwargs) + if self.method is None: + out = self.module(*tensors, **kwargs) + else: + out = getattr(self.module, self.method)(*tensors, **kwargs) + return out + + @dispatch(auto_batch_size=False) + @_set_skip_existing_None() + def forward( + self, + tensordict: TensorDictBase, + *args, + tensordict_out: TensorDictBase | None = None, + **kwargs: Any, + ) -> TensorDictBase: + """When the tensordict parameter is not set, kwargs are used to create an instance of TensorDict.""" + try: + if len(args): + raise ValueError( + "Got a non-empty list of extra agruments, when none was expected." + ) + default = None if not self.strict else NO_DEFAULT + if self._kwargs is not None: + kwargs.update( + { + kwarg: tensordict._get_tuple_maybe_non_tensor( + _unravel_key_to_tuple(in_key), default=default + ) + for kwarg, in_key in _zip_strict(self._kwargs, self.in_keys) + } + ) + tensors = () + else: + tensors = tuple( # type: ignore[unreachable] + tensordict._get_tuple_maybe_non_tensor( + _unravel_key_to_tuple(in_key), + default, + **self._get_kwargs, + ) + for in_key in self.in_keys + ) + try: + tensors_out = self._call_module(tensors, **kwargs) + if tensors_out is None: + tensors_out = () + except Exception as err: + if any(tensor is None for tensor in tensors) and "None" in str(err): + none_set = { + key + for key, tensor in _zip_strict(self.in_keys, tensors) + if tensor is None + } + raise KeyError( + "Some tensors that are necessary for the module call may " + "not have not been found in the input tensordict: " + f"the following inputs are None: {none_set}." + ) from err + else: + raise err + if isinstance(tensors_out, (dict, TensorDictBase)) and all( + key in tensors_out for key in self.out_keys + ): + if isinstance(tensors_out, dict): + keys = unravel_key_list(list(tensors_out.keys())) + values = tensors_out.values() + tensors_out = dict(_zip_strict(keys, values)) + tensors_out = tuple(tensors_out.get(key) for key in self.out_keys) + if not isinstance(tensors_out, tuple): + tensors_out = (tensors_out,) + tensordict_out = self._write_to_tensordict( + tensordict, tensors_out, tensordict_out + ) + return tensordict_out + except Exception as err: + module = self.module + if not isinstance(module, nn.Module): + try: + import inspect + + module = inspect.getsource(module) + except Exception: + # then we can't print the source code + pass + module = indent(str(module), 4 * " ") + in_keys = indent(f"in_keys={self.in_keys}", 4 * " ") + out_keys = indent(f"out_keys={self.out_keys}", 4 * " ") + raise err from RuntimeError( + f"TensorDictModule failed with operation\n{module}\n{in_keys}\n{out_keys}." + ) + + @property + def device(self) -> torch.device: + for p in self.parameters(): + return p.device + return torch.device("cpu") + + def __repr__(self) -> str: + fields = indent( + f"module={self.module},\n" + f"device={self.device},\n" + f"in_keys={self.in_keys},\n" + f"out_keys={self.out_keys}", + 4 * " ", + ) + + return f"{type(self).__name__}(\n{fields})" + + def __getattr__(self, name: str) -> Any: + if not is_compiling(): + __dict__ = self.__dict__ + _parameters = __dict__.get("_parameters") + if _parameters: + # A param can be None so we use False instead to check if the key + # is in the _parameters dict once and only once. + # We use False but any non-None, non-Parameter value would do. + # The alternative `if name in _parameters: return _parameters[name]` + # accesses the value of `name` twice when only one is required + result = _parameters.get(name, False) + if result is not False: + return result + _buffers = __dict__.get("_buffers") + if _buffers: + result = _buffers.get(name, False) + if result is not False: + return result + _modules = __dict__.get("_modules") + if _modules: + result = _modules.get(name, False) + if result is not False: + return result + # TODO: find a way to make this check work with dynamo + # elif hasattr(self, "_parameters"): + else: + _parameters = self._parameters + result = _parameters.get(name, False) + if result is not False: + return result + _buffers = self._buffers + result = _buffers.get(name, False) + if result is not False: + return result + _modules = self._modules + result = _modules.get(name, False) + if result is not False: + return result + + if not name.startswith("_"): + # no fallback for private attributes + return getattr(super().__getattr__("module"), name) + raise AttributeError( + f"module {type(self).__name__} has no attribute named {name}." + ) + + def __getstate__(self): + state = self.__dict__.copy() + if not isinstance(self.module, nn.Module): + state["module"] = cloudpickle_dumps(state["module"]) + return state + + def __setstate__(self, state): + if "module" in state: + state["module"] = cloudpickle_loads(state["module"]) + self.__dict__ = state + + +class TensorDictModuleWrapper(TensorDictModuleBase): + """Wrapper class for TensorDictModule objects. + + Once created, a TensorDictModuleWrapper will behave exactly as the + TensorDictModule it contains except for the methods that are + overwritten. + + Args: + td_module (TensorDictModule): operator to be wrapped. + + """ + + def __init__(self, td_module: TensorDictModuleBase) -> None: + super().__init__() + self.td_module = td_module + if len(self.td_module._forward_hooks): + for pre_hook in self.td_module._forward_hooks: + self.register_forward_hook(self.td_module._forward_hooks[pre_hook]) + + def __getattr__(self, name: str) -> Any: + if not is_compiling(): + __dict__ = self.__dict__ + _parameters = __dict__.get("_parameters") + if _parameters: + # A param can be None so we use False instead to check if the key + # is in the _parameters dict once and only once. + # We use False but any non-None, non-Parameter value would do. + # The alternative `if name in _parameters: return _parameters[name]` + # accesses the value of `name` twice when only one is required + result = _parameters.get(name, False) + if result is not False: + return result + _buffers = __dict__.get("_buffers") + if _buffers: + result = _buffers.get(name, False) + if result is not False: + return result + _modules = __dict__.get("_modules") + if _modules: + result = _modules.get(name, False) + if result is not False: + return result + # TODO: find a way to make this check work with dynamo + # elif hasattr(self, "_parameters"): + else: + _parameters = self._parameters + result = _parameters.get(name, False) + if result is not False: + return result + _buffers = self._buffers + result = _buffers.get(name, False) + if result is not False: + return result + _modules = self._modules + result = _modules.get(name, False) + if result is not False: + return result + if name not in self.__dict__ and not name.startswith("__"): + return getattr(self._modules["td_module"], name) + raise AttributeError( + f"attribute {name} not recognised in {type(self).__name__}" + ) + + def forward(self, *args: Any, **kwargs: Any) -> TensorDictBase: + return self.td_module(*args, **kwargs) + + +class WrapModule(TensorDictModuleBase): + """A wrapper around any callable that processes TensorDict instances. + + This wrapper is useful when building :class:`~tensordict.nn.TensorDictSequential` stacks and when a transform + requires the entire TensorDict instance to be visible. + + Args: + func (Callable[[TensorDictBase], TensorDictBase]): A callable function that takes in a TensorDictBase instance + and returns a transformed TensorDictBase instance. + + Keyword Args: + inplace (bool, optional): If ``True``, the input TensorDict will be modified in-place. Otherwise, a new TensorDict + will be returned (if the function does not modify it in-place and returns it). Defaults to ``False``. + in_keys (list of NestedKey, optional): if provided, indicates what entries are read by the module. + This will not be checked and is provided just for the purpose of informing :class:`~tensordict.nn.TensorDictSequential` + about the input keys of the wrapped module. Defaults to `[]`. + out_keys (list of NestedKey, optional): if provided, indicates what entries are written by the module. + This will not be checked and is provided just for the purpose of informing :class:`~tensordict.nn.TensorDictSequential` + about the output keys of the wrapped module. Defaults to `[]`. + + Examples: + >>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod, WrapModule + >>> seq = Seq( + ... Mod(lambda x: x * 2, in_keys=["x"], out_keys=["y"]), + ... WrapModule(lambda td: td.reshape(-1)), + ... ) + >>> td = TensorDict(x=torch.ones(3, 4, 5), batch_size=[3, 4]) + >>> td = Seq(td) + >>> assert td.shape == (12,) + >>> assert (td["y"] == 2).all() + >>> assert td["y"].shape == (12, 5) + + """ + + in_keys = [] + out_keys = [] + + def __init__( + self, + func: Callable[[TensorDictBase], TensorDictBase], + *, + inplace: bool = False, + in_keys: List[NestedKey] | None = None, + out_keys: List[NestedKey] | None = None, + ) -> None: + super().__init__() + self.func = func + self.inplace = inplace + if in_keys is not None: + self.in_keys = in_keys + if out_keys is not None: + self.out_keys = out_keys + + def forward(self, data: TensorDictBase) -> TensorDictBase: + result = self.func(data) + if self.inplace and result is not data: + return data.update(result) + return result + + +class as_tensordict_module: + """A decorator that converts a function into a TensorDictModule. + + Args: + in_keys (List[NestedKey] | NestedKey | None, optional): The input keys of the resulting TensorDictModule. + out_keys (List[NestedKey] | NestedKey | None, optional): The output keys of the resulting TensorDictModule. + + Returns: + Callable: A decorator that can be applied to a function to convert it into a TensorDictModule. + + Examples: + >>> class MyClass: + ... @as_tensordict_module(in_keys="c", out_keys="d") + ... def my_method(self, c): + ... return c + 1 + >>> obj = MyClass() + >>> result = obj.my_method(TensorDict(c=0)) + >>> print(result["d"]) # prints: 1 + >>> @as_tensordict_module(in_keys="c", out_keys="d") + ... def my_function(c): + ... return c + 1 + >>> result = my_function(TensorDict(c=0)) + >>> print(result["d"]) # prints: 1 + """ + + def __init__( + self, + *, + in_keys: List[NestedKey] | NestedKey, + out_keys: List[NestedKey] | NestedKey, + ) -> None: + if isinstance(in_keys, NestedKey): + in_keys = [in_keys] + if isinstance(out_keys, NestedKey): + out_keys = [out_keys] + self.in_keys = in_keys + self.out_keys = out_keys + + def __call__(self, func): + tdmodule = None + is_bound_method = inspect.ismethod(func) or ( + inspect.isfunction(func) and "self" in inspect.signature(func).parameters + ) + + if is_bound_method: + + @wraps(func) + def wrapped(_self, *args, **kwargs): + nonlocal tdmodule + if tdmodule is None: + + def newfunc(*args, **kwargs): + return func(_self, *args, **kwargs) + + tdmodule = TensorDictModule( + newfunc, in_keys=self.in_keys, out_keys=self.out_keys + ) + return tdmodule(*args, **kwargs) + + else: + + @wraps(func) + def wrapped(*args, **kwargs): + nonlocal tdmodule + if tdmodule is None: + tdmodule = TensorDictModule( + func, in_keys=self.in_keys, out_keys=self.out_keys + ) + return tdmodule(*args, **kwargs) + + return wrapped diff --git a/lib/python3.12/site-packages/tensordict/nn/cudagraphs.py b/lib/python3.12/site-packages/tensordict/nn/cudagraphs.py new file mode 100644 index 0000000000000000000000000000000000000000..86c73512a5324d07d3d1535f31a5d037af8900bc --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/cudagraphs.py @@ -0,0 +1,492 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import contextlib +import functools +import os +import warnings +from functools import partial +from textwrap import indent +from typing import Any, Callable, List + +import torch + +from tensordict._nestedkey import NestedKey +from tensordict.base import is_tensor_collection, TensorDictBase +from tensordict.nn.common import dispatch +from tensordict.nn.functional_modules import ( + _exclude_td_from_pytree, + PYTREE_REGISTERED_LAZY_TDS, + PYTREE_REGISTERED_TDS, +) +from tensordict.utils import _zip_strict, logger as tensordict_logger, strtobool +from torch import Tensor + +from torch.utils._pytree import SUPPORTED_NODES, tree_map + +try: + from torch.utils._pytree import tree_flatten, tree_leaves, tree_unflatten +except ImportError: + from torch.utils._pytree import tree_flatten, tree_unflatten + + def tree_leaves(pytree): + """Torch 2.0 compatible version of tree_leaves.""" + return tree_flatten(pytree)[0] + + +class CudaGraphModule: + """A cudagraph wrapper for PyTorch callables. + + ``CudaGraphModule`` is a wrapper class that provides a user-friendly interface to CUDA graphs for PyTorch callables. + + .. warning:: + ``CudaGraphModule`` is a prototype feature and its API restrictions are likely to change in the future. + + This class provides a user-friendly interface to cudagraphs, allowing for a fast, CPU-overhead free execution of + operations on GPU. + It runs essential checks for the inputs to the function, and gives an nn.Module-like API to run + + .. warning:: + This module requires the wrapped function to meet a few requirements. It is the user responsibility + to make sure that all of these are fullfilled. + + - The function cannot have dynamic control flow. For instance, the following code snippet will fail to be + wrapped in `CudaGraphModule`: + + >>> def func(x): + ... if x.norm() > 1: + ... return x + 1 + ... else: + ... return x - 1 + + Fortunately, PyTorch offers solutions in most cases: + + >>> def func(x): + ... return torch.where(x.norm() > 1, x + 1, x - 1) + + - The function must execute a code that can be exactly re-run using the same buffers. This means that + dynamic shapes (changing shape in the input or during the code execution) is not supported. In other words, + the input must have a constant shape. + - The output of the function must be detached. If a call to the optimizers is required, put it in the input + function. For instance, the following function is a valid operator: + + >>> def func(x, y): + ... optim.zero_grad() + ... loss_val = loss_fn(x, y) + ... loss_val.backward() + ... optim.step() + ... return loss_val.detach() + + - The input should not be differntiable. If you need to use `nn.Parameters` (or differentiable tensors in general), + just write a function that uses them as global values rather than passing them as input: + + >>> x = nn.Parameter(torch.randn(())) + >>> optim = Adam([x], lr=1) + >>> def func(): # right + ... optim.zero_grad() + ... (x+1).backward() + ... optim.step() + >>> def func(x): # wrong + ... optim.zero_grad() + ... (x+1).backward() + ... optim.step() + + - Args and kwargs that are tensors or tensordict may change (provided that device and shape match), but non-tensor + args and kwargs should not change. For instance, if the function receives a string input and the input is changed + at any point, the module will silently execute the code with the string used during the capture of the cudagraph. + The only supported keyword argument is `tensordict_out` in case the input is a tensordict. + + - If the module is a :class:`~tensordict.nn.TensorDictModuleBase` instance and the output id matches the input + id, then this identity will be preserved during a call to ``CudaGraphModule``. In all other cases, the output + will be cloned, irrespective of whether its elements match or do not match one of the inputs. + + .. warning:: + ``CudaGraphModule`` is not an :class:`~torch.nn.Module` by design, to discourage gathering parameters + of the input module and passing them to an optimizer. + + Args: + module (Callable): a function that receives tensors (or tensordict) as input and outputs a + PyTreeable collection of tensors. If a tensordict is provided, the module can be run with keyword arguments + too (see example below). + warmup (int, optional): the number of warmup steps in case the module is compiled (compiled modules should be + run a couple of times before being captured by cudagraphs). Defaults to ``2`` for all modules. + in_keys (list of NestedKeys): the input keys, if the module takes a TensorDict as input. + Defaults to ``module.in_keys`` if this value exists, otherwise ``None``. + + .. note:: + If ``in_keys`` is provided but empty, the module is assumed to receive a tensordict as input. + This is sufficient to make ``CudaGraphModule`` aware that the function should be treated as a + `TensorDictModule`, but keyword arguments will not be dispatched. See below for some examples. + + out_keys (list of NestedKeys): the output keys, if the module takes and outputs TensorDict as output. + Defaults to ``module.out_keys`` if this value exists, otherwise ``None``. + device (torch.device, optional): the device of the stream to use. + + Examples: + >>> # Wrap a simple function + >>> def func(x): + ... return x + 1 + >>> func = CudaGraphModule(func) + >>> x = torch.rand((), device='cuda') + >>> out = func(x) + >>> assert isinstance(out, torch.Tensor) + >>> assert out == x+1 + >>> # Wrap a tensordict module + >>> func = TensorDictModule(lambda x: x+1, in_keys=["x"], out_keys=["y"]) + >>> func = CudaGraphModule(func) + >>> # This can be called either with a TensorDict or regular keyword arguments alike + >>> y = func(x=x) + >>> td = TensorDict(x=x) + >>> td = func(td) + + .. note:: Tips on debugging CudaGraphModule errors: + + - Errors of the sort of `operation would make the legacy stream depend on a capturing blocking stream` are + to be debugged using a non-compiled version of your code first (compiled code will hide the part of the + code that is responsible for the cross-stream dependency). It can happen because you are + doing inter-device operations which make the capturing stream dependent on other streams. + - Errors like `Cannot call CUDAGeneratorImpl::current_seed during CUDA graph capture` or other errors orignating + from the compiler (not the compiled code!) probably point to a capture of the graph while recompiling. + Capture the recompiles with `TORCH_LOGS="+recompiles" python myscrip.py` and try to fix these. + In general, make sure you are using enough warmup steps. + Please file an issue if you're struggling with these. + + """ + + _REQUIRES_GRAD_ERROR = ( + "CudaGraphModule cannot be part of a graph (or leaves variables). If you need tensors to " + "require gradients, please pass them as constant inputs to the module instead (e.g. " + "`def func(a, b, c=c): ...` where c is the tensor requiring gradients)." + ) + + def __init__( + self, + module: Callable[[List[Tensor] | TensorDictBase], None], + warmup: int = 2, + in_keys: List[NestedKey] = None, + out_keys: List[NestedKey] = None, + device: torch.device | None = None, + ): + self._has_cuda = torch.cuda.is_available() + if not self._has_cuda: + warnings.warn( + "This module is instantiated on a machine without CUDA device available. " + "We permit this usage, but calls to this instance will just pass through it and execute " + "the provided function without additional performance gain." + ) + self.module = module + self.counter = 0 + if not isinstance(warmup, int) or warmup < 1: + raise ValueError("warmup must be an integer greater than 0.") + self._warmup = warmup + if torch.cuda.is_available(): + self._warmup_stream = torch.cuda.Stream(device) + self._capture_stream = torch.cuda.Stream(device) + self._warmup_stream_cm = partial(torch.cuda.stream, self._warmup_stream) + else: + self._warmup_stream = None + self._warmup_stream_cm = contextlib.nullcontext + self.device = device + tensordict_logger.info( + f"Setting up CudaGraphModule with stream {self._warmup_stream} on device {self.device}." + ) + + if hasattr(module, "in_keys"): + self.in_keys = module.in_keys + else: + self.in_keys = in_keys + if hasattr(module, "out_keys"): + self.out_keys = module.out_keys + else: + if out_keys is None and self.in_keys is not None: + out_keys = [] + self.out_keys = out_keys + self._is_tensordict_module = self.in_keys is not None + self._out_matches_in = None + for tdtype in PYTREE_REGISTERED_TDS + PYTREE_REGISTERED_LAZY_TDS: + if tdtype in SUPPORTED_NODES: + if not strtobool(os.environ.get("EXCLUDE_TD_FROM_PYTREE", "0")): + warnings.warn( + f"Tensordict is registered in PyTree. This is incompatible with {self.__class__.__name__}. " + f"Removing TDs from PyTree. To silence this warning, call tensordict.nn.functional_modules._exclude_td_from_pytree().set() " + f"or set the environment variable `EXCLUDE_TD_FROM_PYTREE=1`. " + f"This operation is irreversible." + ) + _exclude_td_from_pytree().set() + + if self._is_tensordict_module: + + @dispatch(source=self.in_keys, dest=self.out_keys, auto_batch_size=False) + def _call( + tensordict: TensorDictBase, + *args, + tensordict_out: TensorDictBase | None = None, + **kwargs: Any, + ) -> Any: + if self.counter >= self._warmup: + self._tensordict.update_(tensordict, non_blocking=True) # type: ignore[attr-defined] + torch.cuda.synchronize(self.device) + self.graph.replay() + if self._out_matches_in: + result = tensordict.update( # type-ignore[unreachable] + self._out, keys_to_update=self._selected_keys + ) + elif tensordict_out is not None: + result = tensordict_out.update(self._out, clone=True) + else: + result = self._out.clone() if self._out is not None else None + return result + + if not self._has_cuda or self.counter < self._warmup - 1: + # We must clone the data because providing non-contiguous data will fail later when we clone + if self._has_cuda: + self._warmup_stream.wait_stream(torch.cuda.current_stream()) + with self._warmup_stream_cm(): + tensordict.apply(self._clone, out=tensordict) + if tensordict_out is not None: + kwargs["tensordict_out"] = tensordict_out + out = self.module(tensordict, *args, **kwargs) + if self._out_matches_in is None: + self._out_matches_in = out is tensordict + if self._has_cuda: + torch.cuda.current_stream().wait_stream(self._warmup_stream) + self.counter += self._has_cuda + return out + else: + tensordict_logger.info("Registering CUDA graph...") + if tensordict.device is None: + tensordict.apply(self._check_device_and_grad, filter_empty=True) + elif tensordict.device.type != "cuda": + raise ValueError( + "The input tensordict device must be of the 'cuda' type." + ) + + tree_map(self._check_non_tensor, (args, kwargs)) + + torch.cuda.synchronize(self.device) + kwargs = kwargs.copy() + self._warmup_stream.wait_stream( + torch.cuda.current_stream(self.device) + ) + with self._warmup_stream_cm(): + tensordict.apply(self._clone, out=tensordict) + self._tensordict = tensordict.copy() + if tensordict_out is not None: + td_out_save = tensordict_out.copy() + kwargs["tensordict_out"] = tensordict_out + this_out = self.module(tensordict, *args, **kwargs) + self._capture_stream.wait_stream(self._warmup_stream) + + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph, stream=self._capture_stream): + if tensordict_out is not None: + kwargs["tensordict_out"] = td_out_save + out = self.module(self._tensordict, *args, **kwargs) + tensordict_logger.info("CUDA graph successfully registered.") + + if not is_tensor_collection(out) and out is not None: + raise RuntimeError( + "The output of the function must be a tensordict, a tensorclass or None. Got " + f"type(out)={type(out)}." + ) + if is_tensor_collection(out): + out.lock_() + self._out = out + self.counter += 1 + if self._out_matches_in: + # We need to know what keys to update in out + if self.out_keys: + self._selected_keys = self.out_keys + else: + # Gather keys that have changed during execution + self._selected_keys = [] + + def check_tensor_id(name, t0, t1): + if t0 is not t1: + self._selected_keys.append(name) + + self._out.named_apply( + check_tensor_id, + tensordict, + default=None, + filter_empty=True, + ) + return this_out + + else: + + def _call(*args: torch.Tensor, **kwargs: torch.Tensor): + if self.counter >= self._warmup: + srcs, dests = [], [] + for arg_src, arg_dest in _zip_strict( + tree_leaves((args, kwargs)), self._flat_tree + ): + self._maybe_copy_onto_(arg_src, arg_dest, srcs, dests) + if dests: + torch._foreach_copy_(dests, srcs) + torch.cuda.synchronize(self.device) + self.graph.replay() + if self._return_unchanged: + result = self._out + else: + result = tree_unflatten( + [ + out.clone() if hasattr(out, "clone") else out + for out in self._out + ], + self._out_struct, + ) + return result + + if not self._has_cuda or self.counter < self._warmup - 1: + args, kwargs = tree_map(self._clone, (args, kwargs)) + if self._has_cuda: + self._warmup_stream.wait_stream( + torch.cuda.current_stream(self.device) + ) + with self._warmup_stream_cm(): + out = self.module(*args, **kwargs) + if self._has_cuda: + torch.cuda.current_stream(self.device).wait_stream( + self._warmup_stream + ) + self.counter += self._has_cuda + return out + else: + tensordict_logger.info("Registering CUDA graph...") + self._flat_tree, self._tree_spec = tree_flatten((args, kwargs)) + + self._flat_tree = tuple( + self._check_device_and_clone(arg) for arg in self._flat_tree + ) + args, kwargs = self._args, self._kwargs = tree_unflatten( + self._flat_tree, self._tree_spec + ) + + self._warmup_stream.wait_stream( + torch.cuda.current_stream(self.device) + ) + with self._warmup_stream_cm(): + this_out = self.module(*args, **kwargs) + self._capture_stream.wait_stream(self._warmup_stream) + + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph, stream=self._capture_stream): + out = self.module(*self._args, **self._kwargs) + tensordict_logger.info("CUDA graph successfully registered.") + self._out, self._out_struct = tree_flatten(out) + self.counter += 1 + # Check that there is not intersection between the indentity of inputs and outputs, otherwise warn + # user. + tree_leaves_input = tree_leaves((args, kwargs)) + tree_leaves_output = tree_leaves(out) + if not isinstance(tree_leaves_output, tuple): + tree_leaves_output = (tree_leaves_output,) + for inp in tree_leaves_input: + if isinstance(inp, torch.Tensor) and inp in tree_leaves_output: + warnings.warn( + "An input tensor is also present in the output of the wrapped function. " + f"{type(self).__name__} will not copy tensors in place: the output will be cloned " + f"and the identity between input and output will not match anymore. " + f"Make sure you don't rely on input-output identity further in the code." + ) + if not self._out: + self._return_unchanged = True + else: + self._out = [ # type: ignore[unreachable] + out.lock_() if is_tensor_collection(out) else out + for out in self._out + ] + self._return_unchanged = False + return this_out + + _call_func = functools.wraps(self.module)(_call) + self._call_func = _call_func + + @staticmethod + def _maybe_copy_onto_(src, dest, srcs, dests): + if isinstance(src, torch.Tensor): + srcs.append(src) + dests.append(dest) + return + if is_tensor_collection(src): + dest.copy_(src) + return + isdiff = False + try: + isdiff = src != dest + except Exception as err: + raise RuntimeError( + "Couldn't assess input value. Make sure your function only takes tensor inputs or that " + "the input value can be easily checked and is constant. For a better efficiency, avoid " + "passing non-tensor inputs to your function." + ) from err + if isdiff: + raise ValueError("Varying inputs must be torch.Tensor subclasses.") + + @classmethod + def _check_device_and_clone(cls, x): + if isinstance(x, torch.Tensor) or is_tensor_collection(x): + if x.requires_grad: + raise RuntimeError(cls._REQUIRES_GRAD_ERROR) + if x.device is None: + # Check device of leaves of tensordict + x.apply(cls._check_device_and_grad, filter_empty=True) + + elif x.device.type != "cuda": + raise ValueError( + f"All tensors must be stored on CUDA. Got {x.device.type}." + ) + + return x.clone() + return x + + @classmethod + def _clone(cls, x): + if isinstance(x, torch.Tensor) or is_tensor_collection(x): + if x.requires_grad: + raise RuntimeError(cls._REQUIRES_GRAD_ERROR) + return x.clone() + return x + + @classmethod + def _check_device_and_grad(cls, x): + if isinstance(x, torch.Tensor): + if x.device.type != "cuda": + raise ValueError( + f"All tensors must be stored on CUDA. Got {x.device.type}." + ) + if x.requires_grad: + raise RuntimeError(cls._REQUIRES_GRAD_ERROR) + + @staticmethod + def _check_non_tensor(arg): + if isinstance(arg, torch.Tensor): + raise ValueError( + "All tensors must be passed in the tensordict, not as arg or kwarg." + ) + + def __call__(self, *args, **kwargs): + return self._call_func(*args, **kwargs) + + def __repr__(self): + module = indent(f"module={self.module}", 4 * " ") + warmup = warmup = {self._warmup} + in_keys = indent(f"in_keys={self.in_keys}", 4 * " ") + out_keys = indent(f"out_keys={self.out_keys}", 4 * " ") + return f"{self.__class__.__name__}(\n{module}, \n{warmup}, \n{in_keys}, \n{out_keys}\n)" + + def state_dict(self): + try: + return self.module.state_dict() + except AttributeError: + return {} + + def load_state_dict(self, state_dict): + try: + self.module.load_state_dict(state_dict) + except AttributeError: + pass diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/__init__.py b/lib/python3.12/site-packages/tensordict/nn/distributions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c385805e125c847c4d18af6288bd9fec4f832fc5 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/distributions/__init__.py @@ -0,0 +1,43 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from tensordict.nn.distributions import continuous, discrete + +from tensordict.nn.distributions.composite import CompositeDistribution +from tensordict.nn.distributions.continuous import ( + AddStateIndependentNormalScale, + Delta, + NormalParamExtractor, +) +from tensordict.nn.distributions.discrete import OneHotCategorical, rand_one_hot +from tensordict.nn.distributions.truncated_normal import TruncatedNormal +from tensordict.nn.probabilistic import InteractionType, set_interaction_type +from tensordict.nn.utils import add_custom_mapping, mappings + +distributions_maps = { + distribution_class.lower(): eval(distribution_class) + for distribution_class in (*continuous.__all__, *discrete.__all__) +} + +__all__ = [ + # Distribution classes + "CompositeDistribution", + "AddStateIndependentNormalScale", + "Delta", + "NormalParamExtractor", + "OneHotCategorical", + "rand_one_hot", + "TruncatedNormal", + # Interaction types + "InteractionType", + "set_interaction_type", + # Submodules + "continuous", + "discrete", + # Utilities + "distributions_maps", + "add_custom_mapping", + "mappings", +] diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d38f84295bf2b2cf160b681564998c6ff7d43e5 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/composite.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/composite.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a31750cc2e32b6f2604ab22e891550936b81d883 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/composite.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/continuous.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/continuous.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c3d77094e8534a1a946e628de94d27d88f614a5e Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/continuous.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/discrete.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/discrete.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a14c6a52cb6aea318c32afc980ed74f954fb3fe1 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/discrete.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/truncated_normal.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/truncated_normal.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a47b7919a87b740ca405f495322bf7eb487cde2d Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/truncated_normal.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e26839c7c0fd16cf31946e9cbfad95d636ec7957 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/nn/distributions/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/composite.py b/lib/python3.12/site-packages/tensordict/nn/distributions/composite.py new file mode 100644 index 0000000000000000000000000000000000000000..b4100dedc70bc2ef654fbf4a58bcaf4a70a4641a --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/distributions/composite.py @@ -0,0 +1,566 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import warnings +from collections.abc import Mapping +from typing import Any, Dict, TYPE_CHECKING + +import torch +from tensordict._td import TensorDict +from tensordict.base import TensorDictBase +from tensordict.nn.utils import composite_lp_aggregate, set_composite_lp_aggregate +from tensordict.utils import IndexType, NestedKey, unravel_key, unravel_keys +from torch import distributions as d + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +class CompositeDistribution(d.Distribution, Mapping): + """A composite distribution that groups multiple distributions together using the TensorDict interface. + + This class allows for operations such as `log_prob_composite`, `entropy_composite`, `cdf`, `icdf`, `rsample`, and `sample` + to be performed on a collection of distributions, returning a TensorDict. The input TensorDict may be modified in-place. + + Args: + params (TensorDictBase): A nested key-tensor map where the root entries correspond to sample names, and the leaves + are the distribution parameters. Entry names must match those specified in `distribution_map`. + distribution_map (Dict[NestedKey, Type[torch.distribution.Distribution]]): Specifies the distribution types to be used. + The names of the distributions should match the sample names in the `TensorDict`. + + Keyword Arguments: + name_map (Dict[NestedKey, NestedKey], optional): A mapping of where each sample should be written. If not provided, + the key names from `distribution_map` will be used. + extra_kwargs (Dict[NestedKey, Dict], optional): A dictionary of additional keyword arguments for constructing the distributions. + log_prob_key (NestedKey, optional): The key where the aggregated log probability will be stored. + Defaults to `'sample_log_prob'`. + + .. note:: if :func:`tensordict.nn.probabilistic.composite_lp_aggregate` returns ``False``, tbe log-probabilities will + be written under `("path", "to", "leaf", "_log_prob")` + where `("path", "to", "leaf", "")` is the :class:`~tensordict.NestedKey` corresponding to + the leaf tensor being sampled. In that case, the ``log_prob_key`` argument will be ignored. + + entropy_key (NestedKey, optional): The key where the entropy will be stored. Defaults to `'entropy'` + + .. note:: if :func:`tensordict.nn.probabilistic.composite_lp_aggregate` returns ``False``, tbe entropies will + be written under `("path", "to", "leaf", "_entropy")` + where `("path", "to", "leaf", "")` is the :class:`~tensordict.NestedKey` corresponding to + the leaf tensor being sampled. In that case, the ``entropy_key`` argument will be ignored. + + .. note:: The batch size of the input TensorDict containing the parameters (`params`) determines the batch shape of + the distribution. For example, the `"sample_log_prob"` entry resulting from a call to `log_prob` will have the + shape of the parameters plus any additional batch dimensions. + + .. seealso:: :class:`~tensordict.nn.ProbabilisticTensorDictModule` and :class:`~tensordict.nn.ProbabilisticTensorDictSequential` + to learn how to use this class as part of a model. + + .. seealso:: :class:`~tensordict.nn.set_composite_lp_aggregate` to control the aggregation of the log-probabilities. + + Examples: + >>> params = TensorDict({ + ... "cont": {"loc": torch.randn(3, 4), "scale": torch.rand(3, 4)}, + ... ("nested", "disc"): {"logits": torch.randn(3, 10)} + ... }, [3]) + >>> dist = CompositeDistribution(params, + ... distribution_map={"cont": d.Normal, ("nested", "disc"): d.Categorical}) + >>> sample = dist.sample((4,)) + >>> with set_composite_lp_aggregate(False): + ... sample = dist.log_prob(sample) + ... print(sample) + TensorDict( + fields={ + cont: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + cont_log_prob: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + nested: TensorDict( + fields={ + disc: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.int64, is_shared=False), + disc_log_prob: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4]), + device=None, + is_shared=False) + + """ + + def __init__( + self, + params: TensorDictBase, + distribution_map: dict, + *, + name_map: dict | None = None, + extra_kwargs=None, + log_prob_key: NestedKey | None = None, + entropy_key: NestedKey | None = None, + ): + self._batch_shape = params.shape + if extra_kwargs is None: + extra_kwargs = {} + dists = {} + if name_map is not None: + name_map = { + unravel_key(key): unravel_key(other_key) + for key, other_key in name_map.items() + } + for name, dist_class in distribution_map.items(): + name_unravel = unravel_key(name) + if name_map: + try: + write_name = unravel_key(name_map.get(name, name_unravel)) + except KeyError: + raise KeyError( + f"Failed to retrieve the key {name} from the name_map with keys {name_map.keys()}." + ) + else: + write_name = name_unravel + name = name_unravel + dist_params = params.get(name) + kwargs = extra_kwargs.get(name, {}) + if dist_params is None: + raise KeyError( + f"no param {name} found in params with keys {params.keys(True, True)}" + ) + dist = dist_class(**dist_params, **kwargs) + dists[write_name] = dist + self.dists = dists + self.log_prob_key = log_prob_key + self.entropy_key = entropy_key + + def __iter__(self): + yield from self.dists + + def __getitem__(self, item: IndexType) -> Self: + return self.dists[item] + + def __len__(self): + return len(self.dists) + + @property + def log_prob_key(self): + log_prob_key = self._log_prob_key + if log_prob_key is None: + log_prob_key = "sample_log_prob" + return log_prob_key + + @log_prob_key.setter + def log_prob_key(self, value): + self._log_prob_key = value + + @property + def entropy_key(self): + entropy_key = self._entropy_key + if entropy_key is None: + entropy_key = "entropy" + return entropy_key + + @entropy_key.setter + def entropy_key(self, value): + self._entropy_key = value + + @classmethod + def from_distributions( + cls, + params, + distributions: Dict[NestedKey, d.Distribution], + *, + name_map: dict | None = None, + log_prob_key: NestedKey | None = None, + entropy_key: NestedKey | None = None, + ) -> CompositeDistribution: + """Create a `CompositeDistribution` instance from existing distribution objects. + + This class method allows for the creation of a `CompositeDistribution` by directly providing + a dictionary of distribution instances, rather than specifying distribution types and parameters separately. + + Args: + params (TensorDictBase): A TensorDict that defines the batch shape for the composite distribution. + The params will not be used by this method, but the tensordict will be used to gather the key names of + the distributions. + distributions (Dict[NestedKey, d.Distribution]): A dictionary mapping nested keys to distribution instances. + These distributions will be used directly in the composite distribution. + + Keyword Args: + name_map (Dict[NestedKey, NestedKey], optional): A mapping of where each sample should be written. If not provided, + the key names from `distribution_map` will be used. + log_prob_key (NestedKey, optional): The key where the log probability will be stored. + Defaults to `'sample_log_prob'`. + entropy_key (NestedKey, optional): The key where the entropy will be stored. Defaults to `'entropy'`. + + Returns: + CompositeDistribution: An instance of `CompositeDistribution` initialized with the provided distributions. + + Raises: + KeyError: If a key in `name_map` cannot be found in the provided distributions. + + .. note:: The batch size of the `params` TensorDict determines the batch shape of the composite distribution. + + Example: + >>> from tensordict.nn import CompositeDistribution, ProbabilisticTensorDictSequential, ProbabilisticTensorDictModule, TensorDictModule + >>> import torch + >>> from tensordict import TensorDict + >>> + >>> # Values are not used to build the dists + >>> params = TensorDict({("0", "loc"): None, ("1", "loc"): None, ("0", "scale"): None, ("1", "scale"): None}) + >>> d0 = torch.distributions.Normal(0, 1) + >>> d1 = torch.distributions.Normal(torch.zeros(1, 2), torch.ones(1, 2)) + >>> + >>> d = CompositeDistribution.from_distributions(params, {"0": d0, "1": d1}) + >>> print(d.sample()) + TensorDict( + fields={ + 0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False), + 1: Tensor(shape=torch.Size([1, 2]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + """ + self = cls.__new__(cls) + self._batch_shape = params.shape + dists = {} + if name_map is not None: + name_map = { + unravel_key(key): unravel_key(other_key) + for key, other_key in name_map.items() + } + for name, dist in distributions.items(): + name_unravel = unravel_key(name) + if name_map: + try: + write_name = unravel_key(name_map.get(name, name_unravel)) + except KeyError: + raise KeyError( + f"Failed to retrieve the key {name} from the name_map with keys {name_map.keys()}." + ) + else: + write_name = name_unravel + dists[write_name] = dist + self.dists = dists + self.log_prob_key = log_prob_key + self.entropy_key = entropy_key + self.aggregate_probabilities = None + + return self + + def sample(self, shape=None) -> TensorDictBase: + if shape is None: + shape = torch.Size([]) + samples = {name: dist.sample(shape) for name, dist in self.dists.items()} + return TensorDict( + samples, + shape + self.batch_shape, + ) + + @property + def mode(self) -> TensorDictBase: + samples = {name: dist.mode for name, dist in self.dists.items()} + return TensorDict( + samples, + self.batch_shape, + ) + + @property + def mean(self) -> TensorDictBase: + samples = {name: dist.mean for name, dist in self.dists.items()} + return TensorDict( + samples, + self.batch_shape, + ) + + @property + def deterministic_sample(self) -> TensorDictBase: + def maybe_deterministic_sample(dist): + if hasattr(dist, "deterministic_sample"): + return dist.deterministic_sample + else: + from tensordict.nn.probabilistic import DETERMINISTIC_REGISTER + + # Fallbacks + tdist = type(dist) + if issubclass(tdist, d.Independent): + tdist = type(dist.base_dist) + interaction_type = DETERMINISTIC_REGISTER.get(tdist) + if interaction_type == "mode": + return dist.mode + if interaction_type == "mean": + return dist.mean + if interaction_type == "random": + return dist.rsample() if dist.has_rsample else dist.sample() + if interaction_type is None: + try: + support = dist.support + fallback = ( + "mean" + if isinstance(support, d.constraints._Real) + else "mode" + ) + except NotImplementedError: + # Some custom dists don't have a support + # We arbitrarily fall onto 'mean' in these cases + fallback = "mean" + else: + raise RuntimeError( + f"InteractionType {interaction_type} is unaccounted for." + ) + try: + if fallback == "mean": + return dist.mean + elif fallback == "mode": + return dist.mode + else: + raise AttributeError + except AttributeError: + raise NotImplementedError( + f"method {type(dist)}.deterministic_sample is not implemented, no replacement found." + ) + finally: + warnings.warn( + f"deterministic_sample wasn't found when queried on {type(dist)}. " + f"{type(self).__name__} is falling back on {fallback} instead. " + f"For better code quality and efficiency, make sure to either " + f"provide a distribution with a deterministic_sample attribute or " + f"to change the InteractionMode to the desired value.", + category=UserWarning, + ) + + samples = { + name: maybe_deterministic_sample(dist) for name, dist in self.dists.items() + } + return TensorDict( + samples, + self.batch_shape, + ) + + def __repr__(self): + return f"{type(self).__name__}({self.dists})" + + def rsample(self, shape=None) -> TensorDictBase: + if shape is None: + shape = torch.Size([]) + return TensorDict( + {name: dist.rsample(shape) for name, dist in self.dists.items()}, + shape + self.batch_shape, + ) + + def log_prob( + self, sample: TensorDictBase + ) -> torch.Tensor | TensorDictBase: # noqa: D417 + """Compute the summed log-probability of a given sample. + + Args: + sample (TensorDictBase): The input sample to compute the log probability for. + + If ``self.aggregate_probabilities`` is ``True``, this method will return a single tensor with + the summed log-probabilities. If ``self.aggregate_probabilities`` is ``False``, this method will + call the `:meth:`~.log_prob_composite` method and return a tensordict with the log-probabilities + of each sample in the input tensordict along with a ``sample_log_prob`` entry with the summed + log-prob. In both cases, the output shape will be the shape of the input tensordict. + """ + aggregate_probabilities = composite_lp_aggregate() + if not aggregate_probabilities: + with set_composite_lp_aggregate(False): + return self.log_prob_composite(sample) + slp = 0.0 + for name, dist in self.dists.items(): + local_sample = sample.get(name) + lp = dist.log_prob(local_sample) + if lp.ndim > sample.ndim: + lp = lp.flatten(sample.ndim, -1).sum(-1) + slp = slp + lp + return slp + + def log_prob_composite( + self, + sample: TensorDictBase, + *, + include_sum: bool | None = None, + ) -> TensorDictBase: + """Computes the log-probability of each component in the input sample and return a TensorDict with individual log-probabilities. + + Args: + sample (TensorDictBase): The input sample to compute the log probabilities for. + + Keyword Args: + include_sum (bool, optional): Whether to include the summed log-probability in the output TensorDict. + Defaults to ``composite_lp_aggregate()`` (``False`` by default). + + Returns: + TensorDictBase: A TensorDict containing the individual log-probabilities for each component in the input sample, + along with a "sample_log_prob" entry containing the summed log-probability if `include_sum` is True. + """ + if include_sum is None: + include_sum = composite_lp_aggregate() + inplace = composite_lp_aggregate() + + if include_sum: + slp = 0.0 + d = {} + for name, dist in self.dists.items(): + try: + d[_add_suffix(name, "_log_prob")] = lp = dist.log_prob(sample.get(name)) + except AttributeError: + raise RuntimeError( + f"Expected a tensordict sample, but got a {type(sample).__name__} instead." + ) + if include_sum: + if lp.ndim > sample.ndim: + lp = lp.flatten(sample.ndim, -1).sum(-1) + slp = slp + lp + if include_sum: + d[self.log_prob_key] = slp + if inplace: + return sample.update(d) + else: + return sample.empty(recurse=True).update(d).filter_empty_() + + def entropy( + self, + samples_mc: int = 1, + *, + aggregate_probabilities: bool | None = None, + include_sum: bool | None = None, + ) -> torch.Tensor | TensorDictBase: # noqa: D417 + """Computes and returns the entropy of the composite distribution. + + This method calculates the entropy for each component distribution and optionally sums them. + + Args: + samples_mc (int): The number of samples to draw if the entropy does not have a closed-form solution. + Defaults to `1`. + + Keyword Args: + include_sum (bool, optional): Whether to include the summed entropy in the output TensorDict. + Defaults to `composite_lp_aggregate()`, which defaults to `False`. + + Returns: + torch.Tensor or TensorDictBase: If `aggregate_probabilities` is `True`, returns a single tensor with + the summed entropies. If `aggregate_probabilities` is `False`, returns a TensorDict with the entropies + of each component distribution. + + .. note:: If a distribution does not implement a closed-form solution for entropy, Monte Carlo sampling is used + to estimate it. + """ + aggregate_probabilities = composite_lp_aggregate() + if include_sum is None: + include_sum = composite_lp_aggregate() + + if not aggregate_probabilities: + return self.entropy_composite(samples_mc, include_sum=include_sum) + se = 0.0 + for dist in self.dists.values(): + try: + e = dist.entropy() + except NotImplementedError: + x = dist.rsample((samples_mc,)) + e = -dist.log_prob(x).mean(0) + if e.ndim > len(self.batch_shape): + e = e.flatten(len(self.batch_shape), -1).sum(-1) + se = se + e + return se + + def entropy_composite( + self, + samples_mc=1, + *, + include_sum: bool | None = None, + ) -> TensorDictBase: + """Computes the entropy for each component distribution and returns a TensorDict with individual entropies. + + This method is used by the `entropy` method when `self.aggregate_probabilities` is `False`. + + Args: + samples_mc (int): The number of samples to draw if the entropy does not have a closed-form solution. + Defaults to `1`. + + Keyword Args: + include_sum (bool, optional): Whether to include the summed entropy in the output TensorDict. + Defaults to `composite_lp_aggregate()`, which is defaults to `False`. + + Returns: + TensorDictBase: A TensorDict containing the individual entropies for each component distribution, + along with an "entropy" entry containing the summed entropies if `include_sum` is `True`. + + .. note:: If a distribution does not implement a closed-form solution for entropy, Monte Carlo sampling is used + to estimate it. + """ + if include_sum is None: + include_sum = composite_lp_aggregate() + + se = 0.0 + d = {} + for name, dist in self.dists.items(): + try: + e = dist.entropy() + except NotImplementedError: + x = dist.rsample((samples_mc,)) + e = -dist.log_prob(x).mean(0) + d[_add_suffix(name, "_entropy")] = e + if include_sum: + if e.ndim > len(self.batch_shape): + e = e.flatten(len(self.batch_shape), -1).sum(-1) + se = se + e + if include_sum: + d[self.entropy_key] = se + return TensorDict( + d, + self.batch_shape, + ) + + def cdf(self, sample: TensorDictBase) -> TensorDictBase: + """Computes the cumulative distribution function (CDF) for each component distribution in the composite distribution. + + This method calculates the CDF for each component distribution and updates the input TensorDict with the results. + + Args: + sample (TensorDictBase): A TensorDict containing samples for which to compute the CDF. + + Returns: + TensorDictBase: The input TensorDict updated with `_cdf` entries for each component distribution. + """ + cdfs = { + _add_suffix(name, "_cdf"): dist.cdf(sample.get(name)) + for name, dist in self.dists.items() + } + sample.update(cdfs) + return sample + + def icdf(self, sample: TensorDictBase) -> TensorDictBase: + """Computes the inverse cumulative distribution function (inverse CDF) for each component distribution. + + This method requires the input TensorDict to have either a `_cdf` entry or a `` entry + for each component distribution. It calculates the inverse CDF and updates the TensorDict with the results. + + Args: + sample (TensorDictBase): A TensorDict containing either `_cdf` or `` entries + for each component distribution. + + Returns: + TensorDictBase: The input TensorDict updated with `_icdf` entries for each component distribution. + + Raises: + KeyError: If neither `` nor `_cdf` can be found in the input TensorDict for a component distribution. + """ + for name, dist in self.dists.items(): + prob = sample.get(_add_suffix(name, "_cdf")) + if prob is None: + try: + prob = self.cdf(sample.get(name)) + except KeyError: + raise KeyError( + f"Neither {name} nor {name + '_cdf'} could be found in the sampled tensordict. Make sure one of these is available to icdf." + ) + icdf = dist.icdf(prob) + sample.set(_add_suffix(name, "_icdf"), icdf) + return sample + + +def _add_suffix(key: NestedKey, suffix: str): + key = unravel_keys(key) + if isinstance(key, str): + return key + suffix + return key[:-1] + (key[-1] + suffix,) diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/continuous.py b/lib/python3.12/site-packages/tensordict/nn/distributions/continuous.py new file mode 100644 index 0000000000000000000000000000000000000000..b2c7cb2a21160ddc21851092b71e2bcf0f501b5b --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/distributions/continuous.py @@ -0,0 +1,286 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from numbers import Number +from typing import Sequence + +import numpy as np + +import torch + +from tensordict.nn.utils import mappings +from torch import distributions as D, nn + +# We need this to build the distribution maps +__all__ = [ + "NormalParamExtractor", + "AddStateIndependentNormalScale", + "Delta", +] + +# speeds up distribution construction +# D.Distribution.set_default_validate_args(False) + + +class NormalParamWrapper(nn.Module): + def __init__( + self, + operator: nn.Module, + scale_mapping: str = "biased_softplus_1.0", + scale_lb: Number = 1e-4, + ) -> None: + raise RuntimeError( + "NormalParamWrapper has been deprecated in favor of `tensordict.nn.NormalParamExtractor`. Use this class instead." + ) + + +class NormalParamExtractor(nn.Module): + """A non-parametric nn.Module that splits its input into loc and scale parameters. + + The scale parameters are mapped onto positive values using the specified ``scale_mapping``. + + Args: + scale_mapping (str, optional): positive mapping function to be used with the std. + default = ``"biased_softplus_1.0"`` (i.e. softplus map with bias such that fn(0.0) = 1.0) + choices: ``"softplus"``, ``"exp"``, ``"relu"``, ``"biased_softplus_1"`` or ``"none"`` (no mapping). + See :func:`~tensordict.nn.mappings` for more details. + scale_lb (Number, optional): The minimum value that the variance can take. Default is 1e-4. + + Examples: + >>> import torch + >>> from tensordict.nn.distributions import NormalParamExtractor + >>> from torch import nn + >>> module = nn.Linear(3, 4) + >>> normal_params = NormalParamExtractor() + >>> tensor = torch.randn(3) + >>> loc, scale = normal_params(module(tensor)) + >>> print(loc.shape, scale.shape) + torch.Size([2]) torch.Size([2]) + >>> assert (scale > 0).all() + >>> # with modules that return more than one tensor + >>> module = nn.LSTM(3, 4) + >>> tensor = torch.randn(4, 2, 3) + >>> loc, scale, others = normal_params(*module(tensor)) + >>> print(loc.shape, scale.shape) + torch.Size([4, 2, 2]) torch.Size([4, 2, 2]) + >>> assert (scale > 0).all() + + """ + + def __init__( + self, + scale_mapping: str = "biased_softplus_1.0", + scale_lb: Number = 1e-4, + ) -> None: + super().__init__() + self.scale_mapping = mappings(scale_mapping) + self.scale_lb = scale_lb + + def forward(self, *tensors: torch.Tensor) -> tuple[torch.Tensor, ...]: + tensor, *others = tensors + loc, scale = tensor.chunk(2, -1) + scale = self.scale_mapping(scale).clamp_min(self.scale_lb) + return (loc, scale, *others) + + +class AddStateIndependentNormalScale(torch.nn.Module): + """A nn.Module that adds trainable state-independent scale parameters. + + The scale parameters are mapped onto positive values using the specified ``scale_mapping``. + + Args: + scale_shape (torch.Size or equivalent, optional): the shape of the scale parameter. + Defaults to ``torch.Size(())``. + + Keyword Args: + scale_mapping (str, optional): positive mapping function to be used with the std. + Defaults to ``"exp"``, + choices: ``"softplus"``, ``"exp"``, ``"relu"``, ``"biased_softplus_1"``. + scale_lb (Number, optional): The minimum value that the variance can take. + Defaults to ``1e-4``. + device (torch.device, optional): the device of the module. + make_param (bool, optional): whether the scale should be a parameter (``True``) + or a buffer (``False``). + Defaults to ``True``. + init_value (float, optional): Initial value of state independent scale. + Defaults to 0.0. + + Examples: + >>> from torch import nn + >>> import torch + >>> num_outputs = 4 + >>> module = nn.Linear(3, num_outputs) + >>> module_normal = AddStateIndependentNormalScale(num_outputs) + >>> tensor = torch.randn(3) + >>> loc, scale = module_normal(module(tensor)) + >>> print(loc.shape, scale.shape) + torch.Size([4]) torch.Size([4]) + >>> assert (scale > 0).all() + >>> # with modules that return more than one tensor + >>> module = nn.LSTM(3, num_outputs) + >>> module_normal = AddStateIndependentNormalScale(num_outputs) + >>> tensor = torch.randn(4, 2, 3) + >>> loc, scale, others = module_normal(*module(tensor)) + >>> print(loc.shape, scale.shape) + torch.Size([4, 2, 4]) torch.Size([4, 2, 4]) + >>> assert (scale > 0).all() + """ + + def __init__( + self, + scale_shape: torch.Size | int | tuple = None, + *, + scale_mapping: str = "exp", + scale_lb: Number = 1e-4, + device: torch.device | None = None, + make_param: bool = True, + init_value: float = 0.0, + ) -> None: + + super().__init__() + if scale_shape is None: + scale_shape = torch.Size(()) + self.scale_lb = scale_lb + if isinstance(scale_shape, int): + scale_shape = (scale_shape,) + self.scale_shape = torch.Size(scale_shape) + self.scale_mapping = scale_mapping + if make_param: + self.state_independent_scale = torch.nn.Parameter( + torch.full( + scale_shape, + init_value, + dtype=torch.get_default_dtype(), + device=device, + ) + ) + else: + self.state_independent_scale = torch.nn.Buffer( + torch.full( + scale_shape, + init_value, + dtype=torch.get_default_dtype(), + device=device, + ) + ) + + def forward( + self, loc: torch.Tensor, *others: torch.Tensor + ) -> tuple[torch.Tensor, ...]: + """Forward of AddStateIndependentNormalScale. + + Args: + loc (torch.Tensor): a location parameter. + *others: other unused parameters. + + Returns: + a tuple of two or more tensors containing the ``(loc, scale, *others)`` values. + """ + if self.scale_shape != loc.shape[-len(self.scale_shape) :]: + raise RuntimeError( + f"Last dimensions of loc ({loc.shape[-len(self.scale_shape):]}) do not match the number of dimensions " + f"in scale ({self.state_independent_scale.shape})" + ) + + scale = self.state_independent_scale.expand_as(loc) + scale = mappings(self.scale_mapping)(scale).clamp_min(self.scale_lb) + + return (loc, scale, *others) + + +class Delta(D.Distribution): + """Delta distribution. + + Args: + param (torch.Tensor): parameter of the delta distribution; + atol (number, optional): absolute tolerance to consider that a tensor matches the distribution parameter; + Default is 1e-6 + rtol (number, optional): relative tolerance to consider that a tensor matches the distribution parameter; + Default is 1e-6 + batch_shape (torch.Size, optional): batch shape; + event_shape (torch.Size, optional): shape of the outcome. + + """ + + arg_constraints: dict = {} + + def __init__( + self, + param: torch.Tensor, + atol: float = 1e-6, + rtol: float = 1e-6, + batch_shape: torch.Size | Sequence[int] | None = None, + event_shape: torch.Size | Sequence[int] | None = None, + ) -> None: + if batch_shape is None: + batch_shape = torch.Size([]) + if event_shape is None: + event_shape = torch.Size([]) + self.update(param) + self.atol = atol + self.rtol = rtol + if not len(batch_shape) and not len(event_shape): + batch_shape = param.shape[:-1] + event_shape = param.shape[-1:] + super().__init__(batch_shape=batch_shape, event_shape=event_shape) + + def update(self, param: torch.Tensor) -> None: + self.param = param + + def _is_equal(self, value: torch.Tensor) -> torch.Tensor: + param = self.param.expand_as(value) + is_equal = abs(value - param) < self.atol + self.rtol * abs(param) + for i in range(-1, -len(self.event_shape) - 1, -1): + is_equal = is_equal.all(i) + return is_equal + + def log_prob(self, value: torch.Tensor) -> torch.Tensor: + is_equal = self._is_equal(value) + out = torch.zeros_like(is_equal, dtype=value.dtype) + out.masked_fill_(is_equal, np.inf) + out.masked_fill_(~is_equal, -np.inf) + return out + + @torch.no_grad() + def sample( + self, + sample_shape: torch.Size | Sequence[int] | None = None, + ) -> torch.Tensor: + if sample_shape is None: + sample_shape = torch.Size([]) + return self.param.expand((*sample_shape, *self.param.shape)) + + def rsample( + self, + sample_shape: torch.Size | Sequence[int] | None = None, + ) -> torch.Tensor: + if sample_shape is None: + sample_shape = torch.Size([]) + return self.param.expand((*sample_shape, *self.param.shape)) + + @property + def mode(self) -> torch.Tensor: + return self.param + + @property + def mean(self) -> torch.Tensor: + return self.param + + @property + def deterministic_sample(self) -> torch.Tensor: + return self.param + + +@property +def _logistic_deterministic_sample(self): + s = self.loc + for t in self.transforms: + s = t(s) + return s + + +D.LogisticNormal.deterministic_sample = _logistic_deterministic_sample diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/discrete.py b/lib/python3.12/site-packages/tensordict/nn/distributions/discrete.py new file mode 100644 index 0000000000000000000000000000000000000000..c057dd2f33ed8aa438e6979cedf6cedf2f1c2237 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/distributions/discrete.py @@ -0,0 +1,103 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from typing import Sequence + +import torch +from torch import distributions as D + +# We need this to build the distribution maps +__all__ = [ + "OneHotCategorical", +] + + +def _treat_categorical_params( + params: torch.Tensor | None = None, +) -> torch.Tensor | None: + if params is None: + return None + if params.shape[-1] == 1: + params = params[..., 0] + return params + + +def rand_one_hot(values: torch.Tensor, do_softmax: bool = True) -> torch.Tensor: + if do_softmax: + values = values.softmax(-1) + out = values.cumsum(-1) > torch.rand_like(values[..., :1]) + out = (out.cumsum(-1) == 1).to(torch.long) + return out + + +class OneHotCategorical(D.Categorical): + """One-hot categorical distribution. + + This class behaves excacly as torch.distributions.Categorical except that it reads and produces one-hot encodings + of the discrete tensors. + + """ + + num_params: int = 1 + + def __init__( + self, + logits: torch.Tensor | None = None, + probs: torch.Tensor | None = None, + **kwargs, + ) -> None: + logits = _treat_categorical_params(logits) + probs = _treat_categorical_params(probs) + super().__init__(probs=probs, logits=logits, **kwargs) + + def log_prob(self, value: torch.Tensor) -> torch.Tensor: + return super().log_prob(value.argmax(dim=-1)) + + @property + def mode(self) -> torch.Tensor: + if hasattr(self, "logits"): + return (self.logits == self.logits.max(-1, True)[0]).to(torch.long) + else: + return (self.probs == self.probs.max(-1, True)[0]).to(torch.long) + + deterministic_sample = mode + + def sample( + self, + sample_shape: torch.Size | Sequence[int] | None = None, + ) -> torch.Tensor: + if sample_shape is None: + sample_shape = torch.Size([]) + out = super().sample(sample_shape=sample_shape) + out = torch.nn.functional.one_hot(out, self.logits.shape[-1]).to(torch.long) + return out + + def rsample( + self, + sample_shape: torch.Size | Sequence[int] | None = None, + ) -> torch.Tensor: + if sample_shape is None: + sample_shape = torch.Size([]) + d = D.relaxed_categorical.RelaxedOneHotCategorical( + 1.0, probs=self.probs, logits=self.logits + ) + out = d.rsample(sample_shape) + out.data.copy_((out == out.max(-1)[0].unsqueeze(-1)).to(out.dtype)) + return out + + +D.RelaxedBernoulli.deterministic_sample = D.Bernoulli.mode + + +@property +def _relaxed_onehot_mode(self) -> torch.Tensor: + probs = self.base_dist.probs + mode = probs.argmax(dim=-1) + return torch.nn.functional.one_hot(mode, num_classes=probs.shape[-1]).to(probs) + + +D.RelaxedOneHotCategorical.deterministic_sample = _relaxed_onehot_mode diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/truncated_normal.py b/lib/python3.12/site-packages/tensordict/nn/distributions/truncated_normal.py new file mode 100644 index 0000000000000000000000000000000000000000..449da5bd2f9855d4bee074fa1aa32def13a5ff23 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/distributions/truncated_normal.py @@ -0,0 +1,191 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +# from https://github.com/toshas/torch_truncnorm + +from __future__ import annotations + +import math +from numbers import Number +from typing import Sequence + +import torch +from torch.distributions import constraints, Distribution +from torch.distributions.utils import broadcast_all + + +CONST_SQRT_2 = math.sqrt(2) +CONST_INV_SQRT_2PI = 1 / math.sqrt(2 * math.pi) +CONST_INV_SQRT_2 = 1 / math.sqrt(2) +CONST_LOG_INV_SQRT_2PI = math.log(CONST_INV_SQRT_2PI) +CONST_LOG_SQRT_2PI_E = 0.5 * math.log(2 * math.pi * math.e) + + +class TruncatedStandardNormal(Distribution): + """Truncated Standard Normal distribution. + + Source: https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + """ + + arg_constraints = { + "a": constraints.real, + "b": constraints.real, + } + has_rsample = True + eps = 1e-6 + + def __init__( + self, + a: Number | torch.Tensor, + b: Number | torch.Tensor, + validate_args: bool | None = None, + ) -> None: + self.a, self.b = broadcast_all(a, b) + if isinstance(a, Number) and isinstance(b, Number): + batch_shape = torch.Size() + else: + batch_shape = self.a.size() + super().__init__(batch_shape, validate_args=validate_args) + if self.a.dtype != self.b.dtype: + raise ValueError("Truncation bounds types are different") + if any((self.a >= self.b).view(-1).tolist()): + raise ValueError("Incorrect truncation range") + # eps = torch.finfo(self.a.dtype).eps * 10 + eps = self.eps + self._dtype_min_gt_0 = eps + self._dtype_max_lt_1 = 1 - eps + self._little_phi_a = self._little_phi(self.a) + self._little_phi_b = self._little_phi(self.b) + self._big_phi_a = self._big_phi(self.a) + self._big_phi_b = self._big_phi(self.b) + self._Z = (self._big_phi_b - self._big_phi_a).clamp(eps, 1 - eps) + self._log_Z = self._Z.log() + little_phi_coeff_a = torch.nan_to_num(self.a, nan=math.nan) + little_phi_coeff_b = torch.nan_to_num(self.b, nan=math.nan) + self._lpbb_m_lpaa_d_Z = ( + self._little_phi_b * little_phi_coeff_b + - self._little_phi_a * little_phi_coeff_a + ) / self._Z + self._mean = -(self._little_phi_b - self._little_phi_a) / self._Z + self._variance = ( + 1 + - self._lpbb_m_lpaa_d_Z + - ((self._little_phi_b - self._little_phi_a) / self._Z) ** 2 + ) + self._entropy = CONST_LOG_SQRT_2PI_E + self._log_Z - 0.5 * self._lpbb_m_lpaa_d_Z + + @constraints.dependent_property + def support(self) -> constraints.Constraints: + return constraints.interval(self.a, self.b) + + @property + def mean(self) -> torch.Tensor: + return self._mean + + @property + def variance(self) -> torch.Tensor: + return self._variance + + @property + def entropy(self) -> torch.Tensor: + return self._entropy + + @property + def auc(self) -> torch.Tensor: + return self._Z + + @staticmethod + def _little_phi(x: torch.Tensor) -> torch.Tensor: + return (-(x**2) * 0.5).exp() * CONST_INV_SQRT_2PI + + def _big_phi(self, x: torch.Tensor) -> torch.Tensor: + phi = 0.5 * (1 + (x * CONST_INV_SQRT_2).erf()) + return phi.clamp(self.eps, 1 - self.eps) + + @staticmethod + def _inv_big_phi(x: torch.Tensor) -> torch.Tensor: + return CONST_SQRT_2 * (2 * x - 1).erfinv() + + def cdf(self, value: torch.Tensor) -> torch.Tensor: + if self._validate_args: + self._validate_sample(value) + return ((self._big_phi(value) - self._big_phi_a) / self._Z).clamp(0, 1) + + def icdf(self, value: torch.Tensor) -> torch.Tensor: + y = self._big_phi_a + value * self._Z + y = y.clamp(self.eps, 1 - self.eps) + return self._inv_big_phi(y) + + def log_prob(self, value: torch.Tensor) -> torch.Tensor: + if self._validate_args: + self._validate_sample(value) + return CONST_LOG_INV_SQRT_2PI - self._log_Z - (value**2) * 0.5 + + def rsample( + self, + sample_shape: torch.Size | Sequence[int] | None = None, + ) -> torch.Tensor: + if sample_shape is None: + sample_shape = torch.Size([]) + shape = self._extended_shape(sample_shape) + p = torch.empty(shape, device=self.a.device).uniform_( + self._dtype_min_gt_0, self._dtype_max_lt_1 + ) + return self.icdf(p) + + +class TruncatedNormal(TruncatedStandardNormal): + """Truncated Normal distribution. + + https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf + """ + + has_rsample = True + + def __init__( + self, + loc: Number | torch.Tensor, + scale: Number | torch.Tensor, + a: Number | torch.Tensor, + b: Number | torch.Tensor, + validate_args: bool | None = None, + ) -> None: + scale = scale.clamp_min(self.eps) + self.loc, self.scale, a, b = broadcast_all(loc, scale, a, b) + self._non_std_a = a + self._non_std_b = b + a = (a - self.loc) / self.scale + b = (b - self.loc) / self.scale + super().__init__(a, b, validate_args=validate_args) + self._log_scale = self.scale.log() + self._mean = self._mean * self.scale + self.loc + self._variance = self._variance * self.scale**2 + self._entropy += self._log_scale + + def _to_std_rv(self, value: torch.Tensor) -> torch.Tensor: + return (value - self.loc) / self.scale + + def _from_std_rv(self, value: torch.Tensor) -> torch.Tensor: + return value * self.scale + self.loc + + def cdf(self, value: torch.Tensor) -> torch.Tensor: + return super().cdf(self._to_std_rv(value)) + + def icdf(self, value: torch.Tensor) -> torch.Tensor: + sample = self._from_std_rv(super().icdf(value)) + + # clamp data but keep gradients + sample_clip = torch.stack( + [sample.detach(), self._non_std_a.detach().expand_as(sample)], 0 + ).max(0)[0] + sample_clip = torch.stack( + [sample_clip, self._non_std_b.detach().expand_as(sample)], 0 + ).min(0)[0] + sample.data.copy_(sample_clip) + return sample + + def log_prob(self, value: torch.Tensor) -> torch.Tensor: + value = self._to_std_rv(value) + return super().log_prob(value) - self._log_scale diff --git a/lib/python3.12/site-packages/tensordict/nn/distributions/utils.py b/lib/python3.12/site-packages/tensordict/nn/distributions/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..21170e3177281bff787211e9a11b6a5b0b5e24b4 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/distributions/utils.py @@ -0,0 +1,40 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import torch +from tensordict.utils import DeviceType +from torch import distributions as D + + +def _cast_device( + elt: torch.Tensor | float, + device: DeviceType, +) -> torch.Tensor | float: + if isinstance(elt, torch.Tensor): + return elt.to(device) + return elt + + +def _cast_transform_device( + transform: D.Transform | None, + device: DeviceType, +) -> D.Transform | None: + if transform is None: + return transform + elif isinstance(transform, D.ComposeTransform): + for i, t in enumerate(transform.parts): + transform.parts[i] = _cast_transform_device(t, device) + elif isinstance(transform, D.Transform): + for attribute in dir(transform): + value = getattr(transform, attribute) + if isinstance(value, torch.Tensor): + setattr(transform, attribute, value.to(device)) + return transform + else: + raise TypeError( + f"Cannot perform device casting for transform of type {type(transform)}" + ) diff --git a/lib/python3.12/site-packages/tensordict/nn/ensemble.py b/lib/python3.12/site-packages/tensordict/nn/ensemble.py new file mode 100644 index 0000000000000000000000000000000000000000..0b3cb6520325911c8569f955a310c9cf7d5b4364 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/ensemble.py @@ -0,0 +1,131 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +import warnings + +import torch +from tensordict._lazy import LazyStackedTensorDict +from tensordict._td import TensorDict +from tensordict.base import TensorDictBase +from tensordict.nn.common import TensorDictModuleBase + +from tensordict.nn.params import TensorDictParams + + +class EnsembleModule(TensorDictModuleBase): + """Module that wraps a module and repeats it to form an ensemble. + + Args: + module (nn.Module): The nn.module to duplicate and wrap. + num_copies (int): The number of copies of module to make. + parameter_init_function (Callable): A function that takes a module copy and initializes its parameters. + expand_input (bool): Whether to expand the input TensorDict to match the number of copies. This should be + True unless you are chaining ensemble modules together, e.g. EnsembleModule(cnn) -> EnsembleModule(mlp). + If False, EnsembleModule(mlp) will expected the previous module(s) to have already expanded the input. + + Examples: + >>> import torch + >>> from torch import nn + >>> from tensordict.nn import TensorDictModule, EnsembleModule + >>> from tensordict import TensorDict + >>> net = nn.Sequential(nn.Linear(4, 32), nn.ReLU(), nn.Linear(32, 2)) + >>> mod = TensorDictModule(net, in_keys=['a'], out_keys=['b']) + >>> ensemble = EnsembleModule(mod, num_copies=3) + >>> data = TensorDict({'a': torch.randn(10, 4)}, batch_size=[10]) + >>> ensemble(data) + TensorDict( + fields={ + a: Tensor(shape=torch.Size([3, 10, 4]), device=cpu, dtype=torch.float32, is_shared=False), + b: Tensor(shape=torch.Size([3, 10, 2]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3, 10]), + device=None, + is_shared=False) + + To stack EnsembleModules together, we should be mindful of turning off `expand_input` from the second module and on. + + Examples: + >>> import torch + >>> from tensordict.nn import TensorDictModule, TensorDictSequential, EnsembleModule + >>> from tensordict import TensorDict + >>> module = TensorDictModule(torch.nn.Linear(2,3), in_keys=['bork'], out_keys=['dork']) + >>> next_module = TensorDictModule(torch.nn.Linear(3,1), in_keys=['dork'], out_keys=['spork']) + >>> e0 = EnsembleModule(module, num_copies=4, expand_input=True) + >>> e1 = EnsembleModule(next_module, num_copies=4, expand_input=False) + >>> seq = TensorDictSequential(e0, e1) + >>> data = TensorDict({'bork': torch.randn(5,2)}, batch_size=[5]) + >>> seq(data) + TensorDict( + fields={ + bork: Tensor(shape=torch.Size([4, 5, 2]), device=cpu, dtype=torch.float32, is_shared=False), + dork: Tensor(shape=torch.Size([4, 5, 3]), device=cpu, dtype=torch.float32, is_shared=False), + spork: Tensor(shape=torch.Size([4, 5, 1]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 5]), + device=None, + is_shared=False) + """ + + def __init__( + self, + module: TensorDictModuleBase, + num_copies: int, + expand_input: bool = True, + ): + super().__init__() + self.in_keys = module.in_keys + self.out_keys = module.out_keys + params_td = TensorDict.from_module(module).expand(num_copies).to_tensordict() + + self.module = module + if expand_input: + self.vmapped_forward = torch.vmap(self._func_module_call, (None, 0)) + else: + self.vmapped_forward = torch.vmap(self._func_module_call, 0) + + self.reset_parameters_recursive(params_td) + self.params_td = TensorDictParams(params_td) + + def _func_module_call(self, input, params): + with params.to_module(self.module): + return self.module(input) + + def forward(self, tensordict: TensorDict) -> TensorDict: + return self.vmapped_forward(tensordict, self.params_td) + + def reset_parameters_recursive( + self, parameters: TensorDictBase = None + ) -> TensorDictBase: + """Resets the parameters of all the copies of the module. + + Args: + parameters (TensorDict): A TensorDict of parameters for self.module. The batch dimension(s) of the tensordict + denote the number of module copies to reset. + + Returns: + A TensorDict of pointers to the reset parameters. + """ + if parameters is None: + raise ValueError( + "Ensembles are functional and require passing a TensorDict of parameters to reset_parameters_recursive" + ) + if parameters.ndim: + params_pointers = [] + for params_copy in parameters.unbind(0): + self.reset_parameters_recursive(params_copy) + params_pointers.append(params_copy) + return LazyStackedTensorDict.lazy_stack(params_pointers, -1) + else: + # In case the user has added other neural networks to the EnsembleModule + # besides those in self.module + child_mods = [ + mod + for name, mod in self.named_children() + if name != "module" and name != "ensemble_parameters" + ] + if child_mods: + warnings.warn( + "EnsembleModule.reset_parameters_recursive() only resets parameters of self.module, but other parameters were detected. These parameters will not be reset." + ) + # Reset all self.module descendant parameters + return self.module.reset_parameters_recursive(parameters) diff --git a/lib/python3.12/site-packages/tensordict/nn/functional_modules.py b/lib/python3.12/site-packages/tensordict/nn/functional_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..549b4d12a93a0a85f89ebc4dc875c66d65eb35e6 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/functional_modules.py @@ -0,0 +1,342 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import os +from inspect import signature +from typing import Any, Callable, Iterable + +import torch +import torch.utils._pytree +from tensordict._pytree import PYTREE_REGISTERED_LAZY_TDS, PYTREE_REGISTERED_TDS + +from tensordict._td import TensorDict +from tensordict.base import is_tensor_collection + +from tensordict.utils import implement_for, strtobool +from torch import nn +from torch.utils._pytree import SUPPORTED_NODES + +try: + from torch.nn.modules.module import _global_parameter_registration_hooks +except ImportError: + # old torch version, passing + pass + +__base__setattr__ = nn.Module.__setattr__ + +PYTREE_HAS_ISLEAF = "is_leaf" in signature(torch.utils._pytree.tree_map).parameters + + +@implement_for("torch", "2.0", None) +def _register_params(self, name, param): + """A simplified version of register_param where checks are skipped.""" + for hook in _global_parameter_registration_hooks.values(): + output = hook(self, name, param) + if output is not None: + param = output + self._parameters[name] = param + + +@implement_for("torch", None, "2.0") +def _register_params(self, name, param): # noqa: F811 + self.register_parameter(name, param) + + +def set_tensor(module: "torch.nn.Module", name: str, tensor: torch.Tensor) -> None: + """Simplified version of torch.nn.utils._named_member_accessor.""" + if name in module._parameters: + del module._parameters[name] # type: ignore[assignment] + was_buffer = name in module._buffers + if was_buffer: + del module._buffers[name] + if isinstance(tensor, nn.Parameter): + module.__dict__.pop(name, None) + # module.register_parameter(name, tensor) + _register_params(module, name, tensor) + elif was_buffer and isinstance(tensor, Tensor): + module._buffers[name] = tensor + else: + module.__dict__[name] = tensor + + +@implement_for("torch", "2.0", None) +def set_tensor_dict( # noqa: F811 + module_dict, module, name: str, tensor: torch.Tensor +) -> None: + """Simplified version of torch.nn.utils._named_member_accessor.""" + if name in module_dict["_parameters"]: + del module_dict["_parameters"][name] # type: ignore[assignment] + was_buffer = name in module_dict["_buffers"] + if was_buffer: + del module_dict["_buffers"][name] + if isinstance(tensor, nn.Parameter): + module_dict.pop(name, None) + # module.register_parameter(name, tensor) + for hook in _global_parameter_registration_hooks.values(): + output = hook(module, name, tensor) + if output is not None: + tensor = output + module_dict["_parameters"][name] = tensor + elif was_buffer and isinstance(tensor, Tensor): + module_dict["_buffers"][name] = tensor + else: + module_dict[name] = tensor + + +@implement_for("torch", None, "2.0") +def set_tensor_dict( # noqa: F811 + module_dict, module, name: str, tensor: torch.Tensor +) -> None: + """Simplified version of torch.nn.utils._named_member_accessor.""" + if name in module_dict["_parameters"]: + del module_dict["_parameters"][name] # type: ignore[assignment] + was_buffer = name in module_dict["_buffers"] + if was_buffer: + del module_dict["_buffers"][name] + if isinstance(tensor, nn.Parameter): + module_dict.pop(name, None) + module.register_parameter(name, tensor) + elif was_buffer and isinstance(tensor, Tensor): + module_dict["_buffers"][name] = tensor + else: + module_dict[name] = tensor + + +_RESET_OLD_TENSORDICT = True +import torch._functorch.vmap as vmap_src # @manual=fbcode//caffe2:torch +from torch._functorch.vmap import ( # @manual=fbcode//caffe2:torch + _add_batch_dim, + _broadcast_to_and_flatten, + _get_name, + _maybe_remove_batch_dim, + _validate_and_get_batch_size, + Tensor, + tree_flatten, + tree_unflatten, +) + + +class _exclude_td_from_pytree: + def __init__(self): + self.tdnodes = {} + + def __enter__(self): + for tdtype in PYTREE_REGISTERED_TDS + PYTREE_REGISTERED_LAZY_TDS: + node = SUPPORTED_NODES.pop(tdtype, None) + if node is None: + continue + self.tdnodes[tdtype] = node + + def __exit__(self, exc_type, exc_val, exc_tb): + for tdtype, node in self.tdnodes.items(): + SUPPORTED_NODES[tdtype] = node + + def set(self): + self.__enter__() + + def unset(self): + self.__exit__(None, None, None) + + +if not strtobool(os.getenv("PYTORCH_TENSORDICT_IMPORT_VMAP", "False")): + # Monkey-patches + + def _process_batched_inputs( + in_dims: int | tuple[int, ...], args: Any, func: Callable + ) -> tuple[Any, ...]: + if not isinstance(in_dims, int) and not isinstance(in_dims, tuple): + raise ValueError( + f"""vmap({_get_name(func)}, in_dims={in_dims}, ...)(): +expected `in_dims` to be int or a (potentially nested) tuple +matching the structure of inputs, got: {type(in_dims)}.""" + ) + if len(args) == 0: + raise ValueError( + f"""vmap({_get_name(func)})(): got no inputs. Maybe you forgot to add +inputs, or you are trying to vmap over a function with no inputs. +The latter is unsupported.""" + ) + + # we want to escape TensorDicts as they take care of adding the batch dimension + if PYTREE_HAS_ISLEAF: + flat_args, args_spec = tree_flatten(args, is_leaf=is_tensor_collection) + flat_in_dims = _broadcast_to_and_flatten(in_dims, args_spec) + if flat_in_dims is None: + raise ValueError( + f"""vmap({_get_name(func)}, in_dims={in_dims}, ...)(): + in_dims is not compatible with the structure of `inputs`. + in_dims has structure {tree_flatten(in_dims)[1]} but inputs + has structure {args_spec}.""" + ) + else: + with _exclude_td_from_pytree(): + flat_args, args_spec = tree_flatten(args) + flat_in_dims = _broadcast_to_and_flatten(in_dims, args_spec) + if flat_in_dims is None: + raise ValueError( + f"""vmap({_get_name(func)}, in_dims={in_dims}, ...)(): + in_dims is not compatible with the structure of `inputs`. + in_dims has structure {tree_flatten(in_dims)[1]} but inputs + has structure {args_spec}.""" + ) + + for i, (arg, in_dim) in enumerate(zip(flat_args, flat_in_dims)): + if not isinstance(in_dim, int) and in_dim is not None: + raise ValueError( + f"""vmap({_get_name(func)}, in_dims={in_dims}, ...)(): +Got in_dim={in_dim} for an input but in_dim must be either +an integer dimension or None.""" + ) + if ( + isinstance(in_dim, int) + and not isinstance(arg, Tensor) + and not is_tensor_collection(arg) + ): + raise ValueError( + f"""vmap({_get_name(func)}, in_dims={in_dims}, ...)(): +Got in_dim={in_dim} for an input but the input is of type +{type(arg)}. We cannot vmap over non-Tensor arguments, +please use None as the respective in_dim""" + ) + if in_dim is not None and (in_dim < -arg.dim() or in_dim >= arg.dim()): + raise ValueError( + f"""vmap({_get_name(func)}, in_dims={in_dims}, ...)(): +Got in_dim={in_dim} for some input, but that input is a Tensor +of dimensionality {arg.dim()} so expected in_dim to satisfy +-{arg.dim()} <= in_dim < {arg.dim()}.""" + ) + if in_dim is not None and in_dim < 0: + flat_in_dims[i] = in_dim % arg.dim() + + return ( + _validate_and_get_batch_size(flat_in_dims, flat_args), + flat_in_dims, + flat_args, + args_spec, + ) + + vmap_src._process_batched_inputs = _process_batched_inputs + + def _create_batched_inputs( + flat_in_dims: list[int], flat_args: list[Any], vmap_level: int, args_spec + ) -> Any: + # See NOTE [Ignored _remove_batch_dim, _add_batch_dim] + # If tensordict, we remove the dim at batch_size[in_dim] such that the TensorDict can accept + # the batched tensors. This will be added in _unwrap_batched + + batched_inputs = [] + for in_dim, arg in zip(flat_in_dims, flat_args): + if in_dim is None: + if is_tensor_collection(arg): + # this may be a perf bottleneck and could benefit from caching + # arg = cache(arg.clone)(False) + arg = arg.clone(False) + + batched_input = arg + else: + if is_tensor_collection(arg): + batched_input = arg._add_batch_dim( + in_dim=in_dim, vmap_level=vmap_level + ) + else: + batched_input = _add_batch_dim(arg, in_dim, vmap_level) + batched_inputs.append(batched_input) + return tree_unflatten(batched_inputs, args_spec) + + vmap_src._create_batched_inputs = _create_batched_inputs + + def _unwrap_batched( + batched_outputs: Any, + out_dims: int | tuple[int, ...], + vmap_level: int, + batch_size: int, + func: Callable, + ) -> Any: + if PYTREE_HAS_ISLEAF: + flat_batched_outputs, output_spec = tree_flatten( + batched_outputs, is_leaf=is_tensor_collection + ) + else: + with _exclude_td_from_pytree(): + flat_batched_outputs, output_spec = tree_flatten(batched_outputs) + + def incompatible_error(): + raise ValueError( + f"vmap({_get_name(func)}, ..., out_dims={out_dims})(): " + f"out_dims is not compatible with the structure of `outputs`. " + f"out_dims has structure {tree_flatten(out_dims)[1]} but outputs " + f"has structure {output_spec}." + ) + + if isinstance(batched_outputs, torch.Tensor) or is_tensor_collection( + batched_outputs + ): + # Some weird edge case requires us to spell out the following + # see test_out_dims_edge_case + if isinstance(out_dims, int): + flat_out_dims = [out_dims] + elif isinstance(out_dims, tuple) and len(out_dims) == 1: + flat_out_dims = out_dims + elif out_dims is None: + flat_out_dims = [out_dims] + else: + incompatible_error() + else: + flat_out_dims = _broadcast_to_and_flatten(out_dims, output_spec) + if flat_out_dims is None: + incompatible_error() + flat_outputs = [] + for batched_output, out_dim in zip(flat_batched_outputs, flat_out_dims): + if not is_tensor_collection(batched_output): + out = _maybe_remove_batch_dim( + _get_name(func), batched_output, vmap_level, batch_size, out_dim + ) + else: + out = batched_output._maybe_remove_batch_dim( + _get_name(func), + vmap_level=vmap_level, + batch_size=batch_size, + out_dim=out_dim, + ) + flat_outputs.append(out) + return tree_unflatten(flat_outputs, output_spec) + + vmap_src._unwrap_batched = _unwrap_batched + + +def extract_weights_and_buffers( + model: nn.Module, +) -> TensorDict: # noqa + raise RuntimeError("extract_weights_and_buffers has been removed from tensordict.") + + +def is_functional(module: nn.Module): # noqa + raise RuntimeError("is_functional has been removed from tensordict.") + + +def make_functional( + module: nn.Module, + funs_to_decorate: Iterable[str] | None = None, + keep_params: bool = False, + return_params: bool = True, +) -> TensorDict: # noqa + raise RuntimeError("make_functional has been removed from tensordict.") + + +def get_functional( + module: nn.Module, + funs_to_decorate: Iterable[str] | None = None, +) -> nn.Module: # noqa + raise RuntimeError("get_functional has been removed from tensordict.") + + +def repopulate_module(model: nn.Module, tensordict: TensorDict) -> nn.Module: # noqa + raise RuntimeError("repopulate_module has been removed from tensordict.") + + +if strtobool(os.environ.get("EXCLUDE_TD_FROM_PYTREE", "0")): + _exclude_td_from_pytree().set() diff --git a/lib/python3.12/site-packages/tensordict/nn/params.py b/lib/python3.12/site-packages/tensordict/nn/params.py new file mode 100644 index 0000000000000000000000000000000000000000..a3b379a18209c622f319f419fd6813886fa3b51c --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/params.py @@ -0,0 +1,1392 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from __future__ import annotations + +import functools +import inspect +import re +import weakref +from concurrent.futures import Future, ThreadPoolExecutor +from contextlib import nullcontext +from copy import copy +from functools import wraps +from typing import ( + Any, + Callable, + Dict, + Iterator, + List, + OrderedDict, + Sequence, + Type, + TYPE_CHECKING, +) + +import torch + +from tensordict._lazy import _CustomOpTensorDict, LazyStackedTensorDict +from tensordict._nestedkey import NestedKey +from tensordict._td import _SubTensorDict, TensorDict +from tensordict._tensorcollection import TensorCollection +from tensordict._torch_func import TD_HANDLED_FUNCTIONS + +from tensordict.base import ( + _default_is_leaf, + _is_tensor_collection, + _register_tensor_class, + CompatibleType, + NO_DEFAULT, + T, + TensorDictBase, +) + +from tensordict.memmap import MemoryMappedTensor +from tensordict.utils import ( + _LOCK_ERROR, + _zip_strict, + BufferLegacy, + erase_cache, + implement_for, + IndexType, + is_batchedtensor, + lock_blocked, +) +from torch import multiprocessing as mp, nn, Tensor +from torch.utils._pytree import tree_map + +try: + from functorch import dim as ftdim + + _has_funcdim = True +except ImportError: + from tensordict.utils import _ftdim_mock as ftdim + + _has_funcdim = False + +try: + from torch.nn.parameter import Buffer +except ImportError: + from tensordict.utils import Buffer + + +try: + from torch.compiler import is_compiling +except ImportError: + from torch._dynamo import is_compiling + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +def _apply_leaves(data, fn): + if isinstance(data, TensorDict): + with data.unlock_(): + for key, val in list(data.items()): + data._set_str( + key, + _apply_leaves(val, fn), + validated=True, + inplace=False, + non_blocking=False, + ) + return data + elif isinstance(data, LazyStackedTensorDict): + # this is currently not implemented as the registration of params will only work + # with plain TensorDict. The solution will be using pytree to get each independent + # leaf + raise RuntimeError( + "Using a LazyStackedTensorDict within a TensorDictParams isn't permitted." + ) + # for _data in data.tensordicts: + # _apply_leaves(_data, fn) + # return data + elif isinstance(data, _CustomOpTensorDict): + _apply_leaves(data._source, fn) + return data + elif isinstance(data, _SubTensorDict): + raise RuntimeError( + "Using a _SubTensorDict within a TensorDictParams isn't permitted." + ) + else: + return fn(data) + + +def _get_args_dict(func, args, kwargs): + signature = inspect.signature(func) + bound_arguments = signature.bind(*args, **kwargs) + bound_arguments.apply_defaults() + + args_dict = dict(bound_arguments.arguments) + return args_dict + + +def _maybe_make_param(tensor): + if isinstance(tensor, (Tensor, ftdim.Tensor)) and not isinstance( + tensor, (nn.Parameter, Buffer, BufferLegacy) + ): + if tensor.dtype in (torch.float, torch.double, torch.half): + tensor = nn.Parameter(tensor) + elif not is_batchedtensor(tensor): + # convert all non-parameters to buffers + # dataptr = tensor.data.data_ptr() + tensor = Buffer(tensor) + else: + # We want to keep the grad_fn of tensors, e.g. param.expand(10) should point to the original param + tensor = BufferLegacy(tensor) + return tensor + + +def _maybe_make_param_or_buffer(tensor): + if isinstance(tensor, (Tensor, ftdim.Tensor)) and not isinstance( + tensor, (nn.Parameter, Buffer) + ): + if not tensor.requires_grad and not is_batchedtensor(tensor): + # convert all non-parameters to buffers + # dataptr = tensor.data.data_ptr() + tensor = Buffer(tensor) + else: + # We want to keep the grad_fn of tensors, e.g. param.expand(10) should point to the original param + tensor = BufferLegacy(tensor) + + # assert tensor.data.data_ptr() == dataptr + return tensor + + +class _unlock_and_set: + # temporarily unlocks the nested tensordict to execute a function + def __new__(cls, *args, **kwargs): + if len(args) and callable(args[0]): + return cls(**kwargs)(args[0]) + return super().__new__(cls) + + def __init__(self, **only_for_kwargs): + self.only_for_kwargs = only_for_kwargs + + def __call__(self, func): + name = func.__name__ + + @wraps(func) + def new_func(_self, *args, **kwargs): # type: ignore[misc] + if self.only_for_kwargs: + arg_dict = _get_args_dict(func, (_self, *args), kwargs) + for kwarg, exp_value in self.only_for_kwargs.items(): + cur_val = arg_dict.get(kwarg, NO_DEFAULT) + if cur_val != exp_value: + # escape + meth = getattr(_self._param_td, name) + out = meth(*args, **kwargs) + return out + if not _self.no_convert: + args = tree_map(_maybe_make_param, args) + kwargs = tree_map(_maybe_make_param, kwargs) + else: + args = tree_map(_maybe_make_param_or_buffer, args) + kwargs = tree_map(_maybe_make_param_or_buffer, kwargs) + if _self.is_locked: + # if the root (TensorDictParams) is locked, we still want to raise an exception + raise RuntimeError(_LOCK_ERROR) + with ( + _self._param_td.unlock_() + if _self._param_td.is_locked + else nullcontext() + ): + meth = getattr(_self._param_td, name) + out = meth(*args, **kwargs) + _self._reset_params() + if out is _self._param_td: + return _self + return out + + return new_func + + +def _get_post_hook(func): + @wraps(func) + def new_func(self, *args, **kwargs): # type: ignore[misc] + out = func(self, *args, **kwargs) + return self._apply_get_post_hook(out) + + return new_func + + +def _fallback(func): + """Calls the method on the nested tensordict.""" + name = func.__name__ + + @wraps(func) + def new_func(self, *args, **kwargs): # type: ignore[misc] + out = getattr(self._param_td, name)(*args, **kwargs) + if out is self._param_td: + # if the output does not change, return the wrapper + return self + return out + + return new_func + + +def _fallback_property(func): + name = func.__name__ + + @wraps(func) + def new_func(self): # type: ignore[misc] + out = getattr(self._param_td, name) + if out is self._param_td: + return self + return out + + def setter(self, value): # type: ignore[misc] + return getattr(type(self._param_td), name).fset(self._param_td, value) + + return property(new_func, setter) + + +def _replace(func): + name = func.__name__ + + @wraps(func) + def new_func(self, *args, **kwargs): # type: ignore[misc] + out = getattr(self._param_td, name)(*args, **kwargs) + if out is self._param_td: + return self + self._param_td = out + return self + + return new_func + + +def _carry_over(func): + name = func.__name__ + + @wraps(func) + def new_func(self, *args, **kwargs): # type: ignore[misc] + out = getattr(self._param_td, name)(*args, **kwargs) + if out is self._param_td: + return self + if not isinstance(out, TensorDictParams): + out = TensorDictParams(out, no_convert="skip") + out.no_convert = self.no_convert + return out + + return new_func + + +def _apply_on_data(func): + @wraps(func) + def new_func(self, *args, **kwargs): # type: ignore[misc] + getattr(self.data, func.__name__)(*args, **kwargs) + return self + + return new_func + + +class TensorDictParams(TensorDictBase, nn.Module): # type: ignore[override,misc,attr-defined] + r"""A Wrapper for TensorDictBase with Parameter Exposure. + + This class is designed to hold a `TensorDictBase` instance that contains parameters, making them accessible to a + parent :class:`~torch.nn.Module`. This allows for seamless integration of tensordict parameters into PyTorch modules, + enabling operations like parameter iteration and optimization. + + Key Features: + + - Parameter Exposure: Parameters within the tensordict are exposed to the parent module, allowing them to be included + in operations like `named_parameters()`. + - Indexing: Indexing works similarly to the wrapped tensordict. However, parameter names (in :meth:`~.named_parameters`) are registered using + `TensorDict.flatten_keys("_")`, which may result in different key names compared to the tensordict content. + - Automatic Conversion: Any tensor set in the tensordict is automatically converted to a :class:`torch.nn.Parameter`, + unless specified otherwise through the :attr:`no_convert` keyword argument. + + Args + parameters (TensorDictBase or dict): The tensordict to represent as parameters. Values are converted to + parameters unless `no_convert=True`. If a `dict` is provided, it is wrapped in a `TensorDict` instance. + Keyword arguments can also be used. + + Keyword Args: + no_convert (bool): If `True`, no conversion to `nn.Parameter` occurs and all non-parameter, non-buffer tensors + will be converted to a :class:`~torch.nn.Buffer` instance. + If ``False``, all tensors with non-integer dtypes will be converted to :class:`~torch.nn.Parameter` + whereas integer dtypes will be converted to :class:`~torch.nn.Buffer` instances. + Defaults to `False`. + lock (bool): If `True`, the tensordict hosted by `TensorDictParams` is locked, preventing modifications and + potentially impacting performance when `unlock_()` is required. + Defaults to `False`. + + .. warning:: Because the inner tensordict isn't copied or locked by default, registering the tensordict + in a ``TensorDictParams`` and modifying its content afterwards will __not__ update the values within + the ``TensorDictParams`` :meth:`.parameters` and :meth:`~.buffers` sequences. + + **kwargs: Key-value pairs to populate the `TensorDictParams`. Exclusive with the `parameters` input. + + Examples + >>> from torch import nn + >>> from tensordict import TensorDict + >>> module = nn.Sequential(nn.Linear(3, 4), nn.Linear(4, 4)) + >>> params = TensorDict.from_module(module) + >>> params.lock_() + >>> p = TensorDictParams(params) + >>> print(p) + TensorDictParams(params=TensorDict( + fields={ + 0: TensorDict( + fields={ + bias: Parameter(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Parameter(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False), + 1: TensorDict( + fields={ + bias: Parameter(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), + weight: Parameter(shape=torch.Size([4, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False)) + >>> class CustomModule(nn.Module): + ... def __init__(self, params): + ... super().__init__() + ... self.params = params + >>> m = CustomModule(p) + >>> # The wrapper supports assignment, and values are converted to Parameters + >>> m.params['other'] = torch.randn(3) + >>> assert isinstance(m.params['other'], nn.Parameter) + + """ + + def __init__( + self, + parameters: TensorDictBase | dict | None = None, + *, + no_convert=False, + lock: bool = False, + **kwargs, + ): + nn.Module.__init__(self) + if parameters is None: + parameters = kwargs + elif kwargs: + raise TypeError( + f"parameters cannot be passed along with extra keyword arguments, but got {kwargs.keys()} extra args." + ) + + params = None + buffers = None + if isinstance(parameters, dict): + parameters = TensorDict(parameters) + elif isinstance(parameters, TensorDictParams): + params = dict(parameters._parameters) + buffers = dict(parameters._buffers) + parameters = parameters._param_td.copy().lock_() + no_convert = "skip" + + self.no_convert = no_convert + if no_convert != "skip": + if not no_convert: + func = _maybe_make_param + else: + func = _maybe_make_param_or_buffer + self._param_td = _apply_leaves(parameters, lambda x: func(x)) + else: + self._param_td = parameters + + self._lock_content = lock + if lock: + self._param_td.lock_() + self._reset_params(params=params, buffers=buffers) + self._is_locked = False + self._locked_tensordicts = [] + self._get_post_hook = [] + + @classmethod + def _new_unsafe( + cls, + parameters: TensorDictBase, + *, + no_convert=None, + lock: bool = False, + params: dict | None = None, + buffers: dict | None = None, + **kwargs, + ): + if is_compiling(): + return TensorDictParams(parameters, no_convert="skip", lock=lock) + + if parameters is None: + parameters = kwargs + + if isinstance(parameters, dict): + parameters = TensorDict._new_unsafe(parameters, **kwargs) + if no_convert is None: + # Then _new_unsafe is called from somewhere that doesn't know + # that it's a TDParams and we return a TensorDict (eg, torch.gather) + return parameters + elif isinstance(parameters, TensorDictParams): + if kwargs: + raise TypeError( + f"parameters cannot be passed along with extra keyword arguments, but got {kwargs.keys()} extra args." + ) + params = dict(parameters._parameters) + buffers = dict(parameters._buffers) + parameters = parameters._param_td + no_convert = "skip" + + self = TensorDictParams.__new__(cls) + nn.Module.__init__(self) + + self._param_td = parameters + self.no_convert = no_convert + if no_convert != "skip": + raise RuntimeError("_new_unsafe requires no_convert to be set to 'skip'") + self._lock_content = lock + if lock: + self._param_td.lock_() + self._reset_params(params=params, buffers=buffers) + self._is_locked = False + self._locked_tensordicts = [] + self._get_post_hook = [] + return self + + def __iter__(self): + yield from self._param_td.__iter__() + + def register_get_post_hook(self, hook): + """Register a hook to be called after any get operation on leaf tensors.""" + if not callable(hook): + raise ValueError("Hooks must be callables.") + self._get_post_hook.append(hook) + + def _apply_get_post_hook(self, val): + if not _is_tensor_collection(type(val)): + for hook in self._get_post_hook: + new_val = hook(self, val) + if new_val is not None: + val = new_val + return val + + def _reset_params(self, params: dict | None = None, buffers: dict | None = None): + parameters = self._param_td + + self._parameters.clear() + self._buffers.clear() + + if (params is not None) ^ (buffers is not None): + raise RuntimeError("both params and buffers must either be None or not.") + elif params is None: + param_keys = [] + params = [] + buffer_keys = [] + buffers = [] + for key, value in parameters.items(True, True): + # flatten key + if isinstance(key, tuple): + key = ".".join(key) + if isinstance(value, nn.Parameter): + param_keys.append(key) + params.append(value) + else: + buffer_keys.append(key) + buffers.append(value) + + self._parameters.update(dict(_zip_strict(param_keys, params))) + self._buffers.update(dict(_zip_strict(buffer_keys, buffers))) + else: + self._parameters.update(params) + self._buffers.update(buffers) + + @classmethod + def __torch_function__( + cls, + func: Callable, + types: tuple[type, ...], + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Callable: + if kwargs is None: + kwargs = {} + if func not in TDPARAM_HANDLED_FUNCTIONS or not all( + issubclass(t, (Tensor, ftdim.Tensor, TensorDictBase)) for t in types + ): + return NotImplemented + return TDPARAM_HANDLED_FUNCTIONS[func](*args, **kwargs) + + @classmethod + def _flatten_key(cls, key): + def make_valid_identifier(s): + # Replace invalid characters with underscores + s = re.sub(r"\W|^(?=\d)", "_", s) + + # Ensure the string starts with a letter or underscore + if not s[0].isalpha() and s[0] != "_": + s = "_" + s + + return s + + key_flat = "_".join(key) + if not key_flat.isidentifier(): + key_flat = make_valid_identifier(key_flat) + return key_flat + + @lock_blocked + @_unlock_and_set + def __setitem__( # type: ignore[misc] + self, + index: IndexType, + value: Any, + ) -> None: ... + + @lock_blocked + @_unlock_and_set + def set( + self, key: NestedKey, item: CompatibleType, inplace: bool = False, **kwargs: Any + ) -> TensorDictBase: ... + + @lock_blocked + def update( + self, + input_dict_or_td: dict[str, CompatibleType] | TensorDictBase, + clone: bool = False, + inplace: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + is_leaf: Callable[[Type], bool] | None = None, + update_batch_size: bool = False, + ignore_lock: bool = False, + ) -> TensorDictBase: + # Deprecating this since _set_tuple will do it thx to the decorator + # if not self.no_convert: + # func = _maybe_make_param + # else: + # func = _maybe_make_param_or_buffer + # if _is_tensor_collection(type(input_dict_or_td)): + # input_dict_or_td = input_dict_or_td.apply(func) + # else: + # input_dict_or_td = tree_map(func, input_dict_or_td) + with self._param_td.unlock_(): + TensorDictBase.update( + self, + input_dict_or_td, + clone=clone, + inplace=inplace, + keys_to_update=keys_to_update, + non_blocking=non_blocking, + is_leaf=is_leaf, + ) + self._reset_params() + return self + + @lock_blocked + @_unlock_and_set + def pop(self, key: NestedKey, default: Any = NO_DEFAULT) -> CompatibleType: ... + + @lock_blocked + @_unlock_and_set + def popitem(self): ... + + @lock_blocked + @_unlock_and_set + def rename_key_( + self, old_key: NestedKey, new_key: NestedKey, safe: bool = False + ) -> TensorDictBase: ... + + def map( + self, + fn: Callable, + dim: int = 0, + num_workers: int | None = None, + chunksize: int | None = None, + num_chunks: int | None = None, + pool: mp.Pool = None, + generator: torch.Generator | None = None, + max_tasks_per_child: int | None = None, + worker_threads: int = 1, + mp_start_method: str | None = None, + ): + raise RuntimeError( + "Cannot call map on a TensorDictParams object. Convert it " + "to a detached tensordict first (through ``tensordict.data`` or ``tensordict.to_tensordict()``) and call " + "map in a second time." + ) + + @_unlock_and_set(inplace=True) + def apply( + self, + fn: Callable, + *others: TensorDictBase, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + default: Any = NO_DEFAULT, + filter_empty: bool | None = None, + call_on_nested: bool = False, + **constructor_kwargs, + ) -> TensorDictBase | None: ... + + @_unlock_and_set(inplace=True) + def named_apply( + self, + fn: Callable, + *others: TensorDictBase, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + default: Any = NO_DEFAULT, + filter_empty: bool | None = None, + call_on_nested: bool = False, + **constructor_kwargs, + ) -> TensorDictBase | None: ... + + @_unlock_and_set(inplace=True) + def _apply_nest(*args, **kwargs): ... + + @_fallback + def _multithread_apply_flat( + self, + fn: Callable, + *others: T, + call_on_nested: bool = False, + default: Any = NO_DEFAULT, + named: bool = False, + nested_keys: bool = False, + prefix: tuple = (), + is_leaf: Callable[[Type], bool] | None = None, + executor: ThreadPoolExecutor, + futures: List[Future], + local_futures: List, + ) -> None: ... + + @_fallback + def _multithread_rebuild( + self, + *, + batch_size: Sequence[int] | None = None, + device: torch.device | None = NO_DEFAULT, + names: Sequence[str] | None = NO_DEFAULT, + inplace: bool = False, + checked: bool = False, + out: TensorDictBase | None = None, + filter_empty: bool = False, + executor: ThreadPoolExecutor, + futures: List[Future], + local_futures: List, + subs_results: Dict[Future, Any] | None = None, + multithread_set: bool = False, # Experimental + **constructor_kwargs, + ) -> None: ... + + @_get_post_hook + @_fallback + def get(self, key: NestedKey, default: Any = None) -> CompatibleType: ... + + @_get_post_hook + @_fallback + def __getitem__( + self, index: IndexType + ) -> Self | Tensor | TensorCollection | Any: ... + + @_fallback + def _set_device(self, device: torch.device) -> Self: ... # type: ignore[misc] + + @_fallback + def _set_names(self, names: Sequence[str] | None) -> None: ... + + @_fallback + def auto_device_(self) -> Self: ... + + __getitems__ = __getitem__ + + def to(self, *args, **kwargs) -> TensorDictBase: + params = self._param_td.to(*args, **kwargs) + if params is self._param_td: + return self + return TensorDictParams(params) + + def cpu(self): + params = self._param_td.cpu() + if params is self._param_td: + return self + return TensorDictParams(params) + + def cuda(self, device=None): + params = self._param_td.cuda(device=device) + if params is self._param_td: + return self + return TensorDictParams(params) + + def _clone(self, recurse: bool = True) -> TensorDictBase: + """Clones the TensorDictParams. + + .. warning:: + The effect of this call is different from a regular torch.Tensor.clone call + in that it will create a TensorDictParams instance with a new copy of the + parameters and buffers __detached__ from the current graph. For a + regular clone (ie, cloning leaf parameters onto a new tensor that + is part of the graph), simply call + + >>> params.apply(torch.clone) + + .. note:: + If a parameter is duplicated in the tree, ``clone`` will preserve this + identity (ie, parameter tying is preserved). + + See :meth:`tensordict.TensorDictBase.clone` for more info on the clone + method. + + """ + if not recurse: + return TensorDictParams._new_unsafe( + self._param_td._clone(False), + no_convert="skip", + params=dict(self._parameters), + buffers=dict(self._buffers), + ) + + memo = {} + + def _clone(tensor, memo=memo): + result = memo.get(tensor) + if result is not None: + return result + + if isinstance(tensor, nn.Parameter): + result = nn.Parameter( + tensor.data.clone(), requires_grad=tensor.requires_grad + ) + else: + result = Buffer(tensor.data.clone()) + memo[tensor] = result + return result + + return TensorDictParams(self._param_td.apply(_clone), no_convert="skip") + + @_fallback + def chunk(self, chunks: int, dim: int = 0) -> tuple[TensorDictBase, ...]: ... + + @_fallback + def _unbind(self, dim: int) -> tuple[TensorDictBase, ...]: ... + + @classmethod + def from_dict(cls, *args, **kwargs): + td = TensorDict.from_dict(*args, **kwargs) + return TensorDictParams(td) + + @_fallback + def to_tensordict(self, *, retain_none: bool | None = None): ... + + @_fallback + def to_h5( + self, + filename, + **kwargs, + ): ... + + def __hash__(self): + return hash((id(self), id(self.__dict__.get("_param_td")))) + + @_fallback + def __eq__(self, other: object) -> TensorDictBase: ... + + @_fallback + def __ne__(self, other: object) -> TensorDictBase: ... + + @_fallback + def __xor__(self, other: object) -> TensorDictBase: ... + + @_fallback + def __or__(self, other: object) -> TensorDictBase: ... + + @_fallback + def __ge__(self, other: object) -> TensorDictBase: ... + + @_fallback + def __gt__(self, other: object) -> TensorDictBase: ... + + @_fallback + def __le__(self, other: object) -> TensorDictBase: ... + + @_fallback + def __lt__(self, other: object) -> TensorDictBase: ... + + def __getattr__(self, item: str) -> Any: + if not item.startswith("_"): + try: + return getattr(self.__dict__["_param_td"], item) + except AttributeError: + try: + return super().__getattr__(item) + except AttributeError as e: + # During some state-dict loads, we may encounter cases where pytorch does a getattr + # with the module name + if item in self.keys(): + return TensorDictParams(self[item]) + raise e + else: + return super().__getattr__(item) + + @_fallback + def _change_batch_size(self, *args, **kwargs): ... + + @_fallback + def _erase_names(self, *args, **kwargs): ... + + @_get_post_hook + @_fallback + def _get_str(self, *args, **kwargs): ... + + @_get_post_hook + @_fallback + def _get_tuple(self, *args, **kwargs): ... + + @_get_post_hook + @_fallback + def _get_at_str(self, key, idx, default, **kwargs): ... + + @_get_post_hook + @_fallback + def _get_at_tuple(self, key, idx, default, **kwargs): ... + + @_fallback + def _add_batch_dim(self, *args, **kwargs): ... + + @_fallback + def _convert_to_tensordict(self, *args, **kwargs): ... + + @_fallback + def _get_names_idx(self, *args, **kwargs): ... + + @_fallback + def _index_tensordict(self, *args, **kwargs): ... + + @_fallback + def _remove_batch_dim(self, *args, **kwargs): ... + + @_fallback + def _maybe_remove_batch_dim(self, *args, **kwargs): ... + + @_fallback + def _has_names(self, *args, **kwargs): ... + + @_unlock_and_set + def _rename_subtds(self, *args, **kwargs): ... + + @_unlock_and_set + def _set_at_str(self, *args, **kwargs): ... + + @_fallback + def _set_at_tuple(self, *args, **kwargs): ... + + @_unlock_and_set + def _set_str(self, *args, **kwargs): ... + + @_unlock_and_set + def _set_tuple(self, *args, **kwargs): ... + + @_unlock_and_set + def _create_nested_str(self, *args, **kwargs): ... + + @_fallback_property + def batch_size(self) -> torch.Size: ... + + @_fallback + def contiguous(self, *args, **kwargs): ... + + @lock_blocked + @_unlock_and_set + def del_(self, *args, **kwargs): ... + + @_fallback + def detach_(self, *args, **kwargs): ... + + @_fallback_property + def device(self): ... + + @_fallback + def entry_class(self, *args, **kwargs): ... + + @_fallback + def is_contiguous(self, *args, **kwargs): ... + + @_fallback + def keys(self, *args, **kwargs): ... + + @_fallback + def masked_fill(self, *args, **kwargs): ... + + @_fallback + def masked_fill_(self, *args, **kwargs): ... + + def memmap_( + self, + prefix: str | None = None, + copy_existing: bool = False, + num_threads: int = 0, + ) -> TensorDictBase: + raise RuntimeError( + "Cannot build a memmap TensorDict in-place. Use memmap or memmap_like instead." + ) + + _memmap_ = TensorDict._memmap_ + + _load_memmap = TensorDict._load_memmap + + def make_memmap( + self, + key: NestedKey, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for TensorDictParams." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_storage( + self, + key: NestedKey, + storage: torch.UntypedStorage, + shape: torch.Size | torch.Tensor, + *, + dtype: torch.dtype | None = None, + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for TensorDictParams." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + def make_memmap_from_tensor( + self, key: NestedKey, tensor: torch.Tensor, *, copy_data: bool = True + ) -> MemoryMappedTensor: + raise RuntimeError( + "Making a memory-mapped tensor after instantiation isn't currently allowed for TensorDictParams." + "If this feature is required, open an issue on GitHub to trigger a discussion on the topic!" + ) + + @_fallback_property + def names(self): ... + + def pin_memory(self, *args, **kwargs): + if kwargs.get("inplace", False): + raise RuntimeError( + f"Cannot pin_memory in-place with {type(self).__name__}." + ) + return _fallback(self.pin_memory)(self, *args, **kwargs) + + @_unlock_and_set + def _select(self, *args, **kwargs): ... + + @_fallback + def share_memory_(self, *args, **kwargs): ... + + @property + def is_locked(self) -> bool: + # Cannot be locked + return self._is_locked + + @is_locked.setter + def is_locked(self, value): + self._is_locked = bool(value) + + @_fallback_property + def is_shared(self) -> bool: ... + + @_fallback_property + def is_memmap(self) -> bool: ... + + @property + def _is_shared(self) -> bool: + return self._param_td._is_shared + + @property + def _is_memmap(self) -> bool: + return self._param_td._is_memmap + + @_fallback_property + def shape(self) -> torch.Size: ... + + def _propagate_lock(self, _lock_parents_weakrefs=None, *, is_compiling): + """Registers the parent tensordict that handles the lock.""" + self._is_locked = True + if not is_compiling: + if _lock_parents_weakrefs is None: + _lock_parents_weakrefs = [] + self._lock_parents_weakrefs += _lock_parents_weakrefs + _lock_parents_weakrefs.append(weakref.ref(self)) + # we don't want to double-lock the _param_td attrbute which is locked by default + if not self._param_td.is_locked: + self._param_td._propagate_lock( + _lock_parents_weakrefs, is_compiling=is_compiling + ) + + @erase_cache + def _propagate_unlock(self): + # if we end up here, we can clear the graph associated with this td + self._is_locked = False + + if not self._lock_content: + return self._param_td._propagate_unlock() + return [] + + unlock_ = TensorDict.unlock_ + lock_ = TensorDict.lock_ + + @property + def data(self) -> Self: + return self._param_td._data() + + @property + def grad(self): + return self._param_td._grad() + + @_unlock_and_set(inplace=True) + def flatten_keys( + self, separator: str = ".", inplace: bool = False + ) -> TensorDictBase: ... + + @_unlock_and_set(inplace=True) + def unflatten_keys( + self, separator: str = ".", inplace: bool = False + ) -> TensorDictBase: ... + + @_unlock_and_set(inplace=True) + def _exclude( + self, *keys: NestedKey, inplace: bool = False, set_shared: bool = True + ) -> TensorDictBase: ... + + @_carry_over + def from_dict_instance( + self, + input_dict, + *, + auto_batch_size: bool = False, + batch_size=None, + device=None, + batch_dims=None, + ): ... + + @_carry_over + def _legacy_transpose(self, dim0, dim1): ... + + @_fallback + def _transpose(self, dim0, dim1): ... + + @_fallback + def where( + self, + condition: Tensor, + other: Tensor | TensorDictBase, + *, + out: TensorDictBase | None = None, + pad: int | bool = None, + update_batch_size: bool = False, + ): ... + + @_fallback + def _permute( + self, + *dims_list: int, + dims: list[int] | None = None, + ) -> TensorDictBase: ... + + @_carry_over + def _legacy_permute( + self, + *dims_list: int, + dims: list[int] | None = None, + ) -> TensorDictBase: ... + + @_fallback + def _squeeze(self, dim: int | None = None) -> TensorDictBase: ... + + @_carry_over + def _legacy_squeeze(self, dim: int | None = None) -> TensorDictBase: ... + + @_fallback + def _unsqueeze(self, dim: int) -> TensorDictBase: ... + + @_carry_over + def _legacy_unsqueeze(self, dim: int) -> TensorDictBase: ... + + _check_device = TensorDict._check_device + _check_is_shared = TensorDict._check_is_shared + + @_fallback + def _cast_reduction(self, **kwargs): ... + + @_fallback + def all(self, dim: int | None = None) -> bool | TensorDictBase: ... + + @_fallback + def any(self, dim: int | None = None) -> bool | TensorDictBase: ... + + @_fallback + def expand(self, *args, **kwargs) -> Self: ... + + @_fallback + def masked_select(self, mask: Tensor) -> Self: ... + + @_fallback + def memmap_like( + self, + prefix: str | None = None, + copy_existing: bool = False, + num_threads: int = 0, + ) -> Self: ... + + @_fallback + def reshape(self, *shape: int): ... + + @_fallback + def repeat_interleave(self, *shape: int): ... + + @_fallback + def _repeat(self, *repeats: int): ... + + @_fallback + def split( + self, split_size: int | list[int], dim: int = 0 + ) -> list[TensorDictBase]: ... + + @_fallback + def _to_module( + self, + module, + *, + inplace: bool = False, + return_swap: bool = True, + swap_dest=None, + memo=None, + use_state_dict: bool = False, + non_blocking: bool = False, + ): ... + + @_fallback + def _view(self, *args, **kwargs): ... + + @_carry_over + def _legacy_view(self, *args, **kwargs): ... + + @_unlock_and_set + def create_nested(self, key): ... + + def __repr__(self): + return f"TensorDictParams(params={self._param_td})" + + def values( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> Iterator[CompatibleType]: + if is_leaf is None: + is_leaf = _default_is_leaf + for v in self._param_td.values(include_nested, leaves_only, sort=sort): + if not is_leaf(type(v)): + yield v + continue + yield self._apply_get_post_hook(v) + + def state_dict( + self, *args, destination=None, prefix="", keep_vars=False, flatten=True + ): + # flatten must be True by default to comply with module's state-dict API + # since we want all params to be visible at root + return self._param_td.state_dict( + destination=destination, + prefix=prefix, + keep_vars=keep_vars, + flatten=flatten, + ) + + def load_state_dict( + self, state_dict: OrderedDict[str, Any], strict=True, assign=False + ): + # The state-dict is presumably the result of a call to TensorDictParams.state_dict + # but can't be sure. + + state_dict_tensors = {} + state_dict = dict(state_dict) + for k, v in list(state_dict.items()): + if isinstance(v, torch.Tensor): + del state_dict[k] + state_dict_tensors[k] = v + state_dict_tensors = dict( + TensorDict(state_dict_tensors, []).unflatten_keys(".") + ) + state_dict.update(state_dict_tensors) + self.data.load_state_dict(state_dict, strict=True, assign=False) + return self + + def _load_from_state_dict( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + data = TensorDict( + { + key: val + for key, val in state_dict.items() + if key.startswith(prefix) and val is not None + }, + [], + ).unflatten_keys(".") + prefix = tuple(key for key in prefix.split(".") if key) + if prefix: + data = data.get(prefix) + self.data.load_state_dict(data) + + def items( + self, + include_nested: bool = False, + leaves_only: bool = False, + is_leaf: Callable[[Type], bool] | None = None, + *, + sort: bool = False, + ) -> Iterator[CompatibleType]: + if is_leaf is None: + is_leaf = _default_is_leaf + for k, v in self._param_td.items(include_nested, leaves_only, sort=sort): + if not is_leaf(type(v)): + yield k, v + continue + yield k, self._apply_get_post_hook(v) + + @_apply_on_data + def zero_(self) -> Self: ... + + @_apply_on_data + def fill_(self, key: NestedKey, value: float | bool) -> Self: ... + + @_apply_on_data + def copy_(self, tensordict: T, non_blocking: bool | None = None) -> T: ... + + @_apply_on_data + def set_at_( + self, key: NestedKey, value: CompatibleType, index: IndexType + ) -> Self: ... + + @_apply_on_data + def set_( + self, + key: NestedKey, + item: CompatibleType, + ) -> Self: ... + + @_apply_on_data + def _stack_onto_( + self, + list_item: list[CompatibleType], + dim: int, + ) -> Self: ... + + @_apply_on_data + def _stack_onto_at_( + self, + key: NestedKey, + list_item: list[CompatibleType], + dim: int, + idx: IndexType, + ) -> Self: ... + + @_apply_on_data + def update_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> T: ... + + @_apply_on_data + def update_at_( + self, + input_dict_or_td: dict[str, CompatibleType] | T, + idx: IndexType, + clone: bool = False, + *, + non_blocking: bool = False, + keys_to_update: Sequence[NestedKey] | None = None, + ) -> T: ... + + @_apply_on_data + def apply_(self, fn: Callable, *others, **kwargs) -> Self: ... + + @implement_for("torch", "2.1") + def _apply(self, fn, recurse=True): + self._param_td._erase_cache() + param_td = self._param_td + self._param_td = param_td.copy() + # Keep a list of buffers to update .data only + bufs = dict(self._buffers) + out: TensorDictBase = super()._apply(fn, recurse=recurse) + for key, val in bufs.items(): + val.data = self._buffers[key].data + self._buffers[key] = val + # Check device and shape + cbs = out._check_batch_size(raise_exception=False) + if not cbs: + out.auto_batch_size_() + cd = out._check_device(raise_exception=False) + if not cd: + out.auto_device_() + return out + + @implement_for("torch", None, "2.1") + def _apply(self, fn): # noqa: F811 + self._param_td._erase_cache() + param_td = self._param_td + self._param_td = param_td.copy() + # Keep a list of buffers to update .data only + bufs = dict(self._buffers) + out: TensorDictBase = super()._apply(fn) + for key, val in bufs.items(): + val.data = self._buffers[key].data + self._buffers[key] = val + # Check device and shape + cbs = out._check_batch_size(raise_exception=False) + if not cbs: + out.auto_batch_size_() + cd = out._check_device(raise_exception=False) + if not cd: + out.auto_device_() + return out + + +TDPARAM_HANDLED_FUNCTIONS = copy(TD_HANDLED_FUNCTIONS) + + +def implements_for_tdparam(torch_function: Callable) -> Callable[[Callable], Callable]: + """Register a torch function override for TensorDictParams.""" + + @functools.wraps(torch_function) + def decorator(func: Callable) -> Callable: + TDPARAM_HANDLED_FUNCTIONS[torch_function] = func + return func + + return decorator + + +@implements_for_tdparam(torch.empty_like) +def _empty_like(td: TensorDictBase, *args, **kwargs) -> TensorDictBase: + return td.apply( + lambda x: torch.empty_like(x, *args, **kwargs), + device=kwargs.pop("device", NO_DEFAULT), + ) + + +_register_tensor_class(TensorDictParams) diff --git a/lib/python3.12/site-packages/tensordict/nn/probabilistic.py b/lib/python3.12/site-packages/tensordict/nn/probabilistic.py new file mode 100644 index 0000000000000000000000000000000000000000..7ed87558e3f5ba2cdef4913343edd4d96df71fc6 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/probabilistic.py @@ -0,0 +1,1429 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import collections +import re +import warnings +from collections.abc import MutableSequence +from textwrap import indent +from typing import Any, Dict, List, OrderedDict, overload, TYPE_CHECKING + +import torch + +from tensordict._nestedkey import NestedKey +from tensordict._td import TensorDict +from tensordict.base import is_tensor_collection +from tensordict.nn.common import dispatch, TensorDictModuleBase + +from tensordict.nn.distributions.composite import _add_suffix, CompositeDistribution +from tensordict.nn.distributions.continuous import Delta +from tensordict.nn.distributions.discrete import OneHotCategorical +from tensordict.nn.distributions.truncated_normal import ( + TruncatedNormal, + TruncatedStandardNormal, +) +from tensordict.nn.sequence import TensorDictSequential +from tensordict.nn.utils import ( + _set_skip_existing_None, + composite_lp_aggregate, + set_composite_lp_aggregate, +) +from tensordict.tensorclass import is_non_tensor +from tensordict.tensordict import TensorDictBase +from tensordict.utils import _ContextManager, _zip_strict, unravel_key +from torch import distributions as D, Tensor +from torch.utils._contextlib import _DecoratorContextManager + +try: + from torch.compiler import is_compiling +except ImportError: + from torch._dynamo import is_compiling + +try: + from enum import StrEnum +except ImportError: + from .utils import StrEnum + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + +__all__ = ["ProbabilisticTensorDictModule", "ProbabilisticTensorDictSequential"] + + +class InteractionType(StrEnum): + """A list of possible interaction types with a distribution. + + MODE, MEDIAN and MEAN point to the property / attribute with the same name. + RANDOM points to ``rsample()`` if that method exists or ``sample()`` if not. + + DETERMINISTIC can be used as a generic fallback if ``MEAN`` or ``MODE`` are not guaranteed to + be analytically tractable. In such cases, a rude deterministic estimate can be used + in some cases even if it lacks a true algebraic meaning. + This value will trigger a query to the ``deterministic_sample`` attribute in the distribution + and if it does not exist, the ``mean`` will be used. + + """ + + MODE = "mode" + MEDIAN = "median" + MEAN = "mean" + RANDOM = "random" + DETERMINISTIC = "deterministic" + + @classmethod + def from_str(cls, type_str: str) -> InteractionType: + """Return the interaction_type with name matched to the provided string (case insensitive).""" + return cls(type_str.lower()) + + @classmethod + def _missing_(cls, value): + # Normalize input: if a string, lower it + if isinstance(value, str): + value = value.lower() + # Try to match with the enum values + for member in cls: + if member.value == value: + return member + return None # Default behavior if not found + + +_interaction_type = _ContextManager() + + +DETERMINISTIC_REGISTER = {} + +dist_has_enum_support = {} +# Iterate over all distribution classes in torch.distributions +for dist_name in dir(D): + dist_cls = getattr(D, dist_name) + + # Check if it's a class (not a function or variable) and is a subclass of Distribution + if isinstance(dist_cls, type) and issubclass(dist_cls, D.Distribution): + if dist_cls is D.LogisticNormal: + DETERMINISTIC_REGISTER[dist_cls] = InteractionType.DETERMINISTIC + elif getattr(dist_cls, "has_enumerate_support", False): + DETERMINISTIC_REGISTER[dist_cls] = InteractionType.MODE + else: + DETERMINISTIC_REGISTER[dist_cls] = InteractionType.MEAN + + +DETERMINISTIC_REGISTER[Delta] = InteractionType.DETERMINISTIC +DETERMINISTIC_REGISTER[OneHotCategorical] = InteractionType.MODE + +DETERMINISTIC_REGISTER[TruncatedNormal] = InteractionType.MEAN +DETERMINISTIC_REGISTER[TruncatedStandardNormal] = InteractionType.MEAN + + +def interaction_type() -> InteractionType | None: + """Returns the current sampling type.""" + return _interaction_type.get_mode() + + +class set_interaction_type(_DecoratorContextManager): + """Sets all ProbabilisticTDModules sampling to the desired type. + + Args: + type (InteractionType or str): sampling type to use when the policy is being called. + + """ + + def __init__( + self, type: InteractionType | str | None = InteractionType.DETERMINISTIC + ) -> None: + super().__init__() + if not isinstance(type, InteractionType) and type is not None: + if isinstance(type, str): + type = InteractionType(type.lower()) + else: + raise ValueError(f"{type} is not a valid InteractionType") + self.type = type + + def clone(self) -> set_interaction_type: + # override this method if your children class takes __init__ parameters + return type(self)(self.type) + + def __enter__(self) -> None: + self.prev = _interaction_type.get_mode() + _interaction_type.set_mode(self.type) + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + _interaction_type.set_mode(self.prev) + + +class ProbabilisticTensorDictModule(TensorDictModuleBase): + """A probabilistic TD Module. + + `ProbabilisticTensorDictModule` is a non-parametric module embedding a + probability distribution constructor. It reads the distribution parameters from an input + TensorDict using the specified `in_keys` and outputs a sample (loosely speaking) of the + distribution. + + The output "sample" is produced given some rule, specified by the input ``default_interaction_type`` + argument and the ``interaction_type()`` global function. + + `ProbabilisticTensorDictModule` can be used to construct the distribution + (through the :meth:`~.get_dist` method) and/or sampling from this distribution + (through a regular :meth:`~.__call__` to the module). + + A `ProbabilisticTensorDictModule` instance has two main features: + + - It reads and writes from and to TensorDict objects; + - It uses a real mapping R^n -> R^m to create a distribution in R^d from + which values can be sampled or computed. + + When the :meth:`~.__call__` and :meth:`~.forward` method are called, a distribution is + created, and a value computed (depending on the ``interaction_type`` value, 'dist.mean', + 'dist.mode', 'dist.median' attributes could be used, as well as + the 'dist.rsample', 'dist.sample' method). The sampling step is skipped if the supplied + TensorDict has all the desired key-value pairs already. + + By default, `ProbabilisticTensorDictModule` distribution class is a :class:`~torchrl.modules.distributions.Delta` + distribution, making `ProbabilisticTensorDictModule` a simple wrapper around + a deterministic mapping function. + + + Args: + in_keys (NestedKey | List[NestedKey] | Dict[str, NestedKey]): key(s) that will be read from the input TensorDict + and used to build the distribution. + Importantly, if it's a list of NestedKey or a NestedKey, the leaf (last element) of those keys must match the keywords used by + the distribution class of interest, e.g. ``"loc"`` and ``"scale"`` for + the :class:`~torch.distributions.Normal` distribution and similar. + If in_keys is a dictionary, the keys are the keys of the distribution and the values are the keys in the + tensordict that will get match to the corresponding distribution keys. + out_keys (NestedKey | List[NestedKey] | None): key(s) where the sampled values will be written. + Importantly, if these keys are found in the input TensorDict, the sampling step will be skipped. + + Keyword Args: + default_interaction_type (InteractionType, optional): keyword-only argument. + Default method to be used to retrieve + the output value. Should be one of InteractionType: MODE, MEDIAN, MEAN or RANDOM + (in which case the value is sampled randomly from the distribution). Default + is MODE. + + .. note:: When a sample is drawn, the + :class:`ProbabilisticTensorDictModule` instance will + first look for the interaction mode dictated by the + :func:`~tensordict.nn.probabilistic.interaction_type` + global function. If this returns `None` (its default value), then the + `default_interaction_type` of the `ProbabilisticTDModule` + instance will be used. Note that + :class:`~torchrl.collectors.collectors.DataCollectorBase` + instances will use `set_interaction_type` to + :class:`tensordict.nn.InteractionType.RANDOM` by default. + + .. note:: + In some cases, the mode, median or mean value may not be + readily available through the corresponding attribute. + To paliate this, :class:`~ProbabilisticTensorDictModule` will first attempt + to get the value through a call to ``get_mode()``, ``get_median()`` or ``get_mean()`` + if the method exists. + + distribution_class (Type or Callable[[Any], Distribution], optional): keyword-only argument. + A :class:`torch.distributions.Distribution` class to + be used for sampling. + Default is :class:`~tensordict.nn.distributions.Delta`. + + .. note:: + If the distribution class is of type + :class:`~tensordict.nn.distributions.CompositeDistribution`, the ``out_keys`` + can be inferred directly form the ``"distribution_map"`` or ``"name_map"`` + keywork arguments provided through this class' ``distribution_kwargs`` + keyword argument, making the ``out_keys`` optional in such cases. + + distribution_kwargs (dict, optional): keyword-only argument. + Keyword-argument pairs to be passed to the distribution. + + .. note:: if your kwargs contain tensors that you would like to transfer to device with the module, or + tensors that should see their dtype modified when calling `module.to(dtype)`, you can wrap the kwargs + in a :class:`~tensordict.nn.TensorDictParams` to do this automatically. + + return_log_prob (bool, optional): keyword-only argument. + If ``True``, the log-probability of the + distribution sample will be written in the tensordict with the key + `log_prob_key`. Default is ``False``. + log_prob_keys (List[NestedKey], optional): keys where to write the log_prob if ``return_log_prob=True``. + Defaults to `'_log_prob'`, where `` is each of the :attr:`out_keys`. + + .. note:: This is only available when :func:`~tensordict.nn.probabilistic.composite_lp_aggregate` is set to ``False``. + + log_prob_key (NestedKey, optional): key where to write the log_prob if ``return_log_prob=True``. + Defaults to `'sample_log_prob'` when :func:`~tensordict.nn.probabilistic.composite_lp_aggregate` is set to `True` + or `'_log_prob'` otherwise. + + .. note:: When there is more than one sample, this is only available when :func:`~tensordict.nn.probabilistic.composite_lp_aggregate` is set to ``True``. + + cache_dist (bool, optional): keyword-only argument. + EXPERIMENTAL: if ``True``, the parameters of the + distribution (i.e. the output of the module) will be written to the + tensordict along with the sample. Those parameters can be used to re-compute + the original distribution later on (e.g. to compute the divergence between + the distribution used to sample the action and the updated distribution in + PPO). Default is ``False``. + n_empirical_estimate (int, optional): keyword-only argument. + Number of samples to compute the empirical + mean when it is not available. Defaults to 1000. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> from tensordict.nn import ( + ... ProbabilisticTensorDictModule, + ... ProbabilisticTensorDictSequential, + ... TensorDictModule, + ... ) + >>> from tensordict.nn.distributions import NormalParamExtractor + >>> from tensordict.nn.functional_modules import make_functional + >>> from torch.distributions import Normal, Independent + >>> td = TensorDict( + ... {"input": torch.randn(3, 4), "hidden": torch.randn(3, 8)}, [3] + ... ) + >>> net = torch.nn.GRUCell(4, 8) + >>> module = TensorDictModule( + ... net, in_keys=["input", "hidden"], out_keys=["params"] + ... ) + >>> normal_params = TensorDictModule( + ... NormalParamExtractor(), in_keys=["params"], out_keys=["loc", "scale"] + ... ) + >>> def IndepNormal(**kwargs): + ... return Independent(Normal(**kwargs), 1) + >>> prob_module = ProbabilisticTensorDictModule( + ... in_keys=["loc", "scale"], + ... out_keys=["action"], + ... distribution_class=IndepNormal, + ... return_log_prob=True, + ... ) + >>> td_module = ProbabilisticTensorDictSequential( + ... module, normal_params, prob_module + ... ) + >>> params = TensorDict.from_module(td_module) + >>> with params.to_module(td_module): + ... _ = td_module(td) + >>> print(td) + TensorDict( + fields={ + action: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + hidden: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + input: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + loc: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + params: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + sample_log_prob: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + scale: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + >>> with params.to_module(td_module): + ... dist = td_module.get_dist(td) + >>> print(dist) + Independent(Normal(loc: torch.Size([3, 4]), scale: torch.Size([3, 4])), 1) + >>> # we can also apply the module to the TensorDict with vmap + >>> from torch import vmap + >>> params = params.expand(4) + >>> def func(td, params): + ... with params.to_module(td_module): + ... return td_module(td) + >>> td_vmap = vmap(func, (None, 0))(td, params) + >>> print(td_vmap) + TensorDict( + fields={ + action: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + hidden: Tensor(shape=torch.Size([4, 3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + input: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + loc: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + params: Tensor(shape=torch.Size([4, 3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + sample_log_prob: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False), + scale: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 3]), + device=None, + is_shared=False) + + """ + + def __init__( + self, + in_keys: NestedKey | List[NestedKey] | Dict[str, NestedKey], + out_keys: NestedKey | List[NestedKey] | None = None, + *, + default_interaction_type: InteractionType = InteractionType.DETERMINISTIC, + distribution_class: type = Delta, + distribution_kwargs: dict | None = None, + return_log_prob: bool = False, + log_prob_keys: List[NestedKey] | None = None, + log_prob_key: NestedKey | None = None, + cache_dist: bool = False, + n_empirical_estimate: int = 1000, + num_samples: int | torch.Size | None = None, + ) -> None: + super().__init__() + distribution_kwargs = ( + distribution_kwargs if distribution_kwargs is not None else {} + ) + if isinstance(in_keys, (str, tuple)): + in_keys = [in_keys] + if isinstance(out_keys, (str, tuple)): + out_keys = [out_keys] + elif out_keys is None: + if distribution_class is CompositeDistribution: + distribution_map = distribution_kwargs.get("distribution_map") + if distribution_map is None: + raise KeyError( + "'distribution_map' must be provided within " + "distribution_kwargs whenever the distribution is of type CompositeDistribution." + ) + name_map = distribution_kwargs.get("name_map") + if name_map is not None: + out_keys = list(name_map.values()) + else: + out_keys = list(distribution_map.keys()) + else: + out_keys = ["_"] + if isinstance(in_keys, dict): + dist_keys, in_keys = zip(*in_keys.items()) + if set(map(type, dist_keys)) != {str}: + raise ValueError( + f"If in_keys is dict, its keys must be strings matching to the distribution kwargs." + f"{type(self).__name__} got {dist_keys}" + ) + else: + dist_keys = in_keys + + self._out_keys = [unravel_key(k) for k in out_keys] + self.in_keys = in_keys + self.dist_keys = dist_keys + if log_prob_keys is None: + if len(out_keys) == 1 and log_prob_key is not None: + log_prob_keys = [log_prob_key] + elif composite_lp_aggregate(nowarn=True): + if len(out_keys) == 1: + log_prob_keys = ["sample_log_prob"] + else: + log_prob_keys = None + else: + log_prob_keys = [ + _add_suffix(key, "_log_prob") for key in self._out_keys + ] + elif composite_lp_aggregate(nowarn=True): + raise RuntimeError( + "composite_lp_aggregate is set to True but log_prob_keys were passed. " + "When composite_lp_aggregate() returns ``True``, log_prob_key must be used instead." + ) + self._trigger_warning_lpk = len(self._out_keys) > 1 + if log_prob_key is None: + if composite_lp_aggregate(nowarn=True): + log_prob_key = "sample_log_prob" + self._trigger_warning_lpk = True + elif len(out_keys) == 1: + log_prob_key = _add_suffix(out_keys[0], "_log_prob") + elif len(out_keys) > 1 and not composite_lp_aggregate(nowarn=True): + raise RuntimeError( + "composite_lp_aggregate is set to `False` but a `log_prob_key` was passed. " + "When composite_lp_aggregate() returns ``False``, log_prob_keys must be used instead." + ) + self._log_prob_key = log_prob_key + self._log_prob_keys = log_prob_keys + + self.default_interaction_type = InteractionType(default_interaction_type) + + if isinstance(distribution_class, str): + from tensordict.nn.distributions import distributions_maps + + distribution_class = distributions_maps.get(distribution_class.lower()) + self.distribution_class = distribution_class + self.distribution_kwargs = distribution_kwargs + self.n_empirical_estimate = n_empirical_estimate + self._dist = None + self.cache_dist = cache_dist if hasattr(distribution_class, "update") else False + self.return_log_prob = return_log_prob + if isinstance(num_samples, (int, torch.SymInt)): + num_samples = torch.Size((num_samples,)) + self.num_samples = num_samples + self._composite_lp_aggreate_at_init = composite_lp_aggregate() + + def __repr__(self): + return ( + f"{type(self).__name__}(" + f"\n{indent(f'in_keys={self.in_keys}', 4 * ' ')}," + f"\n{indent(f'out_keys={self.out_keys}', 4 * ' ')}," + f"\n{indent(f'distribution_class={self.distribution_class}', 4 * ' ')}, " + f"\n{indent(f'distribution_kwargs={self.distribution_kwargs})', 4 * ' ')}," + f"\n{indent(f'default_interaction_type={self.default_interaction_type})', 4 * ' ')}," + f"\n{indent(f'num_samples={self.num_samples})', 4 * ' ')})" + ) + + @property + def out_keys(self) -> List[NestedKey]: + out_keys = list(self._out_keys) + if self.return_log_prob: + if not composite_lp_aggregate(): + if out_keys[-len(self.log_prob_keys) :] != self.log_prob_keys: + out_keys.extend(self.log_prob_keys) + elif self.log_prob_key not in out_keys: + out_keys.append(self.log_prob_key) + return out_keys + + @property + def log_prob_key(self): + clpa = composite_lp_aggregate(nowarn=True) + if clpa != self._composite_lp_aggreate_at_init: + raise RuntimeError( + f"composite_lp_aggregate is set to `{clpa}`, but the class was instantiated with `{self._composite_lp_aggreate_at_init}` " + f"which may affect the log_prob_key property." + ) + if not clpa: + if len(self.log_prob_keys) == 1: + return self.log_prob_keys[0] + raise RuntimeError( + f"composite_lp_aggregate is set to ``False``, hence {type(self).__name__}.log_prob_key cannot be accessed " + f"unless there is one and only one element in log_prob_keys (got log_prob_keys={self.log_prob_keys}). " + f"When composite_lp_aggregate() returns ``False``, try to use {type(self).__name__}.log_prob_keys instead." + ) + return self._log_prob_key + + @property + def log_prob_keys(self): + clpa = composite_lp_aggregate(nowarn=True) + if clpa != self._composite_lp_aggreate_at_init: + raise RuntimeError( + f"composite_lp_aggregate is set to `{clpa}`, but the class was instantiated with `{self._composite_lp_aggreate_at_init}` " + f"which may affect the log_prob_keys property." + ) + if clpa: + return [self.log_prob_key] + return self._log_prob_keys + + @property + def dist_params_keys(self) -> List[NestedKey]: + """Returns all the keys pointing at the distribution params.""" + return list(self.in_keys) + + @property + def dist_sample_keys(self) -> List[NestedKey]: + """Returns all the keys pointing at the distribution samples.""" + return list(self._out_keys) + + def get_dist(self, tensordict: TensorDictBase) -> D.Distribution: + """Creates a :class:`torch.distribution.Distribution` instance with the parameters provided in the input tensordict. + + Args: + tensordict (TensorDictBase): The input tensordict containing the distribution parameters. + + Returns: + A :class:`torch.distribution.Distribution` instance created from the input tensordict. + + Raises: + TypeError: If the input tensordict does not match the distribution keywords. + """ + try: + dist_kwargs = {} + for dist_key, td_key in _zip_strict(self.dist_keys, self.in_keys): + if isinstance(dist_key, tuple): + dist_key = dist_key[-1] + dist_kwargs[dist_key] = tensordict.get(td_key, None) + dist = self.distribution_class( + **dist_kwargs, **_dynamo_friendly_to_dict(self.distribution_kwargs) + ) + except TypeError as err: + if "an unexpected keyword argument" in str(err): + raise TypeError( + "distribution keywords and tensordict keys indicated by ProbabilisticTensorDictModule.dist_keys must match. " + f"Got this error message: \n{indent(str(err), 4 * ' ')}\nwith dist_keys={self.dist_keys}" + ) + elif re.search(r"missing.*required positional arguments", str(err)): + raise TypeError( + f"TensorDict with keys {tensordict.keys()} does not match the distribution {self.distribution_class} keywords." + ) + else: + raise err + return dist + + build_dist_from_params = get_dist + + _CHANGE_IN_C_LP_A = "The value returned by composite_lp_aggregate changed between init of {} and its execution ({} -> {}). Make sure the mode matches." + + def log_prob( + self, + tensordict, + *, + dist: torch.distributions.Distribution | None = None, + ): + """Computes the log-probability of the distribution sample. + + Args: + tensordict (TensorDictBase): The input tensordict containing the distribution parameters. + dist (torch.distributions.Distribution, optional): The distribution instance. Defaults to ``None``. + If ``None``, the distribution will be computed using the `get_dist` method. + + Returns: + A tensor representing the log-probability of the distribution sample. + """ + if dist is None: + dist = self.get_dist(tensordict) + if isinstance(dist, CompositeDistribution): + clpa = composite_lp_aggregate() + if clpa != self._composite_lp_aggreate_at_init: + raise RuntimeError( + self._CHANGE_IN_C_LP_A.format( + type(self).__name__, self._composite_lp_aggreate_at_init, clpa + ) + ) + if clpa: + # Old behaviour - discouraged + with set_composite_lp_aggregate(False): + td = dist.log_prob(tensordict) + tensordict.update(td) + lp = sum(td.sum(dim="feature").values(True, True)) + return lp + else: + # Check the values within the dist - if not set, choose defaults + lp = dist.log_prob(tensordict) + self._update_td_lp(lp) + return lp + else: + return dist.log_prob(tensordict.get(self.out_keys[0])) + + def _update_td_lp(self, lp): + for out_key, lp_key in _zip_strict(self.dist_sample_keys, self.log_prob_keys): + lp_key_expected = _add_suffix(out_key, "_log_prob") + if lp_key != lp_key_expected: + lp.rename_key_(lp_key_expected, lp_key) + + @property + def SAMPLE_LOG_PROB_KEY(self): + raise RuntimeError( + "SAMPLE_LOG_PROB_KEY is fully deprecated. Use `obj.log_prob_key` instead." + ) + + @dispatch(auto_batch_size=False) + @_set_skip_existing_None() + def forward( + self, + tensordict: TensorDictBase, + tensordict_out: TensorDictBase | None = None, + _requires_sample: bool = True, + ) -> TensorDictBase: + if tensordict_out is None: + tensordict_out = tensordict + + dist = self.get_dist(tensordict) + if _requires_sample: + out_tensors = self._dist_sample(dist, interaction_type=interaction_type()) + if self.num_samples is not None: + # TODO: capture contiguous error here + tensordict_out = tensordict_out.expand( + self.num_samples + tensordict_out.shape + ) + if isinstance(out_tensors, TensorDictBase): + if self.return_log_prob: + clpa = composite_lp_aggregate() + if clpa != self._composite_lp_aggreate_at_init: + raise RuntimeError( + self._CHANGE_IN_C_LP_A.format( + type(self).__name__, + self._composite_lp_aggreate_at_init, + clpa, + ) + ) + if clpa: + with set_composite_lp_aggregate(False): + # We want the tensordict to do the sum and such + log_prob = dist.log_prob(out_tensors) + if log_prob is not out_tensors: + # Composite dists return the tensordict_out directly when aggrgate_prob is False + out_tensors.update(log_prob) + out_tensors.set( + self.log_prob_key, + sum(log_prob.sum(dim="feature").values(True, True)), + ) + else: + log_prob = dist.log_prob(out_tensors) + out_tensors.update(log_prob) + self._update_td_lp(log_prob) + tensordict_out.update(out_tensors) + else: + if isinstance(out_tensors, Tensor): + out_tensors = (out_tensors,) + tensordict_out.update( + dict(_zip_strict(self.dist_sample_keys, out_tensors)) + ) + if self.return_log_prob: + log_prob = dist.log_prob(*out_tensors) + tensordict_out.set(self.log_prob_key, log_prob) + elif self.return_log_prob: + out_tensors = [ + tensordict.get(key) + for key in self.out_keys + if key not in self.log_prob_keys + ] + log_prob = dist.log_prob(*out_tensors) + tensordict_out.set(self.log_prob_key, log_prob) + # raise RuntimeError( + # "ProbabilisticTensorDictModule.return_log_prob = True is incompatible with settings in which " + # "the submodule is responsible for sampling. To manually gather the log-probability, call first " + # "\n>>> dist, tensordict = tensordict_module.get_dist(tensordict)" + # "\n>>> tensordict.set('sample_log_prob', dist.log_prob(tensordict.get(sample_key))" + # ) + return tensordict_out + + def _dist_sample( + self, + dist: D.Distribution, + interaction_type: InteractionType | None = None, + ) -> tuple[Tensor, ...] | Tensor: + if not isinstance(dist, D.Distribution): + raise TypeError("Expected Distribution, but got {}".format(type(dist))) + if interaction_type is None: + interaction_type = self.default_interaction_type + if isinstance(dist, D.LKJCholesky) and interaction_type in ( + InteractionType.DETERMINISTIC, + InteractionType.MEAN, + InteractionType.MODE, + ): + raise RuntimeError( + f"DETERMINISTIC, MEAN and MODE are not implemented for {type(dist).__name__}." + ) + if interaction_type is InteractionType.DETERMINISTIC: + if hasattr(dist, "deterministic_sample"): + return dist.deterministic_sample + else: + # Fallbacks + tdist = type(dist) + if issubclass(tdist, D.Independent): + tdist = type(dist.base_dist) + interaction_type = DETERMINISTIC_REGISTER.get(tdist) + if interaction_type is None: + try: + support = dist.support + fallback = ( + "mean" + if isinstance(support, D.constraints._Real) + else "mode" + ) + except NotImplementedError: + # Some custom dists don't have a support + # We arbitrarily fall onto 'mean' in these cases + fallback = "mean" + try: + if fallback == "mean": + interaction_type = InteractionType.MEAN + elif fallback == "mode": + # Categorical dists don't have an average + interaction_type = InteractionType.MODE + else: + raise AttributeError + except AttributeError: + raise NotImplementedError( + f"method {type(dist)}.deterministic_sample is not implemented, no replacement found." + ) + finally: + warnings.warn( + f"deterministic_sample wasn't found when queried on {type(dist)}. " + f"{type(self).__name__} is falling back on {fallback} instead. " + f"For better code quality and efficiency, make sure to either " + f"provide a distribution with a deterministic_sample attribute or " + f"to change the InteractionMode to the desired value.", + category=UserWarning, + ) + + if interaction_type is InteractionType.MODE: + try: + return dist.mode + except AttributeError: + raise NotImplementedError( + f"method {type(dist)}.mode is not implemented" + ) + + elif interaction_type is InteractionType.MEDIAN: + try: + return dist.median + except AttributeError: + raise NotImplementedError( + f"method {type(dist)}.median is not implemented" + ) + + elif interaction_type is InteractionType.MEAN: + if hasattr(dist, "mean"): + try: + return dist.mean + except NotImplementedError: + pass + if dist.has_rsample: + return dist.rsample((self.n_empirical_estimate,)).mean(0) + else: + return dist.sample((self.n_empirical_estimate,)).mean(0) + + elif interaction_type is InteractionType.RANDOM: + num_samples = self.num_samples + if num_samples is None: + num_samples = torch.Size(()) + if dist.has_rsample: + return dist.rsample(num_samples) + else: + return dist.sample(num_samples) + else: + raise NotImplementedError(f"unknown interaction_type {interaction_type}") + + +class ProbabilisticTensorDictSequential(TensorDictSequential): + """A sequence of :class:`~tensordict.nn.TensorDictModules` containing at least one :class:`~tensordict.nn.ProbabilisticTensorDictModule`. + + This class extends :class:`~tensordict.nn.TensorDictSequential` and is typically configured with a sequence of + modules where the final module is an instance of :class:`~tensordict.nn.ProbabilisticTensorDictModule`. + However, it also supports configurations where one or more intermediate modules are instances of + :class:`~tensordict.nn.ProbabilisticTensorDictModule`, while the last module may or may not be probabilistic. + In all cases, it exposes the :meth:`~.get_dist` method to recover the distribution object from the + :class:`~tensordict.nn.ProbabilisticTensorDictModule` instances in the sequence. + + Multiple probabilistic modules can co-exist in a single ``ProbabilisticTensorDictSequential``. + If `return_composite` is ``False`` (default), only the last one will produce a distribution and the others + will be executed as regular :class:`~tensordict.nn.TensorDictModule` instances. + However, if a `ProbabilisticTensorDictModule` is not the last module in the sequence and `return_composite=False`, + a `ValueError` will be raised when trying to query the module. If `return_composite=True`, + all intermediate `ProbabilisticTensorDictModule` instances will contribute to a single + :class:`~tensordict.nn.CompositeDistribution` instance. + + Resulting log-probabilities will be conditional probabilities if samples are interdependent: + whenever + + .. math:: + Z = F(X, Y) + + then the log-probability of Z will be + + .. math:: + log(p(z | x, y)) + + Args: + *modules (sequence or OrderedDict of TensorDictModuleBase or ProbabilisticTensorDictModule): An ordered sequence of + :class:`~tensordict.nn.TensorDictModule` instances, usually terminating in a :class:`~tensordict.nn.ProbabilisticTensorDictModule`, + to be run sequentially. + The modules can be instances of TensorDictModuleBase or any other function that matches this signature. + Note that if a non-TensorDictModuleBase callable is used, its input and output keys will not be tracked, + and thus will not affect the `in_keys` and `out_keys` attributes of the TensorDictSequential. + + Keyword Args: + partial_tolerant (bool, optional): If ``True``, the input tensordict can miss some + of the input keys. If so, only the modules that can be executed given the + keys that are present will be executed. Also, if the input tensordict is a + lazy stack of tensordicts AND if partial_tolerant is ``True`` AND if the stack + does not have the required keys, then TensorDictSequential will scan through + the sub-tensordicts looking for those that have the required keys, if any. + Defaults to ``False``. + return_composite (bool, optional): If `True` and multiple + :class:`~tensordict.nn.ProbabilisticTensorDictModule` or + :class:`~tensordict.nn.ProbabilisticTensorDictSequential` instances are found, + a :class:`~tensordict.nn.CompositeDistribution` instance will be used. + Otherwise, only the last module will be used to build the distribution. + Defaults to ``True`` whenever there are more than one probabilistic modules or the last module is not probabilistic. + Errors if `return_composite` is `False` and the neither of the above conditions are met. + selected_out_keys (iterable of NestedKeys, optional): the list of out-keys to select. If not provided, all + ``out_keys`` will be written. + inplace (bool, optional): if `True`, the input tensordict is modified in-place. If `False`, a new empty + :class:`~tensordict.TensorDict` instance is created. If `"empty"`, `input.empty()` is used instead (ie, the + output preserves type, device and batch-size). Defaults to `None` (relies on sub-modules). + + Raises: + ValueError: If the input sequence of modules is empty. + TypeError: If the final module is not an instance of + :obj:`ProbabilisticTensorDictModule` or + :obj:`ProbabilisticTensorDictSequential`. + + Examples: + >>> from tensordict.nn import ProbabilisticTensorDictModule as Prob, ProbabilisticTensorDictSequential as Seq + >>> import torch + >>> # Typical usage: a single distribution is computed last in the sequence + >>> import torch + >>> from tensordict import TensorDict + >>> from tensordict.nn import ProbabilisticTensorDictModule as Prob, ProbabilisticTensorDictSequential as Seq, \ + ... TensorDictModule as Mod + >>> torch.manual_seed(0) + >>> + >>> module = Seq( + ... Mod(lambda x: x + 1, in_keys=["x"], out_keys=["loc"]), + ... Prob(in_keys=["loc"], out_keys=["sample"], distribution_class=torch.distributions.Normal, + ... distribution_kwargs={"scale": 1}), + ... ) + >>> input = TensorDict(x=torch.ones(3)) + >>> td = module(input.copy()) + >>> print(td) + TensorDict( + fields={ + loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + sample: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + x: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> print(module.get_dist(input)) + Normal(loc: torch.Size([3]), scale: torch.Size([3])) + >>> print(module.log_prob(td)) + tensor([-0.9189, -0.9189, -0.9189]) + >>> # Intermediate distributions are ignored when return_composite=False + >>> module = Seq( + ... Mod(lambda x: x + 1, in_keys=["x"], out_keys=["loc"]), + ... Prob(in_keys=["loc"], out_keys=["sample0"], distribution_class=torch.distributions.Normal, + ... distribution_kwargs={"scale": 1}), + ... Mod(lambda x: x + 1, in_keys=["sample0"], out_keys=["loc2"]), + ... Prob(in_keys={"loc": "loc2"}, out_keys=["sample1"], distribution_class=torch.distributions.Normal, + ... distribution_kwargs={"scale": 1}), + ... return_composite=False, + ... ) + >>> td = module(TensorDict(x=torch.ones(3))) + >>> print(td) + TensorDict( + fields={ + loc2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + sample0: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + sample1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + x: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> print(module.get_dist(input)) + Normal(loc: torch.Size([3]), scale: torch.Size([3])) + >>> print(module.log_prob(td)) + tensor([-0.9189, -0.9189, -0.9189]) + >>> # Intermediate distributions produce a CompositeDistribution when return_composite=True + >>> module = Seq( + ... Mod(lambda x: x + 1, in_keys=["x"], out_keys=["loc"]), + ... Prob(in_keys=["loc"], out_keys=["sample0"], distribution_class=torch.distributions.Normal, + ... distribution_kwargs={"scale": 1}), + ... Mod(lambda x: x + 1, in_keys=["sample0"], out_keys=["loc2"]), + ... Prob(in_keys={"loc": "loc2"}, out_keys=["sample1"], distribution_class=torch.distributions.Normal, + ... distribution_kwargs={"scale": 1}), + ... return_composite=True, + ... ) + >>> input = TensorDict(x=torch.ones(3)) + >>> td = module(input.copy()) + >>> print(td) + TensorDict( + fields={ + loc2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + sample0: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + sample1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + x: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> print(module.get_dist(input)) + CompositeDistribution({'sample0': Normal(loc: torch.Size([3]), scale: torch.Size([3])), 'sample1': Normal(loc: torch.Size([3]), scale: torch.Size([3]))}) + >>> print(module.log_prob(td)) + TensorDict( + fields={ + sample0_log_prob: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + sample1_log_prob: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> # Even a single intermediate distribution is wrapped in a CompositeDistribution when + >>> # return_composite=True + >>> module = Seq( + ... Mod(lambda x: x + 1, in_keys=["x"], out_keys=["loc"]), + ... Prob(in_keys=["loc"], out_keys=["sample0"], distribution_class=torch.distributions.Normal, + ... distribution_kwargs={"scale": 1}), + ... Mod(lambda x: x + 1, in_keys=["sample0"], out_keys=["y"]), + ... return_composite=True, + ... ) + >>> td = module(TensorDict(x=torch.ones(3))) + >>> print(td) + TensorDict( + fields={ + loc: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + sample0: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + x: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + y: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> print(module.get_dist(input)) + CompositeDistribution({'sample0': Normal(loc: torch.Size([3]), scale: torch.Size([3]))}) + >>> print(module.log_prob(td)) + TensorDict( + fields={ + sample0_log_prob: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + """ + + @overload + def __init__( + self, + modules: OrderedDict[str, TensorDictModuleBase | ProbabilisticTensorDictModule], + partial_tolerant: bool = False, + return_composite: bool | None = None, + *, + selected_out_keys: List[NestedKey] | None = None, + inplace: bool | None = None, + ) -> None: ... + + @overload + def __init__( + self, + modules: List[TensorDictModuleBase | ProbabilisticTensorDictModule], + partial_tolerant: bool = False, + return_composite: bool | None = None, + *, + selected_out_keys: List[NestedKey] | None = None, + inplace: bool | None = None, + ) -> None: ... + + def __init__( + self, + *modules: TensorDictModuleBase | ProbabilisticTensorDictModule, + partial_tolerant: bool = False, + return_composite: bool | None = None, + selected_out_keys: List[NestedKey] | None = None, + inplace: bool | None = None, + ) -> None: + if len(modules) == 0: + raise ValueError( + "ProbabilisticTensorDictSequential must consist of zero or more " + "TensorDictModules followed by a ProbabilisticTensorDictModule" + ) + self._ordered_dict = False + if len(modules) == 1 and isinstance( + modules[0], (collections.OrderedDict, MutableSequence) + ): + if isinstance(modules[0], collections.OrderedDict): + modules_list = list(modules[0].values()) + self._ordered_dict = True + else: + modules = modules_list = list(modules[0]) + elif len(modules) == 1 and isinstance(modules[0], dict): + modules = [collections.OrderedDict(modules[0])] + return self.__init__( + *modules, + partial_tolerant=partial_tolerant, + return_composite=return_composite, + selected_out_keys=selected_out_keys, + inplace=inplace, + ) + else: + modules_list = list(modules) + modules_list = self._convert_modules(modules_list) + + if return_composite is False and not isinstance( + modules_list[-1], + (ProbabilisticTensorDictModule, ProbabilisticTensorDictSequential), + ): + raise TypeError( + "The final module passed to ProbabilisticTensorDictSequential must be " + "an instance of ProbabilisticTensorDictModule or another " + "ProbabilisticTensorDictSequential (unless return_composite is set to ``True``)." + ) + if return_composite is None: + return_composite = not isinstance( + modules_list[-1], + (ProbabilisticTensorDictModule, ProbabilisticTensorDictSequential), + ) or ( + sum(isinstance(m, ProbabilisticTensorDictModule) for m in modules_list) + > 1 + ) + + if not any( + isinstance( + m, (ProbabilisticTensorDictModule, ProbabilisticTensorDictSequential) + ) + for m in modules_list + ): + # No probabilistic modules - initialize as a regular TensorDictSequential + kwargs = { + "partial_tolerant": partial_tolerant, + "inplace": inplace, + } + if selected_out_keys is not None: + kwargs["selected_out_keys"] = selected_out_keys + TensorDictSequential.__init__(self, *modules, **kwargs) + self._requires_sample = False + self.__dict__["_det_part"] = None + # Use return_composite=True so forward() just iterates through modules + self.return_composite = True + return + + # if the modules not including the final probabilistic module return the sampled + # key we won't be sampling it again, in that case + # ProbabilisticTensorDictSequential is presumably used to return the + # distribution using `get_dist` or to sample log_probabilities + _, out_keys = self._compute_in_and_out_keys(modules_list[:-1]) + self._requires_sample = any( + key not in set(out_keys) + for key in getattr(modules_list[-1], "dist_sample_keys", [None]) + ) + if self._ordered_dict: + self.__dict__["_det_part"] = TensorDictSequential( + collections.OrderedDict(list(modules[0].items())[:-1]) + ) + else: + self.__dict__["_det_part"] = TensorDictSequential(*modules[:-1]) + + kwargs = { + "partial_tolerant": partial_tolerant, + "inplace": inplace, + } + if selected_out_keys is not None: + kwargs["selected_out_keys"] = selected_out_keys + super().__init__(*modules, **kwargs) + self.return_composite = return_composite + + def __getitem__(self, index: int | slice | str) -> Self | TensorDictModuleBase: + if isinstance(index, (int, str)): + return self.module.__getitem__(index) + else: + mods = self.module.__getitem__(index) + if self.return_composite and any( + isinstance( + item, + (ProbabilisticTensorDictModule, ProbabilisticTensorDictSequential), + ) + for item in mods + ): + return type(self)(*mods, return_composite=self.return_composite) + elif isinstance( + mods[-1], + (ProbabilisticTensorDictModule, ProbabilisticTensorDictSequential), + ): + return type(self)(*mods) + else: + return TensorDictSequential(*mods) + + _dist_sample = ProbabilisticTensorDictModule._dist_sample + + @property + def det_part(self): + return self._det_part + + def get_dist_params( + self, + tensordict: TensorDictBase, + tensordict_out: TensorDictBase | None = None, + **kwargs, + ) -> tuple[D.Distribution, TensorDictBase]: + """Returns the distribution parameters and output tensordict. + + This method runs the deterministic part of the :class:`~tensordict.nn.ProbabilisticTensorDictSequential` + module to obtain the distribution parameters. The interaction type is set to the current global + interaction type if available, otherwise it defaults to the interaction type of the last module. + + Args: + tensordict (TensorDictBase): The input tensordict. + tensordict_out (TensorDictBase, optional): The output tensordict. If ``None``, a new tensordict will be created. + Defaults to ``None``. + + Keyword Args: + **kwargs: Additional keyword arguments passed to the deterministic part of the module. + + Returns: + tuple[D.Distribution, TensorDictBase]: A tuple containing the distribution object and the output tensordict. + + .. note:: The interaction type is temporarily set to the specified value during the execution of this method. + """ + tds = self.det_part + type = interaction_type() + if type is None: + for m in reversed(list(self._module_iter())): + if hasattr(m, "default_interaction_type"): + type = m.default_interaction_type + break + else: + raise ValueError("Could not find a default interaction in the modules.") + with set_interaction_type(type): + return tds(tensordict, tensordict_out, **kwargs) + + @property + def log_prob_keys(self): + lpks = [] + for m in reversed(list(self._module_iter())): + lpks.extend(getattr(m, "log_prob_keys", [])) + return lpks + + log_prob_key = ProbabilisticTensorDictModule.log_prob_key + + @property + def dist_params_keys(self) -> List[NestedKey]: + """Returns all the keys pointing at the distribution params.""" + result = [] + for m in reversed(list(self._module_iter())): + result.extend(getattr(m, "dist_params_keys", [])) + return result + + @property + def dist_sample_keys(self) -> List[NestedKey]: + """Returns all the keys pointing at the distribution samples.""" + result = [] + for m in reversed(list(self._module_iter())): + result.extend(getattr(m, "dist_sample_keys", [])) + return result + + @property + def num_samples(self): + num_samples = () + for tdm in self._module_iter(): + if isinstance( + tdm, (ProbabilisticTensorDictModule, ProbabilisticTensorDictSequential) + ): + local_num_samples = tdm.num_samples + if local_num_samples is None: + local_num_samples = () + num_samples = local_num_samples + num_samples + return num_samples + + def get_dist( + self, + tensordict: TensorDictBase, + tensordict_out: TensorDictBase | None = None, + **kwargs, + ) -> D.Distribution: + """Returns the distribution resulting from passing the input tensordict through the sequence. + + If `return_composite` is ``False`` (default), this method will only consider the last probabilistic module in the sequence. + + Otherwise, it will return a :class:`~tensordict.nn.CompositeDistribution` instance containing the distributions of all probabilistic modules. + + Args: + tensordict (TensorDictBase): The input tensordict. + tensordict_out (TensorDictBase, optional): The output tensordict. If ``None``, a new tensordict will be created. + Defaults to ``None``. + + Keyword Args: + **kwargs: Additional keyword arguments passed to the underlying modules. + + Returns: + D.Distribution: The resulting distribution object. + + Raises: + RuntimeError: If no probabilistic module is found in the sequence. + + .. note:: + When `return_composite` is ``True``, the distributions are conditioned on the previous samples in the sequence. + This means that if a module depends on the output of a previous probabilistic module, its distribution will be conditional. + + """ + if not self.return_composite: + tensordict_out = self.get_dist_params(tensordict, tensordict_out, **kwargs) + return self.build_dist_from_params(tensordict_out) + + td_copy = tensordict.copy() + dists = {} + for i, tdm in enumerate(self._module_iter()): + if isinstance( + tdm, (ProbabilisticTensorDictModule, ProbabilisticTensorDictSequential) + ): + dist = tdm.get_dist(td_copy) + if i < len(self.module) - 1: + sample = tdm._dist_sample(dist, interaction_type=interaction_type()) + if tdm.num_samples not in ((), None): + td_copy = td_copy.expand(tdm.num_samples + td_copy.shape) + if isinstance(tdm, ProbabilisticTensorDictModule): + if isinstance(sample, torch.Tensor): + sample = [sample] + td_copy.update(dict(_zip_strict(tdm.dist_sample_keys, sample))) + else: + td_copy.update(sample) + if isinstance(dist, CompositeDistribution): + dists.update(dict(dist)) + else: + dists[tdm.out_keys[0]] = dist + else: + td_copy = tdm(td_copy) + if len(dists) == 0: + raise RuntimeError(f"No distribution module found in {self}.") + # elif len(dists) == 1: + # return dist + return CompositeDistribution.from_distributions( + td_copy, + dists, + ) + + @property + def default_interaction_type(self): + """Returns the `default_interaction_type` of the module using an iterative heuristic. + + This property iterates over all modules in reverse order, attempting to retrieve the + `default_interaction_type` attribute from any child module. The first non-None value + encountered is returned. If no such value is found, a default `interaction_type()` is returned. + + """ + for m in reversed(list(self._module_iter())): + interaction = getattr(m, "default_interaction_type", None) + if interaction is not None: + return interaction + return interaction_type() + + @property + def _last_module(self): + if not self._ordered_dict: + return self.module[-1] + mod = None + for mod in self._module_iter(): # noqa: B007 + continue + return mod + + def log_prob( + self, + tensordict, + tensordict_out: TensorDictBase | None = None, + *, + dist: torch.distributions.Distribution | None = None, + **kwargs, + ) -> TensorDictBase | torch.Tensor: + """Returns the log-probability of the input tensordict. + + If `self.return_composite` is ``True`` and the distribution is a :class:`~tensordict.nn.CompositeDistribution`, + this method will return the log-probability of the entire composite distribution. + + Otherwise, it will only consider the last probabilistic module in the sequence. + + Args: + tensordict (TensorDictBase): The input tensordict. + tensordict_out (TensorDictBase, optional): The output tensordict. If ``None``, a new tensordict will be created. + Defaults to ``None``. + + Keyword Args: + dist (torch.distributions.Distribution, optional): The distribution object. If ``None``, it will be computed using `get_dist`. + Defaults to ``None``. + + Returns: + TensorDictBase or torch.Tensor: The log-probability of the input tensordict. + + """ + if tensordict_out is not None: + tensordict_inp = tensordict.copy() + else: + tensordict_inp = tensordict + if dist is None: + dist = self.get_dist(tensordict_inp) + return_composite = self.return_composite + if return_composite and isinstance(dist, CompositeDistribution): + # Check the values within the dist - if not set, choose defaults + return dist.log_prob( + tensordict, + **kwargs, + ) + last_module: ProbabilisticTensorDictModule = self._last_module + out = last_module.log_prob(tensordict_inp, dist=dist, **kwargs) + if is_tensor_collection(out): + if tensordict_out is not None: + if out is tensordict_inp: + tensordict_out.update( + tensordict_inp.apply( + lambda x, y: x if x is not y else None, + tensordict, + filter_empty=True, + ) + ) + else: + tensordict_out.update(out) + else: + tensordict_out = out + return tensordict_out + return out + + def build_dist_from_params(self, tensordict: TensorDictBase) -> D.Distribution: + """Constructs a distribution from the input parameters without evaluating other modules in the sequence. + + This method searches for the last :class:`~tensordict.nn.ProbabilisticTensorDictModule` in the sequence and uses it to build the distribution. + + Args: + tensordict (TensorDictBase): The input tensordict containing the distribution parameters. + + Returns: + D.Distribution: The constructed distribution object. + + Raises: + RuntimeError: If no :class:`~tensordict.nn.ProbabilisticTensorDictModule` is found in the sequence. + """ + dest_module = None + for module in reversed(list(self.modules())): + if isinstance(module, ProbabilisticTensorDictModule): + dest_module = module + break + if dest_module is None: + raise RuntimeError( + "Could not find any ProbabilisticTensorDictModule in the sequence." + ) + return dest_module.get_dist(tensordict) + + @dispatch(auto_batch_size=False) + @_set_skip_existing_None() + def forward( + self, + tensordict: TensorDictBase, + tensordict_out: TensorDictBase | None = None, + **kwargs, + ) -> TensorDictBase: + if (tensordict_out is None and self._select_before_return) or ( + tensordict_out is not None + ): + tensordict_exec = tensordict.copy() + else: + tensordict_exec = tensordict + if self.return_composite: + for m in self._module_iter(): + try: + if isinstance( + m, + (ProbabilisticTensorDictModule, ProbabilisticTensorDictModule), + ): + tensordict_exec = m( + tensordict_exec, _requires_sample=self._requires_sample + ) + else: + tensordict_exec = m(tensordict_exec, **kwargs) + except Exception as e: + module_num_or_key = self._get_module_num_or_key(m) + raise RuntimeError( + f"Failed while executing module '{module_num_or_key}'. Scroll up for more info." + ) from e + else: + tensordict_exec = self.get_dist_params(tensordict_exec, **kwargs) + tensordict_exec = self._last_module( + tensordict_exec, _requires_sample=self._requires_sample + ) + + if self.inplace is True: + tensordict_out = tensordict + elif self.inplace is False: + tensordict_out = TensorDict() + elif self.inplace == "empty": + tensordict_out = tensordict.empty() + + if tensordict_out is not None: + result = tensordict_out + result.update(tensordict_exec, keys_to_update=self.out_keys) + else: + result = tensordict_exec + if self._select_before_return: + # We must also update any value that has been updated during the course of execution + # from the input data. + if is_compiling(): + keys = [ # noqa: C416 + k + for k in {k for k in self.out_keys}.union( # noqa: C416 + {k for k in tensordict.keys(True, True)} # noqa: C416 + ) + ] + else: + keys = list(set(self.out_keys + list(tensordict.keys(True, True)))) + return tensordict.update(result, keys_to_update=keys) + return result + + +def _dynamo_friendly_to_dict(data): + if not is_compiling(): + return data + if isinstance(data, TensorDictBase): + # to_dict is recursive and we don't want that + items = dict(data.items()) + for k, v in items.items(): + if is_non_tensor(v): + items[k] = v.data + return items + return data diff --git a/lib/python3.12/site-packages/tensordict/nn/sequence.py b/lib/python3.12/site-packages/tensordict/nn/sequence.py new file mode 100644 index 0000000000000000000000000000000000000000..aae7ba9e05bb04e792d547ae44988bf4ac4610ac --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/sequence.py @@ -0,0 +1,691 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import collections +import logging +import sys +from copy import deepcopy +from typing import Any, Callable, Iterable, List, OrderedDict, overload, TYPE_CHECKING + +from tensordict._nestedkey import NestedKey +from tensordict._td import TensorDict + +from tensordict.nn.common import ( + dispatch, + TensorDictModule, + TensorDictModuleBase, + WrapModule, +) +from tensordict.nn.utils import _set_skip_existing_None +from tensordict.tensordict import LazyStackedTensorDict, TensorDictBase +from tensordict.utils import _zip_strict, unravel_key_list +from torch import nn + +_has_functorch = False +try: + import functorch + + _has_functorch = True +except ImportError: + logging.info( + "failed to import functorch. TensorDict's features that do not require " + "functional programming should work, but functionality and performance " + "may be affected. Consider installing functorch and/or upgrating pytorch." + ) + FUNCTORCH_ERROR = "functorch not installed. Consider installing functorch to use this functionality." + +try: + from torch.compiler import is_compiling +except ImportError: + from torch._dynamo import is_compiling + +_has_py311_or_greater = sys.version_info >= (3, 11) + +if TYPE_CHECKING: + from typing import Self +else: + Self = Any + + +__all__ = ["TensorDictSequential"] + + +class TensorDictSequential(TensorDictModule): + """A sequence of TensorDictModules. + + Similarly to :obj:`nn.Sequence` which passes a tensor through a chain of mappings that read and write a single tensor + each, this module will read and write over a tensordict by querying each of the input modules. + When calling a :obj:`TensorDictSequencial` instance with a functional module, it is expected that the parameter lists (and + buffers) will be concatenated in a single list. + + Args: + modules (OrderedDict[str, Callable[[TensorDictBase], TensorDictBase]] | List[Callable[[TensorDictBase], TensorDictBase]]): + ordered sequence of callables that take a TensorDictBase as input and return a TensorDictBase. + These can be instances of TensorDictModuleBase or any other function that matches this signature. + Note that if a non-TensorDictModuleBase callable is used, its input and output keys will not be tracked, + and thus will not affect the `in_keys` and `out_keys` attributes of the TensorDictSequential. + Regular ``dict`` inputs will be converted to ``OrderedDict`` if necessary. + + Keyword Args: + partial_tolerant (bool, optional): if True, the input tensordict can miss some of the input keys. + If so, the only module that will be executed are those who can be executed given the keys that + are present. + Also, if the input tensordict is a lazy stack of tensordicts AND if partial_tolerant is :obj:`True` AND if the + stack does not have the required keys, then TensorDictSequential will scan through the sub-tensordicts + looking for those that have the required keys, if any. Defaults to False. + selected_out_keys (iterable of NestedKeys, optional): the list of out-keys to select. If not provided, all + ``out_keys`` will be written. + inplace (bool or str, optional): if `True`, the input tensordict is modified in-place. If `False`, a new empty + :class:`~tensordict.TensorDict` instance is created. If `"empty"`, `input.empty()` is used instead (ie, the + output preserves type, device and batch-size). Defaults to `None` (relies on sub-modules). + + .. note:: + A :class:`TensorDictSequential` instance may have a long list of output keys, and one may wish to remove + some of them after execution for clarity or memory purposes. If this is the case, the method :meth:`~.select_out_keys` + can be used after instantiation, or `selected_out_keys` may be passed to the constructor. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> from tensordict.nn import TensorDictModule, TensorDictSequential + >>> torch.manual_seed(0) + >>> module = TensorDictSequential( + ... TensorDictModule(lambda x: x+1, in_keys=["x"], out_keys=["x+1"]), + ... TensorDictModule(nn.Linear(3, 4), in_keys=["x+1"], out_keys=["w*(x+1)+b"]), + ... ) + >>> # with tensordict input + >>> print(module(TensorDict({"x": torch.zeros(3)}, []))) + TensorDict( + fields={ + w*(x+1)+b: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False), + x+1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False), + x: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> # with tensor input: returns all the output keys in the order of the modules, ie "x+1" and "w*(x+1)+b" + >>> module(x=torch.zeros(3)) + (tensor([1., 1., 1.]), tensor([-0.7214, -0.8748, 0.1571, -0.1138], grad_fn=)) + >>> module(torch.zeros(3)) + (tensor([1., 1., 1.]), tensor([-0.7214, -0.8748, 0.1571, -0.1138], grad_fn=)) + + TensorDictSequence supports functional, modular and vmap coding. + + Examples: + >>> import torch + >>> from tensordict import TensorDict + >>> from tensordict.nn import ( + ... ProbabilisticTensorDictModule, + ... ProbabilisticTensorDictSequential, + ... TensorDictModule, + ... TensorDictSequential, + ... ) + >>> from tensordict.nn.distributions import NormalParamExtractor + >>> from tensordict.nn.functional_modules import make_functional + >>> from torch.distributions import Normal + >>> td = TensorDict({"input": torch.randn(3, 4)}, [3,]) + >>> net1 = torch.nn.Linear(4, 8) + >>> module1 = TensorDictModule(net1, in_keys=["input"], out_keys=["params"]) + >>> normal_params = TensorDictModule( + ... NormalParamExtractor(), in_keys=["params"], out_keys=["loc", "scale"] + ... ) + >>> td_module1 = ProbabilisticTensorDictSequential( + ... module1, + ... normal_params, + ... ProbabilisticTensorDictModule( + ... in_keys=["loc", "scale"], + ... out_keys=["hidden"], + ... distribution_class=Normal, + ... return_log_prob=True, + ... ) + ... ) + >>> module2 = torch.nn.Linear(4, 8) + >>> td_module2 = TensorDictModule( + ... module=module2, in_keys=["hidden"], out_keys=["output"] + ... ) + >>> td_module = TensorDictSequential(td_module1, td_module2) + >>> params = TensorDict.from_module(td_module) + >>> with params.to_module(td_module): + ... _ = td_module(td) + >>> print(td) + TensorDict( + fields={ + hidden: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + input: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + loc: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + output: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + params: Tensor(shape=torch.Size([3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + sample_log_prob: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + scale: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([3]), + device=None, + is_shared=False) + + In the vmap case: + >>> from torch import vmap + >>> params = params.expand(4) + >>> def func(td, params): + ... with params.to_module(td_module): + ... return td_module(td) + >>> td_vmap = vmap(func, (None, 0))(td, params) + >>> print(td_vmap) + TensorDict( + fields={ + hidden: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + input: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + loc: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + output: Tensor(shape=torch.Size([4, 3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + params: Tensor(shape=torch.Size([4, 3, 8]), device=cpu, dtype=torch.float32, is_shared=False), + sample_log_prob: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + scale: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 3]), + device=None, + is_shared=False) + + """ + + module: nn.ModuleList + _select_before_return = False + + @overload + def __init__( + self, + modules: OrderedDict[str, Callable[[TensorDictBase], TensorDictBase]], + *, + partial_tolerant: bool = False, + selected_out_keys: List[NestedKey] | None = None, + inplace: bool | str | None = None, + ) -> None: ... + + @overload + def __init__( + self, + modules: List[Callable[[TensorDictBase], TensorDictBase]], + *, + partial_tolerant: bool = False, + selected_out_keys: List[NestedKey] | None = None, + inplace: bool | str | None = None, + ) -> None: ... + + @overload + def __init__( + self, + *modules: List[Callable[[TensorDictBase], TensorDictBase]], + partial_tolerant: bool = False, + selected_out_keys: List[NestedKey] | None = None, + inplace: bool | str | None = None, + ) -> None: ... + + def __init__( + self, + *modules: Callable[[TensorDictBase], TensorDictBase], + partial_tolerant: bool = False, + selected_out_keys: List[NestedKey] | None = None, + inplace: bool | str | None = None, + ) -> None: + + if len(modules) == 1 and isinstance(modules[0], collections.OrderedDict): + modules_vals = self._convert_modules(modules[0].values()) # type: ignore[unreachable] + in_keys, out_keys = self._compute_in_and_out_keys(modules_vals) + self._complete_out_keys = list(out_keys) + modules = collections.OrderedDict( + **{key: val for key, val in _zip_strict(modules[0], modules_vals)} + ) + super().__init__( + module=nn.ModuleDict(modules), + in_keys=in_keys, + out_keys=out_keys, + ) + elif len(modules) == 1 and isinstance( + modules[0], collections.abc.MutableSequence + ): + modules = self._convert_modules(modules[0]) # type: ignore[unreachable] + in_keys, out_keys = self._compute_in_and_out_keys(modules) + self._complete_out_keys = list(out_keys) + super().__init__( + module=nn.ModuleList(modules), + in_keys=in_keys, + out_keys=out_keys, + ) + elif len(modules) == 1 and isinstance(modules[0], dict): + return self.__init__( # type: ignore[unreachable] + collections.OrderedDict(modules[0]), + partial_tolerant=partial_tolerant, + selected_out_keys=selected_out_keys, + inplace=inplace, + ) + else: + modules = self._convert_modules(modules) + in_keys, out_keys = self._compute_in_and_out_keys(modules) + self._complete_out_keys = list(out_keys) + super().__init__( + module=nn.ModuleList(list(modules)), + in_keys=in_keys, + out_keys=out_keys, + ) + + self.inplace = inplace + self.partial_tolerant = partial_tolerant + if selected_out_keys: + self._select_before_return = True + selected_out_keys = unravel_key_list(selected_out_keys) + if not all(key in self.out_keys for key in selected_out_keys): + raise ValueError("All keys in selected_out_keys must be in out_keys.") + self.out_keys = selected_out_keys + else: + self._select_before_return = False + + def reset_out_keys(self): + self.out_keys = list(self._complete_out_keys) + return self + + @staticmethod + def _convert_modules(modules): + return [ + ( + WrapModule(module) + if not TensorDictModuleBase.is_tdmodule_compatible(module) + else module + ) + for module in modules + ] + + def _compute_in_and_out_keys( + self, modules: list[TensorDictModule] + ) -> tuple[list[NestedKey], list[NestedKey]]: + in_keys = [] # type: ignore[var-annotated] + out_keys = [] # type: ignore[var-annotated] + for module in modules: + # we sometimes use in_keys to select keys of a tensordict that are + # necessary to run a TensorDictModule. If a key is an intermediary in + # the chain, there is no reason why it should belong to the input + # TensorDict. + for in_key in module.in_keys: + if in_key not in (out_keys + in_keys): + in_keys.append(in_key) + out_keys += module.out_keys + + out_keys = [ + out_key + for i, out_key in enumerate(out_keys) + if out_key not in out_keys[i + 1 :] + ] + return in_keys, out_keys + + @staticmethod + def _find_functional_module(module: TensorDictModuleBase) -> nn.Module: + if not _has_functorch: + raise ImportError(FUNCTORCH_ERROR) + fmodule = module + while not isinstance( + fmodule, (functorch.FunctionalModule, functorch.FunctionalModuleWithBuffers) + ): + try: + fmodule = fmodule.module + except AttributeError: + raise AttributeError( + f"couldn't find a functional module in module of type {type(module)}" + ) + return fmodule + + def select_out_keys(self, *selected_out_keys) -> TensorDictSequential: + self._select_before_return = True + selected_out_keys = unravel_key_list(selected_out_keys) + if not all(key in self.out_keys for key in selected_out_keys): + raise ValueError("All keys in selected_out_keys must be in out_keys.") + self.out_keys = selected_out_keys + return self + + def select_subsequence( + self, + in_keys: Iterable[NestedKey] | None = None, + out_keys: Iterable[NestedKey] | None = None, + ) -> TensorDictSequential: + """Returns a new TensorDictSequential with only the modules that are necessary to compute the given output keys with the given input keys. + + Args: + in_keys: input keys of the subsequence we want to select. + All the keys absent from ``in_keys`` will be considered as + non-relevant, and modules that *just* take these keys as inputs + will be discarded. + The resulting sequential module will follow the pattern "all + the modules which output will be affected by a different value + for any in ". + If none is provided, the module's ``in_keys`` are assumed. + out_keys: output keys of the subsequence we want to select. + Only the modules that are necessary to get the ``out_keys`` + will be found in the resulting sequence. + The resulting sequential module will follow the pattern "all + the modules that condition the value or entries." + If none is provided, the module's ``out_keys`` are assumed. + + Returns: + A new TensorDictSequential with only the modules that are necessary acording to the given input and output keys. + + Examples: + >>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod + >>> idn = lambda x: x + >>> module = Seq( + ... Mod(idn, in_keys=["a"], out_keys=["b"]), + ... Mod(idn, in_keys=["b"], out_keys=["c"]), + ... Mod(idn, in_keys=["c"], out_keys=["d"]), + ... Mod(idn, in_keys=["a"], out_keys=["e"]), + ... ) + >>> # select all modules whose output depend on "a" + >>> module.select_subsequence(in_keys=["a"]) + TensorDictSequential( + module=ModuleList( + (0): TensorDictModule( + module= at 0x126ed1ca0>, + device=cpu, + in_keys=['a'], + out_keys=['b']) + (1): TensorDictModule( + module= at 0x126ed1ca0>, + device=cpu, + in_keys=['b'], + out_keys=['c']) + (2): TensorDictModule( + module= at 0x126ed1ca0>, + device=cpu, + in_keys=['c'], + out_keys=['d']) + (3): TensorDictModule( + module= at 0x126ed1ca0>, + device=cpu, + in_keys=['a'], + out_keys=['e']) + ), + device=cpu, + in_keys=['a'], + out_keys=['b', 'c', 'd', 'e']) + >>> # select all modules whose output depend on "c" + >>> module.select_subsequence(in_keys=["c"]) + TensorDictSequential( + module=ModuleList( + (0): TensorDictModule( + module= at 0x126ed1ca0>, + device=cpu, + in_keys=['c'], + out_keys=['d']) + ), + device=cpu, + in_keys=['c'], + out_keys=['d']) + >>> # select all modules that affect the value of "c" + >>> module.select_subsequence(out_keys=["c"]) + TensorDictSequential( + module=ModuleList( + (0): TensorDictModule( + module= at 0x126ed1ca0>, + device=cpu, + in_keys=['a'], + out_keys=['b']) + (1): TensorDictModule( + module= at 0x126ed1ca0>, + device=cpu, + in_keys=['b'], + out_keys=['c']) + ), + device=cpu, + in_keys=['a'], + out_keys=['b', 'c']) + >>> # select all modules that affect the value of "e" + >>> module.select_subsequence(out_keys=["e"]) + TensorDictSequential( + module=ModuleList( + (0): TensorDictModule( + module= at 0x126ed1ca0>, + device=cpu, + in_keys=['a'], + out_keys=['e']) + ), + device=cpu, + in_keys=['a'], + out_keys=['e']) + + This method propagates to nested sequential: + + >>> module = Seq( + ... Seq( + ... Mod(idn, in_keys=["a"], out_keys=["b"]), + ... Mod(idn, in_keys=["b"], out_keys=["c"]), + ... ), + ... Seq( + ... Mod(idn, in_keys=["b"], out_keys=["d"]), + ... Mod(idn, in_keys=["d"], out_keys=["e"]), + ... ), + ... ) + >>> # select submodules whose output will be affected by a change in "b" or "d" AND which output is "e" + >>> module.select_subsequence(in_keys=["b", "d"], out_keys=["e"]) + TensorDictSequential( + module=ModuleList( + (0): TensorDictSequential( + module=ModuleList( + (0): TensorDictModule( + module= at 0x129efae50>, + device=cpu, + in_keys=['b'], + out_keys=['d']) + (1): TensorDictModule( + module= at 0x129efae50>, + device=cpu, + in_keys=['d'], + out_keys=['e']) + ), + device=cpu, + in_keys=['b'], + out_keys=['d', 'e']) + ), + device=cpu, + in_keys=['b'], + out_keys=['d', 'e']) + + The `inplace` argument allows for a fine-grained control over the output type, allowing for instance to write + the result of the computational graph in the input object without tracking the intermediate tensors. + + Example: + >>> import torch + >>> from tensordict import TensorClass + >>> from tensordict.nn import TensorDictModule as Mod, TensorDictSequential as Seq + >>> + >>> class MyClass(TensorClass): + ... input: torch.Tensor + ... output: torch.Tensor | None = None + >>> + >>> obj = MyClass(torch.randn(2, 3), batch_size=(2,)) + >>> + >>> model = Seq( + ... Mod( + ... lambda x: (x + 1, x - 1), + ... in_keys=["input"], + ... out_keys=[("intermediate", "0"), ("intermediate", "1")], + ... inplace=False + ... ), + ... Mod( + ... lambda y0, y1: y0 * y1, + ... in_keys=[("intermediate", "0"), ("intermediate", "1")], + ... out_keys=["output"], + ... inplace=False + ... ), + ... inplace=True, ) + >>> print(model(obj)) + MyClass( + input=Tensor(shape=torch.Size([2, 3]), device=cpu, dtype=torch.float32, is_shared=False), + output=Tensor(shape=torch.Size([2, 3]), device=cpu, dtype=torch.float32, is_shared=False), + output=None, + batch_size=torch.Size([2]), + device=None, + is_shared=False) + + """ + if in_keys is None: + in_keys = deepcopy(self.in_keys) + in_keys = unravel_key_list(in_keys) + if out_keys is None: + out_keys = deepcopy(self.out_keys) + out_keys = unravel_key_list(out_keys) + + module_list = list(self._module_iter()) + id_to_keep = set(range(len(module_list))) + for i, module in enumerate(module_list): + if ( + type(module) is TensorDictSequential + ): # no isinstance because we don't want to mess up subclasses + try: + module = module_list[i] = module.select_subsequence(in_keys=in_keys) + except ValueError: + # then the module can be removed + id_to_keep.remove(i) + continue + + if all(key in in_keys for key in module.in_keys): + in_keys.extend(module.out_keys) + else: + id_to_keep.remove(i) + for i, module in reversed(list(enumerate(module_list))): + if i in id_to_keep: + if any(key in out_keys for key in module.out_keys): + if ( + type(module) is TensorDictSequential + ): # no isinstance because we don't want to mess up subclasses + module = module_list[i] = module.select_subsequence( + out_keys=out_keys + ) + out_keys.extend(module.in_keys) + else: + id_to_keep.remove(i) + id_to_keep = sorted(id_to_keep) + + modules = [module_list[i] for i in id_to_keep] + + if modules == []: + raise ValueError( + "No modules left after selection. Make sure that in_keys and out_keys are coherent." + ) + if isinstance(self.module, nn.ModuleList): + return type(self)(*modules) + else: + keys = [key for key in self.module if self.module[key] in modules] + modules_dict = collections.OrderedDict( + **{key: val for key, val in _zip_strict(keys, modules)} + ) + return type(self)(modules_dict) + + def _run_module( + self, + module: TensorDictModuleBase, + tensordict: TensorDictBase, + **kwargs: Any, + ) -> Any: + if not self.partial_tolerant or all( + key in tensordict.keys(include_nested=True) for key in module.in_keys + ): + tensordict = module(tensordict, **kwargs) + elif self.partial_tolerant and isinstance(tensordict, LazyStackedTensorDict): + for sub_td in tensordict.tensordicts: + if all( + key in sub_td.keys(include_nested=True) for key in module.in_keys + ): + module(sub_td, **kwargs) + return tensordict + + def _module_iter(self): + if isinstance(self.module, nn.ModuleDict): + yield from self.module.children() + else: + yield from self.module + + def _get_module_num_or_key(self, mod: nn.Module) -> int | str: + if isinstance(self.module, nn.ModuleDict): + for name, m in self.module.named_children(): + if m is mod: + return name + else: + raise RuntimeError("module not found.") + else: + for i, m in enumerate(self.module): + if m is mod: + return i + else: + raise RuntimeError("module not found.") + + @dispatch(auto_batch_size=False) + @_set_skip_existing_None() + def forward( + self, + tensordict: TensorDictBase, + tensordict_out: TensorDictBase | None = None, + **kwargs: Any, + ) -> TensorDictBase: + if (tensordict_out is None and self._select_before_return) or ( + tensordict_out is not None + ): + tensordict_exec = tensordict.copy() + else: + tensordict_exec = tensordict + if tensordict_out is None: + if self.inplace is True: + tensordict_out = tensordict + elif self.inplace is False: + tensordict_out = TensorDict() + elif self.inplace == "empty": + tensordict_out = tensordict.empty() + + if not len(kwargs): + for module in self._module_iter(): + try: + tensordict_exec = self._run_module( + module, tensordict_exec, **kwargs + ) + except Exception as e: + if _has_py311_or_greater: + module_num_or_key = self._get_module_num_or_key(module) + e.add_note( + f"Failed while executing module '{module_num_or_key}'." + ) + raise + else: + raise RuntimeError( + f"TensorDictSequential does not support keyword arguments other than 'tensordict_out' or in_keys: {self.in_keys}. Got {kwargs.keys()} instead." + ) + if tensordict_out is not None: + result = tensordict_out + result.update(tensordict_exec, keys_to_update=self.out_keys) + else: + result = tensordict_exec + if self._select_before_return: + # We must also update any value that has been updated during the course of execution + # from the input data. + if is_compiling(): + keys = [ # noqa: C416 + k + for k in {k for k in self.out_keys}.union( # noqa: C416 + {k for k in tensordict.keys(True, True)} # noqa: C416 + ) + ] + else: + keys = list(set(self.out_keys + list(tensordict.keys(True, True)))) + return tensordict.update(result, keys_to_update=keys) + return result + + def __len__(self) -> int: + return len(self.module) + + def __getitem__(self, index: int | slice | str) -> Self | TensorDictModuleBase: + if isinstance(index, (int, str)): + return self.module.__getitem__(index) + else: + return type(self)(*self.module.__getitem__(index)) + + def __setitem__( + self, index: int | slice | str, tensordict_module: TensorDictModuleBase + ) -> None: + return self.module.__setitem__(idx=index, module=tensordict_module) + + def __delitem__(self, index: int | slice | str) -> None: + self.module.__delitem__(idx=index) diff --git a/lib/python3.12/site-packages/tensordict/nn/utils.py b/lib/python3.12/site-packages/tensordict/nn/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e41411dd361eb7e83d81db2716eda49062ce22fe --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/nn/utils.py @@ -0,0 +1,556 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import functools +import inspect +import os +from enum import Enum +from typing import Any, Callable + +import torch +from tensordict.utils import _ContextManager, strtobool +from torch import nn + +from torch.utils._contextlib import _DecoratorContextManager + +try: + from torch.compiler import is_compiling +except ImportError: # torch 2.0 + from torch._dynamo import is_compiling + + +_dispatch_tdnn_modules = _ContextManager( + default=strtobool(os.environ.get("DISPATCH_TDNN_MODULES", "True")) +) + +__all__ = ["mappings", "inv_softplus", "biased_softplus"] + +_skip_existing = _ContextManager(default=False) + + +def inv_softplus(bias: float | torch.Tensor) -> float | torch.Tensor: + """Inverse softplus function. + + Args: + bias (float or tensor): the value to be softplus-inverted. + """ + is_tensor = True + if not isinstance(bias, torch.Tensor): + is_tensor = False + bias = torch.tensor(bias) + out = bias.expm1().clamp_min(1e-6).log() + if not is_tensor and out.numel() == 1: + return out.item() + return out + + +class biased_softplus(nn.Module): + """A biased softplus module. + + The bias indicates the value that is to be returned when a zero-tensor is + passed through the transform. + + Args: + bias (scalar): 'bias' of the softplus transform. If bias=1.0, then a _bias shift will be computed such that + softplus(0.0 + _bias) = bias. + min_val (scalar): minimum value of the transform. + default: 0.1 + """ + + def __init__(self, bias: float, min_val: float = 0.01) -> None: + super().__init__() + self.bias = inv_softplus(bias - min_val) + self.min_val = min_val + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.softplus(x + self.bias) + self.min_val + + +def expln(x): + """A smooth, continuous positive mapping presented in "State-Dependent Exploration for Policy Gradient Methods". + + https://people.idsia.ch/~juergen/ecml2008rueckstiess.pdf + + """ + out = torch.empty_like(x) + idx_neg = x <= 0 + out[idx_neg] = x[idx_neg].exp() + out[~idx_neg] = x[~idx_neg].log1p() + 1 + return out + + +_MAPPINGS: dict[str, Callable[[torch.Tensor], torch.Tensor]] = { + "softplus": torch.nn.functional.softplus, + "exp": torch.exp, + "relu": torch.relu, + "biased_softplus": biased_softplus(1.0), + "none": lambda x: x, + "expln": expln, +} + + +def mappings(key: str) -> Callable: + """Given an input string, returns a surjective function f(x): R -> R^+. + + Args: + key (str): one of `"softplus"`, `"exp"`, `"relu"`, `"expln"`, + `"biased_softplus"` or `"none"` (no mapping). + + .. note:: + If the key begins with `"biased_softplus"`, then it needs to take the following form: + ```"biased_softplus_{bias}"``` where ```bias``` can be converted to a floating point number that will be + used to bias the softplus function. + Alternatively, the ```"biased_softplus_{bias}_{min_val}"``` syntax can be used. + In that case, the additional ```min_val``` term is a floating point + number that will be used to encode the minimum value of the softplus transform. + In practice, the equation used is `softplus(x + bias) + min_val`, where bias and min_val are values computed + such that the conditions above are met. + + .. note:: + Custom mappings can be added through ``tensordict.nn.add_custom_mapping``. + + Returns: + a Callable + + """ + if key in _MAPPINGS: + return _MAPPINGS[key] + elif key.startswith("biased_softplus"): + stripped_key = key.split("_") + if len(stripped_key) == 3: + return biased_softplus(float(stripped_key[-1])) + elif len(stripped_key) == 4: + return biased_softplus( + float(stripped_key[-2]), min_val=float(stripped_key[-1]) + ) + else: + raise ValueError(f"Invalid number of args in {key}") + + else: + raise NotImplementedError(f"Unknown mapping {key}") + + +def add_custom_mapping(name: str, mapping: Callable[[torch.Tensor], torch.Tensor]): + """Adds a custom mapping to be used in mapping classes. + + Args: + name (str): a mapping name. + mapping (callable): a callable that takes a tensor as input and outputs a tensor + with the same shape. + + Examples: + >>> from tensordict.nn import add_custom_mapping, NormalParamExtractor + >>> add_custom_mapping("my_mapping", lambda x: torch.zeros_like(x)) + >>> npe = NormalParamExtractor(scale_mapping="my_mapping", scale_lb=0.0) + >>> assert (npe(torch.randn(10))[1] == torch.zeros(5)).all() + """ + _MAPPINGS[name] = mapping + + +class set_skip_existing(_DecoratorContextManager): + """A context manager for skipping existing nodes in a TensorDict graph. + + When used as a context manager, it will set the `skip_existing()` value + to the ``mode`` indicated, leaving the user able to code up methods that + will check the global value and execute the code accordingly. + + When used as a method decorator, it will check the tensordict input keys + and if the ``skip_existing()`` call returns ``True``, it will skip the method + if all the output keys are already present. + This not not expected to be used as a decorator for methods that do not + respect the following signature: ``def fun(self, tensordict, *args, **kwargs)``. + + Args: + mode (bool, optional): + If ``True``, it indicates that existing entries in the graph + won't be overwritten, unless they are only partially present. :func:`~.skip_existing` + will return ``True``. + If ``False``, no check will be performed. + If ``None``, the value of :func:`~.skip_existing` will not be + changed. This is intended to be used exclusively for decorating + methods and allow their behaviour to depend on the same class + when used as a context manager (see example below). + Defaults to ``True``. + in_key_attr (str, optional): the name of the input key list attribute + in the module's method being decorated. Defaults to ``in_keys``. + out_key_attr (str, optional): the name of the output key list attribute + in the module's method being decorated. Defaults to ``out_keys``. + + Examples: + >>> with set_skip_existing(): + ... if skip_existing(): + ... print("True") + ... else: + ... print("False") + ... + True + >>> print("calling from outside:", skip_existing()) + calling from outside: False + + This class can also be used as a decorator: + + Examples: + >>> from tensordict import TensorDict + >>> from tensordict.nn import set_skip_existing, skip_existing, TensorDictModuleBase + >>> class MyModule(TensorDictModuleBase): + ... in_keys = [] + ... out_keys = ["out"] + ... @set_skip_existing() + ... def forward(self, tensordict): + ... print("hello") + ... tensordict.set("out", torch.zeros(())) + ... return tensordict + >>> module = MyModule() + >>> module(TensorDict({"out": torch.zeros(())}, [])) # does not print anything + TensorDict( + fields={ + out: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + >>> module(TensorDict()) # prints hello + hello + TensorDict( + fields={ + out: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + Decorating a method with the mode set to ``None`` is useful whenever one + wants ot let the context manager take care of skipping things from the outside: + + Examples: + >>> from tensordict import TensorDict + >>> from tensordict.nn import set_skip_existing, skip_existing, TensorDictModuleBase + >>> class MyModule(TensorDictModuleBase): + ... in_keys = [] + ... out_keys = ["out"] + ... @set_skip_existing(None) + ... def forward(self, tensordict): + ... print("hello") + ... tensordict.set("out", torch.zeros(())) + ... return tensordict + >>> module = MyModule() + >>> _ = module(TensorDict({"out": torch.zeros(())}, [])) # prints "hello" + hello + >>> with set_skip_existing(True): + ... _ = module(TensorDict({"out": torch.zeros(())}, [])) # no print + + + .. note:: + To allow for modules to have the same input and output keys and not + mistakenly ignoring subgraphs, ``@set_skip_existing(True)`` will be + deactivated whenever the output keys are also the input keys: + + >>> class MyModule(TensorDictModuleBase): + ... in_keys = ["out"] + ... out_keys = ["out"] + ... @set_skip_existing() + ... def forward(self, tensordict): + ... print("calling the method!") + ... return tensordict + ... + >>> module = MyModule() + >>> module(TensorDict({"out": torch.zeros(())}, [])) # does not print anything + calling the method! + TensorDict( + fields={ + out: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([]), + device=None, + is_shared=False) + + + """ + + def __init__( + self, mode: bool | None = True, in_key_attr="in_keys", out_key_attr="out_keys" + ): + self.mode = mode + self.in_key_attr = in_key_attr + self.out_key_attr = out_key_attr + self._called = False + + def clone(self) -> set_skip_existing: + # override this method if your children class takes __init__ parameters + out = type(self)(self.mode) + out._called = self._called + return out + + def __call__(self, func: Callable): + + self._called = True + + # sanity check + for i, key in enumerate(inspect.signature(func).parameters): + if i == 0: + # skip self + continue + if key != "tensordict": + raise RuntimeError( + "the first argument of the wrapped function must be " + "named 'tensordict'." + ) + break + + @functools.wraps(func) + def wrapper(_self, tensordict, *args: Any, **kwargs: Any) -> Any: + in_keys = getattr(_self, self.in_key_attr) + out_keys = getattr(_self, self.out_key_attr) + # we use skip_existing to allow users to override the mode internally + if ( + skip_existing() + and all(key in tensordict.keys(True) for key in out_keys) + and not any(key in out_keys for key in in_keys) + ): + return tensordict + return func(_self, tensordict, *args, **kwargs) + + return super().__call__(wrapper) + + def __enter__(self) -> None: + if self.mode and is_compiling(): + raise RuntimeError("skip_existing is not compatible with TorchDynamo.") + self.prev = _skip_existing.get_mode() + if self.mode is not None: + _skip_existing.set_mode(self.mode) + elif not self._called: + raise RuntimeError( + f"It seems you are using {type(self).__name__} as a context manager with ``None`` input. " + f"This behaviour is not allowed." + ) + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + _skip_existing.set_mode(self.prev) + + +class _set_skip_existing_None(set_skip_existing): + """A version of skip_existing that is constant wrt init inputs (for torch.compile compatibility). + + This class should only be used as a decorator, not a context manager. + """ + + def __call__(self, func: Callable): + self._called = True + + # sanity check + for i, key in enumerate(inspect.signature(func).parameters): + if i == 0: + # skip self + continue + if key != "tensordict": + raise RuntimeError( + "the first argument of the wrapped function must be " + "named 'tensordict'." + ) + break + + @functools.wraps(func) + def wrapper(_self, tensordict, *args: Any, **kwargs: Any) -> Any: + if skip_existing() and is_compiling(): + raise RuntimeError( + "skip_existing is not compatible with torch.compile." + ) + in_keys = getattr(_self, self.in_key_attr) + out_keys = getattr(_self, self.out_key_attr) + # we use skip_existing to allow users to override the mode internally + if ( + skip_existing() + and all(key in tensordict.keys(True) for key in out_keys) + and not any(key in out_keys for key in in_keys) + ): + return tensordict + if is_compiling(): + return func(_self, tensordict, *args, **kwargs) + self.prev = _skip_existing.get_mode() + try: + result = func(_self, tensordict, *args, **kwargs) + finally: + _skip_existing.set_mode(self.prev) + return result + + return wrapper + + in_key_attr = "in_keys" + out_key_attr = "out_keys" + __init__ = object.__init__ + + def clone(self) -> _set_skip_existing_None: + # override this method if your children class takes __init__ parameters + out = type(self)() + return out + + +def skip_existing(): + """Returns whether or not existing entries in a tensordict should be re-computed by a module.""" + return _skip_existing.get_mode() + + +def _rebuild_buffer(data, requires_grad, backward_hooks): + buffer = Buffer(data, requires_grad) + # NB: This line exists only for backwards compatibility; the + # general expectation is that backward_hooks is an empty + # OrderedDict. See Note [Don't serialize hooks] + buffer._backward_hooks = backward_hooks + + return buffer + + +# For backward compatibility in imports +try: + from torch.nn.parameter import Buffer # noqa +except ImportError: + from tensordict.utils import Buffer # noqa + + +def _dispatch_td_nn_modules(): + """Returns ``True`` if @dispatch should be used. Not using dispatch is faster and also better compatible with torch.compile.""" + return _dispatch_tdnn_modules.get_mode() + + +class _set_dispatch_td_nn_modules(_DecoratorContextManager): + """Controls whether @dispatch should be used. Not using dispatch is faster and also better compatible with torch.compile.""" + + def __init__(self, mode): + self.mode = mode + self._saved_mode = None + + def clone(self): + return type(self)(self.mode) + + def __enter__(self): + # We want to avoid changing global variables because compile puts guards on them + if _dispatch_tdnn_modules.get_mode() != self.mode: + self._saved_mode = _dispatch_tdnn_modules + _dispatch_tdnn_modules.set_mode(self.mode) + + def __exit__(self, exc_type, exc_val, exc_tb): + if self._saved_mode is None: + return + _dispatch_tdnn_modules.set_mode(self._saved_mode) + + +# Reproduce StrEnum for python<3.11 + + +class StrEnum(str, Enum): # noqa + def __new__(cls, *values): + if len(values) > 3: + raise TypeError("too many arguments for str(): %r" % (values,)) + if len(values) == 1: + # it must be a string + if not isinstance(values[0], str): + raise TypeError("%r is not a string" % (values[0],)) + if len(values) >= 2: + # check that encoding argument is a string + if not isinstance(values[1], str): + raise TypeError("encoding must be a string, not %r" % (values[1],)) + if len(values) == 3: + # check that errors argument is a string + if not isinstance(values[2], str): + raise TypeError("errors must be a string, not %r" % (values[2])) + value = str(*values) + member = str.__new__(cls, value) + member._value_ = value + return member + + def _generate_next_value_(name, start, count, last_values): + return name.lower() + + +_composite_lp_aggregate = _ContextManager( + default=( + strtobool(os.getenv("COMPOSITE_LP_AGGREGATE")) + if os.getenv("COMPOSITE_LP_AGGREGATE") is not None + else None + ) +) + + +def composite_lp_aggregate(nowarn: bool = False) -> bool | None: + """Returns whether a :class:`~tensordict.nn.CompositeDistribution` log-probabilities and entropies will be aggregated in a single tensor. + + Args: + nowarn (bool, optional): whether to ignore warnings. Defaults to False. + + .. seealso:: :func:`~tensordict.nn.set_composite_lp_aggregate` + + """ + mode = _composite_lp_aggregate.get_mode() + return bool(mode) + + +class set_composite_lp_aggregate(_DecoratorContextManager): + """Controls whether :class:`~tensordict.nn.CompositeDistribution` log-probabilities and entropies will be aggregated in a single tensor. + + When :func:`~tensordict.nn.composite_lp_aggregate` returns ``True``, the log-probs / entropies of :class:`~tensordict.nn.CompositeDistribution` + will be summed into a single tensor with the shape of the root tensordict. This behaviour is being deprecated in favor of + non-aggregated log-probs, which offer more flexibility and a somewhat more natural API (tensordict samples, tensordict log-probs, tensordict entropies). + + The value of composite_lp_aggregate can also be controlled through the `COMPOSITE_LP_AGGREGATE` environment variable. + + Example: + >>> _ = torch.manual_seed(0) + >>> from tensordict import TensorDict + >>> from tensordict.nn import CompositeDistribution, set_composite_lp_aggregate + >>> import torch + >>> from torch import distributions as d + >>> params = TensorDict({ + ... "cont": {"loc": torch.randn(3, 4), "scale": torch.rand(3, 4)}, + ... ("nested", "disc"): {"logits": torch.randn(3, 10)} + ... }, [3]) + >>> dist = CompositeDistribution(params, + ... distribution_map={"cont": d.Normal, ("nested", "disc"): d.Categorical}) + >>> sample = dist.sample((4,)) + >>> with set_composite_lp_aggregate(False): + ... lp = dist.log_prob(sample) + ... print(lp) + TensorDict( + fields={ + cont_log_prob: Tensor(shape=torch.Size([4, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False), + nested: TensorDict( + fields={ + disc_log_prob: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, + batch_size=torch.Size([4, 3]), + device=None, + is_shared=False)}, + batch_size=torch.Size([4, 3]), + device=None, + is_shared=False) + >>> with set_composite_lp_aggregate(True): + ... lp = dist.log_prob(sample) + ... print(lp) + tensor([[-2.0886, -1.2155, -0.0414], + [-2.8973, -5.5165, 2.4402], + [-0.2806, -1.2799, 3.1733], + [-3.0407, -4.3593, 0.5763]]) + """ + + def __init__( + self, + mode: bool = True, + ) -> None: + super().__init__() + self.mode = mode + + def clone(self) -> set_composite_lp_aggregate: + # override this method if your children class takes __init__ parameters + return type(self)(self.mode) + + def __enter__(self) -> None: + self.prev = _composite_lp_aggregate.get_mode() + _composite_lp_aggregate.set_mode(self.mode) + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + _composite_lp_aggregate.set_mode(self.prev) + + def set(self): + self.__enter__() + + def unset(self): + return self.__exit__(None, None, None) diff --git a/lib/python3.12/site-packages/tensordict/prototype/__init__.py b/lib/python3.12/site-packages/tensordict/prototype/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a30550412bcc930a34716ee3f56cf8822e680640 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/prototype/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. +from tensordict.prototype.fx import symbolic_trace + +__all__ = [ + "symbolic_trace", +] diff --git a/lib/python3.12/site-packages/tensordict/prototype/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/prototype/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ffa38f57010a9e061dea1f91f8715a3e95f877e8 Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/prototype/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/prototype/__pycache__/fx.cpython-312.pyc b/lib/python3.12/site-packages/tensordict/prototype/__pycache__/fx.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5f219ea70f2299577ba805976d5ec5929995dcfb Binary files /dev/null and b/lib/python3.12/site-packages/tensordict/prototype/__pycache__/fx.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/tensordict/prototype/fx.py b/lib/python3.12/site-packages/tensordict/prototype/fx.py new file mode 100644 index 0000000000000000000000000000000000000000..090d296efd2a438234b48eb6e8311b4bd7aafb88 --- /dev/null +++ b/lib/python3.12/site-packages/tensordict/prototype/fx.py @@ -0,0 +1,199 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import operator +from itertools import filterfalse, tee +from typing import Any, Callable, Iterable + +from tensordict._td import _unravel_key_to_tuple + +from tensordict.nn import TensorDictModule, TensorDictModuleBase, TensorDictSequential +from tensordict.tensordict import TensorDictBase +from tensordict.utils import _zip_strict, NestedKey +from torch import fx, nn + +__all__ = ["symbolic_trace"] + + +class TDGraphModule(nn.Module): + """A graph module for TensorDict.""" + + def __init__( + self, + graph_module: fx.GraphModule, + out_keys: list[NestedKey], + ) -> None: + super().__init__() + self.out_keys = [_unravel_key_to_tuple(ok) for ok in out_keys] + self._gm = graph_module + + def forward( + self, + tensordict: TensorDictBase, + tensordict_out: TensorDictBase | None = None, + **kwargs, + ) -> TensorDictBase: + outputs = self._gm(tensordict, **kwargs) + + if tensordict_out is None: + tensordict_out = tensordict + + for out_key, output in _zip_strict(self.out_keys, outputs): + if out_key != "_": + tensordict_out._set_tuple( + out_key, output, inplace=False, validated=True, non_blocking=False + ) + + return tensordict_out + + def __getattr__(self, name: str) -> Any: + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self._gm, name) + + +def symbolic_trace(td_module: TensorDictModuleBase) -> TDGraphModule: + """A symbolic tracer for TensorDictModule.""" + if isinstance(td_module, TensorDictSequential): + return _trace_tensordictsequential(td_module) + elif isinstance(td_module, TensorDictModule): + return _trace_tensordictmodule(td_module) + raise TypeError(f"Unsupported type {type(td_module)}") + + +# cf. https://docs.python.org/3/library/itertools.html#itertools-recipes +def _partition( + pred: Callable[..., bool], iterable: Iterable[Any] +) -> tuple[Iterable[Any], Iterable[Any]]: + """Use a predicate to partition entries into false entries and true entries.""" + # partition(is_odd, range(10)) --> 0 2 4 6 8 and 1 3 5 7 9 + t1, t2 = tee(iterable) + return filterfalse(pred, t1), filter(pred, t2) + + +def _parse_input_nodes( + in_keys: list[NestedKey], nodes, td: TensorDictBase, inputs: tuple[Any, ...], env +): + for in_key, node in zip(in_keys, nodes): + if in_key in inputs: + new_node = inputs[in_key] + else: + output_proxy = operator.getitem(td, in_key) + new_node = output_proxy.node + inputs[in_key] = new_node + env[node.name] = new_node + + +def _trace_tensordictmodule(td_module: TensorDictModuleBase) -> TDGraphModule: + # this graph manipulation is based heavily on example in the PyTorch docs + # https://pytorch.org/docs/stable/fx.html#proxy-retracing + + # trace the graph of the underlying module + graph = fx.Tracer().trace(td_module.module) + + # create a new graph which we will populate from the old one + new_graph = fx.Graph() + env = {} + + # create a new placeholder for the input tensordict + td = fx.Proxy(new_graph.placeholder("tensordict")) + + node_iter = iter(graph.nodes) + + # the first nodes, in order, are placeholders for the in_keys. We consume them and + # convert them to "call_function" nodes with target=operator.getitem. + _parse_input_nodes(td_module.in_keys, node_iter, td, {}, env) + + # the remaining nodes we simply clone, pulling any arguments from the env + for node in node_iter: + new_node = new_graph.node_copy(node, lambda x: env[x.name]) + env[node.name] = new_node + + return TDGraphModule( + fx.GraphModule(td_module.module, new_graph), td_module.out_keys + ) + + +def _trace_tensordictsequential(td_sequential: TensorDictSequential) -> TDGraphModule: + # we track values previously read from / written to the tensordict by storing the + # nodes / proxy values in the inputs / outputs dictionaries + inputs = {} + outputs = {} + # env is a lookup for nodes in the new graph using names from the old graph + env = {} + + new_graph = fx.Graph() + td = fx.Proxy(new_graph.placeholder("tensordict")) + + for i, td_module in enumerate(td_sequential.module): + # trace the submodule + if isinstance(td_module, TensorDictSequential): + graph = _trace_tensordictsequential(td_module).graph + node_iter = iter(graph.nodes) + _td = next(node_iter) # tensordict placeholder from submodule graph + + # in the graph of TensorDictSequential, the getitem calls to the tensordict + # need not come first, so we partition nodes into getitem calls on the + # placeholder tensordict (input_nodes) and the remaining nodes + node_iter, input_nodes = _partition( + lambda node, _td=_td: ( + node.op == "call_function" + and node.target == operator.getitem + and node.args[0] == _td + ), + node_iter, + ) + _parse_input_nodes(td_module.in_keys, input_nodes, td, inputs, env) + + else: + graph = fx.Tracer().trace(td_module.module) + # in the trace of a regular nn.Module the placeholder nodes all come first, + # so we just consume them in order + node_iter = iter(graph.nodes) + _parse_input_nodes(td_module.in_keys, node_iter, td, inputs, env) + + # clone the remaining nodes + for node in node_iter: + if node.op == "output": + # capture the outputs but don't clone the output node (this would + # result in prematurely returning intermediate values) + + # need to unpack the args in the case that the submodule is itself a + # TensorDictSequential that returns a tuple of arguments + args = ( + node.args[0] + if isinstance(td_module, TensorDictSequential) + else node.args + ) + + # if the submodule has multiple outputs, args has structure + # ((out1, out2,),), so we need to do some extra unpacking + args = args[0] if isinstance(args[0], tuple) else args + + for out_key, arg in _zip_strict(td_module.out_keys, args): + # any outputs of submodules will need to be returned at the end + outputs[out_key] = env[arg.name] + # we also need to make outputs of submodules available as inputs + # to subsequent submodules + inputs[out_key] = env[arg.name] + else: + new_node = new_graph.node_copy(node, lambda x: env[x.name]) + if new_node.op in ("call_module", "get_attr"): + # since we traced the submodule in isolation, we need to patch the + # targets of any calls to methods on the module or attribute access + new_node.target = f"{i}.module.{new_node.target}" + new_node.name = f"_{i}_{new_node.name}" + env[node.name] = new_node + + # finally we add a new output node that collects all of the output values from + # submodules in the graph and returns them together + new_graph.output(tuple(outputs.values())) + + return TDGraphModule( + fx.GraphModule(td_sequential.module, new_graph), tuple(outputs.keys()) + ) diff --git a/lib/python3.12/site-packages/wandb/__init__.py b/lib/python3.12/site-packages/wandb/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..848e8e8ae1cf819ff307e5173c5eb31bf3a45927 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/__init__.py @@ -0,0 +1,246 @@ +"""Use wandb to track machine learning work. + +Train and fine-tune models, manage models from experimentation to production. + +For guides and examples, see https://docs.wandb.ai. + +For scripts and interactive notebooks, see https://github.com/wandb/examples. + +For reference documentation, see https://docs.wandb.com/ref/python. +""" +from __future__ import annotations + +__version__ = "0.21.0" + + +from wandb.errors import Error + +# This needs to be early as other modules call it. +from wandb.errors.term import termsetup, termlog, termerror, termwarn + +# Configure the logger as early as possible for consistent behavior. +from wandb.sdk.lib import wb_logging as _wb_logging +_wb_logging.configure_wandb_logger() + +from wandb import sdk as wandb_sdk + +import wandb + +wandb.wandb_lib = wandb_sdk.lib # type: ignore + +init = wandb_sdk.init +setup = wandb_sdk.setup +attach = _attach = wandb_sdk._attach +_sync = wandb_sdk._sync +teardown = _teardown = wandb_sdk.teardown +finish = wandb_sdk.finish +join = finish +login = wandb_sdk.login +helper = wandb_sdk.helper +sweep = wandb_sdk.sweep +controller = wandb_sdk.controller +require = wandb_sdk.require +Artifact = wandb_sdk.Artifact +AlertLevel = wandb_sdk.AlertLevel +Settings = wandb_sdk.Settings +Config = wandb_sdk.Config + +from wandb.apis import InternalApi, PublicApi +from wandb.errors import CommError, UsageError + +_preinit = wandb.wandb_lib.preinit # type: ignore +_lazyloader = wandb.wandb_lib.lazyloader # type: ignore + +from wandb.integration.torch import wandb_torch + +# Move this (keras.__init__ expects it at top level) +from wandb.sdk.data_types._private import _cleanup_media_tmp_dir + +_cleanup_media_tmp_dir() + +from wandb.data_types import Graph +from wandb.data_types import Image +from wandb.data_types import Plotly + +# from wandb.data_types import Bokeh # keeping out of top level for now since Bokeh plots have poor UI +from wandb.data_types import Video +from wandb.data_types import Audio +from wandb.data_types import Table +from wandb.data_types import Html +from wandb.data_types import box3d +from wandb.data_types import Object3D +from wandb.data_types import Molecule +from wandb.data_types import Histogram +from wandb.data_types import Classes +from wandb.data_types import JoinedTable + +from wandb.wandb_agent import agent + +from wandb.plot import visualize, plot_table +from wandb.integration.sagemaker import sagemaker_auth +from wandb.sdk.internal import profiler + +# Artifact import types +from wandb.sdk.artifacts.artifact_ttl import ArtifactTTL + +# Used to make sure we don't use some code in the incorrect process context +_IS_INTERNAL_PROCESS = False + + +def _set_internal_process(disable=False): + global _IS_INTERNAL_PROCESS + if _IS_INTERNAL_PROCESS is None: + return + if disable: + _IS_INTERNAL_PROCESS = None + return + _IS_INTERNAL_PROCESS = True + + +def _assert_is_internal_process(): + if _IS_INTERNAL_PROCESS is None: + return + assert _IS_INTERNAL_PROCESS + + +def _assert_is_user_process(): + if _IS_INTERNAL_PROCESS is None: + return + assert not _IS_INTERNAL_PROCESS + + +# globals +Api = PublicApi +api = InternalApi() +run: wandb_sdk.wandb_run.Run | None = None +config = _preinit.PreInitObject("wandb.config", wandb_sdk.wandb_config.Config) +summary = _preinit.PreInitObject("wandb.summary", wandb_sdk.wandb_summary.Summary) +log = _preinit.PreInitCallable("wandb.log", wandb_sdk.wandb_run.Run.log) # type: ignore +watch = _preinit.PreInitCallable("wandb.watch", wandb_sdk.wandb_run.Run.watch) # type: ignore +unwatch = _preinit.PreInitCallable("wandb.unwatch", wandb_sdk.wandb_run.Run.unwatch) # type: ignore +save = _preinit.PreInitCallable("wandb.save", wandb_sdk.wandb_run.Run.save) # type: ignore +restore = wandb_sdk.wandb_run.restore +use_artifact = _preinit.PreInitCallable( + "wandb.use_artifact", wandb_sdk.wandb_run.Run.use_artifact # type: ignore +) +log_artifact = _preinit.PreInitCallable( + "wandb.log_artifact", wandb_sdk.wandb_run.Run.log_artifact # type: ignore +) +log_model = _preinit.PreInitCallable( + "wandb.log_model", wandb_sdk.wandb_run.Run.log_model # type: ignore +) +use_model = _preinit.PreInitCallable( + "wandb.use_model", wandb_sdk.wandb_run.Run.use_model # type: ignore +) +link_model = _preinit.PreInitCallable( + "wandb.link_model", wandb_sdk.wandb_run.Run.link_model # type: ignore +) +define_metric = _preinit.PreInitCallable( + "wandb.define_metric", wandb_sdk.wandb_run.Run.define_metric # type: ignore +) + +mark_preempting = _preinit.PreInitCallable( + "wandb.mark_preempting", wandb_sdk.wandb_run.Run.mark_preempting # type: ignore +) + +alert = _preinit.PreInitCallable("wandb.alert", wandb_sdk.wandb_run.Run.alert) # type: ignore + +# record of patched libraries +patched = {"tensorboard": [], "keras": [], "gym": []} # type: ignore + +keras = _lazyloader.LazyLoader("wandb.keras", globals(), "wandb.integration.keras") +sklearn = _lazyloader.LazyLoader("wandb.sklearn", globals(), "wandb.sklearn") +tensorflow = _lazyloader.LazyLoader( + "wandb.tensorflow", globals(), "wandb.integration.tensorflow" +) +xgboost = _lazyloader.LazyLoader( + "wandb.xgboost", globals(), "wandb.integration.xgboost" +) +catboost = _lazyloader.LazyLoader( + "wandb.catboost", globals(), "wandb.integration.catboost" +) +tensorboard = _lazyloader.LazyLoader( + "wandb.tensorboard", globals(), "wandb.integration.tensorboard" +) +gym = _lazyloader.LazyLoader("wandb.gym", globals(), "wandb.integration.gym") +lightgbm = _lazyloader.LazyLoader( + "wandb.lightgbm", globals(), "wandb.integration.lightgbm" +) +jupyter = _lazyloader.LazyLoader("wandb.jupyter", globals(), "wandb.jupyter") +sacred = _lazyloader.LazyLoader("wandb.sacred", globals(), "wandb.integration.sacred") + + +def ensure_configured(): + global api + api = InternalApi() + + +def set_trace(): + import pdb # TODO: support other debuggers + + # frame = sys._getframe().f_back + pdb.set_trace() # TODO: pass the parent stack... + + +def load_ipython_extension(ipython): + ipython.register_magics(wandb.jupyter.WandBMagics) + + +if wandb_sdk.lib.ipython.in_notebook(): + from IPython import get_ipython # type: ignore[import-not-found] + + load_ipython_extension(get_ipython()) + + +from .analytics import Sentry as _Sentry + +if "dev" in __version__: + import wandb.env + import os + + # Disable error reporting in dev versions. + os.environ[wandb.env.ERROR_REPORTING] = os.environ.get( + wandb.env.ERROR_REPORTING, + "false", + ) + +_sentry = _Sentry() +_sentry.setup() + + +__all__ = ( + "__version__", + "init", + "finish", + "setup", + "save", + "sweep", + "controller", + "agent", + "config", + "log", + "summary", + "join", + "Api", + "Graph", + "Image", + "Plotly", + "Video", + "Audio", + "Table", + "Html", + "box3d", + "Object3D", + "Molecule", + "Histogram", + "ArtifactTTL", + "log_artifact", + "use_artifact", + "log_model", + "use_model", + "link_model", + "define_metric", + "watch", + "unwatch", + "plot_table", +) diff --git a/lib/python3.12/site-packages/wandb/__init__.pyi b/lib/python3.12/site-packages/wandb/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..de172c5732b341d1f7b5734ffbdf87bfe76a1adb --- /dev/null +++ b/lib/python3.12/site-packages/wandb/__init__.pyi @@ -0,0 +1,1253 @@ +"""Use wandb to track machine learning work. + +Train and fine-tune models, manage models from experimentation to production. + +For guides and examples, see https://docs.wandb.ai. + +For scripts and interactive notebooks, see https://github.com/wandb/examples. + +For reference documentation, see https://docs.wandb.com/ref/python. +""" + +from __future__ import annotations + +__all__ = ( + "__version__", # doc:exclude + "init", + "finish", + "setup", + "login", + "save", # doc:exclude + "sweep", + "controller", + "agent", + "config", # doc:exclude + "log", # doc:exclude + "summary", # doc:exclude + "Api", + "Graph", # doc:exclude + "Image", + "Plotly", + "Video", + "Audio", + "Table", + "Html", + "box3d", + "Object3D", + "Molecule", + "Histogram", + "ArtifactTTL", # doc:exclude + "log_artifact", # doc:exclude + "use_artifact", # doc:exclude + "log_model", # doc:exclude + "use_model", # doc:exclude + "link_model", # doc:exclude + "define_metric", # doc:exclude + "Error", # doc:exclude + "termsetup", # doc:exclude + "termlog", # doc:exclude + "termerror", # doc:exclude + "termwarn", # doc:exclude + "Artifact", + "Settings", + "teardown", + "watch", # doc:exclude + "unwatch", # doc:exclude + "plot", # doc:exclude + "plot_table", + "restore", + "Run", +) + +import os +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Iterable, + List, + Literal, + Optional, + Sequence, + TextIO, + Union, +) + +import wandb.plot as plot +from wandb.analytics import Sentry +from wandb.apis import InternalApi +from wandb.apis import PublicApi as Api +from wandb.data_types import ( + Audio, + Graph, + Histogram, + Html, + Image, + Molecule, + Object3D, + Plotly, + Table, + Video, + box3d, +) +from wandb.errors import Error +from wandb.errors.term import termerror, termlog, termsetup, termwarn +from wandb.sdk import Artifact, Settings, wandb_config, wandb_metric, wandb_summary +from wandb.sdk.artifacts.artifact_ttl import ArtifactTTL +from wandb.sdk.interface.interface import PolicyName +from wandb.sdk.lib.paths import FilePathStr, StrPath +from wandb.sdk.wandb_run import Run +from wandb.sdk.wandb_setup import _WandbSetup +from wandb.wandb_controller import _WandbController + +if TYPE_CHECKING: + import torch # type: ignore [import-not-found] + + import wandb + from wandb.plot import CustomChart + +__version__: str = "0.21.0" + +run: Run | None +config: wandb_config.Config +summary: wandb_summary.Summary + +# private attributes +_sentry: Sentry +api: InternalApi +patched: Dict[str, List[Callable]] + +def require( + requirement: str | Iterable[str] | None = None, + experiment: str | Iterable[str] | None = None, +) -> None: + """Indicate which experimental features are used by the script. + + This should be called before any other `wandb` functions, ideally right + after importing `wandb`. + + Args: + requirement: The name of a feature to require or an iterable of + feature names. + experiment: An alias for `requirement`. + + Raises: + wandb.errors.UnsupportedError: If a feature name is unknown. + """ + ... + +def setup(settings: Settings | None = None) -> _WandbSetup: + """Prepares W&B for use in the current process and its children. + + You can usually ignore this as it is implicitly called by `wandb.init()`. + + When using wandb in multiple processes, calling `wandb.setup()` + in the parent process before starting child processes may improve + performance and resource utilization. + + Note that `wandb.setup()` modifies `os.environ`, and it is important + that child processes inherit the modified environment variables. + + See also `wandb.teardown()`. + + Args: + settings: Configuration settings to apply globally. These can be + overridden by subsequent `wandb.init()` calls. + + Example: + ```python + import multiprocessing + + import wandb + + def run_experiment(params): + with wandb.init(config=params): + # Run experiment + pass + + if __name__ == "__main__": + # Start backend and set global config + wandb.setup(settings={"project": "my_project"}) + + # Define experiment parameters + experiment_params = [ + {"learning_rate": 0.01, "epochs": 10}, + {"learning_rate": 0.001, "epochs": 20}, + ] + + # Start multiple processes, each running a separate experiment + processes = [] + for params in experiment_params: + p = multiprocessing.Process(target=run_experiment, args=(params,)) + p.start() + processes.append(p) + + # Wait for all processes to complete + for p in processes: + p.join() + + # Optional: Explicitly shut down the backend + wandb.teardown() + ``` + """ + ... + +def teardown(exit_code: int | None = None) -> None: + """Waits for wandb to finish and frees resources. + + Completes any runs that were not explicitly finished + using `run.finish()` and waits for all data to be uploaded. + + It is recommended to call this at the end of a session + that used `wandb.setup()`. It is invoked automatically + in an `atexit` hook, but this is not reliable in certain setups + such as when using Python's `multiprocessing` module. + """ + ... + +def init( + entity: str | None = None, + project: str | None = None, + dir: StrPath | None = None, + id: str | None = None, + name: str | None = None, + notes: str | None = None, + tags: Sequence[str] | None = None, + config: dict[str, Any] | str | None = None, + config_exclude_keys: list[str] | None = None, + config_include_keys: list[str] | None = None, + allow_val_change: bool | None = None, + group: str | None = None, + job_type: str | None = None, + mode: Literal["online", "offline", "disabled"] | None = None, + force: bool | None = None, + anonymous: Literal["never", "allow", "must"] | None = None, + reinit: ( + bool + | Literal[ + None, + "default", + "return_previous", + "finish_previous", + "create_new", + ] + ) = None, + resume: bool | Literal["allow", "never", "must", "auto"] | None = None, + resume_from: str | None = None, + fork_from: str | None = None, + save_code: bool | None = None, + tensorboard: bool | None = None, + sync_tensorboard: bool | None = None, + monitor_gym: bool | None = None, + settings: Settings | dict[str, Any] | None = None, +) -> Run: + r"""Start a new run to track and log to W&B. + + In an ML training pipeline, you could add `wandb.init()` to the beginning of + your training script as well as your evaluation script, and each piece would + be tracked as a run in W&B. + + `wandb.init()` spawns a new background process to log data to a run, and it + also syncs data to https://wandb.ai by default, so you can see your results + in real-time. + + Call `wandb.init()` to start a run before logging data with `wandb.log()`. + When you're done logging data, call `wandb.finish()` to end the run. If you + don't call `wandb.finish()`, the run will end when your script exits. + + For more on using `wandb.init()`, including detailed examples, check out our + [guide and FAQs](https://docs.wandb.ai/guides/track/launch). + + Examples: + ### Explicitly set the entity and project and choose a name for the run: + + ```python + import wandb + + run = wandb.init( + entity="geoff", + project="capsules", + name="experiment-2021-10-31", + ) + + # ... your training code here ... + + run.finish() + ``` + + ### Add metadata about the run using the `config` argument: + + ```python + import wandb + + config = {"lr": 0.01, "batch_size": 32} + with wandb.init(config=config) as run: + run.config.update({"architecture": "resnet", "depth": 34}) + + # ... your training code here ... + ``` + + Note that you can use `wandb.init()` as a context manager to automatically + call `wandb.finish()` at the end of the block. + + Args: + entity: The username or team name under which the runs will be logged. + The entity must already exist, so ensure you’ve created your account + or team in the UI before starting to log runs. If not specified, the + run will default your default entity. To change the default entity, + go to [your settings](https://wandb.ai/settings) and update the + "Default location to create new projects" under "Default team". + project: The name of the project under which this run will be logged. + If not specified, we use a heuristic to infer the project name based + on the system, such as checking the git root or the current program + file. If we can't infer the project name, the project will default to + `"uncategorized"`. + dir: The absolute path to the directory where experiment logs and + metadata files are stored. If not specified, this defaults + to the `./wandb` directory. Note that this does not affect the + location where artifacts are stored when calling `download()`. + id: A unique identifier for this run, used for resuming. It must be unique + within the project and cannot be reused once a run is deleted. The + identifier must not contain any of the following special characters: + `/ \ # ? % :`. For a short descriptive name, use the `name` field, + or for saving hyperparameters to compare across runs, use `config`. + name: A short display name for this run, which appears in the UI to help + you identify it. By default, we generate a random two-word name + allowing easy cross-reference runs from table to charts. Keeping these + run names brief enhances readability in chart legends and tables. For + saving hyperparameters, we recommend using the `config` field. + notes: A detailed description of the run, similar to a commit message in + Git. Use this argument to capture any context or details that may + help you recall the purpose or setup of this run in the future. + tags: A list of tags to label this run in the UI. Tags are helpful for + organizing runs or adding temporary identifiers like "baseline" or + "production." You can easily add, remove tags, or filter by tags in + the UI. + If resuming a run, the tags provided here will replace any existing + tags. To add tags to a resumed run without overwriting the current + tags, use `run.tags += ("new_tag",)` after calling `run = wandb.init()`. + config: Sets `wandb.config`, a dictionary-like object for storing input + parameters to your run, such as model hyperparameters or data + preprocessing settings. + The config appears in the UI in an overview page, allowing you to + group, filter, and sort runs based on these parameters. + Keys should not contain periods (`.`), and values should be + smaller than 10 MB. + If a dictionary, `argparse.Namespace`, or `absl.flags.FLAGS` is + provided, the key-value pairs will be loaded directly into + `wandb.config`. + If a string is provided, it is interpreted as a path to a YAML file, + from which configuration values will be loaded into `wandb.config`. + config_exclude_keys: A list of specific keys to exclude from `wandb.config`. + config_include_keys: A list of specific keys to include in `wandb.config`. + allow_val_change: Controls whether config values can be modified after their + initial set. By default, an exception is raised if a config value is + overwritten. For tracking variables that change during training, such as + a learning rate, consider using `wandb.log()` instead. By default, this + is `False` in scripts and `True` in Notebook environments. + group: Specify a group name to organize individual runs as part of a larger + experiment. This is useful for cases like cross-validation or running + multiple jobs that train and evaluate a model on different test sets. + Grouping allows you to manage related runs collectively in the UI, + making it easy to toggle and review results as a unified experiment. + For more information, refer to our + [guide to grouping runs](https://docs.wandb.com/guides/runs/grouping). + job_type: Specify the type of run, especially helpful when organizing runs + within a group as part of a larger experiment. For example, in a group, + you might label runs with job types such as "train" and "eval". + Defining job types enables you to easily filter and group similar runs + in the UI, facilitating direct comparisons. + mode: Specifies how run data is managed, with the following options: + - `"online"` (default): Enables live syncing with W&B when a network + connection is available, with real-time updates to visualizations. + - `"offline"`: Suitable for air-gapped or offline environments; data + is saved locally and can be synced later. Ensure the run folder + is preserved to enable future syncing. + - `"disabled"`: Disables all W&B functionality, making the run’s methods + no-ops. Typically used in testing to bypass W&B operations. + force: Determines if a W&B login is required to run the script. If `True`, + the user must be logged in to W&B; otherwise, the script will not + proceed. If `False` (default), the script can proceed without a login, + switching to offline mode if the user is not logged in. + anonymous: Specifies the level of control over anonymous data logging. + Available options are: + - `"never"` (default): Requires you to link your W&B account before + tracking the run. This prevents unintentional creation of anonymous + runs by ensuring each run is associated with an account. + - `"allow"`: Enables a logged-in user to track runs with their account, + but also allows someone running the script without a W&B account + to view the charts and data in the UI. + - `"must"`: Forces the run to be logged to an anonymous account, even + if the user is logged in. + reinit: Shorthand for the "reinit" setting. Determines the behavior of + `wandb.init()` when a run is active. + resume: Controls the behavior when resuming a run with the specified `id`. + Available options are: + - `"allow"`: If a run with the specified `id` exists, it will resume + from the last step; otherwise, a new run will be created. + - `"never"`: If a run with the specified `id` exists, an error will + be raised. If no such run is found, a new run will be created. + - `"must"`: If a run with the specified `id` exists, it will resume + from the last step. If no run is found, an error will be raised. + - `"auto"`: Automatically resumes the previous run if it crashed on + this machine; otherwise, starts a new run. + - `True`: Deprecated. Use `"auto"` instead. + - `False`: Deprecated. Use the default behavior (leaving `resume` + unset) to always start a new run. + Note: If `resume` is set, `fork_from` and `resume_from` cannot be + used. When `resume` is unset, the system will always start a new run. + For more details, see our + [guide to resuming runs](https://docs.wandb.com/guides/runs/resuming). + resume_from: Specifies a moment in a previous run to resume a run from, + using the format `{run_id}?_step={step}`. This allows users to truncate + the history logged to a run at an intermediate step and resume logging + from that step. The target run must be in the same project. + If an `id` argument is also provided, the `resume_from` argument will + take precedence. + `resume`, `resume_from` and `fork_from` cannot be used together, only + one of them can be used at a time. + Note: This feature is in beta and may change in the future. + fork_from: Specifies a point in a previous run from which to fork a new + run, using the format `{id}?_step={step}`. This creates a new run that + resumes logging from the specified step in the target run’s history. + The target run must be part of the current project. + If an `id` argument is also provided, it must be different from the + `fork_from` argument, an error will be raised if they are the same. + `resume`, `resume_from` and `fork_from` cannot be used together, only + one of them can be used at a time. + Note: This feature is in beta and may change in the future. + save_code: Enables saving the main script or notebook to W&B, aiding in + experiment reproducibility and allowing code comparisons across runs in + the UI. By default, this is disabled, but you can change the default to + enable on your [settings page](https://wandb.ai/settings). + tensorboard: Deprecated. Use `sync_tensorboard` instead. + sync_tensorboard: Enables automatic syncing of W&B logs from TensorBoard + or TensorBoardX, saving relevant event files for viewing in the W&B UI. + saving relevant event files for viewing in the W&B UI. (Default: `False`) + monitor_gym: Enables automatic logging of videos of the environment when + using OpenAI Gym. For additional details, see our + [guide for gym integration](https://docs.wandb.com/guides/integrations/openai-gym). + settings: Specifies a dictionary or `wandb.Settings` object with advanced + settings for the run. + + Returns: + A `Run` object, which is a handle to the current run. Use this object + to perform operations like logging data, saving files, and finishing + the run. See the [Run API](https://docs.wandb.ai/ref/python/run) for + more details. + + Raises: + Error: If some unknown or internal error happened during the run + initialization. + AuthenticationError: If the user failed to provide valid credentials. + CommError: If there was a problem communicating with the W&B server. + UsageError: If the user provided invalid arguments to the function. + KeyboardInterrupt: If the user interrupts the run initialization process. + If the user interrupts the run initialization process. + """ + ... + +def finish( + exit_code: int | None = None, + quiet: bool | None = None, +) -> None: + """Finish a run and upload any remaining data. + + Marks the completion of a W&B run and ensures all data is synced to the server. + The run's final state is determined by its exit conditions and sync status. + + Run States: + - Running: Active run that is logging data and/or sending heartbeats. + - Crashed: Run that stopped sending heartbeats unexpectedly. + - Finished: Run completed successfully (`exit_code=0`) with all data synced. + - Failed: Run completed with errors (`exit_code!=0`). + + Args: + exit_code: Integer indicating the run's exit status. Use 0 for success, + any other value marks the run as failed. + quiet: Deprecated. Configure logging verbosity using `wandb.Settings(quiet=...)`. + """ + ... + +def login( + anonymous: Optional[Literal["must", "allow", "never"]] = None, + key: Optional[str] = None, + relogin: Optional[bool] = None, + host: Optional[str] = None, + force: Optional[bool] = None, + timeout: Optional[int] = None, + verify: bool = False, + referrer: Optional[str] = None, +) -> bool: + """Set up W&B login credentials. + + By default, this will only store credentials locally without + verifying them with the W&B server. To verify credentials, pass + `verify=True`. + + Args: + anonymous: (string, optional) Can be "must", "allow", or "never". + If set to "must", always log a user in anonymously. If set to + "allow", only create an anonymous user if the user + isn't already logged in. If set to "never", never log a + user anonymously. Default set to "never". + key: (string, optional) The API key to use. + relogin: (bool, optional) If true, will re-prompt for API key. + host: (string, optional) The host to connect to. + force: (bool, optional) If true, will force a relogin. + timeout: (int, optional) Number of seconds to wait for user input. + verify: (bool) Verify the credentials with the W&B server. + referrer: (string, optional) The referrer to use in the URL login request. + + Returns: + bool: if key is configured + + Raises: + AuthenticationError - if api_key fails verification with the server + UsageError - if api_key cannot be configured and no tty + """ + ... + +def log( + data: dict[str, Any], + step: int | None = None, + commit: bool | None = None, +) -> None: + """Upload run data. + + Use `log` to log data from runs, such as scalars, images, video, + histograms, plots, and tables. + + See our [guides to logging](https://docs.wandb.ai/guides/track/log) for + live examples, code snippets, best practices, and more. + + The most basic usage is `run.log({"train-loss": 0.5, "accuracy": 0.9})`. + This will save the loss and accuracy to the run's history and update + the summary values for these metrics. + + Visualize logged data in the workspace at [wandb.ai](https://wandb.ai), + or locally on a [self-hosted instance](https://docs.wandb.ai/guides/hosting) + of the W&B app, or export data to visualize and explore locally, e.g. in + Jupyter notebooks, with [our API](https://docs.wandb.ai/guides/track/public-api-guide). + + Logged values don't have to be scalars. Logging any wandb object is supported. + For example `run.log({"example": wandb.Image("myimage.jpg")})` will log an + example image which will be displayed nicely in the W&B UI. + See the [reference documentation](https://docs.wandb.com/ref/python/data-types) + for all of the different supported types or check out our + [guides to logging](https://docs.wandb.ai/guides/track/log) for examples, + from 3D molecular structures and segmentation masks to PR curves and histograms. + You can use `wandb.Table` to log structured data. See our + [guide to logging tables](https://docs.wandb.ai/guides/models/tables/tables-walkthrough) + for details. + + The W&B UI organizes metrics with a forward slash (`/`) in their name + into sections named using the text before the final slash. For example, + the following results in two sections named "train" and "validate": + + ``` + run.log( + { + "train/accuracy": 0.9, + "train/loss": 30, + "validate/accuracy": 0.8, + "validate/loss": 20, + } + ) + ``` + + Only one level of nesting is supported; `run.log({"a/b/c": 1})` + produces a section named "a/b". + + `run.log` is not intended to be called more than a few times per second. + For optimal performance, limit your logging to once every N iterations, + or collect data over multiple iterations and log it in a single step. + + ### The W&B step + + With basic usage, each call to `log` creates a new "step". + The step must always increase, and it is not possible to log + to a previous step. + + Note that you can use any metric as the X axis in charts. + In many cases, it is better to treat the W&B step like + you'd treat a timestamp rather than a training step. + + ``` + # Example: log an "epoch" metric for use as an X axis. + run.log({"epoch": 40, "train-loss": 0.5}) + ``` + + See also [define_metric](https://docs.wandb.ai/ref/python/run#define_metric). + + It is possible to use multiple `log` invocations to log to + the same step with the `step` and `commit` parameters. + The following are all equivalent: + + ``` + # Normal usage: + run.log({"train-loss": 0.5, "accuracy": 0.8}) + run.log({"train-loss": 0.4, "accuracy": 0.9}) + + # Implicit step without auto-incrementing: + run.log({"train-loss": 0.5}, commit=False) + run.log({"accuracy": 0.8}) + run.log({"train-loss": 0.4}, commit=False) + run.log({"accuracy": 0.9}) + + # Explicit step: + run.log({"train-loss": 0.5}, step=current_step) + run.log({"accuracy": 0.8}, step=current_step) + current_step += 1 + run.log({"train-loss": 0.4}, step=current_step) + run.log({"accuracy": 0.9}, step=current_step) + ``` + + Args: + data: A `dict` with `str` keys and values that are serializable + Python objects including: `int`, `float` and `string`; + any of the `wandb.data_types`; lists, tuples and NumPy arrays + of serializable Python objects; other `dict`s of this + structure. + step: The step number to log. If `None`, then an implicit + auto-incrementing step is used. See the notes in + the description. + commit: If true, finalize and upload the step. If false, then + accumulate data for the step. See the notes in the description. + If `step` is `None`, then the default is `commit=True`; + otherwise, the default is `commit=False`. + + Examples: + For more and more detailed examples, see + [our guides to logging](https://docs.wandb.com/guides/track/log). + + ### Basic usage + ```python + import wandb + + run = wandb.init() + run.log({"accuracy": 0.9, "epoch": 5}) + ``` + + ### Incremental logging + ```python + import wandb + + run = wandb.init() + run.log({"loss": 0.2}, commit=False) + # Somewhere else when I'm ready to report this step: + run.log({"accuracy": 0.8}) + ``` + + ### Histogram + ```python + import numpy as np + import wandb + + # sample gradients at random from normal distribution + gradients = np.random.randn(100, 100) + run = wandb.init() + run.log({"gradients": wandb.Histogram(gradients)}) + ``` + + ### Image from numpy + ```python + import numpy as np + import wandb + + run = wandb.init() + examples = [] + for i in range(3): + pixels = np.random.randint(low=0, high=256, size=(100, 100, 3)) + image = wandb.Image(pixels, caption=f"random field {i}") + examples.append(image) + run.log({"examples": examples}) + ``` + + ### Image from PIL + ```python + import numpy as np + from PIL import Image as PILImage + import wandb + + run = wandb.init() + examples = [] + for i in range(3): + pixels = np.random.randint( + low=0, + high=256, + size=(100, 100, 3), + dtype=np.uint8, + ) + pil_image = PILImage.fromarray(pixels, mode="RGB") + image = wandb.Image(pil_image, caption=f"random field {i}") + examples.append(image) + run.log({"examples": examples}) + ``` + + ### Video from numpy + ```python + import numpy as np + import wandb + + run = wandb.init() + # axes are (time, channel, height, width) + frames = np.random.randint( + low=0, + high=256, + size=(10, 3, 100, 100), + dtype=np.uint8, + ) + run.log({"video": wandb.Video(frames, fps=4)}) + ``` + + ### Matplotlib Plot + ```python + from matplotlib import pyplot as plt + import numpy as np + import wandb + + run = wandb.init() + fig, ax = plt.subplots() + x = np.linspace(0, 10) + y = x * x + ax.plot(x, y) # plot y = x^2 + run.log({"chart": fig}) + ``` + + ### PR Curve + ```python + import wandb + + run = wandb.init() + run.log({"pr": wandb.plot.pr_curve(y_test, y_probas, labels)}) + ``` + + ### 3D Object + ```python + import wandb + + run = wandb.init() + run.log( + { + "generated_samples": [ + wandb.Object3D(open("sample.obj")), + wandb.Object3D(open("sample.gltf")), + wandb.Object3D(open("sample.glb")), + ] + } + ) + ``` + + Raises: + wandb.Error: if called before `wandb.init` + ValueError: if invalid data is passed + """ + ... + +def save( + glob_str: str | os.PathLike, + base_path: str | os.PathLike | None = None, + policy: PolicyName = "live", +) -> bool | list[str]: + """Sync one or more files to W&B. + + Relative paths are relative to the current working directory. + + A Unix glob, such as "myfiles/*", is expanded at the time `save` is + called regardless of the `policy`. In particular, new files are not + picked up automatically. + + A `base_path` may be provided to control the directory structure of + uploaded files. It should be a prefix of `glob_str`, and the directory + structure beneath it is preserved. It's best understood through + examples: + + ``` + wandb.save("these/are/myfiles/*") + # => Saves files in a "these/are/myfiles/" folder in the run. + + wandb.save("these/are/myfiles/*", base_path="these") + # => Saves files in an "are/myfiles/" folder in the run. + + wandb.save("/User/username/Documents/run123/*.txt") + # => Saves files in a "run123/" folder in the run. See note below. + + wandb.save("/User/username/Documents/run123/*.txt", base_path="/User") + # => Saves files in a "username/Documents/run123/" folder in the run. + + wandb.save("files/*/saveme.txt") + # => Saves each "saveme.txt" file in an appropriate subdirectory + # of "files/". + ``` + + Note: when given an absolute path or glob and no `base_path`, one + directory level is preserved as in the example above. + + Args: + glob_str: A relative or absolute path or Unix glob. + base_path: A path to use to infer a directory structure; see examples. + policy: One of `live`, `now`, or `end`. + * live: upload the file as it changes, overwriting the previous version + * now: upload the file once now + * end: upload file when the run ends + + Returns: + Paths to the symlinks created for the matched files. + + For historical reasons, this may return a boolean in legacy code. + """ + ... + +def sweep( + sweep: Union[dict, Callable], + entity: Optional[str] = None, + project: Optional[str] = None, + prior_runs: Optional[List[str]] = None, +) -> str: + """Initialize a hyperparameter sweep. + + Search for hyperparameters that optimizes a cost function + of a machine learning model by testing various combinations. + + Make note the unique identifier, `sweep_id`, that is returned. + At a later step provide the `sweep_id` to a sweep agent. + + Args: + sweep: The configuration of a hyperparameter search. + (or configuration generator). See + [Sweep configuration structure](https://docs.wandb.ai/guides/sweeps/define-sweep-configuration) + for information on how to define your sweep. + If you provide a callable, ensure that the callable does + not take arguments and that it returns a dictionary that + conforms to the W&B sweep config spec. + entity: The username or team name where you want to send W&B + runs created by the sweep to. Ensure that the entity you + specify already exists. If you don't specify an entity, + the run will be sent to your default entity, + which is usually your username. + project: The name of the project where W&B runs created from + the sweep are sent to. If the project is not specified, the + run is sent to a project labeled 'Uncategorized'. + prior_runs: The run IDs of existing runs to add to this sweep. + + Returns: + sweep_id: str. A unique identifier for the sweep. + """ + ... + +def controller( + sweep_id_or_config: Optional[Union[str, Dict]] = None, + entity: Optional[str] = None, + project: Optional[str] = None, +) -> _WandbController: + """Public sweep controller constructor. + + Usage: + ```python + import wandb + + tuner = wandb.controller(...) + print(tuner.sweep_config) + print(tuner.sweep_id) + tuner.configure_search(...) + tuner.configure_stopping(...) + ``` + """ + ... + +def agent( + sweep_id: str, + function: Optional[Callable] = None, + entity: Optional[str] = None, + project: Optional[str] = None, + count: Optional[int] = None, +) -> None: + """Start one or more sweep agents. + + The sweep agent uses the `sweep_id` to know which sweep it + is a part of, what function to execute, and (optionally) how + many agents to run. + + Args: + sweep_id: The unique identifier for a sweep. A sweep ID + is generated by W&B CLI or Python SDK. + function: A function to call instead of the "program" + specified in the sweep config. + entity: The username or team name where you want to send W&B + runs created by the sweep to. Ensure that the entity you + specify already exists. If you don't specify an entity, + the run will be sent to your default entity, + which is usually your username. + project: The name of the project where W&B runs created from + the sweep are sent to. If the project is not specified, the + run is sent to a project labeled "Uncategorized". + count: The number of sweep config trials to try. + """ + ... + +def define_metric( + name: str, + step_metric: str | wandb_metric.Metric | None = None, + step_sync: bool | None = None, + hidden: bool | None = None, + summary: str | None = None, + goal: str | None = None, + overwrite: bool | None = None, +) -> wandb_metric.Metric: + """Customize metrics logged with `wandb.log()`. + + Args: + name: The name of the metric to customize. + step_metric: The name of another metric to serve as the X-axis + for this metric in automatically generated charts. + step_sync: Automatically insert the last value of step_metric into + `run.log()` if it is not provided explicitly. Defaults to True + if step_metric is specified. + hidden: Hide this metric from automatic plots. + summary: Specify aggregate metrics added to summary. + Supported aggregations include "min", "max", "mean", "last", + "best", "copy" and "none". "best" is used together with the + goal parameter. "none" prevents a summary from being generated. + "copy" is deprecated and should not be used. + goal: Specify how to interpret the "best" summary type. + Supported options are "minimize" and "maximize". + overwrite: If false, then this call is merged with previous + `define_metric` calls for the same metric by using their + values for any unspecified parameters. If true, then + unspecified parameters overwrite values specified by + previous calls. + + Returns: + An object that represents this call but can otherwise be discarded. + """ + ... + +def log_artifact( + artifact_or_path: Artifact | StrPath, + name: str | None = None, + type: str | None = None, + aliases: list[str] | None = None, + tags: list[str] | None = None, +) -> Artifact: + """Declare an artifact as an output of a run. + + Args: + artifact_or_path: (str or Artifact) A path to the contents of this artifact, + can be in the following forms: + - `/local/directory` + - `/local/directory/file.txt` + - `s3://bucket/path` + You can also pass an Artifact object created by calling + `wandb.Artifact`. + name: (str, optional) An artifact name. Valid names can be in the following forms: + - name:version + - name:alias + - digest + This will default to the basename of the path prepended with the current + run id if not specified. + type: (str) The type of artifact to log, examples include `dataset`, `model` + aliases: (list, optional) Aliases to apply to this artifact, + defaults to `["latest"]` + tags: (list, optional) Tags to apply to this artifact, if any. + + Returns: + An `Artifact` object. + """ + ... + +def use_artifact( + artifact_or_name: str | Artifact, + type: str | None = None, + aliases: list[str] | None = None, + use_as: str | None = None, +) -> Artifact: + """Declare an artifact as an input to a run. + + Call `download` or `file` on the returned object to get the contents locally. + + Args: + artifact_or_name: (str or Artifact) An artifact name. + May be prefixed with project/ or entity/project/. + If no entity is specified in the name, the Run or API setting's entity is used. + Valid names can be in the following forms: + - name:version + - name:alias + You can also pass an Artifact object created by calling `wandb.Artifact` + type: (str, optional) The type of artifact to use. + aliases: (list, optional) Aliases to apply to this artifact + use_as: This argument is deprecated and does nothing. + + Returns: + An `Artifact` object. + """ + ... + +def log_model( + path: StrPath, + name: str | None = None, + aliases: list[str] | None = None, +) -> None: + """Logs a model artifact containing the contents inside the 'path' to a run and marks it as an output to this run. + + Args: + path: (str) A path to the contents of this model, + can be in the following forms: + - `/local/directory` + - `/local/directory/file.txt` + - `s3://bucket/path` + name: (str, optional) A name to assign to the model artifact that the file contents will be added to. + The string must contain only the following alphanumeric characters: dashes, underscores, and dots. + This will default to the basename of the path prepended with the current + run id if not specified. + aliases: (list, optional) Aliases to apply to the created model artifact, + defaults to `["latest"]` + + Examples: + ```python + run.log_model( + path="/local/directory", + name="my_model_artifact", + aliases=["production"], + ) + ``` + + Invalid usage + ```python + run.log_model( + path="/local/directory", + name="my_entity/my_project/my_model_artifact", + aliases=["production"], + ) + ``` + + Raises: + ValueError: if name has invalid special characters + + Returns: + None + """ + ... + +def use_model(name: str) -> FilePathStr: + """Download the files logged in a model artifact 'name'. + + Args: + name: (str) A model artifact name. 'name' must match the name of an existing logged + model artifact. + May be prefixed with entity/project/. Valid names + can be in the following forms: + - model_artifact_name:version + - model_artifact_name:alias + + Examples: + ```python + run.use_model( + name="my_model_artifact:latest", + ) + + run.use_model( + name="my_project/my_model_artifact:v0", + ) + + run.use_model( + name="my_entity/my_project/my_model_artifact:", + ) + ``` + + Invalid usage + ```python + run.use_model( + name="my_entity/my_project/my_model_artifact", + ) + ``` + + Raises: + AssertionError: if model artifact 'name' is of a type that does not contain the substring 'model'. + + Returns: + path: (str) path to downloaded model artifact file(s). + """ + ... + +def link_model( + path: StrPath, + registered_model_name: str, + name: str | None = None, + aliases: list[str] | None = None, +) -> Artifact | None: + """Log a model artifact version and link it to a registered model in the model registry. + + The linked model version will be visible in the UI for the specified registered model. + + Steps: + - Check if 'name' model artifact has been logged. If so, use the artifact version that matches the files + located at 'path' or log a new version. Otherwise log files under 'path' as a new model artifact, 'name' + of type 'model'. + - Check if registered model with name 'registered_model_name' exists in the 'model-registry' project. + If not, create a new registered model with name 'registered_model_name'. + - Link version of model artifact 'name' to registered model, 'registered_model_name'. + - Attach aliases from 'aliases' list to the newly linked model artifact version. + + Args: + path: (str) A path to the contents of this model, + can be in the following forms: + - `/local/directory` + - `/local/directory/file.txt` + - `s3://bucket/path` + registered_model_name: (str) - the name of the registered model that the model is to be linked to. + A registered model is a collection of model versions linked to the model registry, typically representing a + team's specific ML Task. The entity that this registered model belongs to will be derived from the run + name: (str, optional) - the name of the model artifact that files in 'path' will be logged to. This will + default to the basename of the path prepended with the current run id if not specified. + aliases: (List[str], optional) - alias(es) that will only be applied on this linked artifact + inside the registered model. + The alias "latest" will always be applied to the latest version of an artifact that is linked. + + Examples: + ```python + run.link_model( + path="/local/directory", + registered_model_name="my_reg_model", + name="my_model_artifact", + aliases=["production"], + ) + ``` + + Invalid usage + ```python + run.link_model( + path="/local/directory", + registered_model_name="my_entity/my_project/my_reg_model", + name="my_model_artifact", + aliases=["production"], + ) + + run.link_model( + path="/local/directory", + registered_model_name="my_reg_model", + name="my_entity/my_project/my_model_artifact", + aliases=["production"], + ) + ``` + + Raises: + AssertionError: if registered_model_name is a path or + if model artifact 'name' is of a type that does not contain the substring 'model' + ValueError: if name has invalid special characters + + Returns: + The linked artifact if linking was successful, otherwise None. + """ + ... + +def plot_table( + vega_spec_name: str, + data_table: wandb.Table, + fields: dict[str, Any], + string_fields: dict[str, Any] | None = None, + split_table: bool = False, +) -> CustomChart: + """Creates a custom charts using a Vega-Lite specification and a `wandb.Table`. + + This function creates a custom chart based on a Vega-Lite specification and + a data table represented by a `wandb.Table` object. The specification needs + to be predefined and stored in the W&B backend. The function returns a custom + chart object that can be logged to W&B using `wandb.log()`. + + Args: + vega_spec_name (str): The name or identifier of the Vega-Lite spec + that defines the visualization structure. + data_table (wandb.Table): A `wandb.Table` object containing the data to be + visualized. + fields (dict[str, Any]): A mapping between the fields in the Vega-Lite spec and the + corresponding columns in the data table to be visualized. + string_fields (dict[str, Any] | None): A dictionary for providing values for any string constants + required by the custom visualization. + split_table (bool): Whether the table should be split into a separate section + in the W&B UI. If `True`, the table will be displayed in a section named + "Custom Chart Tables". Default is `False`. + + Returns: + CustomChart: A custom chart object that can be logged to W&B. To log the + chart, pass it to `wandb.log()`. + + Raises: + wandb.Error: If `data_table` is not a `wandb.Table` object. + """ + ... + +def watch( + models: torch.nn.Module | Sequence[torch.nn.Module], + criterion: torch.F | None = None, + log: Literal["gradients", "parameters", "all"] | None = "gradients", + log_freq: int = 1000, + idx: int | None = None, + log_graph: bool = False, +) -> None: + """Hooks into the given PyTorch model(s) to monitor gradients and the model's computational graph. + + This function can track parameters, gradients, or both during training. It should be + extended to support arbitrary machine learning models in the future. + + Args: + models (Union[torch.nn.Module, Sequence[torch.nn.Module]]): + A single model or a sequence of models to be monitored. + criterion (Optional[torch.F]): + The loss function being optimized (optional). + log (Optional[Literal["gradients", "parameters", "all"]]): + Specifies whether to log "gradients", "parameters", or "all". + Set to None to disable logging. (default="gradients") + log_freq (int): + Frequency (in batches) to log gradients and parameters. (default=1000) + idx (Optional[int]): + Index used when tracking multiple models with `wandb.watch`. (default=None) + log_graph (bool): + Whether to log the model's computational graph. (default=False) + + Raises: + ValueError: + If `wandb.init` has not been called or if any of the models are not instances + of `torch.nn.Module`. + """ + ... + +def unwatch( + models: torch.nn.Module | Sequence[torch.nn.Module] | None = None, +) -> None: + """Remove pytorch model topology, gradient and parameter hooks. + + Args: + models (torch.nn.Module | Sequence[torch.nn.Module]): + Optional list of pytorch models that have had watch called on them + """ + ... + +def restore( + name: str, + run_path: str | None = None, + replace: bool = False, + root: str | None = None, +) -> None | TextIO: + """Download the specified file from cloud storage. + + File is placed into the current directory or run directory. + By default, will only download the file if it doesn't already exist. + + Args: + name: the name of the file + run_path: optional path to a run to pull files from, i.e. `username/project_name/run_id` + if wandb.init has not been called, this is required. + replace: whether to download the file even if it already exists locally + root: the directory to download the file to. Defaults to the current + directory or the run directory if wandb.init was called. + + Returns: + None if it can't find the file, otherwise a file object open for reading + + Raises: + wandb.CommError: if we can't connect to the wandb backend + ValueError: if the file is not found or can't find run_path + """ + ... diff --git a/lib/python3.12/site-packages/wandb/__main__.py b/lib/python3.12/site-packages/wandb/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..e32b1ef386cb475c83b09fb527a9621a7370847a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/__main__.py @@ -0,0 +1,3 @@ +from wandb.cli import cli + +cli.cli(prog_name="python -m wandb") diff --git a/lib/python3.12/site-packages/wandb/_iterutils.py b/lib/python3.12/site-packages/wandb/_iterutils.py new file mode 100644 index 0000000000000000000000000000000000000000..a9249fa863f499aee5462e1cfc786dade8f36f53 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/_iterutils.py @@ -0,0 +1,65 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Iterable, TypeVar, Union, overload + +if TYPE_CHECKING: + T = TypeVar("T") + ClassInfo = Union[type[T], tuple[type[T], ...]] + + +@overload +def always_list(obj: Iterable[T], base_type: ClassInfo = ...) -> list[T]: ... +@overload +def always_list(obj: T, base_type: ClassInfo = ...) -> list[T]: ... +def always_list(obj: Any, base_type: Any = (str, bytes)) -> list[T]: + """Return a guaranteed list of objects from a single instance OR iterable of such objects. + + By default, assume the returned list should have string-like elements (i.e. `str`/`bytes`). + + Adapted from `more_itertools.always_iterable`, but simplified for internal use. See: + https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.always_iterable + """ + return [obj] if isinstance(obj, base_type) else list(obj) + + +def one( + iterable: Iterable[T], + too_short: type[Exception] | Exception | None = None, + too_long: type[Exception] | Exception | None = None, +) -> T: + """Return the only item in the iterable. + + Note: + This is intended **only** as an internal helper/convenience function, + and its implementation is directly adapted from `more_itertools.one`. + Users needing similar functionality are strongly encouraged to use + that library instead: + https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.one + + Args: + iterable: The iterable to get the only item from. + too_short: Custom exception to raise if the iterable has no items. + too_long: Custom exception to raise if the iterable has multiple items. + + Raises: + ValueError or `too_short`: If the iterable has no items. + ValueError or `too_long`: If the iterable has multiple items. + """ + # For a general iterable, avoid inadvertently iterating through all values, + # which may be costly or impossible (e.g. if infinite). Only check that: + + # ... the first item exists + it = iter(iterable) + try: + obj = next(it) + except StopIteration: + raise (too_short or ValueError("Expected 1 item in iterable, got 0")) from None + + # ...the second item doesn't + try: + _ = next(it) + except StopIteration: + return obj + raise ( + too_long or ValueError("Expected 1 item in iterable, got multiple") + ) from None diff --git a/lib/python3.12/site-packages/wandb/_pydantic/__init__.py b/lib/python3.12/site-packages/wandb/_pydantic/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aa418931a1e3ea2c2559f3c0685d57b54e6bdc5a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/_pydantic/__init__.py @@ -0,0 +1,30 @@ +"""Internal utilities for working with pydantic.""" + +from .base import ( + CompatBaseModel, + GQLBase, + GQLId, + SerializedToJson, + Typename, + ensure_json, +) +from .utils import IS_PYDANTIC_V2, from_json, gql_typename, pydantic_isinstance, to_json +from .v1_compat import AliasChoices, computed_field, field_validator, model_validator + +__all__ = [ + "IS_PYDANTIC_V2", + "CompatBaseModel", + "GQLBase", + "Typename", + "GQLId", + "SerializedToJson", + "AliasChoices", + "computed_field", + "field_validator", + "model_validator", + "pydantic_isinstance", + "to_json", + "from_json", + "ensure_json", + "gql_typename", +] diff --git a/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3f1361b8e17f64b2b9f93163a73599ce82ab9609 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57822c9c716bc0c0277b3a05ea9a66ddf32eb83b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..365823b0eee34e188b338fbc650d151404668b1b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/v1_compat.cpython-312.pyc b/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/v1_compat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7f547b1df217ae883bcb87779243c9df83504527 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/_pydantic/__pycache__/v1_compat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/_pydantic/base.py b/lib/python3.12/site-packages/wandb/_pydantic/base.py new file mode 100644 index 0000000000000000000000000000000000000000..1c80e1a92d208955af0d5566bc1461cdc0f041cc --- /dev/null +++ b/lib/python3.12/site-packages/wandb/_pydantic/base.py @@ -0,0 +1,128 @@ +"""Base classes and other customizations for generated pydantic types.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Callable, Literal, TypeVar + +from pydantic import BaseModel, ConfigDict, Field, Json, StrictStr +from typing_extensions import Annotated, TypedDict, Unpack, override + +from .utils import IS_PYDANTIC_V2, to_json +from .v1_compat import PydanticCompatMixin + +if TYPE_CHECKING: + from pydantic.main import IncEx + + +class ModelDumpKwargs(TypedDict, total=False): + """Shared keyword arguments for `BaseModel.model_{dump,dump_json}`.""" + + include: IncEx | None + exclude: IncEx | None + context: dict[str, Any] | None + by_alias: bool | None + exclude_unset: bool + exclude_defaults: bool + exclude_none: bool + round_trip: bool + warnings: bool | Literal["none", "warn", "error"] + fallback: Callable[[Any], Any] | None + serialize_as_any: bool + + +#: Custom overrides of default kwargs for `BaseModel.model_{dump,dump_json}`. +MODEL_DUMP_DEFAULTS = ModelDumpKwargs( + by_alias=True, # Always serialize with aliases (e.g. camelCase names) + round_trip=True, # Ensure serialized values remain valid inputs for deserialization +) + + +# v1-compatible base class for pydantic types. +class CompatBaseModel(PydanticCompatMixin, BaseModel): + __doc__ = None # Prevent subclasses from inheriting the BaseModel docstring + + +# Base class for all GraphQL-generated types. +# Omitted from docstring to avoid inclusion in generated docs. +class GQLBase(CompatBaseModel): + model_config = ConfigDict( + populate_by_name=True, # Discouraged in pydantic v2.11+, will be deprecated in v3 + validate_by_name=True, # Introduced in pydantic v2.11 + validate_by_alias=True, # Introduced in pydantic v2.11 + serialize_by_alias=True, # Introduced in pydantic v2.11 + validate_assignment=True, + validate_default=True, + use_attribute_docstrings=True, + from_attributes=True, + revalidate_instances="always", + protected_namespaces=(), # Some GraphQL fields may begin with "model_" + ) + + @override + def model_dump( + self, + *, + mode: Literal["json", "python"] | str = "json", # NOTE: changed default + **kwargs: Unpack[ModelDumpKwargs], + ) -> dict[str, Any]: + kwargs = {**MODEL_DUMP_DEFAULTS, **kwargs} + return super().model_dump(mode=mode, **kwargs) + + @override + def model_dump_json( + self, + *, + indent: int | None = None, + **kwargs: Unpack[ModelDumpKwargs], + ) -> str: + kwargs = {**MODEL_DUMP_DEFAULTS, **kwargs} + return super().model_dump_json(indent=indent, **kwargs) + + +# ------------------------------------------------------------------------------ +# Reusable annotations for field types +T = TypeVar("T") + +if IS_PYDANTIC_V2 or TYPE_CHECKING: + GQLId = Annotated[ + StrictStr, + Field(repr=False, frozen=True), + ] +else: + # FIXME: Find a way to fix this for pydantic v1, which doesn't like when + # `Field(...)` used in the field assignment AND `Annotated[...]`. + # This is a problem for codegen, which can currently outputs e.g. + # + # class MyModel(GQLBase): + # my_id: GQLId = Field(alias="myID") + # + GQLId = StrictStr # type: ignore[misc] + +Typename = Annotated[ + T, + Field(repr=False, frozen=True, alias="__typename"), +] + + +def ensure_json(v: Any) -> Any: + """In case the incoming value isn't serialized JSON, reserialize it. + + This lets us use `Json[...]` fields with values that are already deserialized. + """ + # NOTE: Assumes that the deserialized type is not itself a string. + # Revisit this if we need to support deserialized types that are str/bytes. + return v if isinstance(v, (str, bytes)) else to_json(v) + + +if IS_PYDANTIC_V2 or TYPE_CHECKING: + from pydantic import BeforeValidator, PlainSerializer + + SerializedToJson = Annotated[ + Json[T], + # Allow lenient instantiation/validation: incoming data may already be deserialized. + BeforeValidator(ensure_json), + PlainSerializer(to_json), + ] +else: + # FIXME: Restore, modify, or replace this later after ensuring pydantic v1 compatibility. + SerializedToJson = Json[T] # type: ignore[misc] diff --git a/lib/python3.12/site-packages/wandb/_pydantic/utils.py b/lib/python3.12/site-packages/wandb/_pydantic/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0418c5e8b03da7eba5477fbde2d1f8e193986950 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/_pydantic/utils.py @@ -0,0 +1,80 @@ +"""Internal utilities for working with Pydantic types and data.""" + +from __future__ import annotations + +import json +import sys +from contextlib import suppress +from typing import Any, Type + +import pydantic +from pydantic import BaseModel, ValidationError +from typing_extensions import TypeAlias + +PYTHON_VERSION = sys.version_info + +pydantic_major, *_ = pydantic.VERSION.split(".") +IS_PYDANTIC_V2: bool = int(pydantic_major) >= 2 + + +BaseModelType: TypeAlias = Type[BaseModel] + + +def gql_typename(cls: type[BaseModel]) -> str: + """Get the GraphQL typename for a Pydantic model.""" + if (field := cls.model_fields.get("typename__")) and (typename := field.default): + return typename + raise TypeError(f"Cannot extract GraphQL typename from: {cls.__qualname__!r}.") + + +if IS_PYDANTIC_V2: + import pydantic_core # pydantic_core is only installed by pydantic v2 + + def from_json(s: str) -> Any: + """Quickly deserialize a JSON string to a Python object.""" + return pydantic_core.from_json(s) + + def to_json(v: Any) -> str: + """Quickly serialize a (possibly Pydantic) object to a JSON string.""" + return pydantic_core.to_json(v, by_alias=True, round_trip=True).decode("utf-8") + + def pydantic_isinstance( + v: Any, classinfo: BaseModelType | tuple[BaseModelType, ...] + ) -> bool: + """Return True if the object could be parsed into the given Pydantic type. + + This is like a more lenient version of `isinstance()` for use with Pydantic. + In Pydantic v2, should be fast since the underlying implementation is in Rust, + and it may be preferable over `try:...except ValidationError:...`. + + See: https://docs.pydantic.dev/latest/api/pydantic_core/#pydantic_core.SchemaValidator.isinstance_python + """ + if isinstance(classinfo, tuple): + return any( + cls.__pydantic_validator__.isinstance_python(v) for cls in classinfo + ) + cls = classinfo + return cls.__pydantic_validator__.isinstance_python(v) + +else: + # Pydantic v1 fallback implementations. + # These may be noticeably slower, but their primary goal is to ensure + # compatibility with Pydantic v1 so long as we need to support it. + + from pydantic.json import pydantic_encoder # Only valid in pydantic v1 + + def from_json(s: str) -> Any: + return json.loads(s) + + def to_json(v: Any) -> str: + return json.dumps(v, default=pydantic_encoder) + + def pydantic_isinstance( + v: Any, classinfo: BaseModelType | tuple[BaseModelType, ...] + ) -> bool: + classes = classinfo if isinstance(classinfo, tuple) else (classinfo,) + for cls in classes: + with suppress(ValidationError): + cls.model_validate(v) + return True + return False diff --git a/lib/python3.12/site-packages/wandb/_pydantic/v1_compat.py b/lib/python3.12/site-packages/wandb/_pydantic/v1_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..8cd830a148bf86a531089620010afec049ae24fe --- /dev/null +++ b/lib/python3.12/site-packages/wandb/_pydantic/v1_compat.py @@ -0,0 +1,284 @@ +"""Provides partial support for compatibility with Pydantic v1.""" + +from __future__ import annotations + +import json +from functools import lru_cache +from inspect import signature +from operator import attrgetter +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Literal, overload + +import pydantic + +from .utils import IS_PYDANTIC_V2, to_json + +if TYPE_CHECKING: + from typing import Protocol + + class V1Model(Protocol): + # ------------------------------------------------------------------------------ + # NOTE: These aren't part of the original v1 BaseModel spec, but were added as + # internal helpers and are (re-)declared here to satisfy mypy checks. + @classmethod + def _dump_json_vals(cls, values: dict, by_alias: bool) -> dict: ... + + # ------------------------------------------------------------------------------ + # These methods are part of the original v1 BaseModel spec. + + __config__: ClassVar[type] + __fields__: ClassVar[dict[str, Any]] + __fields_set__: set[str] + + @classmethod + def update_forward_refs(cls, *args: Any, **kwargs: Any) -> None: ... + @classmethod + def construct(cls, *args: Any, **kwargs: Any) -> V1Model: ... + @classmethod + def parse_obj(cls, *args: Any, **kwargs: Any) -> V1Model: ... + @classmethod + def parse_raw(cls, *args: Any, **kwargs: Any) -> V1Model: ... + + def dict(self, **kwargs: Any) -> dict[str, Any]: ... + def json(self, **kwargs: Any) -> str: ... + def copy(self, **kwargs: Any) -> V1Model: ... + + +# Maps {v2 -> v1} model config keys that were renamed in v2. +# See: https://docs.pydantic.dev/latest/migration/#changes-to-config +_V1_CONFIG_KEYS = { + "populate_by_name": "allow_population_by_field_name", + "str_to_lower": "anystr_lower", + "str_strip_whitespace": "anystr_strip_whitespace", + "str_to_upper": "anystr_upper", + "ignored_types": "keep_untouched", + "str_max_length": "max_anystr_length", + "str_min_length": "min_anystr_length", + "from_attributes": "orm_mode", + "json_schema_extra": "schema_extra", + "validate_default": "validate_all", +} + + +def convert_v2_config(v2_config: dict[str, Any]) -> dict[str, Any]: + """Internal helper: Return a copy of the v2 ConfigDict with renamed v1 keys.""" + return {_V1_CONFIG_KEYS.get(k, k): v for k, v in v2_config.items()} + + +@lru_cache(maxsize=None) # Reduce repeat introspection via `signature()` +def allowed_arg_names(func: Callable) -> set[str]: + """Internal helper: Return the names of args accepted by the given function.""" + return set(signature(func).parameters) + + +# Pydantic BaseModels are defined with a custom metaclass, but its namespace +# has changed between pydantic versions. +# +# In v1, it can be imported as `from pydantic.main import ModelMetaclass` +# In v2, it's defined in an internal module so we avoid directly importing it. +PydanticModelMetaclass: type = type(pydantic.BaseModel) + + +class V1MixinMetaclass(PydanticModelMetaclass): + def __new__( + cls, + name: str, + bases: tuple[type, ...], + namespace: dict[str, Any], + **kwargs: Any, + ): + # In the class definition, convert the model config, if any: + # # BEFORE + # class MyModel(BaseModel): # v2 model with `ConfigDict` + # model_config = ConfigDict(populate_by_name=True) + # + # # AFTER + # class MyModel(BaseModel): # v1 model with inner `Config` class + # class Config: + # allow_population_by_field_name = True + if config_dict := namespace.pop("model_config", None): + namespace["Config"] = type("Config", (), convert_v2_config(config_dict)) + return super().__new__(cls, name, bases, namespace, **kwargs) + + @property + def model_fields(self) -> dict[str, Any]: + return self.__fields__ # type: ignore[deprecated] + + +# Mixin to ensure compatibility of Pydantic models if Pydantic v1 is detected. +# These are "best effort" implementations and cannot guarantee complete +# compatibility in v1 environments. +# +# Whenever possible, users should strongly prefer upgrading to Pydantic v2 to +# ensure full compatibility. +class V1Mixin(metaclass=V1MixinMetaclass): + # Internal compat helpers + @classmethod + def _dump_json_vals(cls, values: dict[str, Any], by_alias: bool) -> dict[str, Any]: + """Reserialize values from `Json`-typed fields after dumping the model to dict.""" + # Get the expected keys (after `.model_dump()`) for `Json`-typed fields. + # Note: In v1, `Json` fields have `ModelField.parse_json == True` + json_fields = (f for f in cls.__fields__.values() if f.parse_json) # type: ignore[deprecated] + get_key = attrgetter("alias" if by_alias else "name") + json_field_keys = set(map(get_key, json_fields)) + + return { + # Only serialize `Json` fields with non-null values. + k: to_json(v) if ((v is not None) and (k in json_field_keys)) else v + for k, v in values.items() + } + + # ------------------------------------------------------------------------------ + @classmethod + def __try_update_forward_refs__(cls: type[V1Model], **localns: Any) -> None: + if hasattr(sup := super(), "__try_update_forward_refs__"): + sup.__try_update_forward_refs__(**localns) + + @classmethod + def model_rebuild(cls, *args: Any, **kwargs: Any) -> None: + return cls.update_forward_refs(*args, **kwargs) + + @classmethod + def model_construct(cls, *args: Any, **kwargs: Any) -> V1Model: + return cls.construct(*args, **kwargs) + + @classmethod + def model_validate(cls, *args: Any, **kwargs: Any) -> V1Model: + return cls.parse_obj(*args, **kwargs) + + @classmethod + def model_validate_json(cls, *args: Any, **kwargs: Any) -> V1Model: + return cls.parse_raw(*args, **kwargs) + + def model_dump(self: V1Model, **kwargs: Any) -> dict[str, Any]: + # Pass only kwargs that are allowed in the V1 method. + allowed_keys = allowed_arg_names(self.dict) & kwargs.keys() + dict_ = self.dict(**{k: kwargs[k] for k in allowed_keys}) + + # Ugly hack: Try to serialize `Json` fields correctly when `round_trip=True` in pydantic v1 + if kwargs.get("round_trip", False): + by_alias: bool = kwargs.get("by_alias", False) + return self._dump_json_vals(dict_, by_alias=by_alias) + + return dict_ + + def model_dump_json(self: V1Model, **kwargs: Any) -> str: + # Pass only kwargs that are allowed in the V1 method. + allowed_keys = allowed_arg_names(self.json) & kwargs.keys() + json_ = self.json(**{k: kwargs[k] for k in allowed_keys}) + + # Ugly hack: Try to serialize `Json` fields correctly when `round_trip=True` in pydantic v1 + if kwargs.get("round_trip", False): + by_alias: bool = kwargs.get("by_alias", False) + dict_ = json.loads(json_) + return json.dumps(self._dump_json_vals(dict_, by_alias=by_alias)) + + return json_ + + def model_copy(self: V1Model, **kwargs: Any) -> V1Model: + # Pass only kwargs that are allowed in the V1 method. + allowed_keys = allowed_arg_names(self.copy) & kwargs.keys() + return self.copy(**{k: kwargs[k] for k in allowed_keys}) + + @property + def model_fields_set(self: V1Model) -> set[str]: + return self.__fields_set__ + + +# Placeholder. Pydantic v2 is already compatible with itself, so no need for extra mixins. +class V2Mixin: + pass + + +# Pick the mixin type based on the detected Pydantic version. +PydanticCompatMixin: type = V2Mixin if IS_PYDANTIC_V2 else V1Mixin + + +# ---------------------------------------------------------------------------- +# Decorators and other pydantic helpers +# ---------------------------------------------------------------------------- +if IS_PYDANTIC_V2: + field_validator = pydantic.field_validator + model_validator = pydantic.model_validator + AliasChoices = pydantic.AliasChoices + computed_field = pydantic.computed_field + +else: + # Redefines `@field_validator` with a v2-like signature + # to call `@validator` from v1 instead. + def field_validator( + field: str, + /, + *fields: str, + mode: Literal["before", "after", "wrap", "plain"] = "after", + check_fields: bool | None = None, + **_: Any, + ) -> Callable: + return pydantic.validator( # type: ignore[deprecated] + field, + *fields, + pre=(mode == "before"), + always=True, + check_fields=bool(check_fields), + allow_reuse=True, + ) + + # Redefines `@model_validator` with a v2-like signature + # to call `@root_validator` from v1 instead. + def model_validator( + *, + mode: Literal["before", "after", "wrap", "plain"], + **_: Any, + ) -> Callable: + if mode == "after": + # Patch the behavior for `@model_validator(mode="after")` in v1. This is + # necessarily complicated because: + # - `@model_validator(mode="after")` decorates an instance method in pydantic v2 + # - `@root_validator(pre=False)` always decorates a classmethod in pydantic v1 + def _decorator(v2_method: Callable) -> Any: + def v1_method( + cls: type[V1Model], values: dict[str, Any] + ) -> dict[str, Any]: + # Note: Since this is an "after" validator, the values should already be + # validated, so `.construct()` in v1 (`.model_construct()` in v2) + # should create a valid object to pass to the **original** decorated instance method. + validated = v2_method(cls.construct(**values)) + + # Pydantic v1 expects the validator to return a dict of {field_name -> value} + return { + name: getattr(validated, name) for name in validated.__fields__ + } + + return pydantic.root_validator(pre=False, allow_reuse=True)( # type: ignore[call-overload] + classmethod(v1_method) + ) + + return _decorator + else: + return pydantic.root_validator(pre=(mode == "before"), allow_reuse=True) # type: ignore[call-overload] + + @overload # type: ignore[no-redef] + def computed_field(func: Callable | property, /) -> property: ... + @overload + def computed_field( + func: None, /, **_: Any + ) -> Callable[[Callable | property], property]: ... + + def computed_field( + func: Callable | property | None = None, /, **_: Any + ) -> property | Callable[[Callable | property], property]: + """Compatibility wrapper for Pydantic v2's `computed_field` in v1.""" + + def always_property(f: Callable | property) -> property: + # Convert the method to a property only if needed + return f if isinstance(f, property) else property(f) + + # Handle both decorator styles + return always_property if (func is None) else always_property(func) + + class AliasChoices: # type: ignore [no-redef] + """Placeholder class for Pydantic v2's AliasChoices for partial v1 compatibility.""" + + aliases: list[str] + + def __init__(self, *aliases: str): + self.aliases = list(aliases) diff --git a/lib/python3.12/site-packages/wandb/data_types.py b/lib/python3.12/site-packages/wandb/data_types.py new file mode 100644 index 0000000000000000000000000000000000000000..f28e283d8436d1b366cb59ba66d4171df88bcda7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/data_types.py @@ -0,0 +1,66 @@ +"""This module defines data types for logging rich, interactive visualizations to W&B. + +Data types include common media types, like images, audio, and videos, +flexible containers for information, like tables and HTML, and more. + +For more on logging media, see [our guide](https://docs.wandb.com/guides/track/log/media) + +For more on logging structured data for interactive dataset and model analysis, +see [our guide to W&B Tables](https://docs.wandb.com/guides/models/tables/). + +All of these special data types are subclasses of WBValue. All the data types +serialize to JSON, since that is what wandb uses to save the objects locally +and upload them to the W&B server. +""" + +from .sdk.data_types.audio import Audio +from .sdk.data_types.base_types.media import BatchableMedia, Media +from .sdk.data_types.base_types.wb_value import WBValue +from .sdk.data_types.bokeh import Bokeh +from .sdk.data_types.graph import Graph, Node +from .sdk.data_types.helper_types.bounding_boxes_2d import BoundingBoxes2D +from .sdk.data_types.helper_types.classes import Classes +from .sdk.data_types.helper_types.image_mask import ImageMask +from .sdk.data_types.histogram import Histogram +from .sdk.data_types.html import Html +from .sdk.data_types.image import Image +from .sdk.data_types.molecule import Molecule +from .sdk.data_types.object_3d import Object3D, box3d +from .sdk.data_types.plotly import Plotly +from .sdk.data_types.saved_model import _SavedModel +from .sdk.data_types.table import JoinedTable, PartitionedTable, Table +from .sdk.data_types.trace_tree import WBTraceTree +from .sdk.data_types.video import Video + +# Note: we are importing everything from the sdk/data_types to maintain a namespace for now. +# Once we fully type this file and move it all into sdk, then we will need to clean up the +# other internal imports + +__all__ = [ + # Untyped Exports + "Audio", + "Table", + "JoinedTable", + "PartitionedTable", + "Bokeh", + "Node", + "Graph", + # Typed Exports + "Histogram", + "Html", + "Image", + "Molecule", + "box3d", + "Object3D", + "Plotly", + "Video", + "WBTraceTree", + "_SavedModel", + "WBValue", + "Media", + "BatchableMedia", + # Typed Legacy Exports (I'd like to remove these) + "ImageMask", + "BoundingBoxes2D", + "Classes", +] diff --git a/lib/python3.12/site-packages/wandb/env.py b/lib/python3.12/site-packages/wandb/env.py new file mode 100644 index 0000000000000000000000000000000000000000..96d2bf822f3d3b58e4e43e8a3e7b47aabce6dabb --- /dev/null +++ b/lib/python3.12/site-packages/wandb/env.py @@ -0,0 +1,527 @@ +"""All of W&B's environment variables. + +Getters and putters for all of them should go here. That way it'll be easier to +avoid typos with names and be consistent about environment variables' semantics. + +Environment variables are not the authoritative source for these values in many +cases. +""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path +from typing import MutableMapping + +import platformdirs + +CONFIG_PATHS = "WANDB_CONFIG_PATHS" +SWEEP_PARAM_PATH = "WANDB_SWEEP_PARAM_PATH" +SHOW_RUN = "WANDB_SHOW_RUN" +DEBUG = "WANDB_DEBUG" +SILENT = "WANDB_SILENT" +QUIET = "WANDB_QUIET" +INITED = "WANDB_INITED" +DIR = "WANDB_DIR" +# Deprecate DESCRIPTION in a future release +DESCRIPTION = "WANDB_DESCRIPTION" +NAME = "WANDB_NAME" +NOTEBOOK_NAME = "WANDB_NOTEBOOK_NAME" +NOTES = "WANDB_NOTES" +USERNAME = "WANDB_USERNAME" +USER_EMAIL = "WANDB_USER_EMAIL" +PROJECT = "WANDB_PROJECT" +ENTITY = "WANDB_ENTITY" +ORGANIZATION = "WANDB_ORGANIZATION" +BASE_URL = "WANDB_BASE_URL" +APP_URL = "WANDB_APP_URL" +PROGRAM = "WANDB_PROGRAM" +ARGS = "WANDB_ARGS" +MODE = "WANDB_MODE" +START_METHOD = "WANDB_START_METHOD" +RESUME = "WANDB_RESUME" +RUN_ID = "WANDB_RUN_ID" +RUN_STORAGE_ID = "WANDB_RUN_STORAGE_ID" +RUN_GROUP = "WANDB_RUN_GROUP" +RUN_DIR = "WANDB_RUN_DIR" +SWEEP_ID = "WANDB_SWEEP_ID" +HTTP_TIMEOUT = "WANDB_HTTP_TIMEOUT" +FILE_PUSHER_TIMEOUT = "WANDB_FILE_PUSHER_TIMEOUT" +API_KEY = "WANDB_API_KEY" +IDENTITY_TOKEN_FILE = "WANDB_IDENTITY_TOKEN_FILE" +CREDENTIALS_FILE = "WANDB_CREDENTIALS_FILE" +JOB_TYPE = "WANDB_JOB_TYPE" +DISABLE_CODE = "WANDB_DISABLE_CODE" +DISABLE_GIT = "WANDB_DISABLE_GIT" +GIT_ROOT = "WANDB_GIT_ROOT" +SAVE_CODE = "WANDB_SAVE_CODE" +TAGS = "WANDB_TAGS" +IGNORE = "WANDB_IGNORE_GLOBS" +ERROR_REPORTING = "WANDB_ERROR_REPORTING" +CORE_DEBUG = "WANDB_CORE_DEBUG" +DOCKER = "WANDB_DOCKER" +AGENT_REPORT_INTERVAL = "WANDB_AGENT_REPORT_INTERVAL" +AGENT_KILL_DELAY = "WANDB_AGENT_KILL_DELAY" +AGENT_DISABLE_FLAPPING = "WANDB_AGENT_DISABLE_FLAPPING" +AGENT_MAX_INITIAL_FAILURES = "WANDB_AGENT_MAX_INITIAL_FAILURES" +CRASH_NOSYNC_TIME = "WANDB_CRASH_NOSYNC_TIME" +MAGIC = "WANDB_MAGIC" +HOST = "WANDB_HOST" +ANONYMOUS = "WANDB_ANONYMOUS" +JUPYTER = "WANDB_JUPYTER" +CONFIG_DIR = "WANDB_CONFIG_DIR" +DATA_DIR = "WANDB_DATA_DIR" +ARTIFACT_DIR = "WANDB_ARTIFACT_DIR" +ARTIFACT_FETCH_FILE_URL_BATCH_SIZE = "WANDB_ARTIFACT_FETCH_FILE_URL_BATCH_SIZE" +CACHE_DIR = "WANDB_CACHE_DIR" +DISABLE_SSL = "WANDB_INSECURE_DISABLE_SSL" +SERVICE = "WANDB_SERVICE" +SENTRY_DSN = "WANDB_SENTRY_DSN" +INIT_TIMEOUT = "WANDB_INIT_TIMEOUT" +GIT_COMMIT = "WANDB_GIT_COMMIT" +GIT_REMOTE_URL = "WANDB_GIT_REMOTE_URL" +_EXECUTABLE = "WANDB_X_EXECUTABLE" +LAUNCH_QUEUE_NAME = "WANDB_LAUNCH_QUEUE_NAME" +LAUNCH_QUEUE_ENTITY = "WANDB_LAUNCH_QUEUE_ENTITY" +LAUNCH_TRACE_ID = "WANDB_LAUNCH_TRACE_ID" +ENABLE_DCGM_PROFILING = "WANDB_ENABLE_DCGM_PROFILING" + +# For testing, to be removed in future version +USE_V1_ARTIFACTS = "_WANDB_USE_V1_ARTIFACTS" + + +def immutable_keys() -> list[str]: + """These are env keys that shouldn't change within a single process. + + We use this to maintain certain values between multiple calls to wandb.init within a single process. + """ + return [ + DIR, + ENTITY, + PROJECT, + API_KEY, + IGNORE, + DISABLE_CODE, + DISABLE_GIT, + DOCKER, + MODE, + BASE_URL, + ERROR_REPORTING, + CRASH_NOSYNC_TIME, + MAGIC, + USERNAME, + USER_EMAIL, + DIR, + SILENT, + CONFIG_PATHS, + ANONYMOUS, + RUN_GROUP, + JOB_TYPE, + TAGS, + RESUME, + AGENT_REPORT_INTERVAL, + HTTP_TIMEOUT, + HOST, + DATA_DIR, + ARTIFACT_DIR, + ARTIFACT_FETCH_FILE_URL_BATCH_SIZE, + CACHE_DIR, + USE_V1_ARTIFACTS, + DISABLE_SSL, + IDENTITY_TOKEN_FILE, + CREDENTIALS_FILE, + ] + + +def _env_as_bool( + var: str, default: str | None = None, env: MutableMapping | None = None +) -> bool: + if env is None: + env = os.environ + val = env.get(var, default) + if not isinstance(val, str): + return False + try: + return strtobool(val) + except ValueError: + return False + + +def is_debug(default: str | None = None, env: MutableMapping | None = None) -> bool: + return _env_as_bool(DEBUG, default=default, env=env) + + +def is_offline(env: MutableMapping | None = None) -> bool: + if env is None: + env = os.environ + return env.get(MODE) == "offline" + + +def error_reporting_enabled() -> bool: + return _env_as_bool(ERROR_REPORTING, default="True") + + +def core_debug(default: str | None = None) -> bool: + return _env_as_bool(CORE_DEBUG, default=default) or is_debug() + + +def ssl_disabled() -> bool: + return _env_as_bool(DISABLE_SSL, default="False") + + +def dcgm_profiling_enabled() -> bool: + """Checks whether collecting profiling metrics for Nvidia GPUs using DCGM is requested. + + Note: Enabling this feature can lead to increased resource usage + compared to standard monitoring. + Requires the `nvidia-dcgm` service to be running on the machine. + """ + return _env_as_bool(ENABLE_DCGM_PROFILING, default="False") + + +def get_error_reporting( + default: bool | str = True, + env: MutableMapping | None = None, +) -> bool | str: + if env is None: + env = os.environ + + return env.get(ERROR_REPORTING, default) + + +def get_run( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(RUN_ID, default) + + +def get_args( + default: list[str] | None = None, env: MutableMapping | None = None +) -> list[str] | None: + if env is None: + env = os.environ + if env.get(ARGS): + try: + return json.loads(env.get(ARGS, "[]")) # type: ignore + except ValueError: + return None + else: + return default or sys.argv[1:] + + +def get_docker( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(DOCKER, default) + + +def get_http_timeout(default: int = 20, env: MutableMapping | None = None) -> int: + if env is None: + env = os.environ + + return int(env.get(HTTP_TIMEOUT, default)) + + +def get_file_pusher_timeout( + default: int | None = None, + env: MutableMapping | None = None, +) -> int | None: + if env is None: + env = os.environ + + timeout = env.get(FILE_PUSHER_TIMEOUT, default) + return int(timeout) if timeout else None + + +def get_ignore( + default: list[str] | None = None, env: MutableMapping | None = None +) -> list[str] | None: + if env is None: + env = os.environ + ignore = env.get(IGNORE) + if ignore is not None: + return ignore.split(",") + else: + return default + + +def get_project( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(PROJECT, default) + + +def get_username( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(USERNAME, default) + + +def get_user_email( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(USER_EMAIL, default) + + +def get_entity( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(ENTITY, default) + + +def get_organization( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(ORGANIZATION, default) + + +def get_base_url( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(BASE_URL, default) + + +def get_app_url( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(APP_URL, default) + + +def get_show_run(default: str | None = None, env: MutableMapping | None = None) -> bool: + if env is None: + env = os.environ + + return bool(env.get(SHOW_RUN, default)) + + +def get_description( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + + return env.get(DESCRIPTION, default) + + +def get_tags(default: str = "", env: MutableMapping | None = None) -> list[str]: + if env is None: + env = os.environ + + return [tag for tag in env.get(TAGS, default).split(",") if tag] + + +def get_dir( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + return env.get(DIR, default) + + +def get_config_paths( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + return env.get(CONFIG_PATHS, default) + + +def get_agent_report_interval( + default: str | None = None, env: MutableMapping | None = None +) -> int | None: + if env is None: + env = os.environ + val = env.get(AGENT_REPORT_INTERVAL, default) + try: + val = int(val) # type: ignore + except ValueError: + val = None # silently ignore env format errors, caller should handle. + return val + + +def get_agent_kill_delay( + default: str | None = None, env: MutableMapping | None = None +) -> int | None: + if env is None: + env = os.environ + val = env.get(AGENT_KILL_DELAY, default) + try: + val = int(val) # type: ignore + except ValueError: + val = None # silently ignore env format errors, caller should handle. + return val + + +def get_crash_nosync_time( + default: str | None = None, env: MutableMapping | None = None +) -> int | None: + if env is None: + env = os.environ + val = env.get(CRASH_NOSYNC_TIME, default) + try: + val = int(val) # type: ignore + except ValueError: + val = None # silently ignore env format errors, caller should handle. + return val + + +def get_magic( + default: str | None = None, env: MutableMapping | None = None +) -> str | None: + if env is None: + env = os.environ + val = env.get(MAGIC, default) + return val + + +def get_data_dir(env: MutableMapping | None = None) -> str: + default_dir = platformdirs.user_data_dir("wandb") + if env is None: + env = os.environ + val = env.get(DATA_DIR, default_dir) + return val + + +def get_artifact_dir(env: MutableMapping | None = None) -> str: + default_dir = os.path.join(".", "artifacts") + if env is None: + env = os.environ + val = env.get(ARTIFACT_DIR, default_dir) + return os.path.abspath(str(val)) + + +def get_artifact_fetch_file_url_batch_size(env: MutableMapping | None = None) -> int: + default_batch_size = 5000 + if env is None: + env = os.environ + val = int(env.get(ARTIFACT_FETCH_FILE_URL_BATCH_SIZE, default_batch_size)) + return val + + +def get_cache_dir(env: MutableMapping | None = None) -> Path: + env = env or os.environ + return Path(env.get(CACHE_DIR, platformdirs.user_cache_dir("wandb"))) + + +def get_use_v1_artifacts(env: MutableMapping | None = None) -> bool: + if env is None: + env = os.environ + val = bool(env.get(USE_V1_ARTIFACTS, False)) + return val + + +def get_agent_max_initial_failures( + default: int | None = None, env: MutableMapping | None = None +) -> int | None: + if env is None: + env = os.environ + val = env.get(AGENT_MAX_INITIAL_FAILURES, default) + try: + val = int(val) # type: ignore + except ValueError: + val = default + return val + + +def set_entity(value: str, env: MutableMapping | None = None) -> None: + if env is None: + env = os.environ + env[ENTITY] = value + + +def set_project(value: str, env: MutableMapping | None = None) -> None: + if env is None: + env = os.environ + env[PROJECT] = value or "uncategorized" + + +def should_save_code() -> bool: + save_code = _env_as_bool(SAVE_CODE, default="False") + code_disabled = _env_as_bool(DISABLE_CODE, default="False") + return save_code and not code_disabled + + +def disable_git(env: MutableMapping | None = None) -> bool: + if env is None: + env = os.environ + val = env.get(DISABLE_GIT, default="False") + if isinstance(val, str): + val = False if val.lower() == "false" else True + return val + + +def get_launch_queue_name(env: MutableMapping | None = None) -> str | None: + if env is None: + env = os.environ + val = env.get(LAUNCH_QUEUE_NAME, None) + return val + + +def get_launch_queue_entity(env: MutableMapping | None = None) -> str | None: + if env is None: + env = os.environ + val = env.get(LAUNCH_QUEUE_ENTITY, None) + return val + + +def get_launch_trace_id(env: MutableMapping | None = None) -> str | None: + if env is None: + env = os.environ + val = env.get(LAUNCH_TRACE_ID, None) + return val + + +def get_credentials_file(default: str, env: MutableMapping | None = None) -> Path: + """Retrieve the path for the credentials file used to save access tokens. + + The credentials file path can be set via an environment variable, otherwise + the default path is used. + """ + if env is None: + env = os.environ + credentials_file = env.get(CREDENTIALS_FILE, default) + return Path(credentials_file) + + +def strtobool(val: str) -> bool: + """Convert a string representation of truth to true or false. + + Copied from distutils. distutils was removed in Python 3.12. + """ + val = val.lower() + + if val in ("y", "yes", "t", "true", "on", "1"): + return True + elif val in ("n", "no", "f", "false", "off", "0"): + return False + else: + raise ValueError(f"invalid truth value {val!r}") diff --git a/lib/python3.12/site-packages/wandb/jupyter.py b/lib/python3.12/site-packages/wandb/jupyter.py new file mode 100644 index 0000000000000000000000000000000000000000..4b5141c4dff05305b937be6b10c4315d7e618c91 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/jupyter.py @@ -0,0 +1,538 @@ +from __future__ import annotations + +import json +import logging +import os +import re +import shutil +import sys +import traceback +from base64 import b64encode +from typing import Any + +import IPython +import IPython.display +import requests +from IPython.core.magic import Magics, line_cell_magic, magics_class +from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring +from requests.compat import urljoin + +import wandb +import wandb.util +from wandb.sdk import wandb_run, wandb_setup +from wandb.sdk.lib import filesystem + +logger = logging.getLogger(__name__) + + +def display_if_magic_is_used(run: wandb_run.Run) -> bool: + """Display a run's page if the cell has the %%wandb cell magic. + + Args: + run: The run to display. + + Returns: + Whether the %%wandb cell magic was present. + """ + if not _current_cell_wandb_magic: + return False + + _current_cell_wandb_magic.display_if_allowed(run) + return True + + +class _WandbCellMagicState: + """State for a cell with the %%wandb cell magic.""" + + def __init__(self, *, height: int) -> None: + """Initializes the %%wandb cell magic state. + + Args: + height: The desired height for displayed iframes. + """ + self._height = height + self._already_displayed = False + + def display_if_allowed(self, run: wandb_run.Run) -> None: + """Display a run's iframe if one is not already displayed. + + Args: + run: The run to display. + """ + if self._already_displayed: + return + self._already_displayed = True + + _display_wandb_run(run, height=self._height) + + +_current_cell_wandb_magic: _WandbCellMagicState | None = None + + +def _display_by_wandb_path(path: str, *, height: int) -> None: + """Display a wandb object (usually in an iframe) given its URI. + + Args: + path: A path to a run, sweep, project, report, etc. + height: Height of the iframe in pixels. + """ + api = wandb.Api() + + try: + obj = api.from_path(path) + + IPython.display.display_html( + obj.to_html(height=height), + raw=True, + ) + except wandb.Error: + traceback.print_exc() + IPython.display.display_html( + f"Path {path!r} does not refer to a W&B object you can access.", + raw=True, + ) + + +def _display_wandb_run(run: wandb_run.Run, *, height: int) -> None: + """Display a run (usually in an iframe). + + Args: + run: The run to display. + height: Height of the iframe in pixels. + """ + IPython.display.display_html( + run.to_html(height=height), + raw=True, + ) + + +@magics_class +class WandBMagics(Magics): + def __init__(self, shell): + super().__init__(shell) + + @magic_arguments() + @argument( + "path", + default=None, + nargs="?", + help="The path to a resource you want to display.", + ) + @argument( + "-h", + "--height", + default=420, + type=int, + help="The height of the iframe in pixels.", + ) + @line_cell_magic + def wandb(self, line: str, cell: str | None = None) -> None: + """Display wandb resources in Jupyter. + + This can be used as a line magic: + + %wandb USERNAME/PROJECT/runs/RUN_ID + + Or as a cell magic: + + %%wandb -h 1024 + with wandb.init() as run: + run.log({"loss": 1}) + """ + global _current_cell_wandb_magic + + args = parse_argstring(self.wandb, line) + path: str | None = args.path + height: int = args.height + + if path: + _display_by_wandb_path(path, height=height) + displayed = True + elif run := wandb_setup.singleton().most_recent_active_run: + _display_wandb_run(run, height=height) + displayed = True + else: + displayed = False + + # If this is being used as a line magic ("%wandb"), we are done. + # When used as a cell magic ("%%wandb"), we must run the cell. + if cell is None: + return + + if not displayed: + _current_cell_wandb_magic = _WandbCellMagicState(height=height) + + try: + IPython.get_ipython().run_cell(cell) + finally: + _current_cell_wandb_magic = None + + +def notebook_metadata_from_jupyter_servers_and_kernel_id(): + servers, kernel_id = jupyter_servers_and_kernel_id() + for s in servers: + if s.get("password"): + raise ValueError("Can't query password protected kernel") + res = requests.get( + urljoin(s["url"], "api/sessions"), params={"token": s.get("token", "")} + ).json() + for nn in res: + if isinstance(nn, dict) and nn.get("kernel") and "notebook" in nn: + if nn["kernel"]["id"] == kernel_id: + return { + "root": s.get("root_dir", s.get("notebook_dir", os.getcwd())), + "path": nn["notebook"]["path"], + "name": nn["notebook"]["name"], + } + + if not kernel_id: + return None + + # Built-in notebook server in VS Code + try: + from IPython import get_ipython + + ipython = get_ipython() + notebook_path = ipython.kernel.shell.user_ns.get("__vsc_ipynb_file__") + if notebook_path: + return { + "root": os.path.dirname(notebook_path), + "path": notebook_path, + "name": os.path.basename(notebook_path), + } + except Exception: + return None + + +def notebook_metadata(silent: bool) -> dict[str, str]: + """Attempt to query jupyter for the path and name of the notebook file. + + This can handle different jupyter environments, specifically: + + 1. Colab + 2. Kaggle + 3. JupyterLab + 4. Notebooks + 5. Other? + """ + error_message = ( + "Failed to detect the name of this notebook. You can set it manually" + " with the WANDB_NOTEBOOK_NAME environment variable to enable code" + " saving." + ) + try: + jupyter_metadata = notebook_metadata_from_jupyter_servers_and_kernel_id() + + # Colab: + # request the most recent contents + ipynb = attempt_colab_load_ipynb() + if ipynb is not None and jupyter_metadata is not None: + return { + "root": "/content", + "path": jupyter_metadata["path"], + "name": jupyter_metadata["name"], + } + + # Kaggle: + if wandb.util._is_kaggle(): + # request the most recent contents + ipynb = attempt_kaggle_load_ipynb() + if ipynb: + return { + "root": "/kaggle/working", + "path": ipynb["metadata"]["name"], + "name": ipynb["metadata"]["name"], + } + + if jupyter_metadata: + return jupyter_metadata + except Exception: + logger.exception(error_message) + + wandb.termerror(error_message) + return {} + + +def jupyter_servers_and_kernel_id(): + """Return a list of servers and the current kernel_id. + + Used to query for the name of the notebook. + """ + try: + import ipykernel # type: ignore + + kernel_id = re.search( + "kernel-(.*).json", ipykernel.connect.get_connection_file() + ).group(1) + # We're either in jupyterlab or a notebook, lets prefer the newer jupyter_server package + serverapp = wandb.util.get_module("jupyter_server.serverapp") + notebookapp = wandb.util.get_module("notebook.notebookapp") + servers = [] + if serverapp is not None: + servers.extend(list(serverapp.list_running_servers())) + if notebookapp is not None: + servers.extend(list(notebookapp.list_running_servers())) + except (AttributeError, ValueError, ImportError): + return [], None + + return servers, kernel_id + + +def attempt_colab_load_ipynb(): + colab = wandb.util.get_module("google.colab") + if colab: + # This isn't thread safe, never call in a thread + response = colab._message.blocking_request("get_ipynb", timeout_sec=5) + if response: + return response["ipynb"] + + +def attempt_kaggle_load_ipynb(): + kaggle = wandb.util.get_module("kaggle_session") + if not kaggle: + return None + + try: + client = kaggle.UserSessionClient() + parsed = json.loads(client.get_exportable_ipynb()["source"]) + # TODO: couldn't find a way to get the name of the notebook... + parsed["metadata"]["name"] = "kaggle.ipynb" + except Exception: + wandb.termerror("Unable to load kaggle notebook.") + logger.exception("Unable to load kaggle notebook.") + return None + + return parsed + + +def attempt_colab_login( + app_url: str, + referrer: str | None = None, +): + """This renders an iframe to wandb in the hopes it posts back an api key.""" + from google.colab import output # type: ignore + from google.colab._message import MessageError # type: ignore + from IPython import display + + display.display( + display.Javascript( + """ + window._wandbApiKey = new Promise((resolve, reject) => {{ + function loadScript(url) {{ + return new Promise(function(resolve, reject) {{ + let newScript = document.createElement("script"); + newScript.onerror = reject; + newScript.onload = resolve; + document.body.appendChild(newScript); + newScript.src = url; + }}); + }} + loadScript("https://cdn.jsdelivr.net/npm/postmate/build/postmate.min.js").then(() => {{ + const iframe = document.createElement('iframe') + iframe.style.cssText = "width:0;height:0;border:none" + document.body.appendChild(iframe) + const handshake = new Postmate({{ + container: iframe, + url: '{}/authorize{}' + }}); + const timeout = setTimeout(() => reject("Couldn't auto authenticate"), 5000) + handshake.then(function(child) {{ + child.on('authorize', data => {{ + clearTimeout(timeout) + resolve(data) + }}); + }}); + }}) + }}); + """.format( + app_url.replace("http:", "https:"), + f"?ref={referrer}" if referrer else "", + ) + ) + ) + try: + return output.eval_js("_wandbApiKey") + except MessageError: + return None + + +class Notebook: + def __init__(self, settings: wandb.Settings) -> None: + self.outputs: dict[int, Any] = {} + self.settings = settings + self.shell = IPython.get_ipython() + + def save_display(self, exc_count, data_with_metadata): + self.outputs[exc_count] = self.outputs.get(exc_count, []) + + # byte values such as images need to be encoded in base64 + # otherwise nbformat.v4.new_output will throw a NotebookValidationError + data = data_with_metadata["data"] + b64_data = {} + for key in data: + val = data[key] + if isinstance(val, bytes): + b64_data[key] = b64encode(val).decode("utf-8") + else: + b64_data[key] = val + + self.outputs[exc_count].append( + {"data": b64_data, "metadata": data_with_metadata["metadata"]} + ) + + def probe_ipynb(self): + """Return notebook as dict or None.""" + relpath = self.settings.x_jupyter_path + if relpath: + if os.path.exists(relpath): + with open(relpath) as json_file: + data = json.load(json_file) + return data + + colab_ipynb = attempt_colab_load_ipynb() + if colab_ipynb: + return colab_ipynb + + kaggle_ipynb = attempt_kaggle_load_ipynb() + if kaggle_ipynb and len(kaggle_ipynb["cells"]) > 0: + return kaggle_ipynb + + return + + def save_ipynb(self) -> bool: + if not self.settings.save_code: + logger.info("not saving jupyter notebook") + return False + ret = False + try: + ret = self._save_ipynb() + except Exception: + wandb.termerror("Failed to save notebook.") + logger.exception("Problem saving notebook.") + return ret + + def _save_ipynb(self) -> bool: + relpath = self.settings.x_jupyter_path + logger.info("looking for notebook: %s", relpath) + if relpath: + if os.path.exists(relpath): + shutil.copy( + relpath, + os.path.join( + self.settings._tmp_code_dir, os.path.basename(relpath) + ), + ) + return True + + # TODO: likely only save if the code has changed + colab_ipynb = attempt_colab_load_ipynb() + if colab_ipynb: + try: + jupyter_metadata = ( + notebook_metadata_from_jupyter_servers_and_kernel_id() + ) + nb_name = jupyter_metadata["name"] + except Exception: + nb_name = "colab.ipynb" + if not nb_name.endswith(".ipynb"): + nb_name += ".ipynb" + with open( + os.path.join( + self.settings._tmp_code_dir, + nb_name, + ), + "w", + encoding="utf-8", + ) as f: + f.write(json.dumps(colab_ipynb)) + return True + + kaggle_ipynb = attempt_kaggle_load_ipynb() + if kaggle_ipynb and len(kaggle_ipynb["cells"]) > 0: + with open( + os.path.join( + self.settings._tmp_code_dir, kaggle_ipynb["metadata"]["name"] + ), + "w", + encoding="utf-8", + ) as f: + f.write(json.dumps(kaggle_ipynb)) + return True + + return False + + def save_history(self, run: wandb_run.Run): + """This saves all cell executions in the current session as a new notebook.""" + try: + from nbformat import v4, validator, write # type: ignore + except ImportError: + wandb.termerror( + "The nbformat package was not found." + " It is required to save notebook history." + ) + return + # TODO: some tests didn't patch ipython properly? + if self.shell is None: + return + cells = [] + hist = list(self.shell.history_manager.get_range(output=True)) + if len(hist) <= 1 or not self.settings.save_code: + logger.info("not saving jupyter history") + return + try: + for _, execution_count, exc in hist: + if exc[1]: + # TODO: capture stderr? + outputs = [ + v4.new_output(output_type="stream", name="stdout", text=exc[1]) + ] + else: + outputs = [] + if self.outputs.get(execution_count): + for out in self.outputs[execution_count]: + outputs.append( + v4.new_output( + output_type="display_data", + data=out["data"], + metadata=out["metadata"] or {}, + ) + ) + cells.append( + v4.new_code_cell( + execution_count=execution_count, source=exc[0], outputs=outputs + ) + ) + if hasattr(self.shell, "kernel"): + language_info = self.shell.kernel.language_info + else: + language_info = {"name": "python", "version": sys.version} + logger.info("saving %i cells to _session_history.ipynb", len(cells)) + nb = v4.new_notebook( + cells=cells, + metadata={ + "kernelspec": { + "display_name": f"Python {sys.version_info[0]}", + "name": f"python{sys.version_info[0]}", + "language": "python", + }, + "language_info": language_info, + }, + ) + state_path = os.path.join("code", "_session_history.ipynb") + run._set_config_wandb("session_history", state_path) + filesystem.mkdir_exists_ok(os.path.join(self.settings.files_dir, "code")) + with open( + os.path.join(self.settings._tmp_code_dir, "_session_history.ipynb"), + "w", + encoding="utf-8", + ) as f: + write(nb, f, version=4) + with open( + os.path.join(self.settings.files_dir, state_path), + "w", + encoding="utf-8", + ) as f: + write(nb, f, version=4) + except (OSError, validator.NotebookValidationError): + wandb.termerror("Unable to save notebook session history.") + logger.exception("Unable to save notebook session history.") diff --git a/lib/python3.12/site-packages/wandb/proto/__init__.py b/lib/python3.12/site-packages/wandb/proto/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e99e267446203683a15f97d6471382cd98eb3e18 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_base_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_base_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b86cfb3f441581e283bdbab96edfe712289e9d89 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_base_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_deprecated.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_deprecated.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9ab00a0422d06877ac879512c528a7a53075e4f2 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_deprecated.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_generate_deprecated.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_generate_deprecated.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf65ba9e5cf8372c1f458579dfc9e965f3771652 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_generate_deprecated.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_generate_proto.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_generate_proto.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..eb370b88823df32deb9900e437dffe777e37f2ca Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_generate_proto.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_internal_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_internal_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b35eaaec8dbfc56b1c8ebac54482e8afd998d2ab Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_internal_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_server_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_server_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57f5fed4eda8c8a2cc8553fc021f61f7289eddba Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_server_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_settings_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_settings_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..78231142cc89dc9f0d169290fcd6910ae49db11d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_settings_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_telemetry_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_telemetry_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..55564b8c43a9caeca03e1635149e9779eefaff7a Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/__pycache__/wandb_telemetry_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v3/__init__.py b/lib/python3.12/site-packages/wandb/proto/v3/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..53dc317c7ece8f014d742685dae263272c50559e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_base_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_base_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e0f7f10013cd2d9ebae48758261dedc54c2cea5 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_base_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_internal_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_internal_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6dc3b96a7334c3255442eb1f9ccfcf8e22e0600f Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_internal_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_server_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_server_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..80a611ded32dcca2ea02f554b3a81de32357d435 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_server_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_settings_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_settings_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4fa24093e23ae69c0e8286f08ef413399d8c5a2b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_settings_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_telemetry_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_telemetry_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52032e453046661e04a8f6beb0e1b5c4b356a654 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v3/__pycache__/wandb_telemetry_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v3/wandb_base_pb2.py b/lib/python3.12/site-packages/wandb/proto/v3/wandb_base_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..39db329f85b1bcaee10ef3120c377b6b011035a7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v3/wandb_base_pb2.py @@ -0,0 +1,55 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_base.proto +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cwandb/proto/wandb_base.proto\x12\x0ewandb_internal\"6\n\x0b_RecordInfo\x12\x11\n\tstream_id\x18\x01 \x01(\t\x12\x14\n\x0c_tracelog_id\x18\x64 \x01(\t\"!\n\x0c_RequestInfo\x12\x11\n\tstream_id\x18\x01 \x01(\t\"#\n\x0b_ResultInfo\x12\x14\n\x0c_tracelog_id\x18\x64 \x01(\tB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + + + +__RECORDINFO = DESCRIPTOR.message_types_by_name['_RecordInfo'] +__REQUESTINFO = DESCRIPTOR.message_types_by_name['_RequestInfo'] +__RESULTINFO = DESCRIPTOR.message_types_by_name['_ResultInfo'] +_RecordInfo = _reflection.GeneratedProtocolMessageType('_RecordInfo', (_message.Message,), { + 'DESCRIPTOR' : __RECORDINFO, + '__module__' : 'wandb.proto.wandb_base_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal._RecordInfo) + }) +_sym_db.RegisterMessage(_RecordInfo) + +_RequestInfo = _reflection.GeneratedProtocolMessageType('_RequestInfo', (_message.Message,), { + 'DESCRIPTOR' : __REQUESTINFO, + '__module__' : 'wandb.proto.wandb_base_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal._RequestInfo) + }) +_sym_db.RegisterMessage(_RequestInfo) + +_ResultInfo = _reflection.GeneratedProtocolMessageType('_ResultInfo', (_message.Message,), { + 'DESCRIPTOR' : __RESULTINFO, + '__module__' : 'wandb.proto.wandb_base_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal._ResultInfo) + }) +_sym_db.RegisterMessage(_ResultInfo) + +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + __RECORDINFO._serialized_start=48 + __RECORDINFO._serialized_end=102 + __REQUESTINFO._serialized_start=104 + __REQUESTINFO._serialized_end=137 + __RESULTINFO._serialized_start=139 + __RESULTINFO._serialized_end=174 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v3/wandb_internal_pb2.py b/lib/python3.12/site-packages/wandb/proto/v3/wandb_internal_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..fa373a222f68d1495944f80bb317516c98f4ad71 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v3/wandb_internal_pb2.py @@ -0,0 +1,1727 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_internal.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import enum_type_wrapper +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 +from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 +from wandb.proto import wandb_telemetry_pb2 as wandb_dot_proto_dot_wandb__telemetry__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n wandb/proto/wandb_internal.proto\x12\x0ewandb_internal\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1cwandb/proto/wandb_base.proto\x1a!wandb/proto/wandb_telemetry.proto\"\xcf\t\n\x06Record\x12\x0b\n\x03num\x18\x01 \x01(\x03\x12\x30\n\x07history\x18\x02 \x01(\x0b\x32\x1d.wandb_internal.HistoryRecordH\x00\x12\x30\n\x07summary\x18\x03 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecordH\x00\x12.\n\x06output\x18\x04 \x01(\x0b\x32\x1c.wandb_internal.OutputRecordH\x00\x12.\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x1c.wandb_internal.ConfigRecordH\x00\x12,\n\x05\x66iles\x18\x06 \x01(\x0b\x32\x1b.wandb_internal.FilesRecordH\x00\x12,\n\x05stats\x18\x07 \x01(\x0b\x32\x1b.wandb_internal.StatsRecordH\x00\x12\x32\n\x08\x61rtifact\x18\x08 \x01(\x0b\x32\x1e.wandb_internal.ArtifactRecordH\x00\x12,\n\x08tbrecord\x18\t \x01(\x0b\x32\x18.wandb_internal.TBRecordH\x00\x12,\n\x05\x61lert\x18\n \x01(\x0b\x32\x1b.wandb_internal.AlertRecordH\x00\x12\x34\n\ttelemetry\x18\x0b \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecordH\x00\x12.\n\x06metric\x18\x0c \x01(\x0b\x32\x1c.wandb_internal.MetricRecordH\x00\x12\x35\n\noutput_raw\x18\r \x01(\x0b\x32\x1f.wandb_internal.OutputRawRecordH\x00\x12(\n\x03run\x18\x11 \x01(\x0b\x32\x19.wandb_internal.RunRecordH\x00\x12-\n\x04\x65xit\x18\x12 \x01(\x0b\x32\x1d.wandb_internal.RunExitRecordH\x00\x12,\n\x05\x66inal\x18\x14 \x01(\x0b\x32\x1b.wandb_internal.FinalRecordH\x00\x12.\n\x06header\x18\x15 \x01(\x0b\x32\x1c.wandb_internal.HeaderRecordH\x00\x12.\n\x06\x66ooter\x18\x16 \x01(\x0b\x32\x1c.wandb_internal.FooterRecordH\x00\x12\x39\n\npreempting\x18\x17 \x01(\x0b\x32#.wandb_internal.RunPreemptingRecordH\x00\x12\x34\n\x12noop_link_artifact\x18\x18 \x01(\x0b\x32\x16.google.protobuf.EmptyH\x00\x12\x39\n\x0cuse_artifact\x18\x19 \x01(\x0b\x32!.wandb_internal.UseArtifactRecordH\x00\x12\x38\n\x0b\x65nvironment\x18\x1a \x01(\x0b\x32!.wandb_internal.EnvironmentRecordH\x00\x12*\n\x07request\x18\x64 \x01(\x0b\x32\x17.wandb_internal.RequestH\x00\x12(\n\x07\x63ontrol\x18\x10 \x01(\x0b\x32\x17.wandb_internal.Control\x12\x0c\n\x04uuid\x18\x13 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfoB\r\n\x0brecord_type\"\xa8\x01\n\x07\x43ontrol\x12\x10\n\x08req_resp\x18\x01 \x01(\x08\x12\r\n\x05local\x18\x02 \x01(\x08\x12\x10\n\x08relay_id\x18\x03 \x01(\t\x12\x14\n\x0cmailbox_slot\x18\x04 \x01(\t\x12\x13\n\x0b\x61lways_send\x18\x05 \x01(\x08\x12\x14\n\x0c\x66low_control\x18\x06 \x01(\x08\x12\x12\n\nend_offset\x18\x07 \x01(\x03\x12\x15\n\rconnection_id\x18\x08 \x01(\t\"\xf3\x03\n\x06Result\x12\x35\n\nrun_result\x18\x11 \x01(\x0b\x32\x1f.wandb_internal.RunUpdateResultH\x00\x12\x34\n\x0b\x65xit_result\x18\x12 \x01(\x0b\x32\x1d.wandb_internal.RunExitResultH\x00\x12\x33\n\nlog_result\x18\x14 \x01(\x0b\x32\x1d.wandb_internal.HistoryResultH\x00\x12\x37\n\x0esummary_result\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.SummaryResultH\x00\x12\x35\n\routput_result\x18\x16 \x01(\x0b\x32\x1c.wandb_internal.OutputResultH\x00\x12\x35\n\rconfig_result\x18\x17 \x01(\x0b\x32\x1c.wandb_internal.ConfigResultH\x00\x12,\n\x08response\x18\x64 \x01(\x0b\x32\x18.wandb_internal.ResponseH\x00\x12(\n\x07\x63ontrol\x18\x10 \x01(\x0b\x32\x17.wandb_internal.Control\x12\x0c\n\x04uuid\x18\x18 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._ResultInfoB\r\n\x0bresult_type\":\n\x0b\x46inalRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"b\n\x0bVersionInfo\x12\x10\n\x08producer\x18\x01 \x01(\t\x12\x14\n\x0cmin_consumer\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"n\n\x0cHeaderRecord\x12\x31\n\x0cversion_info\x18\x01 \x01(\x0b\x32\x1b.wandb_internal.VersionInfo\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\";\n\x0c\x46ooterRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"9\n\x0b\x42ranchPoint\x12\x0b\n\x03run\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01\x12\x0e\n\x06metric\x18\x03 \x01(\t\"\x91\x05\n\tRunRecord\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0e\n\x06\x65ntity\x18\x02 \x01(\t\x12\x0f\n\x07project\x18\x03 \x01(\t\x12,\n\x06\x63onfig\x18\x04 \x01(\x0b\x32\x1c.wandb_internal.ConfigRecord\x12.\n\x07summary\x18\x05 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecord\x12\x11\n\trun_group\x18\x06 \x01(\t\x12\x10\n\x08job_type\x18\x07 \x01(\t\x12\x14\n\x0c\x64isplay_name\x18\x08 \x01(\t\x12\r\n\x05notes\x18\t \x01(\t\x12\x0c\n\x04tags\x18\n \x03(\t\x12\x30\n\x08settings\x18\x0b \x01(\x0b\x32\x1e.wandb_internal.SettingsRecord\x12\x10\n\x08sweep_id\x18\x0c \x01(\t\x12\x0c\n\x04host\x18\r \x01(\t\x12\x15\n\rstarting_step\x18\x0e \x01(\x03\x12\x12\n\nstorage_id\x18\x10 \x01(\t\x12.\n\nstart_time\x18\x11 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07resumed\x18\x12 \x01(\x08\x12\x32\n\ttelemetry\x18\x13 \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecord\x12\x0f\n\x07runtime\x18\x14 \x01(\x05\x12*\n\x03git\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.GitRepoRecord\x12\x0e\n\x06\x66orked\x18\x16 \x01(\x08\x12\x31\n\x0c\x62ranch_point\x18\x17 \x01(\x0b\x32\x1b.wandb_internal.BranchPoint\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\";\n\rGitRepoRecord\x12\x1a\n\nremote_url\x18\x01 \x01(\tR\x06remote\x12\x0e\n\x06\x63ommit\x18\x02 \x01(\t\"c\n\x0fRunUpdateResult\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"\xac\x01\n\tErrorInfo\x12\x0f\n\x07message\x18\x01 \x01(\t\x12\x31\n\x04\x63ode\x18\x02 \x01(\x0e\x32#.wandb_internal.ErrorInfo.ErrorCode\"[\n\tErrorCode\x12\x0b\n\x07UNKNOWN\x10\x00\x12\x11\n\rCOMMUNICATION\x10\x01\x12\x12\n\x0e\x41UTHENTICATION\x10\x02\x12\t\n\x05USAGE\x10\x03\x12\x0f\n\x0bUNSUPPORTED\x10\x04\"`\n\rRunExitRecord\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12\x0f\n\x07runtime\x18\x02 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x0f\n\rRunExitResult\"B\n\x13RunPreemptingRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x15\n\x13RunPreemptingResult\"i\n\x0eSettingsRecord\x12*\n\x04item\x18\x01 \x03(\x0b\x32\x1c.wandb_internal.SettingsItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"/\n\x0cSettingsItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x1a\n\x0bHistoryStep\x12\x0b\n\x03num\x18\x01 \x01(\x03\"\x92\x01\n\rHistoryRecord\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.HistoryItem\x12)\n\x04step\x18\x02 \x01(\x0b\x32\x1b.wandb_internal.HistoryStep\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\x0bHistoryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0f\n\rHistoryResult\"\xdc\x01\n\x0cOutputRecord\x12<\n\x0boutput_type\x18\x01 \x01(\x0e\x32\'.wandb_internal.OutputRecord.OutputType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0c\n\x04line\x18\x03 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"$\n\nOutputType\x12\n\n\x06STDERR\x10\x00\x12\n\n\x06STDOUT\x10\x01\"\x0e\n\x0cOutputResult\"\xe2\x01\n\x0fOutputRawRecord\x12?\n\x0boutput_type\x18\x01 \x01(\x0e\x32*.wandb_internal.OutputRawRecord.OutputType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0c\n\x04line\x18\x03 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"$\n\nOutputType\x12\n\n\x06STDERR\x10\x00\x12\n\n\x06STDOUT\x10\x01\"\x11\n\x0fOutputRawResult\"\xb4\x03\n\x0cMetricRecord\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x11\n\tglob_name\x18\x02 \x01(\t\x12\x13\n\x0bstep_metric\x18\x04 \x01(\t\x12\x19\n\x11step_metric_index\x18\x05 \x01(\x05\x12.\n\x07options\x18\x06 \x01(\x0b\x32\x1d.wandb_internal.MetricOptions\x12.\n\x07summary\x18\x07 \x01(\x0b\x32\x1d.wandb_internal.MetricSummary\x12\x35\n\x04goal\x18\x08 \x01(\x0e\x32\'.wandb_internal.MetricRecord.MetricGoal\x12/\n\x08_control\x18\t \x01(\x0b\x32\x1d.wandb_internal.MetricControl\x12\x1a\n\x12\x65xpanded_from_glob\x18\n \x01(\x08\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\nMetricGoal\x12\x0e\n\nGOAL_UNSET\x10\x00\x12\x11\n\rGOAL_MINIMIZE\x10\x01\x12\x11\n\rGOAL_MAXIMIZE\x10\x02\"\x0e\n\x0cMetricResult\"C\n\rMetricOptions\x12\x11\n\tstep_sync\x18\x01 \x01(\x08\x12\x0e\n\x06hidden\x18\x02 \x01(\x08\x12\x0f\n\x07\x64\x65\x66ined\x18\x03 \x01(\x08\"\"\n\rMetricControl\x12\x11\n\toverwrite\x18\x01 \x01(\x08\"o\n\rMetricSummary\x12\x0b\n\x03min\x18\x01 \x01(\x08\x12\x0b\n\x03max\x18\x02 \x01(\x08\x12\x0c\n\x04mean\x18\x03 \x01(\x08\x12\x0c\n\x04\x62\x65st\x18\x04 \x01(\x08\x12\x0c\n\x04last\x18\x05 \x01(\x08\x12\x0c\n\x04none\x18\x06 \x01(\x08\x12\x0c\n\x04\x63opy\x18\x07 \x01(\x08\"\x93\x01\n\x0c\x43onfigRecord\x12*\n\x06update\x18\x01 \x03(\x0b\x32\x1a.wandb_internal.ConfigItem\x12*\n\x06remove\x18\x02 \x03(\x0b\x32\x1a.wandb_internal.ConfigItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"A\n\nConfigItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0e\n\x0c\x43onfigResult\"\x96\x01\n\rSummaryRecord\x12+\n\x06update\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\x12+\n\x06remove\x18\x02 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\x0bSummaryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0f\n\rSummaryResult\"d\n\x0b\x46ilesRecord\x12(\n\x05\x66iles\x18\x01 \x03(\x0b\x32\x19.wandb_internal.FilesItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\xec\x01\n\tFilesItem\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x34\n\x06policy\x18\x02 \x01(\x0e\x32$.wandb_internal.FilesItem.PolicyType\x12\x30\n\x04type\x18\x03 \x01(\x0e\x32\".wandb_internal.FilesItem.FileType\"(\n\nPolicyType\x12\x07\n\x03NOW\x10\x00\x12\x07\n\x03\x45ND\x10\x01\x12\x08\n\x04LIVE\x10\x02\"9\n\x08\x46ileType\x12\t\n\x05OTHER\x10\x00\x12\t\n\x05WANDB\x10\x01\x12\t\n\x05MEDIA\x10\x02\x12\x0c\n\x08\x41RTIFACT\x10\x03J\x04\x08\x10\x10\x11\"\r\n\x0b\x46ilesResult\"\xe6\x01\n\x0bStatsRecord\x12\x39\n\nstats_type\x18\x01 \x01(\x0e\x32%.wandb_internal.StatsRecord.StatsType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\'\n\x04item\x18\x03 \x03(\x0b\x32\x19.wandb_internal.StatsItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x17\n\tStatsType\x12\n\n\x06SYSTEM\x10\x00\",\n\tStatsItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\xe7\x03\n\x0e\x41rtifactRecord\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0f\n\x07project\x18\x02 \x01(\t\x12\x0e\n\x06\x65ntity\x18\x03 \x01(\t\x12\x0c\n\x04type\x18\x04 \x01(\t\x12\x0c\n\x04name\x18\x05 \x01(\t\x12\x0e\n\x06\x64igest\x18\x06 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x07 \x01(\t\x12\x10\n\x08metadata\x18\x08 \x01(\t\x12\x14\n\x0cuser_created\x18\t \x01(\x08\x12\x18\n\x10use_after_commit\x18\n \x01(\x08\x12\x0f\n\x07\x61liases\x18\x0b \x03(\t\x12\x32\n\x08manifest\x18\x0c \x01(\x0b\x32 .wandb_internal.ArtifactManifest\x12\x16\n\x0e\x64istributed_id\x18\r \x01(\t\x12\x10\n\x08\x66inalize\x18\x0e \x01(\x08\x12\x11\n\tclient_id\x18\x0f \x01(\t\x12\x1a\n\x12sequence_client_id\x18\x10 \x01(\t\x12\x0f\n\x07\x62\x61se_id\x18\x11 \x01(\t\x12\x1c\n\x14ttl_duration_seconds\x18\x12 \x01(\x03\x12\x0c\n\x04tags\x18\x13 \x03(\t\x12\x19\n\x11incremental_beta1\x18\x64 \x01(\x08\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\xd8\x01\n\x10\x41rtifactManifest\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x16\n\x0estorage_policy\x18\x02 \x01(\t\x12\x46\n\x15storage_policy_config\x18\x03 \x03(\x0b\x32\'.wandb_internal.StoragePolicyConfigItem\x12\x37\n\x08\x63ontents\x18\x04 \x03(\x0b\x32%.wandb_internal.ArtifactManifestEntry\x12\x1a\n\x12manifest_file_path\x18\x05 \x01(\t\"\xcf\x01\n\x15\x41rtifactManifestEntry\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x0e\n\x06\x64igest\x18\x02 \x01(\t\x12\x0b\n\x03ref\x18\x03 \x01(\t\x12\x0c\n\x04size\x18\x04 \x01(\x03\x12\x10\n\x08mimetype\x18\x05 \x01(\t\x12\x12\n\nlocal_path\x18\x06 \x01(\t\x12\x19\n\x11\x62irth_artifact_id\x18\x07 \x01(\t\x12\x12\n\nskip_cache\x18\x08 \x01(\x08\x12(\n\x05\x65xtra\x18\x10 \x03(\x0b\x32\x19.wandb_internal.ExtraItem\",\n\tExtraItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x02 \x01(\t\":\n\x17StoragePolicyConfigItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x02 \x01(\t\"\x10\n\x0e\x41rtifactResult\"\x14\n\x12LinkArtifactResult\"\xf0\x01\n\x13LinkArtifactRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x11\n\tserver_id\x18\x02 \x01(\t\x12\x16\n\x0eportfolio_name\x18\x03 \x01(\t\x12\x18\n\x10portfolio_entity\x18\x04 \x01(\t\x12\x19\n\x11portfolio_project\x18\x05 \x01(\t\x12\x19\n\x11portfolio_aliases\x18\x06 \x03(\t\x12\x1e\n\x16portfolio_organization\x18\x07 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"[\n\x14LinkArtifactResponse\x12\x15\n\rerror_message\x18\x01 \x01(\t\x12\x1a\n\rversion_index\x18\x02 \x01(\x05H\x00\x88\x01\x01\x42\x10\n\x0e_version_index\"h\n\x08TBRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\x12\x0f\n\x07log_dir\x18\x01 \x01(\t\x12\x10\n\x08root_dir\x18\x03 \x01(\t\x12\x0c\n\x04save\x18\x02 \x01(\x08\"\n\n\x08TBResult\"}\n\x0b\x41lertRecord\x12\r\n\x05title\x18\x01 \x01(\t\x12\x0c\n\x04text\x18\x02 \x01(\t\x12\r\n\x05level\x18\x03 \x01(\t\x12\x15\n\rwait_duration\x18\x04 \x01(\x03\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\r\n\x0b\x41lertResult\"\xb3\x11\n\x07Request\x12\x38\n\x0bstop_status\x18\x01 \x01(\x0b\x32!.wandb_internal.StopStatusRequestH\x00\x12>\n\x0enetwork_status\x18\x02 \x01(\x0b\x32$.wandb_internal.NetworkStatusRequestH\x00\x12-\n\x05\x64\x65\x66\x65r\x18\x03 \x01(\x0b\x32\x1c.wandb_internal.DeferRequestH\x00\x12\x38\n\x0bget_summary\x18\x04 \x01(\x0b\x32!.wandb_internal.GetSummaryRequestH\x00\x12-\n\x05login\x18\x05 \x01(\x0b\x32\x1c.wandb_internal.LoginRequestH\x00\x12-\n\x05pause\x18\x06 \x01(\x0b\x32\x1c.wandb_internal.PauseRequestH\x00\x12/\n\x06resume\x18\x07 \x01(\x0b\x32\x1d.wandb_internal.ResumeRequestH\x00\x12\x34\n\tpoll_exit\x18\x08 \x01(\x0b\x32\x1f.wandb_internal.PollExitRequestH\x00\x12@\n\x0fsampled_history\x18\t \x01(\x0b\x32%.wandb_internal.SampledHistoryRequestH\x00\x12@\n\x0fpartial_history\x18\n \x01(\x0b\x32%.wandb_internal.PartialHistoryRequestH\x00\x12\x34\n\trun_start\x18\x0b \x01(\x0b\x32\x1f.wandb_internal.RunStartRequestH\x00\x12<\n\rcheck_version\x18\x0c \x01(\x0b\x32#.wandb_internal.CheckVersionRequestH\x00\x12:\n\x0clog_artifact\x18\r \x01(\x0b\x32\".wandb_internal.LogArtifactRequestH\x00\x12\x44\n\x11\x64ownload_artifact\x18\x0e \x01(\x0b\x32\'.wandb_internal.DownloadArtifactRequestH\x00\x12\x35\n\tkeepalive\x18\x11 \x01(\x0b\x32 .wandb_internal.KeepaliveRequestH\x00\x12>\n\x0eserver_feature\x18\x12 \x01(\x0b\x32$.wandb_internal.ServerFeatureRequestH\x00\x12\x36\n\nrun_status\x18\x14 \x01(\x0b\x32 .wandb_internal.RunStatusRequestH\x00\x12/\n\x06\x63\x61ncel\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.CancelRequestH\x00\x12\x44\n\x11internal_messages\x18\x17 \x01(\x0b\x32\'.wandb_internal.InternalMessagesRequestH\x00\x12@\n\x0fpython_packages\x18\x18 \x01(\x0b\x32%.wandb_internal.PythonPackagesRequestH\x00\x12\x33\n\x08shutdown\x18@ \x01(\x0b\x32\x1f.wandb_internal.ShutdownRequestH\x00\x12/\n\x06\x61ttach\x18\x41 \x01(\x0b\x32\x1d.wandb_internal.AttachRequestH\x00\x12/\n\x06status\x18\x42 \x01(\x0b\x32\x1d.wandb_internal.StatusRequestH\x00\x12\x38\n\x0bserver_info\x18\x43 \x01(\x0b\x32!.wandb_internal.ServerInfoRequestH\x00\x12\x38\n\x0bsender_mark\x18\x44 \x01(\x0b\x32!.wandb_internal.SenderMarkRequestH\x00\x12\x38\n\x0bsender_read\x18\x45 \x01(\x0b\x32!.wandb_internal.SenderReadRequestH\x00\x12<\n\rstatus_report\x18\x46 \x01(\x0b\x32#.wandb_internal.StatusReportRequestH\x00\x12>\n\x0esummary_record\x18G \x01(\x0b\x32$.wandb_internal.SummaryRecordRequestH\x00\x12\x42\n\x10telemetry_record\x18H \x01(\x0b\x32&.wandb_internal.TelemetryRecordRequestH\x00\x12\x32\n\x08job_info\x18I \x01(\x0b\x32\x1e.wandb_internal.JobInfoRequestH\x00\x12\x45\n\x12get_system_metrics\x18J \x01(\x0b\x32\'.wandb_internal.GetSystemMetricsRequestH\x00\x12\x34\n\tjob_input\x18M \x01(\x0b\x32\x1f.wandb_internal.JobInputRequestH\x00\x12<\n\rlink_artifact\x18N \x01(\x0b\x32#.wandb_internal.LinkArtifactRequestH\x00\x12N\n\x17run_finish_without_exit\x18O \x01(\x0b\x32+.wandb_internal.RunFinishWithoutExitRequestH\x00\x12\x38\n\x0bsync_finish\x18Q \x01(\x0b\x32!.wandb_internal.SyncFinishRequestH\x00\x12;\n\noperations\x18R \x01(\x0b\x32%.wandb_internal.OperationStatsRequestH\x00\x12\x39\n\x0btest_inject\x18\xe8\x07 \x01(\x0b\x32!.wandb_internal.TestInjectRequestH\x00\x42\x0e\n\x0crequest_typeJ\x04\x08\x16\x10\x17J\x04\x08K\x10LJ\x04\x08L\x10MJ\x04\x08P\x10Q\"\x9b\x0e\n\x08Response\x12?\n\x12keepalive_response\x18\x12 \x01(\x0b\x32!.wandb_internal.KeepaliveResponseH\x00\x12\x42\n\x14stop_status_response\x18\x13 \x01(\x0b\x32\".wandb_internal.StopStatusResponseH\x00\x12H\n\x17network_status_response\x18\x14 \x01(\x0b\x32%.wandb_internal.NetworkStatusResponseH\x00\x12\x37\n\x0elogin_response\x18\x18 \x01(\x0b\x32\x1d.wandb_internal.LoginResponseH\x00\x12\x42\n\x14get_summary_response\x18\x19 \x01(\x0b\x32\".wandb_internal.GetSummaryResponseH\x00\x12>\n\x12poll_exit_response\x18\x1a \x01(\x0b\x32 .wandb_internal.PollExitResponseH\x00\x12J\n\x18sampled_history_response\x18\x1b \x01(\x0b\x32&.wandb_internal.SampledHistoryResponseH\x00\x12>\n\x12run_start_response\x18\x1c \x01(\x0b\x32 .wandb_internal.RunStartResponseH\x00\x12\x46\n\x16\x63heck_version_response\x18\x1d \x01(\x0b\x32$.wandb_internal.CheckVersionResponseH\x00\x12\x44\n\x15log_artifact_response\x18\x1e \x01(\x0b\x32#.wandb_internal.LogArtifactResponseH\x00\x12N\n\x1a\x64ownload_artifact_response\x18\x1f \x01(\x0b\x32(.wandb_internal.DownloadArtifactResponseH\x00\x12H\n\x17server_feature_response\x18 \x01(\x0b\x32%.wandb_internal.ServerFeatureResponseH\x00\x12@\n\x13run_status_response\x18# \x01(\x0b\x32!.wandb_internal.RunStatusResponseH\x00\x12\x39\n\x0f\x63\x61ncel_response\x18$ \x01(\x0b\x32\x1e.wandb_internal.CancelResponseH\x00\x12N\n\x1ainternal_messages_response\x18% \x01(\x0b\x32(.wandb_internal.InternalMessagesResponseH\x00\x12=\n\x11shutdown_response\x18@ \x01(\x0b\x32 .wandb_internal.ShutdownResponseH\x00\x12\x39\n\x0f\x61ttach_response\x18\x41 \x01(\x0b\x32\x1e.wandb_internal.AttachResponseH\x00\x12\x39\n\x0fstatus_response\x18\x42 \x01(\x0b\x32\x1e.wandb_internal.StatusResponseH\x00\x12\x42\n\x14server_info_response\x18\x43 \x01(\x0b\x32\".wandb_internal.ServerInfoResponseH\x00\x12<\n\x11job_info_response\x18\x44 \x01(\x0b\x32\x1f.wandb_internal.JobInfoResponseH\x00\x12O\n\x1bget_system_metrics_response\x18\x45 \x01(\x0b\x32(.wandb_internal.GetSystemMetricsResponseH\x00\x12\x46\n\x16link_artifact_response\x18G \x01(\x0b\x32$.wandb_internal.LinkArtifactResponseH\x00\x12\x35\n\rsync_response\x18\x46 \x01(\x0b\x32\x1c.wandb_internal.SyncResponseH\x00\x12X\n run_finish_without_exit_response\x18H \x01(\x0b\x32,.wandb_internal.RunFinishWithoutExitResponseH\x00\x12\x45\n\x13operations_response\x18J \x01(\x0b\x32&.wandb_internal.OperationStatsResponseH\x00\x12\x43\n\x14test_inject_response\x18\xe8\x07 \x01(\x0b\x32\".wandb_internal.TestInjectResponseH\x00\x42\x0f\n\rresponse_typeJ\x04\x08I\x10J\"\xc0\x02\n\x0c\x44\x65\x66\x65rRequest\x12\x36\n\x05state\x18\x01 \x01(\x0e\x32\'.wandb_internal.DeferRequest.DeferState\"\xf7\x01\n\nDeferState\x12\t\n\x05\x42\x45GIN\x10\x00\x12\r\n\tFLUSH_RUN\x10\x01\x12\x0f\n\x0b\x46LUSH_STATS\x10\x02\x12\x19\n\x15\x46LUSH_PARTIAL_HISTORY\x10\x03\x12\x0c\n\x08\x46LUSH_TB\x10\x04\x12\r\n\tFLUSH_SUM\x10\x05\x12\x13\n\x0f\x46LUSH_DEBOUNCER\x10\x06\x12\x10\n\x0c\x46LUSH_OUTPUT\x10\x07\x12\r\n\tFLUSH_JOB\x10\x08\x12\r\n\tFLUSH_DIR\x10\t\x12\x0c\n\x08\x46LUSH_FP\x10\n\x12\x0b\n\x07JOIN_FP\x10\x0b\x12\x0c\n\x08\x46LUSH_FS\x10\x0c\x12\x0f\n\x0b\x46LUSH_FINAL\x10\r\x12\x07\n\x03\x45ND\x10\x0e\"<\n\x0cPauseRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x0f\n\rPauseResponse\"=\n\rResumeRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x10\n\x0eResumeResponse\"M\n\x0cLoginRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"&\n\rLoginResponse\x12\x15\n\ractive_entity\x18\x01 \x01(\t\"A\n\x11GetSummaryRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"?\n\x12GetSummaryResponse\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\"G\n\x17GetSystemMetricsRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"R\n\x12SystemMetricSample\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\r\n\x05value\x18\x02 \x01(\x02\"I\n\x13SystemMetricsBuffer\x12\x32\n\x06record\x18\x01 \x03(\x0b\x32\".wandb_internal.SystemMetricSample\"\xca\x01\n\x18GetSystemMetricsResponse\x12S\n\x0esystem_metrics\x18\x01 \x03(\x0b\x32;.wandb_internal.GetSystemMetricsResponse.SystemMetricsEntry\x1aY\n\x12SystemMetricsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x32\n\x05value\x18\x02 \x01(\x0b\x32#.wandb_internal.SystemMetricsBuffer:\x02\x38\x01\"=\n\rStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\")\n\x0eStatusResponse\x12\x17\n\x0frun_should_stop\x18\x01 \x01(\x08\"A\n\x11StopStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"-\n\x12StopStatusResponse\x12\x17\n\x0frun_should_stop\x18\x01 \x01(\x08\"D\n\x14NetworkStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"P\n\x15NetworkStatusResponse\x12\x37\n\x11network_responses\x18\x01 \x03(\x0b\x32\x1c.wandb_internal.HttpResponse\"D\n\x0cHttpResponse\x12\x18\n\x10http_status_code\x18\x01 \x01(\x05\x12\x1a\n\x12http_response_text\x18\x02 \x01(\t\"G\n\x17InternalMessagesRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"N\n\x18InternalMessagesResponse\x12\x32\n\x08messages\x18\x01 \x01(\x0b\x32 .wandb_internal.InternalMessages\"#\n\x10InternalMessages\x12\x0f\n\x07warning\x18\x01 \x03(\t\"?\n\x0fPollExitRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\xf5\x01\n\x10PollExitResponse\x12\x0c\n\x04\x64one\x18\x01 \x01(\x08\x12\x32\n\x0b\x65xit_result\x18\x02 \x01(\x0b\x32\x1d.wandb_internal.RunExitResult\x12\x35\n\x0cpusher_stats\x18\x03 \x01(\x0b\x32\x1f.wandb_internal.FilePusherStats\x12/\n\x0b\x66ile_counts\x18\x04 \x01(\x0b\x32\x1a.wandb_internal.FileCounts\x12\x37\n\x0foperation_stats\x18\x05 \x01(\x0b\x32\x1e.wandb_internal.OperationStats\"E\n\x15OperationStatsRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"Q\n\x16OperationStatsResponse\x12\x37\n\x0foperation_stats\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.OperationStats\"Y\n\x0eOperationStats\x12-\n\noperations\x18\x01 \x03(\x0b\x32\x19.wandb_internal.Operation\x12\x18\n\x10total_operations\x18\x02 \x01(\x03\"\x87\x01\n\tOperation\x12\x0c\n\x04\x64\x65sc\x18\x01 \x01(\t\x12\x17\n\x0fruntime_seconds\x18\x02 \x01(\x01\x12\x10\n\x08progress\x18\x03 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x04 \x01(\t\x12+\n\x08subtasks\x18\x05 \x03(\x0b\x32\x19.wandb_internal.Operation\"\x13\n\x11SenderMarkRequest\"\x13\n\x11SyncFinishRequest\"E\n\x0cSyncResponse\x12\x0b\n\x03url\x18\x01 \x01(\t\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"?\n\x11SenderReadRequest\x12\x14\n\x0cstart_offset\x18\x01 \x01(\x03\x12\x14\n\x0c\x66inal_offset\x18\x02 \x01(\x03\"m\n\x13StatusReportRequest\x12\x12\n\nrecord_num\x18\x01 \x01(\x03\x12\x13\n\x0bsent_offset\x18\x02 \x01(\x03\x12-\n\tsync_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"F\n\x14SummaryRecordRequest\x12.\n\x07summary\x18\x01 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecord\"L\n\x16TelemetryRecordRequest\x12\x32\n\ttelemetry\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecord\"A\n\x11ServerInfoRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"|\n\x12ServerInfoResponse\x12-\n\nlocal_info\x18\x01 \x01(\x0b\x32\x19.wandb_internal.LocalInfo\x12\x37\n\x0fserver_messages\x18\x02 \x01(\x0b\x32\x1e.wandb_internal.ServerMessages\"=\n\x0eServerMessages\x12+\n\x04item\x18\x01 \x03(\x0b\x32\x1d.wandb_internal.ServerMessage\"e\n\rServerMessage\x12\x12\n\nplain_text\x18\x01 \x01(\t\x12\x10\n\x08utf_text\x18\x02 \x01(\t\x12\x11\n\thtml_text\x18\x03 \x01(\t\x12\x0c\n\x04type\x18\x04 \x01(\t\x12\r\n\x05level\x18\x05 \x01(\x05\"c\n\nFileCounts\x12\x13\n\x0bwandb_count\x18\x01 \x01(\x05\x12\x13\n\x0bmedia_count\x18\x02 \x01(\x05\x12\x16\n\x0e\x61rtifact_count\x18\x03 \x01(\x05\x12\x13\n\x0bother_count\x18\x04 \x01(\x05\"U\n\x0f\x46ilePusherStats\x12\x16\n\x0euploaded_bytes\x18\x01 \x01(\x03\x12\x13\n\x0btotal_bytes\x18\x02 \x01(\x03\x12\x15\n\rdeduped_bytes\x18\x03 \x01(\x03\"\x1e\n\rFilesUploaded\x12\r\n\x05\x66iles\x18\x01 \x03(\t\"\xf4\x01\n\x17\x46ileTransferInfoRequest\x12\x42\n\x04type\x18\x01 \x01(\x0e\x32\x34.wandb_internal.FileTransferInfoRequest.TransferType\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\x0b\n\x03url\x18\x03 \x01(\t\x12\x0c\n\x04size\x18\x04 \x01(\x03\x12\x11\n\tprocessed\x18\x05 \x01(\x03\x12/\n\x0b\x66ile_counts\x18\x06 \x01(\x0b\x32\x1a.wandb_internal.FileCounts\"(\n\x0cTransferType\x12\n\n\x06Upload\x10\x00\x12\x0c\n\x08\x44ownload\x10\x01\"1\n\tLocalInfo\x12\x0f\n\x07version\x18\x01 \x01(\t\x12\x13\n\x0bout_of_date\x18\x02 \x01(\x08\"?\n\x0fShutdownRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x12\n\x10ShutdownResponse\"P\n\rAttachRequest\x12\x11\n\tattach_id\x18\x14 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"b\n\x0e\x41ttachResponse\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"\xd5\x02\n\x11TestInjectRequest\x12\x13\n\x0bhandler_exc\x18\x01 \x01(\x08\x12\x14\n\x0chandler_exit\x18\x02 \x01(\x08\x12\x15\n\rhandler_abort\x18\x03 \x01(\x08\x12\x12\n\nsender_exc\x18\x04 \x01(\x08\x12\x13\n\x0bsender_exit\x18\x05 \x01(\x08\x12\x14\n\x0csender_abort\x18\x06 \x01(\x08\x12\x0f\n\x07req_exc\x18\x07 \x01(\x08\x12\x10\n\x08req_exit\x18\x08 \x01(\x08\x12\x11\n\treq_abort\x18\t \x01(\x08\x12\x10\n\x08resp_exc\x18\n \x01(\x08\x12\x11\n\tresp_exit\x18\x0b \x01(\x08\x12\x12\n\nresp_abort\x18\x0c \x01(\x08\x12\x10\n\x08msg_drop\x18\r \x01(\x08\x12\x10\n\x08msg_hang\x18\x0e \x01(\x08\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x14\n\x12TestInjectResponse\"\x1e\n\rHistoryAction\x12\r\n\x05\x66lush\x18\x01 \x01(\x08\"\xca\x01\n\x15PartialHistoryRequest\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.HistoryItem\x12)\n\x04step\x18\x02 \x01(\x0b\x32\x1b.wandb_internal.HistoryStep\x12-\n\x06\x61\x63tion\x18\x03 \x01(\x0b\x32\x1d.wandb_internal.HistoryAction\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x18\n\x16PartialHistoryResponse\"E\n\x15SampledHistoryRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"_\n\x12SampledHistoryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x14\n\x0cvalues_float\x18\x03 \x03(\x02\x12\x12\n\nvalues_int\x18\x04 \x03(\x03\"J\n\x16SampledHistoryResponse\x12\x30\n\x04item\x18\x01 \x03(\x0b\x32\".wandb_internal.SampledHistoryItem\"@\n\x10RunStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"x\n\x11RunStatusResponse\x12\x18\n\x10sync_items_total\x18\x01 \x01(\x03\x12\x1a\n\x12sync_items_pending\x18\x02 \x01(\x03\x12-\n\tsync_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"g\n\x0fRunStartRequest\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x12\n\x10RunStartResponse\"K\n\x1bRunFinishWithoutExitRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x1e\n\x1cRunFinishWithoutExitResponse\"\\\n\x13\x43heckVersionRequest\x12\x17\n\x0f\x63urrent_version\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"]\n\x14\x43heckVersionResponse\x12\x17\n\x0fupgrade_message\x18\x01 \x01(\t\x12\x14\n\x0cyank_message\x18\x02 \x01(\t\x12\x16\n\x0e\x64\x65lete_message\x18\x03 \x01(\t\">\n\x0eJobInfoRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"6\n\x0fJobInfoResponse\x12\x12\n\nsequenceId\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"\x9f\x01\n\x12LogArtifactRequest\x12\x30\n\x08\x61rtifact\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.ArtifactRecord\x12\x14\n\x0chistory_step\x18\x02 \x01(\x03\x12\x13\n\x0bstaging_dir\x18\x03 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"A\n\x13LogArtifactResponse\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x15\n\rerror_message\x18\x02 \x01(\t\"\xbe\x01\n\x17\x44ownloadArtifactRequest\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x15\n\rdownload_root\x18\x02 \x01(\t\x12 \n\x18\x61llow_missing_references\x18\x04 \x01(\x08\x12\x12\n\nskip_cache\x18\x05 \x01(\x08\x12\x13\n\x0bpath_prefix\x18\x06 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"1\n\x18\x44ownloadArtifactResponse\x12\x15\n\rerror_message\x18\x01 \x01(\t\"@\n\x10KeepaliveRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x13\n\x11KeepaliveResponse\"q\n\x0c\x41rtifactInfo\x12\x10\n\x08\x61rtifact\x18\x01 \x01(\t\x12\x12\n\nentrypoint\x18\x02 \x03(\t\x12\x10\n\x08notebook\x18\x03 \x01(\x08\x12\x15\n\rbuild_context\x18\x04 \x01(\t\x12\x12\n\ndockerfile\x18\x05 \x01(\t\")\n\x07GitInfo\x12\x0e\n\x06remote\x18\x01 \x01(\t\x12\x0e\n\x06\x63ommit\x18\x02 \x01(\t\"\x87\x01\n\tGitSource\x12)\n\x08git_info\x18\x01 \x01(\x0b\x32\x17.wandb_internal.GitInfo\x12\x12\n\nentrypoint\x18\x02 \x03(\t\x12\x10\n\x08notebook\x18\x03 \x01(\x08\x12\x15\n\rbuild_context\x18\x04 \x01(\t\x12\x12\n\ndockerfile\x18\x05 \x01(\t\"\x1c\n\x0bImageSource\x12\r\n\x05image\x18\x01 \x01(\t\"\x8c\x01\n\x06Source\x12&\n\x03git\x18\x01 \x01(\x0b\x32\x19.wandb_internal.GitSource\x12.\n\x08\x61rtifact\x18\x02 \x01(\x0b\x32\x1c.wandb_internal.ArtifactInfo\x12*\n\x05image\x18\x03 \x01(\x0b\x32\x1b.wandb_internal.ImageSource\"k\n\tJobSource\x12\x10\n\x08_version\x18\x01 \x01(\t\x12\x13\n\x0bsource_type\x18\x02 \x01(\t\x12&\n\x06source\x18\x03 \x01(\x0b\x32\x16.wandb_internal.Source\x12\x0f\n\x07runtime\x18\x04 \x01(\t\"V\n\x12PartialJobArtifact\x12\x10\n\x08job_name\x18\x01 \x01(\t\x12.\n\x0bsource_info\x18\x02 \x01(\x0b\x32\x19.wandb_internal.JobSource\"\x9d\x01\n\x11UseArtifactRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x33\n\x07partial\x18\x04 \x01(\x0b\x32\".wandb_internal.PartialJobArtifact\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x13\n\x11UseArtifactResult\"R\n\rCancelRequest\x12\x13\n\x0b\x63\x61ncel_slot\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x10\n\x0e\x43\x61ncelResponse\"\'\n\x08\x44iskInfo\x12\r\n\x05total\x18\x01 \x01(\x04\x12\x0c\n\x04used\x18\x02 \x01(\x04\"\x1b\n\nMemoryInfo\x12\r\n\x05total\x18\x01 \x01(\x04\"/\n\x07\x43puInfo\x12\r\n\x05\x63ount\x18\x01 \x01(\r\x12\x15\n\rcount_logical\x18\x02 \x01(\r\"\x9a\x01\n\tAppleInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\necpu_cores\x18\x02 \x01(\r\x12\x12\n\npcpu_cores\x18\x03 \x01(\r\x12\x11\n\tgpu_cores\x18\x04 \x01(\r\x12\x11\n\tmemory_gb\x18\x05 \x01(\r\x12\x18\n\x10swap_total_bytes\x18\x06 \x01(\x04\x12\x17\n\x0fram_total_bytes\x18\x07 \x01(\x04\"k\n\rGpuNvidiaInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x14\n\x0cmemory_total\x18\x02 \x01(\x04\x12\x12\n\ncuda_cores\x18\x03 \x01(\r\x12\x14\n\x0c\x61rchitecture\x18\x04 \x01(\t\x12\x0c\n\x04uuid\x18\x05 \x01(\t\"\x89\x02\n\nGpuAmdInfo\x12\n\n\x02id\x18\x01 \x01(\t\x12\x11\n\tunique_id\x18\x02 \x01(\t\x12\x15\n\rvbios_version\x18\x03 \x01(\t\x12\x19\n\x11performance_level\x18\x04 \x01(\t\x12\x15\n\rgpu_overdrive\x18\x05 \x01(\t\x12\x1c\n\x14gpu_memory_overdrive\x18\x06 \x01(\t\x12\x11\n\tmax_power\x18\x07 \x01(\t\x12\x0e\n\x06series\x18\x08 \x01(\t\x12\r\n\x05model\x18\t \x01(\t\x12\x0e\n\x06vendor\x18\n \x01(\t\x12\x0b\n\x03sku\x18\x0b \x01(\t\x12\x12\n\nsclk_range\x18\x0c \x01(\t\x12\x12\n\nmclk_range\x18\r \x01(\t\"n\n\x0cTrainiumInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06vendor\x18\x02 \x01(\t\x12\x1b\n\x13neuron_device_count\x18\x03 \x01(\r\x12#\n\x1bneuroncore_per_device_count\x18\x04 \x01(\r\"Q\n\x07TPUInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07hbm_gib\x18\x02 \x01(\r\x12\x18\n\x10\x64\x65vices_per_chip\x18\x03 \x01(\r\x12\r\n\x05\x63ount\x18\x04 \x01(\r\"E\n\rCoreWeaveInfo\x12\x14\n\x0c\x63luster_name\x18\x01 \x01(\t\x12\x0e\n\x06org_id\x18\x02 \x01(\t\x12\x0e\n\x06region\x18\x03 \x01(\t\"\xa8\t\n\x11\x45nvironmentRecord\x12\n\n\x02os\x18\x01 \x01(\t\x12\x0e\n\x06python\x18\x02 \x01(\t\x12\x39\n\nstarted_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\tstartedAt\x12\x0e\n\x06\x64ocker\x18\x04 \x01(\t\x12\x0c\n\x04\x61rgs\x18\x05 \x03(\t\x12\x0f\n\x07program\x18\x06 \x01(\t\x12\x1b\n\tcode_path\x18\x07 \x01(\tR\x08\x63odePath\x12&\n\x0f\x63ode_path_local\x18\x08 \x01(\tR\rcodePathLocal\x12*\n\x03git\x18\t \x01(\x0b\x32\x1d.wandb_internal.GitRepoRecord\x12\r\n\x05\x65mail\x18\n \x01(\t\x12\x0c\n\x04root\x18\x0b \x01(\t\x12\x0c\n\x04host\x18\x0c \x01(\t\x12\x10\n\x08username\x18\r \x01(\t\x12\x12\n\nexecutable\x18\x0e \x01(\t\x12\r\n\x05\x63olab\x18\x0f \x01(\t\x12\x1c\n\tcpu_count\x18\x10 \x01(\rR\tcpu_count\x12,\n\x11\x63pu_count_logical\x18\x11 \x01(\rR\x11\x63pu_count_logical\x12\x15\n\x08gpu_type\x18\x12 \x01(\tR\x03gpu\x12\x1c\n\tgpu_count\x18\x13 \x01(\rR\tgpu_count\x12\x39\n\x04\x64isk\x18\x14 \x03(\x0b\x32+.wandb_internal.EnvironmentRecord.DiskEntry\x12*\n\x06memory\x18\x15 \x01(\x0b\x32\x1a.wandb_internal.MemoryInfo\x12$\n\x03\x63pu\x18\x16 \x01(\x0b\x32\x17.wandb_internal.CpuInfo\x12(\n\x05\x61pple\x18\x17 \x01(\x0b\x32\x19.wandb_internal.AppleInfo\x12=\n\ngpu_nvidia\x18\x18 \x03(\x0b\x32\x1d.wandb_internal.GpuNvidiaInfoR\ngpu_nvidia\x12\x14\n\x0c\x63uda_version\x18\x19 \x01(\t\x12\x34\n\x07gpu_amd\x18\x1a \x03(\x0b\x32\x1a.wandb_internal.GpuAmdInfoR\x07gpu_amd\x12;\n\x05slurm\x18\x1b \x03(\x0b\x32,.wandb_internal.EnvironmentRecord.SlurmEntry\x12.\n\x08trainium\x18\x1c \x01(\x0b\x32\x1c.wandb_internal.TrainiumInfo\x12$\n\x03tpu\x18\x1d \x01(\x0b\x32\x17.wandb_internal.TPUInfo\x12\x30\n\tcoreweave\x18\x1e \x01(\x0b\x32\x1d.wandb_internal.CoreWeaveInfo\x12\x12\n\twriter_id\x18\xc7\x01 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\x1a\x45\n\tDiskEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.wandb_internal.DiskInfo:\x02\x38\x01\x1a,\n\nSlurmEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\x15PythonPackagesRequest\x12\x44\n\x07package\x18\x01 \x03(\x0b\x32\x33.wandb_internal.PythonPackagesRequest.PythonPackage\x1a.\n\rPythonPackage\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"\x1c\n\x0cJobInputPath\x12\x0c\n\x04path\x18\x01 \x03(\t\"\xd6\x01\n\x0eJobInputSource\x12\x44\n\nrun_config\x18\x01 \x01(\x0b\x32..wandb_internal.JobInputSource.RunConfigSourceH\x00\x12?\n\x04\x66ile\x18\x02 \x01(\x0b\x32/.wandb_internal.JobInputSource.ConfigFileSourceH\x00\x1a\x11\n\x0fRunConfigSource\x1a \n\x10\x43onfigFileSource\x12\x0c\n\x04path\x18\x01 \x01(\tB\x08\n\x06source\"\xc7\x01\n\x0fJobInputRequest\x12\x34\n\x0cinput_source\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.JobInputSource\x12\x33\n\rinclude_paths\x18\x02 \x03(\x0b\x32\x1c.wandb_internal.JobInputPath\x12\x33\n\rexclude_paths\x18\x03 \x03(\x0b\x32\x1c.wandb_internal.JobInputPath\x12\x14\n\x0cinput_schema\x18\x04 \x01(\t\"t\n\x14ServerFeatureRequest\x12.\n\x07\x66\x65\x61ture\x18\x01 \x01(\x0e\x32\x1d.wandb_internal.ServerFeature\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"K\n\x15ServerFeatureResponse\x12\x32\n\x07\x66\x65\x61ture\x18\x01 \x01(\x0b\x32!.wandb_internal.ServerFeatureItem\"2\n\x11ServerFeatureItem\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07\x65nabled\x18\x02 \x01(\x08*\x80\x04\n\rServerFeature\x12\x13\n\x0fLARGE_FILENAMES\x10\x00\x12\x11\n\rARTIFACT_TAGS\x10\x01\x12\x0e\n\nCLIENT_IDS\x10\x02\x12\x1c\n\x18\x41RTIFACT_REGISTRY_SEARCH\x10\x03\x12\x1b\n\x17STRUCTURED_CONSOLE_LOGS\x10\x04\x12(\n$ARTIFACT_COLLECTION_MEMBERSHIP_FILES\x10\x05\x12\x38\n4ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER\x10\x06\x12\x34\n0USE_ARTIFACT_WITH_ENTITY_AND_PROJECT_INFORMATION\x10\x07\x12\x1f\n\x1b\x45XPAND_DEFINED_METRIC_GLOBS\x10\x08\x12\x1f\n\x1b\x41UTOMATION_EVENT_RUN_METRIC\x10\t\x12&\n\"AUTOMATION_EVENT_RUN_METRIC_CHANGE\x10\n\x12\x1b\n\x17\x41UTOMATION_ACTION_NO_OP\x10\x0b\x12/\n+INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION\x10\x0c\x12*\n&PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP\x10\rB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_SERVERFEATURE = DESCRIPTOR.enum_types_by_name['ServerFeature'] +ServerFeature = enum_type_wrapper.EnumTypeWrapper(_SERVERFEATURE) +LARGE_FILENAMES = 0 +ARTIFACT_TAGS = 1 +CLIENT_IDS = 2 +ARTIFACT_REGISTRY_SEARCH = 3 +STRUCTURED_CONSOLE_LOGS = 4 +ARTIFACT_COLLECTION_MEMBERSHIP_FILES = 5 +ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER = 6 +USE_ARTIFACT_WITH_ENTITY_AND_PROJECT_INFORMATION = 7 +EXPAND_DEFINED_METRIC_GLOBS = 8 +AUTOMATION_EVENT_RUN_METRIC = 9 +AUTOMATION_EVENT_RUN_METRIC_CHANGE = 10 +AUTOMATION_ACTION_NO_OP = 11 +INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION = 12 +PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP = 13 + + +_RECORD = DESCRIPTOR.message_types_by_name['Record'] +_CONTROL = DESCRIPTOR.message_types_by_name['Control'] +_RESULT = DESCRIPTOR.message_types_by_name['Result'] +_FINALRECORD = DESCRIPTOR.message_types_by_name['FinalRecord'] +_VERSIONINFO = DESCRIPTOR.message_types_by_name['VersionInfo'] +_HEADERRECORD = DESCRIPTOR.message_types_by_name['HeaderRecord'] +_FOOTERRECORD = DESCRIPTOR.message_types_by_name['FooterRecord'] +_BRANCHPOINT = DESCRIPTOR.message_types_by_name['BranchPoint'] +_RUNRECORD = DESCRIPTOR.message_types_by_name['RunRecord'] +_GITREPORECORD = DESCRIPTOR.message_types_by_name['GitRepoRecord'] +_RUNUPDATERESULT = DESCRIPTOR.message_types_by_name['RunUpdateResult'] +_ERRORINFO = DESCRIPTOR.message_types_by_name['ErrorInfo'] +_RUNEXITRECORD = DESCRIPTOR.message_types_by_name['RunExitRecord'] +_RUNEXITRESULT = DESCRIPTOR.message_types_by_name['RunExitResult'] +_RUNPREEMPTINGRECORD = DESCRIPTOR.message_types_by_name['RunPreemptingRecord'] +_RUNPREEMPTINGRESULT = DESCRIPTOR.message_types_by_name['RunPreemptingResult'] +_SETTINGSRECORD = DESCRIPTOR.message_types_by_name['SettingsRecord'] +_SETTINGSITEM = DESCRIPTOR.message_types_by_name['SettingsItem'] +_HISTORYSTEP = DESCRIPTOR.message_types_by_name['HistoryStep'] +_HISTORYRECORD = DESCRIPTOR.message_types_by_name['HistoryRecord'] +_HISTORYITEM = DESCRIPTOR.message_types_by_name['HistoryItem'] +_HISTORYRESULT = DESCRIPTOR.message_types_by_name['HistoryResult'] +_OUTPUTRECORD = DESCRIPTOR.message_types_by_name['OutputRecord'] +_OUTPUTRESULT = DESCRIPTOR.message_types_by_name['OutputResult'] +_OUTPUTRAWRECORD = DESCRIPTOR.message_types_by_name['OutputRawRecord'] +_OUTPUTRAWRESULT = DESCRIPTOR.message_types_by_name['OutputRawResult'] +_METRICRECORD = DESCRIPTOR.message_types_by_name['MetricRecord'] +_METRICRESULT = DESCRIPTOR.message_types_by_name['MetricResult'] +_METRICOPTIONS = DESCRIPTOR.message_types_by_name['MetricOptions'] +_METRICCONTROL = DESCRIPTOR.message_types_by_name['MetricControl'] +_METRICSUMMARY = DESCRIPTOR.message_types_by_name['MetricSummary'] +_CONFIGRECORD = DESCRIPTOR.message_types_by_name['ConfigRecord'] +_CONFIGITEM = DESCRIPTOR.message_types_by_name['ConfigItem'] +_CONFIGRESULT = DESCRIPTOR.message_types_by_name['ConfigResult'] +_SUMMARYRECORD = DESCRIPTOR.message_types_by_name['SummaryRecord'] +_SUMMARYITEM = DESCRIPTOR.message_types_by_name['SummaryItem'] +_SUMMARYRESULT = DESCRIPTOR.message_types_by_name['SummaryResult'] +_FILESRECORD = DESCRIPTOR.message_types_by_name['FilesRecord'] +_FILESITEM = DESCRIPTOR.message_types_by_name['FilesItem'] +_FILESRESULT = DESCRIPTOR.message_types_by_name['FilesResult'] +_STATSRECORD = DESCRIPTOR.message_types_by_name['StatsRecord'] +_STATSITEM = DESCRIPTOR.message_types_by_name['StatsItem'] +_ARTIFACTRECORD = DESCRIPTOR.message_types_by_name['ArtifactRecord'] +_ARTIFACTMANIFEST = DESCRIPTOR.message_types_by_name['ArtifactManifest'] +_ARTIFACTMANIFESTENTRY = DESCRIPTOR.message_types_by_name['ArtifactManifestEntry'] +_EXTRAITEM = DESCRIPTOR.message_types_by_name['ExtraItem'] +_STORAGEPOLICYCONFIGITEM = DESCRIPTOR.message_types_by_name['StoragePolicyConfigItem'] +_ARTIFACTRESULT = DESCRIPTOR.message_types_by_name['ArtifactResult'] +_LINKARTIFACTRESULT = DESCRIPTOR.message_types_by_name['LinkArtifactResult'] +_LINKARTIFACTREQUEST = DESCRIPTOR.message_types_by_name['LinkArtifactRequest'] +_LINKARTIFACTRESPONSE = DESCRIPTOR.message_types_by_name['LinkArtifactResponse'] +_TBRECORD = DESCRIPTOR.message_types_by_name['TBRecord'] +_TBRESULT = DESCRIPTOR.message_types_by_name['TBResult'] +_ALERTRECORD = DESCRIPTOR.message_types_by_name['AlertRecord'] +_ALERTRESULT = DESCRIPTOR.message_types_by_name['AlertResult'] +_REQUEST = DESCRIPTOR.message_types_by_name['Request'] +_RESPONSE = DESCRIPTOR.message_types_by_name['Response'] +_DEFERREQUEST = DESCRIPTOR.message_types_by_name['DeferRequest'] +_PAUSEREQUEST = DESCRIPTOR.message_types_by_name['PauseRequest'] +_PAUSERESPONSE = DESCRIPTOR.message_types_by_name['PauseResponse'] +_RESUMEREQUEST = DESCRIPTOR.message_types_by_name['ResumeRequest'] +_RESUMERESPONSE = DESCRIPTOR.message_types_by_name['ResumeResponse'] +_LOGINREQUEST = DESCRIPTOR.message_types_by_name['LoginRequest'] +_LOGINRESPONSE = DESCRIPTOR.message_types_by_name['LoginResponse'] +_GETSUMMARYREQUEST = DESCRIPTOR.message_types_by_name['GetSummaryRequest'] +_GETSUMMARYRESPONSE = DESCRIPTOR.message_types_by_name['GetSummaryResponse'] +_GETSYSTEMMETRICSREQUEST = DESCRIPTOR.message_types_by_name['GetSystemMetricsRequest'] +_SYSTEMMETRICSAMPLE = DESCRIPTOR.message_types_by_name['SystemMetricSample'] +_SYSTEMMETRICSBUFFER = DESCRIPTOR.message_types_by_name['SystemMetricsBuffer'] +_GETSYSTEMMETRICSRESPONSE = DESCRIPTOR.message_types_by_name['GetSystemMetricsResponse'] +_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY = _GETSYSTEMMETRICSRESPONSE.nested_types_by_name['SystemMetricsEntry'] +_STATUSREQUEST = DESCRIPTOR.message_types_by_name['StatusRequest'] +_STATUSRESPONSE = DESCRIPTOR.message_types_by_name['StatusResponse'] +_STOPSTATUSREQUEST = DESCRIPTOR.message_types_by_name['StopStatusRequest'] +_STOPSTATUSRESPONSE = DESCRIPTOR.message_types_by_name['StopStatusResponse'] +_NETWORKSTATUSREQUEST = DESCRIPTOR.message_types_by_name['NetworkStatusRequest'] +_NETWORKSTATUSRESPONSE = DESCRIPTOR.message_types_by_name['NetworkStatusResponse'] +_HTTPRESPONSE = DESCRIPTOR.message_types_by_name['HttpResponse'] +_INTERNALMESSAGESREQUEST = DESCRIPTOR.message_types_by_name['InternalMessagesRequest'] +_INTERNALMESSAGESRESPONSE = DESCRIPTOR.message_types_by_name['InternalMessagesResponse'] +_INTERNALMESSAGES = DESCRIPTOR.message_types_by_name['InternalMessages'] +_POLLEXITREQUEST = DESCRIPTOR.message_types_by_name['PollExitRequest'] +_POLLEXITRESPONSE = DESCRIPTOR.message_types_by_name['PollExitResponse'] +_OPERATIONSTATSREQUEST = DESCRIPTOR.message_types_by_name['OperationStatsRequest'] +_OPERATIONSTATSRESPONSE = DESCRIPTOR.message_types_by_name['OperationStatsResponse'] +_OPERATIONSTATS = DESCRIPTOR.message_types_by_name['OperationStats'] +_OPERATION = DESCRIPTOR.message_types_by_name['Operation'] +_SENDERMARKREQUEST = DESCRIPTOR.message_types_by_name['SenderMarkRequest'] +_SYNCFINISHREQUEST = DESCRIPTOR.message_types_by_name['SyncFinishRequest'] +_SYNCRESPONSE = DESCRIPTOR.message_types_by_name['SyncResponse'] +_SENDERREADREQUEST = DESCRIPTOR.message_types_by_name['SenderReadRequest'] +_STATUSREPORTREQUEST = DESCRIPTOR.message_types_by_name['StatusReportRequest'] +_SUMMARYRECORDREQUEST = DESCRIPTOR.message_types_by_name['SummaryRecordRequest'] +_TELEMETRYRECORDREQUEST = DESCRIPTOR.message_types_by_name['TelemetryRecordRequest'] +_SERVERINFOREQUEST = DESCRIPTOR.message_types_by_name['ServerInfoRequest'] +_SERVERINFORESPONSE = DESCRIPTOR.message_types_by_name['ServerInfoResponse'] +_SERVERMESSAGES = DESCRIPTOR.message_types_by_name['ServerMessages'] +_SERVERMESSAGE = DESCRIPTOR.message_types_by_name['ServerMessage'] +_FILECOUNTS = DESCRIPTOR.message_types_by_name['FileCounts'] +_FILEPUSHERSTATS = DESCRIPTOR.message_types_by_name['FilePusherStats'] +_FILESUPLOADED = DESCRIPTOR.message_types_by_name['FilesUploaded'] +_FILETRANSFERINFOREQUEST = DESCRIPTOR.message_types_by_name['FileTransferInfoRequest'] +_LOCALINFO = DESCRIPTOR.message_types_by_name['LocalInfo'] +_SHUTDOWNREQUEST = DESCRIPTOR.message_types_by_name['ShutdownRequest'] +_SHUTDOWNRESPONSE = DESCRIPTOR.message_types_by_name['ShutdownResponse'] +_ATTACHREQUEST = DESCRIPTOR.message_types_by_name['AttachRequest'] +_ATTACHRESPONSE = DESCRIPTOR.message_types_by_name['AttachResponse'] +_TESTINJECTREQUEST = DESCRIPTOR.message_types_by_name['TestInjectRequest'] +_TESTINJECTRESPONSE = DESCRIPTOR.message_types_by_name['TestInjectResponse'] +_HISTORYACTION = DESCRIPTOR.message_types_by_name['HistoryAction'] +_PARTIALHISTORYREQUEST = DESCRIPTOR.message_types_by_name['PartialHistoryRequest'] +_PARTIALHISTORYRESPONSE = DESCRIPTOR.message_types_by_name['PartialHistoryResponse'] +_SAMPLEDHISTORYREQUEST = DESCRIPTOR.message_types_by_name['SampledHistoryRequest'] +_SAMPLEDHISTORYITEM = DESCRIPTOR.message_types_by_name['SampledHistoryItem'] +_SAMPLEDHISTORYRESPONSE = DESCRIPTOR.message_types_by_name['SampledHistoryResponse'] +_RUNSTATUSREQUEST = DESCRIPTOR.message_types_by_name['RunStatusRequest'] +_RUNSTATUSRESPONSE = DESCRIPTOR.message_types_by_name['RunStatusResponse'] +_RUNSTARTREQUEST = DESCRIPTOR.message_types_by_name['RunStartRequest'] +_RUNSTARTRESPONSE = DESCRIPTOR.message_types_by_name['RunStartResponse'] +_RUNFINISHWITHOUTEXITREQUEST = DESCRIPTOR.message_types_by_name['RunFinishWithoutExitRequest'] +_RUNFINISHWITHOUTEXITRESPONSE = DESCRIPTOR.message_types_by_name['RunFinishWithoutExitResponse'] +_CHECKVERSIONREQUEST = DESCRIPTOR.message_types_by_name['CheckVersionRequest'] +_CHECKVERSIONRESPONSE = DESCRIPTOR.message_types_by_name['CheckVersionResponse'] +_JOBINFOREQUEST = DESCRIPTOR.message_types_by_name['JobInfoRequest'] +_JOBINFORESPONSE = DESCRIPTOR.message_types_by_name['JobInfoResponse'] +_LOGARTIFACTREQUEST = DESCRIPTOR.message_types_by_name['LogArtifactRequest'] +_LOGARTIFACTRESPONSE = DESCRIPTOR.message_types_by_name['LogArtifactResponse'] +_DOWNLOADARTIFACTREQUEST = DESCRIPTOR.message_types_by_name['DownloadArtifactRequest'] +_DOWNLOADARTIFACTRESPONSE = DESCRIPTOR.message_types_by_name['DownloadArtifactResponse'] +_KEEPALIVEREQUEST = DESCRIPTOR.message_types_by_name['KeepaliveRequest'] +_KEEPALIVERESPONSE = DESCRIPTOR.message_types_by_name['KeepaliveResponse'] +_ARTIFACTINFO = DESCRIPTOR.message_types_by_name['ArtifactInfo'] +_GITINFO = DESCRIPTOR.message_types_by_name['GitInfo'] +_GITSOURCE = DESCRIPTOR.message_types_by_name['GitSource'] +_IMAGESOURCE = DESCRIPTOR.message_types_by_name['ImageSource'] +_SOURCE = DESCRIPTOR.message_types_by_name['Source'] +_JOBSOURCE = DESCRIPTOR.message_types_by_name['JobSource'] +_PARTIALJOBARTIFACT = DESCRIPTOR.message_types_by_name['PartialJobArtifact'] +_USEARTIFACTRECORD = DESCRIPTOR.message_types_by_name['UseArtifactRecord'] +_USEARTIFACTRESULT = DESCRIPTOR.message_types_by_name['UseArtifactResult'] +_CANCELREQUEST = DESCRIPTOR.message_types_by_name['CancelRequest'] +_CANCELRESPONSE = DESCRIPTOR.message_types_by_name['CancelResponse'] +_DISKINFO = DESCRIPTOR.message_types_by_name['DiskInfo'] +_MEMORYINFO = DESCRIPTOR.message_types_by_name['MemoryInfo'] +_CPUINFO = DESCRIPTOR.message_types_by_name['CpuInfo'] +_APPLEINFO = DESCRIPTOR.message_types_by_name['AppleInfo'] +_GPUNVIDIAINFO = DESCRIPTOR.message_types_by_name['GpuNvidiaInfo'] +_GPUAMDINFO = DESCRIPTOR.message_types_by_name['GpuAmdInfo'] +_TRAINIUMINFO = DESCRIPTOR.message_types_by_name['TrainiumInfo'] +_TPUINFO = DESCRIPTOR.message_types_by_name['TPUInfo'] +_COREWEAVEINFO = DESCRIPTOR.message_types_by_name['CoreWeaveInfo'] +_ENVIRONMENTRECORD = DESCRIPTOR.message_types_by_name['EnvironmentRecord'] +_ENVIRONMENTRECORD_DISKENTRY = _ENVIRONMENTRECORD.nested_types_by_name['DiskEntry'] +_ENVIRONMENTRECORD_SLURMENTRY = _ENVIRONMENTRECORD.nested_types_by_name['SlurmEntry'] +_PYTHONPACKAGESREQUEST = DESCRIPTOR.message_types_by_name['PythonPackagesRequest'] +_PYTHONPACKAGESREQUEST_PYTHONPACKAGE = _PYTHONPACKAGESREQUEST.nested_types_by_name['PythonPackage'] +_JOBINPUTPATH = DESCRIPTOR.message_types_by_name['JobInputPath'] +_JOBINPUTSOURCE = DESCRIPTOR.message_types_by_name['JobInputSource'] +_JOBINPUTSOURCE_RUNCONFIGSOURCE = _JOBINPUTSOURCE.nested_types_by_name['RunConfigSource'] +_JOBINPUTSOURCE_CONFIGFILESOURCE = _JOBINPUTSOURCE.nested_types_by_name['ConfigFileSource'] +_JOBINPUTREQUEST = DESCRIPTOR.message_types_by_name['JobInputRequest'] +_SERVERFEATUREREQUEST = DESCRIPTOR.message_types_by_name['ServerFeatureRequest'] +_SERVERFEATURERESPONSE = DESCRIPTOR.message_types_by_name['ServerFeatureResponse'] +_SERVERFEATUREITEM = DESCRIPTOR.message_types_by_name['ServerFeatureItem'] +_ERRORINFO_ERRORCODE = _ERRORINFO.enum_types_by_name['ErrorCode'] +_OUTPUTRECORD_OUTPUTTYPE = _OUTPUTRECORD.enum_types_by_name['OutputType'] +_OUTPUTRAWRECORD_OUTPUTTYPE = _OUTPUTRAWRECORD.enum_types_by_name['OutputType'] +_METRICRECORD_METRICGOAL = _METRICRECORD.enum_types_by_name['MetricGoal'] +_FILESITEM_POLICYTYPE = _FILESITEM.enum_types_by_name['PolicyType'] +_FILESITEM_FILETYPE = _FILESITEM.enum_types_by_name['FileType'] +_STATSRECORD_STATSTYPE = _STATSRECORD.enum_types_by_name['StatsType'] +_DEFERREQUEST_DEFERSTATE = _DEFERREQUEST.enum_types_by_name['DeferState'] +_FILETRANSFERINFOREQUEST_TRANSFERTYPE = _FILETRANSFERINFOREQUEST.enum_types_by_name['TransferType'] +Record = _reflection.GeneratedProtocolMessageType('Record', (_message.Message,), { + 'DESCRIPTOR' : _RECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Record) + }) +_sym_db.RegisterMessage(Record) + +Control = _reflection.GeneratedProtocolMessageType('Control', (_message.Message,), { + 'DESCRIPTOR' : _CONTROL, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Control) + }) +_sym_db.RegisterMessage(Control) + +Result = _reflection.GeneratedProtocolMessageType('Result', (_message.Message,), { + 'DESCRIPTOR' : _RESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Result) + }) +_sym_db.RegisterMessage(Result) + +FinalRecord = _reflection.GeneratedProtocolMessageType('FinalRecord', (_message.Message,), { + 'DESCRIPTOR' : _FINALRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FinalRecord) + }) +_sym_db.RegisterMessage(FinalRecord) + +VersionInfo = _reflection.GeneratedProtocolMessageType('VersionInfo', (_message.Message,), { + 'DESCRIPTOR' : _VERSIONINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.VersionInfo) + }) +_sym_db.RegisterMessage(VersionInfo) + +HeaderRecord = _reflection.GeneratedProtocolMessageType('HeaderRecord', (_message.Message,), { + 'DESCRIPTOR' : _HEADERRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.HeaderRecord) + }) +_sym_db.RegisterMessage(HeaderRecord) + +FooterRecord = _reflection.GeneratedProtocolMessageType('FooterRecord', (_message.Message,), { + 'DESCRIPTOR' : _FOOTERRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FooterRecord) + }) +_sym_db.RegisterMessage(FooterRecord) + +BranchPoint = _reflection.GeneratedProtocolMessageType('BranchPoint', (_message.Message,), { + 'DESCRIPTOR' : _BRANCHPOINT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.BranchPoint) + }) +_sym_db.RegisterMessage(BranchPoint) + +RunRecord = _reflection.GeneratedProtocolMessageType('RunRecord', (_message.Message,), { + 'DESCRIPTOR' : _RUNRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunRecord) + }) +_sym_db.RegisterMessage(RunRecord) + +GitRepoRecord = _reflection.GeneratedProtocolMessageType('GitRepoRecord', (_message.Message,), { + 'DESCRIPTOR' : _GITREPORECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GitRepoRecord) + }) +_sym_db.RegisterMessage(GitRepoRecord) + +RunUpdateResult = _reflection.GeneratedProtocolMessageType('RunUpdateResult', (_message.Message,), { + 'DESCRIPTOR' : _RUNUPDATERESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunUpdateResult) + }) +_sym_db.RegisterMessage(RunUpdateResult) + +ErrorInfo = _reflection.GeneratedProtocolMessageType('ErrorInfo', (_message.Message,), { + 'DESCRIPTOR' : _ERRORINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ErrorInfo) + }) +_sym_db.RegisterMessage(ErrorInfo) + +RunExitRecord = _reflection.GeneratedProtocolMessageType('RunExitRecord', (_message.Message,), { + 'DESCRIPTOR' : _RUNEXITRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunExitRecord) + }) +_sym_db.RegisterMessage(RunExitRecord) + +RunExitResult = _reflection.GeneratedProtocolMessageType('RunExitResult', (_message.Message,), { + 'DESCRIPTOR' : _RUNEXITRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunExitResult) + }) +_sym_db.RegisterMessage(RunExitResult) + +RunPreemptingRecord = _reflection.GeneratedProtocolMessageType('RunPreemptingRecord', (_message.Message,), { + 'DESCRIPTOR' : _RUNPREEMPTINGRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunPreemptingRecord) + }) +_sym_db.RegisterMessage(RunPreemptingRecord) + +RunPreemptingResult = _reflection.GeneratedProtocolMessageType('RunPreemptingResult', (_message.Message,), { + 'DESCRIPTOR' : _RUNPREEMPTINGRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunPreemptingResult) + }) +_sym_db.RegisterMessage(RunPreemptingResult) + +SettingsRecord = _reflection.GeneratedProtocolMessageType('SettingsRecord', (_message.Message,), { + 'DESCRIPTOR' : _SETTINGSRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SettingsRecord) + }) +_sym_db.RegisterMessage(SettingsRecord) + +SettingsItem = _reflection.GeneratedProtocolMessageType('SettingsItem', (_message.Message,), { + 'DESCRIPTOR' : _SETTINGSITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SettingsItem) + }) +_sym_db.RegisterMessage(SettingsItem) + +HistoryStep = _reflection.GeneratedProtocolMessageType('HistoryStep', (_message.Message,), { + 'DESCRIPTOR' : _HISTORYSTEP, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.HistoryStep) + }) +_sym_db.RegisterMessage(HistoryStep) + +HistoryRecord = _reflection.GeneratedProtocolMessageType('HistoryRecord', (_message.Message,), { + 'DESCRIPTOR' : _HISTORYRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.HistoryRecord) + }) +_sym_db.RegisterMessage(HistoryRecord) + +HistoryItem = _reflection.GeneratedProtocolMessageType('HistoryItem', (_message.Message,), { + 'DESCRIPTOR' : _HISTORYITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.HistoryItem) + }) +_sym_db.RegisterMessage(HistoryItem) + +HistoryResult = _reflection.GeneratedProtocolMessageType('HistoryResult', (_message.Message,), { + 'DESCRIPTOR' : _HISTORYRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.HistoryResult) + }) +_sym_db.RegisterMessage(HistoryResult) + +OutputRecord = _reflection.GeneratedProtocolMessageType('OutputRecord', (_message.Message,), { + 'DESCRIPTOR' : _OUTPUTRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.OutputRecord) + }) +_sym_db.RegisterMessage(OutputRecord) + +OutputResult = _reflection.GeneratedProtocolMessageType('OutputResult', (_message.Message,), { + 'DESCRIPTOR' : _OUTPUTRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.OutputResult) + }) +_sym_db.RegisterMessage(OutputResult) + +OutputRawRecord = _reflection.GeneratedProtocolMessageType('OutputRawRecord', (_message.Message,), { + 'DESCRIPTOR' : _OUTPUTRAWRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.OutputRawRecord) + }) +_sym_db.RegisterMessage(OutputRawRecord) + +OutputRawResult = _reflection.GeneratedProtocolMessageType('OutputRawResult', (_message.Message,), { + 'DESCRIPTOR' : _OUTPUTRAWRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.OutputRawResult) + }) +_sym_db.RegisterMessage(OutputRawResult) + +MetricRecord = _reflection.GeneratedProtocolMessageType('MetricRecord', (_message.Message,), { + 'DESCRIPTOR' : _METRICRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MetricRecord) + }) +_sym_db.RegisterMessage(MetricRecord) + +MetricResult = _reflection.GeneratedProtocolMessageType('MetricResult', (_message.Message,), { + 'DESCRIPTOR' : _METRICRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MetricResult) + }) +_sym_db.RegisterMessage(MetricResult) + +MetricOptions = _reflection.GeneratedProtocolMessageType('MetricOptions', (_message.Message,), { + 'DESCRIPTOR' : _METRICOPTIONS, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MetricOptions) + }) +_sym_db.RegisterMessage(MetricOptions) + +MetricControl = _reflection.GeneratedProtocolMessageType('MetricControl', (_message.Message,), { + 'DESCRIPTOR' : _METRICCONTROL, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MetricControl) + }) +_sym_db.RegisterMessage(MetricControl) + +MetricSummary = _reflection.GeneratedProtocolMessageType('MetricSummary', (_message.Message,), { + 'DESCRIPTOR' : _METRICSUMMARY, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MetricSummary) + }) +_sym_db.RegisterMessage(MetricSummary) + +ConfigRecord = _reflection.GeneratedProtocolMessageType('ConfigRecord', (_message.Message,), { + 'DESCRIPTOR' : _CONFIGRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ConfigRecord) + }) +_sym_db.RegisterMessage(ConfigRecord) + +ConfigItem = _reflection.GeneratedProtocolMessageType('ConfigItem', (_message.Message,), { + 'DESCRIPTOR' : _CONFIGITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ConfigItem) + }) +_sym_db.RegisterMessage(ConfigItem) + +ConfigResult = _reflection.GeneratedProtocolMessageType('ConfigResult', (_message.Message,), { + 'DESCRIPTOR' : _CONFIGRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ConfigResult) + }) +_sym_db.RegisterMessage(ConfigResult) + +SummaryRecord = _reflection.GeneratedProtocolMessageType('SummaryRecord', (_message.Message,), { + 'DESCRIPTOR' : _SUMMARYRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SummaryRecord) + }) +_sym_db.RegisterMessage(SummaryRecord) + +SummaryItem = _reflection.GeneratedProtocolMessageType('SummaryItem', (_message.Message,), { + 'DESCRIPTOR' : _SUMMARYITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SummaryItem) + }) +_sym_db.RegisterMessage(SummaryItem) + +SummaryResult = _reflection.GeneratedProtocolMessageType('SummaryResult', (_message.Message,), { + 'DESCRIPTOR' : _SUMMARYRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SummaryResult) + }) +_sym_db.RegisterMessage(SummaryResult) + +FilesRecord = _reflection.GeneratedProtocolMessageType('FilesRecord', (_message.Message,), { + 'DESCRIPTOR' : _FILESRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FilesRecord) + }) +_sym_db.RegisterMessage(FilesRecord) + +FilesItem = _reflection.GeneratedProtocolMessageType('FilesItem', (_message.Message,), { + 'DESCRIPTOR' : _FILESITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FilesItem) + }) +_sym_db.RegisterMessage(FilesItem) + +FilesResult = _reflection.GeneratedProtocolMessageType('FilesResult', (_message.Message,), { + 'DESCRIPTOR' : _FILESRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FilesResult) + }) +_sym_db.RegisterMessage(FilesResult) + +StatsRecord = _reflection.GeneratedProtocolMessageType('StatsRecord', (_message.Message,), { + 'DESCRIPTOR' : _STATSRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.StatsRecord) + }) +_sym_db.RegisterMessage(StatsRecord) + +StatsItem = _reflection.GeneratedProtocolMessageType('StatsItem', (_message.Message,), { + 'DESCRIPTOR' : _STATSITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.StatsItem) + }) +_sym_db.RegisterMessage(StatsItem) + +ArtifactRecord = _reflection.GeneratedProtocolMessageType('ArtifactRecord', (_message.Message,), { + 'DESCRIPTOR' : _ARTIFACTRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ArtifactRecord) + }) +_sym_db.RegisterMessage(ArtifactRecord) + +ArtifactManifest = _reflection.GeneratedProtocolMessageType('ArtifactManifest', (_message.Message,), { + 'DESCRIPTOR' : _ARTIFACTMANIFEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ArtifactManifest) + }) +_sym_db.RegisterMessage(ArtifactManifest) + +ArtifactManifestEntry = _reflection.GeneratedProtocolMessageType('ArtifactManifestEntry', (_message.Message,), { + 'DESCRIPTOR' : _ARTIFACTMANIFESTENTRY, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ArtifactManifestEntry) + }) +_sym_db.RegisterMessage(ArtifactManifestEntry) + +ExtraItem = _reflection.GeneratedProtocolMessageType('ExtraItem', (_message.Message,), { + 'DESCRIPTOR' : _EXTRAITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ExtraItem) + }) +_sym_db.RegisterMessage(ExtraItem) + +StoragePolicyConfigItem = _reflection.GeneratedProtocolMessageType('StoragePolicyConfigItem', (_message.Message,), { + 'DESCRIPTOR' : _STORAGEPOLICYCONFIGITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.StoragePolicyConfigItem) + }) +_sym_db.RegisterMessage(StoragePolicyConfigItem) + +ArtifactResult = _reflection.GeneratedProtocolMessageType('ArtifactResult', (_message.Message,), { + 'DESCRIPTOR' : _ARTIFACTRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ArtifactResult) + }) +_sym_db.RegisterMessage(ArtifactResult) + +LinkArtifactResult = _reflection.GeneratedProtocolMessageType('LinkArtifactResult', (_message.Message,), { + 'DESCRIPTOR' : _LINKARTIFACTRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.LinkArtifactResult) + }) +_sym_db.RegisterMessage(LinkArtifactResult) + +LinkArtifactRequest = _reflection.GeneratedProtocolMessageType('LinkArtifactRequest', (_message.Message,), { + 'DESCRIPTOR' : _LINKARTIFACTREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.LinkArtifactRequest) + }) +_sym_db.RegisterMessage(LinkArtifactRequest) + +LinkArtifactResponse = _reflection.GeneratedProtocolMessageType('LinkArtifactResponse', (_message.Message,), { + 'DESCRIPTOR' : _LINKARTIFACTRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.LinkArtifactResponse) + }) +_sym_db.RegisterMessage(LinkArtifactResponse) + +TBRecord = _reflection.GeneratedProtocolMessageType('TBRecord', (_message.Message,), { + 'DESCRIPTOR' : _TBRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TBRecord) + }) +_sym_db.RegisterMessage(TBRecord) + +TBResult = _reflection.GeneratedProtocolMessageType('TBResult', (_message.Message,), { + 'DESCRIPTOR' : _TBRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TBResult) + }) +_sym_db.RegisterMessage(TBResult) + +AlertRecord = _reflection.GeneratedProtocolMessageType('AlertRecord', (_message.Message,), { + 'DESCRIPTOR' : _ALERTRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.AlertRecord) + }) +_sym_db.RegisterMessage(AlertRecord) + +AlertResult = _reflection.GeneratedProtocolMessageType('AlertResult', (_message.Message,), { + 'DESCRIPTOR' : _ALERTRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.AlertResult) + }) +_sym_db.RegisterMessage(AlertResult) + +Request = _reflection.GeneratedProtocolMessageType('Request', (_message.Message,), { + 'DESCRIPTOR' : _REQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Request) + }) +_sym_db.RegisterMessage(Request) + +Response = _reflection.GeneratedProtocolMessageType('Response', (_message.Message,), { + 'DESCRIPTOR' : _RESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Response) + }) +_sym_db.RegisterMessage(Response) + +DeferRequest = _reflection.GeneratedProtocolMessageType('DeferRequest', (_message.Message,), { + 'DESCRIPTOR' : _DEFERREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.DeferRequest) + }) +_sym_db.RegisterMessage(DeferRequest) + +PauseRequest = _reflection.GeneratedProtocolMessageType('PauseRequest', (_message.Message,), { + 'DESCRIPTOR' : _PAUSEREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PauseRequest) + }) +_sym_db.RegisterMessage(PauseRequest) + +PauseResponse = _reflection.GeneratedProtocolMessageType('PauseResponse', (_message.Message,), { + 'DESCRIPTOR' : _PAUSERESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PauseResponse) + }) +_sym_db.RegisterMessage(PauseResponse) + +ResumeRequest = _reflection.GeneratedProtocolMessageType('ResumeRequest', (_message.Message,), { + 'DESCRIPTOR' : _RESUMEREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ResumeRequest) + }) +_sym_db.RegisterMessage(ResumeRequest) + +ResumeResponse = _reflection.GeneratedProtocolMessageType('ResumeResponse', (_message.Message,), { + 'DESCRIPTOR' : _RESUMERESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ResumeResponse) + }) +_sym_db.RegisterMessage(ResumeResponse) + +LoginRequest = _reflection.GeneratedProtocolMessageType('LoginRequest', (_message.Message,), { + 'DESCRIPTOR' : _LOGINREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.LoginRequest) + }) +_sym_db.RegisterMessage(LoginRequest) + +LoginResponse = _reflection.GeneratedProtocolMessageType('LoginResponse', (_message.Message,), { + 'DESCRIPTOR' : _LOGINRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.LoginResponse) + }) +_sym_db.RegisterMessage(LoginResponse) + +GetSummaryRequest = _reflection.GeneratedProtocolMessageType('GetSummaryRequest', (_message.Message,), { + 'DESCRIPTOR' : _GETSUMMARYREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GetSummaryRequest) + }) +_sym_db.RegisterMessage(GetSummaryRequest) + +GetSummaryResponse = _reflection.GeneratedProtocolMessageType('GetSummaryResponse', (_message.Message,), { + 'DESCRIPTOR' : _GETSUMMARYRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GetSummaryResponse) + }) +_sym_db.RegisterMessage(GetSummaryResponse) + +GetSystemMetricsRequest = _reflection.GeneratedProtocolMessageType('GetSystemMetricsRequest', (_message.Message,), { + 'DESCRIPTOR' : _GETSYSTEMMETRICSREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GetSystemMetricsRequest) + }) +_sym_db.RegisterMessage(GetSystemMetricsRequest) + +SystemMetricSample = _reflection.GeneratedProtocolMessageType('SystemMetricSample', (_message.Message,), { + 'DESCRIPTOR' : _SYSTEMMETRICSAMPLE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SystemMetricSample) + }) +_sym_db.RegisterMessage(SystemMetricSample) + +SystemMetricsBuffer = _reflection.GeneratedProtocolMessageType('SystemMetricsBuffer', (_message.Message,), { + 'DESCRIPTOR' : _SYSTEMMETRICSBUFFER, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SystemMetricsBuffer) + }) +_sym_db.RegisterMessage(SystemMetricsBuffer) + +GetSystemMetricsResponse = _reflection.GeneratedProtocolMessageType('GetSystemMetricsResponse', (_message.Message,), { + + 'SystemMetricsEntry' : _reflection.GeneratedProtocolMessageType('SystemMetricsEntry', (_message.Message,), { + 'DESCRIPTOR' : _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GetSystemMetricsResponse.SystemMetricsEntry) + }) + , + 'DESCRIPTOR' : _GETSYSTEMMETRICSRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GetSystemMetricsResponse) + }) +_sym_db.RegisterMessage(GetSystemMetricsResponse) +_sym_db.RegisterMessage(GetSystemMetricsResponse.SystemMetricsEntry) + +StatusRequest = _reflection.GeneratedProtocolMessageType('StatusRequest', (_message.Message,), { + 'DESCRIPTOR' : _STATUSREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.StatusRequest) + }) +_sym_db.RegisterMessage(StatusRequest) + +StatusResponse = _reflection.GeneratedProtocolMessageType('StatusResponse', (_message.Message,), { + 'DESCRIPTOR' : _STATUSRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.StatusResponse) + }) +_sym_db.RegisterMessage(StatusResponse) + +StopStatusRequest = _reflection.GeneratedProtocolMessageType('StopStatusRequest', (_message.Message,), { + 'DESCRIPTOR' : _STOPSTATUSREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.StopStatusRequest) + }) +_sym_db.RegisterMessage(StopStatusRequest) + +StopStatusResponse = _reflection.GeneratedProtocolMessageType('StopStatusResponse', (_message.Message,), { + 'DESCRIPTOR' : _STOPSTATUSRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.StopStatusResponse) + }) +_sym_db.RegisterMessage(StopStatusResponse) + +NetworkStatusRequest = _reflection.GeneratedProtocolMessageType('NetworkStatusRequest', (_message.Message,), { + 'DESCRIPTOR' : _NETWORKSTATUSREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.NetworkStatusRequest) + }) +_sym_db.RegisterMessage(NetworkStatusRequest) + +NetworkStatusResponse = _reflection.GeneratedProtocolMessageType('NetworkStatusResponse', (_message.Message,), { + 'DESCRIPTOR' : _NETWORKSTATUSRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.NetworkStatusResponse) + }) +_sym_db.RegisterMessage(NetworkStatusResponse) + +HttpResponse = _reflection.GeneratedProtocolMessageType('HttpResponse', (_message.Message,), { + 'DESCRIPTOR' : _HTTPRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.HttpResponse) + }) +_sym_db.RegisterMessage(HttpResponse) + +InternalMessagesRequest = _reflection.GeneratedProtocolMessageType('InternalMessagesRequest', (_message.Message,), { + 'DESCRIPTOR' : _INTERNALMESSAGESREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.InternalMessagesRequest) + }) +_sym_db.RegisterMessage(InternalMessagesRequest) + +InternalMessagesResponse = _reflection.GeneratedProtocolMessageType('InternalMessagesResponse', (_message.Message,), { + 'DESCRIPTOR' : _INTERNALMESSAGESRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.InternalMessagesResponse) + }) +_sym_db.RegisterMessage(InternalMessagesResponse) + +InternalMessages = _reflection.GeneratedProtocolMessageType('InternalMessages', (_message.Message,), { + 'DESCRIPTOR' : _INTERNALMESSAGES, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.InternalMessages) + }) +_sym_db.RegisterMessage(InternalMessages) + +PollExitRequest = _reflection.GeneratedProtocolMessageType('PollExitRequest', (_message.Message,), { + 'DESCRIPTOR' : _POLLEXITREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PollExitRequest) + }) +_sym_db.RegisterMessage(PollExitRequest) + +PollExitResponse = _reflection.GeneratedProtocolMessageType('PollExitResponse', (_message.Message,), { + 'DESCRIPTOR' : _POLLEXITRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PollExitResponse) + }) +_sym_db.RegisterMessage(PollExitResponse) + +OperationStatsRequest = _reflection.GeneratedProtocolMessageType('OperationStatsRequest', (_message.Message,), { + 'DESCRIPTOR' : _OPERATIONSTATSREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.OperationStatsRequest) + }) +_sym_db.RegisterMessage(OperationStatsRequest) + +OperationStatsResponse = _reflection.GeneratedProtocolMessageType('OperationStatsResponse', (_message.Message,), { + 'DESCRIPTOR' : _OPERATIONSTATSRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.OperationStatsResponse) + }) +_sym_db.RegisterMessage(OperationStatsResponse) + +OperationStats = _reflection.GeneratedProtocolMessageType('OperationStats', (_message.Message,), { + 'DESCRIPTOR' : _OPERATIONSTATS, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.OperationStats) + }) +_sym_db.RegisterMessage(OperationStats) + +Operation = _reflection.GeneratedProtocolMessageType('Operation', (_message.Message,), { + 'DESCRIPTOR' : _OPERATION, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Operation) + }) +_sym_db.RegisterMessage(Operation) + +SenderMarkRequest = _reflection.GeneratedProtocolMessageType('SenderMarkRequest', (_message.Message,), { + 'DESCRIPTOR' : _SENDERMARKREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SenderMarkRequest) + }) +_sym_db.RegisterMessage(SenderMarkRequest) + +SyncFinishRequest = _reflection.GeneratedProtocolMessageType('SyncFinishRequest', (_message.Message,), { + 'DESCRIPTOR' : _SYNCFINISHREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SyncFinishRequest) + }) +_sym_db.RegisterMessage(SyncFinishRequest) + +SyncResponse = _reflection.GeneratedProtocolMessageType('SyncResponse', (_message.Message,), { + 'DESCRIPTOR' : _SYNCRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SyncResponse) + }) +_sym_db.RegisterMessage(SyncResponse) + +SenderReadRequest = _reflection.GeneratedProtocolMessageType('SenderReadRequest', (_message.Message,), { + 'DESCRIPTOR' : _SENDERREADREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SenderReadRequest) + }) +_sym_db.RegisterMessage(SenderReadRequest) + +StatusReportRequest = _reflection.GeneratedProtocolMessageType('StatusReportRequest', (_message.Message,), { + 'DESCRIPTOR' : _STATUSREPORTREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.StatusReportRequest) + }) +_sym_db.RegisterMessage(StatusReportRequest) + +SummaryRecordRequest = _reflection.GeneratedProtocolMessageType('SummaryRecordRequest', (_message.Message,), { + 'DESCRIPTOR' : _SUMMARYRECORDREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SummaryRecordRequest) + }) +_sym_db.RegisterMessage(SummaryRecordRequest) + +TelemetryRecordRequest = _reflection.GeneratedProtocolMessageType('TelemetryRecordRequest', (_message.Message,), { + 'DESCRIPTOR' : _TELEMETRYRECORDREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TelemetryRecordRequest) + }) +_sym_db.RegisterMessage(TelemetryRecordRequest) + +ServerInfoRequest = _reflection.GeneratedProtocolMessageType('ServerInfoRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFOREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInfoRequest) + }) +_sym_db.RegisterMessage(ServerInfoRequest) + +ServerInfoResponse = _reflection.GeneratedProtocolMessageType('ServerInfoResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInfoResponse) + }) +_sym_db.RegisterMessage(ServerInfoResponse) + +ServerMessages = _reflection.GeneratedProtocolMessageType('ServerMessages', (_message.Message,), { + 'DESCRIPTOR' : _SERVERMESSAGES, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerMessages) + }) +_sym_db.RegisterMessage(ServerMessages) + +ServerMessage = _reflection.GeneratedProtocolMessageType('ServerMessage', (_message.Message,), { + 'DESCRIPTOR' : _SERVERMESSAGE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerMessage) + }) +_sym_db.RegisterMessage(ServerMessage) + +FileCounts = _reflection.GeneratedProtocolMessageType('FileCounts', (_message.Message,), { + 'DESCRIPTOR' : _FILECOUNTS, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FileCounts) + }) +_sym_db.RegisterMessage(FileCounts) + +FilePusherStats = _reflection.GeneratedProtocolMessageType('FilePusherStats', (_message.Message,), { + 'DESCRIPTOR' : _FILEPUSHERSTATS, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FilePusherStats) + }) +_sym_db.RegisterMessage(FilePusherStats) + +FilesUploaded = _reflection.GeneratedProtocolMessageType('FilesUploaded', (_message.Message,), { + 'DESCRIPTOR' : _FILESUPLOADED, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FilesUploaded) + }) +_sym_db.RegisterMessage(FilesUploaded) + +FileTransferInfoRequest = _reflection.GeneratedProtocolMessageType('FileTransferInfoRequest', (_message.Message,), { + 'DESCRIPTOR' : _FILETRANSFERINFOREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.FileTransferInfoRequest) + }) +_sym_db.RegisterMessage(FileTransferInfoRequest) + +LocalInfo = _reflection.GeneratedProtocolMessageType('LocalInfo', (_message.Message,), { + 'DESCRIPTOR' : _LOCALINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.LocalInfo) + }) +_sym_db.RegisterMessage(LocalInfo) + +ShutdownRequest = _reflection.GeneratedProtocolMessageType('ShutdownRequest', (_message.Message,), { + 'DESCRIPTOR' : _SHUTDOWNREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ShutdownRequest) + }) +_sym_db.RegisterMessage(ShutdownRequest) + +ShutdownResponse = _reflection.GeneratedProtocolMessageType('ShutdownResponse', (_message.Message,), { + 'DESCRIPTOR' : _SHUTDOWNRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ShutdownResponse) + }) +_sym_db.RegisterMessage(ShutdownResponse) + +AttachRequest = _reflection.GeneratedProtocolMessageType('AttachRequest', (_message.Message,), { + 'DESCRIPTOR' : _ATTACHREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.AttachRequest) + }) +_sym_db.RegisterMessage(AttachRequest) + +AttachResponse = _reflection.GeneratedProtocolMessageType('AttachResponse', (_message.Message,), { + 'DESCRIPTOR' : _ATTACHRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.AttachResponse) + }) +_sym_db.RegisterMessage(AttachResponse) + +TestInjectRequest = _reflection.GeneratedProtocolMessageType('TestInjectRequest', (_message.Message,), { + 'DESCRIPTOR' : _TESTINJECTREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TestInjectRequest) + }) +_sym_db.RegisterMessage(TestInjectRequest) + +TestInjectResponse = _reflection.GeneratedProtocolMessageType('TestInjectResponse', (_message.Message,), { + 'DESCRIPTOR' : _TESTINJECTRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TestInjectResponse) + }) +_sym_db.RegisterMessage(TestInjectResponse) + +HistoryAction = _reflection.GeneratedProtocolMessageType('HistoryAction', (_message.Message,), { + 'DESCRIPTOR' : _HISTORYACTION, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.HistoryAction) + }) +_sym_db.RegisterMessage(HistoryAction) + +PartialHistoryRequest = _reflection.GeneratedProtocolMessageType('PartialHistoryRequest', (_message.Message,), { + 'DESCRIPTOR' : _PARTIALHISTORYREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PartialHistoryRequest) + }) +_sym_db.RegisterMessage(PartialHistoryRequest) + +PartialHistoryResponse = _reflection.GeneratedProtocolMessageType('PartialHistoryResponse', (_message.Message,), { + 'DESCRIPTOR' : _PARTIALHISTORYRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PartialHistoryResponse) + }) +_sym_db.RegisterMessage(PartialHistoryResponse) + +SampledHistoryRequest = _reflection.GeneratedProtocolMessageType('SampledHistoryRequest', (_message.Message,), { + 'DESCRIPTOR' : _SAMPLEDHISTORYREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SampledHistoryRequest) + }) +_sym_db.RegisterMessage(SampledHistoryRequest) + +SampledHistoryItem = _reflection.GeneratedProtocolMessageType('SampledHistoryItem', (_message.Message,), { + 'DESCRIPTOR' : _SAMPLEDHISTORYITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SampledHistoryItem) + }) +_sym_db.RegisterMessage(SampledHistoryItem) + +SampledHistoryResponse = _reflection.GeneratedProtocolMessageType('SampledHistoryResponse', (_message.Message,), { + 'DESCRIPTOR' : _SAMPLEDHISTORYRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.SampledHistoryResponse) + }) +_sym_db.RegisterMessage(SampledHistoryResponse) + +RunStatusRequest = _reflection.GeneratedProtocolMessageType('RunStatusRequest', (_message.Message,), { + 'DESCRIPTOR' : _RUNSTATUSREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunStatusRequest) + }) +_sym_db.RegisterMessage(RunStatusRequest) + +RunStatusResponse = _reflection.GeneratedProtocolMessageType('RunStatusResponse', (_message.Message,), { + 'DESCRIPTOR' : _RUNSTATUSRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunStatusResponse) + }) +_sym_db.RegisterMessage(RunStatusResponse) + +RunStartRequest = _reflection.GeneratedProtocolMessageType('RunStartRequest', (_message.Message,), { + 'DESCRIPTOR' : _RUNSTARTREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunStartRequest) + }) +_sym_db.RegisterMessage(RunStartRequest) + +RunStartResponse = _reflection.GeneratedProtocolMessageType('RunStartResponse', (_message.Message,), { + 'DESCRIPTOR' : _RUNSTARTRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunStartResponse) + }) +_sym_db.RegisterMessage(RunStartResponse) + +RunFinishWithoutExitRequest = _reflection.GeneratedProtocolMessageType('RunFinishWithoutExitRequest', (_message.Message,), { + 'DESCRIPTOR' : _RUNFINISHWITHOUTEXITREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunFinishWithoutExitRequest) + }) +_sym_db.RegisterMessage(RunFinishWithoutExitRequest) + +RunFinishWithoutExitResponse = _reflection.GeneratedProtocolMessageType('RunFinishWithoutExitResponse', (_message.Message,), { + 'DESCRIPTOR' : _RUNFINISHWITHOUTEXITRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunFinishWithoutExitResponse) + }) +_sym_db.RegisterMessage(RunFinishWithoutExitResponse) + +CheckVersionRequest = _reflection.GeneratedProtocolMessageType('CheckVersionRequest', (_message.Message,), { + 'DESCRIPTOR' : _CHECKVERSIONREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.CheckVersionRequest) + }) +_sym_db.RegisterMessage(CheckVersionRequest) + +CheckVersionResponse = _reflection.GeneratedProtocolMessageType('CheckVersionResponse', (_message.Message,), { + 'DESCRIPTOR' : _CHECKVERSIONRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.CheckVersionResponse) + }) +_sym_db.RegisterMessage(CheckVersionResponse) + +JobInfoRequest = _reflection.GeneratedProtocolMessageType('JobInfoRequest', (_message.Message,), { + 'DESCRIPTOR' : _JOBINFOREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.JobInfoRequest) + }) +_sym_db.RegisterMessage(JobInfoRequest) + +JobInfoResponse = _reflection.GeneratedProtocolMessageType('JobInfoResponse', (_message.Message,), { + 'DESCRIPTOR' : _JOBINFORESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.JobInfoResponse) + }) +_sym_db.RegisterMessage(JobInfoResponse) + +LogArtifactRequest = _reflection.GeneratedProtocolMessageType('LogArtifactRequest', (_message.Message,), { + 'DESCRIPTOR' : _LOGARTIFACTREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.LogArtifactRequest) + }) +_sym_db.RegisterMessage(LogArtifactRequest) + +LogArtifactResponse = _reflection.GeneratedProtocolMessageType('LogArtifactResponse', (_message.Message,), { + 'DESCRIPTOR' : _LOGARTIFACTRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.LogArtifactResponse) + }) +_sym_db.RegisterMessage(LogArtifactResponse) + +DownloadArtifactRequest = _reflection.GeneratedProtocolMessageType('DownloadArtifactRequest', (_message.Message,), { + 'DESCRIPTOR' : _DOWNLOADARTIFACTREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.DownloadArtifactRequest) + }) +_sym_db.RegisterMessage(DownloadArtifactRequest) + +DownloadArtifactResponse = _reflection.GeneratedProtocolMessageType('DownloadArtifactResponse', (_message.Message,), { + 'DESCRIPTOR' : _DOWNLOADARTIFACTRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.DownloadArtifactResponse) + }) +_sym_db.RegisterMessage(DownloadArtifactResponse) + +KeepaliveRequest = _reflection.GeneratedProtocolMessageType('KeepaliveRequest', (_message.Message,), { + 'DESCRIPTOR' : _KEEPALIVEREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.KeepaliveRequest) + }) +_sym_db.RegisterMessage(KeepaliveRequest) + +KeepaliveResponse = _reflection.GeneratedProtocolMessageType('KeepaliveResponse', (_message.Message,), { + 'DESCRIPTOR' : _KEEPALIVERESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.KeepaliveResponse) + }) +_sym_db.RegisterMessage(KeepaliveResponse) + +ArtifactInfo = _reflection.GeneratedProtocolMessageType('ArtifactInfo', (_message.Message,), { + 'DESCRIPTOR' : _ARTIFACTINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ArtifactInfo) + }) +_sym_db.RegisterMessage(ArtifactInfo) + +GitInfo = _reflection.GeneratedProtocolMessageType('GitInfo', (_message.Message,), { + 'DESCRIPTOR' : _GITINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GitInfo) + }) +_sym_db.RegisterMessage(GitInfo) + +GitSource = _reflection.GeneratedProtocolMessageType('GitSource', (_message.Message,), { + 'DESCRIPTOR' : _GITSOURCE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GitSource) + }) +_sym_db.RegisterMessage(GitSource) + +ImageSource = _reflection.GeneratedProtocolMessageType('ImageSource', (_message.Message,), { + 'DESCRIPTOR' : _IMAGESOURCE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ImageSource) + }) +_sym_db.RegisterMessage(ImageSource) + +Source = _reflection.GeneratedProtocolMessageType('Source', (_message.Message,), { + 'DESCRIPTOR' : _SOURCE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Source) + }) +_sym_db.RegisterMessage(Source) + +JobSource = _reflection.GeneratedProtocolMessageType('JobSource', (_message.Message,), { + 'DESCRIPTOR' : _JOBSOURCE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.JobSource) + }) +_sym_db.RegisterMessage(JobSource) + +PartialJobArtifact = _reflection.GeneratedProtocolMessageType('PartialJobArtifact', (_message.Message,), { + 'DESCRIPTOR' : _PARTIALJOBARTIFACT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PartialJobArtifact) + }) +_sym_db.RegisterMessage(PartialJobArtifact) + +UseArtifactRecord = _reflection.GeneratedProtocolMessageType('UseArtifactRecord', (_message.Message,), { + 'DESCRIPTOR' : _USEARTIFACTRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.UseArtifactRecord) + }) +_sym_db.RegisterMessage(UseArtifactRecord) + +UseArtifactResult = _reflection.GeneratedProtocolMessageType('UseArtifactResult', (_message.Message,), { + 'DESCRIPTOR' : _USEARTIFACTRESULT, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.UseArtifactResult) + }) +_sym_db.RegisterMessage(UseArtifactResult) + +CancelRequest = _reflection.GeneratedProtocolMessageType('CancelRequest', (_message.Message,), { + 'DESCRIPTOR' : _CANCELREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.CancelRequest) + }) +_sym_db.RegisterMessage(CancelRequest) + +CancelResponse = _reflection.GeneratedProtocolMessageType('CancelResponse', (_message.Message,), { + 'DESCRIPTOR' : _CANCELRESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.CancelResponse) + }) +_sym_db.RegisterMessage(CancelResponse) + +DiskInfo = _reflection.GeneratedProtocolMessageType('DiskInfo', (_message.Message,), { + 'DESCRIPTOR' : _DISKINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.DiskInfo) + }) +_sym_db.RegisterMessage(DiskInfo) + +MemoryInfo = _reflection.GeneratedProtocolMessageType('MemoryInfo', (_message.Message,), { + 'DESCRIPTOR' : _MEMORYINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MemoryInfo) + }) +_sym_db.RegisterMessage(MemoryInfo) + +CpuInfo = _reflection.GeneratedProtocolMessageType('CpuInfo', (_message.Message,), { + 'DESCRIPTOR' : _CPUINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.CpuInfo) + }) +_sym_db.RegisterMessage(CpuInfo) + +AppleInfo = _reflection.GeneratedProtocolMessageType('AppleInfo', (_message.Message,), { + 'DESCRIPTOR' : _APPLEINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.AppleInfo) + }) +_sym_db.RegisterMessage(AppleInfo) + +GpuNvidiaInfo = _reflection.GeneratedProtocolMessageType('GpuNvidiaInfo', (_message.Message,), { + 'DESCRIPTOR' : _GPUNVIDIAINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GpuNvidiaInfo) + }) +_sym_db.RegisterMessage(GpuNvidiaInfo) + +GpuAmdInfo = _reflection.GeneratedProtocolMessageType('GpuAmdInfo', (_message.Message,), { + 'DESCRIPTOR' : _GPUAMDINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.GpuAmdInfo) + }) +_sym_db.RegisterMessage(GpuAmdInfo) + +TrainiumInfo = _reflection.GeneratedProtocolMessageType('TrainiumInfo', (_message.Message,), { + 'DESCRIPTOR' : _TRAINIUMINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TrainiumInfo) + }) +_sym_db.RegisterMessage(TrainiumInfo) + +TPUInfo = _reflection.GeneratedProtocolMessageType('TPUInfo', (_message.Message,), { + 'DESCRIPTOR' : _TPUINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TPUInfo) + }) +_sym_db.RegisterMessage(TPUInfo) + +CoreWeaveInfo = _reflection.GeneratedProtocolMessageType('CoreWeaveInfo', (_message.Message,), { + 'DESCRIPTOR' : _COREWEAVEINFO, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.CoreWeaveInfo) + }) +_sym_db.RegisterMessage(CoreWeaveInfo) + +EnvironmentRecord = _reflection.GeneratedProtocolMessageType('EnvironmentRecord', (_message.Message,), { + + 'DiskEntry' : _reflection.GeneratedProtocolMessageType('DiskEntry', (_message.Message,), { + 'DESCRIPTOR' : _ENVIRONMENTRECORD_DISKENTRY, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.EnvironmentRecord.DiskEntry) + }) + , + + 'SlurmEntry' : _reflection.GeneratedProtocolMessageType('SlurmEntry', (_message.Message,), { + 'DESCRIPTOR' : _ENVIRONMENTRECORD_SLURMENTRY, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.EnvironmentRecord.SlurmEntry) + }) + , + 'DESCRIPTOR' : _ENVIRONMENTRECORD, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.EnvironmentRecord) + }) +_sym_db.RegisterMessage(EnvironmentRecord) +_sym_db.RegisterMessage(EnvironmentRecord.DiskEntry) +_sym_db.RegisterMessage(EnvironmentRecord.SlurmEntry) + +PythonPackagesRequest = _reflection.GeneratedProtocolMessageType('PythonPackagesRequest', (_message.Message,), { + + 'PythonPackage' : _reflection.GeneratedProtocolMessageType('PythonPackage', (_message.Message,), { + 'DESCRIPTOR' : _PYTHONPACKAGESREQUEST_PYTHONPACKAGE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PythonPackagesRequest.PythonPackage) + }) + , + 'DESCRIPTOR' : _PYTHONPACKAGESREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.PythonPackagesRequest) + }) +_sym_db.RegisterMessage(PythonPackagesRequest) +_sym_db.RegisterMessage(PythonPackagesRequest.PythonPackage) + +JobInputPath = _reflection.GeneratedProtocolMessageType('JobInputPath', (_message.Message,), { + 'DESCRIPTOR' : _JOBINPUTPATH, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.JobInputPath) + }) +_sym_db.RegisterMessage(JobInputPath) + +JobInputSource = _reflection.GeneratedProtocolMessageType('JobInputSource', (_message.Message,), { + + 'RunConfigSource' : _reflection.GeneratedProtocolMessageType('RunConfigSource', (_message.Message,), { + 'DESCRIPTOR' : _JOBINPUTSOURCE_RUNCONFIGSOURCE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.JobInputSource.RunConfigSource) + }) + , + + 'ConfigFileSource' : _reflection.GeneratedProtocolMessageType('ConfigFileSource', (_message.Message,), { + 'DESCRIPTOR' : _JOBINPUTSOURCE_CONFIGFILESOURCE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.JobInputSource.ConfigFileSource) + }) + , + 'DESCRIPTOR' : _JOBINPUTSOURCE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.JobInputSource) + }) +_sym_db.RegisterMessage(JobInputSource) +_sym_db.RegisterMessage(JobInputSource.RunConfigSource) +_sym_db.RegisterMessage(JobInputSource.ConfigFileSource) + +JobInputRequest = _reflection.GeneratedProtocolMessageType('JobInputRequest', (_message.Message,), { + 'DESCRIPTOR' : _JOBINPUTREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.JobInputRequest) + }) +_sym_db.RegisterMessage(JobInputRequest) + +ServerFeatureRequest = _reflection.GeneratedProtocolMessageType('ServerFeatureRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERFEATUREREQUEST, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerFeatureRequest) + }) +_sym_db.RegisterMessage(ServerFeatureRequest) + +ServerFeatureResponse = _reflection.GeneratedProtocolMessageType('ServerFeatureResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERFEATURERESPONSE, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerFeatureResponse) + }) +_sym_db.RegisterMessage(ServerFeatureResponse) + +ServerFeatureItem = _reflection.GeneratedProtocolMessageType('ServerFeatureItem', (_message.Message,), { + 'DESCRIPTOR' : _SERVERFEATUREITEM, + '__module__' : 'wandb.proto.wandb_internal_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerFeatureItem) + }) +_sym_db.RegisterMessage(ServerFeatureItem) + +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY._options = None + _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY._serialized_options = b'8\001' + _ENVIRONMENTRECORD_DISKENTRY._options = None + _ENVIRONMENTRECORD_DISKENTRY._serialized_options = b'8\001' + _ENVIRONMENTRECORD_SLURMENTRY._options = None + _ENVIRONMENTRECORD_SLURMENTRY._serialized_options = b'8\001' + _SERVERFEATURE._serialized_start=22228 + _SERVERFEATURE._serialized_end=22740 + _RECORD._serialized_start=180 + _RECORD._serialized_end=1411 + _CONTROL._serialized_start=1414 + _CONTROL._serialized_end=1582 + _RESULT._serialized_start=1585 + _RESULT._serialized_end=2084 + _FINALRECORD._serialized_start=2086 + _FINALRECORD._serialized_end=2144 + _VERSIONINFO._serialized_start=2146 + _VERSIONINFO._serialized_end=2244 + _HEADERRECORD._serialized_start=2246 + _HEADERRECORD._serialized_end=2356 + _FOOTERRECORD._serialized_start=2358 + _FOOTERRECORD._serialized_end=2417 + _BRANCHPOINT._serialized_start=2419 + _BRANCHPOINT._serialized_end=2476 + _RUNRECORD._serialized_start=2479 + _RUNRECORD._serialized_end=3136 + _GITREPORECORD._serialized_start=3138 + _GITREPORECORD._serialized_end=3197 + _RUNUPDATERESULT._serialized_start=3199 + _RUNUPDATERESULT._serialized_end=3298 + _ERRORINFO._serialized_start=3301 + _ERRORINFO._serialized_end=3473 + _ERRORINFO_ERRORCODE._serialized_start=3382 + _ERRORINFO_ERRORCODE._serialized_end=3473 + _RUNEXITRECORD._serialized_start=3475 + _RUNEXITRECORD._serialized_end=3571 + _RUNEXITRESULT._serialized_start=3573 + _RUNEXITRESULT._serialized_end=3588 + _RUNPREEMPTINGRECORD._serialized_start=3590 + _RUNPREEMPTINGRECORD._serialized_end=3656 + _RUNPREEMPTINGRESULT._serialized_start=3658 + _RUNPREEMPTINGRESULT._serialized_end=3679 + _SETTINGSRECORD._serialized_start=3681 + _SETTINGSRECORD._serialized_end=3786 + _SETTINGSITEM._serialized_start=3788 + _SETTINGSITEM._serialized_end=3835 + _HISTORYSTEP._serialized_start=3837 + _HISTORYSTEP._serialized_end=3863 + _HISTORYRECORD._serialized_start=3866 + _HISTORYRECORD._serialized_end=4012 + _HISTORYITEM._serialized_start=4014 + _HISTORYITEM._serialized_end=4080 + _HISTORYRESULT._serialized_start=4082 + _HISTORYRESULT._serialized_end=4097 + _OUTPUTRECORD._serialized_start=4100 + _OUTPUTRECORD._serialized_end=4320 + _OUTPUTRECORD_OUTPUTTYPE._serialized_start=4284 + _OUTPUTRECORD_OUTPUTTYPE._serialized_end=4320 + _OUTPUTRESULT._serialized_start=4322 + _OUTPUTRESULT._serialized_end=4336 + _OUTPUTRAWRECORD._serialized_start=4339 + _OUTPUTRAWRECORD._serialized_end=4565 + _OUTPUTRAWRECORD_OUTPUTTYPE._serialized_start=4284 + _OUTPUTRAWRECORD_OUTPUTTYPE._serialized_end=4320 + _OUTPUTRAWRESULT._serialized_start=4567 + _OUTPUTRAWRESULT._serialized_end=4584 + _METRICRECORD._serialized_start=4587 + _METRICRECORD._serialized_end=5023 + _METRICRECORD_METRICGOAL._serialized_start=4957 + _METRICRECORD_METRICGOAL._serialized_end=5023 + _METRICRESULT._serialized_start=5025 + _METRICRESULT._serialized_end=5039 + _METRICOPTIONS._serialized_start=5041 + _METRICOPTIONS._serialized_end=5108 + _METRICCONTROL._serialized_start=5110 + _METRICCONTROL._serialized_end=5144 + _METRICSUMMARY._serialized_start=5146 + _METRICSUMMARY._serialized_end=5257 + _CONFIGRECORD._serialized_start=5260 + _CONFIGRECORD._serialized_end=5407 + _CONFIGITEM._serialized_start=5409 + _CONFIGITEM._serialized_end=5474 + _CONFIGRESULT._serialized_start=5476 + _CONFIGRESULT._serialized_end=5490 + _SUMMARYRECORD._serialized_start=5493 + _SUMMARYRECORD._serialized_end=5643 + _SUMMARYITEM._serialized_start=5645 + _SUMMARYITEM._serialized_end=5711 + _SUMMARYRESULT._serialized_start=5713 + _SUMMARYRESULT._serialized_end=5728 + _FILESRECORD._serialized_start=5730 + _FILESRECORD._serialized_end=5830 + _FILESITEM._serialized_start=5833 + _FILESITEM._serialized_end=6069 + _FILESITEM_POLICYTYPE._serialized_start=5964 + _FILESITEM_POLICYTYPE._serialized_end=6004 + _FILESITEM_FILETYPE._serialized_start=6006 + _FILESITEM_FILETYPE._serialized_end=6063 + _FILESRESULT._serialized_start=6071 + _FILESRESULT._serialized_end=6084 + _STATSRECORD._serialized_start=6087 + _STATSRECORD._serialized_end=6317 + _STATSRECORD_STATSTYPE._serialized_start=6294 + _STATSRECORD_STATSTYPE._serialized_end=6317 + _STATSITEM._serialized_start=6319 + _STATSITEM._serialized_end=6363 + _ARTIFACTRECORD._serialized_start=6366 + _ARTIFACTRECORD._serialized_end=6853 + _ARTIFACTMANIFEST._serialized_start=6856 + _ARTIFACTMANIFEST._serialized_end=7072 + _ARTIFACTMANIFESTENTRY._serialized_start=7075 + _ARTIFACTMANIFESTENTRY._serialized_end=7282 + _EXTRAITEM._serialized_start=7284 + _EXTRAITEM._serialized_end=7328 + _STORAGEPOLICYCONFIGITEM._serialized_start=7330 + _STORAGEPOLICYCONFIGITEM._serialized_end=7388 + _ARTIFACTRESULT._serialized_start=7390 + _ARTIFACTRESULT._serialized_end=7406 + _LINKARTIFACTRESULT._serialized_start=7408 + _LINKARTIFACTRESULT._serialized_end=7428 + _LINKARTIFACTREQUEST._serialized_start=7431 + _LINKARTIFACTREQUEST._serialized_end=7671 + _LINKARTIFACTRESPONSE._serialized_start=7673 + _LINKARTIFACTRESPONSE._serialized_end=7764 + _TBRECORD._serialized_start=7766 + _TBRECORD._serialized_end=7870 + _TBRESULT._serialized_start=7872 + _TBRESULT._serialized_end=7882 + _ALERTRECORD._serialized_start=7884 + _ALERTRECORD._serialized_end=8009 + _ALERTRESULT._serialized_start=8011 + _ALERTRESULT._serialized_end=8024 + _REQUEST._serialized_start=8027 + _REQUEST._serialized_end=10254 + _RESPONSE._serialized_start=10257 + _RESPONSE._serialized_end=12076 + _DEFERREQUEST._serialized_start=12079 + _DEFERREQUEST._serialized_end=12399 + _DEFERREQUEST_DEFERSTATE._serialized_start=12152 + _DEFERREQUEST_DEFERSTATE._serialized_end=12399 + _PAUSEREQUEST._serialized_start=12401 + _PAUSEREQUEST._serialized_end=12461 + _PAUSERESPONSE._serialized_start=12463 + _PAUSERESPONSE._serialized_end=12478 + _RESUMEREQUEST._serialized_start=12480 + _RESUMEREQUEST._serialized_end=12541 + _RESUMERESPONSE._serialized_start=12543 + _RESUMERESPONSE._serialized_end=12559 + _LOGINREQUEST._serialized_start=12561 + _LOGINREQUEST._serialized_end=12638 + _LOGINRESPONSE._serialized_start=12640 + _LOGINRESPONSE._serialized_end=12678 + _GETSUMMARYREQUEST._serialized_start=12680 + _GETSUMMARYREQUEST._serialized_end=12745 + _GETSUMMARYRESPONSE._serialized_start=12747 + _GETSUMMARYRESPONSE._serialized_end=12810 + _GETSYSTEMMETRICSREQUEST._serialized_start=12812 + _GETSYSTEMMETRICSREQUEST._serialized_end=12883 + _SYSTEMMETRICSAMPLE._serialized_start=12885 + _SYSTEMMETRICSAMPLE._serialized_end=12967 + _SYSTEMMETRICSBUFFER._serialized_start=12969 + _SYSTEMMETRICSBUFFER._serialized_end=13042 + _GETSYSTEMMETRICSRESPONSE._serialized_start=13045 + _GETSYSTEMMETRICSRESPONSE._serialized_end=13247 + _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY._serialized_start=13158 + _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY._serialized_end=13247 + _STATUSREQUEST._serialized_start=13249 + _STATUSREQUEST._serialized_end=13310 + _STATUSRESPONSE._serialized_start=13312 + _STATUSRESPONSE._serialized_end=13353 + _STOPSTATUSREQUEST._serialized_start=13355 + _STOPSTATUSREQUEST._serialized_end=13420 + _STOPSTATUSRESPONSE._serialized_start=13422 + _STOPSTATUSRESPONSE._serialized_end=13467 + _NETWORKSTATUSREQUEST._serialized_start=13469 + _NETWORKSTATUSREQUEST._serialized_end=13537 + _NETWORKSTATUSRESPONSE._serialized_start=13539 + _NETWORKSTATUSRESPONSE._serialized_end=13619 + _HTTPRESPONSE._serialized_start=13621 + _HTTPRESPONSE._serialized_end=13689 + _INTERNALMESSAGESREQUEST._serialized_start=13691 + _INTERNALMESSAGESREQUEST._serialized_end=13762 + _INTERNALMESSAGESRESPONSE._serialized_start=13764 + _INTERNALMESSAGESRESPONSE._serialized_end=13842 + _INTERNALMESSAGES._serialized_start=13844 + _INTERNALMESSAGES._serialized_end=13879 + _POLLEXITREQUEST._serialized_start=13881 + _POLLEXITREQUEST._serialized_end=13944 + _POLLEXITRESPONSE._serialized_start=13947 + _POLLEXITRESPONSE._serialized_end=14192 + _OPERATIONSTATSREQUEST._serialized_start=14194 + _OPERATIONSTATSREQUEST._serialized_end=14263 + _OPERATIONSTATSRESPONSE._serialized_start=14265 + _OPERATIONSTATSRESPONSE._serialized_end=14346 + _OPERATIONSTATS._serialized_start=14348 + _OPERATIONSTATS._serialized_end=14437 + _OPERATION._serialized_start=14440 + _OPERATION._serialized_end=14575 + _SENDERMARKREQUEST._serialized_start=14577 + _SENDERMARKREQUEST._serialized_end=14596 + _SYNCFINISHREQUEST._serialized_start=14598 + _SYNCFINISHREQUEST._serialized_end=14617 + _SYNCRESPONSE._serialized_start=14619 + _SYNCRESPONSE._serialized_end=14688 + _SENDERREADREQUEST._serialized_start=14690 + _SENDERREADREQUEST._serialized_end=14753 + _STATUSREPORTREQUEST._serialized_start=14755 + _STATUSREPORTREQUEST._serialized_end=14864 + _SUMMARYRECORDREQUEST._serialized_start=14866 + _SUMMARYRECORDREQUEST._serialized_end=14936 + _TELEMETRYRECORDREQUEST._serialized_start=14938 + _TELEMETRYRECORDREQUEST._serialized_end=15014 + _SERVERINFOREQUEST._serialized_start=15016 + _SERVERINFOREQUEST._serialized_end=15081 + _SERVERINFORESPONSE._serialized_start=15083 + _SERVERINFORESPONSE._serialized_end=15207 + _SERVERMESSAGES._serialized_start=15209 + _SERVERMESSAGES._serialized_end=15270 + _SERVERMESSAGE._serialized_start=15272 + _SERVERMESSAGE._serialized_end=15373 + _FILECOUNTS._serialized_start=15375 + _FILECOUNTS._serialized_end=15474 + _FILEPUSHERSTATS._serialized_start=15476 + _FILEPUSHERSTATS._serialized_end=15561 + _FILESUPLOADED._serialized_start=15563 + _FILESUPLOADED._serialized_end=15593 + _FILETRANSFERINFOREQUEST._serialized_start=15596 + _FILETRANSFERINFOREQUEST._serialized_end=15840 + _FILETRANSFERINFOREQUEST_TRANSFERTYPE._serialized_start=15800 + _FILETRANSFERINFOREQUEST_TRANSFERTYPE._serialized_end=15840 + _LOCALINFO._serialized_start=15842 + _LOCALINFO._serialized_end=15891 + _SHUTDOWNREQUEST._serialized_start=15893 + _SHUTDOWNREQUEST._serialized_end=15956 + _SHUTDOWNRESPONSE._serialized_start=15958 + _SHUTDOWNRESPONSE._serialized_end=15976 + _ATTACHREQUEST._serialized_start=15978 + _ATTACHREQUEST._serialized_end=16058 + _ATTACHRESPONSE._serialized_start=16060 + _ATTACHRESPONSE._serialized_end=16158 + _TESTINJECTREQUEST._serialized_start=16161 + _TESTINJECTREQUEST._serialized_end=16502 + _TESTINJECTRESPONSE._serialized_start=16504 + _TESTINJECTRESPONSE._serialized_end=16524 + _HISTORYACTION._serialized_start=16526 + _HISTORYACTION._serialized_end=16556 + _PARTIALHISTORYREQUEST._serialized_start=16559 + _PARTIALHISTORYREQUEST._serialized_end=16761 + _PARTIALHISTORYRESPONSE._serialized_start=16763 + _PARTIALHISTORYRESPONSE._serialized_end=16787 + _SAMPLEDHISTORYREQUEST._serialized_start=16789 + _SAMPLEDHISTORYREQUEST._serialized_end=16858 + _SAMPLEDHISTORYITEM._serialized_start=16860 + _SAMPLEDHISTORYITEM._serialized_end=16955 + _SAMPLEDHISTORYRESPONSE._serialized_start=16957 + _SAMPLEDHISTORYRESPONSE._serialized_end=17031 + _RUNSTATUSREQUEST._serialized_start=17033 + _RUNSTATUSREQUEST._serialized_end=17097 + _RUNSTATUSRESPONSE._serialized_start=17099 + _RUNSTATUSRESPONSE._serialized_end=17219 + _RUNSTARTREQUEST._serialized_start=17221 + _RUNSTARTREQUEST._serialized_end=17324 + _RUNSTARTRESPONSE._serialized_start=17326 + _RUNSTARTRESPONSE._serialized_end=17344 + _RUNFINISHWITHOUTEXITREQUEST._serialized_start=17346 + _RUNFINISHWITHOUTEXITREQUEST._serialized_end=17421 + _RUNFINISHWITHOUTEXITRESPONSE._serialized_start=17423 + _RUNFINISHWITHOUTEXITRESPONSE._serialized_end=17453 + _CHECKVERSIONREQUEST._serialized_start=17455 + _CHECKVERSIONREQUEST._serialized_end=17547 + _CHECKVERSIONRESPONSE._serialized_start=17549 + _CHECKVERSIONRESPONSE._serialized_end=17642 + _JOBINFOREQUEST._serialized_start=17644 + _JOBINFOREQUEST._serialized_end=17706 + _JOBINFORESPONSE._serialized_start=17708 + _JOBINFORESPONSE._serialized_end=17762 + _LOGARTIFACTREQUEST._serialized_start=17765 + _LOGARTIFACTREQUEST._serialized_end=17924 + _LOGARTIFACTRESPONSE._serialized_start=17926 + _LOGARTIFACTRESPONSE._serialized_end=17991 + _DOWNLOADARTIFACTREQUEST._serialized_start=17994 + _DOWNLOADARTIFACTREQUEST._serialized_end=18184 + _DOWNLOADARTIFACTRESPONSE._serialized_start=18186 + _DOWNLOADARTIFACTRESPONSE._serialized_end=18235 + _KEEPALIVEREQUEST._serialized_start=18237 + _KEEPALIVEREQUEST._serialized_end=18301 + _KEEPALIVERESPONSE._serialized_start=18303 + _KEEPALIVERESPONSE._serialized_end=18322 + _ARTIFACTINFO._serialized_start=18324 + _ARTIFACTINFO._serialized_end=18437 + _GITINFO._serialized_start=18439 + _GITINFO._serialized_end=18480 + _GITSOURCE._serialized_start=18483 + _GITSOURCE._serialized_end=18618 + _IMAGESOURCE._serialized_start=18620 + _IMAGESOURCE._serialized_end=18648 + _SOURCE._serialized_start=18651 + _SOURCE._serialized_end=18791 + _JOBSOURCE._serialized_start=18793 + _JOBSOURCE._serialized_end=18900 + _PARTIALJOBARTIFACT._serialized_start=18902 + _PARTIALJOBARTIFACT._serialized_end=18988 + _USEARTIFACTRECORD._serialized_start=18991 + _USEARTIFACTRECORD._serialized_end=19148 + _USEARTIFACTRESULT._serialized_start=19150 + _USEARTIFACTRESULT._serialized_end=19169 + _CANCELREQUEST._serialized_start=19171 + _CANCELREQUEST._serialized_end=19253 + _CANCELRESPONSE._serialized_start=19255 + _CANCELRESPONSE._serialized_end=19271 + _DISKINFO._serialized_start=19273 + _DISKINFO._serialized_end=19312 + _MEMORYINFO._serialized_start=19314 + _MEMORYINFO._serialized_end=19341 + _CPUINFO._serialized_start=19343 + _CPUINFO._serialized_end=19390 + _APPLEINFO._serialized_start=19393 + _APPLEINFO._serialized_end=19547 + _GPUNVIDIAINFO._serialized_start=19549 + _GPUNVIDIAINFO._serialized_end=19656 + _GPUAMDINFO._serialized_start=19659 + _GPUAMDINFO._serialized_end=19924 + _TRAINIUMINFO._serialized_start=19926 + _TRAINIUMINFO._serialized_end=20036 + _TPUINFO._serialized_start=20038 + _TPUINFO._serialized_end=20119 + _COREWEAVEINFO._serialized_start=20121 + _COREWEAVEINFO._serialized_end=20190 + _ENVIRONMENTRECORD._serialized_start=20193 + _ENVIRONMENTRECORD._serialized_end=21385 + _ENVIRONMENTRECORD_DISKENTRY._serialized_start=21270 + _ENVIRONMENTRECORD_DISKENTRY._serialized_end=21339 + _ENVIRONMENTRECORD_SLURMENTRY._serialized_start=21341 + _ENVIRONMENTRECORD_SLURMENTRY._serialized_end=21385 + _PYTHONPACKAGESREQUEST._serialized_start=21388 + _PYTHONPACKAGESREQUEST._serialized_end=21529 + _PYTHONPACKAGESREQUEST_PYTHONPACKAGE._serialized_start=21483 + _PYTHONPACKAGESREQUEST_PYTHONPACKAGE._serialized_end=21529 + _JOBINPUTPATH._serialized_start=21531 + _JOBINPUTPATH._serialized_end=21559 + _JOBINPUTSOURCE._serialized_start=21562 + _JOBINPUTSOURCE._serialized_end=21776 + _JOBINPUTSOURCE_RUNCONFIGSOURCE._serialized_start=21715 + _JOBINPUTSOURCE_RUNCONFIGSOURCE._serialized_end=21732 + _JOBINPUTSOURCE_CONFIGFILESOURCE._serialized_start=21734 + _JOBINPUTSOURCE_CONFIGFILESOURCE._serialized_end=21766 + _JOBINPUTREQUEST._serialized_start=21779 + _JOBINPUTREQUEST._serialized_end=21978 + _SERVERFEATUREREQUEST._serialized_start=21980 + _SERVERFEATUREREQUEST._serialized_end=22096 + _SERVERFEATURERESPONSE._serialized_start=22098 + _SERVERFEATURERESPONSE._serialized_end=22173 + _SERVERFEATUREITEM._serialized_start=22175 + _SERVERFEATUREITEM._serialized_end=22225 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v3/wandb_server_pb2.py b/lib/python3.12/site-packages/wandb/proto/v3/wandb_server_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..0e7b3b33f0388762b349b263314257d4673ef24d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v3/wandb_server_pb2.py @@ -0,0 +1,228 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_server.proto +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 +from wandb.proto import wandb_internal_pb2 as wandb_dot_proto_dot_wandb__internal__pb2 +from wandb.proto import wandb_settings_pb2 as wandb_dot_proto_dot_wandb__settings__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1ewandb/proto/wandb_server.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\x1a wandb/proto/wandb_internal.proto\x1a wandb/proto/wandb_settings.proto\"k\n\x19ServerAuthenticateRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12\x10\n\x08\x62\x61se_url\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"w\n\x1aServerAuthenticateResponse\x12\x16\n\x0e\x64\x65\x66\x61ult_entity\x18\x01 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"D\n\x15ServerShutdownRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x18\n\x16ServerShutdownResponse\"B\n\x13ServerStatusRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x16\n\x14ServerStatusResponse\"r\n\x17ServerInformInitRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1a\n\x18ServerInformInitResponse\"s\n\x18ServerInformStartRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1b\n\x19ServerInformStartResponse\"H\n\x19ServerInformFinishRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformFinishResponse\"H\n\x19ServerInformAttachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"u\n\x1aServerInformAttachResponse\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"H\n\x19ServerInformDetachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformDetachResponse\"]\n\x1bServerInformTeardownRequest\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1e\n\x1cServerInformTeardownResponse\"\xfb\x04\n\rServerRequest\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x30\n\x0erecord_publish\x18\x01 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12\x34\n\x12record_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12>\n\x0binform_init\x18\x03 \x01(\x0b\x32\'.wandb_internal.ServerInformInitRequestH\x00\x12\x42\n\rinform_finish\x18\x04 \x01(\x0b\x32).wandb_internal.ServerInformFinishRequestH\x00\x12\x42\n\rinform_attach\x18\x05 \x01(\x0b\x32).wandb_internal.ServerInformAttachRequestH\x00\x12\x42\n\rinform_detach\x18\x06 \x01(\x0b\x32).wandb_internal.ServerInformDetachRequestH\x00\x12\x46\n\x0finform_teardown\x18\x07 \x01(\x0b\x32+.wandb_internal.ServerInformTeardownRequestH\x00\x12@\n\x0cinform_start\x18\x08 \x01(\x0b\x32(.wandb_internal.ServerInformStartRequestH\x00\x12\x41\n\x0c\x61uthenticate\x18\t \x01(\x0b\x32).wandb_internal.ServerAuthenticateRequestH\x00\x42\x15\n\x13server_request_type\"\x91\x05\n\x0eServerResponse\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x34\n\x12result_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.ResultH\x00\x12H\n\x14inform_init_response\x18\x03 \x01(\x0b\x32(.wandb_internal.ServerInformInitResponseH\x00\x12L\n\x16inform_finish_response\x18\x04 \x01(\x0b\x32*.wandb_internal.ServerInformFinishResponseH\x00\x12L\n\x16inform_attach_response\x18\x05 \x01(\x0b\x32*.wandb_internal.ServerInformAttachResponseH\x00\x12L\n\x16inform_detach_response\x18\x06 \x01(\x0b\x32*.wandb_internal.ServerInformDetachResponseH\x00\x12P\n\x18inform_teardown_response\x18\x07 \x01(\x0b\x32,.wandb_internal.ServerInformTeardownResponseH\x00\x12J\n\x15inform_start_response\x18\x08 \x01(\x0b\x32).wandb_internal.ServerInformStartResponseH\x00\x12K\n\x15\x61uthenticate_response\x18\t \x01(\x0b\x32*.wandb_internal.ServerAuthenticateResponseH\x00\x42\x16\n\x14server_response_typeB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + + + +_SERVERAUTHENTICATEREQUEST = DESCRIPTOR.message_types_by_name['ServerAuthenticateRequest'] +_SERVERAUTHENTICATERESPONSE = DESCRIPTOR.message_types_by_name['ServerAuthenticateResponse'] +_SERVERSHUTDOWNREQUEST = DESCRIPTOR.message_types_by_name['ServerShutdownRequest'] +_SERVERSHUTDOWNRESPONSE = DESCRIPTOR.message_types_by_name['ServerShutdownResponse'] +_SERVERSTATUSREQUEST = DESCRIPTOR.message_types_by_name['ServerStatusRequest'] +_SERVERSTATUSRESPONSE = DESCRIPTOR.message_types_by_name['ServerStatusResponse'] +_SERVERINFORMINITREQUEST = DESCRIPTOR.message_types_by_name['ServerInformInitRequest'] +_SERVERINFORMINITRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformInitResponse'] +_SERVERINFORMSTARTREQUEST = DESCRIPTOR.message_types_by_name['ServerInformStartRequest'] +_SERVERINFORMSTARTRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformStartResponse'] +_SERVERINFORMFINISHREQUEST = DESCRIPTOR.message_types_by_name['ServerInformFinishRequest'] +_SERVERINFORMFINISHRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformFinishResponse'] +_SERVERINFORMATTACHREQUEST = DESCRIPTOR.message_types_by_name['ServerInformAttachRequest'] +_SERVERINFORMATTACHRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformAttachResponse'] +_SERVERINFORMDETACHREQUEST = DESCRIPTOR.message_types_by_name['ServerInformDetachRequest'] +_SERVERINFORMDETACHRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformDetachResponse'] +_SERVERINFORMTEARDOWNREQUEST = DESCRIPTOR.message_types_by_name['ServerInformTeardownRequest'] +_SERVERINFORMTEARDOWNRESPONSE = DESCRIPTOR.message_types_by_name['ServerInformTeardownResponse'] +_SERVERREQUEST = DESCRIPTOR.message_types_by_name['ServerRequest'] +_SERVERRESPONSE = DESCRIPTOR.message_types_by_name['ServerResponse'] +ServerAuthenticateRequest = _reflection.GeneratedProtocolMessageType('ServerAuthenticateRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERAUTHENTICATEREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerAuthenticateRequest) + }) +_sym_db.RegisterMessage(ServerAuthenticateRequest) + +ServerAuthenticateResponse = _reflection.GeneratedProtocolMessageType('ServerAuthenticateResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERAUTHENTICATERESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerAuthenticateResponse) + }) +_sym_db.RegisterMessage(ServerAuthenticateResponse) + +ServerShutdownRequest = _reflection.GeneratedProtocolMessageType('ServerShutdownRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERSHUTDOWNREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerShutdownRequest) + }) +_sym_db.RegisterMessage(ServerShutdownRequest) + +ServerShutdownResponse = _reflection.GeneratedProtocolMessageType('ServerShutdownResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERSHUTDOWNRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerShutdownResponse) + }) +_sym_db.RegisterMessage(ServerShutdownResponse) + +ServerStatusRequest = _reflection.GeneratedProtocolMessageType('ServerStatusRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERSTATUSREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerStatusRequest) + }) +_sym_db.RegisterMessage(ServerStatusRequest) + +ServerStatusResponse = _reflection.GeneratedProtocolMessageType('ServerStatusResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERSTATUSRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerStatusResponse) + }) +_sym_db.RegisterMessage(ServerStatusResponse) + +ServerInformInitRequest = _reflection.GeneratedProtocolMessageType('ServerInformInitRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMINITREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformInitRequest) + }) +_sym_db.RegisterMessage(ServerInformInitRequest) + +ServerInformInitResponse = _reflection.GeneratedProtocolMessageType('ServerInformInitResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMINITRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformInitResponse) + }) +_sym_db.RegisterMessage(ServerInformInitResponse) + +ServerInformStartRequest = _reflection.GeneratedProtocolMessageType('ServerInformStartRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMSTARTREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformStartRequest) + }) +_sym_db.RegisterMessage(ServerInformStartRequest) + +ServerInformStartResponse = _reflection.GeneratedProtocolMessageType('ServerInformStartResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMSTARTRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformStartResponse) + }) +_sym_db.RegisterMessage(ServerInformStartResponse) + +ServerInformFinishRequest = _reflection.GeneratedProtocolMessageType('ServerInformFinishRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMFINISHREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformFinishRequest) + }) +_sym_db.RegisterMessage(ServerInformFinishRequest) + +ServerInformFinishResponse = _reflection.GeneratedProtocolMessageType('ServerInformFinishResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMFINISHRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformFinishResponse) + }) +_sym_db.RegisterMessage(ServerInformFinishResponse) + +ServerInformAttachRequest = _reflection.GeneratedProtocolMessageType('ServerInformAttachRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMATTACHREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformAttachRequest) + }) +_sym_db.RegisterMessage(ServerInformAttachRequest) + +ServerInformAttachResponse = _reflection.GeneratedProtocolMessageType('ServerInformAttachResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMATTACHRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformAttachResponse) + }) +_sym_db.RegisterMessage(ServerInformAttachResponse) + +ServerInformDetachRequest = _reflection.GeneratedProtocolMessageType('ServerInformDetachRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMDETACHREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformDetachRequest) + }) +_sym_db.RegisterMessage(ServerInformDetachRequest) + +ServerInformDetachResponse = _reflection.GeneratedProtocolMessageType('ServerInformDetachResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMDETACHRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformDetachResponse) + }) +_sym_db.RegisterMessage(ServerInformDetachResponse) + +ServerInformTeardownRequest = _reflection.GeneratedProtocolMessageType('ServerInformTeardownRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMTEARDOWNREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformTeardownRequest) + }) +_sym_db.RegisterMessage(ServerInformTeardownRequest) + +ServerInformTeardownResponse = _reflection.GeneratedProtocolMessageType('ServerInformTeardownResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERINFORMTEARDOWNRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerInformTeardownResponse) + }) +_sym_db.RegisterMessage(ServerInformTeardownResponse) + +ServerRequest = _reflection.GeneratedProtocolMessageType('ServerRequest', (_message.Message,), { + 'DESCRIPTOR' : _SERVERREQUEST, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerRequest) + }) +_sym_db.RegisterMessage(ServerRequest) + +ServerResponse = _reflection.GeneratedProtocolMessageType('ServerResponse', (_message.Message,), { + 'DESCRIPTOR' : _SERVERRESPONSE, + '__module__' : 'wandb.proto.wandb_server_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ServerResponse) + }) +_sym_db.RegisterMessage(ServerResponse) + +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + _SERVERAUTHENTICATEREQUEST._serialized_start=148 + _SERVERAUTHENTICATEREQUEST._serialized_end=255 + _SERVERAUTHENTICATERESPONSE._serialized_start=257 + _SERVERAUTHENTICATERESPONSE._serialized_end=376 + _SERVERSHUTDOWNREQUEST._serialized_start=378 + _SERVERSHUTDOWNREQUEST._serialized_end=446 + _SERVERSHUTDOWNRESPONSE._serialized_start=448 + _SERVERSHUTDOWNRESPONSE._serialized_end=472 + _SERVERSTATUSREQUEST._serialized_start=474 + _SERVERSTATUSREQUEST._serialized_end=540 + _SERVERSTATUSRESPONSE._serialized_start=542 + _SERVERSTATUSRESPONSE._serialized_end=564 + _SERVERINFORMINITREQUEST._serialized_start=566 + _SERVERINFORMINITREQUEST._serialized_end=680 + _SERVERINFORMINITRESPONSE._serialized_start=682 + _SERVERINFORMINITRESPONSE._serialized_end=708 + _SERVERINFORMSTARTREQUEST._serialized_start=710 + _SERVERINFORMSTARTREQUEST._serialized_end=825 + _SERVERINFORMSTARTRESPONSE._serialized_start=827 + _SERVERINFORMSTARTRESPONSE._serialized_end=854 + _SERVERINFORMFINISHREQUEST._serialized_start=856 + _SERVERINFORMFINISHREQUEST._serialized_end=928 + _SERVERINFORMFINISHRESPONSE._serialized_start=930 + _SERVERINFORMFINISHRESPONSE._serialized_end=958 + _SERVERINFORMATTACHREQUEST._serialized_start=960 + _SERVERINFORMATTACHREQUEST._serialized_end=1032 + _SERVERINFORMATTACHRESPONSE._serialized_start=1034 + _SERVERINFORMATTACHRESPONSE._serialized_end=1151 + _SERVERINFORMDETACHREQUEST._serialized_start=1153 + _SERVERINFORMDETACHREQUEST._serialized_end=1225 + _SERVERINFORMDETACHRESPONSE._serialized_start=1227 + _SERVERINFORMDETACHRESPONSE._serialized_end=1255 + _SERVERINFORMTEARDOWNREQUEST._serialized_start=1257 + _SERVERINFORMTEARDOWNREQUEST._serialized_end=1350 + _SERVERINFORMTEARDOWNRESPONSE._serialized_start=1352 + _SERVERINFORMTEARDOWNRESPONSE._serialized_end=1382 + _SERVERREQUEST._serialized_start=1385 + _SERVERREQUEST._serialized_end=2020 + _SERVERRESPONSE._serialized_start=2023 + _SERVERRESPONSE._serialized_end=2680 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v3/wandb_settings_pb2.py b/lib/python3.12/site-packages/wandb/proto/v3/wandb_settings_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..748547a9fb9c0bdde9909d8aef4ad853f0d0c14e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v3/wandb_settings_pb2.py @@ -0,0 +1,122 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_settings.proto +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import wrappers_pb2 as google_dot_protobuf_dot_wrappers__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n wandb/proto/wandb_settings.proto\x12\x0ewandb_internal\x1a\x1egoogle/protobuf/wrappers.proto\" \n\x0fListStringValue\x12\r\n\x05value\x18\x01 \x03(\t\"\x1d\n\x0cListIntValue\x12\r\n\x05value\x18\x01 \x03(\x05\"\x8a\x01\n\x17MapStringKeyStringValue\x12\x41\n\x05value\x18\x01 \x03(\x0b\x32\x32.wandb_internal.MapStringKeyStringValue.ValueEntry\x1a,\n\nValueEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\xcb\x01\n#MapStringKeyMapStringKeyStringValue\x12M\n\x05value\x18\x01 \x03(\x0b\x32>.wandb_internal.MapStringKeyMapStringKeyStringValue.ValueEntry\x1aU\n\nValueEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x36\n\x05value\x18\x02 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue:\x02\x38\x01\"\x9a\x01\n\x12OpenMetricsFilters\x12\x33\n\x08sequence\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValueH\x00\x12\x46\n\x07mapping\x18\x02 \x01(\x0b\x32\x33.wandb_internal.MapStringKeyMapStringKeyStringValueH\x00\x42\x07\n\x05value\"7\n\tRunMoment\x12\x0b\n\x03run\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01\x12\x0e\n\x06metric\x18\x03 \x01(\t\"\x99N\n\x08Settings\x12-\n\x07\x61pi_key\x18\x37 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x13identity_token_file\x18\xaa\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10\x63redentials_file\x18\xab\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x39\n\x14insecure_disable_ssl\x18\xb9\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08_offline\x18\x1e \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06x_sync\x18\x1f \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\tsync_file\x18\x86\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x07_shared\x18\xa2\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x06run_id\x18k \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07run_url\x18q \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07project\x18\x61 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06\x65ntity\x18\x45 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\x0corganization\x18\xbc\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0cx_start_time\x18) \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12.\n\x08root_dir\x18i \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07log_dir\x18U \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0clog_internal\x18V \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\tfiles_dir\x18\x46 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0bx_files_dir\x18\xb4\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0cignore_globs\x18N \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12.\n\x08\x62\x61se_url\x18\x39 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12=\n\x17x_file_stream_max_bytes\x18\xac\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x46\n\x1fx_file_stream_transmit_interval\x18\xaf\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x45\n\x14x_extra_http_headers\x18\x0e \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12=\n\x17x_file_stream_retry_max\x18\x93\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12K\n$x_file_stream_retry_wait_min_seconds\x18\x94\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12K\n$x_file_stream_retry_wait_max_seconds\x18\x95\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x43\n\x1dx_file_stream_timeout_seconds\x18\x0f \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x42\n\x1cx_file_stream_max_line_bytes\x18\xb2\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12?\n\x19x_file_transfer_retry_max\x18\x96\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12M\n&x_file_transfer_retry_wait_min_seconds\x18\x97\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12M\n&x_file_transfer_retry_wait_max_seconds\x18\x98\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x46\n\x1fx_file_transfer_timeout_seconds\x18\x99\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x39\n\x13x_graphql_retry_max\x18\x9a\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12G\n x_graphql_retry_wait_min_seconds\x18\x9b\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12G\n x_graphql_retry_wait_max_seconds\x18\x9c\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12@\n\x19x_graphql_timeout_seconds\x18\x9d\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x31\n\nhttp_proxy\x18\xa8\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0bhttps_proxy\x18\xa9\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12;\n\tx_proxies\x18\xc8\x01 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12-\n\x07program\x18_ \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0fprogram_relpath\x18` \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10_code_path_local\x18\xa3\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x0fprogram_abspath\x18\x9f\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x05_args\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12)\n\x03_os\x18 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06\x64ocker\x18\x43 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0cx_executable\x18\r \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07_python\x18\" \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\tcolab_url\x18\xa0\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x04host\x18M \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x08username\x18\x8d\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x05\x65mail\x18\x44 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06resume\x18\x66 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0bresume_from\x18\xa7\x01 \x01(\x0b\x32\x19.wandb_internal.RunMoment\x12-\n\tfork_from\x18\xa4\x01 \x01(\x0b\x32\x19.wandb_internal.RunMoment\x12\x38\n\x14\x64isable_job_creation\x18\x41 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\tsweep_url\x18\x83\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12;\n\x16x_disable_update_check\x18\xa5\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0ex_disable_meta\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\tsave_code\x18s \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x0b\x64isable_git\x18? \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x16x_disable_machine_info\x18\x9e\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0fx_disable_stats\x18\n \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x13x_stats_buffer_size\x18\xa1\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12@\n\x19x_stats_sampling_interval\x18\xae\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x30\n\x0bx_stats_pid\x18* \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12<\n\x12x_stats_disk_paths\x18\x92\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12H\n\"x_stats_neuron_monitor_config_path\x18. \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12<\n\x15x_stats_dcgm_exporter\x18\xbb\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12O\n\x1ex_stats_open_metrics_endpoints\x18/ \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12H\n\x1cx_stats_open_metrics_filters\x18\x30 \x01(\x0b\x32\".wandb_internal.OpenMetricsFilters\x12S\n!x_stats_open_metrics_http_headers\x18\xb8\x01 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12=\n\x16x_stats_gpu_device_ids\x18\xba\x01 \x01(\x0b\x32\x1c.wandb_internal.ListIntValue\x12\x37\n\x11x_stats_cpu_count\x18\xc2\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12?\n\x19x_stats_cpu_logical_count\x18\xc3\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x37\n\x11x_stats_gpu_count\x18\xc4\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x37\n\x10x_stats_gpu_type\x18\xc5\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12?\n\x1ax_stats_track_process_tree\x18\xc6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12.\n\x07x_label\x18\xb5\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\tx_primary\x18\xb6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12:\n\x15x_update_finish_state\x18\xb7\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12<\n\x17\x61llow_offline_artifacts\x18\xb1\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\x07\x63onsole\x18< \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x11\x63onsole_multipart\x18\xa6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x35\n\x10sync_tensorboard\x18\xb3\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x42\n\x1dx_server_side_derived_summary\x18\xbd\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x46\n!x_server_side_expand_glob_metrics\x18\xbe\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x16x_skip_transaction_log\x18\xbf\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12J\n#x_stats_coreweave_metadata_base_url\x18\xc0\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12J\n#x_stats_coreweave_metadata_endpoint\x18\xc1\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0b_aws_lambda\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0fx_cli_only_mode\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06_colab\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10x_disable_viewer\x18\x0b \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x15x_flow_control_custom\x18\x10 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x17x_flow_control_disabled\x18\x11 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12>\n\x18x_internal_check_process\x18\x12 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12,\n\x08_ipython\x18\x14 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08_jupyter\x18\x15 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x0ex_jupyter_root\x18\x16 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x07_kaggle\x18\x17 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12=\n\x18x_live_policy_rate_limit\x18\x18 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12<\n\x17x_live_policy_wait_time\x18\x19 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x30\n\x0bx_log_level\x18\x1a \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x35\n\x10x_network_buffer\x18\x1b \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12)\n\x05_noop\x18\x1c \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\t_notebook\x18\x1d \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\t_platform\x18! \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x38\n\x12x_runqueue_item_id\x18# \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x13x_save_requirements\x18% \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x13x_service_transport\x18& \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x34\n\x0ex_service_wait\x18\' \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x35\n\x0f_start_datetime\x18( \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\r_tmp_code_dir\x18\x31 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x08_windows\x18\x34 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10\x61llow_val_change\x18\x35 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\tanonymous\x18\x36 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12P\n\x1f\x61zure_account_url_to_access_key\x18\x38 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12.\n\x08\x63ode_dir\x18: \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0c\x63onfig_paths\x18; \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12\x30\n\ndeployment\x18= \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\x0c\x64isable_code\x18> \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x31\n\rdisable_hints\x18@ \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08\x64isabled\x18\x42 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12)\n\x05\x66orce\x18G \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\ngit_commit\x18H \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\ngit_remote\x18I \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x34\n\x0egit_remote_url\x18J \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08git_root\x18K \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x11heartbeat_seconds\x18L \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x32\n\x0cinit_timeout\x18O \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12,\n\x08is_local\x18P \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\njob_source\x18Q \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\rlabel_disable\x18R \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06launch\x18S \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x38\n\x12launch_config_path\x18T \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x14log_symlink_internal\x18W \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x10log_symlink_user\x18X \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08log_user\x18Y \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\rlogin_timeout\x18Z \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12*\n\x04mode\x18\\ \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\rnotebook_name\x18] \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\x0bproject_url\x18\x62 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12)\n\x05quiet\x18\x63 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12+\n\x07relogin\x18\x65 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0cresume_fname\x18g \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x07resumed\x18h \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\trun_group\x18j \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0crun_job_type\x18l \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08run_mode\x18m \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08run_name\x18n \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\trun_notes\x18o \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\x08run_tags\x18p \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12\x35\n\x11sagemaker_disable\x18r \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x35\n\x0fsettings_system\x18t \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x38\n\x12settings_workspace\x18u \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0bshow_colors\x18v \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12.\n\nshow_emoji\x18w \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x0bshow_errors\x18x \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\tshow_info\x18y \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x31\n\rshow_warnings\x18z \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06silent\x18{ \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0cstart_method\x18| \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x06strict\x18} \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0esummary_errors\x18~ \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x34\n\x0fsummary_timeout\x18\x7f \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x36\n\x10summary_warnings\x18\x80\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12/\n\x08sweep_id\x18\x81\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10sweep_param_path\x18\x82\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x07symlink\x18\x84\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x08sync_dir\x18\x85\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x13sync_symlink_latest\x18\x87\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12J\n%table_raise_on_max_row_limit_exceeded\x18\x8a\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x08timespec\x18\x8b\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x07tmp_dir\x18\x8c\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\twandb_dir\x18\x8e\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0ex_jupyter_name\x18\x8f\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0ex_jupyter_path\x18\x90\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x08job_name\x18\x91\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValueJ\x04\x08\x03\x10\x04J\x04\x08\x06\x10\x07J\x04\x08\x08\x10\tJ\x04\x08\t\x10\nJ\x04\x08\x0c\x10\rJ\x04\x08\x13\x10\x14J\x04\x08$\x10%J\x04\x08+\x10,J\x04\x08,\x10-J\x04\x08-\x10.J\x04\x08\x32\x10\x33J\x04\x08\x33\x10\x34J\x04\x08[\x10\\J\x04\x08^\x10_J\x04\x08\x64\x10\x65J\x06\x08\x88\x01\x10\x89\x01J\x06\x08\x89\x01\x10\x8a\x01J\x06\x08\xad\x01\x10\xae\x01J\x06\x08\xb0\x01\x10\xb1\x01\x42\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + + + +_LISTSTRINGVALUE = DESCRIPTOR.message_types_by_name['ListStringValue'] +_LISTINTVALUE = DESCRIPTOR.message_types_by_name['ListIntValue'] +_MAPSTRINGKEYSTRINGVALUE = DESCRIPTOR.message_types_by_name['MapStringKeyStringValue'] +_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY = _MAPSTRINGKEYSTRINGVALUE.nested_types_by_name['ValueEntry'] +_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE = DESCRIPTOR.message_types_by_name['MapStringKeyMapStringKeyStringValue'] +_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY = _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE.nested_types_by_name['ValueEntry'] +_OPENMETRICSFILTERS = DESCRIPTOR.message_types_by_name['OpenMetricsFilters'] +_RUNMOMENT = DESCRIPTOR.message_types_by_name['RunMoment'] +_SETTINGS = DESCRIPTOR.message_types_by_name['Settings'] +ListStringValue = _reflection.GeneratedProtocolMessageType('ListStringValue', (_message.Message,), { + 'DESCRIPTOR' : _LISTSTRINGVALUE, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ListStringValue) + }) +_sym_db.RegisterMessage(ListStringValue) + +ListIntValue = _reflection.GeneratedProtocolMessageType('ListIntValue', (_message.Message,), { + 'DESCRIPTOR' : _LISTINTVALUE, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.ListIntValue) + }) +_sym_db.RegisterMessage(ListIntValue) + +MapStringKeyStringValue = _reflection.GeneratedProtocolMessageType('MapStringKeyStringValue', (_message.Message,), { + + 'ValueEntry' : _reflection.GeneratedProtocolMessageType('ValueEntry', (_message.Message,), { + 'DESCRIPTOR' : _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MapStringKeyStringValue.ValueEntry) + }) + , + 'DESCRIPTOR' : _MAPSTRINGKEYSTRINGVALUE, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MapStringKeyStringValue) + }) +_sym_db.RegisterMessage(MapStringKeyStringValue) +_sym_db.RegisterMessage(MapStringKeyStringValue.ValueEntry) + +MapStringKeyMapStringKeyStringValue = _reflection.GeneratedProtocolMessageType('MapStringKeyMapStringKeyStringValue', (_message.Message,), { + + 'ValueEntry' : _reflection.GeneratedProtocolMessageType('ValueEntry', (_message.Message,), { + 'DESCRIPTOR' : _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MapStringKeyMapStringKeyStringValue.ValueEntry) + }) + , + 'DESCRIPTOR' : _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.MapStringKeyMapStringKeyStringValue) + }) +_sym_db.RegisterMessage(MapStringKeyMapStringKeyStringValue) +_sym_db.RegisterMessage(MapStringKeyMapStringKeyStringValue.ValueEntry) + +OpenMetricsFilters = _reflection.GeneratedProtocolMessageType('OpenMetricsFilters', (_message.Message,), { + 'DESCRIPTOR' : _OPENMETRICSFILTERS, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.OpenMetricsFilters) + }) +_sym_db.RegisterMessage(OpenMetricsFilters) + +RunMoment = _reflection.GeneratedProtocolMessageType('RunMoment', (_message.Message,), { + 'DESCRIPTOR' : _RUNMOMENT, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.RunMoment) + }) +_sym_db.RegisterMessage(RunMoment) + +Settings = _reflection.GeneratedProtocolMessageType('Settings', (_message.Message,), { + 'DESCRIPTOR' : _SETTINGS, + '__module__' : 'wandb.proto.wandb_settings_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Settings) + }) +_sym_db.RegisterMessage(Settings) + +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY._options = None + _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_options = b'8\001' + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY._options = None + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_options = b'8\001' + _LISTSTRINGVALUE._serialized_start=84 + _LISTSTRINGVALUE._serialized_end=116 + _LISTINTVALUE._serialized_start=118 + _LISTINTVALUE._serialized_end=147 + _MAPSTRINGKEYSTRINGVALUE._serialized_start=150 + _MAPSTRINGKEYSTRINGVALUE._serialized_end=288 + _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_start=244 + _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_end=288 + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE._serialized_start=291 + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE._serialized_end=494 + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_start=409 + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_end=494 + _OPENMETRICSFILTERS._serialized_start=497 + _OPENMETRICSFILTERS._serialized_end=651 + _RUNMOMENT._serialized_start=653 + _RUNMOMENT._serialized_end=708 + _SETTINGS._serialized_start=711 + _SETTINGS._serialized_end=10720 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v3/wandb_telemetry_pb2.py b/lib/python3.12/site-packages/wandb/proto/v3/wandb_telemetry_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..93fc1a266fc43582efc855fb71e6506e4a995246 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v3/wandb_telemetry_pb2.py @@ -0,0 +1,106 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_telemetry.proto +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import message as _message +from google.protobuf import reflection as _reflection +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n!wandb/proto/wandb_telemetry.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\"\xdb\x03\n\x0fTelemetryRecord\x12-\n\x0cimports_init\x18\x01 \x01(\x0b\x32\x17.wandb_internal.Imports\x12/\n\x0eimports_finish\x18\x02 \x01(\x0b\x32\x17.wandb_internal.Imports\x12(\n\x07\x66\x65\x61ture\x18\x03 \x01(\x0b\x32\x17.wandb_internal.Feature\x12\x16\n\x0epython_version\x18\x04 \x01(\t\x12\x13\n\x0b\x63li_version\x18\x05 \x01(\t\x12\x1b\n\x13huggingface_version\x18\x06 \x01(\t\x12 \n\x03\x65nv\x18\x08 \x01(\x0b\x32\x13.wandb_internal.Env\x12%\n\x05label\x18\t \x01(\x0b\x32\x16.wandb_internal.Labels\x12.\n\ndeprecated\x18\n \x01(\x0b\x32\x1a.wandb_internal.Deprecated\x12&\n\x06issues\x18\x0b \x01(\x0b\x32\x16.wandb_internal.Issues\x12\x14\n\x0c\x63ore_version\x18\x0c \x01(\t\x12\x10\n\x08platform\x18\r \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x11\n\x0fTelemetryResult\"\x9a\x0e\n\x07Imports\x12\r\n\x05torch\x18\x01 \x01(\x08\x12\r\n\x05keras\x18\x02 \x01(\x08\x12\x12\n\ntensorflow\x18\x03 \x01(\x08\x12\x0e\n\x06\x66\x61stai\x18\x04 \x01(\x08\x12\x0f\n\x07sklearn\x18\x05 \x01(\x08\x12\x0f\n\x07xgboost\x18\x06 \x01(\x08\x12\x10\n\x08\x63\x61tboost\x18\x07 \x01(\x08\x12\x10\n\x08lightgbm\x18\x08 \x01(\x08\x12\x19\n\x11pytorch_lightning\x18\t \x01(\x08\x12\x0e\n\x06ignite\x18\n \x01(\x08\x12\x14\n\x0ctransformers\x18\x0b \x01(\x08\x12\x0b\n\x03jax\x18\x0c \x01(\x08\x12\x10\n\x08metaflow\x18\r \x01(\x08\x12\x10\n\x08\x61llennlp\x18\x0e \x01(\x08\x12\x11\n\tautogluon\x18\x0f \x01(\x08\x12\x11\n\tautokeras\x18\x10 \x01(\x08\x12\x10\n\x08\x63\x61talyst\x18\x12 \x01(\x08\x12\x10\n\x08\x64\x65\x65pchem\x18\x15 \x01(\x08\x12\x0f\n\x07\x64\x65\x65pctr\x18\x16 \x01(\x08\x12\x0f\n\x07pycaret\x18\x1c \x01(\x08\x12\x14\n\x0cpytorchvideo\x18\x1d \x01(\x08\x12\x0b\n\x03ray\x18\x1e \x01(\x08\x12\x1a\n\x12simpletransformers\x18\x1f \x01(\x08\x12\x0e\n\x06skorch\x18 \x01(\x08\x12\r\n\x05spacy\x18! \x01(\x08\x12\r\n\x05\x66lash\x18\" \x01(\x08\x12\x0e\n\x06optuna\x18# \x01(\x08\x12\x0f\n\x07recbole\x18$ \x01(\x08\x12\x0c\n\x04mmcv\x18% \x01(\x08\x12\r\n\x05mmdet\x18& \x01(\x08\x12\x11\n\ttorchdrug\x18\' \x01(\x08\x12\x11\n\ttorchtext\x18( \x01(\x08\x12\x13\n\x0btorchvision\x18) \x01(\x08\x12\r\n\x05\x65legy\x18* \x01(\x08\x12\x12\n\ndetectron2\x18+ \x01(\x08\x12\r\n\x05\x66lair\x18, \x01(\x08\x12\x0c\n\x04\x66lax\x18- \x01(\x08\x12\x0c\n\x04syft\x18. \x01(\x08\x12\x0b\n\x03TTS\x18/ \x01(\x08\x12\r\n\x05monai\x18\x30 \x01(\x08\x12\x17\n\x0fhuggingface_hub\x18\x31 \x01(\x08\x12\r\n\x05hydra\x18\x32 \x01(\x08\x12\x10\n\x08\x64\x61tasets\x18\x33 \x01(\x08\x12\x0e\n\x06sacred\x18\x34 \x01(\x08\x12\x0e\n\x06joblib\x18\x35 \x01(\x08\x12\x0c\n\x04\x64\x61sk\x18\x36 \x01(\x08\x12\x11\n\tpaddleocr\x18\x38 \x01(\x08\x12\r\n\x05ppdet\x18\x39 \x01(\x08\x12\x11\n\tpaddleseg\x18: \x01(\x08\x12\x11\n\tpaddlenlp\x18; \x01(\x08\x12\r\n\x05mmseg\x18< \x01(\x08\x12\r\n\x05mmocr\x18= \x01(\x08\x12\r\n\x05mmcls\x18> \x01(\x08\x12\x0c\n\x04timm\x18? \x01(\x08\x12\x0f\n\x07\x66\x61irseq\x18@ \x01(\x08\x12\x12\n\ndeepchecks\x18\x41 \x01(\x08\x12\x10\n\x08\x63omposer\x18\x42 \x01(\x08\x12\x10\n\x08sparseml\x18\x43 \x01(\x08\x12\x10\n\x08\x61nomalib\x18\x44 \x01(\x08\x12\r\n\x05zenml\x18\x45 \x01(\x08\x12\x12\n\ncolossalai\x18\x46 \x01(\x08\x12\x12\n\naccelerate\x18G \x01(\x08\x12\x0e\n\x06merlin\x18H \x01(\x08\x12\x0f\n\x07nanodet\x18I \x01(\x08\x12#\n\x1bsegmentation_models_pytorch\x18J \x01(\x08\x12\x1d\n\x15sentence_transformers\x18K \x01(\x08\x12\x0b\n\x03\x64gl\x18L \x01(\x08\x12\x17\n\x0ftorch_geometric\x18M \x01(\x08\x12\x0c\n\x04jina\x18N \x01(\x08\x12\x0e\n\x06kornia\x18O \x01(\x08\x12\x16\n\x0e\x61lbumentations\x18P \x01(\x08\x12\x10\n\x08keras_cv\x18Q \x01(\x08\x12\x10\n\x08mmengine\x18R \x01(\x08\x12\x11\n\tdiffusers\x18S \x01(\x08\x12\x0b\n\x03trl\x18T \x01(\x08\x12\x0c\n\x04trlx\x18U \x01(\x08\x12\x11\n\tlangchain\x18V \x01(\x08\x12\x13\n\x0bllama_index\x18W \x01(\x08\x12\x15\n\rstability_sdk\x18X \x01(\x08\x12\x0f\n\x07prefect\x18Y \x01(\x08\x12\x13\n\x0bprefect_ray\x18Z \x01(\x08\x12\x10\n\x08pinecone\x18[ \x01(\x08\x12\x10\n\x08\x63hromadb\x18\\ \x01(\x08\x12\x10\n\x08weaviate\x18] \x01(\x08\x12\x13\n\x0bpromptlayer\x18^ \x01(\x08\x12\x0e\n\x06openai\x18_ \x01(\x08\x12\x0e\n\x06\x63ohere\x18` \x01(\x08\x12\x11\n\tanthropic\x18\x61 \x01(\x08\x12\x0c\n\x04peft\x18\x62 \x01(\x08\x12\x0f\n\x07optimum\x18\x63 \x01(\x08\x12\x10\n\x08\x65valuate\x18\x64 \x01(\x08\x12\x10\n\x08langflow\x18\x65 \x01(\x08\x12\x12\n\nkeras_core\x18\x66 \x01(\x08\x12\x18\n\x10lightning_fabric\x18g \x01(\x08\x12\x1c\n\x14\x63urated_transformers\x18h \x01(\x08\x12\x0e\n\x06orjson\x18i \x01(\x08\x12\x11\n\tlightning\x18j \x01(\x08J\x04\x08\x11\x10\x12J\x04\x08\x13\x10\x14J\x04\x08\x14\x10\x15J\x04\x08\x17\x10\x18J\x04\x08\x18\x10\x19J\x04\x08\x19\x10\x1aJ\x04\x08\x1a\x10\x1bJ\x04\x08\x1b\x10\x1cJ\x04\x08\x37\x10\x38\"\x95\r\n\x07\x46\x65\x61ture\x12\r\n\x05watch\x18\x01 \x01(\x08\x12\x0e\n\x06\x66inish\x18\x02 \x01(\x08\x12\x0c\n\x04save\x18\x03 \x01(\x08\x12\x0f\n\x07offline\x18\x04 \x01(\x08\x12\x0f\n\x07resumed\x18\x05 \x01(\x08\x12\x0c\n\x04grpc\x18\x06 \x01(\x08\x12\x0e\n\x06metric\x18\x07 \x01(\x08\x12\r\n\x05keras\x18\x08 \x01(\x08\x12\x11\n\tsagemaker\x18\t \x01(\x08\x12\x1c\n\x14\x61rtifact_incremental\x18\n \x01(\x08\x12\x10\n\x08metaflow\x18\x0b \x01(\x08\x12\x0f\n\x07prodigy\x18\x0c \x01(\x08\x12\x15\n\rset_init_name\x18\r \x01(\x08\x12\x13\n\x0bset_init_id\x18\x0e \x01(\x08\x12\x15\n\rset_init_tags\x18\x0f \x01(\x08\x12\x17\n\x0fset_init_config\x18\x10 \x01(\x08\x12\x14\n\x0cset_run_name\x18\x11 \x01(\x08\x12\x14\n\x0cset_run_tags\x18\x12 \x01(\x08\x12\x17\n\x0fset_config_item\x18\x13 \x01(\x08\x12\x0e\n\x06launch\x18\x14 \x01(\x08\x12\x1c\n\x14torch_profiler_trace\x18\x15 \x01(\x08\x12\x0b\n\x03sb3\x18\x16 \x01(\x08\x12\x17\n\x0finit_return_run\x18\x18 \x01(\x08\x12\x1f\n\x17lightgbm_wandb_callback\x18\x19 \x01(\x08\x12\x1c\n\x14lightgbm_log_summary\x18\x1a \x01(\x08\x12\x1f\n\x17\x63\x61tboost_wandb_callback\x18\x1b \x01(\x08\x12\x1c\n\x14\x63\x61tboost_log_summary\x18\x1c \x01(\x08\x12\x17\n\x0ftensorboard_log\x18\x1d \x01(\x08\x12\x16\n\x0e\x65stimator_hook\x18\x1e \x01(\x08\x12\x1e\n\x16xgboost_wandb_callback\x18\x1f \x01(\x08\x12\"\n\x1axgboost_old_wandb_callback\x18 \x01(\x08\x12\x0e\n\x06\x61ttach\x18! \x01(\x08\x12\x19\n\x11tensorboard_patch\x18\" \x01(\x08\x12\x18\n\x10tensorboard_sync\x18# \x01(\x08\x12\x15\n\rkfp_wandb_log\x18$ \x01(\x08\x12\x1b\n\x13maybe_run_overwrite\x18% \x01(\x08\x12\x1c\n\x14keras_metrics_logger\x18& \x01(\x08\x12\x1e\n\x16keras_model_checkpoint\x18\' \x01(\x08\x12!\n\x19keras_wandb_eval_callback\x18( \x01(\x08\x12\x1d\n\x15\x66low_control_overflow\x18) \x01(\x08\x12\x0c\n\x04sync\x18* \x01(\x08\x12\x1d\n\x15\x66low_control_disabled\x18+ \x01(\x08\x12\x1b\n\x13\x66low_control_custom\x18, \x01(\x08\x12\x14\n\x0copen_metrics\x18. \x01(\x08\x12\x1a\n\x12ultralytics_yolov8\x18/ \x01(\x08\x12\x17\n\x0fimporter_mlflow\x18\x30 \x01(\x08\x12\x15\n\rsync_tfevents\x18\x31 \x01(\x08\x12\x15\n\rasync_uploads\x18\x32 \x01(\x08\x12\x16\n\x0eopenai_autolog\x18\x33 \x01(\x08\x12\x18\n\x10langchain_tracer\x18\x34 \x01(\x08\x12\x16\n\x0e\x63ohere_autolog\x18\x35 \x01(\x08\x12\x1b\n\x13hf_pipeline_autolog\x18\x36 \x01(\x08\x12\r\n\x05lib_c\x18\x38 \x01(\x08\x12\x0f\n\x07lib_cpp\x18\x39 \x01(\x08\x12\x19\n\x11openai_finetuning\x18: \x01(\x08\x12\x19\n\x11\x64iffusers_autolog\x18; \x01(\x08\x12\x1f\n\x17lightning_fabric_logger\x18< \x01(\x08\x12\x14\n\x0cset_step_log\x18= \x01(\x08\x12\x13\n\x0bset_summary\x18> \x01(\x08\x12\x16\n\x0emetric_summary\x18? \x01(\x08\x12\x13\n\x0bmetric_goal\x18@ \x01(\x08\x12\x15\n\rmetric_hidden\x18\x41 \x01(\x08\x12\x18\n\x10metric_step_sync\x18\x42 \x01(\x08\x12\x13\n\x0bshared_mode\x18\x43 \x01(\x08\x12#\n\x1bserver_side_derived_summary\x18\x44 \x01(\x08\x12\x1b\n\x13user_provided_label\x18\x45 \x01(\x08\x12\x1e\n\x16\x64\x63gm_profiling_enabled\x18\x46 \x01(\x08\x12\x11\n\tfork_mode\x18G \x01(\x08\x12\x13\n\x0brewind_mode\x18H \x01(\x08J\x04\x08\x17\x10\x18J\x04\x08-\x10.J\x04\x08\x37\x10\x38\"\xc9\x01\n\x03\x45nv\x12\x0f\n\x07jupyter\x18\x01 \x01(\x08\x12\x0e\n\x06kaggle\x18\x02 \x01(\x08\x12\x0f\n\x07windows\x18\x03 \x01(\x08\x12\x0e\n\x06m1_gpu\x18\x04 \x01(\x08\x12\x10\n\x08trainium\x18\n \x01(\x08\x12\x0b\n\x03pex\x18\x0b \x01(\x08\x12\r\n\x05\x63olab\x18\x0c \x01(\x08\x12\x0f\n\x07ipython\x18\r \x01(\x08\x12\x12\n\naws_lambda\x18\x0e \x01(\x08\x12\x0f\n\x07\x61md_gpu\x18\x0f \x01(\x08J\x04\x08\x05\x10\x06J\x04\x08\x06\x10\x07J\x04\x08\x07\x10\x08J\x04\x08\x08\x10\tJ\x04\x08\t\x10\n\"H\n\x06Labels\x12\x13\n\x0b\x63ode_string\x18\x01 \x01(\t\x12\x13\n\x0brepo_string\x18\x02 \x01(\t\x12\x14\n\x0c\x63ode_version\x18\x03 \x01(\t\"\xd5\x05\n\nDeprecated\x12!\n\x19keras_callback__data_type\x18\x01 \x01(\x08\x12\r\n\x05plots\x18\x05 \x01(\x08\x12!\n\x19init__config_include_keys\x18\x07 \x01(\x08\x12!\n\x19init__config_exclude_keys\x18\x08 \x01(\x08\x12\"\n\x1akeras_callback__save_model\x18\t \x01(\x08\x12\x18\n\x10langchain_tracer\x18\n \x01(\x08\x12\x1a\n\x12\x61rtifact__get_path\x18\x0b \x01(\x08\x12#\n\x1b\x61rtifactmanifestentry__name\x18\x0c \x01(\x08\x12\x1e\n\x16\x61pi__artifact_versions\x18\r \x01(\x08\x12(\n artifact_collection__change_type\x18\x0e \x01(\x08\x12\x1f\n\x17run__define_metric_copy\x18\x0f \x01(\x08\x12\x14\n\x0crun_disabled\x18\x10 \x01(\x08\x12\x16\n\x0ekeras_callback\x18\x11 \x01(\x08\x12$\n\x1crun__define_metric_best_goal\x18\x12 \x01(\x08\x12\x19\n\x11run__finish_quiet\x18\x13 \x01(\x08\x12\x18\n\x10run__reinit_bool\x18\x14 \x01(\x08\x12\x14\n\x0crun__get_url\x18\x15 \x01(\x08\x12\x19\n\x11run__project_name\x18\x16 \x01(\x08\x12\x1c\n\x14run__get_project_url\x18\x17 \x01(\x08\x12\x1a\n\x12run__get_sweep_url\x18\x18 \x01(\x08\x12 \n\x18run__use_artifact_use_as\x18\x19 \x01(\x08\x12\x18\n\x10\x61rtifact__use_as\x18\x1a \x01(\x08\x12\x1d\n\x15\x61rtifact__init_use_as\x18\x1b \x01(\x08J\x04\x08\x02\x10\x03J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05J\x04\x08\x06\x10\x07\"|\n\x06Issues\x12%\n\x1dsettings__validation_warnings\x18\x01 \x01(\x08\x12!\n\x19settings__unexpected_args\x18\x02 \x01(\x08\x12(\n settings__preprocessing_warnings\x18\x03 \x01(\x08\x42\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + + + +_TELEMETRYRECORD = DESCRIPTOR.message_types_by_name['TelemetryRecord'] +_TELEMETRYRESULT = DESCRIPTOR.message_types_by_name['TelemetryResult'] +_IMPORTS = DESCRIPTOR.message_types_by_name['Imports'] +_FEATURE = DESCRIPTOR.message_types_by_name['Feature'] +_ENV = DESCRIPTOR.message_types_by_name['Env'] +_LABELS = DESCRIPTOR.message_types_by_name['Labels'] +_DEPRECATED = DESCRIPTOR.message_types_by_name['Deprecated'] +_ISSUES = DESCRIPTOR.message_types_by_name['Issues'] +TelemetryRecord = _reflection.GeneratedProtocolMessageType('TelemetryRecord', (_message.Message,), { + 'DESCRIPTOR' : _TELEMETRYRECORD, + '__module__' : 'wandb.proto.wandb_telemetry_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TelemetryRecord) + }) +_sym_db.RegisterMessage(TelemetryRecord) + +TelemetryResult = _reflection.GeneratedProtocolMessageType('TelemetryResult', (_message.Message,), { + 'DESCRIPTOR' : _TELEMETRYRESULT, + '__module__' : 'wandb.proto.wandb_telemetry_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.TelemetryResult) + }) +_sym_db.RegisterMessage(TelemetryResult) + +Imports = _reflection.GeneratedProtocolMessageType('Imports', (_message.Message,), { + 'DESCRIPTOR' : _IMPORTS, + '__module__' : 'wandb.proto.wandb_telemetry_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Imports) + }) +_sym_db.RegisterMessage(Imports) + +Feature = _reflection.GeneratedProtocolMessageType('Feature', (_message.Message,), { + 'DESCRIPTOR' : _FEATURE, + '__module__' : 'wandb.proto.wandb_telemetry_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Feature) + }) +_sym_db.RegisterMessage(Feature) + +Env = _reflection.GeneratedProtocolMessageType('Env', (_message.Message,), { + 'DESCRIPTOR' : _ENV, + '__module__' : 'wandb.proto.wandb_telemetry_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Env) + }) +_sym_db.RegisterMessage(Env) + +Labels = _reflection.GeneratedProtocolMessageType('Labels', (_message.Message,), { + 'DESCRIPTOR' : _LABELS, + '__module__' : 'wandb.proto.wandb_telemetry_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Labels) + }) +_sym_db.RegisterMessage(Labels) + +Deprecated = _reflection.GeneratedProtocolMessageType('Deprecated', (_message.Message,), { + 'DESCRIPTOR' : _DEPRECATED, + '__module__' : 'wandb.proto.wandb_telemetry_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Deprecated) + }) +_sym_db.RegisterMessage(Deprecated) + +Issues = _reflection.GeneratedProtocolMessageType('Issues', (_message.Message,), { + 'DESCRIPTOR' : _ISSUES, + '__module__' : 'wandb.proto.wandb_telemetry_pb2' + # @@protoc_insertion_point(class_scope:wandb_internal.Issues) + }) +_sym_db.RegisterMessage(Issues) + +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + _TELEMETRYRECORD._serialized_start=84 + _TELEMETRYRECORD._serialized_end=559 + _TELEMETRYRESULT._serialized_start=561 + _TELEMETRYRESULT._serialized_end=578 + _IMPORTS._serialized_start=581 + _IMPORTS._serialized_end=2399 + _FEATURE._serialized_start=2402 + _FEATURE._serialized_end=4087 + _ENV._serialized_start=4090 + _ENV._serialized_end=4291 + _LABELS._serialized_start=4293 + _LABELS._serialized_end=4365 + _DEPRECATED._serialized_start=4368 + _DEPRECATED._serialized_end=5093 + _ISSUES._serialized_start=5095 + _ISSUES._serialized_end=5219 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v4/__init__.py b/lib/python3.12/site-packages/wandb/proto/v4/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b0ff704deea6476d57d77cc87674bed2250d5c9d Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_base_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_base_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f74e7d87075eaf7b6ff2f228eabab7543b39c9ae Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_base_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_internal_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_internal_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f7a5c13fc01cbc99de7a81a851bb799453fa5da6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_internal_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_server_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_server_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b090d73ddfc770e2500b9cd60c85beff0f838ae6 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_server_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_settings_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_settings_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ab0938ebff6bd2ac70a207672ca7e41e5b987a48 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_settings_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_telemetry_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_telemetry_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8322ed138b2a1ae429fe410119aa0cac4bd93389 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v4/__pycache__/wandb_telemetry_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v4/wandb_base_pb2.py b/lib/python3.12/site-packages/wandb/proto/v4/wandb_base_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..e1df6042049011a6f2b2277569e2d6f011979a90 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v4/wandb_base_pb2.py @@ -0,0 +1,30 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_base.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cwandb/proto/wandb_base.proto\x12\x0ewandb_internal\"6\n\x0b_RecordInfo\x12\x11\n\tstream_id\x18\x01 \x01(\t\x12\x14\n\x0c_tracelog_id\x18\x64 \x01(\t\"!\n\x0c_RequestInfo\x12\x11\n\tstream_id\x18\x01 \x01(\t\"#\n\x0b_ResultInfo\x12\x14\n\x0c_tracelog_id\x18\x64 \x01(\tB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_base_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + __RECORDINFO._serialized_start=48 + __RECORDINFO._serialized_end=102 + __REQUESTINFO._serialized_start=104 + __REQUESTINFO._serialized_end=137 + __RESULTINFO._serialized_start=139 + __RESULTINFO._serialized_end=174 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v4/wandb_internal_pb2.py b/lib/python3.12/site-packages/wandb/proto/v4/wandb_internal_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..61eb6111a30410e2f0527fe5cf55065c9ca43d37 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v4/wandb_internal_pb2.py @@ -0,0 +1,382 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_internal.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 +from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 +from wandb.proto import wandb_telemetry_pb2 as wandb_dot_proto_dot_wandb__telemetry__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n wandb/proto/wandb_internal.proto\x12\x0ewandb_internal\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1cwandb/proto/wandb_base.proto\x1a!wandb/proto/wandb_telemetry.proto\"\xcf\t\n\x06Record\x12\x0b\n\x03num\x18\x01 \x01(\x03\x12\x30\n\x07history\x18\x02 \x01(\x0b\x32\x1d.wandb_internal.HistoryRecordH\x00\x12\x30\n\x07summary\x18\x03 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecordH\x00\x12.\n\x06output\x18\x04 \x01(\x0b\x32\x1c.wandb_internal.OutputRecordH\x00\x12.\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x1c.wandb_internal.ConfigRecordH\x00\x12,\n\x05\x66iles\x18\x06 \x01(\x0b\x32\x1b.wandb_internal.FilesRecordH\x00\x12,\n\x05stats\x18\x07 \x01(\x0b\x32\x1b.wandb_internal.StatsRecordH\x00\x12\x32\n\x08\x61rtifact\x18\x08 \x01(\x0b\x32\x1e.wandb_internal.ArtifactRecordH\x00\x12,\n\x08tbrecord\x18\t \x01(\x0b\x32\x18.wandb_internal.TBRecordH\x00\x12,\n\x05\x61lert\x18\n \x01(\x0b\x32\x1b.wandb_internal.AlertRecordH\x00\x12\x34\n\ttelemetry\x18\x0b \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecordH\x00\x12.\n\x06metric\x18\x0c \x01(\x0b\x32\x1c.wandb_internal.MetricRecordH\x00\x12\x35\n\noutput_raw\x18\r \x01(\x0b\x32\x1f.wandb_internal.OutputRawRecordH\x00\x12(\n\x03run\x18\x11 \x01(\x0b\x32\x19.wandb_internal.RunRecordH\x00\x12-\n\x04\x65xit\x18\x12 \x01(\x0b\x32\x1d.wandb_internal.RunExitRecordH\x00\x12,\n\x05\x66inal\x18\x14 \x01(\x0b\x32\x1b.wandb_internal.FinalRecordH\x00\x12.\n\x06header\x18\x15 \x01(\x0b\x32\x1c.wandb_internal.HeaderRecordH\x00\x12.\n\x06\x66ooter\x18\x16 \x01(\x0b\x32\x1c.wandb_internal.FooterRecordH\x00\x12\x39\n\npreempting\x18\x17 \x01(\x0b\x32#.wandb_internal.RunPreemptingRecordH\x00\x12\x34\n\x12noop_link_artifact\x18\x18 \x01(\x0b\x32\x16.google.protobuf.EmptyH\x00\x12\x39\n\x0cuse_artifact\x18\x19 \x01(\x0b\x32!.wandb_internal.UseArtifactRecordH\x00\x12\x38\n\x0b\x65nvironment\x18\x1a \x01(\x0b\x32!.wandb_internal.EnvironmentRecordH\x00\x12*\n\x07request\x18\x64 \x01(\x0b\x32\x17.wandb_internal.RequestH\x00\x12(\n\x07\x63ontrol\x18\x10 \x01(\x0b\x32\x17.wandb_internal.Control\x12\x0c\n\x04uuid\x18\x13 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfoB\r\n\x0brecord_type\"\xa8\x01\n\x07\x43ontrol\x12\x10\n\x08req_resp\x18\x01 \x01(\x08\x12\r\n\x05local\x18\x02 \x01(\x08\x12\x10\n\x08relay_id\x18\x03 \x01(\t\x12\x14\n\x0cmailbox_slot\x18\x04 \x01(\t\x12\x13\n\x0b\x61lways_send\x18\x05 \x01(\x08\x12\x14\n\x0c\x66low_control\x18\x06 \x01(\x08\x12\x12\n\nend_offset\x18\x07 \x01(\x03\x12\x15\n\rconnection_id\x18\x08 \x01(\t\"\xf3\x03\n\x06Result\x12\x35\n\nrun_result\x18\x11 \x01(\x0b\x32\x1f.wandb_internal.RunUpdateResultH\x00\x12\x34\n\x0b\x65xit_result\x18\x12 \x01(\x0b\x32\x1d.wandb_internal.RunExitResultH\x00\x12\x33\n\nlog_result\x18\x14 \x01(\x0b\x32\x1d.wandb_internal.HistoryResultH\x00\x12\x37\n\x0esummary_result\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.SummaryResultH\x00\x12\x35\n\routput_result\x18\x16 \x01(\x0b\x32\x1c.wandb_internal.OutputResultH\x00\x12\x35\n\rconfig_result\x18\x17 \x01(\x0b\x32\x1c.wandb_internal.ConfigResultH\x00\x12,\n\x08response\x18\x64 \x01(\x0b\x32\x18.wandb_internal.ResponseH\x00\x12(\n\x07\x63ontrol\x18\x10 \x01(\x0b\x32\x17.wandb_internal.Control\x12\x0c\n\x04uuid\x18\x18 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._ResultInfoB\r\n\x0bresult_type\":\n\x0b\x46inalRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"b\n\x0bVersionInfo\x12\x10\n\x08producer\x18\x01 \x01(\t\x12\x14\n\x0cmin_consumer\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"n\n\x0cHeaderRecord\x12\x31\n\x0cversion_info\x18\x01 \x01(\x0b\x32\x1b.wandb_internal.VersionInfo\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\";\n\x0c\x46ooterRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"9\n\x0b\x42ranchPoint\x12\x0b\n\x03run\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01\x12\x0e\n\x06metric\x18\x03 \x01(\t\"\x91\x05\n\tRunRecord\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0e\n\x06\x65ntity\x18\x02 \x01(\t\x12\x0f\n\x07project\x18\x03 \x01(\t\x12,\n\x06\x63onfig\x18\x04 \x01(\x0b\x32\x1c.wandb_internal.ConfigRecord\x12.\n\x07summary\x18\x05 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecord\x12\x11\n\trun_group\x18\x06 \x01(\t\x12\x10\n\x08job_type\x18\x07 \x01(\t\x12\x14\n\x0c\x64isplay_name\x18\x08 \x01(\t\x12\r\n\x05notes\x18\t \x01(\t\x12\x0c\n\x04tags\x18\n \x03(\t\x12\x30\n\x08settings\x18\x0b \x01(\x0b\x32\x1e.wandb_internal.SettingsRecord\x12\x10\n\x08sweep_id\x18\x0c \x01(\t\x12\x0c\n\x04host\x18\r \x01(\t\x12\x15\n\rstarting_step\x18\x0e \x01(\x03\x12\x12\n\nstorage_id\x18\x10 \x01(\t\x12.\n\nstart_time\x18\x11 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07resumed\x18\x12 \x01(\x08\x12\x32\n\ttelemetry\x18\x13 \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecord\x12\x0f\n\x07runtime\x18\x14 \x01(\x05\x12*\n\x03git\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.GitRepoRecord\x12\x0e\n\x06\x66orked\x18\x16 \x01(\x08\x12\x31\n\x0c\x62ranch_point\x18\x17 \x01(\x0b\x32\x1b.wandb_internal.BranchPoint\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\";\n\rGitRepoRecord\x12\x1a\n\nremote_url\x18\x01 \x01(\tR\x06remote\x12\x0e\n\x06\x63ommit\x18\x02 \x01(\t\"c\n\x0fRunUpdateResult\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"\xac\x01\n\tErrorInfo\x12\x0f\n\x07message\x18\x01 \x01(\t\x12\x31\n\x04\x63ode\x18\x02 \x01(\x0e\x32#.wandb_internal.ErrorInfo.ErrorCode\"[\n\tErrorCode\x12\x0b\n\x07UNKNOWN\x10\x00\x12\x11\n\rCOMMUNICATION\x10\x01\x12\x12\n\x0e\x41UTHENTICATION\x10\x02\x12\t\n\x05USAGE\x10\x03\x12\x0f\n\x0bUNSUPPORTED\x10\x04\"`\n\rRunExitRecord\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12\x0f\n\x07runtime\x18\x02 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x0f\n\rRunExitResult\"B\n\x13RunPreemptingRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x15\n\x13RunPreemptingResult\"i\n\x0eSettingsRecord\x12*\n\x04item\x18\x01 \x03(\x0b\x32\x1c.wandb_internal.SettingsItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"/\n\x0cSettingsItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x1a\n\x0bHistoryStep\x12\x0b\n\x03num\x18\x01 \x01(\x03\"\x92\x01\n\rHistoryRecord\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.HistoryItem\x12)\n\x04step\x18\x02 \x01(\x0b\x32\x1b.wandb_internal.HistoryStep\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\x0bHistoryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0f\n\rHistoryResult\"\xdc\x01\n\x0cOutputRecord\x12<\n\x0boutput_type\x18\x01 \x01(\x0e\x32\'.wandb_internal.OutputRecord.OutputType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0c\n\x04line\x18\x03 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"$\n\nOutputType\x12\n\n\x06STDERR\x10\x00\x12\n\n\x06STDOUT\x10\x01\"\x0e\n\x0cOutputResult\"\xe2\x01\n\x0fOutputRawRecord\x12?\n\x0boutput_type\x18\x01 \x01(\x0e\x32*.wandb_internal.OutputRawRecord.OutputType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0c\n\x04line\x18\x03 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"$\n\nOutputType\x12\n\n\x06STDERR\x10\x00\x12\n\n\x06STDOUT\x10\x01\"\x11\n\x0fOutputRawResult\"\xb4\x03\n\x0cMetricRecord\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x11\n\tglob_name\x18\x02 \x01(\t\x12\x13\n\x0bstep_metric\x18\x04 \x01(\t\x12\x19\n\x11step_metric_index\x18\x05 \x01(\x05\x12.\n\x07options\x18\x06 \x01(\x0b\x32\x1d.wandb_internal.MetricOptions\x12.\n\x07summary\x18\x07 \x01(\x0b\x32\x1d.wandb_internal.MetricSummary\x12\x35\n\x04goal\x18\x08 \x01(\x0e\x32\'.wandb_internal.MetricRecord.MetricGoal\x12/\n\x08_control\x18\t \x01(\x0b\x32\x1d.wandb_internal.MetricControl\x12\x1a\n\x12\x65xpanded_from_glob\x18\n \x01(\x08\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\nMetricGoal\x12\x0e\n\nGOAL_UNSET\x10\x00\x12\x11\n\rGOAL_MINIMIZE\x10\x01\x12\x11\n\rGOAL_MAXIMIZE\x10\x02\"\x0e\n\x0cMetricResult\"C\n\rMetricOptions\x12\x11\n\tstep_sync\x18\x01 \x01(\x08\x12\x0e\n\x06hidden\x18\x02 \x01(\x08\x12\x0f\n\x07\x64\x65\x66ined\x18\x03 \x01(\x08\"\"\n\rMetricControl\x12\x11\n\toverwrite\x18\x01 \x01(\x08\"o\n\rMetricSummary\x12\x0b\n\x03min\x18\x01 \x01(\x08\x12\x0b\n\x03max\x18\x02 \x01(\x08\x12\x0c\n\x04mean\x18\x03 \x01(\x08\x12\x0c\n\x04\x62\x65st\x18\x04 \x01(\x08\x12\x0c\n\x04last\x18\x05 \x01(\x08\x12\x0c\n\x04none\x18\x06 \x01(\x08\x12\x0c\n\x04\x63opy\x18\x07 \x01(\x08\"\x93\x01\n\x0c\x43onfigRecord\x12*\n\x06update\x18\x01 \x03(\x0b\x32\x1a.wandb_internal.ConfigItem\x12*\n\x06remove\x18\x02 \x03(\x0b\x32\x1a.wandb_internal.ConfigItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"A\n\nConfigItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0e\n\x0c\x43onfigResult\"\x96\x01\n\rSummaryRecord\x12+\n\x06update\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\x12+\n\x06remove\x18\x02 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\x0bSummaryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0f\n\rSummaryResult\"d\n\x0b\x46ilesRecord\x12(\n\x05\x66iles\x18\x01 \x03(\x0b\x32\x19.wandb_internal.FilesItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\xec\x01\n\tFilesItem\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x34\n\x06policy\x18\x02 \x01(\x0e\x32$.wandb_internal.FilesItem.PolicyType\x12\x30\n\x04type\x18\x03 \x01(\x0e\x32\".wandb_internal.FilesItem.FileType\"(\n\nPolicyType\x12\x07\n\x03NOW\x10\x00\x12\x07\n\x03\x45ND\x10\x01\x12\x08\n\x04LIVE\x10\x02\"9\n\x08\x46ileType\x12\t\n\x05OTHER\x10\x00\x12\t\n\x05WANDB\x10\x01\x12\t\n\x05MEDIA\x10\x02\x12\x0c\n\x08\x41RTIFACT\x10\x03J\x04\x08\x10\x10\x11\"\r\n\x0b\x46ilesResult\"\xe6\x01\n\x0bStatsRecord\x12\x39\n\nstats_type\x18\x01 \x01(\x0e\x32%.wandb_internal.StatsRecord.StatsType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\'\n\x04item\x18\x03 \x03(\x0b\x32\x19.wandb_internal.StatsItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x17\n\tStatsType\x12\n\n\x06SYSTEM\x10\x00\",\n\tStatsItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\xe7\x03\n\x0e\x41rtifactRecord\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0f\n\x07project\x18\x02 \x01(\t\x12\x0e\n\x06\x65ntity\x18\x03 \x01(\t\x12\x0c\n\x04type\x18\x04 \x01(\t\x12\x0c\n\x04name\x18\x05 \x01(\t\x12\x0e\n\x06\x64igest\x18\x06 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x07 \x01(\t\x12\x10\n\x08metadata\x18\x08 \x01(\t\x12\x14\n\x0cuser_created\x18\t \x01(\x08\x12\x18\n\x10use_after_commit\x18\n \x01(\x08\x12\x0f\n\x07\x61liases\x18\x0b \x03(\t\x12\x32\n\x08manifest\x18\x0c \x01(\x0b\x32 .wandb_internal.ArtifactManifest\x12\x16\n\x0e\x64istributed_id\x18\r \x01(\t\x12\x10\n\x08\x66inalize\x18\x0e \x01(\x08\x12\x11\n\tclient_id\x18\x0f \x01(\t\x12\x1a\n\x12sequence_client_id\x18\x10 \x01(\t\x12\x0f\n\x07\x62\x61se_id\x18\x11 \x01(\t\x12\x1c\n\x14ttl_duration_seconds\x18\x12 \x01(\x03\x12\x0c\n\x04tags\x18\x13 \x03(\t\x12\x19\n\x11incremental_beta1\x18\x64 \x01(\x08\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\xd8\x01\n\x10\x41rtifactManifest\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x16\n\x0estorage_policy\x18\x02 \x01(\t\x12\x46\n\x15storage_policy_config\x18\x03 \x03(\x0b\x32\'.wandb_internal.StoragePolicyConfigItem\x12\x37\n\x08\x63ontents\x18\x04 \x03(\x0b\x32%.wandb_internal.ArtifactManifestEntry\x12\x1a\n\x12manifest_file_path\x18\x05 \x01(\t\"\xcf\x01\n\x15\x41rtifactManifestEntry\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x0e\n\x06\x64igest\x18\x02 \x01(\t\x12\x0b\n\x03ref\x18\x03 \x01(\t\x12\x0c\n\x04size\x18\x04 \x01(\x03\x12\x10\n\x08mimetype\x18\x05 \x01(\t\x12\x12\n\nlocal_path\x18\x06 \x01(\t\x12\x19\n\x11\x62irth_artifact_id\x18\x07 \x01(\t\x12\x12\n\nskip_cache\x18\x08 \x01(\x08\x12(\n\x05\x65xtra\x18\x10 \x03(\x0b\x32\x19.wandb_internal.ExtraItem\",\n\tExtraItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x02 \x01(\t\":\n\x17StoragePolicyConfigItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x02 \x01(\t\"\x10\n\x0e\x41rtifactResult\"\x14\n\x12LinkArtifactResult\"\xf0\x01\n\x13LinkArtifactRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x11\n\tserver_id\x18\x02 \x01(\t\x12\x16\n\x0eportfolio_name\x18\x03 \x01(\t\x12\x18\n\x10portfolio_entity\x18\x04 \x01(\t\x12\x19\n\x11portfolio_project\x18\x05 \x01(\t\x12\x19\n\x11portfolio_aliases\x18\x06 \x03(\t\x12\x1e\n\x16portfolio_organization\x18\x07 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"[\n\x14LinkArtifactResponse\x12\x15\n\rerror_message\x18\x01 \x01(\t\x12\x1a\n\rversion_index\x18\x02 \x01(\x05H\x00\x88\x01\x01\x42\x10\n\x0e_version_index\"h\n\x08TBRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\x12\x0f\n\x07log_dir\x18\x01 \x01(\t\x12\x10\n\x08root_dir\x18\x03 \x01(\t\x12\x0c\n\x04save\x18\x02 \x01(\x08\"\n\n\x08TBResult\"}\n\x0b\x41lertRecord\x12\r\n\x05title\x18\x01 \x01(\t\x12\x0c\n\x04text\x18\x02 \x01(\t\x12\r\n\x05level\x18\x03 \x01(\t\x12\x15\n\rwait_duration\x18\x04 \x01(\x03\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\r\n\x0b\x41lertResult\"\xb3\x11\n\x07Request\x12\x38\n\x0bstop_status\x18\x01 \x01(\x0b\x32!.wandb_internal.StopStatusRequestH\x00\x12>\n\x0enetwork_status\x18\x02 \x01(\x0b\x32$.wandb_internal.NetworkStatusRequestH\x00\x12-\n\x05\x64\x65\x66\x65r\x18\x03 \x01(\x0b\x32\x1c.wandb_internal.DeferRequestH\x00\x12\x38\n\x0bget_summary\x18\x04 \x01(\x0b\x32!.wandb_internal.GetSummaryRequestH\x00\x12-\n\x05login\x18\x05 \x01(\x0b\x32\x1c.wandb_internal.LoginRequestH\x00\x12-\n\x05pause\x18\x06 \x01(\x0b\x32\x1c.wandb_internal.PauseRequestH\x00\x12/\n\x06resume\x18\x07 \x01(\x0b\x32\x1d.wandb_internal.ResumeRequestH\x00\x12\x34\n\tpoll_exit\x18\x08 \x01(\x0b\x32\x1f.wandb_internal.PollExitRequestH\x00\x12@\n\x0fsampled_history\x18\t \x01(\x0b\x32%.wandb_internal.SampledHistoryRequestH\x00\x12@\n\x0fpartial_history\x18\n \x01(\x0b\x32%.wandb_internal.PartialHistoryRequestH\x00\x12\x34\n\trun_start\x18\x0b \x01(\x0b\x32\x1f.wandb_internal.RunStartRequestH\x00\x12<\n\rcheck_version\x18\x0c \x01(\x0b\x32#.wandb_internal.CheckVersionRequestH\x00\x12:\n\x0clog_artifact\x18\r \x01(\x0b\x32\".wandb_internal.LogArtifactRequestH\x00\x12\x44\n\x11\x64ownload_artifact\x18\x0e \x01(\x0b\x32\'.wandb_internal.DownloadArtifactRequestH\x00\x12\x35\n\tkeepalive\x18\x11 \x01(\x0b\x32 .wandb_internal.KeepaliveRequestH\x00\x12>\n\x0eserver_feature\x18\x12 \x01(\x0b\x32$.wandb_internal.ServerFeatureRequestH\x00\x12\x36\n\nrun_status\x18\x14 \x01(\x0b\x32 .wandb_internal.RunStatusRequestH\x00\x12/\n\x06\x63\x61ncel\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.CancelRequestH\x00\x12\x44\n\x11internal_messages\x18\x17 \x01(\x0b\x32\'.wandb_internal.InternalMessagesRequestH\x00\x12@\n\x0fpython_packages\x18\x18 \x01(\x0b\x32%.wandb_internal.PythonPackagesRequestH\x00\x12\x33\n\x08shutdown\x18@ \x01(\x0b\x32\x1f.wandb_internal.ShutdownRequestH\x00\x12/\n\x06\x61ttach\x18\x41 \x01(\x0b\x32\x1d.wandb_internal.AttachRequestH\x00\x12/\n\x06status\x18\x42 \x01(\x0b\x32\x1d.wandb_internal.StatusRequestH\x00\x12\x38\n\x0bserver_info\x18\x43 \x01(\x0b\x32!.wandb_internal.ServerInfoRequestH\x00\x12\x38\n\x0bsender_mark\x18\x44 \x01(\x0b\x32!.wandb_internal.SenderMarkRequestH\x00\x12\x38\n\x0bsender_read\x18\x45 \x01(\x0b\x32!.wandb_internal.SenderReadRequestH\x00\x12<\n\rstatus_report\x18\x46 \x01(\x0b\x32#.wandb_internal.StatusReportRequestH\x00\x12>\n\x0esummary_record\x18G \x01(\x0b\x32$.wandb_internal.SummaryRecordRequestH\x00\x12\x42\n\x10telemetry_record\x18H \x01(\x0b\x32&.wandb_internal.TelemetryRecordRequestH\x00\x12\x32\n\x08job_info\x18I \x01(\x0b\x32\x1e.wandb_internal.JobInfoRequestH\x00\x12\x45\n\x12get_system_metrics\x18J \x01(\x0b\x32\'.wandb_internal.GetSystemMetricsRequestH\x00\x12\x34\n\tjob_input\x18M \x01(\x0b\x32\x1f.wandb_internal.JobInputRequestH\x00\x12<\n\rlink_artifact\x18N \x01(\x0b\x32#.wandb_internal.LinkArtifactRequestH\x00\x12N\n\x17run_finish_without_exit\x18O \x01(\x0b\x32+.wandb_internal.RunFinishWithoutExitRequestH\x00\x12\x38\n\x0bsync_finish\x18Q \x01(\x0b\x32!.wandb_internal.SyncFinishRequestH\x00\x12;\n\noperations\x18R \x01(\x0b\x32%.wandb_internal.OperationStatsRequestH\x00\x12\x39\n\x0btest_inject\x18\xe8\x07 \x01(\x0b\x32!.wandb_internal.TestInjectRequestH\x00\x42\x0e\n\x0crequest_typeJ\x04\x08\x16\x10\x17J\x04\x08K\x10LJ\x04\x08L\x10MJ\x04\x08P\x10Q\"\x9b\x0e\n\x08Response\x12?\n\x12keepalive_response\x18\x12 \x01(\x0b\x32!.wandb_internal.KeepaliveResponseH\x00\x12\x42\n\x14stop_status_response\x18\x13 \x01(\x0b\x32\".wandb_internal.StopStatusResponseH\x00\x12H\n\x17network_status_response\x18\x14 \x01(\x0b\x32%.wandb_internal.NetworkStatusResponseH\x00\x12\x37\n\x0elogin_response\x18\x18 \x01(\x0b\x32\x1d.wandb_internal.LoginResponseH\x00\x12\x42\n\x14get_summary_response\x18\x19 \x01(\x0b\x32\".wandb_internal.GetSummaryResponseH\x00\x12>\n\x12poll_exit_response\x18\x1a \x01(\x0b\x32 .wandb_internal.PollExitResponseH\x00\x12J\n\x18sampled_history_response\x18\x1b \x01(\x0b\x32&.wandb_internal.SampledHistoryResponseH\x00\x12>\n\x12run_start_response\x18\x1c \x01(\x0b\x32 .wandb_internal.RunStartResponseH\x00\x12\x46\n\x16\x63heck_version_response\x18\x1d \x01(\x0b\x32$.wandb_internal.CheckVersionResponseH\x00\x12\x44\n\x15log_artifact_response\x18\x1e \x01(\x0b\x32#.wandb_internal.LogArtifactResponseH\x00\x12N\n\x1a\x64ownload_artifact_response\x18\x1f \x01(\x0b\x32(.wandb_internal.DownloadArtifactResponseH\x00\x12H\n\x17server_feature_response\x18 \x01(\x0b\x32%.wandb_internal.ServerFeatureResponseH\x00\x12@\n\x13run_status_response\x18# \x01(\x0b\x32!.wandb_internal.RunStatusResponseH\x00\x12\x39\n\x0f\x63\x61ncel_response\x18$ \x01(\x0b\x32\x1e.wandb_internal.CancelResponseH\x00\x12N\n\x1ainternal_messages_response\x18% \x01(\x0b\x32(.wandb_internal.InternalMessagesResponseH\x00\x12=\n\x11shutdown_response\x18@ \x01(\x0b\x32 .wandb_internal.ShutdownResponseH\x00\x12\x39\n\x0f\x61ttach_response\x18\x41 \x01(\x0b\x32\x1e.wandb_internal.AttachResponseH\x00\x12\x39\n\x0fstatus_response\x18\x42 \x01(\x0b\x32\x1e.wandb_internal.StatusResponseH\x00\x12\x42\n\x14server_info_response\x18\x43 \x01(\x0b\x32\".wandb_internal.ServerInfoResponseH\x00\x12<\n\x11job_info_response\x18\x44 \x01(\x0b\x32\x1f.wandb_internal.JobInfoResponseH\x00\x12O\n\x1bget_system_metrics_response\x18\x45 \x01(\x0b\x32(.wandb_internal.GetSystemMetricsResponseH\x00\x12\x46\n\x16link_artifact_response\x18G \x01(\x0b\x32$.wandb_internal.LinkArtifactResponseH\x00\x12\x35\n\rsync_response\x18\x46 \x01(\x0b\x32\x1c.wandb_internal.SyncResponseH\x00\x12X\n run_finish_without_exit_response\x18H \x01(\x0b\x32,.wandb_internal.RunFinishWithoutExitResponseH\x00\x12\x45\n\x13operations_response\x18J \x01(\x0b\x32&.wandb_internal.OperationStatsResponseH\x00\x12\x43\n\x14test_inject_response\x18\xe8\x07 \x01(\x0b\x32\".wandb_internal.TestInjectResponseH\x00\x42\x0f\n\rresponse_typeJ\x04\x08I\x10J\"\xc0\x02\n\x0c\x44\x65\x66\x65rRequest\x12\x36\n\x05state\x18\x01 \x01(\x0e\x32\'.wandb_internal.DeferRequest.DeferState\"\xf7\x01\n\nDeferState\x12\t\n\x05\x42\x45GIN\x10\x00\x12\r\n\tFLUSH_RUN\x10\x01\x12\x0f\n\x0b\x46LUSH_STATS\x10\x02\x12\x19\n\x15\x46LUSH_PARTIAL_HISTORY\x10\x03\x12\x0c\n\x08\x46LUSH_TB\x10\x04\x12\r\n\tFLUSH_SUM\x10\x05\x12\x13\n\x0f\x46LUSH_DEBOUNCER\x10\x06\x12\x10\n\x0c\x46LUSH_OUTPUT\x10\x07\x12\r\n\tFLUSH_JOB\x10\x08\x12\r\n\tFLUSH_DIR\x10\t\x12\x0c\n\x08\x46LUSH_FP\x10\n\x12\x0b\n\x07JOIN_FP\x10\x0b\x12\x0c\n\x08\x46LUSH_FS\x10\x0c\x12\x0f\n\x0b\x46LUSH_FINAL\x10\r\x12\x07\n\x03\x45ND\x10\x0e\"<\n\x0cPauseRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x0f\n\rPauseResponse\"=\n\rResumeRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x10\n\x0eResumeResponse\"M\n\x0cLoginRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"&\n\rLoginResponse\x12\x15\n\ractive_entity\x18\x01 \x01(\t\"A\n\x11GetSummaryRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"?\n\x12GetSummaryResponse\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\"G\n\x17GetSystemMetricsRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"R\n\x12SystemMetricSample\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\r\n\x05value\x18\x02 \x01(\x02\"I\n\x13SystemMetricsBuffer\x12\x32\n\x06record\x18\x01 \x03(\x0b\x32\".wandb_internal.SystemMetricSample\"\xca\x01\n\x18GetSystemMetricsResponse\x12S\n\x0esystem_metrics\x18\x01 \x03(\x0b\x32;.wandb_internal.GetSystemMetricsResponse.SystemMetricsEntry\x1aY\n\x12SystemMetricsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x32\n\x05value\x18\x02 \x01(\x0b\x32#.wandb_internal.SystemMetricsBuffer:\x02\x38\x01\"=\n\rStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\")\n\x0eStatusResponse\x12\x17\n\x0frun_should_stop\x18\x01 \x01(\x08\"A\n\x11StopStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"-\n\x12StopStatusResponse\x12\x17\n\x0frun_should_stop\x18\x01 \x01(\x08\"D\n\x14NetworkStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"P\n\x15NetworkStatusResponse\x12\x37\n\x11network_responses\x18\x01 \x03(\x0b\x32\x1c.wandb_internal.HttpResponse\"D\n\x0cHttpResponse\x12\x18\n\x10http_status_code\x18\x01 \x01(\x05\x12\x1a\n\x12http_response_text\x18\x02 \x01(\t\"G\n\x17InternalMessagesRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"N\n\x18InternalMessagesResponse\x12\x32\n\x08messages\x18\x01 \x01(\x0b\x32 .wandb_internal.InternalMessages\"#\n\x10InternalMessages\x12\x0f\n\x07warning\x18\x01 \x03(\t\"?\n\x0fPollExitRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\xf5\x01\n\x10PollExitResponse\x12\x0c\n\x04\x64one\x18\x01 \x01(\x08\x12\x32\n\x0b\x65xit_result\x18\x02 \x01(\x0b\x32\x1d.wandb_internal.RunExitResult\x12\x35\n\x0cpusher_stats\x18\x03 \x01(\x0b\x32\x1f.wandb_internal.FilePusherStats\x12/\n\x0b\x66ile_counts\x18\x04 \x01(\x0b\x32\x1a.wandb_internal.FileCounts\x12\x37\n\x0foperation_stats\x18\x05 \x01(\x0b\x32\x1e.wandb_internal.OperationStats\"E\n\x15OperationStatsRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"Q\n\x16OperationStatsResponse\x12\x37\n\x0foperation_stats\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.OperationStats\"Y\n\x0eOperationStats\x12-\n\noperations\x18\x01 \x03(\x0b\x32\x19.wandb_internal.Operation\x12\x18\n\x10total_operations\x18\x02 \x01(\x03\"\x87\x01\n\tOperation\x12\x0c\n\x04\x64\x65sc\x18\x01 \x01(\t\x12\x17\n\x0fruntime_seconds\x18\x02 \x01(\x01\x12\x10\n\x08progress\x18\x03 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x04 \x01(\t\x12+\n\x08subtasks\x18\x05 \x03(\x0b\x32\x19.wandb_internal.Operation\"\x13\n\x11SenderMarkRequest\"\x13\n\x11SyncFinishRequest\"E\n\x0cSyncResponse\x12\x0b\n\x03url\x18\x01 \x01(\t\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"?\n\x11SenderReadRequest\x12\x14\n\x0cstart_offset\x18\x01 \x01(\x03\x12\x14\n\x0c\x66inal_offset\x18\x02 \x01(\x03\"m\n\x13StatusReportRequest\x12\x12\n\nrecord_num\x18\x01 \x01(\x03\x12\x13\n\x0bsent_offset\x18\x02 \x01(\x03\x12-\n\tsync_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"F\n\x14SummaryRecordRequest\x12.\n\x07summary\x18\x01 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecord\"L\n\x16TelemetryRecordRequest\x12\x32\n\ttelemetry\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecord\"A\n\x11ServerInfoRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"|\n\x12ServerInfoResponse\x12-\n\nlocal_info\x18\x01 \x01(\x0b\x32\x19.wandb_internal.LocalInfo\x12\x37\n\x0fserver_messages\x18\x02 \x01(\x0b\x32\x1e.wandb_internal.ServerMessages\"=\n\x0eServerMessages\x12+\n\x04item\x18\x01 \x03(\x0b\x32\x1d.wandb_internal.ServerMessage\"e\n\rServerMessage\x12\x12\n\nplain_text\x18\x01 \x01(\t\x12\x10\n\x08utf_text\x18\x02 \x01(\t\x12\x11\n\thtml_text\x18\x03 \x01(\t\x12\x0c\n\x04type\x18\x04 \x01(\t\x12\r\n\x05level\x18\x05 \x01(\x05\"c\n\nFileCounts\x12\x13\n\x0bwandb_count\x18\x01 \x01(\x05\x12\x13\n\x0bmedia_count\x18\x02 \x01(\x05\x12\x16\n\x0e\x61rtifact_count\x18\x03 \x01(\x05\x12\x13\n\x0bother_count\x18\x04 \x01(\x05\"U\n\x0f\x46ilePusherStats\x12\x16\n\x0euploaded_bytes\x18\x01 \x01(\x03\x12\x13\n\x0btotal_bytes\x18\x02 \x01(\x03\x12\x15\n\rdeduped_bytes\x18\x03 \x01(\x03\"\x1e\n\rFilesUploaded\x12\r\n\x05\x66iles\x18\x01 \x03(\t\"\xf4\x01\n\x17\x46ileTransferInfoRequest\x12\x42\n\x04type\x18\x01 \x01(\x0e\x32\x34.wandb_internal.FileTransferInfoRequest.TransferType\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\x0b\n\x03url\x18\x03 \x01(\t\x12\x0c\n\x04size\x18\x04 \x01(\x03\x12\x11\n\tprocessed\x18\x05 \x01(\x03\x12/\n\x0b\x66ile_counts\x18\x06 \x01(\x0b\x32\x1a.wandb_internal.FileCounts\"(\n\x0cTransferType\x12\n\n\x06Upload\x10\x00\x12\x0c\n\x08\x44ownload\x10\x01\"1\n\tLocalInfo\x12\x0f\n\x07version\x18\x01 \x01(\t\x12\x13\n\x0bout_of_date\x18\x02 \x01(\x08\"?\n\x0fShutdownRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x12\n\x10ShutdownResponse\"P\n\rAttachRequest\x12\x11\n\tattach_id\x18\x14 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"b\n\x0e\x41ttachResponse\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"\xd5\x02\n\x11TestInjectRequest\x12\x13\n\x0bhandler_exc\x18\x01 \x01(\x08\x12\x14\n\x0chandler_exit\x18\x02 \x01(\x08\x12\x15\n\rhandler_abort\x18\x03 \x01(\x08\x12\x12\n\nsender_exc\x18\x04 \x01(\x08\x12\x13\n\x0bsender_exit\x18\x05 \x01(\x08\x12\x14\n\x0csender_abort\x18\x06 \x01(\x08\x12\x0f\n\x07req_exc\x18\x07 \x01(\x08\x12\x10\n\x08req_exit\x18\x08 \x01(\x08\x12\x11\n\treq_abort\x18\t \x01(\x08\x12\x10\n\x08resp_exc\x18\n \x01(\x08\x12\x11\n\tresp_exit\x18\x0b \x01(\x08\x12\x12\n\nresp_abort\x18\x0c \x01(\x08\x12\x10\n\x08msg_drop\x18\r \x01(\x08\x12\x10\n\x08msg_hang\x18\x0e \x01(\x08\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x14\n\x12TestInjectResponse\"\x1e\n\rHistoryAction\x12\r\n\x05\x66lush\x18\x01 \x01(\x08\"\xca\x01\n\x15PartialHistoryRequest\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.HistoryItem\x12)\n\x04step\x18\x02 \x01(\x0b\x32\x1b.wandb_internal.HistoryStep\x12-\n\x06\x61\x63tion\x18\x03 \x01(\x0b\x32\x1d.wandb_internal.HistoryAction\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x18\n\x16PartialHistoryResponse\"E\n\x15SampledHistoryRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"_\n\x12SampledHistoryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x14\n\x0cvalues_float\x18\x03 \x03(\x02\x12\x12\n\nvalues_int\x18\x04 \x03(\x03\"J\n\x16SampledHistoryResponse\x12\x30\n\x04item\x18\x01 \x03(\x0b\x32\".wandb_internal.SampledHistoryItem\"@\n\x10RunStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"x\n\x11RunStatusResponse\x12\x18\n\x10sync_items_total\x18\x01 \x01(\x03\x12\x1a\n\x12sync_items_pending\x18\x02 \x01(\x03\x12-\n\tsync_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"g\n\x0fRunStartRequest\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x12\n\x10RunStartResponse\"K\n\x1bRunFinishWithoutExitRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x1e\n\x1cRunFinishWithoutExitResponse\"\\\n\x13\x43heckVersionRequest\x12\x17\n\x0f\x63urrent_version\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"]\n\x14\x43heckVersionResponse\x12\x17\n\x0fupgrade_message\x18\x01 \x01(\t\x12\x14\n\x0cyank_message\x18\x02 \x01(\t\x12\x16\n\x0e\x64\x65lete_message\x18\x03 \x01(\t\">\n\x0eJobInfoRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"6\n\x0fJobInfoResponse\x12\x12\n\nsequenceId\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"\x9f\x01\n\x12LogArtifactRequest\x12\x30\n\x08\x61rtifact\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.ArtifactRecord\x12\x14\n\x0chistory_step\x18\x02 \x01(\x03\x12\x13\n\x0bstaging_dir\x18\x03 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"A\n\x13LogArtifactResponse\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x15\n\rerror_message\x18\x02 \x01(\t\"\xbe\x01\n\x17\x44ownloadArtifactRequest\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x15\n\rdownload_root\x18\x02 \x01(\t\x12 \n\x18\x61llow_missing_references\x18\x04 \x01(\x08\x12\x12\n\nskip_cache\x18\x05 \x01(\x08\x12\x13\n\x0bpath_prefix\x18\x06 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"1\n\x18\x44ownloadArtifactResponse\x12\x15\n\rerror_message\x18\x01 \x01(\t\"@\n\x10KeepaliveRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x13\n\x11KeepaliveResponse\"q\n\x0c\x41rtifactInfo\x12\x10\n\x08\x61rtifact\x18\x01 \x01(\t\x12\x12\n\nentrypoint\x18\x02 \x03(\t\x12\x10\n\x08notebook\x18\x03 \x01(\x08\x12\x15\n\rbuild_context\x18\x04 \x01(\t\x12\x12\n\ndockerfile\x18\x05 \x01(\t\")\n\x07GitInfo\x12\x0e\n\x06remote\x18\x01 \x01(\t\x12\x0e\n\x06\x63ommit\x18\x02 \x01(\t\"\x87\x01\n\tGitSource\x12)\n\x08git_info\x18\x01 \x01(\x0b\x32\x17.wandb_internal.GitInfo\x12\x12\n\nentrypoint\x18\x02 \x03(\t\x12\x10\n\x08notebook\x18\x03 \x01(\x08\x12\x15\n\rbuild_context\x18\x04 \x01(\t\x12\x12\n\ndockerfile\x18\x05 \x01(\t\"\x1c\n\x0bImageSource\x12\r\n\x05image\x18\x01 \x01(\t\"\x8c\x01\n\x06Source\x12&\n\x03git\x18\x01 \x01(\x0b\x32\x19.wandb_internal.GitSource\x12.\n\x08\x61rtifact\x18\x02 \x01(\x0b\x32\x1c.wandb_internal.ArtifactInfo\x12*\n\x05image\x18\x03 \x01(\x0b\x32\x1b.wandb_internal.ImageSource\"k\n\tJobSource\x12\x10\n\x08_version\x18\x01 \x01(\t\x12\x13\n\x0bsource_type\x18\x02 \x01(\t\x12&\n\x06source\x18\x03 \x01(\x0b\x32\x16.wandb_internal.Source\x12\x0f\n\x07runtime\x18\x04 \x01(\t\"V\n\x12PartialJobArtifact\x12\x10\n\x08job_name\x18\x01 \x01(\t\x12.\n\x0bsource_info\x18\x02 \x01(\x0b\x32\x19.wandb_internal.JobSource\"\x9d\x01\n\x11UseArtifactRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x33\n\x07partial\x18\x04 \x01(\x0b\x32\".wandb_internal.PartialJobArtifact\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x13\n\x11UseArtifactResult\"R\n\rCancelRequest\x12\x13\n\x0b\x63\x61ncel_slot\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x10\n\x0e\x43\x61ncelResponse\"\'\n\x08\x44iskInfo\x12\r\n\x05total\x18\x01 \x01(\x04\x12\x0c\n\x04used\x18\x02 \x01(\x04\"\x1b\n\nMemoryInfo\x12\r\n\x05total\x18\x01 \x01(\x04\"/\n\x07\x43puInfo\x12\r\n\x05\x63ount\x18\x01 \x01(\r\x12\x15\n\rcount_logical\x18\x02 \x01(\r\"\x9a\x01\n\tAppleInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\necpu_cores\x18\x02 \x01(\r\x12\x12\n\npcpu_cores\x18\x03 \x01(\r\x12\x11\n\tgpu_cores\x18\x04 \x01(\r\x12\x11\n\tmemory_gb\x18\x05 \x01(\r\x12\x18\n\x10swap_total_bytes\x18\x06 \x01(\x04\x12\x17\n\x0fram_total_bytes\x18\x07 \x01(\x04\"k\n\rGpuNvidiaInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x14\n\x0cmemory_total\x18\x02 \x01(\x04\x12\x12\n\ncuda_cores\x18\x03 \x01(\r\x12\x14\n\x0c\x61rchitecture\x18\x04 \x01(\t\x12\x0c\n\x04uuid\x18\x05 \x01(\t\"\x89\x02\n\nGpuAmdInfo\x12\n\n\x02id\x18\x01 \x01(\t\x12\x11\n\tunique_id\x18\x02 \x01(\t\x12\x15\n\rvbios_version\x18\x03 \x01(\t\x12\x19\n\x11performance_level\x18\x04 \x01(\t\x12\x15\n\rgpu_overdrive\x18\x05 \x01(\t\x12\x1c\n\x14gpu_memory_overdrive\x18\x06 \x01(\t\x12\x11\n\tmax_power\x18\x07 \x01(\t\x12\x0e\n\x06series\x18\x08 \x01(\t\x12\r\n\x05model\x18\t \x01(\t\x12\x0e\n\x06vendor\x18\n \x01(\t\x12\x0b\n\x03sku\x18\x0b \x01(\t\x12\x12\n\nsclk_range\x18\x0c \x01(\t\x12\x12\n\nmclk_range\x18\r \x01(\t\"n\n\x0cTrainiumInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06vendor\x18\x02 \x01(\t\x12\x1b\n\x13neuron_device_count\x18\x03 \x01(\r\x12#\n\x1bneuroncore_per_device_count\x18\x04 \x01(\r\"Q\n\x07TPUInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07hbm_gib\x18\x02 \x01(\r\x12\x18\n\x10\x64\x65vices_per_chip\x18\x03 \x01(\r\x12\r\n\x05\x63ount\x18\x04 \x01(\r\"E\n\rCoreWeaveInfo\x12\x14\n\x0c\x63luster_name\x18\x01 \x01(\t\x12\x0e\n\x06org_id\x18\x02 \x01(\t\x12\x0e\n\x06region\x18\x03 \x01(\t\"\xa8\t\n\x11\x45nvironmentRecord\x12\n\n\x02os\x18\x01 \x01(\t\x12\x0e\n\x06python\x18\x02 \x01(\t\x12\x39\n\nstarted_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\tstartedAt\x12\x0e\n\x06\x64ocker\x18\x04 \x01(\t\x12\x0c\n\x04\x61rgs\x18\x05 \x03(\t\x12\x0f\n\x07program\x18\x06 \x01(\t\x12\x1b\n\tcode_path\x18\x07 \x01(\tR\x08\x63odePath\x12&\n\x0f\x63ode_path_local\x18\x08 \x01(\tR\rcodePathLocal\x12*\n\x03git\x18\t \x01(\x0b\x32\x1d.wandb_internal.GitRepoRecord\x12\r\n\x05\x65mail\x18\n \x01(\t\x12\x0c\n\x04root\x18\x0b \x01(\t\x12\x0c\n\x04host\x18\x0c \x01(\t\x12\x10\n\x08username\x18\r \x01(\t\x12\x12\n\nexecutable\x18\x0e \x01(\t\x12\r\n\x05\x63olab\x18\x0f \x01(\t\x12\x1c\n\tcpu_count\x18\x10 \x01(\rR\tcpu_count\x12,\n\x11\x63pu_count_logical\x18\x11 \x01(\rR\x11\x63pu_count_logical\x12\x15\n\x08gpu_type\x18\x12 \x01(\tR\x03gpu\x12\x1c\n\tgpu_count\x18\x13 \x01(\rR\tgpu_count\x12\x39\n\x04\x64isk\x18\x14 \x03(\x0b\x32+.wandb_internal.EnvironmentRecord.DiskEntry\x12*\n\x06memory\x18\x15 \x01(\x0b\x32\x1a.wandb_internal.MemoryInfo\x12$\n\x03\x63pu\x18\x16 \x01(\x0b\x32\x17.wandb_internal.CpuInfo\x12(\n\x05\x61pple\x18\x17 \x01(\x0b\x32\x19.wandb_internal.AppleInfo\x12=\n\ngpu_nvidia\x18\x18 \x03(\x0b\x32\x1d.wandb_internal.GpuNvidiaInfoR\ngpu_nvidia\x12\x14\n\x0c\x63uda_version\x18\x19 \x01(\t\x12\x34\n\x07gpu_amd\x18\x1a \x03(\x0b\x32\x1a.wandb_internal.GpuAmdInfoR\x07gpu_amd\x12;\n\x05slurm\x18\x1b \x03(\x0b\x32,.wandb_internal.EnvironmentRecord.SlurmEntry\x12.\n\x08trainium\x18\x1c \x01(\x0b\x32\x1c.wandb_internal.TrainiumInfo\x12$\n\x03tpu\x18\x1d \x01(\x0b\x32\x17.wandb_internal.TPUInfo\x12\x30\n\tcoreweave\x18\x1e \x01(\x0b\x32\x1d.wandb_internal.CoreWeaveInfo\x12\x12\n\twriter_id\x18\xc7\x01 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\x1a\x45\n\tDiskEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.wandb_internal.DiskInfo:\x02\x38\x01\x1a,\n\nSlurmEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\x15PythonPackagesRequest\x12\x44\n\x07package\x18\x01 \x03(\x0b\x32\x33.wandb_internal.PythonPackagesRequest.PythonPackage\x1a.\n\rPythonPackage\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"\x1c\n\x0cJobInputPath\x12\x0c\n\x04path\x18\x01 \x03(\t\"\xd6\x01\n\x0eJobInputSource\x12\x44\n\nrun_config\x18\x01 \x01(\x0b\x32..wandb_internal.JobInputSource.RunConfigSourceH\x00\x12?\n\x04\x66ile\x18\x02 \x01(\x0b\x32/.wandb_internal.JobInputSource.ConfigFileSourceH\x00\x1a\x11\n\x0fRunConfigSource\x1a \n\x10\x43onfigFileSource\x12\x0c\n\x04path\x18\x01 \x01(\tB\x08\n\x06source\"\xc7\x01\n\x0fJobInputRequest\x12\x34\n\x0cinput_source\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.JobInputSource\x12\x33\n\rinclude_paths\x18\x02 \x03(\x0b\x32\x1c.wandb_internal.JobInputPath\x12\x33\n\rexclude_paths\x18\x03 \x03(\x0b\x32\x1c.wandb_internal.JobInputPath\x12\x14\n\x0cinput_schema\x18\x04 \x01(\t\"t\n\x14ServerFeatureRequest\x12.\n\x07\x66\x65\x61ture\x18\x01 \x01(\x0e\x32\x1d.wandb_internal.ServerFeature\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"K\n\x15ServerFeatureResponse\x12\x32\n\x07\x66\x65\x61ture\x18\x01 \x01(\x0b\x32!.wandb_internal.ServerFeatureItem\"2\n\x11ServerFeatureItem\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07\x65nabled\x18\x02 \x01(\x08*\x80\x04\n\rServerFeature\x12\x13\n\x0fLARGE_FILENAMES\x10\x00\x12\x11\n\rARTIFACT_TAGS\x10\x01\x12\x0e\n\nCLIENT_IDS\x10\x02\x12\x1c\n\x18\x41RTIFACT_REGISTRY_SEARCH\x10\x03\x12\x1b\n\x17STRUCTURED_CONSOLE_LOGS\x10\x04\x12(\n$ARTIFACT_COLLECTION_MEMBERSHIP_FILES\x10\x05\x12\x38\n4ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER\x10\x06\x12\x34\n0USE_ARTIFACT_WITH_ENTITY_AND_PROJECT_INFORMATION\x10\x07\x12\x1f\n\x1b\x45XPAND_DEFINED_METRIC_GLOBS\x10\x08\x12\x1f\n\x1b\x41UTOMATION_EVENT_RUN_METRIC\x10\t\x12&\n\"AUTOMATION_EVENT_RUN_METRIC_CHANGE\x10\n\x12\x1b\n\x17\x41UTOMATION_ACTION_NO_OP\x10\x0b\x12/\n+INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION\x10\x0c\x12*\n&PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP\x10\rB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_internal_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY._options = None + _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY._serialized_options = b'8\001' + _ENVIRONMENTRECORD_DISKENTRY._options = None + _ENVIRONMENTRECORD_DISKENTRY._serialized_options = b'8\001' + _ENVIRONMENTRECORD_SLURMENTRY._options = None + _ENVIRONMENTRECORD_SLURMENTRY._serialized_options = b'8\001' + _SERVERFEATURE._serialized_start=22228 + _SERVERFEATURE._serialized_end=22740 + _RECORD._serialized_start=180 + _RECORD._serialized_end=1411 + _CONTROL._serialized_start=1414 + _CONTROL._serialized_end=1582 + _RESULT._serialized_start=1585 + _RESULT._serialized_end=2084 + _FINALRECORD._serialized_start=2086 + _FINALRECORD._serialized_end=2144 + _VERSIONINFO._serialized_start=2146 + _VERSIONINFO._serialized_end=2244 + _HEADERRECORD._serialized_start=2246 + _HEADERRECORD._serialized_end=2356 + _FOOTERRECORD._serialized_start=2358 + _FOOTERRECORD._serialized_end=2417 + _BRANCHPOINT._serialized_start=2419 + _BRANCHPOINT._serialized_end=2476 + _RUNRECORD._serialized_start=2479 + _RUNRECORD._serialized_end=3136 + _GITREPORECORD._serialized_start=3138 + _GITREPORECORD._serialized_end=3197 + _RUNUPDATERESULT._serialized_start=3199 + _RUNUPDATERESULT._serialized_end=3298 + _ERRORINFO._serialized_start=3301 + _ERRORINFO._serialized_end=3473 + _ERRORINFO_ERRORCODE._serialized_start=3382 + _ERRORINFO_ERRORCODE._serialized_end=3473 + _RUNEXITRECORD._serialized_start=3475 + _RUNEXITRECORD._serialized_end=3571 + _RUNEXITRESULT._serialized_start=3573 + _RUNEXITRESULT._serialized_end=3588 + _RUNPREEMPTINGRECORD._serialized_start=3590 + _RUNPREEMPTINGRECORD._serialized_end=3656 + _RUNPREEMPTINGRESULT._serialized_start=3658 + _RUNPREEMPTINGRESULT._serialized_end=3679 + _SETTINGSRECORD._serialized_start=3681 + _SETTINGSRECORD._serialized_end=3786 + _SETTINGSITEM._serialized_start=3788 + _SETTINGSITEM._serialized_end=3835 + _HISTORYSTEP._serialized_start=3837 + _HISTORYSTEP._serialized_end=3863 + _HISTORYRECORD._serialized_start=3866 + _HISTORYRECORD._serialized_end=4012 + _HISTORYITEM._serialized_start=4014 + _HISTORYITEM._serialized_end=4080 + _HISTORYRESULT._serialized_start=4082 + _HISTORYRESULT._serialized_end=4097 + _OUTPUTRECORD._serialized_start=4100 + _OUTPUTRECORD._serialized_end=4320 + _OUTPUTRECORD_OUTPUTTYPE._serialized_start=4284 + _OUTPUTRECORD_OUTPUTTYPE._serialized_end=4320 + _OUTPUTRESULT._serialized_start=4322 + _OUTPUTRESULT._serialized_end=4336 + _OUTPUTRAWRECORD._serialized_start=4339 + _OUTPUTRAWRECORD._serialized_end=4565 + _OUTPUTRAWRECORD_OUTPUTTYPE._serialized_start=4284 + _OUTPUTRAWRECORD_OUTPUTTYPE._serialized_end=4320 + _OUTPUTRAWRESULT._serialized_start=4567 + _OUTPUTRAWRESULT._serialized_end=4584 + _METRICRECORD._serialized_start=4587 + _METRICRECORD._serialized_end=5023 + _METRICRECORD_METRICGOAL._serialized_start=4957 + _METRICRECORD_METRICGOAL._serialized_end=5023 + _METRICRESULT._serialized_start=5025 + _METRICRESULT._serialized_end=5039 + _METRICOPTIONS._serialized_start=5041 + _METRICOPTIONS._serialized_end=5108 + _METRICCONTROL._serialized_start=5110 + _METRICCONTROL._serialized_end=5144 + _METRICSUMMARY._serialized_start=5146 + _METRICSUMMARY._serialized_end=5257 + _CONFIGRECORD._serialized_start=5260 + _CONFIGRECORD._serialized_end=5407 + _CONFIGITEM._serialized_start=5409 + _CONFIGITEM._serialized_end=5474 + _CONFIGRESULT._serialized_start=5476 + _CONFIGRESULT._serialized_end=5490 + _SUMMARYRECORD._serialized_start=5493 + _SUMMARYRECORD._serialized_end=5643 + _SUMMARYITEM._serialized_start=5645 + _SUMMARYITEM._serialized_end=5711 + _SUMMARYRESULT._serialized_start=5713 + _SUMMARYRESULT._serialized_end=5728 + _FILESRECORD._serialized_start=5730 + _FILESRECORD._serialized_end=5830 + _FILESITEM._serialized_start=5833 + _FILESITEM._serialized_end=6069 + _FILESITEM_POLICYTYPE._serialized_start=5964 + _FILESITEM_POLICYTYPE._serialized_end=6004 + _FILESITEM_FILETYPE._serialized_start=6006 + _FILESITEM_FILETYPE._serialized_end=6063 + _FILESRESULT._serialized_start=6071 + _FILESRESULT._serialized_end=6084 + _STATSRECORD._serialized_start=6087 + _STATSRECORD._serialized_end=6317 + _STATSRECORD_STATSTYPE._serialized_start=6294 + _STATSRECORD_STATSTYPE._serialized_end=6317 + _STATSITEM._serialized_start=6319 + _STATSITEM._serialized_end=6363 + _ARTIFACTRECORD._serialized_start=6366 + _ARTIFACTRECORD._serialized_end=6853 + _ARTIFACTMANIFEST._serialized_start=6856 + _ARTIFACTMANIFEST._serialized_end=7072 + _ARTIFACTMANIFESTENTRY._serialized_start=7075 + _ARTIFACTMANIFESTENTRY._serialized_end=7282 + _EXTRAITEM._serialized_start=7284 + _EXTRAITEM._serialized_end=7328 + _STORAGEPOLICYCONFIGITEM._serialized_start=7330 + _STORAGEPOLICYCONFIGITEM._serialized_end=7388 + _ARTIFACTRESULT._serialized_start=7390 + _ARTIFACTRESULT._serialized_end=7406 + _LINKARTIFACTRESULT._serialized_start=7408 + _LINKARTIFACTRESULT._serialized_end=7428 + _LINKARTIFACTREQUEST._serialized_start=7431 + _LINKARTIFACTREQUEST._serialized_end=7671 + _LINKARTIFACTRESPONSE._serialized_start=7673 + _LINKARTIFACTRESPONSE._serialized_end=7764 + _TBRECORD._serialized_start=7766 + _TBRECORD._serialized_end=7870 + _TBRESULT._serialized_start=7872 + _TBRESULT._serialized_end=7882 + _ALERTRECORD._serialized_start=7884 + _ALERTRECORD._serialized_end=8009 + _ALERTRESULT._serialized_start=8011 + _ALERTRESULT._serialized_end=8024 + _REQUEST._serialized_start=8027 + _REQUEST._serialized_end=10254 + _RESPONSE._serialized_start=10257 + _RESPONSE._serialized_end=12076 + _DEFERREQUEST._serialized_start=12079 + _DEFERREQUEST._serialized_end=12399 + _DEFERREQUEST_DEFERSTATE._serialized_start=12152 + _DEFERREQUEST_DEFERSTATE._serialized_end=12399 + _PAUSEREQUEST._serialized_start=12401 + _PAUSEREQUEST._serialized_end=12461 + _PAUSERESPONSE._serialized_start=12463 + _PAUSERESPONSE._serialized_end=12478 + _RESUMEREQUEST._serialized_start=12480 + _RESUMEREQUEST._serialized_end=12541 + _RESUMERESPONSE._serialized_start=12543 + _RESUMERESPONSE._serialized_end=12559 + _LOGINREQUEST._serialized_start=12561 + _LOGINREQUEST._serialized_end=12638 + _LOGINRESPONSE._serialized_start=12640 + _LOGINRESPONSE._serialized_end=12678 + _GETSUMMARYREQUEST._serialized_start=12680 + _GETSUMMARYREQUEST._serialized_end=12745 + _GETSUMMARYRESPONSE._serialized_start=12747 + _GETSUMMARYRESPONSE._serialized_end=12810 + _GETSYSTEMMETRICSREQUEST._serialized_start=12812 + _GETSYSTEMMETRICSREQUEST._serialized_end=12883 + _SYSTEMMETRICSAMPLE._serialized_start=12885 + _SYSTEMMETRICSAMPLE._serialized_end=12967 + _SYSTEMMETRICSBUFFER._serialized_start=12969 + _SYSTEMMETRICSBUFFER._serialized_end=13042 + _GETSYSTEMMETRICSRESPONSE._serialized_start=13045 + _GETSYSTEMMETRICSRESPONSE._serialized_end=13247 + _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY._serialized_start=13158 + _GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY._serialized_end=13247 + _STATUSREQUEST._serialized_start=13249 + _STATUSREQUEST._serialized_end=13310 + _STATUSRESPONSE._serialized_start=13312 + _STATUSRESPONSE._serialized_end=13353 + _STOPSTATUSREQUEST._serialized_start=13355 + _STOPSTATUSREQUEST._serialized_end=13420 + _STOPSTATUSRESPONSE._serialized_start=13422 + _STOPSTATUSRESPONSE._serialized_end=13467 + _NETWORKSTATUSREQUEST._serialized_start=13469 + _NETWORKSTATUSREQUEST._serialized_end=13537 + _NETWORKSTATUSRESPONSE._serialized_start=13539 + _NETWORKSTATUSRESPONSE._serialized_end=13619 + _HTTPRESPONSE._serialized_start=13621 + _HTTPRESPONSE._serialized_end=13689 + _INTERNALMESSAGESREQUEST._serialized_start=13691 + _INTERNALMESSAGESREQUEST._serialized_end=13762 + _INTERNALMESSAGESRESPONSE._serialized_start=13764 + _INTERNALMESSAGESRESPONSE._serialized_end=13842 + _INTERNALMESSAGES._serialized_start=13844 + _INTERNALMESSAGES._serialized_end=13879 + _POLLEXITREQUEST._serialized_start=13881 + _POLLEXITREQUEST._serialized_end=13944 + _POLLEXITRESPONSE._serialized_start=13947 + _POLLEXITRESPONSE._serialized_end=14192 + _OPERATIONSTATSREQUEST._serialized_start=14194 + _OPERATIONSTATSREQUEST._serialized_end=14263 + _OPERATIONSTATSRESPONSE._serialized_start=14265 + _OPERATIONSTATSRESPONSE._serialized_end=14346 + _OPERATIONSTATS._serialized_start=14348 + _OPERATIONSTATS._serialized_end=14437 + _OPERATION._serialized_start=14440 + _OPERATION._serialized_end=14575 + _SENDERMARKREQUEST._serialized_start=14577 + _SENDERMARKREQUEST._serialized_end=14596 + _SYNCFINISHREQUEST._serialized_start=14598 + _SYNCFINISHREQUEST._serialized_end=14617 + _SYNCRESPONSE._serialized_start=14619 + _SYNCRESPONSE._serialized_end=14688 + _SENDERREADREQUEST._serialized_start=14690 + _SENDERREADREQUEST._serialized_end=14753 + _STATUSREPORTREQUEST._serialized_start=14755 + _STATUSREPORTREQUEST._serialized_end=14864 + _SUMMARYRECORDREQUEST._serialized_start=14866 + _SUMMARYRECORDREQUEST._serialized_end=14936 + _TELEMETRYRECORDREQUEST._serialized_start=14938 + _TELEMETRYRECORDREQUEST._serialized_end=15014 + _SERVERINFOREQUEST._serialized_start=15016 + _SERVERINFOREQUEST._serialized_end=15081 + _SERVERINFORESPONSE._serialized_start=15083 + _SERVERINFORESPONSE._serialized_end=15207 + _SERVERMESSAGES._serialized_start=15209 + _SERVERMESSAGES._serialized_end=15270 + _SERVERMESSAGE._serialized_start=15272 + _SERVERMESSAGE._serialized_end=15373 + _FILECOUNTS._serialized_start=15375 + _FILECOUNTS._serialized_end=15474 + _FILEPUSHERSTATS._serialized_start=15476 + _FILEPUSHERSTATS._serialized_end=15561 + _FILESUPLOADED._serialized_start=15563 + _FILESUPLOADED._serialized_end=15593 + _FILETRANSFERINFOREQUEST._serialized_start=15596 + _FILETRANSFERINFOREQUEST._serialized_end=15840 + _FILETRANSFERINFOREQUEST_TRANSFERTYPE._serialized_start=15800 + _FILETRANSFERINFOREQUEST_TRANSFERTYPE._serialized_end=15840 + _LOCALINFO._serialized_start=15842 + _LOCALINFO._serialized_end=15891 + _SHUTDOWNREQUEST._serialized_start=15893 + _SHUTDOWNREQUEST._serialized_end=15956 + _SHUTDOWNRESPONSE._serialized_start=15958 + _SHUTDOWNRESPONSE._serialized_end=15976 + _ATTACHREQUEST._serialized_start=15978 + _ATTACHREQUEST._serialized_end=16058 + _ATTACHRESPONSE._serialized_start=16060 + _ATTACHRESPONSE._serialized_end=16158 + _TESTINJECTREQUEST._serialized_start=16161 + _TESTINJECTREQUEST._serialized_end=16502 + _TESTINJECTRESPONSE._serialized_start=16504 + _TESTINJECTRESPONSE._serialized_end=16524 + _HISTORYACTION._serialized_start=16526 + _HISTORYACTION._serialized_end=16556 + _PARTIALHISTORYREQUEST._serialized_start=16559 + _PARTIALHISTORYREQUEST._serialized_end=16761 + _PARTIALHISTORYRESPONSE._serialized_start=16763 + _PARTIALHISTORYRESPONSE._serialized_end=16787 + _SAMPLEDHISTORYREQUEST._serialized_start=16789 + _SAMPLEDHISTORYREQUEST._serialized_end=16858 + _SAMPLEDHISTORYITEM._serialized_start=16860 + _SAMPLEDHISTORYITEM._serialized_end=16955 + _SAMPLEDHISTORYRESPONSE._serialized_start=16957 + _SAMPLEDHISTORYRESPONSE._serialized_end=17031 + _RUNSTATUSREQUEST._serialized_start=17033 + _RUNSTATUSREQUEST._serialized_end=17097 + _RUNSTATUSRESPONSE._serialized_start=17099 + _RUNSTATUSRESPONSE._serialized_end=17219 + _RUNSTARTREQUEST._serialized_start=17221 + _RUNSTARTREQUEST._serialized_end=17324 + _RUNSTARTRESPONSE._serialized_start=17326 + _RUNSTARTRESPONSE._serialized_end=17344 + _RUNFINISHWITHOUTEXITREQUEST._serialized_start=17346 + _RUNFINISHWITHOUTEXITREQUEST._serialized_end=17421 + _RUNFINISHWITHOUTEXITRESPONSE._serialized_start=17423 + _RUNFINISHWITHOUTEXITRESPONSE._serialized_end=17453 + _CHECKVERSIONREQUEST._serialized_start=17455 + _CHECKVERSIONREQUEST._serialized_end=17547 + _CHECKVERSIONRESPONSE._serialized_start=17549 + _CHECKVERSIONRESPONSE._serialized_end=17642 + _JOBINFOREQUEST._serialized_start=17644 + _JOBINFOREQUEST._serialized_end=17706 + _JOBINFORESPONSE._serialized_start=17708 + _JOBINFORESPONSE._serialized_end=17762 + _LOGARTIFACTREQUEST._serialized_start=17765 + _LOGARTIFACTREQUEST._serialized_end=17924 + _LOGARTIFACTRESPONSE._serialized_start=17926 + _LOGARTIFACTRESPONSE._serialized_end=17991 + _DOWNLOADARTIFACTREQUEST._serialized_start=17994 + _DOWNLOADARTIFACTREQUEST._serialized_end=18184 + _DOWNLOADARTIFACTRESPONSE._serialized_start=18186 + _DOWNLOADARTIFACTRESPONSE._serialized_end=18235 + _KEEPALIVEREQUEST._serialized_start=18237 + _KEEPALIVEREQUEST._serialized_end=18301 + _KEEPALIVERESPONSE._serialized_start=18303 + _KEEPALIVERESPONSE._serialized_end=18322 + _ARTIFACTINFO._serialized_start=18324 + _ARTIFACTINFO._serialized_end=18437 + _GITINFO._serialized_start=18439 + _GITINFO._serialized_end=18480 + _GITSOURCE._serialized_start=18483 + _GITSOURCE._serialized_end=18618 + _IMAGESOURCE._serialized_start=18620 + _IMAGESOURCE._serialized_end=18648 + _SOURCE._serialized_start=18651 + _SOURCE._serialized_end=18791 + _JOBSOURCE._serialized_start=18793 + _JOBSOURCE._serialized_end=18900 + _PARTIALJOBARTIFACT._serialized_start=18902 + _PARTIALJOBARTIFACT._serialized_end=18988 + _USEARTIFACTRECORD._serialized_start=18991 + _USEARTIFACTRECORD._serialized_end=19148 + _USEARTIFACTRESULT._serialized_start=19150 + _USEARTIFACTRESULT._serialized_end=19169 + _CANCELREQUEST._serialized_start=19171 + _CANCELREQUEST._serialized_end=19253 + _CANCELRESPONSE._serialized_start=19255 + _CANCELRESPONSE._serialized_end=19271 + _DISKINFO._serialized_start=19273 + _DISKINFO._serialized_end=19312 + _MEMORYINFO._serialized_start=19314 + _MEMORYINFO._serialized_end=19341 + _CPUINFO._serialized_start=19343 + _CPUINFO._serialized_end=19390 + _APPLEINFO._serialized_start=19393 + _APPLEINFO._serialized_end=19547 + _GPUNVIDIAINFO._serialized_start=19549 + _GPUNVIDIAINFO._serialized_end=19656 + _GPUAMDINFO._serialized_start=19659 + _GPUAMDINFO._serialized_end=19924 + _TRAINIUMINFO._serialized_start=19926 + _TRAINIUMINFO._serialized_end=20036 + _TPUINFO._serialized_start=20038 + _TPUINFO._serialized_end=20119 + _COREWEAVEINFO._serialized_start=20121 + _COREWEAVEINFO._serialized_end=20190 + _ENVIRONMENTRECORD._serialized_start=20193 + _ENVIRONMENTRECORD._serialized_end=21385 + _ENVIRONMENTRECORD_DISKENTRY._serialized_start=21270 + _ENVIRONMENTRECORD_DISKENTRY._serialized_end=21339 + _ENVIRONMENTRECORD_SLURMENTRY._serialized_start=21341 + _ENVIRONMENTRECORD_SLURMENTRY._serialized_end=21385 + _PYTHONPACKAGESREQUEST._serialized_start=21388 + _PYTHONPACKAGESREQUEST._serialized_end=21529 + _PYTHONPACKAGESREQUEST_PYTHONPACKAGE._serialized_start=21483 + _PYTHONPACKAGESREQUEST_PYTHONPACKAGE._serialized_end=21529 + _JOBINPUTPATH._serialized_start=21531 + _JOBINPUTPATH._serialized_end=21559 + _JOBINPUTSOURCE._serialized_start=21562 + _JOBINPUTSOURCE._serialized_end=21776 + _JOBINPUTSOURCE_RUNCONFIGSOURCE._serialized_start=21715 + _JOBINPUTSOURCE_RUNCONFIGSOURCE._serialized_end=21732 + _JOBINPUTSOURCE_CONFIGFILESOURCE._serialized_start=21734 + _JOBINPUTSOURCE_CONFIGFILESOURCE._serialized_end=21766 + _JOBINPUTREQUEST._serialized_start=21779 + _JOBINPUTREQUEST._serialized_end=21978 + _SERVERFEATUREREQUEST._serialized_start=21980 + _SERVERFEATUREREQUEST._serialized_end=22096 + _SERVERFEATURERESPONSE._serialized_start=22098 + _SERVERFEATURERESPONSE._serialized_end=22173 + _SERVERFEATUREITEM._serialized_start=22175 + _SERVERFEATUREITEM._serialized_end=22225 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v4/wandb_server_pb2.py b/lib/python3.12/site-packages/wandb/proto/v4/wandb_server_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..9f6f08921222b9f876a6775640a7134b1771c470 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v4/wandb_server_pb2.py @@ -0,0 +1,67 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_server.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 +from wandb.proto import wandb_internal_pb2 as wandb_dot_proto_dot_wandb__internal__pb2 +from wandb.proto import wandb_settings_pb2 as wandb_dot_proto_dot_wandb__settings__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1ewandb/proto/wandb_server.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\x1a wandb/proto/wandb_internal.proto\x1a wandb/proto/wandb_settings.proto\"k\n\x19ServerAuthenticateRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12\x10\n\x08\x62\x61se_url\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"w\n\x1aServerAuthenticateResponse\x12\x16\n\x0e\x64\x65\x66\x61ult_entity\x18\x01 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"D\n\x15ServerShutdownRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x18\n\x16ServerShutdownResponse\"B\n\x13ServerStatusRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x16\n\x14ServerStatusResponse\"r\n\x17ServerInformInitRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1a\n\x18ServerInformInitResponse\"s\n\x18ServerInformStartRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1b\n\x19ServerInformStartResponse\"H\n\x19ServerInformFinishRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformFinishResponse\"H\n\x19ServerInformAttachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"u\n\x1aServerInformAttachResponse\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"H\n\x19ServerInformDetachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformDetachResponse\"]\n\x1bServerInformTeardownRequest\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1e\n\x1cServerInformTeardownResponse\"\xfb\x04\n\rServerRequest\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x30\n\x0erecord_publish\x18\x01 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12\x34\n\x12record_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12>\n\x0binform_init\x18\x03 \x01(\x0b\x32\'.wandb_internal.ServerInformInitRequestH\x00\x12\x42\n\rinform_finish\x18\x04 \x01(\x0b\x32).wandb_internal.ServerInformFinishRequestH\x00\x12\x42\n\rinform_attach\x18\x05 \x01(\x0b\x32).wandb_internal.ServerInformAttachRequestH\x00\x12\x42\n\rinform_detach\x18\x06 \x01(\x0b\x32).wandb_internal.ServerInformDetachRequestH\x00\x12\x46\n\x0finform_teardown\x18\x07 \x01(\x0b\x32+.wandb_internal.ServerInformTeardownRequestH\x00\x12@\n\x0cinform_start\x18\x08 \x01(\x0b\x32(.wandb_internal.ServerInformStartRequestH\x00\x12\x41\n\x0c\x61uthenticate\x18\t \x01(\x0b\x32).wandb_internal.ServerAuthenticateRequestH\x00\x42\x15\n\x13server_request_type\"\x91\x05\n\x0eServerResponse\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x34\n\x12result_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.ResultH\x00\x12H\n\x14inform_init_response\x18\x03 \x01(\x0b\x32(.wandb_internal.ServerInformInitResponseH\x00\x12L\n\x16inform_finish_response\x18\x04 \x01(\x0b\x32*.wandb_internal.ServerInformFinishResponseH\x00\x12L\n\x16inform_attach_response\x18\x05 \x01(\x0b\x32*.wandb_internal.ServerInformAttachResponseH\x00\x12L\n\x16inform_detach_response\x18\x06 \x01(\x0b\x32*.wandb_internal.ServerInformDetachResponseH\x00\x12P\n\x18inform_teardown_response\x18\x07 \x01(\x0b\x32,.wandb_internal.ServerInformTeardownResponseH\x00\x12J\n\x15inform_start_response\x18\x08 \x01(\x0b\x32).wandb_internal.ServerInformStartResponseH\x00\x12K\n\x15\x61uthenticate_response\x18\t \x01(\x0b\x32*.wandb_internal.ServerAuthenticateResponseH\x00\x42\x16\n\x14server_response_typeB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_server_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + _SERVERAUTHENTICATEREQUEST._serialized_start=148 + _SERVERAUTHENTICATEREQUEST._serialized_end=255 + _SERVERAUTHENTICATERESPONSE._serialized_start=257 + _SERVERAUTHENTICATERESPONSE._serialized_end=376 + _SERVERSHUTDOWNREQUEST._serialized_start=378 + _SERVERSHUTDOWNREQUEST._serialized_end=446 + _SERVERSHUTDOWNRESPONSE._serialized_start=448 + _SERVERSHUTDOWNRESPONSE._serialized_end=472 + _SERVERSTATUSREQUEST._serialized_start=474 + _SERVERSTATUSREQUEST._serialized_end=540 + _SERVERSTATUSRESPONSE._serialized_start=542 + _SERVERSTATUSRESPONSE._serialized_end=564 + _SERVERINFORMINITREQUEST._serialized_start=566 + _SERVERINFORMINITREQUEST._serialized_end=680 + _SERVERINFORMINITRESPONSE._serialized_start=682 + _SERVERINFORMINITRESPONSE._serialized_end=708 + _SERVERINFORMSTARTREQUEST._serialized_start=710 + _SERVERINFORMSTARTREQUEST._serialized_end=825 + _SERVERINFORMSTARTRESPONSE._serialized_start=827 + _SERVERINFORMSTARTRESPONSE._serialized_end=854 + _SERVERINFORMFINISHREQUEST._serialized_start=856 + _SERVERINFORMFINISHREQUEST._serialized_end=928 + _SERVERINFORMFINISHRESPONSE._serialized_start=930 + _SERVERINFORMFINISHRESPONSE._serialized_end=958 + _SERVERINFORMATTACHREQUEST._serialized_start=960 + _SERVERINFORMATTACHREQUEST._serialized_end=1032 + _SERVERINFORMATTACHRESPONSE._serialized_start=1034 + _SERVERINFORMATTACHRESPONSE._serialized_end=1151 + _SERVERINFORMDETACHREQUEST._serialized_start=1153 + _SERVERINFORMDETACHREQUEST._serialized_end=1225 + _SERVERINFORMDETACHRESPONSE._serialized_start=1227 + _SERVERINFORMDETACHRESPONSE._serialized_end=1255 + _SERVERINFORMTEARDOWNREQUEST._serialized_start=1257 + _SERVERINFORMTEARDOWNREQUEST._serialized_end=1350 + _SERVERINFORMTEARDOWNRESPONSE._serialized_start=1352 + _SERVERINFORMTEARDOWNRESPONSE._serialized_end=1382 + _SERVERREQUEST._serialized_start=1385 + _SERVERREQUEST._serialized_end=2020 + _SERVERRESPONSE._serialized_start=2023 + _SERVERRESPONSE._serialized_end=2680 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v4/wandb_settings_pb2.py b/lib/python3.12/site-packages/wandb/proto/v4/wandb_settings_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..9f9695b7d3908ddc6e4e7d1c2c5f18677463a305 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v4/wandb_settings_pb2.py @@ -0,0 +1,47 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_settings.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import wrappers_pb2 as google_dot_protobuf_dot_wrappers__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n wandb/proto/wandb_settings.proto\x12\x0ewandb_internal\x1a\x1egoogle/protobuf/wrappers.proto\" \n\x0fListStringValue\x12\r\n\x05value\x18\x01 \x03(\t\"\x1d\n\x0cListIntValue\x12\r\n\x05value\x18\x01 \x03(\x05\"\x8a\x01\n\x17MapStringKeyStringValue\x12\x41\n\x05value\x18\x01 \x03(\x0b\x32\x32.wandb_internal.MapStringKeyStringValue.ValueEntry\x1a,\n\nValueEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\xcb\x01\n#MapStringKeyMapStringKeyStringValue\x12M\n\x05value\x18\x01 \x03(\x0b\x32>.wandb_internal.MapStringKeyMapStringKeyStringValue.ValueEntry\x1aU\n\nValueEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x36\n\x05value\x18\x02 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue:\x02\x38\x01\"\x9a\x01\n\x12OpenMetricsFilters\x12\x33\n\x08sequence\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValueH\x00\x12\x46\n\x07mapping\x18\x02 \x01(\x0b\x32\x33.wandb_internal.MapStringKeyMapStringKeyStringValueH\x00\x42\x07\n\x05value\"7\n\tRunMoment\x12\x0b\n\x03run\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01\x12\x0e\n\x06metric\x18\x03 \x01(\t\"\x99N\n\x08Settings\x12-\n\x07\x61pi_key\x18\x37 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x13identity_token_file\x18\xaa\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10\x63redentials_file\x18\xab\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x39\n\x14insecure_disable_ssl\x18\xb9\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08_offline\x18\x1e \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06x_sync\x18\x1f \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\tsync_file\x18\x86\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x07_shared\x18\xa2\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x06run_id\x18k \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07run_url\x18q \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07project\x18\x61 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06\x65ntity\x18\x45 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\x0corganization\x18\xbc\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0cx_start_time\x18) \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12.\n\x08root_dir\x18i \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07log_dir\x18U \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0clog_internal\x18V \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\tfiles_dir\x18\x46 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0bx_files_dir\x18\xb4\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0cignore_globs\x18N \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12.\n\x08\x62\x61se_url\x18\x39 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12=\n\x17x_file_stream_max_bytes\x18\xac\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x46\n\x1fx_file_stream_transmit_interval\x18\xaf\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x45\n\x14x_extra_http_headers\x18\x0e \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12=\n\x17x_file_stream_retry_max\x18\x93\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12K\n$x_file_stream_retry_wait_min_seconds\x18\x94\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12K\n$x_file_stream_retry_wait_max_seconds\x18\x95\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x43\n\x1dx_file_stream_timeout_seconds\x18\x0f \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x42\n\x1cx_file_stream_max_line_bytes\x18\xb2\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12?\n\x19x_file_transfer_retry_max\x18\x96\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12M\n&x_file_transfer_retry_wait_min_seconds\x18\x97\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12M\n&x_file_transfer_retry_wait_max_seconds\x18\x98\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x46\n\x1fx_file_transfer_timeout_seconds\x18\x99\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x39\n\x13x_graphql_retry_max\x18\x9a\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12G\n x_graphql_retry_wait_min_seconds\x18\x9b\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12G\n x_graphql_retry_wait_max_seconds\x18\x9c\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12@\n\x19x_graphql_timeout_seconds\x18\x9d\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x31\n\nhttp_proxy\x18\xa8\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0bhttps_proxy\x18\xa9\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12;\n\tx_proxies\x18\xc8\x01 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12-\n\x07program\x18_ \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0fprogram_relpath\x18` \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10_code_path_local\x18\xa3\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x0fprogram_abspath\x18\x9f\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x05_args\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12)\n\x03_os\x18 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06\x64ocker\x18\x43 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0cx_executable\x18\r \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07_python\x18\" \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\tcolab_url\x18\xa0\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x04host\x18M \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x08username\x18\x8d\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x05\x65mail\x18\x44 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06resume\x18\x66 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0bresume_from\x18\xa7\x01 \x01(\x0b\x32\x19.wandb_internal.RunMoment\x12-\n\tfork_from\x18\xa4\x01 \x01(\x0b\x32\x19.wandb_internal.RunMoment\x12\x38\n\x14\x64isable_job_creation\x18\x41 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\tsweep_url\x18\x83\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12;\n\x16x_disable_update_check\x18\xa5\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0ex_disable_meta\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\tsave_code\x18s \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x0b\x64isable_git\x18? \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x16x_disable_machine_info\x18\x9e\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0fx_disable_stats\x18\n \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x13x_stats_buffer_size\x18\xa1\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12@\n\x19x_stats_sampling_interval\x18\xae\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x30\n\x0bx_stats_pid\x18* \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12<\n\x12x_stats_disk_paths\x18\x92\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12H\n\"x_stats_neuron_monitor_config_path\x18. \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12<\n\x15x_stats_dcgm_exporter\x18\xbb\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12O\n\x1ex_stats_open_metrics_endpoints\x18/ \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12H\n\x1cx_stats_open_metrics_filters\x18\x30 \x01(\x0b\x32\".wandb_internal.OpenMetricsFilters\x12S\n!x_stats_open_metrics_http_headers\x18\xb8\x01 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12=\n\x16x_stats_gpu_device_ids\x18\xba\x01 \x01(\x0b\x32\x1c.wandb_internal.ListIntValue\x12\x37\n\x11x_stats_cpu_count\x18\xc2\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12?\n\x19x_stats_cpu_logical_count\x18\xc3\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x37\n\x11x_stats_gpu_count\x18\xc4\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x37\n\x10x_stats_gpu_type\x18\xc5\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12?\n\x1ax_stats_track_process_tree\x18\xc6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12.\n\x07x_label\x18\xb5\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\tx_primary\x18\xb6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12:\n\x15x_update_finish_state\x18\xb7\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12<\n\x17\x61llow_offline_artifacts\x18\xb1\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\x07\x63onsole\x18< \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x11\x63onsole_multipart\x18\xa6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x35\n\x10sync_tensorboard\x18\xb3\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x42\n\x1dx_server_side_derived_summary\x18\xbd\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x46\n!x_server_side_expand_glob_metrics\x18\xbe\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x16x_skip_transaction_log\x18\xbf\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12J\n#x_stats_coreweave_metadata_base_url\x18\xc0\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12J\n#x_stats_coreweave_metadata_endpoint\x18\xc1\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0b_aws_lambda\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0fx_cli_only_mode\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06_colab\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10x_disable_viewer\x18\x0b \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x15x_flow_control_custom\x18\x10 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x17x_flow_control_disabled\x18\x11 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12>\n\x18x_internal_check_process\x18\x12 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12,\n\x08_ipython\x18\x14 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08_jupyter\x18\x15 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x0ex_jupyter_root\x18\x16 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x07_kaggle\x18\x17 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12=\n\x18x_live_policy_rate_limit\x18\x18 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12<\n\x17x_live_policy_wait_time\x18\x19 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x30\n\x0bx_log_level\x18\x1a \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x35\n\x10x_network_buffer\x18\x1b \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12)\n\x05_noop\x18\x1c \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\t_notebook\x18\x1d \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\t_platform\x18! \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x38\n\x12x_runqueue_item_id\x18# \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x13x_save_requirements\x18% \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x13x_service_transport\x18& \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x34\n\x0ex_service_wait\x18\' \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x35\n\x0f_start_datetime\x18( \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\r_tmp_code_dir\x18\x31 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x08_windows\x18\x34 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10\x61llow_val_change\x18\x35 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\tanonymous\x18\x36 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12P\n\x1f\x61zure_account_url_to_access_key\x18\x38 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12.\n\x08\x63ode_dir\x18: \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0c\x63onfig_paths\x18; \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12\x30\n\ndeployment\x18= \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\x0c\x64isable_code\x18> \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x31\n\rdisable_hints\x18@ \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08\x64isabled\x18\x42 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12)\n\x05\x66orce\x18G \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\ngit_commit\x18H \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\ngit_remote\x18I \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x34\n\x0egit_remote_url\x18J \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08git_root\x18K \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x11heartbeat_seconds\x18L \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x32\n\x0cinit_timeout\x18O \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12,\n\x08is_local\x18P \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\njob_source\x18Q \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\rlabel_disable\x18R \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06launch\x18S \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x38\n\x12launch_config_path\x18T \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x14log_symlink_internal\x18W \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x10log_symlink_user\x18X \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08log_user\x18Y \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\rlogin_timeout\x18Z \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12*\n\x04mode\x18\\ \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\rnotebook_name\x18] \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\x0bproject_url\x18\x62 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12)\n\x05quiet\x18\x63 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12+\n\x07relogin\x18\x65 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0cresume_fname\x18g \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x07resumed\x18h \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\trun_group\x18j \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0crun_job_type\x18l \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08run_mode\x18m \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08run_name\x18n \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\trun_notes\x18o \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\x08run_tags\x18p \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12\x35\n\x11sagemaker_disable\x18r \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x35\n\x0fsettings_system\x18t \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x38\n\x12settings_workspace\x18u \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0bshow_colors\x18v \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12.\n\nshow_emoji\x18w \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x0bshow_errors\x18x \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\tshow_info\x18y \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x31\n\rshow_warnings\x18z \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06silent\x18{ \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0cstart_method\x18| \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x06strict\x18} \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0esummary_errors\x18~ \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x34\n\x0fsummary_timeout\x18\x7f \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x36\n\x10summary_warnings\x18\x80\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12/\n\x08sweep_id\x18\x81\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10sweep_param_path\x18\x82\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x07symlink\x18\x84\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x08sync_dir\x18\x85\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x13sync_symlink_latest\x18\x87\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12J\n%table_raise_on_max_row_limit_exceeded\x18\x8a\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x08timespec\x18\x8b\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x07tmp_dir\x18\x8c\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\twandb_dir\x18\x8e\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0ex_jupyter_name\x18\x8f\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0ex_jupyter_path\x18\x90\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x08job_name\x18\x91\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValueJ\x04\x08\x03\x10\x04J\x04\x08\x06\x10\x07J\x04\x08\x08\x10\tJ\x04\x08\t\x10\nJ\x04\x08\x0c\x10\rJ\x04\x08\x13\x10\x14J\x04\x08$\x10%J\x04\x08+\x10,J\x04\x08,\x10-J\x04\x08-\x10.J\x04\x08\x32\x10\x33J\x04\x08\x33\x10\x34J\x04\x08[\x10\\J\x04\x08^\x10_J\x04\x08\x64\x10\x65J\x06\x08\x88\x01\x10\x89\x01J\x06\x08\x89\x01\x10\x8a\x01J\x06\x08\xad\x01\x10\xae\x01J\x06\x08\xb0\x01\x10\xb1\x01\x42\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_settings_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY._options = None + _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_options = b'8\001' + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY._options = None + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_options = b'8\001' + _LISTSTRINGVALUE._serialized_start=84 + _LISTSTRINGVALUE._serialized_end=116 + _LISTINTVALUE._serialized_start=118 + _LISTINTVALUE._serialized_end=147 + _MAPSTRINGKEYSTRINGVALUE._serialized_start=150 + _MAPSTRINGKEYSTRINGVALUE._serialized_end=288 + _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_start=244 + _MAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_end=288 + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE._serialized_start=291 + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE._serialized_end=494 + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_start=409 + _MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY._serialized_end=494 + _OPENMETRICSFILTERS._serialized_start=497 + _OPENMETRICSFILTERS._serialized_end=651 + _RUNMOMENT._serialized_start=653 + _RUNMOMENT._serialized_end=708 + _SETTINGS._serialized_start=711 + _SETTINGS._serialized_end=10720 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v4/wandb_telemetry_pb2.py b/lib/python3.12/site-packages/wandb/proto/v4/wandb_telemetry_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..c4e0ac498ac2d8793e92383e0973bd1eee94666f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v4/wandb_telemetry_pb2.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_telemetry.proto +"""Generated protocol buffer code.""" +from google.protobuf.internal import builder as _builder +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n!wandb/proto/wandb_telemetry.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\"\xdb\x03\n\x0fTelemetryRecord\x12-\n\x0cimports_init\x18\x01 \x01(\x0b\x32\x17.wandb_internal.Imports\x12/\n\x0eimports_finish\x18\x02 \x01(\x0b\x32\x17.wandb_internal.Imports\x12(\n\x07\x66\x65\x61ture\x18\x03 \x01(\x0b\x32\x17.wandb_internal.Feature\x12\x16\n\x0epython_version\x18\x04 \x01(\t\x12\x13\n\x0b\x63li_version\x18\x05 \x01(\t\x12\x1b\n\x13huggingface_version\x18\x06 \x01(\t\x12 \n\x03\x65nv\x18\x08 \x01(\x0b\x32\x13.wandb_internal.Env\x12%\n\x05label\x18\t \x01(\x0b\x32\x16.wandb_internal.Labels\x12.\n\ndeprecated\x18\n \x01(\x0b\x32\x1a.wandb_internal.Deprecated\x12&\n\x06issues\x18\x0b \x01(\x0b\x32\x16.wandb_internal.Issues\x12\x14\n\x0c\x63ore_version\x18\x0c \x01(\t\x12\x10\n\x08platform\x18\r \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x11\n\x0fTelemetryResult\"\x9a\x0e\n\x07Imports\x12\r\n\x05torch\x18\x01 \x01(\x08\x12\r\n\x05keras\x18\x02 \x01(\x08\x12\x12\n\ntensorflow\x18\x03 \x01(\x08\x12\x0e\n\x06\x66\x61stai\x18\x04 \x01(\x08\x12\x0f\n\x07sklearn\x18\x05 \x01(\x08\x12\x0f\n\x07xgboost\x18\x06 \x01(\x08\x12\x10\n\x08\x63\x61tboost\x18\x07 \x01(\x08\x12\x10\n\x08lightgbm\x18\x08 \x01(\x08\x12\x19\n\x11pytorch_lightning\x18\t \x01(\x08\x12\x0e\n\x06ignite\x18\n \x01(\x08\x12\x14\n\x0ctransformers\x18\x0b \x01(\x08\x12\x0b\n\x03jax\x18\x0c \x01(\x08\x12\x10\n\x08metaflow\x18\r \x01(\x08\x12\x10\n\x08\x61llennlp\x18\x0e \x01(\x08\x12\x11\n\tautogluon\x18\x0f \x01(\x08\x12\x11\n\tautokeras\x18\x10 \x01(\x08\x12\x10\n\x08\x63\x61talyst\x18\x12 \x01(\x08\x12\x10\n\x08\x64\x65\x65pchem\x18\x15 \x01(\x08\x12\x0f\n\x07\x64\x65\x65pctr\x18\x16 \x01(\x08\x12\x0f\n\x07pycaret\x18\x1c \x01(\x08\x12\x14\n\x0cpytorchvideo\x18\x1d \x01(\x08\x12\x0b\n\x03ray\x18\x1e \x01(\x08\x12\x1a\n\x12simpletransformers\x18\x1f \x01(\x08\x12\x0e\n\x06skorch\x18 \x01(\x08\x12\r\n\x05spacy\x18! \x01(\x08\x12\r\n\x05\x66lash\x18\" \x01(\x08\x12\x0e\n\x06optuna\x18# \x01(\x08\x12\x0f\n\x07recbole\x18$ \x01(\x08\x12\x0c\n\x04mmcv\x18% \x01(\x08\x12\r\n\x05mmdet\x18& \x01(\x08\x12\x11\n\ttorchdrug\x18\' \x01(\x08\x12\x11\n\ttorchtext\x18( \x01(\x08\x12\x13\n\x0btorchvision\x18) \x01(\x08\x12\r\n\x05\x65legy\x18* \x01(\x08\x12\x12\n\ndetectron2\x18+ \x01(\x08\x12\r\n\x05\x66lair\x18, \x01(\x08\x12\x0c\n\x04\x66lax\x18- \x01(\x08\x12\x0c\n\x04syft\x18. \x01(\x08\x12\x0b\n\x03TTS\x18/ \x01(\x08\x12\r\n\x05monai\x18\x30 \x01(\x08\x12\x17\n\x0fhuggingface_hub\x18\x31 \x01(\x08\x12\r\n\x05hydra\x18\x32 \x01(\x08\x12\x10\n\x08\x64\x61tasets\x18\x33 \x01(\x08\x12\x0e\n\x06sacred\x18\x34 \x01(\x08\x12\x0e\n\x06joblib\x18\x35 \x01(\x08\x12\x0c\n\x04\x64\x61sk\x18\x36 \x01(\x08\x12\x11\n\tpaddleocr\x18\x38 \x01(\x08\x12\r\n\x05ppdet\x18\x39 \x01(\x08\x12\x11\n\tpaddleseg\x18: \x01(\x08\x12\x11\n\tpaddlenlp\x18; \x01(\x08\x12\r\n\x05mmseg\x18< \x01(\x08\x12\r\n\x05mmocr\x18= \x01(\x08\x12\r\n\x05mmcls\x18> \x01(\x08\x12\x0c\n\x04timm\x18? \x01(\x08\x12\x0f\n\x07\x66\x61irseq\x18@ \x01(\x08\x12\x12\n\ndeepchecks\x18\x41 \x01(\x08\x12\x10\n\x08\x63omposer\x18\x42 \x01(\x08\x12\x10\n\x08sparseml\x18\x43 \x01(\x08\x12\x10\n\x08\x61nomalib\x18\x44 \x01(\x08\x12\r\n\x05zenml\x18\x45 \x01(\x08\x12\x12\n\ncolossalai\x18\x46 \x01(\x08\x12\x12\n\naccelerate\x18G \x01(\x08\x12\x0e\n\x06merlin\x18H \x01(\x08\x12\x0f\n\x07nanodet\x18I \x01(\x08\x12#\n\x1bsegmentation_models_pytorch\x18J \x01(\x08\x12\x1d\n\x15sentence_transformers\x18K \x01(\x08\x12\x0b\n\x03\x64gl\x18L \x01(\x08\x12\x17\n\x0ftorch_geometric\x18M \x01(\x08\x12\x0c\n\x04jina\x18N \x01(\x08\x12\x0e\n\x06kornia\x18O \x01(\x08\x12\x16\n\x0e\x61lbumentations\x18P \x01(\x08\x12\x10\n\x08keras_cv\x18Q \x01(\x08\x12\x10\n\x08mmengine\x18R \x01(\x08\x12\x11\n\tdiffusers\x18S \x01(\x08\x12\x0b\n\x03trl\x18T \x01(\x08\x12\x0c\n\x04trlx\x18U \x01(\x08\x12\x11\n\tlangchain\x18V \x01(\x08\x12\x13\n\x0bllama_index\x18W \x01(\x08\x12\x15\n\rstability_sdk\x18X \x01(\x08\x12\x0f\n\x07prefect\x18Y \x01(\x08\x12\x13\n\x0bprefect_ray\x18Z \x01(\x08\x12\x10\n\x08pinecone\x18[ \x01(\x08\x12\x10\n\x08\x63hromadb\x18\\ \x01(\x08\x12\x10\n\x08weaviate\x18] \x01(\x08\x12\x13\n\x0bpromptlayer\x18^ \x01(\x08\x12\x0e\n\x06openai\x18_ \x01(\x08\x12\x0e\n\x06\x63ohere\x18` \x01(\x08\x12\x11\n\tanthropic\x18\x61 \x01(\x08\x12\x0c\n\x04peft\x18\x62 \x01(\x08\x12\x0f\n\x07optimum\x18\x63 \x01(\x08\x12\x10\n\x08\x65valuate\x18\x64 \x01(\x08\x12\x10\n\x08langflow\x18\x65 \x01(\x08\x12\x12\n\nkeras_core\x18\x66 \x01(\x08\x12\x18\n\x10lightning_fabric\x18g \x01(\x08\x12\x1c\n\x14\x63urated_transformers\x18h \x01(\x08\x12\x0e\n\x06orjson\x18i \x01(\x08\x12\x11\n\tlightning\x18j \x01(\x08J\x04\x08\x11\x10\x12J\x04\x08\x13\x10\x14J\x04\x08\x14\x10\x15J\x04\x08\x17\x10\x18J\x04\x08\x18\x10\x19J\x04\x08\x19\x10\x1aJ\x04\x08\x1a\x10\x1bJ\x04\x08\x1b\x10\x1cJ\x04\x08\x37\x10\x38\"\x95\r\n\x07\x46\x65\x61ture\x12\r\n\x05watch\x18\x01 \x01(\x08\x12\x0e\n\x06\x66inish\x18\x02 \x01(\x08\x12\x0c\n\x04save\x18\x03 \x01(\x08\x12\x0f\n\x07offline\x18\x04 \x01(\x08\x12\x0f\n\x07resumed\x18\x05 \x01(\x08\x12\x0c\n\x04grpc\x18\x06 \x01(\x08\x12\x0e\n\x06metric\x18\x07 \x01(\x08\x12\r\n\x05keras\x18\x08 \x01(\x08\x12\x11\n\tsagemaker\x18\t \x01(\x08\x12\x1c\n\x14\x61rtifact_incremental\x18\n \x01(\x08\x12\x10\n\x08metaflow\x18\x0b \x01(\x08\x12\x0f\n\x07prodigy\x18\x0c \x01(\x08\x12\x15\n\rset_init_name\x18\r \x01(\x08\x12\x13\n\x0bset_init_id\x18\x0e \x01(\x08\x12\x15\n\rset_init_tags\x18\x0f \x01(\x08\x12\x17\n\x0fset_init_config\x18\x10 \x01(\x08\x12\x14\n\x0cset_run_name\x18\x11 \x01(\x08\x12\x14\n\x0cset_run_tags\x18\x12 \x01(\x08\x12\x17\n\x0fset_config_item\x18\x13 \x01(\x08\x12\x0e\n\x06launch\x18\x14 \x01(\x08\x12\x1c\n\x14torch_profiler_trace\x18\x15 \x01(\x08\x12\x0b\n\x03sb3\x18\x16 \x01(\x08\x12\x17\n\x0finit_return_run\x18\x18 \x01(\x08\x12\x1f\n\x17lightgbm_wandb_callback\x18\x19 \x01(\x08\x12\x1c\n\x14lightgbm_log_summary\x18\x1a \x01(\x08\x12\x1f\n\x17\x63\x61tboost_wandb_callback\x18\x1b \x01(\x08\x12\x1c\n\x14\x63\x61tboost_log_summary\x18\x1c \x01(\x08\x12\x17\n\x0ftensorboard_log\x18\x1d \x01(\x08\x12\x16\n\x0e\x65stimator_hook\x18\x1e \x01(\x08\x12\x1e\n\x16xgboost_wandb_callback\x18\x1f \x01(\x08\x12\"\n\x1axgboost_old_wandb_callback\x18 \x01(\x08\x12\x0e\n\x06\x61ttach\x18! \x01(\x08\x12\x19\n\x11tensorboard_patch\x18\" \x01(\x08\x12\x18\n\x10tensorboard_sync\x18# \x01(\x08\x12\x15\n\rkfp_wandb_log\x18$ \x01(\x08\x12\x1b\n\x13maybe_run_overwrite\x18% \x01(\x08\x12\x1c\n\x14keras_metrics_logger\x18& \x01(\x08\x12\x1e\n\x16keras_model_checkpoint\x18\' \x01(\x08\x12!\n\x19keras_wandb_eval_callback\x18( \x01(\x08\x12\x1d\n\x15\x66low_control_overflow\x18) \x01(\x08\x12\x0c\n\x04sync\x18* \x01(\x08\x12\x1d\n\x15\x66low_control_disabled\x18+ \x01(\x08\x12\x1b\n\x13\x66low_control_custom\x18, \x01(\x08\x12\x14\n\x0copen_metrics\x18. \x01(\x08\x12\x1a\n\x12ultralytics_yolov8\x18/ \x01(\x08\x12\x17\n\x0fimporter_mlflow\x18\x30 \x01(\x08\x12\x15\n\rsync_tfevents\x18\x31 \x01(\x08\x12\x15\n\rasync_uploads\x18\x32 \x01(\x08\x12\x16\n\x0eopenai_autolog\x18\x33 \x01(\x08\x12\x18\n\x10langchain_tracer\x18\x34 \x01(\x08\x12\x16\n\x0e\x63ohere_autolog\x18\x35 \x01(\x08\x12\x1b\n\x13hf_pipeline_autolog\x18\x36 \x01(\x08\x12\r\n\x05lib_c\x18\x38 \x01(\x08\x12\x0f\n\x07lib_cpp\x18\x39 \x01(\x08\x12\x19\n\x11openai_finetuning\x18: \x01(\x08\x12\x19\n\x11\x64iffusers_autolog\x18; \x01(\x08\x12\x1f\n\x17lightning_fabric_logger\x18< \x01(\x08\x12\x14\n\x0cset_step_log\x18= \x01(\x08\x12\x13\n\x0bset_summary\x18> \x01(\x08\x12\x16\n\x0emetric_summary\x18? \x01(\x08\x12\x13\n\x0bmetric_goal\x18@ \x01(\x08\x12\x15\n\rmetric_hidden\x18\x41 \x01(\x08\x12\x18\n\x10metric_step_sync\x18\x42 \x01(\x08\x12\x13\n\x0bshared_mode\x18\x43 \x01(\x08\x12#\n\x1bserver_side_derived_summary\x18\x44 \x01(\x08\x12\x1b\n\x13user_provided_label\x18\x45 \x01(\x08\x12\x1e\n\x16\x64\x63gm_profiling_enabled\x18\x46 \x01(\x08\x12\x11\n\tfork_mode\x18G \x01(\x08\x12\x13\n\x0brewind_mode\x18H \x01(\x08J\x04\x08\x17\x10\x18J\x04\x08-\x10.J\x04\x08\x37\x10\x38\"\xc9\x01\n\x03\x45nv\x12\x0f\n\x07jupyter\x18\x01 \x01(\x08\x12\x0e\n\x06kaggle\x18\x02 \x01(\x08\x12\x0f\n\x07windows\x18\x03 \x01(\x08\x12\x0e\n\x06m1_gpu\x18\x04 \x01(\x08\x12\x10\n\x08trainium\x18\n \x01(\x08\x12\x0b\n\x03pex\x18\x0b \x01(\x08\x12\r\n\x05\x63olab\x18\x0c \x01(\x08\x12\x0f\n\x07ipython\x18\r \x01(\x08\x12\x12\n\naws_lambda\x18\x0e \x01(\x08\x12\x0f\n\x07\x61md_gpu\x18\x0f \x01(\x08J\x04\x08\x05\x10\x06J\x04\x08\x06\x10\x07J\x04\x08\x07\x10\x08J\x04\x08\x08\x10\tJ\x04\x08\t\x10\n\"H\n\x06Labels\x12\x13\n\x0b\x63ode_string\x18\x01 \x01(\t\x12\x13\n\x0brepo_string\x18\x02 \x01(\t\x12\x14\n\x0c\x63ode_version\x18\x03 \x01(\t\"\xd5\x05\n\nDeprecated\x12!\n\x19keras_callback__data_type\x18\x01 \x01(\x08\x12\r\n\x05plots\x18\x05 \x01(\x08\x12!\n\x19init__config_include_keys\x18\x07 \x01(\x08\x12!\n\x19init__config_exclude_keys\x18\x08 \x01(\x08\x12\"\n\x1akeras_callback__save_model\x18\t \x01(\x08\x12\x18\n\x10langchain_tracer\x18\n \x01(\x08\x12\x1a\n\x12\x61rtifact__get_path\x18\x0b \x01(\x08\x12#\n\x1b\x61rtifactmanifestentry__name\x18\x0c \x01(\x08\x12\x1e\n\x16\x61pi__artifact_versions\x18\r \x01(\x08\x12(\n artifact_collection__change_type\x18\x0e \x01(\x08\x12\x1f\n\x17run__define_metric_copy\x18\x0f \x01(\x08\x12\x14\n\x0crun_disabled\x18\x10 \x01(\x08\x12\x16\n\x0ekeras_callback\x18\x11 \x01(\x08\x12$\n\x1crun__define_metric_best_goal\x18\x12 \x01(\x08\x12\x19\n\x11run__finish_quiet\x18\x13 \x01(\x08\x12\x18\n\x10run__reinit_bool\x18\x14 \x01(\x08\x12\x14\n\x0crun__get_url\x18\x15 \x01(\x08\x12\x19\n\x11run__project_name\x18\x16 \x01(\x08\x12\x1c\n\x14run__get_project_url\x18\x17 \x01(\x08\x12\x1a\n\x12run__get_sweep_url\x18\x18 \x01(\x08\x12 \n\x18run__use_artifact_use_as\x18\x19 \x01(\x08\x12\x18\n\x10\x61rtifact__use_as\x18\x1a \x01(\x08\x12\x1d\n\x15\x61rtifact__init_use_as\x18\x1b \x01(\x08J\x04\x08\x02\x10\x03J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05J\x04\x08\x06\x10\x07\"|\n\x06Issues\x12%\n\x1dsettings__validation_warnings\x18\x01 \x01(\x08\x12!\n\x19settings__unexpected_args\x18\x02 \x01(\x08\x12(\n settings__preprocessing_warnings\x18\x03 \x01(\x08\x42\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_telemetry_pb2', globals()) +if _descriptor._USE_C_DESCRIPTORS == False: + + DESCRIPTOR._options = None + DESCRIPTOR._serialized_options = b'Z\031core/pkg/service_go_proto' + _TELEMETRYRECORD._serialized_start=84 + _TELEMETRYRECORD._serialized_end=559 + _TELEMETRYRESULT._serialized_start=561 + _TELEMETRYRESULT._serialized_end=578 + _IMPORTS._serialized_start=581 + _IMPORTS._serialized_end=2399 + _FEATURE._serialized_start=2402 + _FEATURE._serialized_end=4087 + _ENV._serialized_start=4090 + _ENV._serialized_end=4291 + _LABELS._serialized_start=4293 + _LABELS._serialized_end=4365 + _DEPRECATED._serialized_start=4368 + _DEPRECATED._serialized_end=5093 + _ISSUES._serialized_start=5095 + _ISSUES._serialized_end=5219 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v5/wandb_base_pb2.py b/lib/python3.12/site-packages/wandb/proto/v5/wandb_base_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..d02477cff9f03584418a3a316e8cae7db7048f91 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v5/wandb_base_pb2.py @@ -0,0 +1,31 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_base.proto +# Protobuf Python Version: 5.26.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cwandb/proto/wandb_base.proto\x12\x0ewandb_internal\"6\n\x0b_RecordInfo\x12\x11\n\tstream_id\x18\x01 \x01(\t\x12\x14\n\x0c_tracelog_id\x18\x64 \x01(\t\"!\n\x0c_RequestInfo\x12\x11\n\tstream_id\x18\x01 \x01(\t\"#\n\x0b_ResultInfo\x12\x14\n\x0c_tracelog_id\x18\x64 \x01(\tB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_base_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['__RECORDINFO']._serialized_start=48 + _globals['__RECORDINFO']._serialized_end=102 + _globals['__REQUESTINFO']._serialized_start=104 + _globals['__REQUESTINFO']._serialized_end=137 + _globals['__RESULTINFO']._serialized_start=139 + _globals['__RESULTINFO']._serialized_end=174 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v5/wandb_internal_pb2.py b/lib/python3.12/site-packages/wandb/proto/v5/wandb_internal_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..540aade4d40ef73a074d817029df6588936ec0a9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v5/wandb_internal_pb2.py @@ -0,0 +1,383 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_internal.proto +# Protobuf Python Version: 5.26.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 +from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 +from wandb.proto import wandb_telemetry_pb2 as wandb_dot_proto_dot_wandb__telemetry__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n wandb/proto/wandb_internal.proto\x12\x0ewandb_internal\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1cwandb/proto/wandb_base.proto\x1a!wandb/proto/wandb_telemetry.proto\"\xcf\t\n\x06Record\x12\x0b\n\x03num\x18\x01 \x01(\x03\x12\x30\n\x07history\x18\x02 \x01(\x0b\x32\x1d.wandb_internal.HistoryRecordH\x00\x12\x30\n\x07summary\x18\x03 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecordH\x00\x12.\n\x06output\x18\x04 \x01(\x0b\x32\x1c.wandb_internal.OutputRecordH\x00\x12.\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x1c.wandb_internal.ConfigRecordH\x00\x12,\n\x05\x66iles\x18\x06 \x01(\x0b\x32\x1b.wandb_internal.FilesRecordH\x00\x12,\n\x05stats\x18\x07 \x01(\x0b\x32\x1b.wandb_internal.StatsRecordH\x00\x12\x32\n\x08\x61rtifact\x18\x08 \x01(\x0b\x32\x1e.wandb_internal.ArtifactRecordH\x00\x12,\n\x08tbrecord\x18\t \x01(\x0b\x32\x18.wandb_internal.TBRecordH\x00\x12,\n\x05\x61lert\x18\n \x01(\x0b\x32\x1b.wandb_internal.AlertRecordH\x00\x12\x34\n\ttelemetry\x18\x0b \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecordH\x00\x12.\n\x06metric\x18\x0c \x01(\x0b\x32\x1c.wandb_internal.MetricRecordH\x00\x12\x35\n\noutput_raw\x18\r \x01(\x0b\x32\x1f.wandb_internal.OutputRawRecordH\x00\x12(\n\x03run\x18\x11 \x01(\x0b\x32\x19.wandb_internal.RunRecordH\x00\x12-\n\x04\x65xit\x18\x12 \x01(\x0b\x32\x1d.wandb_internal.RunExitRecordH\x00\x12,\n\x05\x66inal\x18\x14 \x01(\x0b\x32\x1b.wandb_internal.FinalRecordH\x00\x12.\n\x06header\x18\x15 \x01(\x0b\x32\x1c.wandb_internal.HeaderRecordH\x00\x12.\n\x06\x66ooter\x18\x16 \x01(\x0b\x32\x1c.wandb_internal.FooterRecordH\x00\x12\x39\n\npreempting\x18\x17 \x01(\x0b\x32#.wandb_internal.RunPreemptingRecordH\x00\x12\x34\n\x12noop_link_artifact\x18\x18 \x01(\x0b\x32\x16.google.protobuf.EmptyH\x00\x12\x39\n\x0cuse_artifact\x18\x19 \x01(\x0b\x32!.wandb_internal.UseArtifactRecordH\x00\x12\x38\n\x0b\x65nvironment\x18\x1a \x01(\x0b\x32!.wandb_internal.EnvironmentRecordH\x00\x12*\n\x07request\x18\x64 \x01(\x0b\x32\x17.wandb_internal.RequestH\x00\x12(\n\x07\x63ontrol\x18\x10 \x01(\x0b\x32\x17.wandb_internal.Control\x12\x0c\n\x04uuid\x18\x13 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfoB\r\n\x0brecord_type\"\xa8\x01\n\x07\x43ontrol\x12\x10\n\x08req_resp\x18\x01 \x01(\x08\x12\r\n\x05local\x18\x02 \x01(\x08\x12\x10\n\x08relay_id\x18\x03 \x01(\t\x12\x14\n\x0cmailbox_slot\x18\x04 \x01(\t\x12\x13\n\x0b\x61lways_send\x18\x05 \x01(\x08\x12\x14\n\x0c\x66low_control\x18\x06 \x01(\x08\x12\x12\n\nend_offset\x18\x07 \x01(\x03\x12\x15\n\rconnection_id\x18\x08 \x01(\t\"\xf3\x03\n\x06Result\x12\x35\n\nrun_result\x18\x11 \x01(\x0b\x32\x1f.wandb_internal.RunUpdateResultH\x00\x12\x34\n\x0b\x65xit_result\x18\x12 \x01(\x0b\x32\x1d.wandb_internal.RunExitResultH\x00\x12\x33\n\nlog_result\x18\x14 \x01(\x0b\x32\x1d.wandb_internal.HistoryResultH\x00\x12\x37\n\x0esummary_result\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.SummaryResultH\x00\x12\x35\n\routput_result\x18\x16 \x01(\x0b\x32\x1c.wandb_internal.OutputResultH\x00\x12\x35\n\rconfig_result\x18\x17 \x01(\x0b\x32\x1c.wandb_internal.ConfigResultH\x00\x12,\n\x08response\x18\x64 \x01(\x0b\x32\x18.wandb_internal.ResponseH\x00\x12(\n\x07\x63ontrol\x18\x10 \x01(\x0b\x32\x17.wandb_internal.Control\x12\x0c\n\x04uuid\x18\x18 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._ResultInfoB\r\n\x0bresult_type\":\n\x0b\x46inalRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"b\n\x0bVersionInfo\x12\x10\n\x08producer\x18\x01 \x01(\t\x12\x14\n\x0cmin_consumer\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"n\n\x0cHeaderRecord\x12\x31\n\x0cversion_info\x18\x01 \x01(\x0b\x32\x1b.wandb_internal.VersionInfo\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\";\n\x0c\x46ooterRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"9\n\x0b\x42ranchPoint\x12\x0b\n\x03run\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01\x12\x0e\n\x06metric\x18\x03 \x01(\t\"\x91\x05\n\tRunRecord\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0e\n\x06\x65ntity\x18\x02 \x01(\t\x12\x0f\n\x07project\x18\x03 \x01(\t\x12,\n\x06\x63onfig\x18\x04 \x01(\x0b\x32\x1c.wandb_internal.ConfigRecord\x12.\n\x07summary\x18\x05 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecord\x12\x11\n\trun_group\x18\x06 \x01(\t\x12\x10\n\x08job_type\x18\x07 \x01(\t\x12\x14\n\x0c\x64isplay_name\x18\x08 \x01(\t\x12\r\n\x05notes\x18\t \x01(\t\x12\x0c\n\x04tags\x18\n \x03(\t\x12\x30\n\x08settings\x18\x0b \x01(\x0b\x32\x1e.wandb_internal.SettingsRecord\x12\x10\n\x08sweep_id\x18\x0c \x01(\t\x12\x0c\n\x04host\x18\r \x01(\t\x12\x15\n\rstarting_step\x18\x0e \x01(\x03\x12\x12\n\nstorage_id\x18\x10 \x01(\t\x12.\n\nstart_time\x18\x11 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07resumed\x18\x12 \x01(\x08\x12\x32\n\ttelemetry\x18\x13 \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecord\x12\x0f\n\x07runtime\x18\x14 \x01(\x05\x12*\n\x03git\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.GitRepoRecord\x12\x0e\n\x06\x66orked\x18\x16 \x01(\x08\x12\x31\n\x0c\x62ranch_point\x18\x17 \x01(\x0b\x32\x1b.wandb_internal.BranchPoint\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\";\n\rGitRepoRecord\x12\x1a\n\nremote_url\x18\x01 \x01(\tR\x06remote\x12\x0e\n\x06\x63ommit\x18\x02 \x01(\t\"c\n\x0fRunUpdateResult\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"\xac\x01\n\tErrorInfo\x12\x0f\n\x07message\x18\x01 \x01(\t\x12\x31\n\x04\x63ode\x18\x02 \x01(\x0e\x32#.wandb_internal.ErrorInfo.ErrorCode\"[\n\tErrorCode\x12\x0b\n\x07UNKNOWN\x10\x00\x12\x11\n\rCOMMUNICATION\x10\x01\x12\x12\n\x0e\x41UTHENTICATION\x10\x02\x12\t\n\x05USAGE\x10\x03\x12\x0f\n\x0bUNSUPPORTED\x10\x04\"`\n\rRunExitRecord\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12\x0f\n\x07runtime\x18\x02 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x0f\n\rRunExitResult\"B\n\x13RunPreemptingRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x15\n\x13RunPreemptingResult\"i\n\x0eSettingsRecord\x12*\n\x04item\x18\x01 \x03(\x0b\x32\x1c.wandb_internal.SettingsItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"/\n\x0cSettingsItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x1a\n\x0bHistoryStep\x12\x0b\n\x03num\x18\x01 \x01(\x03\"\x92\x01\n\rHistoryRecord\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.HistoryItem\x12)\n\x04step\x18\x02 \x01(\x0b\x32\x1b.wandb_internal.HistoryStep\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\x0bHistoryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0f\n\rHistoryResult\"\xdc\x01\n\x0cOutputRecord\x12<\n\x0boutput_type\x18\x01 \x01(\x0e\x32\'.wandb_internal.OutputRecord.OutputType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0c\n\x04line\x18\x03 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"$\n\nOutputType\x12\n\n\x06STDERR\x10\x00\x12\n\n\x06STDOUT\x10\x01\"\x0e\n\x0cOutputResult\"\xe2\x01\n\x0fOutputRawRecord\x12?\n\x0boutput_type\x18\x01 \x01(\x0e\x32*.wandb_internal.OutputRawRecord.OutputType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0c\n\x04line\x18\x03 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"$\n\nOutputType\x12\n\n\x06STDERR\x10\x00\x12\n\n\x06STDOUT\x10\x01\"\x11\n\x0fOutputRawResult\"\xb4\x03\n\x0cMetricRecord\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x11\n\tglob_name\x18\x02 \x01(\t\x12\x13\n\x0bstep_metric\x18\x04 \x01(\t\x12\x19\n\x11step_metric_index\x18\x05 \x01(\x05\x12.\n\x07options\x18\x06 \x01(\x0b\x32\x1d.wandb_internal.MetricOptions\x12.\n\x07summary\x18\x07 \x01(\x0b\x32\x1d.wandb_internal.MetricSummary\x12\x35\n\x04goal\x18\x08 \x01(\x0e\x32\'.wandb_internal.MetricRecord.MetricGoal\x12/\n\x08_control\x18\t \x01(\x0b\x32\x1d.wandb_internal.MetricControl\x12\x1a\n\x12\x65xpanded_from_glob\x18\n \x01(\x08\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\nMetricGoal\x12\x0e\n\nGOAL_UNSET\x10\x00\x12\x11\n\rGOAL_MINIMIZE\x10\x01\x12\x11\n\rGOAL_MAXIMIZE\x10\x02\"\x0e\n\x0cMetricResult\"C\n\rMetricOptions\x12\x11\n\tstep_sync\x18\x01 \x01(\x08\x12\x0e\n\x06hidden\x18\x02 \x01(\x08\x12\x0f\n\x07\x64\x65\x66ined\x18\x03 \x01(\x08\"\"\n\rMetricControl\x12\x11\n\toverwrite\x18\x01 \x01(\x08\"o\n\rMetricSummary\x12\x0b\n\x03min\x18\x01 \x01(\x08\x12\x0b\n\x03max\x18\x02 \x01(\x08\x12\x0c\n\x04mean\x18\x03 \x01(\x08\x12\x0c\n\x04\x62\x65st\x18\x04 \x01(\x08\x12\x0c\n\x04last\x18\x05 \x01(\x08\x12\x0c\n\x04none\x18\x06 \x01(\x08\x12\x0c\n\x04\x63opy\x18\x07 \x01(\x08\"\x93\x01\n\x0c\x43onfigRecord\x12*\n\x06update\x18\x01 \x03(\x0b\x32\x1a.wandb_internal.ConfigItem\x12*\n\x06remove\x18\x02 \x03(\x0b\x32\x1a.wandb_internal.ConfigItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"A\n\nConfigItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0e\n\x0c\x43onfigResult\"\x96\x01\n\rSummaryRecord\x12+\n\x06update\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\x12+\n\x06remove\x18\x02 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\x0bSummaryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0f\n\rSummaryResult\"d\n\x0b\x46ilesRecord\x12(\n\x05\x66iles\x18\x01 \x03(\x0b\x32\x19.wandb_internal.FilesItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\xec\x01\n\tFilesItem\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x34\n\x06policy\x18\x02 \x01(\x0e\x32$.wandb_internal.FilesItem.PolicyType\x12\x30\n\x04type\x18\x03 \x01(\x0e\x32\".wandb_internal.FilesItem.FileType\"(\n\nPolicyType\x12\x07\n\x03NOW\x10\x00\x12\x07\n\x03\x45ND\x10\x01\x12\x08\n\x04LIVE\x10\x02\"9\n\x08\x46ileType\x12\t\n\x05OTHER\x10\x00\x12\t\n\x05WANDB\x10\x01\x12\t\n\x05MEDIA\x10\x02\x12\x0c\n\x08\x41RTIFACT\x10\x03J\x04\x08\x10\x10\x11\"\r\n\x0b\x46ilesResult\"\xe6\x01\n\x0bStatsRecord\x12\x39\n\nstats_type\x18\x01 \x01(\x0e\x32%.wandb_internal.StatsRecord.StatsType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\'\n\x04item\x18\x03 \x03(\x0b\x32\x19.wandb_internal.StatsItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x17\n\tStatsType\x12\n\n\x06SYSTEM\x10\x00\",\n\tStatsItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\xe7\x03\n\x0e\x41rtifactRecord\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0f\n\x07project\x18\x02 \x01(\t\x12\x0e\n\x06\x65ntity\x18\x03 \x01(\t\x12\x0c\n\x04type\x18\x04 \x01(\t\x12\x0c\n\x04name\x18\x05 \x01(\t\x12\x0e\n\x06\x64igest\x18\x06 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x07 \x01(\t\x12\x10\n\x08metadata\x18\x08 \x01(\t\x12\x14\n\x0cuser_created\x18\t \x01(\x08\x12\x18\n\x10use_after_commit\x18\n \x01(\x08\x12\x0f\n\x07\x61liases\x18\x0b \x03(\t\x12\x32\n\x08manifest\x18\x0c \x01(\x0b\x32 .wandb_internal.ArtifactManifest\x12\x16\n\x0e\x64istributed_id\x18\r \x01(\t\x12\x10\n\x08\x66inalize\x18\x0e \x01(\x08\x12\x11\n\tclient_id\x18\x0f \x01(\t\x12\x1a\n\x12sequence_client_id\x18\x10 \x01(\t\x12\x0f\n\x07\x62\x61se_id\x18\x11 \x01(\t\x12\x1c\n\x14ttl_duration_seconds\x18\x12 \x01(\x03\x12\x0c\n\x04tags\x18\x13 \x03(\t\x12\x19\n\x11incremental_beta1\x18\x64 \x01(\x08\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\xd8\x01\n\x10\x41rtifactManifest\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x16\n\x0estorage_policy\x18\x02 \x01(\t\x12\x46\n\x15storage_policy_config\x18\x03 \x03(\x0b\x32\'.wandb_internal.StoragePolicyConfigItem\x12\x37\n\x08\x63ontents\x18\x04 \x03(\x0b\x32%.wandb_internal.ArtifactManifestEntry\x12\x1a\n\x12manifest_file_path\x18\x05 \x01(\t\"\xcf\x01\n\x15\x41rtifactManifestEntry\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x0e\n\x06\x64igest\x18\x02 \x01(\t\x12\x0b\n\x03ref\x18\x03 \x01(\t\x12\x0c\n\x04size\x18\x04 \x01(\x03\x12\x10\n\x08mimetype\x18\x05 \x01(\t\x12\x12\n\nlocal_path\x18\x06 \x01(\t\x12\x19\n\x11\x62irth_artifact_id\x18\x07 \x01(\t\x12\x12\n\nskip_cache\x18\x08 \x01(\x08\x12(\n\x05\x65xtra\x18\x10 \x03(\x0b\x32\x19.wandb_internal.ExtraItem\",\n\tExtraItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x02 \x01(\t\":\n\x17StoragePolicyConfigItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x02 \x01(\t\"\x10\n\x0e\x41rtifactResult\"\x14\n\x12LinkArtifactResult\"\xf0\x01\n\x13LinkArtifactRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x11\n\tserver_id\x18\x02 \x01(\t\x12\x16\n\x0eportfolio_name\x18\x03 \x01(\t\x12\x18\n\x10portfolio_entity\x18\x04 \x01(\t\x12\x19\n\x11portfolio_project\x18\x05 \x01(\t\x12\x19\n\x11portfolio_aliases\x18\x06 \x03(\t\x12\x1e\n\x16portfolio_organization\x18\x07 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"[\n\x14LinkArtifactResponse\x12\x15\n\rerror_message\x18\x01 \x01(\t\x12\x1a\n\rversion_index\x18\x02 \x01(\x05H\x00\x88\x01\x01\x42\x10\n\x0e_version_index\"h\n\x08TBRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\x12\x0f\n\x07log_dir\x18\x01 \x01(\t\x12\x10\n\x08root_dir\x18\x03 \x01(\t\x12\x0c\n\x04save\x18\x02 \x01(\x08\"\n\n\x08TBResult\"}\n\x0b\x41lertRecord\x12\r\n\x05title\x18\x01 \x01(\t\x12\x0c\n\x04text\x18\x02 \x01(\t\x12\r\n\x05level\x18\x03 \x01(\t\x12\x15\n\rwait_duration\x18\x04 \x01(\x03\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\r\n\x0b\x41lertResult\"\xb3\x11\n\x07Request\x12\x38\n\x0bstop_status\x18\x01 \x01(\x0b\x32!.wandb_internal.StopStatusRequestH\x00\x12>\n\x0enetwork_status\x18\x02 \x01(\x0b\x32$.wandb_internal.NetworkStatusRequestH\x00\x12-\n\x05\x64\x65\x66\x65r\x18\x03 \x01(\x0b\x32\x1c.wandb_internal.DeferRequestH\x00\x12\x38\n\x0bget_summary\x18\x04 \x01(\x0b\x32!.wandb_internal.GetSummaryRequestH\x00\x12-\n\x05login\x18\x05 \x01(\x0b\x32\x1c.wandb_internal.LoginRequestH\x00\x12-\n\x05pause\x18\x06 \x01(\x0b\x32\x1c.wandb_internal.PauseRequestH\x00\x12/\n\x06resume\x18\x07 \x01(\x0b\x32\x1d.wandb_internal.ResumeRequestH\x00\x12\x34\n\tpoll_exit\x18\x08 \x01(\x0b\x32\x1f.wandb_internal.PollExitRequestH\x00\x12@\n\x0fsampled_history\x18\t \x01(\x0b\x32%.wandb_internal.SampledHistoryRequestH\x00\x12@\n\x0fpartial_history\x18\n \x01(\x0b\x32%.wandb_internal.PartialHistoryRequestH\x00\x12\x34\n\trun_start\x18\x0b \x01(\x0b\x32\x1f.wandb_internal.RunStartRequestH\x00\x12<\n\rcheck_version\x18\x0c \x01(\x0b\x32#.wandb_internal.CheckVersionRequestH\x00\x12:\n\x0clog_artifact\x18\r \x01(\x0b\x32\".wandb_internal.LogArtifactRequestH\x00\x12\x44\n\x11\x64ownload_artifact\x18\x0e \x01(\x0b\x32\'.wandb_internal.DownloadArtifactRequestH\x00\x12\x35\n\tkeepalive\x18\x11 \x01(\x0b\x32 .wandb_internal.KeepaliveRequestH\x00\x12>\n\x0eserver_feature\x18\x12 \x01(\x0b\x32$.wandb_internal.ServerFeatureRequestH\x00\x12\x36\n\nrun_status\x18\x14 \x01(\x0b\x32 .wandb_internal.RunStatusRequestH\x00\x12/\n\x06\x63\x61ncel\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.CancelRequestH\x00\x12\x44\n\x11internal_messages\x18\x17 \x01(\x0b\x32\'.wandb_internal.InternalMessagesRequestH\x00\x12@\n\x0fpython_packages\x18\x18 \x01(\x0b\x32%.wandb_internal.PythonPackagesRequestH\x00\x12\x33\n\x08shutdown\x18@ \x01(\x0b\x32\x1f.wandb_internal.ShutdownRequestH\x00\x12/\n\x06\x61ttach\x18\x41 \x01(\x0b\x32\x1d.wandb_internal.AttachRequestH\x00\x12/\n\x06status\x18\x42 \x01(\x0b\x32\x1d.wandb_internal.StatusRequestH\x00\x12\x38\n\x0bserver_info\x18\x43 \x01(\x0b\x32!.wandb_internal.ServerInfoRequestH\x00\x12\x38\n\x0bsender_mark\x18\x44 \x01(\x0b\x32!.wandb_internal.SenderMarkRequestH\x00\x12\x38\n\x0bsender_read\x18\x45 \x01(\x0b\x32!.wandb_internal.SenderReadRequestH\x00\x12<\n\rstatus_report\x18\x46 \x01(\x0b\x32#.wandb_internal.StatusReportRequestH\x00\x12>\n\x0esummary_record\x18G \x01(\x0b\x32$.wandb_internal.SummaryRecordRequestH\x00\x12\x42\n\x10telemetry_record\x18H \x01(\x0b\x32&.wandb_internal.TelemetryRecordRequestH\x00\x12\x32\n\x08job_info\x18I \x01(\x0b\x32\x1e.wandb_internal.JobInfoRequestH\x00\x12\x45\n\x12get_system_metrics\x18J \x01(\x0b\x32\'.wandb_internal.GetSystemMetricsRequestH\x00\x12\x34\n\tjob_input\x18M \x01(\x0b\x32\x1f.wandb_internal.JobInputRequestH\x00\x12<\n\rlink_artifact\x18N \x01(\x0b\x32#.wandb_internal.LinkArtifactRequestH\x00\x12N\n\x17run_finish_without_exit\x18O \x01(\x0b\x32+.wandb_internal.RunFinishWithoutExitRequestH\x00\x12\x38\n\x0bsync_finish\x18Q \x01(\x0b\x32!.wandb_internal.SyncFinishRequestH\x00\x12;\n\noperations\x18R \x01(\x0b\x32%.wandb_internal.OperationStatsRequestH\x00\x12\x39\n\x0btest_inject\x18\xe8\x07 \x01(\x0b\x32!.wandb_internal.TestInjectRequestH\x00\x42\x0e\n\x0crequest_typeJ\x04\x08\x16\x10\x17J\x04\x08K\x10LJ\x04\x08L\x10MJ\x04\x08P\x10Q\"\x9b\x0e\n\x08Response\x12?\n\x12keepalive_response\x18\x12 \x01(\x0b\x32!.wandb_internal.KeepaliveResponseH\x00\x12\x42\n\x14stop_status_response\x18\x13 \x01(\x0b\x32\".wandb_internal.StopStatusResponseH\x00\x12H\n\x17network_status_response\x18\x14 \x01(\x0b\x32%.wandb_internal.NetworkStatusResponseH\x00\x12\x37\n\x0elogin_response\x18\x18 \x01(\x0b\x32\x1d.wandb_internal.LoginResponseH\x00\x12\x42\n\x14get_summary_response\x18\x19 \x01(\x0b\x32\".wandb_internal.GetSummaryResponseH\x00\x12>\n\x12poll_exit_response\x18\x1a \x01(\x0b\x32 .wandb_internal.PollExitResponseH\x00\x12J\n\x18sampled_history_response\x18\x1b \x01(\x0b\x32&.wandb_internal.SampledHistoryResponseH\x00\x12>\n\x12run_start_response\x18\x1c \x01(\x0b\x32 .wandb_internal.RunStartResponseH\x00\x12\x46\n\x16\x63heck_version_response\x18\x1d \x01(\x0b\x32$.wandb_internal.CheckVersionResponseH\x00\x12\x44\n\x15log_artifact_response\x18\x1e \x01(\x0b\x32#.wandb_internal.LogArtifactResponseH\x00\x12N\n\x1a\x64ownload_artifact_response\x18\x1f \x01(\x0b\x32(.wandb_internal.DownloadArtifactResponseH\x00\x12H\n\x17server_feature_response\x18 \x01(\x0b\x32%.wandb_internal.ServerFeatureResponseH\x00\x12@\n\x13run_status_response\x18# \x01(\x0b\x32!.wandb_internal.RunStatusResponseH\x00\x12\x39\n\x0f\x63\x61ncel_response\x18$ \x01(\x0b\x32\x1e.wandb_internal.CancelResponseH\x00\x12N\n\x1ainternal_messages_response\x18% \x01(\x0b\x32(.wandb_internal.InternalMessagesResponseH\x00\x12=\n\x11shutdown_response\x18@ \x01(\x0b\x32 .wandb_internal.ShutdownResponseH\x00\x12\x39\n\x0f\x61ttach_response\x18\x41 \x01(\x0b\x32\x1e.wandb_internal.AttachResponseH\x00\x12\x39\n\x0fstatus_response\x18\x42 \x01(\x0b\x32\x1e.wandb_internal.StatusResponseH\x00\x12\x42\n\x14server_info_response\x18\x43 \x01(\x0b\x32\".wandb_internal.ServerInfoResponseH\x00\x12<\n\x11job_info_response\x18\x44 \x01(\x0b\x32\x1f.wandb_internal.JobInfoResponseH\x00\x12O\n\x1bget_system_metrics_response\x18\x45 \x01(\x0b\x32(.wandb_internal.GetSystemMetricsResponseH\x00\x12\x46\n\x16link_artifact_response\x18G \x01(\x0b\x32$.wandb_internal.LinkArtifactResponseH\x00\x12\x35\n\rsync_response\x18\x46 \x01(\x0b\x32\x1c.wandb_internal.SyncResponseH\x00\x12X\n run_finish_without_exit_response\x18H \x01(\x0b\x32,.wandb_internal.RunFinishWithoutExitResponseH\x00\x12\x45\n\x13operations_response\x18J \x01(\x0b\x32&.wandb_internal.OperationStatsResponseH\x00\x12\x43\n\x14test_inject_response\x18\xe8\x07 \x01(\x0b\x32\".wandb_internal.TestInjectResponseH\x00\x42\x0f\n\rresponse_typeJ\x04\x08I\x10J\"\xc0\x02\n\x0c\x44\x65\x66\x65rRequest\x12\x36\n\x05state\x18\x01 \x01(\x0e\x32\'.wandb_internal.DeferRequest.DeferState\"\xf7\x01\n\nDeferState\x12\t\n\x05\x42\x45GIN\x10\x00\x12\r\n\tFLUSH_RUN\x10\x01\x12\x0f\n\x0b\x46LUSH_STATS\x10\x02\x12\x19\n\x15\x46LUSH_PARTIAL_HISTORY\x10\x03\x12\x0c\n\x08\x46LUSH_TB\x10\x04\x12\r\n\tFLUSH_SUM\x10\x05\x12\x13\n\x0f\x46LUSH_DEBOUNCER\x10\x06\x12\x10\n\x0c\x46LUSH_OUTPUT\x10\x07\x12\r\n\tFLUSH_JOB\x10\x08\x12\r\n\tFLUSH_DIR\x10\t\x12\x0c\n\x08\x46LUSH_FP\x10\n\x12\x0b\n\x07JOIN_FP\x10\x0b\x12\x0c\n\x08\x46LUSH_FS\x10\x0c\x12\x0f\n\x0b\x46LUSH_FINAL\x10\r\x12\x07\n\x03\x45ND\x10\x0e\"<\n\x0cPauseRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x0f\n\rPauseResponse\"=\n\rResumeRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x10\n\x0eResumeResponse\"M\n\x0cLoginRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"&\n\rLoginResponse\x12\x15\n\ractive_entity\x18\x01 \x01(\t\"A\n\x11GetSummaryRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"?\n\x12GetSummaryResponse\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\"G\n\x17GetSystemMetricsRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"R\n\x12SystemMetricSample\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\r\n\x05value\x18\x02 \x01(\x02\"I\n\x13SystemMetricsBuffer\x12\x32\n\x06record\x18\x01 \x03(\x0b\x32\".wandb_internal.SystemMetricSample\"\xca\x01\n\x18GetSystemMetricsResponse\x12S\n\x0esystem_metrics\x18\x01 \x03(\x0b\x32;.wandb_internal.GetSystemMetricsResponse.SystemMetricsEntry\x1aY\n\x12SystemMetricsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x32\n\x05value\x18\x02 \x01(\x0b\x32#.wandb_internal.SystemMetricsBuffer:\x02\x38\x01\"=\n\rStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\")\n\x0eStatusResponse\x12\x17\n\x0frun_should_stop\x18\x01 \x01(\x08\"A\n\x11StopStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"-\n\x12StopStatusResponse\x12\x17\n\x0frun_should_stop\x18\x01 \x01(\x08\"D\n\x14NetworkStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"P\n\x15NetworkStatusResponse\x12\x37\n\x11network_responses\x18\x01 \x03(\x0b\x32\x1c.wandb_internal.HttpResponse\"D\n\x0cHttpResponse\x12\x18\n\x10http_status_code\x18\x01 \x01(\x05\x12\x1a\n\x12http_response_text\x18\x02 \x01(\t\"G\n\x17InternalMessagesRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"N\n\x18InternalMessagesResponse\x12\x32\n\x08messages\x18\x01 \x01(\x0b\x32 .wandb_internal.InternalMessages\"#\n\x10InternalMessages\x12\x0f\n\x07warning\x18\x01 \x03(\t\"?\n\x0fPollExitRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\xf5\x01\n\x10PollExitResponse\x12\x0c\n\x04\x64one\x18\x01 \x01(\x08\x12\x32\n\x0b\x65xit_result\x18\x02 \x01(\x0b\x32\x1d.wandb_internal.RunExitResult\x12\x35\n\x0cpusher_stats\x18\x03 \x01(\x0b\x32\x1f.wandb_internal.FilePusherStats\x12/\n\x0b\x66ile_counts\x18\x04 \x01(\x0b\x32\x1a.wandb_internal.FileCounts\x12\x37\n\x0foperation_stats\x18\x05 \x01(\x0b\x32\x1e.wandb_internal.OperationStats\"E\n\x15OperationStatsRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"Q\n\x16OperationStatsResponse\x12\x37\n\x0foperation_stats\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.OperationStats\"Y\n\x0eOperationStats\x12-\n\noperations\x18\x01 \x03(\x0b\x32\x19.wandb_internal.Operation\x12\x18\n\x10total_operations\x18\x02 \x01(\x03\"\x87\x01\n\tOperation\x12\x0c\n\x04\x64\x65sc\x18\x01 \x01(\t\x12\x17\n\x0fruntime_seconds\x18\x02 \x01(\x01\x12\x10\n\x08progress\x18\x03 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x04 \x01(\t\x12+\n\x08subtasks\x18\x05 \x03(\x0b\x32\x19.wandb_internal.Operation\"\x13\n\x11SenderMarkRequest\"\x13\n\x11SyncFinishRequest\"E\n\x0cSyncResponse\x12\x0b\n\x03url\x18\x01 \x01(\t\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"?\n\x11SenderReadRequest\x12\x14\n\x0cstart_offset\x18\x01 \x01(\x03\x12\x14\n\x0c\x66inal_offset\x18\x02 \x01(\x03\"m\n\x13StatusReportRequest\x12\x12\n\nrecord_num\x18\x01 \x01(\x03\x12\x13\n\x0bsent_offset\x18\x02 \x01(\x03\x12-\n\tsync_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"F\n\x14SummaryRecordRequest\x12.\n\x07summary\x18\x01 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecord\"L\n\x16TelemetryRecordRequest\x12\x32\n\ttelemetry\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecord\"A\n\x11ServerInfoRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"|\n\x12ServerInfoResponse\x12-\n\nlocal_info\x18\x01 \x01(\x0b\x32\x19.wandb_internal.LocalInfo\x12\x37\n\x0fserver_messages\x18\x02 \x01(\x0b\x32\x1e.wandb_internal.ServerMessages\"=\n\x0eServerMessages\x12+\n\x04item\x18\x01 \x03(\x0b\x32\x1d.wandb_internal.ServerMessage\"e\n\rServerMessage\x12\x12\n\nplain_text\x18\x01 \x01(\t\x12\x10\n\x08utf_text\x18\x02 \x01(\t\x12\x11\n\thtml_text\x18\x03 \x01(\t\x12\x0c\n\x04type\x18\x04 \x01(\t\x12\r\n\x05level\x18\x05 \x01(\x05\"c\n\nFileCounts\x12\x13\n\x0bwandb_count\x18\x01 \x01(\x05\x12\x13\n\x0bmedia_count\x18\x02 \x01(\x05\x12\x16\n\x0e\x61rtifact_count\x18\x03 \x01(\x05\x12\x13\n\x0bother_count\x18\x04 \x01(\x05\"U\n\x0f\x46ilePusherStats\x12\x16\n\x0euploaded_bytes\x18\x01 \x01(\x03\x12\x13\n\x0btotal_bytes\x18\x02 \x01(\x03\x12\x15\n\rdeduped_bytes\x18\x03 \x01(\x03\"\x1e\n\rFilesUploaded\x12\r\n\x05\x66iles\x18\x01 \x03(\t\"\xf4\x01\n\x17\x46ileTransferInfoRequest\x12\x42\n\x04type\x18\x01 \x01(\x0e\x32\x34.wandb_internal.FileTransferInfoRequest.TransferType\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\x0b\n\x03url\x18\x03 \x01(\t\x12\x0c\n\x04size\x18\x04 \x01(\x03\x12\x11\n\tprocessed\x18\x05 \x01(\x03\x12/\n\x0b\x66ile_counts\x18\x06 \x01(\x0b\x32\x1a.wandb_internal.FileCounts\"(\n\x0cTransferType\x12\n\n\x06Upload\x10\x00\x12\x0c\n\x08\x44ownload\x10\x01\"1\n\tLocalInfo\x12\x0f\n\x07version\x18\x01 \x01(\t\x12\x13\n\x0bout_of_date\x18\x02 \x01(\x08\"?\n\x0fShutdownRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x12\n\x10ShutdownResponse\"P\n\rAttachRequest\x12\x11\n\tattach_id\x18\x14 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"b\n\x0e\x41ttachResponse\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"\xd5\x02\n\x11TestInjectRequest\x12\x13\n\x0bhandler_exc\x18\x01 \x01(\x08\x12\x14\n\x0chandler_exit\x18\x02 \x01(\x08\x12\x15\n\rhandler_abort\x18\x03 \x01(\x08\x12\x12\n\nsender_exc\x18\x04 \x01(\x08\x12\x13\n\x0bsender_exit\x18\x05 \x01(\x08\x12\x14\n\x0csender_abort\x18\x06 \x01(\x08\x12\x0f\n\x07req_exc\x18\x07 \x01(\x08\x12\x10\n\x08req_exit\x18\x08 \x01(\x08\x12\x11\n\treq_abort\x18\t \x01(\x08\x12\x10\n\x08resp_exc\x18\n \x01(\x08\x12\x11\n\tresp_exit\x18\x0b \x01(\x08\x12\x12\n\nresp_abort\x18\x0c \x01(\x08\x12\x10\n\x08msg_drop\x18\r \x01(\x08\x12\x10\n\x08msg_hang\x18\x0e \x01(\x08\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x14\n\x12TestInjectResponse\"\x1e\n\rHistoryAction\x12\r\n\x05\x66lush\x18\x01 \x01(\x08\"\xca\x01\n\x15PartialHistoryRequest\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.HistoryItem\x12)\n\x04step\x18\x02 \x01(\x0b\x32\x1b.wandb_internal.HistoryStep\x12-\n\x06\x61\x63tion\x18\x03 \x01(\x0b\x32\x1d.wandb_internal.HistoryAction\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x18\n\x16PartialHistoryResponse\"E\n\x15SampledHistoryRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"_\n\x12SampledHistoryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x14\n\x0cvalues_float\x18\x03 \x03(\x02\x12\x12\n\nvalues_int\x18\x04 \x03(\x03\"J\n\x16SampledHistoryResponse\x12\x30\n\x04item\x18\x01 \x03(\x0b\x32\".wandb_internal.SampledHistoryItem\"@\n\x10RunStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"x\n\x11RunStatusResponse\x12\x18\n\x10sync_items_total\x18\x01 \x01(\x03\x12\x1a\n\x12sync_items_pending\x18\x02 \x01(\x03\x12-\n\tsync_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"g\n\x0fRunStartRequest\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x12\n\x10RunStartResponse\"K\n\x1bRunFinishWithoutExitRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x1e\n\x1cRunFinishWithoutExitResponse\"\\\n\x13\x43heckVersionRequest\x12\x17\n\x0f\x63urrent_version\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"]\n\x14\x43heckVersionResponse\x12\x17\n\x0fupgrade_message\x18\x01 \x01(\t\x12\x14\n\x0cyank_message\x18\x02 \x01(\t\x12\x16\n\x0e\x64\x65lete_message\x18\x03 \x01(\t\">\n\x0eJobInfoRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"6\n\x0fJobInfoResponse\x12\x12\n\nsequenceId\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"\x9f\x01\n\x12LogArtifactRequest\x12\x30\n\x08\x61rtifact\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.ArtifactRecord\x12\x14\n\x0chistory_step\x18\x02 \x01(\x03\x12\x13\n\x0bstaging_dir\x18\x03 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"A\n\x13LogArtifactResponse\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x15\n\rerror_message\x18\x02 \x01(\t\"\xbe\x01\n\x17\x44ownloadArtifactRequest\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x15\n\rdownload_root\x18\x02 \x01(\t\x12 \n\x18\x61llow_missing_references\x18\x04 \x01(\x08\x12\x12\n\nskip_cache\x18\x05 \x01(\x08\x12\x13\n\x0bpath_prefix\x18\x06 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"1\n\x18\x44ownloadArtifactResponse\x12\x15\n\rerror_message\x18\x01 \x01(\t\"@\n\x10KeepaliveRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x13\n\x11KeepaliveResponse\"q\n\x0c\x41rtifactInfo\x12\x10\n\x08\x61rtifact\x18\x01 \x01(\t\x12\x12\n\nentrypoint\x18\x02 \x03(\t\x12\x10\n\x08notebook\x18\x03 \x01(\x08\x12\x15\n\rbuild_context\x18\x04 \x01(\t\x12\x12\n\ndockerfile\x18\x05 \x01(\t\")\n\x07GitInfo\x12\x0e\n\x06remote\x18\x01 \x01(\t\x12\x0e\n\x06\x63ommit\x18\x02 \x01(\t\"\x87\x01\n\tGitSource\x12)\n\x08git_info\x18\x01 \x01(\x0b\x32\x17.wandb_internal.GitInfo\x12\x12\n\nentrypoint\x18\x02 \x03(\t\x12\x10\n\x08notebook\x18\x03 \x01(\x08\x12\x15\n\rbuild_context\x18\x04 \x01(\t\x12\x12\n\ndockerfile\x18\x05 \x01(\t\"\x1c\n\x0bImageSource\x12\r\n\x05image\x18\x01 \x01(\t\"\x8c\x01\n\x06Source\x12&\n\x03git\x18\x01 \x01(\x0b\x32\x19.wandb_internal.GitSource\x12.\n\x08\x61rtifact\x18\x02 \x01(\x0b\x32\x1c.wandb_internal.ArtifactInfo\x12*\n\x05image\x18\x03 \x01(\x0b\x32\x1b.wandb_internal.ImageSource\"k\n\tJobSource\x12\x10\n\x08_version\x18\x01 \x01(\t\x12\x13\n\x0bsource_type\x18\x02 \x01(\t\x12&\n\x06source\x18\x03 \x01(\x0b\x32\x16.wandb_internal.Source\x12\x0f\n\x07runtime\x18\x04 \x01(\t\"V\n\x12PartialJobArtifact\x12\x10\n\x08job_name\x18\x01 \x01(\t\x12.\n\x0bsource_info\x18\x02 \x01(\x0b\x32\x19.wandb_internal.JobSource\"\x9d\x01\n\x11UseArtifactRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x33\n\x07partial\x18\x04 \x01(\x0b\x32\".wandb_internal.PartialJobArtifact\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x13\n\x11UseArtifactResult\"R\n\rCancelRequest\x12\x13\n\x0b\x63\x61ncel_slot\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x10\n\x0e\x43\x61ncelResponse\"\'\n\x08\x44iskInfo\x12\r\n\x05total\x18\x01 \x01(\x04\x12\x0c\n\x04used\x18\x02 \x01(\x04\"\x1b\n\nMemoryInfo\x12\r\n\x05total\x18\x01 \x01(\x04\"/\n\x07\x43puInfo\x12\r\n\x05\x63ount\x18\x01 \x01(\r\x12\x15\n\rcount_logical\x18\x02 \x01(\r\"\x9a\x01\n\tAppleInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\necpu_cores\x18\x02 \x01(\r\x12\x12\n\npcpu_cores\x18\x03 \x01(\r\x12\x11\n\tgpu_cores\x18\x04 \x01(\r\x12\x11\n\tmemory_gb\x18\x05 \x01(\r\x12\x18\n\x10swap_total_bytes\x18\x06 \x01(\x04\x12\x17\n\x0fram_total_bytes\x18\x07 \x01(\x04\"k\n\rGpuNvidiaInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x14\n\x0cmemory_total\x18\x02 \x01(\x04\x12\x12\n\ncuda_cores\x18\x03 \x01(\r\x12\x14\n\x0c\x61rchitecture\x18\x04 \x01(\t\x12\x0c\n\x04uuid\x18\x05 \x01(\t\"\x89\x02\n\nGpuAmdInfo\x12\n\n\x02id\x18\x01 \x01(\t\x12\x11\n\tunique_id\x18\x02 \x01(\t\x12\x15\n\rvbios_version\x18\x03 \x01(\t\x12\x19\n\x11performance_level\x18\x04 \x01(\t\x12\x15\n\rgpu_overdrive\x18\x05 \x01(\t\x12\x1c\n\x14gpu_memory_overdrive\x18\x06 \x01(\t\x12\x11\n\tmax_power\x18\x07 \x01(\t\x12\x0e\n\x06series\x18\x08 \x01(\t\x12\r\n\x05model\x18\t \x01(\t\x12\x0e\n\x06vendor\x18\n \x01(\t\x12\x0b\n\x03sku\x18\x0b \x01(\t\x12\x12\n\nsclk_range\x18\x0c \x01(\t\x12\x12\n\nmclk_range\x18\r \x01(\t\"n\n\x0cTrainiumInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06vendor\x18\x02 \x01(\t\x12\x1b\n\x13neuron_device_count\x18\x03 \x01(\r\x12#\n\x1bneuroncore_per_device_count\x18\x04 \x01(\r\"Q\n\x07TPUInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07hbm_gib\x18\x02 \x01(\r\x12\x18\n\x10\x64\x65vices_per_chip\x18\x03 \x01(\r\x12\r\n\x05\x63ount\x18\x04 \x01(\r\"E\n\rCoreWeaveInfo\x12\x14\n\x0c\x63luster_name\x18\x01 \x01(\t\x12\x0e\n\x06org_id\x18\x02 \x01(\t\x12\x0e\n\x06region\x18\x03 \x01(\t\"\xa8\t\n\x11\x45nvironmentRecord\x12\n\n\x02os\x18\x01 \x01(\t\x12\x0e\n\x06python\x18\x02 \x01(\t\x12\x39\n\nstarted_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\tstartedAt\x12\x0e\n\x06\x64ocker\x18\x04 \x01(\t\x12\x0c\n\x04\x61rgs\x18\x05 \x03(\t\x12\x0f\n\x07program\x18\x06 \x01(\t\x12\x1b\n\tcode_path\x18\x07 \x01(\tR\x08\x63odePath\x12&\n\x0f\x63ode_path_local\x18\x08 \x01(\tR\rcodePathLocal\x12*\n\x03git\x18\t \x01(\x0b\x32\x1d.wandb_internal.GitRepoRecord\x12\r\n\x05\x65mail\x18\n \x01(\t\x12\x0c\n\x04root\x18\x0b \x01(\t\x12\x0c\n\x04host\x18\x0c \x01(\t\x12\x10\n\x08username\x18\r \x01(\t\x12\x12\n\nexecutable\x18\x0e \x01(\t\x12\r\n\x05\x63olab\x18\x0f \x01(\t\x12\x1c\n\tcpu_count\x18\x10 \x01(\rR\tcpu_count\x12,\n\x11\x63pu_count_logical\x18\x11 \x01(\rR\x11\x63pu_count_logical\x12\x15\n\x08gpu_type\x18\x12 \x01(\tR\x03gpu\x12\x1c\n\tgpu_count\x18\x13 \x01(\rR\tgpu_count\x12\x39\n\x04\x64isk\x18\x14 \x03(\x0b\x32+.wandb_internal.EnvironmentRecord.DiskEntry\x12*\n\x06memory\x18\x15 \x01(\x0b\x32\x1a.wandb_internal.MemoryInfo\x12$\n\x03\x63pu\x18\x16 \x01(\x0b\x32\x17.wandb_internal.CpuInfo\x12(\n\x05\x61pple\x18\x17 \x01(\x0b\x32\x19.wandb_internal.AppleInfo\x12=\n\ngpu_nvidia\x18\x18 \x03(\x0b\x32\x1d.wandb_internal.GpuNvidiaInfoR\ngpu_nvidia\x12\x14\n\x0c\x63uda_version\x18\x19 \x01(\t\x12\x34\n\x07gpu_amd\x18\x1a \x03(\x0b\x32\x1a.wandb_internal.GpuAmdInfoR\x07gpu_amd\x12;\n\x05slurm\x18\x1b \x03(\x0b\x32,.wandb_internal.EnvironmentRecord.SlurmEntry\x12.\n\x08trainium\x18\x1c \x01(\x0b\x32\x1c.wandb_internal.TrainiumInfo\x12$\n\x03tpu\x18\x1d \x01(\x0b\x32\x17.wandb_internal.TPUInfo\x12\x30\n\tcoreweave\x18\x1e \x01(\x0b\x32\x1d.wandb_internal.CoreWeaveInfo\x12\x12\n\twriter_id\x18\xc7\x01 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\x1a\x45\n\tDiskEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.wandb_internal.DiskInfo:\x02\x38\x01\x1a,\n\nSlurmEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\x15PythonPackagesRequest\x12\x44\n\x07package\x18\x01 \x03(\x0b\x32\x33.wandb_internal.PythonPackagesRequest.PythonPackage\x1a.\n\rPythonPackage\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"\x1c\n\x0cJobInputPath\x12\x0c\n\x04path\x18\x01 \x03(\t\"\xd6\x01\n\x0eJobInputSource\x12\x44\n\nrun_config\x18\x01 \x01(\x0b\x32..wandb_internal.JobInputSource.RunConfigSourceH\x00\x12?\n\x04\x66ile\x18\x02 \x01(\x0b\x32/.wandb_internal.JobInputSource.ConfigFileSourceH\x00\x1a\x11\n\x0fRunConfigSource\x1a \n\x10\x43onfigFileSource\x12\x0c\n\x04path\x18\x01 \x01(\tB\x08\n\x06source\"\xc7\x01\n\x0fJobInputRequest\x12\x34\n\x0cinput_source\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.JobInputSource\x12\x33\n\rinclude_paths\x18\x02 \x03(\x0b\x32\x1c.wandb_internal.JobInputPath\x12\x33\n\rexclude_paths\x18\x03 \x03(\x0b\x32\x1c.wandb_internal.JobInputPath\x12\x14\n\x0cinput_schema\x18\x04 \x01(\t\"t\n\x14ServerFeatureRequest\x12.\n\x07\x66\x65\x61ture\x18\x01 \x01(\x0e\x32\x1d.wandb_internal.ServerFeature\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"K\n\x15ServerFeatureResponse\x12\x32\n\x07\x66\x65\x61ture\x18\x01 \x01(\x0b\x32!.wandb_internal.ServerFeatureItem\"2\n\x11ServerFeatureItem\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07\x65nabled\x18\x02 \x01(\x08*\x80\x04\n\rServerFeature\x12\x13\n\x0fLARGE_FILENAMES\x10\x00\x12\x11\n\rARTIFACT_TAGS\x10\x01\x12\x0e\n\nCLIENT_IDS\x10\x02\x12\x1c\n\x18\x41RTIFACT_REGISTRY_SEARCH\x10\x03\x12\x1b\n\x17STRUCTURED_CONSOLE_LOGS\x10\x04\x12(\n$ARTIFACT_COLLECTION_MEMBERSHIP_FILES\x10\x05\x12\x38\n4ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER\x10\x06\x12\x34\n0USE_ARTIFACT_WITH_ENTITY_AND_PROJECT_INFORMATION\x10\x07\x12\x1f\n\x1b\x45XPAND_DEFINED_METRIC_GLOBS\x10\x08\x12\x1f\n\x1b\x41UTOMATION_EVENT_RUN_METRIC\x10\t\x12&\n\"AUTOMATION_EVENT_RUN_METRIC_CHANGE\x10\n\x12\x1b\n\x17\x41UTOMATION_ACTION_NO_OP\x10\x0b\x12/\n+INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION\x10\x0c\x12*\n&PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP\x10\rB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_internal_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY']._loaded_options = None + _globals['_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY']._serialized_options = b'8\001' + _globals['_ENVIRONMENTRECORD_DISKENTRY']._loaded_options = None + _globals['_ENVIRONMENTRECORD_DISKENTRY']._serialized_options = b'8\001' + _globals['_ENVIRONMENTRECORD_SLURMENTRY']._loaded_options = None + _globals['_ENVIRONMENTRECORD_SLURMENTRY']._serialized_options = b'8\001' + _globals['_SERVERFEATURE']._serialized_start=22228 + _globals['_SERVERFEATURE']._serialized_end=22740 + _globals['_RECORD']._serialized_start=180 + _globals['_RECORD']._serialized_end=1411 + _globals['_CONTROL']._serialized_start=1414 + _globals['_CONTROL']._serialized_end=1582 + _globals['_RESULT']._serialized_start=1585 + _globals['_RESULT']._serialized_end=2084 + _globals['_FINALRECORD']._serialized_start=2086 + _globals['_FINALRECORD']._serialized_end=2144 + _globals['_VERSIONINFO']._serialized_start=2146 + _globals['_VERSIONINFO']._serialized_end=2244 + _globals['_HEADERRECORD']._serialized_start=2246 + _globals['_HEADERRECORD']._serialized_end=2356 + _globals['_FOOTERRECORD']._serialized_start=2358 + _globals['_FOOTERRECORD']._serialized_end=2417 + _globals['_BRANCHPOINT']._serialized_start=2419 + _globals['_BRANCHPOINT']._serialized_end=2476 + _globals['_RUNRECORD']._serialized_start=2479 + _globals['_RUNRECORD']._serialized_end=3136 + _globals['_GITREPORECORD']._serialized_start=3138 + _globals['_GITREPORECORD']._serialized_end=3197 + _globals['_RUNUPDATERESULT']._serialized_start=3199 + _globals['_RUNUPDATERESULT']._serialized_end=3298 + _globals['_ERRORINFO']._serialized_start=3301 + _globals['_ERRORINFO']._serialized_end=3473 + _globals['_ERRORINFO_ERRORCODE']._serialized_start=3382 + _globals['_ERRORINFO_ERRORCODE']._serialized_end=3473 + _globals['_RUNEXITRECORD']._serialized_start=3475 + _globals['_RUNEXITRECORD']._serialized_end=3571 + _globals['_RUNEXITRESULT']._serialized_start=3573 + _globals['_RUNEXITRESULT']._serialized_end=3588 + _globals['_RUNPREEMPTINGRECORD']._serialized_start=3590 + _globals['_RUNPREEMPTINGRECORD']._serialized_end=3656 + _globals['_RUNPREEMPTINGRESULT']._serialized_start=3658 + _globals['_RUNPREEMPTINGRESULT']._serialized_end=3679 + _globals['_SETTINGSRECORD']._serialized_start=3681 + _globals['_SETTINGSRECORD']._serialized_end=3786 + _globals['_SETTINGSITEM']._serialized_start=3788 + _globals['_SETTINGSITEM']._serialized_end=3835 + _globals['_HISTORYSTEP']._serialized_start=3837 + _globals['_HISTORYSTEP']._serialized_end=3863 + _globals['_HISTORYRECORD']._serialized_start=3866 + _globals['_HISTORYRECORD']._serialized_end=4012 + _globals['_HISTORYITEM']._serialized_start=4014 + _globals['_HISTORYITEM']._serialized_end=4080 + _globals['_HISTORYRESULT']._serialized_start=4082 + _globals['_HISTORYRESULT']._serialized_end=4097 + _globals['_OUTPUTRECORD']._serialized_start=4100 + _globals['_OUTPUTRECORD']._serialized_end=4320 + _globals['_OUTPUTRECORD_OUTPUTTYPE']._serialized_start=4284 + _globals['_OUTPUTRECORD_OUTPUTTYPE']._serialized_end=4320 + _globals['_OUTPUTRESULT']._serialized_start=4322 + _globals['_OUTPUTRESULT']._serialized_end=4336 + _globals['_OUTPUTRAWRECORD']._serialized_start=4339 + _globals['_OUTPUTRAWRECORD']._serialized_end=4565 + _globals['_OUTPUTRAWRECORD_OUTPUTTYPE']._serialized_start=4284 + _globals['_OUTPUTRAWRECORD_OUTPUTTYPE']._serialized_end=4320 + _globals['_OUTPUTRAWRESULT']._serialized_start=4567 + _globals['_OUTPUTRAWRESULT']._serialized_end=4584 + _globals['_METRICRECORD']._serialized_start=4587 + _globals['_METRICRECORD']._serialized_end=5023 + _globals['_METRICRECORD_METRICGOAL']._serialized_start=4957 + _globals['_METRICRECORD_METRICGOAL']._serialized_end=5023 + _globals['_METRICRESULT']._serialized_start=5025 + _globals['_METRICRESULT']._serialized_end=5039 + _globals['_METRICOPTIONS']._serialized_start=5041 + _globals['_METRICOPTIONS']._serialized_end=5108 + _globals['_METRICCONTROL']._serialized_start=5110 + _globals['_METRICCONTROL']._serialized_end=5144 + _globals['_METRICSUMMARY']._serialized_start=5146 + _globals['_METRICSUMMARY']._serialized_end=5257 + _globals['_CONFIGRECORD']._serialized_start=5260 + _globals['_CONFIGRECORD']._serialized_end=5407 + _globals['_CONFIGITEM']._serialized_start=5409 + _globals['_CONFIGITEM']._serialized_end=5474 + _globals['_CONFIGRESULT']._serialized_start=5476 + _globals['_CONFIGRESULT']._serialized_end=5490 + _globals['_SUMMARYRECORD']._serialized_start=5493 + _globals['_SUMMARYRECORD']._serialized_end=5643 + _globals['_SUMMARYITEM']._serialized_start=5645 + _globals['_SUMMARYITEM']._serialized_end=5711 + _globals['_SUMMARYRESULT']._serialized_start=5713 + _globals['_SUMMARYRESULT']._serialized_end=5728 + _globals['_FILESRECORD']._serialized_start=5730 + _globals['_FILESRECORD']._serialized_end=5830 + _globals['_FILESITEM']._serialized_start=5833 + _globals['_FILESITEM']._serialized_end=6069 + _globals['_FILESITEM_POLICYTYPE']._serialized_start=5964 + _globals['_FILESITEM_POLICYTYPE']._serialized_end=6004 + _globals['_FILESITEM_FILETYPE']._serialized_start=6006 + _globals['_FILESITEM_FILETYPE']._serialized_end=6063 + _globals['_FILESRESULT']._serialized_start=6071 + _globals['_FILESRESULT']._serialized_end=6084 + _globals['_STATSRECORD']._serialized_start=6087 + _globals['_STATSRECORD']._serialized_end=6317 + _globals['_STATSRECORD_STATSTYPE']._serialized_start=6294 + _globals['_STATSRECORD_STATSTYPE']._serialized_end=6317 + _globals['_STATSITEM']._serialized_start=6319 + _globals['_STATSITEM']._serialized_end=6363 + _globals['_ARTIFACTRECORD']._serialized_start=6366 + _globals['_ARTIFACTRECORD']._serialized_end=6853 + _globals['_ARTIFACTMANIFEST']._serialized_start=6856 + _globals['_ARTIFACTMANIFEST']._serialized_end=7072 + _globals['_ARTIFACTMANIFESTENTRY']._serialized_start=7075 + _globals['_ARTIFACTMANIFESTENTRY']._serialized_end=7282 + _globals['_EXTRAITEM']._serialized_start=7284 + _globals['_EXTRAITEM']._serialized_end=7328 + _globals['_STORAGEPOLICYCONFIGITEM']._serialized_start=7330 + _globals['_STORAGEPOLICYCONFIGITEM']._serialized_end=7388 + _globals['_ARTIFACTRESULT']._serialized_start=7390 + _globals['_ARTIFACTRESULT']._serialized_end=7406 + _globals['_LINKARTIFACTRESULT']._serialized_start=7408 + _globals['_LINKARTIFACTRESULT']._serialized_end=7428 + _globals['_LINKARTIFACTREQUEST']._serialized_start=7431 + _globals['_LINKARTIFACTREQUEST']._serialized_end=7671 + _globals['_LINKARTIFACTRESPONSE']._serialized_start=7673 + _globals['_LINKARTIFACTRESPONSE']._serialized_end=7764 + _globals['_TBRECORD']._serialized_start=7766 + _globals['_TBRECORD']._serialized_end=7870 + _globals['_TBRESULT']._serialized_start=7872 + _globals['_TBRESULT']._serialized_end=7882 + _globals['_ALERTRECORD']._serialized_start=7884 + _globals['_ALERTRECORD']._serialized_end=8009 + _globals['_ALERTRESULT']._serialized_start=8011 + _globals['_ALERTRESULT']._serialized_end=8024 + _globals['_REQUEST']._serialized_start=8027 + _globals['_REQUEST']._serialized_end=10254 + _globals['_RESPONSE']._serialized_start=10257 + _globals['_RESPONSE']._serialized_end=12076 + _globals['_DEFERREQUEST']._serialized_start=12079 + _globals['_DEFERREQUEST']._serialized_end=12399 + _globals['_DEFERREQUEST_DEFERSTATE']._serialized_start=12152 + _globals['_DEFERREQUEST_DEFERSTATE']._serialized_end=12399 + _globals['_PAUSEREQUEST']._serialized_start=12401 + _globals['_PAUSEREQUEST']._serialized_end=12461 + _globals['_PAUSERESPONSE']._serialized_start=12463 + _globals['_PAUSERESPONSE']._serialized_end=12478 + _globals['_RESUMEREQUEST']._serialized_start=12480 + _globals['_RESUMEREQUEST']._serialized_end=12541 + _globals['_RESUMERESPONSE']._serialized_start=12543 + _globals['_RESUMERESPONSE']._serialized_end=12559 + _globals['_LOGINREQUEST']._serialized_start=12561 + _globals['_LOGINREQUEST']._serialized_end=12638 + _globals['_LOGINRESPONSE']._serialized_start=12640 + _globals['_LOGINRESPONSE']._serialized_end=12678 + _globals['_GETSUMMARYREQUEST']._serialized_start=12680 + _globals['_GETSUMMARYREQUEST']._serialized_end=12745 + _globals['_GETSUMMARYRESPONSE']._serialized_start=12747 + _globals['_GETSUMMARYRESPONSE']._serialized_end=12810 + _globals['_GETSYSTEMMETRICSREQUEST']._serialized_start=12812 + _globals['_GETSYSTEMMETRICSREQUEST']._serialized_end=12883 + _globals['_SYSTEMMETRICSAMPLE']._serialized_start=12885 + _globals['_SYSTEMMETRICSAMPLE']._serialized_end=12967 + _globals['_SYSTEMMETRICSBUFFER']._serialized_start=12969 + _globals['_SYSTEMMETRICSBUFFER']._serialized_end=13042 + _globals['_GETSYSTEMMETRICSRESPONSE']._serialized_start=13045 + _globals['_GETSYSTEMMETRICSRESPONSE']._serialized_end=13247 + _globals['_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY']._serialized_start=13158 + _globals['_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY']._serialized_end=13247 + _globals['_STATUSREQUEST']._serialized_start=13249 + _globals['_STATUSREQUEST']._serialized_end=13310 + _globals['_STATUSRESPONSE']._serialized_start=13312 + _globals['_STATUSRESPONSE']._serialized_end=13353 + _globals['_STOPSTATUSREQUEST']._serialized_start=13355 + _globals['_STOPSTATUSREQUEST']._serialized_end=13420 + _globals['_STOPSTATUSRESPONSE']._serialized_start=13422 + _globals['_STOPSTATUSRESPONSE']._serialized_end=13467 + _globals['_NETWORKSTATUSREQUEST']._serialized_start=13469 + _globals['_NETWORKSTATUSREQUEST']._serialized_end=13537 + _globals['_NETWORKSTATUSRESPONSE']._serialized_start=13539 + _globals['_NETWORKSTATUSRESPONSE']._serialized_end=13619 + _globals['_HTTPRESPONSE']._serialized_start=13621 + _globals['_HTTPRESPONSE']._serialized_end=13689 + _globals['_INTERNALMESSAGESREQUEST']._serialized_start=13691 + _globals['_INTERNALMESSAGESREQUEST']._serialized_end=13762 + _globals['_INTERNALMESSAGESRESPONSE']._serialized_start=13764 + _globals['_INTERNALMESSAGESRESPONSE']._serialized_end=13842 + _globals['_INTERNALMESSAGES']._serialized_start=13844 + _globals['_INTERNALMESSAGES']._serialized_end=13879 + _globals['_POLLEXITREQUEST']._serialized_start=13881 + _globals['_POLLEXITREQUEST']._serialized_end=13944 + _globals['_POLLEXITRESPONSE']._serialized_start=13947 + _globals['_POLLEXITRESPONSE']._serialized_end=14192 + _globals['_OPERATIONSTATSREQUEST']._serialized_start=14194 + _globals['_OPERATIONSTATSREQUEST']._serialized_end=14263 + _globals['_OPERATIONSTATSRESPONSE']._serialized_start=14265 + _globals['_OPERATIONSTATSRESPONSE']._serialized_end=14346 + _globals['_OPERATIONSTATS']._serialized_start=14348 + _globals['_OPERATIONSTATS']._serialized_end=14437 + _globals['_OPERATION']._serialized_start=14440 + _globals['_OPERATION']._serialized_end=14575 + _globals['_SENDERMARKREQUEST']._serialized_start=14577 + _globals['_SENDERMARKREQUEST']._serialized_end=14596 + _globals['_SYNCFINISHREQUEST']._serialized_start=14598 + _globals['_SYNCFINISHREQUEST']._serialized_end=14617 + _globals['_SYNCRESPONSE']._serialized_start=14619 + _globals['_SYNCRESPONSE']._serialized_end=14688 + _globals['_SENDERREADREQUEST']._serialized_start=14690 + _globals['_SENDERREADREQUEST']._serialized_end=14753 + _globals['_STATUSREPORTREQUEST']._serialized_start=14755 + _globals['_STATUSREPORTREQUEST']._serialized_end=14864 + _globals['_SUMMARYRECORDREQUEST']._serialized_start=14866 + _globals['_SUMMARYRECORDREQUEST']._serialized_end=14936 + _globals['_TELEMETRYRECORDREQUEST']._serialized_start=14938 + _globals['_TELEMETRYRECORDREQUEST']._serialized_end=15014 + _globals['_SERVERINFOREQUEST']._serialized_start=15016 + _globals['_SERVERINFOREQUEST']._serialized_end=15081 + _globals['_SERVERINFORESPONSE']._serialized_start=15083 + _globals['_SERVERINFORESPONSE']._serialized_end=15207 + _globals['_SERVERMESSAGES']._serialized_start=15209 + _globals['_SERVERMESSAGES']._serialized_end=15270 + _globals['_SERVERMESSAGE']._serialized_start=15272 + _globals['_SERVERMESSAGE']._serialized_end=15373 + _globals['_FILECOUNTS']._serialized_start=15375 + _globals['_FILECOUNTS']._serialized_end=15474 + _globals['_FILEPUSHERSTATS']._serialized_start=15476 + _globals['_FILEPUSHERSTATS']._serialized_end=15561 + _globals['_FILESUPLOADED']._serialized_start=15563 + _globals['_FILESUPLOADED']._serialized_end=15593 + _globals['_FILETRANSFERINFOREQUEST']._serialized_start=15596 + _globals['_FILETRANSFERINFOREQUEST']._serialized_end=15840 + _globals['_FILETRANSFERINFOREQUEST_TRANSFERTYPE']._serialized_start=15800 + _globals['_FILETRANSFERINFOREQUEST_TRANSFERTYPE']._serialized_end=15840 + _globals['_LOCALINFO']._serialized_start=15842 + _globals['_LOCALINFO']._serialized_end=15891 + _globals['_SHUTDOWNREQUEST']._serialized_start=15893 + _globals['_SHUTDOWNREQUEST']._serialized_end=15956 + _globals['_SHUTDOWNRESPONSE']._serialized_start=15958 + _globals['_SHUTDOWNRESPONSE']._serialized_end=15976 + _globals['_ATTACHREQUEST']._serialized_start=15978 + _globals['_ATTACHREQUEST']._serialized_end=16058 + _globals['_ATTACHRESPONSE']._serialized_start=16060 + _globals['_ATTACHRESPONSE']._serialized_end=16158 + _globals['_TESTINJECTREQUEST']._serialized_start=16161 + _globals['_TESTINJECTREQUEST']._serialized_end=16502 + _globals['_TESTINJECTRESPONSE']._serialized_start=16504 + _globals['_TESTINJECTRESPONSE']._serialized_end=16524 + _globals['_HISTORYACTION']._serialized_start=16526 + _globals['_HISTORYACTION']._serialized_end=16556 + _globals['_PARTIALHISTORYREQUEST']._serialized_start=16559 + _globals['_PARTIALHISTORYREQUEST']._serialized_end=16761 + _globals['_PARTIALHISTORYRESPONSE']._serialized_start=16763 + _globals['_PARTIALHISTORYRESPONSE']._serialized_end=16787 + _globals['_SAMPLEDHISTORYREQUEST']._serialized_start=16789 + _globals['_SAMPLEDHISTORYREQUEST']._serialized_end=16858 + _globals['_SAMPLEDHISTORYITEM']._serialized_start=16860 + _globals['_SAMPLEDHISTORYITEM']._serialized_end=16955 + _globals['_SAMPLEDHISTORYRESPONSE']._serialized_start=16957 + _globals['_SAMPLEDHISTORYRESPONSE']._serialized_end=17031 + _globals['_RUNSTATUSREQUEST']._serialized_start=17033 + _globals['_RUNSTATUSREQUEST']._serialized_end=17097 + _globals['_RUNSTATUSRESPONSE']._serialized_start=17099 + _globals['_RUNSTATUSRESPONSE']._serialized_end=17219 + _globals['_RUNSTARTREQUEST']._serialized_start=17221 + _globals['_RUNSTARTREQUEST']._serialized_end=17324 + _globals['_RUNSTARTRESPONSE']._serialized_start=17326 + _globals['_RUNSTARTRESPONSE']._serialized_end=17344 + _globals['_RUNFINISHWITHOUTEXITREQUEST']._serialized_start=17346 + _globals['_RUNFINISHWITHOUTEXITREQUEST']._serialized_end=17421 + _globals['_RUNFINISHWITHOUTEXITRESPONSE']._serialized_start=17423 + _globals['_RUNFINISHWITHOUTEXITRESPONSE']._serialized_end=17453 + _globals['_CHECKVERSIONREQUEST']._serialized_start=17455 + _globals['_CHECKVERSIONREQUEST']._serialized_end=17547 + _globals['_CHECKVERSIONRESPONSE']._serialized_start=17549 + _globals['_CHECKVERSIONRESPONSE']._serialized_end=17642 + _globals['_JOBINFOREQUEST']._serialized_start=17644 + _globals['_JOBINFOREQUEST']._serialized_end=17706 + _globals['_JOBINFORESPONSE']._serialized_start=17708 + _globals['_JOBINFORESPONSE']._serialized_end=17762 + _globals['_LOGARTIFACTREQUEST']._serialized_start=17765 + _globals['_LOGARTIFACTREQUEST']._serialized_end=17924 + _globals['_LOGARTIFACTRESPONSE']._serialized_start=17926 + _globals['_LOGARTIFACTRESPONSE']._serialized_end=17991 + _globals['_DOWNLOADARTIFACTREQUEST']._serialized_start=17994 + _globals['_DOWNLOADARTIFACTREQUEST']._serialized_end=18184 + _globals['_DOWNLOADARTIFACTRESPONSE']._serialized_start=18186 + _globals['_DOWNLOADARTIFACTRESPONSE']._serialized_end=18235 + _globals['_KEEPALIVEREQUEST']._serialized_start=18237 + _globals['_KEEPALIVEREQUEST']._serialized_end=18301 + _globals['_KEEPALIVERESPONSE']._serialized_start=18303 + _globals['_KEEPALIVERESPONSE']._serialized_end=18322 + _globals['_ARTIFACTINFO']._serialized_start=18324 + _globals['_ARTIFACTINFO']._serialized_end=18437 + _globals['_GITINFO']._serialized_start=18439 + _globals['_GITINFO']._serialized_end=18480 + _globals['_GITSOURCE']._serialized_start=18483 + _globals['_GITSOURCE']._serialized_end=18618 + _globals['_IMAGESOURCE']._serialized_start=18620 + _globals['_IMAGESOURCE']._serialized_end=18648 + _globals['_SOURCE']._serialized_start=18651 + _globals['_SOURCE']._serialized_end=18791 + _globals['_JOBSOURCE']._serialized_start=18793 + _globals['_JOBSOURCE']._serialized_end=18900 + _globals['_PARTIALJOBARTIFACT']._serialized_start=18902 + _globals['_PARTIALJOBARTIFACT']._serialized_end=18988 + _globals['_USEARTIFACTRECORD']._serialized_start=18991 + _globals['_USEARTIFACTRECORD']._serialized_end=19148 + _globals['_USEARTIFACTRESULT']._serialized_start=19150 + _globals['_USEARTIFACTRESULT']._serialized_end=19169 + _globals['_CANCELREQUEST']._serialized_start=19171 + _globals['_CANCELREQUEST']._serialized_end=19253 + _globals['_CANCELRESPONSE']._serialized_start=19255 + _globals['_CANCELRESPONSE']._serialized_end=19271 + _globals['_DISKINFO']._serialized_start=19273 + _globals['_DISKINFO']._serialized_end=19312 + _globals['_MEMORYINFO']._serialized_start=19314 + _globals['_MEMORYINFO']._serialized_end=19341 + _globals['_CPUINFO']._serialized_start=19343 + _globals['_CPUINFO']._serialized_end=19390 + _globals['_APPLEINFO']._serialized_start=19393 + _globals['_APPLEINFO']._serialized_end=19547 + _globals['_GPUNVIDIAINFO']._serialized_start=19549 + _globals['_GPUNVIDIAINFO']._serialized_end=19656 + _globals['_GPUAMDINFO']._serialized_start=19659 + _globals['_GPUAMDINFO']._serialized_end=19924 + _globals['_TRAINIUMINFO']._serialized_start=19926 + _globals['_TRAINIUMINFO']._serialized_end=20036 + _globals['_TPUINFO']._serialized_start=20038 + _globals['_TPUINFO']._serialized_end=20119 + _globals['_COREWEAVEINFO']._serialized_start=20121 + _globals['_COREWEAVEINFO']._serialized_end=20190 + _globals['_ENVIRONMENTRECORD']._serialized_start=20193 + _globals['_ENVIRONMENTRECORD']._serialized_end=21385 + _globals['_ENVIRONMENTRECORD_DISKENTRY']._serialized_start=21270 + _globals['_ENVIRONMENTRECORD_DISKENTRY']._serialized_end=21339 + _globals['_ENVIRONMENTRECORD_SLURMENTRY']._serialized_start=21341 + _globals['_ENVIRONMENTRECORD_SLURMENTRY']._serialized_end=21385 + _globals['_PYTHONPACKAGESREQUEST']._serialized_start=21388 + _globals['_PYTHONPACKAGESREQUEST']._serialized_end=21529 + _globals['_PYTHONPACKAGESREQUEST_PYTHONPACKAGE']._serialized_start=21483 + _globals['_PYTHONPACKAGESREQUEST_PYTHONPACKAGE']._serialized_end=21529 + _globals['_JOBINPUTPATH']._serialized_start=21531 + _globals['_JOBINPUTPATH']._serialized_end=21559 + _globals['_JOBINPUTSOURCE']._serialized_start=21562 + _globals['_JOBINPUTSOURCE']._serialized_end=21776 + _globals['_JOBINPUTSOURCE_RUNCONFIGSOURCE']._serialized_start=21715 + _globals['_JOBINPUTSOURCE_RUNCONFIGSOURCE']._serialized_end=21732 + _globals['_JOBINPUTSOURCE_CONFIGFILESOURCE']._serialized_start=21734 + _globals['_JOBINPUTSOURCE_CONFIGFILESOURCE']._serialized_end=21766 + _globals['_JOBINPUTREQUEST']._serialized_start=21779 + _globals['_JOBINPUTREQUEST']._serialized_end=21978 + _globals['_SERVERFEATUREREQUEST']._serialized_start=21980 + _globals['_SERVERFEATUREREQUEST']._serialized_end=22096 + _globals['_SERVERFEATURERESPONSE']._serialized_start=22098 + _globals['_SERVERFEATURERESPONSE']._serialized_end=22173 + _globals['_SERVERFEATUREITEM']._serialized_start=22175 + _globals['_SERVERFEATUREITEM']._serialized_end=22225 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v5/wandb_settings_pb2.py b/lib/python3.12/site-packages/wandb/proto/v5/wandb_settings_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..3d858a96870f51b11422aabd6d9e83923f06d8ea --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v5/wandb_settings_pb2.py @@ -0,0 +1,48 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# source: wandb/proto/wandb_settings.proto +# Protobuf Python Version: 5.26.1 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import wrappers_pb2 as google_dot_protobuf_dot_wrappers__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n wandb/proto/wandb_settings.proto\x12\x0ewandb_internal\x1a\x1egoogle/protobuf/wrappers.proto\" \n\x0fListStringValue\x12\r\n\x05value\x18\x01 \x03(\t\"\x1d\n\x0cListIntValue\x12\r\n\x05value\x18\x01 \x03(\x05\"\x8a\x01\n\x17MapStringKeyStringValue\x12\x41\n\x05value\x18\x01 \x03(\x0b\x32\x32.wandb_internal.MapStringKeyStringValue.ValueEntry\x1a,\n\nValueEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\xcb\x01\n#MapStringKeyMapStringKeyStringValue\x12M\n\x05value\x18\x01 \x03(\x0b\x32>.wandb_internal.MapStringKeyMapStringKeyStringValue.ValueEntry\x1aU\n\nValueEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x36\n\x05value\x18\x02 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue:\x02\x38\x01\"\x9a\x01\n\x12OpenMetricsFilters\x12\x33\n\x08sequence\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValueH\x00\x12\x46\n\x07mapping\x18\x02 \x01(\x0b\x32\x33.wandb_internal.MapStringKeyMapStringKeyStringValueH\x00\x42\x07\n\x05value\"7\n\tRunMoment\x12\x0b\n\x03run\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01\x12\x0e\n\x06metric\x18\x03 \x01(\t\"\x99N\n\x08Settings\x12-\n\x07\x61pi_key\x18\x37 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x13identity_token_file\x18\xaa\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10\x63redentials_file\x18\xab\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x39\n\x14insecure_disable_ssl\x18\xb9\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08_offline\x18\x1e \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06x_sync\x18\x1f \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\tsync_file\x18\x86\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x07_shared\x18\xa2\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x06run_id\x18k \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07run_url\x18q \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07project\x18\x61 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06\x65ntity\x18\x45 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\x0corganization\x18\xbc\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0cx_start_time\x18) \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12.\n\x08root_dir\x18i \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07log_dir\x18U \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0clog_internal\x18V \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\tfiles_dir\x18\x46 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0bx_files_dir\x18\xb4\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0cignore_globs\x18N \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12.\n\x08\x62\x61se_url\x18\x39 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12=\n\x17x_file_stream_max_bytes\x18\xac\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x46\n\x1fx_file_stream_transmit_interval\x18\xaf\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x45\n\x14x_extra_http_headers\x18\x0e \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12=\n\x17x_file_stream_retry_max\x18\x93\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12K\n$x_file_stream_retry_wait_min_seconds\x18\x94\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12K\n$x_file_stream_retry_wait_max_seconds\x18\x95\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x43\n\x1dx_file_stream_timeout_seconds\x18\x0f \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x42\n\x1cx_file_stream_max_line_bytes\x18\xb2\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12?\n\x19x_file_transfer_retry_max\x18\x96\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12M\n&x_file_transfer_retry_wait_min_seconds\x18\x97\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12M\n&x_file_transfer_retry_wait_max_seconds\x18\x98\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x46\n\x1fx_file_transfer_timeout_seconds\x18\x99\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x39\n\x13x_graphql_retry_max\x18\x9a\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12G\n x_graphql_retry_wait_min_seconds\x18\x9b\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12G\n x_graphql_retry_wait_max_seconds\x18\x9c\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12@\n\x19x_graphql_timeout_seconds\x18\x9d\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x31\n\nhttp_proxy\x18\xa8\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0bhttps_proxy\x18\xa9\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12;\n\tx_proxies\x18\xc8\x01 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12-\n\x07program\x18_ \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0fprogram_relpath\x18` \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10_code_path_local\x18\xa3\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x0fprogram_abspath\x18\x9f\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x05_args\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12)\n\x03_os\x18 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06\x64ocker\x18\x43 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0cx_executable\x18\r \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07_python\x18\" \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\tcolab_url\x18\xa0\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x04host\x18M \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x08username\x18\x8d\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x05\x65mail\x18\x44 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06resume\x18\x66 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0bresume_from\x18\xa7\x01 \x01(\x0b\x32\x19.wandb_internal.RunMoment\x12-\n\tfork_from\x18\xa4\x01 \x01(\x0b\x32\x19.wandb_internal.RunMoment\x12\x38\n\x14\x64isable_job_creation\x18\x41 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\tsweep_url\x18\x83\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12;\n\x16x_disable_update_check\x18\xa5\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0ex_disable_meta\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\tsave_code\x18s \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x0b\x64isable_git\x18? \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x16x_disable_machine_info\x18\x9e\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0fx_disable_stats\x18\n \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x13x_stats_buffer_size\x18\xa1\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12@\n\x19x_stats_sampling_interval\x18\xae\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x30\n\x0bx_stats_pid\x18* \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12<\n\x12x_stats_disk_paths\x18\x92\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12H\n\"x_stats_neuron_monitor_config_path\x18. \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12<\n\x15x_stats_dcgm_exporter\x18\xbb\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12O\n\x1ex_stats_open_metrics_endpoints\x18/ \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12H\n\x1cx_stats_open_metrics_filters\x18\x30 \x01(\x0b\x32\".wandb_internal.OpenMetricsFilters\x12S\n!x_stats_open_metrics_http_headers\x18\xb8\x01 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12=\n\x16x_stats_gpu_device_ids\x18\xba\x01 \x01(\x0b\x32\x1c.wandb_internal.ListIntValue\x12\x37\n\x11x_stats_cpu_count\x18\xc2\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12?\n\x19x_stats_cpu_logical_count\x18\xc3\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x37\n\x11x_stats_gpu_count\x18\xc4\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x37\n\x10x_stats_gpu_type\x18\xc5\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12?\n\x1ax_stats_track_process_tree\x18\xc6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12.\n\x07x_label\x18\xb5\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\tx_primary\x18\xb6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12:\n\x15x_update_finish_state\x18\xb7\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12<\n\x17\x61llow_offline_artifacts\x18\xb1\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\x07\x63onsole\x18< \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x11\x63onsole_multipart\x18\xa6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x35\n\x10sync_tensorboard\x18\xb3\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x42\n\x1dx_server_side_derived_summary\x18\xbd\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x46\n!x_server_side_expand_glob_metrics\x18\xbe\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x16x_skip_transaction_log\x18\xbf\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12J\n#x_stats_coreweave_metadata_base_url\x18\xc0\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12J\n#x_stats_coreweave_metadata_endpoint\x18\xc1\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0b_aws_lambda\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0fx_cli_only_mode\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06_colab\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10x_disable_viewer\x18\x0b \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x15x_flow_control_custom\x18\x10 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x17x_flow_control_disabled\x18\x11 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12>\n\x18x_internal_check_process\x18\x12 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12,\n\x08_ipython\x18\x14 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08_jupyter\x18\x15 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x0ex_jupyter_root\x18\x16 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x07_kaggle\x18\x17 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12=\n\x18x_live_policy_rate_limit\x18\x18 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12<\n\x17x_live_policy_wait_time\x18\x19 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x30\n\x0bx_log_level\x18\x1a \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x35\n\x10x_network_buffer\x18\x1b \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12)\n\x05_noop\x18\x1c \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\t_notebook\x18\x1d \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\t_platform\x18! \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x38\n\x12x_runqueue_item_id\x18# \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x13x_save_requirements\x18% \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x13x_service_transport\x18& \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x34\n\x0ex_service_wait\x18\' \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x35\n\x0f_start_datetime\x18( \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\r_tmp_code_dir\x18\x31 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x08_windows\x18\x34 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10\x61llow_val_change\x18\x35 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\tanonymous\x18\x36 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12P\n\x1f\x61zure_account_url_to_access_key\x18\x38 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12.\n\x08\x63ode_dir\x18: \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0c\x63onfig_paths\x18; \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12\x30\n\ndeployment\x18= \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\x0c\x64isable_code\x18> \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x31\n\rdisable_hints\x18@ \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08\x64isabled\x18\x42 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12)\n\x05\x66orce\x18G \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\ngit_commit\x18H \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\ngit_remote\x18I \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x34\n\x0egit_remote_url\x18J \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08git_root\x18K \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x11heartbeat_seconds\x18L \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x32\n\x0cinit_timeout\x18O \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12,\n\x08is_local\x18P \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\njob_source\x18Q \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\rlabel_disable\x18R \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06launch\x18S \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x38\n\x12launch_config_path\x18T \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x14log_symlink_internal\x18W \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x10log_symlink_user\x18X \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08log_user\x18Y \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\rlogin_timeout\x18Z \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12*\n\x04mode\x18\\ \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\rnotebook_name\x18] \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\x0bproject_url\x18\x62 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12)\n\x05quiet\x18\x63 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12+\n\x07relogin\x18\x65 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0cresume_fname\x18g \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x07resumed\x18h \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\trun_group\x18j \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0crun_job_type\x18l \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08run_mode\x18m \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08run_name\x18n \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\trun_notes\x18o \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\x08run_tags\x18p \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12\x35\n\x11sagemaker_disable\x18r \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x35\n\x0fsettings_system\x18t \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x38\n\x12settings_workspace\x18u \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0bshow_colors\x18v \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12.\n\nshow_emoji\x18w \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x0bshow_errors\x18x \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\tshow_info\x18y \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x31\n\rshow_warnings\x18z \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06silent\x18{ \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0cstart_method\x18| \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x06strict\x18} \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0esummary_errors\x18~ \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x34\n\x0fsummary_timeout\x18\x7f \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x36\n\x10summary_warnings\x18\x80\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12/\n\x08sweep_id\x18\x81\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10sweep_param_path\x18\x82\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x07symlink\x18\x84\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x08sync_dir\x18\x85\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x13sync_symlink_latest\x18\x87\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12J\n%table_raise_on_max_row_limit_exceeded\x18\x8a\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x08timespec\x18\x8b\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x07tmp_dir\x18\x8c\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\twandb_dir\x18\x8e\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0ex_jupyter_name\x18\x8f\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0ex_jupyter_path\x18\x90\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x08job_name\x18\x91\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValueJ\x04\x08\x03\x10\x04J\x04\x08\x06\x10\x07J\x04\x08\x08\x10\tJ\x04\x08\t\x10\nJ\x04\x08\x0c\x10\rJ\x04\x08\x13\x10\x14J\x04\x08$\x10%J\x04\x08+\x10,J\x04\x08,\x10-J\x04\x08-\x10.J\x04\x08\x32\x10\x33J\x04\x08\x33\x10\x34J\x04\x08[\x10\\J\x04\x08^\x10_J\x04\x08\x64\x10\x65J\x06\x08\x88\x01\x10\x89\x01J\x06\x08\x89\x01\x10\x8a\x01J\x06\x08\xad\x01\x10\xae\x01J\x06\x08\xb0\x01\x10\xb1\x01\x42\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_settings_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._loaded_options = None + _globals['_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_options = b'8\001' + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._loaded_options = None + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_options = b'8\001' + _globals['_LISTSTRINGVALUE']._serialized_start=84 + _globals['_LISTSTRINGVALUE']._serialized_end=116 + _globals['_LISTINTVALUE']._serialized_start=118 + _globals['_LISTINTVALUE']._serialized_end=147 + _globals['_MAPSTRINGKEYSTRINGVALUE']._serialized_start=150 + _globals['_MAPSTRINGKEYSTRINGVALUE']._serialized_end=288 + _globals['_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_start=244 + _globals['_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_end=288 + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE']._serialized_start=291 + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE']._serialized_end=494 + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_start=409 + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_end=494 + _globals['_OPENMETRICSFILTERS']._serialized_start=497 + _globals['_OPENMETRICSFILTERS']._serialized_end=651 + _globals['_RUNMOMENT']._serialized_start=653 + _globals['_RUNMOMENT']._serialized_end=708 + _globals['_SETTINGS']._serialized_start=711 + _globals['_SETTINGS']._serialized_end=10720 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_base_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_base_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08bc9a9d9a4261b267d5d54418eecd81b2ab196b Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_base_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_internal_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_internal_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..726ea908c357a237918d4d01561999cf6fc3eb4e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_internal_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_server_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_server_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c020b34d55d8a3bc6e277b876a59fbd2d727b22e Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_server_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_settings_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_settings_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a98b5b467da6315369c823187c33f15e0a7ab961 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_settings_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_telemetry_pb2.cpython-312.pyc b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_telemetry_pb2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9e29cd8b1df1899e303f583627d97bdd19d9a7aa Binary files /dev/null and b/lib/python3.12/site-packages/wandb/proto/v6/__pycache__/wandb_telemetry_pb2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/proto/v6/wandb_base_pb2.py b/lib/python3.12/site-packages/wandb/proto/v6/wandb_base_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..3c9881ffc9d5f84fec9b4ad70f9fc58642d9f471 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v6/wandb_base_pb2.py @@ -0,0 +1,41 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: wandb/proto/wandb_base.proto +# Protobuf Python Version: 6.30.0 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 30, + 0, + '', + 'wandb/proto/wandb_base.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1cwandb/proto/wandb_base.proto\x12\x0ewandb_internal\"6\n\x0b_RecordInfo\x12\x11\n\tstream_id\x18\x01 \x01(\t\x12\x14\n\x0c_tracelog_id\x18\x64 \x01(\t\"!\n\x0c_RequestInfo\x12\x11\n\tstream_id\x18\x01 \x01(\t\"#\n\x0b_ResultInfo\x12\x14\n\x0c_tracelog_id\x18\x64 \x01(\tB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_base_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['__RECORDINFO']._serialized_start=48 + _globals['__RECORDINFO']._serialized_end=102 + _globals['__REQUESTINFO']._serialized_start=104 + _globals['__REQUESTINFO']._serialized_end=137 + _globals['__RESULTINFO']._serialized_start=139 + _globals['__RESULTINFO']._serialized_end=174 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v6/wandb_internal_pb2.py b/lib/python3.12/site-packages/wandb/proto/v6/wandb_internal_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..04f1900221a7bc30245f9cbe74425c9ad7a55a72 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v6/wandb_internal_pb2.py @@ -0,0 +1,393 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: wandb/proto/wandb_internal.proto +# Protobuf Python Version: 6.30.0 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 30, + 0, + '', + 'wandb/proto/wandb_internal.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import empty_pb2 as google_dot_protobuf_dot_empty__pb2 +from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2 +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 +from wandb.proto import wandb_telemetry_pb2 as wandb_dot_proto_dot_wandb__telemetry__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n wandb/proto/wandb_internal.proto\x12\x0ewandb_internal\x1a\x1bgoogle/protobuf/empty.proto\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1cwandb/proto/wandb_base.proto\x1a!wandb/proto/wandb_telemetry.proto\"\xcf\t\n\x06Record\x12\x0b\n\x03num\x18\x01 \x01(\x03\x12\x30\n\x07history\x18\x02 \x01(\x0b\x32\x1d.wandb_internal.HistoryRecordH\x00\x12\x30\n\x07summary\x18\x03 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecordH\x00\x12.\n\x06output\x18\x04 \x01(\x0b\x32\x1c.wandb_internal.OutputRecordH\x00\x12.\n\x06\x63onfig\x18\x05 \x01(\x0b\x32\x1c.wandb_internal.ConfigRecordH\x00\x12,\n\x05\x66iles\x18\x06 \x01(\x0b\x32\x1b.wandb_internal.FilesRecordH\x00\x12,\n\x05stats\x18\x07 \x01(\x0b\x32\x1b.wandb_internal.StatsRecordH\x00\x12\x32\n\x08\x61rtifact\x18\x08 \x01(\x0b\x32\x1e.wandb_internal.ArtifactRecordH\x00\x12,\n\x08tbrecord\x18\t \x01(\x0b\x32\x18.wandb_internal.TBRecordH\x00\x12,\n\x05\x61lert\x18\n \x01(\x0b\x32\x1b.wandb_internal.AlertRecordH\x00\x12\x34\n\ttelemetry\x18\x0b \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecordH\x00\x12.\n\x06metric\x18\x0c \x01(\x0b\x32\x1c.wandb_internal.MetricRecordH\x00\x12\x35\n\noutput_raw\x18\r \x01(\x0b\x32\x1f.wandb_internal.OutputRawRecordH\x00\x12(\n\x03run\x18\x11 \x01(\x0b\x32\x19.wandb_internal.RunRecordH\x00\x12-\n\x04\x65xit\x18\x12 \x01(\x0b\x32\x1d.wandb_internal.RunExitRecordH\x00\x12,\n\x05\x66inal\x18\x14 \x01(\x0b\x32\x1b.wandb_internal.FinalRecordH\x00\x12.\n\x06header\x18\x15 \x01(\x0b\x32\x1c.wandb_internal.HeaderRecordH\x00\x12.\n\x06\x66ooter\x18\x16 \x01(\x0b\x32\x1c.wandb_internal.FooterRecordH\x00\x12\x39\n\npreempting\x18\x17 \x01(\x0b\x32#.wandb_internal.RunPreemptingRecordH\x00\x12\x34\n\x12noop_link_artifact\x18\x18 \x01(\x0b\x32\x16.google.protobuf.EmptyH\x00\x12\x39\n\x0cuse_artifact\x18\x19 \x01(\x0b\x32!.wandb_internal.UseArtifactRecordH\x00\x12\x38\n\x0b\x65nvironment\x18\x1a \x01(\x0b\x32!.wandb_internal.EnvironmentRecordH\x00\x12*\n\x07request\x18\x64 \x01(\x0b\x32\x17.wandb_internal.RequestH\x00\x12(\n\x07\x63ontrol\x18\x10 \x01(\x0b\x32\x17.wandb_internal.Control\x12\x0c\n\x04uuid\x18\x13 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfoB\r\n\x0brecord_type\"\xa8\x01\n\x07\x43ontrol\x12\x10\n\x08req_resp\x18\x01 \x01(\x08\x12\r\n\x05local\x18\x02 \x01(\x08\x12\x10\n\x08relay_id\x18\x03 \x01(\t\x12\x14\n\x0cmailbox_slot\x18\x04 \x01(\t\x12\x13\n\x0b\x61lways_send\x18\x05 \x01(\x08\x12\x14\n\x0c\x66low_control\x18\x06 \x01(\x08\x12\x12\n\nend_offset\x18\x07 \x01(\x03\x12\x15\n\rconnection_id\x18\x08 \x01(\t\"\xf3\x03\n\x06Result\x12\x35\n\nrun_result\x18\x11 \x01(\x0b\x32\x1f.wandb_internal.RunUpdateResultH\x00\x12\x34\n\x0b\x65xit_result\x18\x12 \x01(\x0b\x32\x1d.wandb_internal.RunExitResultH\x00\x12\x33\n\nlog_result\x18\x14 \x01(\x0b\x32\x1d.wandb_internal.HistoryResultH\x00\x12\x37\n\x0esummary_result\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.SummaryResultH\x00\x12\x35\n\routput_result\x18\x16 \x01(\x0b\x32\x1c.wandb_internal.OutputResultH\x00\x12\x35\n\rconfig_result\x18\x17 \x01(\x0b\x32\x1c.wandb_internal.ConfigResultH\x00\x12,\n\x08response\x18\x64 \x01(\x0b\x32\x18.wandb_internal.ResponseH\x00\x12(\n\x07\x63ontrol\x18\x10 \x01(\x0b\x32\x17.wandb_internal.Control\x12\x0c\n\x04uuid\x18\x18 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._ResultInfoB\r\n\x0bresult_type\":\n\x0b\x46inalRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"b\n\x0bVersionInfo\x12\x10\n\x08producer\x18\x01 \x01(\t\x12\x14\n\x0cmin_consumer\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"n\n\x0cHeaderRecord\x12\x31\n\x0cversion_info\x18\x01 \x01(\x0b\x32\x1b.wandb_internal.VersionInfo\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\";\n\x0c\x46ooterRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"9\n\x0b\x42ranchPoint\x12\x0b\n\x03run\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01\x12\x0e\n\x06metric\x18\x03 \x01(\t\"\x91\x05\n\tRunRecord\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0e\n\x06\x65ntity\x18\x02 \x01(\t\x12\x0f\n\x07project\x18\x03 \x01(\t\x12,\n\x06\x63onfig\x18\x04 \x01(\x0b\x32\x1c.wandb_internal.ConfigRecord\x12.\n\x07summary\x18\x05 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecord\x12\x11\n\trun_group\x18\x06 \x01(\t\x12\x10\n\x08job_type\x18\x07 \x01(\t\x12\x14\n\x0c\x64isplay_name\x18\x08 \x01(\t\x12\r\n\x05notes\x18\t \x01(\t\x12\x0c\n\x04tags\x18\n \x03(\t\x12\x30\n\x08settings\x18\x0b \x01(\x0b\x32\x1e.wandb_internal.SettingsRecord\x12\x10\n\x08sweep_id\x18\x0c \x01(\t\x12\x0c\n\x04host\x18\r \x01(\t\x12\x15\n\rstarting_step\x18\x0e \x01(\x03\x12\x12\n\nstorage_id\x18\x10 \x01(\t\x12.\n\nstart_time\x18\x11 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0f\n\x07resumed\x18\x12 \x01(\x08\x12\x32\n\ttelemetry\x18\x13 \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecord\x12\x0f\n\x07runtime\x18\x14 \x01(\x05\x12*\n\x03git\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.GitRepoRecord\x12\x0e\n\x06\x66orked\x18\x16 \x01(\x08\x12\x31\n\x0c\x62ranch_point\x18\x17 \x01(\x0b\x32\x1b.wandb_internal.BranchPoint\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\";\n\rGitRepoRecord\x12\x1a\n\nremote_url\x18\x01 \x01(\tR\x06remote\x12\x0e\n\x06\x63ommit\x18\x02 \x01(\t\"c\n\x0fRunUpdateResult\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"\xac\x01\n\tErrorInfo\x12\x0f\n\x07message\x18\x01 \x01(\t\x12\x31\n\x04\x63ode\x18\x02 \x01(\x0e\x32#.wandb_internal.ErrorInfo.ErrorCode\"[\n\tErrorCode\x12\x0b\n\x07UNKNOWN\x10\x00\x12\x11\n\rCOMMUNICATION\x10\x01\x12\x12\n\x0e\x41UTHENTICATION\x10\x02\x12\t\n\x05USAGE\x10\x03\x12\x0f\n\x0bUNSUPPORTED\x10\x04\"`\n\rRunExitRecord\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12\x0f\n\x07runtime\x18\x02 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x0f\n\rRunExitResult\"B\n\x13RunPreemptingRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x15\n\x13RunPreemptingResult\"i\n\x0eSettingsRecord\x12*\n\x04item\x18\x01 \x03(\x0b\x32\x1c.wandb_internal.SettingsItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"/\n\x0cSettingsItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x1a\n\x0bHistoryStep\x12\x0b\n\x03num\x18\x01 \x01(\x03\"\x92\x01\n\rHistoryRecord\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.HistoryItem\x12)\n\x04step\x18\x02 \x01(\x0b\x32\x1b.wandb_internal.HistoryStep\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\x0bHistoryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0f\n\rHistoryResult\"\xdc\x01\n\x0cOutputRecord\x12<\n\x0boutput_type\x18\x01 \x01(\x0e\x32\'.wandb_internal.OutputRecord.OutputType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0c\n\x04line\x18\x03 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"$\n\nOutputType\x12\n\n\x06STDERR\x10\x00\x12\n\n\x06STDOUT\x10\x01\"\x0e\n\x0cOutputResult\"\xe2\x01\n\x0fOutputRawRecord\x12?\n\x0boutput_type\x18\x01 \x01(\x0e\x32*.wandb_internal.OutputRawRecord.OutputType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x0c\n\x04line\x18\x03 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"$\n\nOutputType\x12\n\n\x06STDERR\x10\x00\x12\n\n\x06STDOUT\x10\x01\"\x11\n\x0fOutputRawResult\"\xb4\x03\n\x0cMetricRecord\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x11\n\tglob_name\x18\x02 \x01(\t\x12\x13\n\x0bstep_metric\x18\x04 \x01(\t\x12\x19\n\x11step_metric_index\x18\x05 \x01(\x05\x12.\n\x07options\x18\x06 \x01(\x0b\x32\x1d.wandb_internal.MetricOptions\x12.\n\x07summary\x18\x07 \x01(\x0b\x32\x1d.wandb_internal.MetricSummary\x12\x35\n\x04goal\x18\x08 \x01(\x0e\x32\'.wandb_internal.MetricRecord.MetricGoal\x12/\n\x08_control\x18\t \x01(\x0b\x32\x1d.wandb_internal.MetricControl\x12\x1a\n\x12\x65xpanded_from_glob\x18\n \x01(\x08\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\nMetricGoal\x12\x0e\n\nGOAL_UNSET\x10\x00\x12\x11\n\rGOAL_MINIMIZE\x10\x01\x12\x11\n\rGOAL_MAXIMIZE\x10\x02\"\x0e\n\x0cMetricResult\"C\n\rMetricOptions\x12\x11\n\tstep_sync\x18\x01 \x01(\x08\x12\x0e\n\x06hidden\x18\x02 \x01(\x08\x12\x0f\n\x07\x64\x65\x66ined\x18\x03 \x01(\x08\"\"\n\rMetricControl\x12\x11\n\toverwrite\x18\x01 \x01(\x08\"o\n\rMetricSummary\x12\x0b\n\x03min\x18\x01 \x01(\x08\x12\x0b\n\x03max\x18\x02 \x01(\x08\x12\x0c\n\x04mean\x18\x03 \x01(\x08\x12\x0c\n\x04\x62\x65st\x18\x04 \x01(\x08\x12\x0c\n\x04last\x18\x05 \x01(\x08\x12\x0c\n\x04none\x18\x06 \x01(\x08\x12\x0c\n\x04\x63opy\x18\x07 \x01(\x08\"\x93\x01\n\x0c\x43onfigRecord\x12*\n\x06update\x18\x01 \x03(\x0b\x32\x1a.wandb_internal.ConfigItem\x12*\n\x06remove\x18\x02 \x03(\x0b\x32\x1a.wandb_internal.ConfigItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"A\n\nConfigItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0e\n\x0c\x43onfigResult\"\x96\x01\n\rSummaryRecord\x12+\n\x06update\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\x12+\n\x06remove\x18\x02 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"B\n\x0bSummaryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\x0f\n\rSummaryResult\"d\n\x0b\x46ilesRecord\x12(\n\x05\x66iles\x18\x01 \x03(\x0b\x32\x19.wandb_internal.FilesItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\xec\x01\n\tFilesItem\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x34\n\x06policy\x18\x02 \x01(\x0e\x32$.wandb_internal.FilesItem.PolicyType\x12\x30\n\x04type\x18\x03 \x01(\x0e\x32\".wandb_internal.FilesItem.FileType\"(\n\nPolicyType\x12\x07\n\x03NOW\x10\x00\x12\x07\n\x03\x45ND\x10\x01\x12\x08\n\x04LIVE\x10\x02\"9\n\x08\x46ileType\x12\t\n\x05OTHER\x10\x00\x12\t\n\x05WANDB\x10\x01\x12\t\n\x05MEDIA\x10\x02\x12\x0c\n\x08\x41RTIFACT\x10\x03J\x04\x08\x10\x10\x11\"\r\n\x0b\x46ilesResult\"\xe6\x01\n\x0bStatsRecord\x12\x39\n\nstats_type\x18\x01 \x01(\x0e\x32%.wandb_internal.StatsRecord.StatsType\x12-\n\ttimestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\'\n\x04item\x18\x03 \x03(\x0b\x32\x19.wandb_internal.StatsItem\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x17\n\tStatsType\x12\n\n\x06SYSTEM\x10\x00\",\n\tStatsItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x10 \x01(\t\"\xe7\x03\n\x0e\x41rtifactRecord\x12\x0e\n\x06run_id\x18\x01 \x01(\t\x12\x0f\n\x07project\x18\x02 \x01(\t\x12\x0e\n\x06\x65ntity\x18\x03 \x01(\t\x12\x0c\n\x04type\x18\x04 \x01(\t\x12\x0c\n\x04name\x18\x05 \x01(\t\x12\x0e\n\x06\x64igest\x18\x06 \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x07 \x01(\t\x12\x10\n\x08metadata\x18\x08 \x01(\t\x12\x14\n\x0cuser_created\x18\t \x01(\x08\x12\x18\n\x10use_after_commit\x18\n \x01(\x08\x12\x0f\n\x07\x61liases\x18\x0b \x03(\t\x12\x32\n\x08manifest\x18\x0c \x01(\x0b\x32 .wandb_internal.ArtifactManifest\x12\x16\n\x0e\x64istributed_id\x18\r \x01(\t\x12\x10\n\x08\x66inalize\x18\x0e \x01(\x08\x12\x11\n\tclient_id\x18\x0f \x01(\t\x12\x1a\n\x12sequence_client_id\x18\x10 \x01(\t\x12\x0f\n\x07\x62\x61se_id\x18\x11 \x01(\t\x12\x1c\n\x14ttl_duration_seconds\x18\x12 \x01(\x03\x12\x0c\n\x04tags\x18\x13 \x03(\t\x12\x19\n\x11incremental_beta1\x18\x64 \x01(\x08\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\xd8\x01\n\x10\x41rtifactManifest\x12\x0f\n\x07version\x18\x01 \x01(\x05\x12\x16\n\x0estorage_policy\x18\x02 \x01(\t\x12\x46\n\x15storage_policy_config\x18\x03 \x03(\x0b\x32\'.wandb_internal.StoragePolicyConfigItem\x12\x37\n\x08\x63ontents\x18\x04 \x03(\x0b\x32%.wandb_internal.ArtifactManifestEntry\x12\x1a\n\x12manifest_file_path\x18\x05 \x01(\t\"\xcf\x01\n\x15\x41rtifactManifestEntry\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x0e\n\x06\x64igest\x18\x02 \x01(\t\x12\x0b\n\x03ref\x18\x03 \x01(\t\x12\x0c\n\x04size\x18\x04 \x01(\x03\x12\x10\n\x08mimetype\x18\x05 \x01(\t\x12\x12\n\nlocal_path\x18\x06 \x01(\t\x12\x19\n\x11\x62irth_artifact_id\x18\x07 \x01(\t\x12\x12\n\nskip_cache\x18\x08 \x01(\x08\x12(\n\x05\x65xtra\x18\x10 \x03(\x0b\x32\x19.wandb_internal.ExtraItem\",\n\tExtraItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x02 \x01(\t\":\n\x17StoragePolicyConfigItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nvalue_json\x18\x02 \x01(\t\"\x10\n\x0e\x41rtifactResult\"\x14\n\x12LinkArtifactResult\"\xf0\x01\n\x13LinkArtifactRequest\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x11\n\tserver_id\x18\x02 \x01(\t\x12\x16\n\x0eportfolio_name\x18\x03 \x01(\t\x12\x18\n\x10portfolio_entity\x18\x04 \x01(\t\x12\x19\n\x11portfolio_project\x18\x05 \x01(\t\x12\x19\n\x11portfolio_aliases\x18\x06 \x03(\t\x12\x1e\n\x16portfolio_organization\x18\x07 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"[\n\x14LinkArtifactResponse\x12\x15\n\rerror_message\x18\x01 \x01(\t\x12\x1a\n\rversion_index\x18\x02 \x01(\x05H\x00\x88\x01\x01\x42\x10\n\x0e_version_index\"h\n\x08TBRecord\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\x12\x0f\n\x07log_dir\x18\x01 \x01(\t\x12\x10\n\x08root_dir\x18\x03 \x01(\t\x12\x0c\n\x04save\x18\x02 \x01(\x08\"\n\n\x08TBResult\"}\n\x0b\x41lertRecord\x12\r\n\x05title\x18\x01 \x01(\t\x12\x0c\n\x04text\x18\x02 \x01(\t\x12\r\n\x05level\x18\x03 \x01(\t\x12\x15\n\rwait_duration\x18\x04 \x01(\x03\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\r\n\x0b\x41lertResult\"\xb3\x11\n\x07Request\x12\x38\n\x0bstop_status\x18\x01 \x01(\x0b\x32!.wandb_internal.StopStatusRequestH\x00\x12>\n\x0enetwork_status\x18\x02 \x01(\x0b\x32$.wandb_internal.NetworkStatusRequestH\x00\x12-\n\x05\x64\x65\x66\x65r\x18\x03 \x01(\x0b\x32\x1c.wandb_internal.DeferRequestH\x00\x12\x38\n\x0bget_summary\x18\x04 \x01(\x0b\x32!.wandb_internal.GetSummaryRequestH\x00\x12-\n\x05login\x18\x05 \x01(\x0b\x32\x1c.wandb_internal.LoginRequestH\x00\x12-\n\x05pause\x18\x06 \x01(\x0b\x32\x1c.wandb_internal.PauseRequestH\x00\x12/\n\x06resume\x18\x07 \x01(\x0b\x32\x1d.wandb_internal.ResumeRequestH\x00\x12\x34\n\tpoll_exit\x18\x08 \x01(\x0b\x32\x1f.wandb_internal.PollExitRequestH\x00\x12@\n\x0fsampled_history\x18\t \x01(\x0b\x32%.wandb_internal.SampledHistoryRequestH\x00\x12@\n\x0fpartial_history\x18\n \x01(\x0b\x32%.wandb_internal.PartialHistoryRequestH\x00\x12\x34\n\trun_start\x18\x0b \x01(\x0b\x32\x1f.wandb_internal.RunStartRequestH\x00\x12<\n\rcheck_version\x18\x0c \x01(\x0b\x32#.wandb_internal.CheckVersionRequestH\x00\x12:\n\x0clog_artifact\x18\r \x01(\x0b\x32\".wandb_internal.LogArtifactRequestH\x00\x12\x44\n\x11\x64ownload_artifact\x18\x0e \x01(\x0b\x32\'.wandb_internal.DownloadArtifactRequestH\x00\x12\x35\n\tkeepalive\x18\x11 \x01(\x0b\x32 .wandb_internal.KeepaliveRequestH\x00\x12>\n\x0eserver_feature\x18\x12 \x01(\x0b\x32$.wandb_internal.ServerFeatureRequestH\x00\x12\x36\n\nrun_status\x18\x14 \x01(\x0b\x32 .wandb_internal.RunStatusRequestH\x00\x12/\n\x06\x63\x61ncel\x18\x15 \x01(\x0b\x32\x1d.wandb_internal.CancelRequestH\x00\x12\x44\n\x11internal_messages\x18\x17 \x01(\x0b\x32\'.wandb_internal.InternalMessagesRequestH\x00\x12@\n\x0fpython_packages\x18\x18 \x01(\x0b\x32%.wandb_internal.PythonPackagesRequestH\x00\x12\x33\n\x08shutdown\x18@ \x01(\x0b\x32\x1f.wandb_internal.ShutdownRequestH\x00\x12/\n\x06\x61ttach\x18\x41 \x01(\x0b\x32\x1d.wandb_internal.AttachRequestH\x00\x12/\n\x06status\x18\x42 \x01(\x0b\x32\x1d.wandb_internal.StatusRequestH\x00\x12\x38\n\x0bserver_info\x18\x43 \x01(\x0b\x32!.wandb_internal.ServerInfoRequestH\x00\x12\x38\n\x0bsender_mark\x18\x44 \x01(\x0b\x32!.wandb_internal.SenderMarkRequestH\x00\x12\x38\n\x0bsender_read\x18\x45 \x01(\x0b\x32!.wandb_internal.SenderReadRequestH\x00\x12<\n\rstatus_report\x18\x46 \x01(\x0b\x32#.wandb_internal.StatusReportRequestH\x00\x12>\n\x0esummary_record\x18G \x01(\x0b\x32$.wandb_internal.SummaryRecordRequestH\x00\x12\x42\n\x10telemetry_record\x18H \x01(\x0b\x32&.wandb_internal.TelemetryRecordRequestH\x00\x12\x32\n\x08job_info\x18I \x01(\x0b\x32\x1e.wandb_internal.JobInfoRequestH\x00\x12\x45\n\x12get_system_metrics\x18J \x01(\x0b\x32\'.wandb_internal.GetSystemMetricsRequestH\x00\x12\x34\n\tjob_input\x18M \x01(\x0b\x32\x1f.wandb_internal.JobInputRequestH\x00\x12<\n\rlink_artifact\x18N \x01(\x0b\x32#.wandb_internal.LinkArtifactRequestH\x00\x12N\n\x17run_finish_without_exit\x18O \x01(\x0b\x32+.wandb_internal.RunFinishWithoutExitRequestH\x00\x12\x38\n\x0bsync_finish\x18Q \x01(\x0b\x32!.wandb_internal.SyncFinishRequestH\x00\x12;\n\noperations\x18R \x01(\x0b\x32%.wandb_internal.OperationStatsRequestH\x00\x12\x39\n\x0btest_inject\x18\xe8\x07 \x01(\x0b\x32!.wandb_internal.TestInjectRequestH\x00\x42\x0e\n\x0crequest_typeJ\x04\x08\x16\x10\x17J\x04\x08K\x10LJ\x04\x08L\x10MJ\x04\x08P\x10Q\"\x9b\x0e\n\x08Response\x12?\n\x12keepalive_response\x18\x12 \x01(\x0b\x32!.wandb_internal.KeepaliveResponseH\x00\x12\x42\n\x14stop_status_response\x18\x13 \x01(\x0b\x32\".wandb_internal.StopStatusResponseH\x00\x12H\n\x17network_status_response\x18\x14 \x01(\x0b\x32%.wandb_internal.NetworkStatusResponseH\x00\x12\x37\n\x0elogin_response\x18\x18 \x01(\x0b\x32\x1d.wandb_internal.LoginResponseH\x00\x12\x42\n\x14get_summary_response\x18\x19 \x01(\x0b\x32\".wandb_internal.GetSummaryResponseH\x00\x12>\n\x12poll_exit_response\x18\x1a \x01(\x0b\x32 .wandb_internal.PollExitResponseH\x00\x12J\n\x18sampled_history_response\x18\x1b \x01(\x0b\x32&.wandb_internal.SampledHistoryResponseH\x00\x12>\n\x12run_start_response\x18\x1c \x01(\x0b\x32 .wandb_internal.RunStartResponseH\x00\x12\x46\n\x16\x63heck_version_response\x18\x1d \x01(\x0b\x32$.wandb_internal.CheckVersionResponseH\x00\x12\x44\n\x15log_artifact_response\x18\x1e \x01(\x0b\x32#.wandb_internal.LogArtifactResponseH\x00\x12N\n\x1a\x64ownload_artifact_response\x18\x1f \x01(\x0b\x32(.wandb_internal.DownloadArtifactResponseH\x00\x12H\n\x17server_feature_response\x18 \x01(\x0b\x32%.wandb_internal.ServerFeatureResponseH\x00\x12@\n\x13run_status_response\x18# \x01(\x0b\x32!.wandb_internal.RunStatusResponseH\x00\x12\x39\n\x0f\x63\x61ncel_response\x18$ \x01(\x0b\x32\x1e.wandb_internal.CancelResponseH\x00\x12N\n\x1ainternal_messages_response\x18% \x01(\x0b\x32(.wandb_internal.InternalMessagesResponseH\x00\x12=\n\x11shutdown_response\x18@ \x01(\x0b\x32 .wandb_internal.ShutdownResponseH\x00\x12\x39\n\x0f\x61ttach_response\x18\x41 \x01(\x0b\x32\x1e.wandb_internal.AttachResponseH\x00\x12\x39\n\x0fstatus_response\x18\x42 \x01(\x0b\x32\x1e.wandb_internal.StatusResponseH\x00\x12\x42\n\x14server_info_response\x18\x43 \x01(\x0b\x32\".wandb_internal.ServerInfoResponseH\x00\x12<\n\x11job_info_response\x18\x44 \x01(\x0b\x32\x1f.wandb_internal.JobInfoResponseH\x00\x12O\n\x1bget_system_metrics_response\x18\x45 \x01(\x0b\x32(.wandb_internal.GetSystemMetricsResponseH\x00\x12\x46\n\x16link_artifact_response\x18G \x01(\x0b\x32$.wandb_internal.LinkArtifactResponseH\x00\x12\x35\n\rsync_response\x18\x46 \x01(\x0b\x32\x1c.wandb_internal.SyncResponseH\x00\x12X\n run_finish_without_exit_response\x18H \x01(\x0b\x32,.wandb_internal.RunFinishWithoutExitResponseH\x00\x12\x45\n\x13operations_response\x18J \x01(\x0b\x32&.wandb_internal.OperationStatsResponseH\x00\x12\x43\n\x14test_inject_response\x18\xe8\x07 \x01(\x0b\x32\".wandb_internal.TestInjectResponseH\x00\x42\x0f\n\rresponse_typeJ\x04\x08I\x10J\"\xc0\x02\n\x0c\x44\x65\x66\x65rRequest\x12\x36\n\x05state\x18\x01 \x01(\x0e\x32\'.wandb_internal.DeferRequest.DeferState\"\xf7\x01\n\nDeferState\x12\t\n\x05\x42\x45GIN\x10\x00\x12\r\n\tFLUSH_RUN\x10\x01\x12\x0f\n\x0b\x46LUSH_STATS\x10\x02\x12\x19\n\x15\x46LUSH_PARTIAL_HISTORY\x10\x03\x12\x0c\n\x08\x46LUSH_TB\x10\x04\x12\r\n\tFLUSH_SUM\x10\x05\x12\x13\n\x0f\x46LUSH_DEBOUNCER\x10\x06\x12\x10\n\x0c\x46LUSH_OUTPUT\x10\x07\x12\r\n\tFLUSH_JOB\x10\x08\x12\r\n\tFLUSH_DIR\x10\t\x12\x0c\n\x08\x46LUSH_FP\x10\n\x12\x0b\n\x07JOIN_FP\x10\x0b\x12\x0c\n\x08\x46LUSH_FS\x10\x0c\x12\x0f\n\x0b\x46LUSH_FINAL\x10\r\x12\x07\n\x03\x45ND\x10\x0e\"<\n\x0cPauseRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x0f\n\rPauseResponse\"=\n\rResumeRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x10\n\x0eResumeResponse\"M\n\x0cLoginRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"&\n\rLoginResponse\x12\x15\n\ractive_entity\x18\x01 \x01(\t\"A\n\x11GetSummaryRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"?\n\x12GetSummaryResponse\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.SummaryItem\"G\n\x17GetSystemMetricsRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"R\n\x12SystemMetricSample\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\r\n\x05value\x18\x02 \x01(\x02\"I\n\x13SystemMetricsBuffer\x12\x32\n\x06record\x18\x01 \x03(\x0b\x32\".wandb_internal.SystemMetricSample\"\xca\x01\n\x18GetSystemMetricsResponse\x12S\n\x0esystem_metrics\x18\x01 \x03(\x0b\x32;.wandb_internal.GetSystemMetricsResponse.SystemMetricsEntry\x1aY\n\x12SystemMetricsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x32\n\x05value\x18\x02 \x01(\x0b\x32#.wandb_internal.SystemMetricsBuffer:\x02\x38\x01\"=\n\rStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\")\n\x0eStatusResponse\x12\x17\n\x0frun_should_stop\x18\x01 \x01(\x08\"A\n\x11StopStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"-\n\x12StopStatusResponse\x12\x17\n\x0frun_should_stop\x18\x01 \x01(\x08\"D\n\x14NetworkStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"P\n\x15NetworkStatusResponse\x12\x37\n\x11network_responses\x18\x01 \x03(\x0b\x32\x1c.wandb_internal.HttpResponse\"D\n\x0cHttpResponse\x12\x18\n\x10http_status_code\x18\x01 \x01(\x05\x12\x1a\n\x12http_response_text\x18\x02 \x01(\t\"G\n\x17InternalMessagesRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"N\n\x18InternalMessagesResponse\x12\x32\n\x08messages\x18\x01 \x01(\x0b\x32 .wandb_internal.InternalMessages\"#\n\x10InternalMessages\x12\x0f\n\x07warning\x18\x01 \x03(\t\"?\n\x0fPollExitRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\xf5\x01\n\x10PollExitResponse\x12\x0c\n\x04\x64one\x18\x01 \x01(\x08\x12\x32\n\x0b\x65xit_result\x18\x02 \x01(\x0b\x32\x1d.wandb_internal.RunExitResult\x12\x35\n\x0cpusher_stats\x18\x03 \x01(\x0b\x32\x1f.wandb_internal.FilePusherStats\x12/\n\x0b\x66ile_counts\x18\x04 \x01(\x0b\x32\x1a.wandb_internal.FileCounts\x12\x37\n\x0foperation_stats\x18\x05 \x01(\x0b\x32\x1e.wandb_internal.OperationStats\"E\n\x15OperationStatsRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"Q\n\x16OperationStatsResponse\x12\x37\n\x0foperation_stats\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.OperationStats\"Y\n\x0eOperationStats\x12-\n\noperations\x18\x01 \x03(\x0b\x32\x19.wandb_internal.Operation\x12\x18\n\x10total_operations\x18\x02 \x01(\x03\"\x87\x01\n\tOperation\x12\x0c\n\x04\x64\x65sc\x18\x01 \x01(\t\x12\x17\n\x0fruntime_seconds\x18\x02 \x01(\x01\x12\x10\n\x08progress\x18\x03 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x04 \x01(\t\x12+\n\x08subtasks\x18\x05 \x03(\x0b\x32\x19.wandb_internal.Operation\"\x13\n\x11SenderMarkRequest\"\x13\n\x11SyncFinishRequest\"E\n\x0cSyncResponse\x12\x0b\n\x03url\x18\x01 \x01(\t\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"?\n\x11SenderReadRequest\x12\x14\n\x0cstart_offset\x18\x01 \x01(\x03\x12\x14\n\x0c\x66inal_offset\x18\x02 \x01(\x03\"m\n\x13StatusReportRequest\x12\x12\n\nrecord_num\x18\x01 \x01(\x03\x12\x13\n\x0bsent_offset\x18\x02 \x01(\x03\x12-\n\tsync_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"F\n\x14SummaryRecordRequest\x12.\n\x07summary\x18\x01 \x01(\x0b\x32\x1d.wandb_internal.SummaryRecord\"L\n\x16TelemetryRecordRequest\x12\x32\n\ttelemetry\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.TelemetryRecord\"A\n\x11ServerInfoRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"|\n\x12ServerInfoResponse\x12-\n\nlocal_info\x18\x01 \x01(\x0b\x32\x19.wandb_internal.LocalInfo\x12\x37\n\x0fserver_messages\x18\x02 \x01(\x0b\x32\x1e.wandb_internal.ServerMessages\"=\n\x0eServerMessages\x12+\n\x04item\x18\x01 \x03(\x0b\x32\x1d.wandb_internal.ServerMessage\"e\n\rServerMessage\x12\x12\n\nplain_text\x18\x01 \x01(\t\x12\x10\n\x08utf_text\x18\x02 \x01(\t\x12\x11\n\thtml_text\x18\x03 \x01(\t\x12\x0c\n\x04type\x18\x04 \x01(\t\x12\r\n\x05level\x18\x05 \x01(\x05\"c\n\nFileCounts\x12\x13\n\x0bwandb_count\x18\x01 \x01(\x05\x12\x13\n\x0bmedia_count\x18\x02 \x01(\x05\x12\x16\n\x0e\x61rtifact_count\x18\x03 \x01(\x05\x12\x13\n\x0bother_count\x18\x04 \x01(\x05\"U\n\x0f\x46ilePusherStats\x12\x16\n\x0euploaded_bytes\x18\x01 \x01(\x03\x12\x13\n\x0btotal_bytes\x18\x02 \x01(\x03\x12\x15\n\rdeduped_bytes\x18\x03 \x01(\x03\"\x1e\n\rFilesUploaded\x12\r\n\x05\x66iles\x18\x01 \x03(\t\"\xf4\x01\n\x17\x46ileTransferInfoRequest\x12\x42\n\x04type\x18\x01 \x01(\x0e\x32\x34.wandb_internal.FileTransferInfoRequest.TransferType\x12\x0c\n\x04path\x18\x02 \x01(\t\x12\x0b\n\x03url\x18\x03 \x01(\t\x12\x0c\n\x04size\x18\x04 \x01(\x03\x12\x11\n\tprocessed\x18\x05 \x01(\x03\x12/\n\x0b\x66ile_counts\x18\x06 \x01(\x0b\x32\x1a.wandb_internal.FileCounts\"(\n\x0cTransferType\x12\n\n\x06Upload\x10\x00\x12\x0c\n\x08\x44ownload\x10\x01\"1\n\tLocalInfo\x12\x0f\n\x07version\x18\x01 \x01(\t\x12\x13\n\x0bout_of_date\x18\x02 \x01(\x08\"?\n\x0fShutdownRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x12\n\x10ShutdownResponse\"P\n\rAttachRequest\x12\x11\n\tattach_id\x18\x14 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"b\n\x0e\x41ttachResponse\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12(\n\x05\x65rror\x18\x02 \x01(\x0b\x32\x19.wandb_internal.ErrorInfo\"\xd5\x02\n\x11TestInjectRequest\x12\x13\n\x0bhandler_exc\x18\x01 \x01(\x08\x12\x14\n\x0chandler_exit\x18\x02 \x01(\x08\x12\x15\n\rhandler_abort\x18\x03 \x01(\x08\x12\x12\n\nsender_exc\x18\x04 \x01(\x08\x12\x13\n\x0bsender_exit\x18\x05 \x01(\x08\x12\x14\n\x0csender_abort\x18\x06 \x01(\x08\x12\x0f\n\x07req_exc\x18\x07 \x01(\x08\x12\x10\n\x08req_exit\x18\x08 \x01(\x08\x12\x11\n\treq_abort\x18\t \x01(\x08\x12\x10\n\x08resp_exc\x18\n \x01(\x08\x12\x11\n\tresp_exit\x18\x0b \x01(\x08\x12\x12\n\nresp_abort\x18\x0c \x01(\x08\x12\x10\n\x08msg_drop\x18\r \x01(\x08\x12\x10\n\x08msg_hang\x18\x0e \x01(\x08\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x14\n\x12TestInjectResponse\"\x1e\n\rHistoryAction\x12\r\n\x05\x66lush\x18\x01 \x01(\x08\"\xca\x01\n\x15PartialHistoryRequest\x12)\n\x04item\x18\x01 \x03(\x0b\x32\x1b.wandb_internal.HistoryItem\x12)\n\x04step\x18\x02 \x01(\x0b\x32\x1b.wandb_internal.HistoryStep\x12-\n\x06\x61\x63tion\x18\x03 \x01(\x0b\x32\x1d.wandb_internal.HistoryAction\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x18\n\x16PartialHistoryResponse\"E\n\x15SampledHistoryRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"_\n\x12SampledHistoryItem\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x12\n\nnested_key\x18\x02 \x03(\t\x12\x14\n\x0cvalues_float\x18\x03 \x03(\x02\x12\x12\n\nvalues_int\x18\x04 \x03(\x03\"J\n\x16SampledHistoryResponse\x12\x30\n\x04item\x18\x01 \x03(\x0b\x32\".wandb_internal.SampledHistoryItem\"@\n\x10RunStatusRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"x\n\x11RunStatusResponse\x12\x18\n\x10sync_items_total\x18\x01 \x01(\x03\x12\x1a\n\x12sync_items_pending\x18\x02 \x01(\x03\x12-\n\tsync_time\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\"g\n\x0fRunStartRequest\x12&\n\x03run\x18\x01 \x01(\x0b\x32\x19.wandb_internal.RunRecord\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x12\n\x10RunStartResponse\"K\n\x1bRunFinishWithoutExitRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x1e\n\x1cRunFinishWithoutExitResponse\"\\\n\x13\x43heckVersionRequest\x12\x17\n\x0f\x63urrent_version\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"]\n\x14\x43heckVersionResponse\x12\x17\n\x0fupgrade_message\x18\x01 \x01(\t\x12\x14\n\x0cyank_message\x18\x02 \x01(\t\x12\x16\n\x0e\x64\x65lete_message\x18\x03 \x01(\t\">\n\x0eJobInfoRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"6\n\x0fJobInfoResponse\x12\x12\n\nsequenceId\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"\x9f\x01\n\x12LogArtifactRequest\x12\x30\n\x08\x61rtifact\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.ArtifactRecord\x12\x14\n\x0chistory_step\x18\x02 \x01(\x03\x12\x13\n\x0bstaging_dir\x18\x03 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"A\n\x13LogArtifactResponse\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x15\n\rerror_message\x18\x02 \x01(\t\"\xbe\x01\n\x17\x44ownloadArtifactRequest\x12\x13\n\x0b\x61rtifact_id\x18\x01 \x01(\t\x12\x15\n\rdownload_root\x18\x02 \x01(\t\x12 \n\x18\x61llow_missing_references\x18\x04 \x01(\x08\x12\x12\n\nskip_cache\x18\x05 \x01(\x08\x12\x13\n\x0bpath_prefix\x18\x06 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"1\n\x18\x44ownloadArtifactResponse\x12\x15\n\rerror_message\x18\x01 \x01(\t\"@\n\x10KeepaliveRequest\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x13\n\x11KeepaliveResponse\"q\n\x0c\x41rtifactInfo\x12\x10\n\x08\x61rtifact\x18\x01 \x01(\t\x12\x12\n\nentrypoint\x18\x02 \x03(\t\x12\x10\n\x08notebook\x18\x03 \x01(\x08\x12\x15\n\rbuild_context\x18\x04 \x01(\t\x12\x12\n\ndockerfile\x18\x05 \x01(\t\")\n\x07GitInfo\x12\x0e\n\x06remote\x18\x01 \x01(\t\x12\x0e\n\x06\x63ommit\x18\x02 \x01(\t\"\x87\x01\n\tGitSource\x12)\n\x08git_info\x18\x01 \x01(\x0b\x32\x17.wandb_internal.GitInfo\x12\x12\n\nentrypoint\x18\x02 \x03(\t\x12\x10\n\x08notebook\x18\x03 \x01(\x08\x12\x15\n\rbuild_context\x18\x04 \x01(\t\x12\x12\n\ndockerfile\x18\x05 \x01(\t\"\x1c\n\x0bImageSource\x12\r\n\x05image\x18\x01 \x01(\t\"\x8c\x01\n\x06Source\x12&\n\x03git\x18\x01 \x01(\x0b\x32\x19.wandb_internal.GitSource\x12.\n\x08\x61rtifact\x18\x02 \x01(\x0b\x32\x1c.wandb_internal.ArtifactInfo\x12*\n\x05image\x18\x03 \x01(\x0b\x32\x1b.wandb_internal.ImageSource\"k\n\tJobSource\x12\x10\n\x08_version\x18\x01 \x01(\t\x12\x13\n\x0bsource_type\x18\x02 \x01(\t\x12&\n\x06source\x18\x03 \x01(\x0b\x32\x16.wandb_internal.Source\x12\x0f\n\x07runtime\x18\x04 \x01(\t\"V\n\x12PartialJobArtifact\x12\x10\n\x08job_name\x18\x01 \x01(\t\x12.\n\x0bsource_info\x18\x02 \x01(\x0b\x32\x19.wandb_internal.JobSource\"\x9d\x01\n\x11UseArtifactRecord\x12\n\n\x02id\x18\x01 \x01(\t\x12\x0c\n\x04type\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x33\n\x07partial\x18\x04 \x01(\x0b\x32\".wandb_internal.PartialJobArtifact\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x13\n\x11UseArtifactResult\"R\n\rCancelRequest\x12\x13\n\x0b\x63\x61ncel_slot\x18\x01 \x01(\t\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"\x10\n\x0e\x43\x61ncelResponse\"\'\n\x08\x44iskInfo\x12\r\n\x05total\x18\x01 \x01(\x04\x12\x0c\n\x04used\x18\x02 \x01(\x04\"\x1b\n\nMemoryInfo\x12\r\n\x05total\x18\x01 \x01(\x04\"/\n\x07\x43puInfo\x12\r\n\x05\x63ount\x18\x01 \x01(\r\x12\x15\n\rcount_logical\x18\x02 \x01(\r\"\x9a\x01\n\tAppleInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x12\n\necpu_cores\x18\x02 \x01(\r\x12\x12\n\npcpu_cores\x18\x03 \x01(\r\x12\x11\n\tgpu_cores\x18\x04 \x01(\r\x12\x11\n\tmemory_gb\x18\x05 \x01(\r\x12\x18\n\x10swap_total_bytes\x18\x06 \x01(\x04\x12\x17\n\x0fram_total_bytes\x18\x07 \x01(\x04\"k\n\rGpuNvidiaInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x14\n\x0cmemory_total\x18\x02 \x01(\x04\x12\x12\n\ncuda_cores\x18\x03 \x01(\r\x12\x14\n\x0c\x61rchitecture\x18\x04 \x01(\t\x12\x0c\n\x04uuid\x18\x05 \x01(\t\"\x89\x02\n\nGpuAmdInfo\x12\n\n\x02id\x18\x01 \x01(\t\x12\x11\n\tunique_id\x18\x02 \x01(\t\x12\x15\n\rvbios_version\x18\x03 \x01(\t\x12\x19\n\x11performance_level\x18\x04 \x01(\t\x12\x15\n\rgpu_overdrive\x18\x05 \x01(\t\x12\x1c\n\x14gpu_memory_overdrive\x18\x06 \x01(\t\x12\x11\n\tmax_power\x18\x07 \x01(\t\x12\x0e\n\x06series\x18\x08 \x01(\t\x12\r\n\x05model\x18\t \x01(\t\x12\x0e\n\x06vendor\x18\n \x01(\t\x12\x0b\n\x03sku\x18\x0b \x01(\t\x12\x12\n\nsclk_range\x18\x0c \x01(\t\x12\x12\n\nmclk_range\x18\r \x01(\t\"n\n\x0cTrainiumInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0e\n\x06vendor\x18\x02 \x01(\t\x12\x1b\n\x13neuron_device_count\x18\x03 \x01(\r\x12#\n\x1bneuroncore_per_device_count\x18\x04 \x01(\r\"Q\n\x07TPUInfo\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07hbm_gib\x18\x02 \x01(\r\x12\x18\n\x10\x64\x65vices_per_chip\x18\x03 \x01(\r\x12\r\n\x05\x63ount\x18\x04 \x01(\r\"E\n\rCoreWeaveInfo\x12\x14\n\x0c\x63luster_name\x18\x01 \x01(\t\x12\x0e\n\x06org_id\x18\x02 \x01(\t\x12\x0e\n\x06region\x18\x03 \x01(\t\"\xa8\t\n\x11\x45nvironmentRecord\x12\n\n\x02os\x18\x01 \x01(\t\x12\x0e\n\x06python\x18\x02 \x01(\t\x12\x39\n\nstarted_at\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.TimestampR\tstartedAt\x12\x0e\n\x06\x64ocker\x18\x04 \x01(\t\x12\x0c\n\x04\x61rgs\x18\x05 \x03(\t\x12\x0f\n\x07program\x18\x06 \x01(\t\x12\x1b\n\tcode_path\x18\x07 \x01(\tR\x08\x63odePath\x12&\n\x0f\x63ode_path_local\x18\x08 \x01(\tR\rcodePathLocal\x12*\n\x03git\x18\t \x01(\x0b\x32\x1d.wandb_internal.GitRepoRecord\x12\r\n\x05\x65mail\x18\n \x01(\t\x12\x0c\n\x04root\x18\x0b \x01(\t\x12\x0c\n\x04host\x18\x0c \x01(\t\x12\x10\n\x08username\x18\r \x01(\t\x12\x12\n\nexecutable\x18\x0e \x01(\t\x12\r\n\x05\x63olab\x18\x0f \x01(\t\x12\x1c\n\tcpu_count\x18\x10 \x01(\rR\tcpu_count\x12,\n\x11\x63pu_count_logical\x18\x11 \x01(\rR\x11\x63pu_count_logical\x12\x15\n\x08gpu_type\x18\x12 \x01(\tR\x03gpu\x12\x1c\n\tgpu_count\x18\x13 \x01(\rR\tgpu_count\x12\x39\n\x04\x64isk\x18\x14 \x03(\x0b\x32+.wandb_internal.EnvironmentRecord.DiskEntry\x12*\n\x06memory\x18\x15 \x01(\x0b\x32\x1a.wandb_internal.MemoryInfo\x12$\n\x03\x63pu\x18\x16 \x01(\x0b\x32\x17.wandb_internal.CpuInfo\x12(\n\x05\x61pple\x18\x17 \x01(\x0b\x32\x19.wandb_internal.AppleInfo\x12=\n\ngpu_nvidia\x18\x18 \x03(\x0b\x32\x1d.wandb_internal.GpuNvidiaInfoR\ngpu_nvidia\x12\x14\n\x0c\x63uda_version\x18\x19 \x01(\t\x12\x34\n\x07gpu_amd\x18\x1a \x03(\x0b\x32\x1a.wandb_internal.GpuAmdInfoR\x07gpu_amd\x12;\n\x05slurm\x18\x1b \x03(\x0b\x32,.wandb_internal.EnvironmentRecord.SlurmEntry\x12.\n\x08trainium\x18\x1c \x01(\x0b\x32\x1c.wandb_internal.TrainiumInfo\x12$\n\x03tpu\x18\x1d \x01(\x0b\x32\x17.wandb_internal.TPUInfo\x12\x30\n\tcoreweave\x18\x1e \x01(\x0b\x32\x1d.wandb_internal.CoreWeaveInfo\x12\x12\n\twriter_id\x18\xc7\x01 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\x1a\x45\n\tDiskEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.wandb_internal.DiskInfo:\x02\x38\x01\x1a,\n\nSlurmEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\x15PythonPackagesRequest\x12\x44\n\x07package\x18\x01 \x03(\x0b\x32\x33.wandb_internal.PythonPackagesRequest.PythonPackage\x1a.\n\rPythonPackage\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07version\x18\x02 \x01(\t\"\x1c\n\x0cJobInputPath\x12\x0c\n\x04path\x18\x01 \x03(\t\"\xd6\x01\n\x0eJobInputSource\x12\x44\n\nrun_config\x18\x01 \x01(\x0b\x32..wandb_internal.JobInputSource.RunConfigSourceH\x00\x12?\n\x04\x66ile\x18\x02 \x01(\x0b\x32/.wandb_internal.JobInputSource.ConfigFileSourceH\x00\x1a\x11\n\x0fRunConfigSource\x1a \n\x10\x43onfigFileSource\x12\x0c\n\x04path\x18\x01 \x01(\tB\x08\n\x06source\"\xc7\x01\n\x0fJobInputRequest\x12\x34\n\x0cinput_source\x18\x01 \x01(\x0b\x32\x1e.wandb_internal.JobInputSource\x12\x33\n\rinclude_paths\x18\x02 \x03(\x0b\x32\x1c.wandb_internal.JobInputPath\x12\x33\n\rexclude_paths\x18\x03 \x03(\x0b\x32\x1c.wandb_internal.JobInputPath\x12\x14\n\x0cinput_schema\x18\x04 \x01(\t\"t\n\x14ServerFeatureRequest\x12.\n\x07\x66\x65\x61ture\x18\x01 \x01(\x0e\x32\x1d.wandb_internal.ServerFeature\x12,\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1c.wandb_internal._RequestInfo\"K\n\x15ServerFeatureResponse\x12\x32\n\x07\x66\x65\x61ture\x18\x01 \x01(\x0b\x32!.wandb_internal.ServerFeatureItem\"2\n\x11ServerFeatureItem\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x0f\n\x07\x65nabled\x18\x02 \x01(\x08*\x80\x04\n\rServerFeature\x12\x13\n\x0fLARGE_FILENAMES\x10\x00\x12\x11\n\rARTIFACT_TAGS\x10\x01\x12\x0e\n\nCLIENT_IDS\x10\x02\x12\x1c\n\x18\x41RTIFACT_REGISTRY_SEARCH\x10\x03\x12\x1b\n\x17STRUCTURED_CONSOLE_LOGS\x10\x04\x12(\n$ARTIFACT_COLLECTION_MEMBERSHIP_FILES\x10\x05\x12\x38\n4ARTIFACT_COLLECTION_MEMBERSHIP_FILE_DOWNLOAD_HANDLER\x10\x06\x12\x34\n0USE_ARTIFACT_WITH_ENTITY_AND_PROJECT_INFORMATION\x10\x07\x12\x1f\n\x1b\x45XPAND_DEFINED_METRIC_GLOBS\x10\x08\x12\x1f\n\x1b\x41UTOMATION_EVENT_RUN_METRIC\x10\t\x12&\n\"AUTOMATION_EVENT_RUN_METRIC_CHANGE\x10\n\x12\x1b\n\x17\x41UTOMATION_ACTION_NO_OP\x10\x0b\x12/\n+INCLUDE_ARTIFACT_TYPES_IN_REGISTRY_CREATION\x10\x0c\x12*\n&PROJECT_ARTIFACT_COLLECTION_MEMBERSHIP\x10\rB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_internal_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY']._loaded_options = None + _globals['_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY']._serialized_options = b'8\001' + _globals['_ENVIRONMENTRECORD_DISKENTRY']._loaded_options = None + _globals['_ENVIRONMENTRECORD_DISKENTRY']._serialized_options = b'8\001' + _globals['_ENVIRONMENTRECORD_SLURMENTRY']._loaded_options = None + _globals['_ENVIRONMENTRECORD_SLURMENTRY']._serialized_options = b'8\001' + _globals['_SERVERFEATURE']._serialized_start=22228 + _globals['_SERVERFEATURE']._serialized_end=22740 + _globals['_RECORD']._serialized_start=180 + _globals['_RECORD']._serialized_end=1411 + _globals['_CONTROL']._serialized_start=1414 + _globals['_CONTROL']._serialized_end=1582 + _globals['_RESULT']._serialized_start=1585 + _globals['_RESULT']._serialized_end=2084 + _globals['_FINALRECORD']._serialized_start=2086 + _globals['_FINALRECORD']._serialized_end=2144 + _globals['_VERSIONINFO']._serialized_start=2146 + _globals['_VERSIONINFO']._serialized_end=2244 + _globals['_HEADERRECORD']._serialized_start=2246 + _globals['_HEADERRECORD']._serialized_end=2356 + _globals['_FOOTERRECORD']._serialized_start=2358 + _globals['_FOOTERRECORD']._serialized_end=2417 + _globals['_BRANCHPOINT']._serialized_start=2419 + _globals['_BRANCHPOINT']._serialized_end=2476 + _globals['_RUNRECORD']._serialized_start=2479 + _globals['_RUNRECORD']._serialized_end=3136 + _globals['_GITREPORECORD']._serialized_start=3138 + _globals['_GITREPORECORD']._serialized_end=3197 + _globals['_RUNUPDATERESULT']._serialized_start=3199 + _globals['_RUNUPDATERESULT']._serialized_end=3298 + _globals['_ERRORINFO']._serialized_start=3301 + _globals['_ERRORINFO']._serialized_end=3473 + _globals['_ERRORINFO_ERRORCODE']._serialized_start=3382 + _globals['_ERRORINFO_ERRORCODE']._serialized_end=3473 + _globals['_RUNEXITRECORD']._serialized_start=3475 + _globals['_RUNEXITRECORD']._serialized_end=3571 + _globals['_RUNEXITRESULT']._serialized_start=3573 + _globals['_RUNEXITRESULT']._serialized_end=3588 + _globals['_RUNPREEMPTINGRECORD']._serialized_start=3590 + _globals['_RUNPREEMPTINGRECORD']._serialized_end=3656 + _globals['_RUNPREEMPTINGRESULT']._serialized_start=3658 + _globals['_RUNPREEMPTINGRESULT']._serialized_end=3679 + _globals['_SETTINGSRECORD']._serialized_start=3681 + _globals['_SETTINGSRECORD']._serialized_end=3786 + _globals['_SETTINGSITEM']._serialized_start=3788 + _globals['_SETTINGSITEM']._serialized_end=3835 + _globals['_HISTORYSTEP']._serialized_start=3837 + _globals['_HISTORYSTEP']._serialized_end=3863 + _globals['_HISTORYRECORD']._serialized_start=3866 + _globals['_HISTORYRECORD']._serialized_end=4012 + _globals['_HISTORYITEM']._serialized_start=4014 + _globals['_HISTORYITEM']._serialized_end=4080 + _globals['_HISTORYRESULT']._serialized_start=4082 + _globals['_HISTORYRESULT']._serialized_end=4097 + _globals['_OUTPUTRECORD']._serialized_start=4100 + _globals['_OUTPUTRECORD']._serialized_end=4320 + _globals['_OUTPUTRECORD_OUTPUTTYPE']._serialized_start=4284 + _globals['_OUTPUTRECORD_OUTPUTTYPE']._serialized_end=4320 + _globals['_OUTPUTRESULT']._serialized_start=4322 + _globals['_OUTPUTRESULT']._serialized_end=4336 + _globals['_OUTPUTRAWRECORD']._serialized_start=4339 + _globals['_OUTPUTRAWRECORD']._serialized_end=4565 + _globals['_OUTPUTRAWRECORD_OUTPUTTYPE']._serialized_start=4284 + _globals['_OUTPUTRAWRECORD_OUTPUTTYPE']._serialized_end=4320 + _globals['_OUTPUTRAWRESULT']._serialized_start=4567 + _globals['_OUTPUTRAWRESULT']._serialized_end=4584 + _globals['_METRICRECORD']._serialized_start=4587 + _globals['_METRICRECORD']._serialized_end=5023 + _globals['_METRICRECORD_METRICGOAL']._serialized_start=4957 + _globals['_METRICRECORD_METRICGOAL']._serialized_end=5023 + _globals['_METRICRESULT']._serialized_start=5025 + _globals['_METRICRESULT']._serialized_end=5039 + _globals['_METRICOPTIONS']._serialized_start=5041 + _globals['_METRICOPTIONS']._serialized_end=5108 + _globals['_METRICCONTROL']._serialized_start=5110 + _globals['_METRICCONTROL']._serialized_end=5144 + _globals['_METRICSUMMARY']._serialized_start=5146 + _globals['_METRICSUMMARY']._serialized_end=5257 + _globals['_CONFIGRECORD']._serialized_start=5260 + _globals['_CONFIGRECORD']._serialized_end=5407 + _globals['_CONFIGITEM']._serialized_start=5409 + _globals['_CONFIGITEM']._serialized_end=5474 + _globals['_CONFIGRESULT']._serialized_start=5476 + _globals['_CONFIGRESULT']._serialized_end=5490 + _globals['_SUMMARYRECORD']._serialized_start=5493 + _globals['_SUMMARYRECORD']._serialized_end=5643 + _globals['_SUMMARYITEM']._serialized_start=5645 + _globals['_SUMMARYITEM']._serialized_end=5711 + _globals['_SUMMARYRESULT']._serialized_start=5713 + _globals['_SUMMARYRESULT']._serialized_end=5728 + _globals['_FILESRECORD']._serialized_start=5730 + _globals['_FILESRECORD']._serialized_end=5830 + _globals['_FILESITEM']._serialized_start=5833 + _globals['_FILESITEM']._serialized_end=6069 + _globals['_FILESITEM_POLICYTYPE']._serialized_start=5964 + _globals['_FILESITEM_POLICYTYPE']._serialized_end=6004 + _globals['_FILESITEM_FILETYPE']._serialized_start=6006 + _globals['_FILESITEM_FILETYPE']._serialized_end=6063 + _globals['_FILESRESULT']._serialized_start=6071 + _globals['_FILESRESULT']._serialized_end=6084 + _globals['_STATSRECORD']._serialized_start=6087 + _globals['_STATSRECORD']._serialized_end=6317 + _globals['_STATSRECORD_STATSTYPE']._serialized_start=6294 + _globals['_STATSRECORD_STATSTYPE']._serialized_end=6317 + _globals['_STATSITEM']._serialized_start=6319 + _globals['_STATSITEM']._serialized_end=6363 + _globals['_ARTIFACTRECORD']._serialized_start=6366 + _globals['_ARTIFACTRECORD']._serialized_end=6853 + _globals['_ARTIFACTMANIFEST']._serialized_start=6856 + _globals['_ARTIFACTMANIFEST']._serialized_end=7072 + _globals['_ARTIFACTMANIFESTENTRY']._serialized_start=7075 + _globals['_ARTIFACTMANIFESTENTRY']._serialized_end=7282 + _globals['_EXTRAITEM']._serialized_start=7284 + _globals['_EXTRAITEM']._serialized_end=7328 + _globals['_STORAGEPOLICYCONFIGITEM']._serialized_start=7330 + _globals['_STORAGEPOLICYCONFIGITEM']._serialized_end=7388 + _globals['_ARTIFACTRESULT']._serialized_start=7390 + _globals['_ARTIFACTRESULT']._serialized_end=7406 + _globals['_LINKARTIFACTRESULT']._serialized_start=7408 + _globals['_LINKARTIFACTRESULT']._serialized_end=7428 + _globals['_LINKARTIFACTREQUEST']._serialized_start=7431 + _globals['_LINKARTIFACTREQUEST']._serialized_end=7671 + _globals['_LINKARTIFACTRESPONSE']._serialized_start=7673 + _globals['_LINKARTIFACTRESPONSE']._serialized_end=7764 + _globals['_TBRECORD']._serialized_start=7766 + _globals['_TBRECORD']._serialized_end=7870 + _globals['_TBRESULT']._serialized_start=7872 + _globals['_TBRESULT']._serialized_end=7882 + _globals['_ALERTRECORD']._serialized_start=7884 + _globals['_ALERTRECORD']._serialized_end=8009 + _globals['_ALERTRESULT']._serialized_start=8011 + _globals['_ALERTRESULT']._serialized_end=8024 + _globals['_REQUEST']._serialized_start=8027 + _globals['_REQUEST']._serialized_end=10254 + _globals['_RESPONSE']._serialized_start=10257 + _globals['_RESPONSE']._serialized_end=12076 + _globals['_DEFERREQUEST']._serialized_start=12079 + _globals['_DEFERREQUEST']._serialized_end=12399 + _globals['_DEFERREQUEST_DEFERSTATE']._serialized_start=12152 + _globals['_DEFERREQUEST_DEFERSTATE']._serialized_end=12399 + _globals['_PAUSEREQUEST']._serialized_start=12401 + _globals['_PAUSEREQUEST']._serialized_end=12461 + _globals['_PAUSERESPONSE']._serialized_start=12463 + _globals['_PAUSERESPONSE']._serialized_end=12478 + _globals['_RESUMEREQUEST']._serialized_start=12480 + _globals['_RESUMEREQUEST']._serialized_end=12541 + _globals['_RESUMERESPONSE']._serialized_start=12543 + _globals['_RESUMERESPONSE']._serialized_end=12559 + _globals['_LOGINREQUEST']._serialized_start=12561 + _globals['_LOGINREQUEST']._serialized_end=12638 + _globals['_LOGINRESPONSE']._serialized_start=12640 + _globals['_LOGINRESPONSE']._serialized_end=12678 + _globals['_GETSUMMARYREQUEST']._serialized_start=12680 + _globals['_GETSUMMARYREQUEST']._serialized_end=12745 + _globals['_GETSUMMARYRESPONSE']._serialized_start=12747 + _globals['_GETSUMMARYRESPONSE']._serialized_end=12810 + _globals['_GETSYSTEMMETRICSREQUEST']._serialized_start=12812 + _globals['_GETSYSTEMMETRICSREQUEST']._serialized_end=12883 + _globals['_SYSTEMMETRICSAMPLE']._serialized_start=12885 + _globals['_SYSTEMMETRICSAMPLE']._serialized_end=12967 + _globals['_SYSTEMMETRICSBUFFER']._serialized_start=12969 + _globals['_SYSTEMMETRICSBUFFER']._serialized_end=13042 + _globals['_GETSYSTEMMETRICSRESPONSE']._serialized_start=13045 + _globals['_GETSYSTEMMETRICSRESPONSE']._serialized_end=13247 + _globals['_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY']._serialized_start=13158 + _globals['_GETSYSTEMMETRICSRESPONSE_SYSTEMMETRICSENTRY']._serialized_end=13247 + _globals['_STATUSREQUEST']._serialized_start=13249 + _globals['_STATUSREQUEST']._serialized_end=13310 + _globals['_STATUSRESPONSE']._serialized_start=13312 + _globals['_STATUSRESPONSE']._serialized_end=13353 + _globals['_STOPSTATUSREQUEST']._serialized_start=13355 + _globals['_STOPSTATUSREQUEST']._serialized_end=13420 + _globals['_STOPSTATUSRESPONSE']._serialized_start=13422 + _globals['_STOPSTATUSRESPONSE']._serialized_end=13467 + _globals['_NETWORKSTATUSREQUEST']._serialized_start=13469 + _globals['_NETWORKSTATUSREQUEST']._serialized_end=13537 + _globals['_NETWORKSTATUSRESPONSE']._serialized_start=13539 + _globals['_NETWORKSTATUSRESPONSE']._serialized_end=13619 + _globals['_HTTPRESPONSE']._serialized_start=13621 + _globals['_HTTPRESPONSE']._serialized_end=13689 + _globals['_INTERNALMESSAGESREQUEST']._serialized_start=13691 + _globals['_INTERNALMESSAGESREQUEST']._serialized_end=13762 + _globals['_INTERNALMESSAGESRESPONSE']._serialized_start=13764 + _globals['_INTERNALMESSAGESRESPONSE']._serialized_end=13842 + _globals['_INTERNALMESSAGES']._serialized_start=13844 + _globals['_INTERNALMESSAGES']._serialized_end=13879 + _globals['_POLLEXITREQUEST']._serialized_start=13881 + _globals['_POLLEXITREQUEST']._serialized_end=13944 + _globals['_POLLEXITRESPONSE']._serialized_start=13947 + _globals['_POLLEXITRESPONSE']._serialized_end=14192 + _globals['_OPERATIONSTATSREQUEST']._serialized_start=14194 + _globals['_OPERATIONSTATSREQUEST']._serialized_end=14263 + _globals['_OPERATIONSTATSRESPONSE']._serialized_start=14265 + _globals['_OPERATIONSTATSRESPONSE']._serialized_end=14346 + _globals['_OPERATIONSTATS']._serialized_start=14348 + _globals['_OPERATIONSTATS']._serialized_end=14437 + _globals['_OPERATION']._serialized_start=14440 + _globals['_OPERATION']._serialized_end=14575 + _globals['_SENDERMARKREQUEST']._serialized_start=14577 + _globals['_SENDERMARKREQUEST']._serialized_end=14596 + _globals['_SYNCFINISHREQUEST']._serialized_start=14598 + _globals['_SYNCFINISHREQUEST']._serialized_end=14617 + _globals['_SYNCRESPONSE']._serialized_start=14619 + _globals['_SYNCRESPONSE']._serialized_end=14688 + _globals['_SENDERREADREQUEST']._serialized_start=14690 + _globals['_SENDERREADREQUEST']._serialized_end=14753 + _globals['_STATUSREPORTREQUEST']._serialized_start=14755 + _globals['_STATUSREPORTREQUEST']._serialized_end=14864 + _globals['_SUMMARYRECORDREQUEST']._serialized_start=14866 + _globals['_SUMMARYRECORDREQUEST']._serialized_end=14936 + _globals['_TELEMETRYRECORDREQUEST']._serialized_start=14938 + _globals['_TELEMETRYRECORDREQUEST']._serialized_end=15014 + _globals['_SERVERINFOREQUEST']._serialized_start=15016 + _globals['_SERVERINFOREQUEST']._serialized_end=15081 + _globals['_SERVERINFORESPONSE']._serialized_start=15083 + _globals['_SERVERINFORESPONSE']._serialized_end=15207 + _globals['_SERVERMESSAGES']._serialized_start=15209 + _globals['_SERVERMESSAGES']._serialized_end=15270 + _globals['_SERVERMESSAGE']._serialized_start=15272 + _globals['_SERVERMESSAGE']._serialized_end=15373 + _globals['_FILECOUNTS']._serialized_start=15375 + _globals['_FILECOUNTS']._serialized_end=15474 + _globals['_FILEPUSHERSTATS']._serialized_start=15476 + _globals['_FILEPUSHERSTATS']._serialized_end=15561 + _globals['_FILESUPLOADED']._serialized_start=15563 + _globals['_FILESUPLOADED']._serialized_end=15593 + _globals['_FILETRANSFERINFOREQUEST']._serialized_start=15596 + _globals['_FILETRANSFERINFOREQUEST']._serialized_end=15840 + _globals['_FILETRANSFERINFOREQUEST_TRANSFERTYPE']._serialized_start=15800 + _globals['_FILETRANSFERINFOREQUEST_TRANSFERTYPE']._serialized_end=15840 + _globals['_LOCALINFO']._serialized_start=15842 + _globals['_LOCALINFO']._serialized_end=15891 + _globals['_SHUTDOWNREQUEST']._serialized_start=15893 + _globals['_SHUTDOWNREQUEST']._serialized_end=15956 + _globals['_SHUTDOWNRESPONSE']._serialized_start=15958 + _globals['_SHUTDOWNRESPONSE']._serialized_end=15976 + _globals['_ATTACHREQUEST']._serialized_start=15978 + _globals['_ATTACHREQUEST']._serialized_end=16058 + _globals['_ATTACHRESPONSE']._serialized_start=16060 + _globals['_ATTACHRESPONSE']._serialized_end=16158 + _globals['_TESTINJECTREQUEST']._serialized_start=16161 + _globals['_TESTINJECTREQUEST']._serialized_end=16502 + _globals['_TESTINJECTRESPONSE']._serialized_start=16504 + _globals['_TESTINJECTRESPONSE']._serialized_end=16524 + _globals['_HISTORYACTION']._serialized_start=16526 + _globals['_HISTORYACTION']._serialized_end=16556 + _globals['_PARTIALHISTORYREQUEST']._serialized_start=16559 + _globals['_PARTIALHISTORYREQUEST']._serialized_end=16761 + _globals['_PARTIALHISTORYRESPONSE']._serialized_start=16763 + _globals['_PARTIALHISTORYRESPONSE']._serialized_end=16787 + _globals['_SAMPLEDHISTORYREQUEST']._serialized_start=16789 + _globals['_SAMPLEDHISTORYREQUEST']._serialized_end=16858 + _globals['_SAMPLEDHISTORYITEM']._serialized_start=16860 + _globals['_SAMPLEDHISTORYITEM']._serialized_end=16955 + _globals['_SAMPLEDHISTORYRESPONSE']._serialized_start=16957 + _globals['_SAMPLEDHISTORYRESPONSE']._serialized_end=17031 + _globals['_RUNSTATUSREQUEST']._serialized_start=17033 + _globals['_RUNSTATUSREQUEST']._serialized_end=17097 + _globals['_RUNSTATUSRESPONSE']._serialized_start=17099 + _globals['_RUNSTATUSRESPONSE']._serialized_end=17219 + _globals['_RUNSTARTREQUEST']._serialized_start=17221 + _globals['_RUNSTARTREQUEST']._serialized_end=17324 + _globals['_RUNSTARTRESPONSE']._serialized_start=17326 + _globals['_RUNSTARTRESPONSE']._serialized_end=17344 + _globals['_RUNFINISHWITHOUTEXITREQUEST']._serialized_start=17346 + _globals['_RUNFINISHWITHOUTEXITREQUEST']._serialized_end=17421 + _globals['_RUNFINISHWITHOUTEXITRESPONSE']._serialized_start=17423 + _globals['_RUNFINISHWITHOUTEXITRESPONSE']._serialized_end=17453 + _globals['_CHECKVERSIONREQUEST']._serialized_start=17455 + _globals['_CHECKVERSIONREQUEST']._serialized_end=17547 + _globals['_CHECKVERSIONRESPONSE']._serialized_start=17549 + _globals['_CHECKVERSIONRESPONSE']._serialized_end=17642 + _globals['_JOBINFOREQUEST']._serialized_start=17644 + _globals['_JOBINFOREQUEST']._serialized_end=17706 + _globals['_JOBINFORESPONSE']._serialized_start=17708 + _globals['_JOBINFORESPONSE']._serialized_end=17762 + _globals['_LOGARTIFACTREQUEST']._serialized_start=17765 + _globals['_LOGARTIFACTREQUEST']._serialized_end=17924 + _globals['_LOGARTIFACTRESPONSE']._serialized_start=17926 + _globals['_LOGARTIFACTRESPONSE']._serialized_end=17991 + _globals['_DOWNLOADARTIFACTREQUEST']._serialized_start=17994 + _globals['_DOWNLOADARTIFACTREQUEST']._serialized_end=18184 + _globals['_DOWNLOADARTIFACTRESPONSE']._serialized_start=18186 + _globals['_DOWNLOADARTIFACTRESPONSE']._serialized_end=18235 + _globals['_KEEPALIVEREQUEST']._serialized_start=18237 + _globals['_KEEPALIVEREQUEST']._serialized_end=18301 + _globals['_KEEPALIVERESPONSE']._serialized_start=18303 + _globals['_KEEPALIVERESPONSE']._serialized_end=18322 + _globals['_ARTIFACTINFO']._serialized_start=18324 + _globals['_ARTIFACTINFO']._serialized_end=18437 + _globals['_GITINFO']._serialized_start=18439 + _globals['_GITINFO']._serialized_end=18480 + _globals['_GITSOURCE']._serialized_start=18483 + _globals['_GITSOURCE']._serialized_end=18618 + _globals['_IMAGESOURCE']._serialized_start=18620 + _globals['_IMAGESOURCE']._serialized_end=18648 + _globals['_SOURCE']._serialized_start=18651 + _globals['_SOURCE']._serialized_end=18791 + _globals['_JOBSOURCE']._serialized_start=18793 + _globals['_JOBSOURCE']._serialized_end=18900 + _globals['_PARTIALJOBARTIFACT']._serialized_start=18902 + _globals['_PARTIALJOBARTIFACT']._serialized_end=18988 + _globals['_USEARTIFACTRECORD']._serialized_start=18991 + _globals['_USEARTIFACTRECORD']._serialized_end=19148 + _globals['_USEARTIFACTRESULT']._serialized_start=19150 + _globals['_USEARTIFACTRESULT']._serialized_end=19169 + _globals['_CANCELREQUEST']._serialized_start=19171 + _globals['_CANCELREQUEST']._serialized_end=19253 + _globals['_CANCELRESPONSE']._serialized_start=19255 + _globals['_CANCELRESPONSE']._serialized_end=19271 + _globals['_DISKINFO']._serialized_start=19273 + _globals['_DISKINFO']._serialized_end=19312 + _globals['_MEMORYINFO']._serialized_start=19314 + _globals['_MEMORYINFO']._serialized_end=19341 + _globals['_CPUINFO']._serialized_start=19343 + _globals['_CPUINFO']._serialized_end=19390 + _globals['_APPLEINFO']._serialized_start=19393 + _globals['_APPLEINFO']._serialized_end=19547 + _globals['_GPUNVIDIAINFO']._serialized_start=19549 + _globals['_GPUNVIDIAINFO']._serialized_end=19656 + _globals['_GPUAMDINFO']._serialized_start=19659 + _globals['_GPUAMDINFO']._serialized_end=19924 + _globals['_TRAINIUMINFO']._serialized_start=19926 + _globals['_TRAINIUMINFO']._serialized_end=20036 + _globals['_TPUINFO']._serialized_start=20038 + _globals['_TPUINFO']._serialized_end=20119 + _globals['_COREWEAVEINFO']._serialized_start=20121 + _globals['_COREWEAVEINFO']._serialized_end=20190 + _globals['_ENVIRONMENTRECORD']._serialized_start=20193 + _globals['_ENVIRONMENTRECORD']._serialized_end=21385 + _globals['_ENVIRONMENTRECORD_DISKENTRY']._serialized_start=21270 + _globals['_ENVIRONMENTRECORD_DISKENTRY']._serialized_end=21339 + _globals['_ENVIRONMENTRECORD_SLURMENTRY']._serialized_start=21341 + _globals['_ENVIRONMENTRECORD_SLURMENTRY']._serialized_end=21385 + _globals['_PYTHONPACKAGESREQUEST']._serialized_start=21388 + _globals['_PYTHONPACKAGESREQUEST']._serialized_end=21529 + _globals['_PYTHONPACKAGESREQUEST_PYTHONPACKAGE']._serialized_start=21483 + _globals['_PYTHONPACKAGESREQUEST_PYTHONPACKAGE']._serialized_end=21529 + _globals['_JOBINPUTPATH']._serialized_start=21531 + _globals['_JOBINPUTPATH']._serialized_end=21559 + _globals['_JOBINPUTSOURCE']._serialized_start=21562 + _globals['_JOBINPUTSOURCE']._serialized_end=21776 + _globals['_JOBINPUTSOURCE_RUNCONFIGSOURCE']._serialized_start=21715 + _globals['_JOBINPUTSOURCE_RUNCONFIGSOURCE']._serialized_end=21732 + _globals['_JOBINPUTSOURCE_CONFIGFILESOURCE']._serialized_start=21734 + _globals['_JOBINPUTSOURCE_CONFIGFILESOURCE']._serialized_end=21766 + _globals['_JOBINPUTREQUEST']._serialized_start=21779 + _globals['_JOBINPUTREQUEST']._serialized_end=21978 + _globals['_SERVERFEATUREREQUEST']._serialized_start=21980 + _globals['_SERVERFEATUREREQUEST']._serialized_end=22096 + _globals['_SERVERFEATURERESPONSE']._serialized_start=22098 + _globals['_SERVERFEATURERESPONSE']._serialized_end=22173 + _globals['_SERVERFEATUREITEM']._serialized_start=22175 + _globals['_SERVERFEATUREITEM']._serialized_end=22225 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v6/wandb_server_pb2.py b/lib/python3.12/site-packages/wandb/proto/v6/wandb_server_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..5b0f4e91c13f56c7ec327ccd2a34daae6ea8bfe9 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v6/wandb_server_pb2.py @@ -0,0 +1,78 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: wandb/proto/wandb_server.proto +# Protobuf Python Version: 6.30.0 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 30, + 0, + '', + 'wandb/proto/wandb_server.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 +from wandb.proto import wandb_internal_pb2 as wandb_dot_proto_dot_wandb__internal__pb2 +from wandb.proto import wandb_settings_pb2 as wandb_dot_proto_dot_wandb__settings__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1ewandb/proto/wandb_server.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\x1a wandb/proto/wandb_internal.proto\x1a wandb/proto/wandb_settings.proto\"k\n\x19ServerAuthenticateRequest\x12\x0f\n\x07\x61pi_key\x18\x01 \x01(\t\x12\x10\n\x08\x62\x61se_url\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"w\n\x1aServerAuthenticateResponse\x12\x16\n\x0e\x64\x65\x66\x61ult_entity\x18\x01 \x01(\t\x12\x14\n\x0c\x65rror_status\x18\x02 \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"D\n\x15ServerShutdownRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x18\n\x16ServerShutdownResponse\"B\n\x13ServerStatusRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x16\n\x14ServerStatusResponse\"r\n\x17ServerInformInitRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1a\n\x18ServerInformInitResponse\"s\n\x18ServerInformStartRequest\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1b\n\x19ServerInformStartResponse\"H\n\x19ServerInformFinishRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformFinishResponse\"H\n\x19ServerInformAttachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"u\n\x1aServerInformAttachResponse\x12*\n\x08settings\x18\x01 \x01(\x0b\x32\x18.wandb_internal.Settings\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"H\n\x19ServerInformDetachRequest\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1c\n\x1aServerInformDetachResponse\"]\n\x1bServerInformTeardownRequest\x12\x11\n\texit_code\x18\x01 \x01(\x05\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x1e\n\x1cServerInformTeardownResponse\"\xfb\x04\n\rServerRequest\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x30\n\x0erecord_publish\x18\x01 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12\x34\n\x12record_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.RecordH\x00\x12>\n\x0binform_init\x18\x03 \x01(\x0b\x32\'.wandb_internal.ServerInformInitRequestH\x00\x12\x42\n\rinform_finish\x18\x04 \x01(\x0b\x32).wandb_internal.ServerInformFinishRequestH\x00\x12\x42\n\rinform_attach\x18\x05 \x01(\x0b\x32).wandb_internal.ServerInformAttachRequestH\x00\x12\x42\n\rinform_detach\x18\x06 \x01(\x0b\x32).wandb_internal.ServerInformDetachRequestH\x00\x12\x46\n\x0finform_teardown\x18\x07 \x01(\x0b\x32+.wandb_internal.ServerInformTeardownRequestH\x00\x12@\n\x0cinform_start\x18\x08 \x01(\x0b\x32(.wandb_internal.ServerInformStartRequestH\x00\x12\x41\n\x0c\x61uthenticate\x18\t \x01(\x0b\x32).wandb_internal.ServerAuthenticateRequestH\x00\x42\x15\n\x13server_request_type\"\x91\x05\n\x0eServerResponse\x12\x12\n\nrequest_id\x18\n \x01(\t\x12\x34\n\x12result_communicate\x18\x02 \x01(\x0b\x32\x16.wandb_internal.ResultH\x00\x12H\n\x14inform_init_response\x18\x03 \x01(\x0b\x32(.wandb_internal.ServerInformInitResponseH\x00\x12L\n\x16inform_finish_response\x18\x04 \x01(\x0b\x32*.wandb_internal.ServerInformFinishResponseH\x00\x12L\n\x16inform_attach_response\x18\x05 \x01(\x0b\x32*.wandb_internal.ServerInformAttachResponseH\x00\x12L\n\x16inform_detach_response\x18\x06 \x01(\x0b\x32*.wandb_internal.ServerInformDetachResponseH\x00\x12P\n\x18inform_teardown_response\x18\x07 \x01(\x0b\x32,.wandb_internal.ServerInformTeardownResponseH\x00\x12J\n\x15inform_start_response\x18\x08 \x01(\x0b\x32).wandb_internal.ServerInformStartResponseH\x00\x12K\n\x15\x61uthenticate_response\x18\t \x01(\x0b\x32*.wandb_internal.ServerAuthenticateResponseH\x00\x42\x16\n\x14server_response_typeB\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_server_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['_SERVERAUTHENTICATEREQUEST']._serialized_start=148 + _globals['_SERVERAUTHENTICATEREQUEST']._serialized_end=255 + _globals['_SERVERAUTHENTICATERESPONSE']._serialized_start=257 + _globals['_SERVERAUTHENTICATERESPONSE']._serialized_end=376 + _globals['_SERVERSHUTDOWNREQUEST']._serialized_start=378 + _globals['_SERVERSHUTDOWNREQUEST']._serialized_end=446 + _globals['_SERVERSHUTDOWNRESPONSE']._serialized_start=448 + _globals['_SERVERSHUTDOWNRESPONSE']._serialized_end=472 + _globals['_SERVERSTATUSREQUEST']._serialized_start=474 + _globals['_SERVERSTATUSREQUEST']._serialized_end=540 + _globals['_SERVERSTATUSRESPONSE']._serialized_start=542 + _globals['_SERVERSTATUSRESPONSE']._serialized_end=564 + _globals['_SERVERINFORMINITREQUEST']._serialized_start=566 + _globals['_SERVERINFORMINITREQUEST']._serialized_end=680 + _globals['_SERVERINFORMINITRESPONSE']._serialized_start=682 + _globals['_SERVERINFORMINITRESPONSE']._serialized_end=708 + _globals['_SERVERINFORMSTARTREQUEST']._serialized_start=710 + _globals['_SERVERINFORMSTARTREQUEST']._serialized_end=825 + _globals['_SERVERINFORMSTARTRESPONSE']._serialized_start=827 + _globals['_SERVERINFORMSTARTRESPONSE']._serialized_end=854 + _globals['_SERVERINFORMFINISHREQUEST']._serialized_start=856 + _globals['_SERVERINFORMFINISHREQUEST']._serialized_end=928 + _globals['_SERVERINFORMFINISHRESPONSE']._serialized_start=930 + _globals['_SERVERINFORMFINISHRESPONSE']._serialized_end=958 + _globals['_SERVERINFORMATTACHREQUEST']._serialized_start=960 + _globals['_SERVERINFORMATTACHREQUEST']._serialized_end=1032 + _globals['_SERVERINFORMATTACHRESPONSE']._serialized_start=1034 + _globals['_SERVERINFORMATTACHRESPONSE']._serialized_end=1151 + _globals['_SERVERINFORMDETACHREQUEST']._serialized_start=1153 + _globals['_SERVERINFORMDETACHREQUEST']._serialized_end=1225 + _globals['_SERVERINFORMDETACHRESPONSE']._serialized_start=1227 + _globals['_SERVERINFORMDETACHRESPONSE']._serialized_end=1255 + _globals['_SERVERINFORMTEARDOWNREQUEST']._serialized_start=1257 + _globals['_SERVERINFORMTEARDOWNREQUEST']._serialized_end=1350 + _globals['_SERVERINFORMTEARDOWNRESPONSE']._serialized_start=1352 + _globals['_SERVERINFORMTEARDOWNRESPONSE']._serialized_end=1382 + _globals['_SERVERREQUEST']._serialized_start=1385 + _globals['_SERVERREQUEST']._serialized_end=2020 + _globals['_SERVERRESPONSE']._serialized_start=2023 + _globals['_SERVERRESPONSE']._serialized_end=2680 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v6/wandb_settings_pb2.py b/lib/python3.12/site-packages/wandb/proto/v6/wandb_settings_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..a38a7778fd7f0ddccd415091f69981b2498b3f80 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v6/wandb_settings_pb2.py @@ -0,0 +1,58 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: wandb/proto/wandb_settings.proto +# Protobuf Python Version: 6.30.0 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 30, + 0, + '', + 'wandb/proto/wandb_settings.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from google.protobuf import wrappers_pb2 as google_dot_protobuf_dot_wrappers__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n wandb/proto/wandb_settings.proto\x12\x0ewandb_internal\x1a\x1egoogle/protobuf/wrappers.proto\" \n\x0fListStringValue\x12\r\n\x05value\x18\x01 \x03(\t\"\x1d\n\x0cListIntValue\x12\r\n\x05value\x18\x01 \x03(\x05\"\x8a\x01\n\x17MapStringKeyStringValue\x12\x41\n\x05value\x18\x01 \x03(\x0b\x32\x32.wandb_internal.MapStringKeyStringValue.ValueEntry\x1a,\n\nValueEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\xcb\x01\n#MapStringKeyMapStringKeyStringValue\x12M\n\x05value\x18\x01 \x03(\x0b\x32>.wandb_internal.MapStringKeyMapStringKeyStringValue.ValueEntry\x1aU\n\nValueEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\x36\n\x05value\x18\x02 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue:\x02\x38\x01\"\x9a\x01\n\x12OpenMetricsFilters\x12\x33\n\x08sequence\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValueH\x00\x12\x46\n\x07mapping\x18\x02 \x01(\x0b\x32\x33.wandb_internal.MapStringKeyMapStringKeyStringValueH\x00\x42\x07\n\x05value\"7\n\tRunMoment\x12\x0b\n\x03run\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\x01\x12\x0e\n\x06metric\x18\x03 \x01(\t\"\x99N\n\x08Settings\x12-\n\x07\x61pi_key\x18\x37 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x13identity_token_file\x18\xaa\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10\x63redentials_file\x18\xab\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x39\n\x14insecure_disable_ssl\x18\xb9\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08_offline\x18\x1e \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06x_sync\x18\x1f \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\tsync_file\x18\x86\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x07_shared\x18\xa2\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x06run_id\x18k \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07run_url\x18q \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07project\x18\x61 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06\x65ntity\x18\x45 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\x0corganization\x18\xbc\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0cx_start_time\x18) \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12.\n\x08root_dir\x18i \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07log_dir\x18U \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0clog_internal\x18V \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\tfiles_dir\x18\x46 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0bx_files_dir\x18\xb4\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0cignore_globs\x18N \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12.\n\x08\x62\x61se_url\x18\x39 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12=\n\x17x_file_stream_max_bytes\x18\xac\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x46\n\x1fx_file_stream_transmit_interval\x18\xaf\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x45\n\x14x_extra_http_headers\x18\x0e \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12=\n\x17x_file_stream_retry_max\x18\x93\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12K\n$x_file_stream_retry_wait_min_seconds\x18\x94\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12K\n$x_file_stream_retry_wait_max_seconds\x18\x95\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x43\n\x1dx_file_stream_timeout_seconds\x18\x0f \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x42\n\x1cx_file_stream_max_line_bytes\x18\xb2\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12?\n\x19x_file_transfer_retry_max\x18\x96\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12M\n&x_file_transfer_retry_wait_min_seconds\x18\x97\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12M\n&x_file_transfer_retry_wait_max_seconds\x18\x98\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x46\n\x1fx_file_transfer_timeout_seconds\x18\x99\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x39\n\x13x_graphql_retry_max\x18\x9a\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12G\n x_graphql_retry_wait_min_seconds\x18\x9b\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12G\n x_graphql_retry_wait_max_seconds\x18\x9c\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12@\n\x19x_graphql_timeout_seconds\x18\x9d\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x31\n\nhttp_proxy\x18\xa8\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0bhttps_proxy\x18\xa9\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12;\n\tx_proxies\x18\xc8\x01 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12-\n\x07program\x18_ \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0fprogram_relpath\x18` \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10_code_path_local\x18\xa3\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x0fprogram_abspath\x18\x9f\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x05_args\x18\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12)\n\x03_os\x18 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06\x64ocker\x18\x43 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0cx_executable\x18\r \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12-\n\x07_python\x18\" \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\tcolab_url\x18\xa0\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x04host\x18M \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x08username\x18\x8d\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x05\x65mail\x18\x44 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x06resume\x18\x66 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0bresume_from\x18\xa7\x01 \x01(\x0b\x32\x19.wandb_internal.RunMoment\x12-\n\tfork_from\x18\xa4\x01 \x01(\x0b\x32\x19.wandb_internal.RunMoment\x12\x38\n\x14\x64isable_job_creation\x18\x41 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\tsweep_url\x18\x83\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12;\n\x16x_disable_update_check\x18\xa5\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0ex_disable_meta\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\tsave_code\x18s \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x0b\x64isable_git\x18? \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x16x_disable_machine_info\x18\x9e\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0fx_disable_stats\x18\n \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x13x_stats_buffer_size\x18\xa1\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12@\n\x19x_stats_sampling_interval\x18\xae\x01 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x30\n\x0bx_stats_pid\x18* \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12<\n\x12x_stats_disk_paths\x18\x92\x01 \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12H\n\"x_stats_neuron_monitor_config_path\x18. \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12<\n\x15x_stats_dcgm_exporter\x18\xbb\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12O\n\x1ex_stats_open_metrics_endpoints\x18/ \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12H\n\x1cx_stats_open_metrics_filters\x18\x30 \x01(\x0b\x32\".wandb_internal.OpenMetricsFilters\x12S\n!x_stats_open_metrics_http_headers\x18\xb8\x01 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12=\n\x16x_stats_gpu_device_ids\x18\xba\x01 \x01(\x0b\x32\x1c.wandb_internal.ListIntValue\x12\x37\n\x11x_stats_cpu_count\x18\xc2\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12?\n\x19x_stats_cpu_logical_count\x18\xc3\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x37\n\x11x_stats_gpu_count\x18\xc4\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x37\n\x10x_stats_gpu_type\x18\xc5\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12?\n\x1ax_stats_track_process_tree\x18\xc6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12.\n\x07x_label\x18\xb5\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\tx_primary\x18\xb6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12:\n\x15x_update_finish_state\x18\xb7\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12<\n\x17\x61llow_offline_artifacts\x18\xb1\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\x07\x63onsole\x18< \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x11\x63onsole_multipart\x18\xa6\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x35\n\x10sync_tensorboard\x18\xb3\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x42\n\x1dx_server_side_derived_summary\x18\xbd\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x46\n!x_server_side_expand_glob_metrics\x18\xbe\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x16x_skip_transaction_log\x18\xbf\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12J\n#x_stats_coreweave_metadata_base_url\x18\xc0\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12J\n#x_stats_coreweave_metadata_endpoint\x18\xc1\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0b_aws_lambda\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0fx_cli_only_mode\x18\x04 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06_colab\x18\x05 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10x_disable_viewer\x18\x0b \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x15x_flow_control_custom\x18\x10 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12;\n\x17x_flow_control_disabled\x18\x11 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12>\n\x18x_internal_check_process\x18\x12 \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12,\n\x08_ipython\x18\x14 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08_jupyter\x18\x15 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x0ex_jupyter_root\x18\x16 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x07_kaggle\x18\x17 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12=\n\x18x_live_policy_rate_limit\x18\x18 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12<\n\x17x_live_policy_wait_time\x18\x19 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x30\n\x0bx_log_level\x18\x1a \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x35\n\x10x_network_buffer\x18\x1b \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12)\n\x05_noop\x18\x1c \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\t_notebook\x18\x1d \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\t_platform\x18! \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x38\n\x12x_runqueue_item_id\x18# \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x13x_save_requirements\x18% \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x39\n\x13x_service_transport\x18& \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x34\n\x0ex_service_wait\x18\' \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12\x35\n\x0f_start_datetime\x18( \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\r_tmp_code_dir\x18\x31 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x08_windows\x18\x34 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x34\n\x10\x61llow_val_change\x18\x35 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\tanonymous\x18\x36 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12P\n\x1f\x61zure_account_url_to_access_key\x18\x38 \x01(\x0b\x32\'.wandb_internal.MapStringKeyStringValue\x12.\n\x08\x63ode_dir\x18: \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0c\x63onfig_paths\x18; \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12\x30\n\ndeployment\x18= \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\x0c\x64isable_code\x18> \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x31\n\rdisable_hints\x18@ \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12,\n\x08\x64isabled\x18\x42 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12)\n\x05\x66orce\x18G \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\ngit_commit\x18H \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\ngit_remote\x18I \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x34\n\x0egit_remote_url\x18J \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08git_root\x18K \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x11heartbeat_seconds\x18L \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x32\n\x0cinit_timeout\x18O \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12,\n\x08is_local\x18P \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x30\n\njob_source\x18Q \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\rlabel_disable\x18R \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06launch\x18S \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x38\n\x12launch_config_path\x18T \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x14log_symlink_internal\x18W \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x36\n\x10log_symlink_user\x18X \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08log_user\x18Y \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\rlogin_timeout\x18Z \x01(\x0b\x32\x1c.google.protobuf.DoubleValue\x12*\n\x04mode\x18\\ \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x33\n\rnotebook_name\x18] \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\x0bproject_url\x18\x62 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12)\n\x05quiet\x18\x63 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12+\n\x07relogin\x18\x65 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0cresume_fname\x18g \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12+\n\x07resumed\x18h \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\trun_group\x18j \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x32\n\x0crun_job_type\x18l \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08run_mode\x18m \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x08run_name\x18n \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\trun_notes\x18o \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x31\n\x08run_tags\x18p \x01(\x0b\x32\x1f.wandb_internal.ListStringValue\x12\x35\n\x11sagemaker_disable\x18r \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x35\n\x0fsettings_system\x18t \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x38\n\x12settings_workspace\x18u \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x0bshow_colors\x18v \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12.\n\nshow_emoji\x18w \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x0bshow_errors\x18x \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12-\n\tshow_info\x18y \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x31\n\rshow_warnings\x18z \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12*\n\x06silent\x18{ \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x32\n\x0cstart_method\x18| \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12*\n\x06strict\x18} \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12\x33\n\x0esummary_errors\x18~ \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x34\n\x0fsummary_timeout\x18\x7f \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12\x36\n\x10summary_warnings\x18\x80\x01 \x01(\x0b\x32\x1b.google.protobuf.Int32Value\x12/\n\x08sweep_id\x18\x81\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x37\n\x10sweep_param_path\x18\x82\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12,\n\x07symlink\x18\x84\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x08sync_dir\x18\x85\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12:\n\x13sync_symlink_latest\x18\x87\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12J\n%table_raise_on_max_row_limit_exceeded\x18\x8a\x01 \x01(\x0b\x32\x1a.google.protobuf.BoolValue\x12/\n\x08timespec\x18\x8b\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12.\n\x07tmp_dir\x18\x8c\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x30\n\twandb_dir\x18\x8e\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0ex_jupyter_name\x18\x8f\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12\x35\n\x0ex_jupyter_path\x18\x90\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValue\x12/\n\x08job_name\x18\x91\x01 \x01(\x0b\x32\x1c.google.protobuf.StringValueJ\x04\x08\x03\x10\x04J\x04\x08\x06\x10\x07J\x04\x08\x08\x10\tJ\x04\x08\t\x10\nJ\x04\x08\x0c\x10\rJ\x04\x08\x13\x10\x14J\x04\x08$\x10%J\x04\x08+\x10,J\x04\x08,\x10-J\x04\x08-\x10.J\x04\x08\x32\x10\x33J\x04\x08\x33\x10\x34J\x04\x08[\x10\\J\x04\x08^\x10_J\x04\x08\x64\x10\x65J\x06\x08\x88\x01\x10\x89\x01J\x06\x08\x89\x01\x10\x8a\x01J\x06\x08\xad\x01\x10\xae\x01J\x06\x08\xb0\x01\x10\xb1\x01\x42\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_settings_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._loaded_options = None + _globals['_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_options = b'8\001' + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._loaded_options = None + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_options = b'8\001' + _globals['_LISTSTRINGVALUE']._serialized_start=84 + _globals['_LISTSTRINGVALUE']._serialized_end=116 + _globals['_LISTINTVALUE']._serialized_start=118 + _globals['_LISTINTVALUE']._serialized_end=147 + _globals['_MAPSTRINGKEYSTRINGVALUE']._serialized_start=150 + _globals['_MAPSTRINGKEYSTRINGVALUE']._serialized_end=288 + _globals['_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_start=244 + _globals['_MAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_end=288 + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE']._serialized_start=291 + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE']._serialized_end=494 + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_start=409 + _globals['_MAPSTRINGKEYMAPSTRINGKEYSTRINGVALUE_VALUEENTRY']._serialized_end=494 + _globals['_OPENMETRICSFILTERS']._serialized_start=497 + _globals['_OPENMETRICSFILTERS']._serialized_end=651 + _globals['_RUNMOMENT']._serialized_start=653 + _globals['_RUNMOMENT']._serialized_end=708 + _globals['_SETTINGS']._serialized_start=711 + _globals['_SETTINGS']._serialized_end=10720 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/v6/wandb_telemetry_pb2.py b/lib/python3.12/site-packages/wandb/proto/v6/wandb_telemetry_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..994c57a537e1823ef1908a88901e1be5f8754951 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/v6/wandb_telemetry_pb2.py @@ -0,0 +1,52 @@ +# -*- coding: utf-8 -*- +# Generated by the protocol buffer compiler. DO NOT EDIT! +# NO CHECKED-IN PROTOBUF GENCODE +# source: wandb/proto/wandb_telemetry.proto +# Protobuf Python Version: 6.30.0 +"""Generated protocol buffer code.""" +from google.protobuf import descriptor as _descriptor +from google.protobuf import descriptor_pool as _descriptor_pool +from google.protobuf import runtime_version as _runtime_version +from google.protobuf import symbol_database as _symbol_database +from google.protobuf.internal import builder as _builder +_runtime_version.ValidateProtobufRuntimeVersion( + _runtime_version.Domain.PUBLIC, + 6, + 30, + 0, + '', + 'wandb/proto/wandb_telemetry.proto' +) +# @@protoc_insertion_point(imports) + +_sym_db = _symbol_database.Default() + + +from wandb.proto import wandb_base_pb2 as wandb_dot_proto_dot_wandb__base__pb2 + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n!wandb/proto/wandb_telemetry.proto\x12\x0ewandb_internal\x1a\x1cwandb/proto/wandb_base.proto\"\xdb\x03\n\x0fTelemetryRecord\x12-\n\x0cimports_init\x18\x01 \x01(\x0b\x32\x17.wandb_internal.Imports\x12/\n\x0eimports_finish\x18\x02 \x01(\x0b\x32\x17.wandb_internal.Imports\x12(\n\x07\x66\x65\x61ture\x18\x03 \x01(\x0b\x32\x17.wandb_internal.Feature\x12\x16\n\x0epython_version\x18\x04 \x01(\t\x12\x13\n\x0b\x63li_version\x18\x05 \x01(\t\x12\x1b\n\x13huggingface_version\x18\x06 \x01(\t\x12 \n\x03\x65nv\x18\x08 \x01(\x0b\x32\x13.wandb_internal.Env\x12%\n\x05label\x18\t \x01(\x0b\x32\x16.wandb_internal.Labels\x12.\n\ndeprecated\x18\n \x01(\x0b\x32\x1a.wandb_internal.Deprecated\x12&\n\x06issues\x18\x0b \x01(\x0b\x32\x16.wandb_internal.Issues\x12\x14\n\x0c\x63ore_version\x18\x0c \x01(\t\x12\x10\n\x08platform\x18\r \x01(\t\x12+\n\x05_info\x18\xc8\x01 \x01(\x0b\x32\x1b.wandb_internal._RecordInfo\"\x11\n\x0fTelemetryResult\"\x9a\x0e\n\x07Imports\x12\r\n\x05torch\x18\x01 \x01(\x08\x12\r\n\x05keras\x18\x02 \x01(\x08\x12\x12\n\ntensorflow\x18\x03 \x01(\x08\x12\x0e\n\x06\x66\x61stai\x18\x04 \x01(\x08\x12\x0f\n\x07sklearn\x18\x05 \x01(\x08\x12\x0f\n\x07xgboost\x18\x06 \x01(\x08\x12\x10\n\x08\x63\x61tboost\x18\x07 \x01(\x08\x12\x10\n\x08lightgbm\x18\x08 \x01(\x08\x12\x19\n\x11pytorch_lightning\x18\t \x01(\x08\x12\x0e\n\x06ignite\x18\n \x01(\x08\x12\x14\n\x0ctransformers\x18\x0b \x01(\x08\x12\x0b\n\x03jax\x18\x0c \x01(\x08\x12\x10\n\x08metaflow\x18\r \x01(\x08\x12\x10\n\x08\x61llennlp\x18\x0e \x01(\x08\x12\x11\n\tautogluon\x18\x0f \x01(\x08\x12\x11\n\tautokeras\x18\x10 \x01(\x08\x12\x10\n\x08\x63\x61talyst\x18\x12 \x01(\x08\x12\x10\n\x08\x64\x65\x65pchem\x18\x15 \x01(\x08\x12\x0f\n\x07\x64\x65\x65pctr\x18\x16 \x01(\x08\x12\x0f\n\x07pycaret\x18\x1c \x01(\x08\x12\x14\n\x0cpytorchvideo\x18\x1d \x01(\x08\x12\x0b\n\x03ray\x18\x1e \x01(\x08\x12\x1a\n\x12simpletransformers\x18\x1f \x01(\x08\x12\x0e\n\x06skorch\x18 \x01(\x08\x12\r\n\x05spacy\x18! \x01(\x08\x12\r\n\x05\x66lash\x18\" \x01(\x08\x12\x0e\n\x06optuna\x18# \x01(\x08\x12\x0f\n\x07recbole\x18$ \x01(\x08\x12\x0c\n\x04mmcv\x18% \x01(\x08\x12\r\n\x05mmdet\x18& \x01(\x08\x12\x11\n\ttorchdrug\x18\' \x01(\x08\x12\x11\n\ttorchtext\x18( \x01(\x08\x12\x13\n\x0btorchvision\x18) \x01(\x08\x12\r\n\x05\x65legy\x18* \x01(\x08\x12\x12\n\ndetectron2\x18+ \x01(\x08\x12\r\n\x05\x66lair\x18, \x01(\x08\x12\x0c\n\x04\x66lax\x18- \x01(\x08\x12\x0c\n\x04syft\x18. \x01(\x08\x12\x0b\n\x03TTS\x18/ \x01(\x08\x12\r\n\x05monai\x18\x30 \x01(\x08\x12\x17\n\x0fhuggingface_hub\x18\x31 \x01(\x08\x12\r\n\x05hydra\x18\x32 \x01(\x08\x12\x10\n\x08\x64\x61tasets\x18\x33 \x01(\x08\x12\x0e\n\x06sacred\x18\x34 \x01(\x08\x12\x0e\n\x06joblib\x18\x35 \x01(\x08\x12\x0c\n\x04\x64\x61sk\x18\x36 \x01(\x08\x12\x11\n\tpaddleocr\x18\x38 \x01(\x08\x12\r\n\x05ppdet\x18\x39 \x01(\x08\x12\x11\n\tpaddleseg\x18: \x01(\x08\x12\x11\n\tpaddlenlp\x18; \x01(\x08\x12\r\n\x05mmseg\x18< \x01(\x08\x12\r\n\x05mmocr\x18= \x01(\x08\x12\r\n\x05mmcls\x18> \x01(\x08\x12\x0c\n\x04timm\x18? \x01(\x08\x12\x0f\n\x07\x66\x61irseq\x18@ \x01(\x08\x12\x12\n\ndeepchecks\x18\x41 \x01(\x08\x12\x10\n\x08\x63omposer\x18\x42 \x01(\x08\x12\x10\n\x08sparseml\x18\x43 \x01(\x08\x12\x10\n\x08\x61nomalib\x18\x44 \x01(\x08\x12\r\n\x05zenml\x18\x45 \x01(\x08\x12\x12\n\ncolossalai\x18\x46 \x01(\x08\x12\x12\n\naccelerate\x18G \x01(\x08\x12\x0e\n\x06merlin\x18H \x01(\x08\x12\x0f\n\x07nanodet\x18I \x01(\x08\x12#\n\x1bsegmentation_models_pytorch\x18J \x01(\x08\x12\x1d\n\x15sentence_transformers\x18K \x01(\x08\x12\x0b\n\x03\x64gl\x18L \x01(\x08\x12\x17\n\x0ftorch_geometric\x18M \x01(\x08\x12\x0c\n\x04jina\x18N \x01(\x08\x12\x0e\n\x06kornia\x18O \x01(\x08\x12\x16\n\x0e\x61lbumentations\x18P \x01(\x08\x12\x10\n\x08keras_cv\x18Q \x01(\x08\x12\x10\n\x08mmengine\x18R \x01(\x08\x12\x11\n\tdiffusers\x18S \x01(\x08\x12\x0b\n\x03trl\x18T \x01(\x08\x12\x0c\n\x04trlx\x18U \x01(\x08\x12\x11\n\tlangchain\x18V \x01(\x08\x12\x13\n\x0bllama_index\x18W \x01(\x08\x12\x15\n\rstability_sdk\x18X \x01(\x08\x12\x0f\n\x07prefect\x18Y \x01(\x08\x12\x13\n\x0bprefect_ray\x18Z \x01(\x08\x12\x10\n\x08pinecone\x18[ \x01(\x08\x12\x10\n\x08\x63hromadb\x18\\ \x01(\x08\x12\x10\n\x08weaviate\x18] \x01(\x08\x12\x13\n\x0bpromptlayer\x18^ \x01(\x08\x12\x0e\n\x06openai\x18_ \x01(\x08\x12\x0e\n\x06\x63ohere\x18` \x01(\x08\x12\x11\n\tanthropic\x18\x61 \x01(\x08\x12\x0c\n\x04peft\x18\x62 \x01(\x08\x12\x0f\n\x07optimum\x18\x63 \x01(\x08\x12\x10\n\x08\x65valuate\x18\x64 \x01(\x08\x12\x10\n\x08langflow\x18\x65 \x01(\x08\x12\x12\n\nkeras_core\x18\x66 \x01(\x08\x12\x18\n\x10lightning_fabric\x18g \x01(\x08\x12\x1c\n\x14\x63urated_transformers\x18h \x01(\x08\x12\x0e\n\x06orjson\x18i \x01(\x08\x12\x11\n\tlightning\x18j \x01(\x08J\x04\x08\x11\x10\x12J\x04\x08\x13\x10\x14J\x04\x08\x14\x10\x15J\x04\x08\x17\x10\x18J\x04\x08\x18\x10\x19J\x04\x08\x19\x10\x1aJ\x04\x08\x1a\x10\x1bJ\x04\x08\x1b\x10\x1cJ\x04\x08\x37\x10\x38\"\x95\r\n\x07\x46\x65\x61ture\x12\r\n\x05watch\x18\x01 \x01(\x08\x12\x0e\n\x06\x66inish\x18\x02 \x01(\x08\x12\x0c\n\x04save\x18\x03 \x01(\x08\x12\x0f\n\x07offline\x18\x04 \x01(\x08\x12\x0f\n\x07resumed\x18\x05 \x01(\x08\x12\x0c\n\x04grpc\x18\x06 \x01(\x08\x12\x0e\n\x06metric\x18\x07 \x01(\x08\x12\r\n\x05keras\x18\x08 \x01(\x08\x12\x11\n\tsagemaker\x18\t \x01(\x08\x12\x1c\n\x14\x61rtifact_incremental\x18\n \x01(\x08\x12\x10\n\x08metaflow\x18\x0b \x01(\x08\x12\x0f\n\x07prodigy\x18\x0c \x01(\x08\x12\x15\n\rset_init_name\x18\r \x01(\x08\x12\x13\n\x0bset_init_id\x18\x0e \x01(\x08\x12\x15\n\rset_init_tags\x18\x0f \x01(\x08\x12\x17\n\x0fset_init_config\x18\x10 \x01(\x08\x12\x14\n\x0cset_run_name\x18\x11 \x01(\x08\x12\x14\n\x0cset_run_tags\x18\x12 \x01(\x08\x12\x17\n\x0fset_config_item\x18\x13 \x01(\x08\x12\x0e\n\x06launch\x18\x14 \x01(\x08\x12\x1c\n\x14torch_profiler_trace\x18\x15 \x01(\x08\x12\x0b\n\x03sb3\x18\x16 \x01(\x08\x12\x17\n\x0finit_return_run\x18\x18 \x01(\x08\x12\x1f\n\x17lightgbm_wandb_callback\x18\x19 \x01(\x08\x12\x1c\n\x14lightgbm_log_summary\x18\x1a \x01(\x08\x12\x1f\n\x17\x63\x61tboost_wandb_callback\x18\x1b \x01(\x08\x12\x1c\n\x14\x63\x61tboost_log_summary\x18\x1c \x01(\x08\x12\x17\n\x0ftensorboard_log\x18\x1d \x01(\x08\x12\x16\n\x0e\x65stimator_hook\x18\x1e \x01(\x08\x12\x1e\n\x16xgboost_wandb_callback\x18\x1f \x01(\x08\x12\"\n\x1axgboost_old_wandb_callback\x18 \x01(\x08\x12\x0e\n\x06\x61ttach\x18! \x01(\x08\x12\x19\n\x11tensorboard_patch\x18\" \x01(\x08\x12\x18\n\x10tensorboard_sync\x18# \x01(\x08\x12\x15\n\rkfp_wandb_log\x18$ \x01(\x08\x12\x1b\n\x13maybe_run_overwrite\x18% \x01(\x08\x12\x1c\n\x14keras_metrics_logger\x18& \x01(\x08\x12\x1e\n\x16keras_model_checkpoint\x18\' \x01(\x08\x12!\n\x19keras_wandb_eval_callback\x18( \x01(\x08\x12\x1d\n\x15\x66low_control_overflow\x18) \x01(\x08\x12\x0c\n\x04sync\x18* \x01(\x08\x12\x1d\n\x15\x66low_control_disabled\x18+ \x01(\x08\x12\x1b\n\x13\x66low_control_custom\x18, \x01(\x08\x12\x14\n\x0copen_metrics\x18. \x01(\x08\x12\x1a\n\x12ultralytics_yolov8\x18/ \x01(\x08\x12\x17\n\x0fimporter_mlflow\x18\x30 \x01(\x08\x12\x15\n\rsync_tfevents\x18\x31 \x01(\x08\x12\x15\n\rasync_uploads\x18\x32 \x01(\x08\x12\x16\n\x0eopenai_autolog\x18\x33 \x01(\x08\x12\x18\n\x10langchain_tracer\x18\x34 \x01(\x08\x12\x16\n\x0e\x63ohere_autolog\x18\x35 \x01(\x08\x12\x1b\n\x13hf_pipeline_autolog\x18\x36 \x01(\x08\x12\r\n\x05lib_c\x18\x38 \x01(\x08\x12\x0f\n\x07lib_cpp\x18\x39 \x01(\x08\x12\x19\n\x11openai_finetuning\x18: \x01(\x08\x12\x19\n\x11\x64iffusers_autolog\x18; \x01(\x08\x12\x1f\n\x17lightning_fabric_logger\x18< \x01(\x08\x12\x14\n\x0cset_step_log\x18= \x01(\x08\x12\x13\n\x0bset_summary\x18> \x01(\x08\x12\x16\n\x0emetric_summary\x18? \x01(\x08\x12\x13\n\x0bmetric_goal\x18@ \x01(\x08\x12\x15\n\rmetric_hidden\x18\x41 \x01(\x08\x12\x18\n\x10metric_step_sync\x18\x42 \x01(\x08\x12\x13\n\x0bshared_mode\x18\x43 \x01(\x08\x12#\n\x1bserver_side_derived_summary\x18\x44 \x01(\x08\x12\x1b\n\x13user_provided_label\x18\x45 \x01(\x08\x12\x1e\n\x16\x64\x63gm_profiling_enabled\x18\x46 \x01(\x08\x12\x11\n\tfork_mode\x18G \x01(\x08\x12\x13\n\x0brewind_mode\x18H \x01(\x08J\x04\x08\x17\x10\x18J\x04\x08-\x10.J\x04\x08\x37\x10\x38\"\xc9\x01\n\x03\x45nv\x12\x0f\n\x07jupyter\x18\x01 \x01(\x08\x12\x0e\n\x06kaggle\x18\x02 \x01(\x08\x12\x0f\n\x07windows\x18\x03 \x01(\x08\x12\x0e\n\x06m1_gpu\x18\x04 \x01(\x08\x12\x10\n\x08trainium\x18\n \x01(\x08\x12\x0b\n\x03pex\x18\x0b \x01(\x08\x12\r\n\x05\x63olab\x18\x0c \x01(\x08\x12\x0f\n\x07ipython\x18\r \x01(\x08\x12\x12\n\naws_lambda\x18\x0e \x01(\x08\x12\x0f\n\x07\x61md_gpu\x18\x0f \x01(\x08J\x04\x08\x05\x10\x06J\x04\x08\x06\x10\x07J\x04\x08\x07\x10\x08J\x04\x08\x08\x10\tJ\x04\x08\t\x10\n\"H\n\x06Labels\x12\x13\n\x0b\x63ode_string\x18\x01 \x01(\t\x12\x13\n\x0brepo_string\x18\x02 \x01(\t\x12\x14\n\x0c\x63ode_version\x18\x03 \x01(\t\"\xd5\x05\n\nDeprecated\x12!\n\x19keras_callback__data_type\x18\x01 \x01(\x08\x12\r\n\x05plots\x18\x05 \x01(\x08\x12!\n\x19init__config_include_keys\x18\x07 \x01(\x08\x12!\n\x19init__config_exclude_keys\x18\x08 \x01(\x08\x12\"\n\x1akeras_callback__save_model\x18\t \x01(\x08\x12\x18\n\x10langchain_tracer\x18\n \x01(\x08\x12\x1a\n\x12\x61rtifact__get_path\x18\x0b \x01(\x08\x12#\n\x1b\x61rtifactmanifestentry__name\x18\x0c \x01(\x08\x12\x1e\n\x16\x61pi__artifact_versions\x18\r \x01(\x08\x12(\n artifact_collection__change_type\x18\x0e \x01(\x08\x12\x1f\n\x17run__define_metric_copy\x18\x0f \x01(\x08\x12\x14\n\x0crun_disabled\x18\x10 \x01(\x08\x12\x16\n\x0ekeras_callback\x18\x11 \x01(\x08\x12$\n\x1crun__define_metric_best_goal\x18\x12 \x01(\x08\x12\x19\n\x11run__finish_quiet\x18\x13 \x01(\x08\x12\x18\n\x10run__reinit_bool\x18\x14 \x01(\x08\x12\x14\n\x0crun__get_url\x18\x15 \x01(\x08\x12\x19\n\x11run__project_name\x18\x16 \x01(\x08\x12\x1c\n\x14run__get_project_url\x18\x17 \x01(\x08\x12\x1a\n\x12run__get_sweep_url\x18\x18 \x01(\x08\x12 \n\x18run__use_artifact_use_as\x18\x19 \x01(\x08\x12\x18\n\x10\x61rtifact__use_as\x18\x1a \x01(\x08\x12\x1d\n\x15\x61rtifact__init_use_as\x18\x1b \x01(\x08J\x04\x08\x02\x10\x03J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05J\x04\x08\x06\x10\x07\"|\n\x06Issues\x12%\n\x1dsettings__validation_warnings\x18\x01 \x01(\x08\x12!\n\x19settings__unexpected_args\x18\x02 \x01(\x08\x12(\n settings__preprocessing_warnings\x18\x03 \x01(\x08\x42\x1bZ\x19\x63ore/pkg/service_go_protob\x06proto3') + +_globals = globals() +_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'wandb.proto.wandb_telemetry_pb2', _globals) +if not _descriptor._USE_C_DESCRIPTORS: + _globals['DESCRIPTOR']._loaded_options = None + _globals['DESCRIPTOR']._serialized_options = b'Z\031core/pkg/service_go_proto' + _globals['_TELEMETRYRECORD']._serialized_start=84 + _globals['_TELEMETRYRECORD']._serialized_end=559 + _globals['_TELEMETRYRESULT']._serialized_start=561 + _globals['_TELEMETRYRESULT']._serialized_end=578 + _globals['_IMPORTS']._serialized_start=581 + _globals['_IMPORTS']._serialized_end=2399 + _globals['_FEATURE']._serialized_start=2402 + _globals['_FEATURE']._serialized_end=4087 + _globals['_ENV']._serialized_start=4090 + _globals['_ENV']._serialized_end=4291 + _globals['_LABELS']._serialized_start=4293 + _globals['_LABELS']._serialized_end=4365 + _globals['_DEPRECATED']._serialized_start=4368 + _globals['_DEPRECATED']._serialized_end=5093 + _globals['_ISSUES']._serialized_start=5095 + _globals['_ISSUES']._serialized_end=5219 +# @@protoc_insertion_point(module_scope) diff --git a/lib/python3.12/site-packages/wandb/proto/wandb_base_pb2.py b/lib/python3.12/site-packages/wandb/proto/wandb_base_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..8b2ae6b240436e0da0d94ffda7c4a150d24dc54f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/wandb_base_pb2.py @@ -0,0 +1,12 @@ +import google.protobuf + +protobuf_version = google.protobuf.__version__[0] + +if protobuf_version == "3": + from wandb.proto.v3.wandb_base_pb2 import * +elif protobuf_version == "4": + from wandb.proto.v4.wandb_base_pb2 import * +elif protobuf_version == "5": + from wandb.proto.v5.wandb_base_pb2 import * +elif protobuf_version == "6": + from wandb.proto.v6.wandb_base_pb2 import * diff --git a/lib/python3.12/site-packages/wandb/proto/wandb_deprecated.py b/lib/python3.12/site-packages/wandb/proto/wandb_deprecated.py new file mode 100644 index 0000000000000000000000000000000000000000..8b0e37ac535ccc71d1c34dc07eec7925c29b433d --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/wandb_deprecated.py @@ -0,0 +1,53 @@ +# Generated by wandb/proto/wandb_internal_codegen.py. DO NOT EDIT! + +from typing import Literal +DEPRECATED_FEATURES = Literal[ + "keras_callback__data_type", + "plots", + "init__config_include_keys", + "init__config_exclude_keys", + "keras_callback__save_model", + "langchain_tracer", + "artifact__get_path", + "artifactmanifestentry__name", + "api__artifact_versions", + "artifact_collection__change_type", + "run__define_metric_copy", + "run_disabled", + "keras_callback", + "run__define_metric_best_goal", + "run__finish_quiet", + "run__reinit_bool", + "run__get_url", + "run__project_name", + "run__get_project_url", + "run__get_sweep_url", + "run__use_artifact_use_as", + "artifact__use_as", + "artifact__init_use_as", +] + +class Deprecated: + keras_callback__data_type: DEPRECATED_FEATURES = "keras_callback__data_type" + plots: DEPRECATED_FEATURES = "plots" + init__config_include_keys: DEPRECATED_FEATURES = "init__config_include_keys" + init__config_exclude_keys: DEPRECATED_FEATURES = "init__config_exclude_keys" + keras_callback__save_model: DEPRECATED_FEATURES = "keras_callback__save_model" + langchain_tracer: DEPRECATED_FEATURES = "langchain_tracer" + artifact__get_path: DEPRECATED_FEATURES = "artifact__get_path" + artifactmanifestentry__name: DEPRECATED_FEATURES = "artifactmanifestentry__name" + api__artifact_versions: DEPRECATED_FEATURES = "api__artifact_versions" + artifact_collection__change_type: DEPRECATED_FEATURES = "artifact_collection__change_type" + run__define_metric_copy: DEPRECATED_FEATURES = "run__define_metric_copy" + run_disabled: DEPRECATED_FEATURES = "run_disabled" + keras_callback: DEPRECATED_FEATURES = "keras_callback" + run__define_metric_best_goal: DEPRECATED_FEATURES = "run__define_metric_best_goal" + run__finish_quiet: DEPRECATED_FEATURES = "run__finish_quiet" + run__reinit_bool: DEPRECATED_FEATURES = "run__reinit_bool" + run__get_url: DEPRECATED_FEATURES = "run__get_url" + run__project_name: DEPRECATED_FEATURES = "run__project_name" + run__get_project_url: DEPRECATED_FEATURES = "run__get_project_url" + run__get_sweep_url: DEPRECATED_FEATURES = "run__get_sweep_url" + run__use_artifact_use_as: DEPRECATED_FEATURES = "run__use_artifact_use_as" + artifact__use_as: DEPRECATED_FEATURES = "artifact__use_as" + artifact__init_use_as: DEPRECATED_FEATURES = "artifact__init_use_as" diff --git a/lib/python3.12/site-packages/wandb/proto/wandb_generate_deprecated.py b/lib/python3.12/site-packages/wandb/proto/wandb_generate_deprecated.py new file mode 100644 index 0000000000000000000000000000000000000000..23240018155a6bdf1f933439382c4aefc9d479a5 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/wandb_generate_deprecated.py @@ -0,0 +1,30 @@ +#!/usr/bin/env python + + +def generate_deprecated_class_definition() -> None: + """Generate a class definition listing the deprecated features. + This is to allow static checks to ensure that proper field names are used. + """ + from wandb.proto.wandb_telemetry_pb2 import Deprecated # type: ignore[import] + + deprecated_features = Deprecated.DESCRIPTOR.fields_by_name.keys() + + code: str = ( + "# Generated by wandb/proto/wandb_internal_codegen.py. DO NOT EDIT!\n\n" + "from typing import Literal\n" + "DEPRECATED_FEATURES = Literal[\n" + + ",\n".join(f' "{feature}"' for feature in deprecated_features) + + ",\n" + + "]\n\n" + "class Deprecated:\n" + + "".join( + [ + f' {feature}: DEPRECATED_FEATURES = "{feature}"\n' + for feature in deprecated_features + ] + ) + ) + with open("wandb_deprecated.py", "w") as f: + f.write(code) + +generate_deprecated_class_definition() diff --git a/lib/python3.12/site-packages/wandb/proto/wandb_generate_proto.py b/lib/python3.12/site-packages/wandb/proto/wandb_generate_proto.py new file mode 100644 index 0000000000000000000000000000000000000000..7dd8979f0b4929e43b9f529519423589dc2fbba0 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/wandb_generate_proto.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python + +import importlib.metadata +import os +import pathlib + +import grpc_tools # type: ignore +from grpc_tools import protoc # type: ignore +from packaging import version + + +def get_pip_package_version(package_name: str) -> str: + try: + return importlib.metadata.version(package_name) + except importlib.metadata.PackageNotFoundError: + raise ValueError(f"Package `{package_name}` not found") + +protobuf_version = version.Version(get_pip_package_version("protobuf")) + +proto_root = os.path.join(os.path.dirname(grpc_tools.__file__), "_proto") +tmp_out: pathlib.Path = pathlib.Path(f"wandb/proto/v{protobuf_version.major}/") + +os.chdir("../..") +for proto_file in [ + "wandb_base.proto", + "wandb_internal.proto", + "wandb_settings.proto", + "wandb_telemetry.proto", + "wandb_server.proto", +]: + ret = protoc.main( + ( + "", + "-I", + proto_root, + "-I", + ".", + f"--python_out={tmp_out}", + f"--mypy_out={tmp_out}", + f"wandb/proto/{proto_file}", + ) + ) + assert not ret + +# clean up tmp dirs +for p in (tmp_out / "wandb" / "proto").glob("*pb2*"): + p.rename(tmp_out / p.name) +os.rmdir(tmp_out / "wandb" / "proto") +os.rmdir(tmp_out / "wandb") diff --git a/lib/python3.12/site-packages/wandb/proto/wandb_internal_pb2.py b/lib/python3.12/site-packages/wandb/proto/wandb_internal_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..97cbe72e2148955f2aaa357053905e81ccfc5825 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/wandb_internal_pb2.py @@ -0,0 +1,18 @@ +import google.protobuf + +protobuf_version = google.protobuf.__version__[0] + +if protobuf_version == "3": + from wandb.proto.v3.wandb_internal_pb2 import * +elif protobuf_version == "4": + from wandb.proto.v4.wandb_internal_pb2 import * +elif protobuf_version == "5": + from wandb.proto.v5.wandb_internal_pb2 import * +elif protobuf_version == "6": + from wandb.proto.v6.wandb_internal_pb2 import * +else: + raise ImportError( + "Failed to import protobufs for protobuf version" + f" {google.protobuf.__version__}. `wandb` only works with major" + " versions 3, 4, 5, and 6 of the protobuf package.", + ) diff --git a/lib/python3.12/site-packages/wandb/proto/wandb_server_pb2.py b/lib/python3.12/site-packages/wandb/proto/wandb_server_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..d30616eb0902a99ce4180d7760a975b3e4383bc7 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/wandb_server_pb2.py @@ -0,0 +1,12 @@ +import google.protobuf + +protobuf_version = google.protobuf.__version__[0] + +if protobuf_version == "3": + from wandb.proto.v3.wandb_server_pb2 import * +elif protobuf_version == "4": + from wandb.proto.v4.wandb_server_pb2 import * +elif protobuf_version == "5": + from wandb.proto.v5.wandb_server_pb2 import * +elif protobuf_version == "6": + from wandb.proto.v6.wandb_server_pb2 import * diff --git a/lib/python3.12/site-packages/wandb/proto/wandb_settings_pb2.py b/lib/python3.12/site-packages/wandb/proto/wandb_settings_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..0600cc522c3e7dae1222d19779652da44010a100 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/wandb_settings_pb2.py @@ -0,0 +1,12 @@ +import google.protobuf + +protobuf_version = google.protobuf.__version__[0] + +if protobuf_version == "3": + from wandb.proto.v3.wandb_settings_pb2 import * +elif protobuf_version == "4": + from wandb.proto.v4.wandb_settings_pb2 import * +elif protobuf_version == "5": + from wandb.proto.v5.wandb_settings_pb2 import * +elif protobuf_version == "6": + from wandb.proto.v6.wandb_settings_pb2 import * diff --git a/lib/python3.12/site-packages/wandb/proto/wandb_telemetry_pb2.py b/lib/python3.12/site-packages/wandb/proto/wandb_telemetry_pb2.py new file mode 100644 index 0000000000000000000000000000000000000000..19c81ce3f03a8906b4108f3f9f5d602d2b84f792 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/proto/wandb_telemetry_pb2.py @@ -0,0 +1,12 @@ +import google.protobuf + +protobuf_version = google.protobuf.__version__[0] + +if protobuf_version == "3": + from wandb.proto.v3.wandb_telemetry_pb2 import * +elif protobuf_version == "4": + from wandb.proto.v4.wandb_telemetry_pb2 import * +elif protobuf_version == "5": + from wandb.proto.v5.wandb_telemetry_pb2 import * +elif protobuf_version == "6": + from wandb.proto.v6.wandb_telemetry_pb2 import * diff --git a/lib/python3.12/site-packages/wandb/py.typed b/lib/python3.12/site-packages/wandb/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/wandb/sklearn.py b/lib/python3.12/site-packages/wandb/sklearn.py new file mode 100644 index 0000000000000000000000000000000000000000..5c72b8228b5cb342ba2a217104ee5ce9267a2d2f --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sklearn.py @@ -0,0 +1,35 @@ +from wandb.integration.sklearn import ( + plot_calibration_curve, + plot_class_proportions, + plot_classifier, + plot_clusterer, + plot_confusion_matrix, + plot_elbow_curve, + plot_feature_importances, + plot_learning_curve, + plot_outlier_candidates, + plot_precision_recall, + plot_regressor, + plot_residuals, + plot_roc, + plot_silhouette, + plot_summary_metrics, +) + +__all__ = ( + "plot_classifier", + "plot_clusterer", + "plot_regressor", + "plot_summary_metrics", + "plot_learning_curve", + "plot_feature_importances", + "plot_class_proportions", + "plot_calibration_curve", + "plot_roc", + "plot_precision_recall", + "plot_confusion_matrix", + "plot_elbow_curve", + "plot_silhouette", + "plot_residuals", + "plot_outlier_candidates", +) diff --git a/lib/python3.12/site-packages/wandb/sync/__init__.py b/lib/python3.12/site-packages/wandb/sync/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..92213239a7fd075123dfe49928241b99640abc0b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sync/__init__.py @@ -0,0 +1,3 @@ +"""module sync.""" + +from .sync import SyncManager, get_run_from_path, get_runs diff --git a/lib/python3.12/site-packages/wandb/sync/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sync/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..64d6a1b936a2a26a6f8fc58cb5146d430e5e92b1 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sync/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sync/__pycache__/sync.cpython-312.pyc b/lib/python3.12/site-packages/wandb/sync/__pycache__/sync.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5eea95226bbcecb0c71060bc693e7f91d917cff0 Binary files /dev/null and b/lib/python3.12/site-packages/wandb/sync/__pycache__/sync.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/wandb/sync/sync.py b/lib/python3.12/site-packages/wandb/sync/sync.py new file mode 100644 index 0000000000000000000000000000000000000000..51ea6908696ffee0e925b6f518ca5d282420511b --- /dev/null +++ b/lib/python3.12/site-packages/wandb/sync/sync.py @@ -0,0 +1,442 @@ +"""sync.""" + +import atexit +import datetime +import fnmatch +import os +import queue +import sys +import tempfile +import threading +import time +from typing import List, Optional +from urllib.parse import quote as url_quote + +import wandb +from wandb.proto import wandb_internal_pb2 # type: ignore +from wandb.sdk.interface.interface_queue import InterfaceQueue +from wandb.sdk.internal import context, datastore, handler, sender, tb_watcher +from wandb.sdk.internal.settings_static import SettingsStatic +from wandb.sdk.lib import filesystem +from wandb.util import check_and_warn_old + +WANDB_SUFFIX = ".wandb" +SYNCED_SUFFIX = ".synced" +TFEVENT_SUBSTRING = ".tfevents." + + +class _LocalRun: + def __init__(self, path, synced=None): + self.path = path + self.synced = synced + self.offline = os.path.basename(path).startswith("offline-") + self.datetime = datetime.datetime.strptime( + os.path.basename(path).split("run-")[1].split("-")[0], "%Y%m%d_%H%M%S" + ) + + def __str__(self): + return self.path + + +class SyncThread(threading.Thread): + def __init__( + self, + sync_list, + project=None, + entity=None, + run_id=None, + job_type=None, + view=None, + verbose=None, + mark_synced=None, + app_url=None, + sync_tensorboard=None, + log_path=None, + append=None, + skip_console=None, + ): + threading.Thread.__init__(self) + # mark this process as internal + wandb._set_internal_process(disable=True) + self._sync_list = sync_list + self._project = project + self._entity = entity + self._run_id = run_id + self._job_type = job_type + self._view = view + self._verbose = verbose + self._mark_synced = mark_synced + self._app_url = app_url + self._sync_tensorboard = sync_tensorboard + self._log_path = log_path + self._append = append + self._skip_console = skip_console + + self._tmp_dir = tempfile.TemporaryDirectory() + atexit.register(self._tmp_dir.cleanup) + + def _parse_pb(self, data, exit_pb=None): + pb = wandb_internal_pb2.Record() + pb.ParseFromString(data) + record_type = pb.WhichOneof("record_type") + if self._view: + if self._verbose: + print("Record:", pb) # noqa: T201 + else: + print("Record:", record_type) # noqa: T201 + return pb, exit_pb, True + if record_type == "run": + if self._run_id: + pb.run.run_id = self._run_id + if self._project: + pb.run.project = self._project + if self._entity: + pb.run.entity = self._entity + if self._job_type: + pb.run.job_type = self._job_type + pb.control.req_resp = True + elif record_type in ("output", "output_raw") and self._skip_console: + return pb, exit_pb, True + elif record_type == "exit": + exit_pb = pb + return pb, exit_pb, True + elif record_type == "final": + assert exit_pb, "final seen without exit" + pb = exit_pb + exit_pb = None + return pb, exit_pb, False + + def _find_tfevent_files(self, sync_item): + tb_event_files = 0 + tb_logdirs = [] + tb_root = None + if self._sync_tensorboard: + if os.path.isdir(sync_item): + files = [] + for dirpath, _, _files in os.walk(sync_item): + for f in _files: + if TFEVENT_SUBSTRING in f: + files.append(os.path.join(dirpath, f)) + for tfevent in files: + tb_event_files += 1 + tb_dir = os.path.dirname(os.path.abspath(tfevent)) + if tb_dir not in tb_logdirs: + tb_logdirs.append(tb_dir) + if len(tb_logdirs) > 0: + tb_root = os.path.dirname(os.path.commonprefix(tb_logdirs)) + + elif TFEVENT_SUBSTRING in sync_item: + tb_root = os.path.dirname(os.path.abspath(sync_item)) + tb_logdirs.append(tb_root) + tb_event_files = 1 + return tb_event_files, tb_logdirs, tb_root + + def _setup_tensorboard(self, tb_root, tb_logdirs, tb_event_files, sync_item): + """Return true if this sync item can be synced as tensorboard.""" + if tb_root is not None: + if tb_event_files > 0 and sync_item.endswith(WANDB_SUFFIX): + wandb.termwarn("Found .wandb file, not streaming tensorboard metrics.") + else: + print(f"Found {tb_event_files} tfevent files in {tb_root}") # noqa: T201 + if len(tb_logdirs) > 3: + wandb.termwarn( + f"Found {len(tb_logdirs)} directories containing tfevent files. " + "If these represent multiple experiments, sync them " + "individually or pass a list of paths." + ) + return True + return False + + def _send_tensorboard(self, tb_root, tb_logdirs, send_manager): + if self._entity is None: + viewer, _ = send_manager._api.viewer_server_info() + self._entity = viewer.get("entity") + proto_run = wandb_internal_pb2.RunRecord() + proto_run.run_id = self._run_id or wandb.util.generate_id() + proto_run.project = self._project or wandb.util.auto_project_name(None) + proto_run.entity = self._entity + proto_run.telemetry.feature.sync_tfevents = True + + url = ( + f"{self._app_url}" + f"/{url_quote(proto_run.entity)}" + f"/{url_quote(proto_run.project)}" + f"/runs/{url_quote(proto_run.run_id)}" + ) + print(f"Syncing: {url} ...") # noqa: T201 + sys.stdout.flush() + # using a handler here automatically handles the step + # logic, adds summaries to the run, and handles different + # file types (like images)... but we need to remake the send_manager + record_q = queue.Queue() + sender_record_q = queue.Queue() + new_interface = InterfaceQueue(record_q) + context_keeper = context.ContextKeeper() + send_manager = sender.SendManager( + settings=send_manager._settings, + record_q=sender_record_q, + result_q=queue.Queue(), + interface=new_interface, + context_keeper=context_keeper, + ) + record = send_manager._interface._make_record(run=proto_run) + settings = wandb.Settings( + root_dir=self._tmp_dir.name, + run_id=proto_run.run_id, + x_start_time=time.time(), + ) + + settings_static = SettingsStatic(settings.to_proto()) + + handle_manager = handler.HandleManager( + settings=settings_static, + record_q=record_q, + result_q=None, + stopped=False, + writer_q=sender_record_q, + interface=new_interface, + context_keeper=context_keeper, + ) + + filesystem.mkdir_exists_ok(settings.files_dir) + send_manager.send_run(record, file_dir=settings.files_dir) + watcher = tb_watcher.TBWatcher(settings, proto_run, new_interface, True) + + for tb in tb_logdirs: + watcher.add(tb, True, tb_root) + sys.stdout.flush() + watcher.finish() + + # send all of our records like a boss + progress_step = 0 + spinner_states = ["-", "\\", "|", "/"] + line = " Uploading data to wandb\r" + while len(handle_manager) > 0: + data = next(handle_manager) + handle_manager.handle(data) + while len(send_manager) > 0: + data = next(send_manager) + send_manager.send(data) + + print_line = spinner_states[progress_step % 4] + line + wandb.termlog(print_line, newline=False, prefix=True) + progress_step += 1 + + # finish sending any data + while len(send_manager) > 0: + data = next(send_manager) + send_manager.send(data) + sys.stdout.flush() + handle_manager.finish() + send_manager.finish() + + def _robust_scan(self, ds): + """Attempt to scan data, handling incomplete files.""" + try: + return ds.scan_data() + except AssertionError as e: + if ds.in_last_block(): + wandb.termwarn( + f".wandb file is incomplete ({e}), be sure to sync this run again once it's finished" + ) + return None + else: + raise + + def run(self): + if self._log_path is not None: + print(f"Find logs at: {self._log_path}") # noqa: T201 + for sync_item in self._sync_list: + tb_event_files, tb_logdirs, tb_root = self._find_tfevent_files(sync_item) + if os.path.isdir(sync_item): + files = os.listdir(sync_item) + filtered_files = list(filter(lambda f: f.endswith(WANDB_SUFFIX), files)) + if tb_root is None and ( + check_and_warn_old(files) or len(filtered_files) != 1 + ): + print(f"Skipping directory: {sync_item}") # noqa: T201 + continue + if len(filtered_files) > 0: + sync_item = os.path.join(sync_item, filtered_files[0]) + sync_tb = self._setup_tensorboard( + tb_root, tb_logdirs, tb_event_files, sync_item + ) + # If we're syncing tensorboard, let's use a tmp dir for images etc. + root_dir = self._tmp_dir.name if sync_tb else os.path.dirname(sync_item) + + # When appending we are allowing a possible resume, ie the run + # does not have to exist already + resume = "allow" if self._append else None + + sm = sender.SendManager.setup(root_dir, resume=resume) + if sync_tb: + self._send_tensorboard(tb_root, tb_logdirs, sm) + continue + + ds = datastore.DataStore() + try: + ds.open_for_scan(sync_item) + except AssertionError as e: + print(f".wandb file is empty ({e}), skipping: {sync_item}") # noqa: T201 + continue + + # save exit for final send + exit_pb = None + finished = False + shown = False + while True: + data = self._robust_scan(ds) + if data is None: + break + pb, exit_pb, cont = self._parse_pb(data, exit_pb) + if exit_pb is not None: + finished = True + if cont: + continue + sm.send(pb) + # send any records that were added in previous send + while not sm._record_q.empty(): + data = sm._record_q.get(block=True) + sm.send(data) + + if pb.control.req_resp: + result = sm._result_q.get(block=True) + result_type = result.WhichOneof("result_type") + if not shown and result_type == "run_result": + r = result.run_result.run + # TODO(jhr): hardcode until we have settings in sync + url = ( + f"{self._app_url}" + f"/{url_quote(r.entity)}" + f"/{url_quote(r.project)}" + f"/runs/{url_quote(r.run_id)}" + ) + print(f"Syncing: {url} ... ", end="") # noqa: T201 + sys.stdout.flush() + shown = True + sm.finish() + # Only mark synced if the run actually finished + if self._mark_synced and not self._view and finished: + synced_file = f"{sync_item}{SYNCED_SUFFIX}" + with open(synced_file, "w"): + pass + print("done.") # noqa: T201 + + +class SyncManager: + def __init__( + self, + project=None, + entity=None, + run_id=None, + job_type=None, + mark_synced=None, + app_url=None, + view=None, + verbose=None, + sync_tensorboard=None, + log_path=None, + append=None, + skip_console=None, + ): + self._sync_list = [] + self._thread = None + self._project = project + self._entity = entity + self._run_id = run_id + self._job_type = job_type + self._mark_synced = mark_synced + self._app_url = app_url + self._view = view + self._verbose = verbose + self._sync_tensorboard = sync_tensorboard + self._log_path = log_path + self._append = append + self._skip_console = skip_console + + def status(self): + pass + + def add(self, p): + self._sync_list.append(os.path.abspath(str(p))) + + def start(self): + # create a thread for each file? + self._thread = SyncThread( + sync_list=self._sync_list, + project=self._project, + entity=self._entity, + run_id=self._run_id, + job_type=self._job_type, + view=self._view, + verbose=self._verbose, + mark_synced=self._mark_synced, + app_url=self._app_url, + sync_tensorboard=self._sync_tensorboard, + log_path=self._log_path, + append=self._append, + skip_console=self._skip_console, + ) + self._thread.start() + + def is_done(self): + return not self._thread.is_alive() + + def poll(self): + time.sleep(1) + return False + + +def get_runs( + include_offline: bool = True, + include_online: bool = True, + include_synced: bool = False, + include_unsynced: bool = True, + exclude_globs: Optional[List[str]] = None, + include_globs: Optional[List[str]] = None, +): + # TODO(jhr): grab dir info from settings + base = ".wandb" if os.path.exists(".wandb") else "wandb" + if not os.path.exists(base): + return () + all_dirs = os.listdir(base) + dirs = [] + if include_offline: + dirs += filter(lambda _d: _d.startswith("offline-run-"), all_dirs) + if include_online: + dirs += filter(lambda _d: _d.startswith("run-"), all_dirs) + # find run file in each dir + fnames = [] + dirs.sort() + for d in dirs: + paths = os.listdir(os.path.join(base, d)) + if exclude_globs: + paths = set(paths) + for g in exclude_globs: + paths = paths - set(fnmatch.filter(paths, g)) + paths = list(paths) + if include_globs: + new_paths = set() + for g in include_globs: + new_paths = new_paths.union(fnmatch.filter(paths, g)) + paths = list(new_paths) + for f in paths: + if f.endswith(WANDB_SUFFIX): + fnames.append(os.path.join(base, d, f)) + filtered = [] + for f in fnames: + dname = os.path.dirname(f) + # TODO(frz): online runs are assumed to be synced, verify from binary log. + if os.path.exists(f"{f}{SYNCED_SUFFIX}") or os.path.basename(dname).startswith( + "run-" + ): + if include_synced: + filtered.append(_LocalRun(dname, True)) + else: + if include_unsynced: + filtered.append(_LocalRun(dname, False)) + return tuple(filtered) + + +def get_run_from_path(path): + return _LocalRun(path) diff --git a/lib/python3.12/site-packages/wandb/trigger.py b/lib/python3.12/site-packages/wandb/trigger.py new file mode 100644 index 0000000000000000000000000000000000000000..15b52eec9e121a430ae64f1a5b8265f2cb2b8786 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/trigger.py @@ -0,0 +1,29 @@ +"""Module to facilitate adding hooks to wandb actions. + +Usage: + import trigger + trigger.register('on_something', func) + trigger.call('on_something', *args, **kwargs) + trigger.unregister('on_something', func) +""" + +from typing import Any, Callable + +_triggers = {} + + +def reset(): + _triggers.clear() + + +def register(event: str, func: Callable): + _triggers.setdefault(event, []).append(func) + + +def call(event_str: str, *args: Any, **kwargs: Any): + for func in _triggers.get(event_str, []): + func(*args, **kwargs) + + +def unregister(event: str, func: Callable): + _triggers[event].remove(func) diff --git a/lib/python3.12/site-packages/wandb/util.py b/lib/python3.12/site-packages/wandb/util.py new file mode 100644 index 0000000000000000000000000000000000000000..42b62cd43db1c74dde4d5b46f8d82dcf7301dc5e --- /dev/null +++ b/lib/python3.12/site-packages/wandb/util.py @@ -0,0 +1,1983 @@ +import colorsys +import contextlib +import dataclasses +import enum +import functools +import gzip +import importlib +import importlib.util +import itertools +import json +import logging +import math +import numbers +import os +import pathlib +import platform +import queue +import random +import re +import secrets +import shlex +import socket +import string +import sys +import tarfile +import tempfile +import threading +import time +import types +import urllib +from dataclasses import asdict, is_dataclass +from datetime import date, datetime, timedelta +from importlib import import_module +from sys import getsizeof +from types import ModuleType +from typing import ( + IO, + TYPE_CHECKING, + Callable, + Dict, + Iterable, + List, + Mapping, + Optional, + Sequence, + TextIO, + Tuple, + Union, +) + +import requests +import yaml +from typing_extensions import Any, Generator, TypeGuard, TypeVar + +import wandb +import wandb.env +from wandb.errors import ( + AuthenticationError, + CommError, + UsageError, + WandbCoreNotAvailableError, + term, +) +from wandb.sdk.internal.thread_local_settings import _thread_local_api_settings +from wandb.sdk.lib import filesystem, runid +from wandb.sdk.lib.json_util import dump, dumps +from wandb.sdk.lib.paths import FilePathStr, StrPath + +if TYPE_CHECKING: + import wandb.sdk.internal.settings_static + import wandb.sdk.wandb_settings + from wandb.sdk.artifacts.artifact import Artifact + +CheckRetryFnType = Callable[[Exception], Union[bool, timedelta]] +T = TypeVar("T") + + +logger = logging.getLogger(__name__) +_not_importable = set() + +LAUNCH_JOB_ARTIFACT_SLOT_NAME = "_wandb_job" + +MAX_LINE_BYTES = (10 << 20) - (100 << 10) # imposed by back end +IS_GIT = os.path.exists(os.path.join(os.path.dirname(__file__), "..", ".git")) + +# From https://docs.docker.com/engine/reference/commandline/tag/ +# "Name components may contain lowercase letters, digits and separators. +# A separator is defined as a period, one or two underscores, or one or more dashes. +# A name component may not start or end with a separator." +DOCKER_IMAGE_NAME_SEPARATOR = "(?:__|[._]|[-]+)" +RE_DOCKER_IMAGE_NAME_SEPARATOR_START = re.compile("^" + DOCKER_IMAGE_NAME_SEPARATOR) +RE_DOCKER_IMAGE_NAME_SEPARATOR_END = re.compile(DOCKER_IMAGE_NAME_SEPARATOR + "$") +RE_DOCKER_IMAGE_NAME_SEPARATOR_REPEAT = re.compile(DOCKER_IMAGE_NAME_SEPARATOR + "{2,}") +RE_DOCKER_IMAGE_NAME_CHARS = re.compile(r"[^a-z0-9._\-]") + +# these match the environments for gorilla +if IS_GIT: + SENTRY_ENV = "development" +else: + SENTRY_ENV = "production" + + +POW_10_BYTES = [ + ("B", 10**0), + ("KB", 10**3), + ("MB", 10**6), + ("GB", 10**9), + ("TB", 10**12), + ("PB", 10**15), + ("EB", 10**18), +] + +POW_2_BYTES = [ + ("B", 2**0), + ("KiB", 2**10), + ("MiB", 2**20), + ("GiB", 2**30), + ("TiB", 2**40), + ("PiB", 2**50), + ("EiB", 2**60), +] + + +def vendor_setup() -> Callable: + """Create a function that restores user paths after vendor imports. + + This enables us to use the vendor directory for packages we don't depend on. Call + the returned function after imports are complete. If you don't you may modify the + user's path which is never good. + + Usage: + + ```python + reset_path = vendor_setup() + # do any vendor imports... + reset_path() + ``` + """ + original_path = [directory for directory in sys.path] + + def reset_import_path() -> None: + sys.path = original_path + + parent_dir = os.path.abspath(os.path.dirname(__file__)) + vendor_dir = os.path.join(parent_dir, "vendor") + vendor_packages = ( + "gql-0.2.0", + "graphql-core-1.1", + "watchdog_0_9_0", + "promise-2.3.0", + ) + package_dirs = [os.path.join(vendor_dir, p) for p in vendor_packages] + for p in [vendor_dir] + package_dirs: + if p not in sys.path: + sys.path.insert(1, p) + + return reset_import_path + + +def vendor_import(name: str) -> Any: + reset_path = vendor_setup() + module = import_module(name) + reset_path() + return module + + +class LazyModuleState: + def __init__(self, module: types.ModuleType) -> None: + self.module = module + self.load_started = False + self.lock = threading.RLock() + + def load(self) -> None: + with self.lock: + if self.load_started: + return + self.load_started = True + assert self.module.__spec__ is not None + assert self.module.__spec__.loader is not None + self.module.__spec__.loader.exec_module(self.module) + self.module.__class__ = types.ModuleType + + # Set the submodule as an attribute on the parent module + # This enables access to the submodule via normal attribute access. + parent, _, child = self.module.__name__.rpartition(".") + if parent: + parent_module = sys.modules[parent] + setattr(parent_module, child, self.module) + + +class LazyModule(types.ModuleType): + def __getattribute__(self, name: str) -> Any: + state = object.__getattribute__(self, "__lazy_module_state__") + state.load() + return object.__getattribute__(self, name) + + def __setattr__(self, name: str, value: Any) -> None: + state = object.__getattribute__(self, "__lazy_module_state__") + state.load() + object.__setattr__(self, name, value) + + def __delattr__(self, name: str) -> None: + state = object.__getattribute__(self, "__lazy_module_state__") + state.load() + object.__delattr__(self, name) + + +def import_module_lazy(name: str) -> types.ModuleType: + """Import a module lazily, only when it is used. + + Inspired by importlib.util.LazyLoader, but improved so that the module loading is + thread-safe. Circular dependency between modules can lead to a deadlock if the two + modules are loaded from different threads. + + :param (str) name: Dot-separated module path. E.g., 'scipy.stats'. + """ + try: + return sys.modules[name] + except KeyError: + spec = importlib.util.find_spec(name) + if spec is None: + raise ModuleNotFoundError + module = importlib.util.module_from_spec(spec) + module.__lazy_module_state__ = LazyModuleState(module) # type: ignore + module.__class__ = LazyModule + sys.modules[name] = module + return module + + +def get_module( + name: str, + required: Optional[Union[str, bool]] = None, + lazy: bool = True, +) -> Any: + """Return module or None. Absolute import is required. + + :param (str) name: Dot-separated module path. E.g., 'scipy.stats'. + :param (str) required: A string to raise a ValueError if missing + :param (bool) lazy: If True, return a lazy loader for the module. + :return: (module|None) If import succeeds, the module will be returned. + """ + if name not in _not_importable: + try: + if not lazy: + return import_module(name) + else: + return import_module_lazy(name) + except Exception: + _not_importable.add(name) + msg = f"Error importing optional module {name}" + if required: + logger.exception(msg) + if required and name in _not_importable: + raise wandb.Error(required) + + +def get_optional_module(name) -> Optional["importlib.ModuleInterface"]: # type: ignore + return get_module(name) + + +np = get_module("numpy") + +pd_available = False +pandas_spec = importlib.util.find_spec("pandas") +if pandas_spec is not None: + pd_available = True + +# TODO: Revisit these limits +VALUE_BYTES_LIMIT = 100000 + + +def app_url(api_url: str) -> str: + """Return the frontend app url without a trailing slash.""" + # TODO: move me to settings + app_url = wandb.env.get_app_url() + if app_url is not None: + return str(app_url.strip("/")) + if "://api.wandb.test" in api_url: + # dev mode + return api_url.replace("://api.", "://app.").strip("/") + elif "://api.wandb." in api_url: + # cloud + return api_url.replace("://api.", "://").strip("/") + elif "://api." in api_url: + # onprem cloud + return api_url.replace("://api.", "://app.").strip("/") + # wandb/local + return api_url + + +def get_full_typename(o: Any) -> Any: + """Determine types based on type names. + + Avoids needing to to import (and therefore depend on) PyTorch, TensorFlow, etc. + """ + instance_name = o.__class__.__module__ + "." + o.__class__.__name__ + if instance_name in ["builtins.module", "__builtin__.module"]: + return o.__name__ + else: + return instance_name + + +def get_h5_typename(o: Any) -> Any: + typename = get_full_typename(o) + if is_tf_tensor_typename(typename): + return "tensorflow.Tensor" + elif is_pytorch_tensor_typename(typename): + return "torch.Tensor" + else: + return o.__class__.__module__.split(".")[0] + "." + o.__class__.__name__ + + +def is_uri(string: str) -> bool: + parsed_uri = urllib.parse.urlparse(string) + return len(parsed_uri.scheme) > 0 + + +def local_file_uri_to_path(uri: str) -> str: + """Convert URI to local filesystem path. + + No-op if the uri does not have the expected scheme. + """ + path = urllib.parse.urlparse(uri).path if uri.startswith("file:") else uri + return urllib.request.url2pathname(path) + + +def get_local_path_or_none(path_or_uri: str) -> Optional[str]: + """Return path if local, None otherwise. + + Return None if the argument is a local path (not a scheme or file:///). Otherwise + return `path_or_uri`. + """ + parsed_uri = urllib.parse.urlparse(path_or_uri) + if ( + len(parsed_uri.scheme) == 0 + or parsed_uri.scheme == "file" + and len(parsed_uri.netloc) == 0 + ): + return local_file_uri_to_path(path_or_uri) + else: + return None + + +def make_tarfile( + output_filename: str, + source_dir: str, + archive_name: str, + custom_filter: Optional[Callable] = None, +) -> None: + # Helper for filtering out modification timestamps + def _filter_timestamps(tar_info: "tarfile.TarInfo") -> Optional["tarfile.TarInfo"]: + tar_info.mtime = 0 + return tar_info if custom_filter is None else custom_filter(tar_info) + + descriptor, unzipped_filename = tempfile.mkstemp() + try: + with tarfile.open(unzipped_filename, "w") as tar: + tar.add(source_dir, arcname=archive_name, filter=_filter_timestamps) + # When gzipping the tar, don't include the tar's filename or modification time in the + # zipped archive (see https://docs.python.org/3/library/gzip.html#gzip.GzipFile) + with gzip.GzipFile( + filename="", fileobj=open(output_filename, "wb"), mode="wb", mtime=0 + ) as gzipped_tar, open(unzipped_filename, "rb") as tar_file: + gzipped_tar.write(tar_file.read()) + finally: + os.close(descriptor) + os.remove(unzipped_filename) + + +def is_tf_tensor(obj: Any) -> bool: + import tensorflow # type: ignore + + return isinstance(obj, tensorflow.Tensor) + + +def is_tf_tensor_typename(typename: str) -> bool: + return typename.startswith("tensorflow.") and ( + "Tensor" in typename or "Variable" in typename + ) + + +def is_tf_eager_tensor_typename(typename: str) -> bool: + return typename.startswith("tensorflow.") and ("EagerTensor" in typename) + + +def is_pytorch_tensor(obj: Any) -> bool: + import torch # type: ignore + + return isinstance(obj, torch.Tensor) + + +def is_pytorch_tensor_typename(typename: str) -> bool: + return typename.startswith("torch.") and ( + "Tensor" in typename or "Variable" in typename + ) + + +def is_jax_tensor_typename(typename: str) -> bool: + return typename.startswith("jaxlib.") and "Array" in typename + + +def get_jax_tensor(obj: Any) -> Optional[Any]: + import jax # type: ignore + + return jax.device_get(obj) + + +def is_fastai_tensor_typename(typename: str) -> bool: + return typename.startswith("fastai.") and ("Tensor" in typename) + + +def is_pandas_data_frame_typename(typename: str) -> bool: + return typename.startswith("pandas.") and "DataFrame" in typename + + +def is_matplotlib_typename(typename: str) -> bool: + return typename.startswith("matplotlib.") + + +def is_plotly_typename(typename: str) -> bool: + return typename.startswith("plotly.") + + +def is_plotly_figure_typename(typename: str) -> bool: + return typename.startswith("plotly.") and typename.endswith(".Figure") + + +def is_numpy_array(obj: Any) -> bool: + return np and isinstance(obj, np.ndarray) + + +def is_pandas_data_frame(obj: Any) -> bool: + if pd_available: + import pandas as pd + + return isinstance(obj, pd.DataFrame) + else: + return is_pandas_data_frame_typename(get_full_typename(obj)) + + +def ensure_matplotlib_figure(obj: Any) -> Any: + """Extract the current figure from a matplotlib object. + + Return the object itself if it's a figure. + Raises ValueError if the object can't be converted. + """ + import matplotlib # type: ignore + from matplotlib.figure import Figure # type: ignore + + # there are combinations of plotly and matplotlib versions that don't work well together, + # this patches matplotlib to add a removed method that plotly assumes exists + from matplotlib.spines import Spine # type: ignore + + def is_frame_like(self: Any) -> bool: + """Return True if directly on axes frame. + + This is useful for determining if a spine is the edge of an + old style MPL plot. If so, this function will return True. + """ + position = self._position or ("outward", 0.0) + if isinstance(position, str): + if position == "center": + position = ("axes", 0.5) + elif position == "zero": + position = ("data", 0) + if len(position) != 2: + raise ValueError("position should be 2-tuple") + position_type, amount = position # type: ignore + if position_type == "outward" and amount == 0: + return True + else: + return False + + Spine.is_frame_like = is_frame_like + + if obj == matplotlib.pyplot: + obj = obj.gcf() + elif not isinstance(obj, Figure): + if hasattr(obj, "figure"): + obj = obj.figure + # Some matplotlib objects have a figure function + if not isinstance(obj, Figure): + raise ValueError( + "Only matplotlib.pyplot or matplotlib.pyplot.Figure objects are accepted." + ) + return obj + + +def matplotlib_to_plotly(obj: Any) -> Any: + obj = ensure_matplotlib_figure(obj) + tools = get_module( + "plotly.tools", + required=( + "plotly is required to log interactive plots, install with: " + "`pip install plotly` or convert the plot to an image with `wandb.Image(plt)`" + ), + ) + return tools.mpl_to_plotly(obj) + + +def matplotlib_contains_images(obj: Any) -> bool: + obj = ensure_matplotlib_figure(obj) + return any(len(ax.images) > 0 for ax in obj.axes) + + +def _numpy_generic_convert(obj: Any) -> Any: + obj = obj.item() + if isinstance(obj, float) and math.isnan(obj): + obj = None + elif isinstance(obj, np.generic) and ( + obj.dtype.kind == "f" or obj.dtype == "bfloat16" + ): + # obj is a numpy float with precision greater than that of native python float + # (i.e., float96 or float128) or it is of custom type such as bfloat16. + # in these cases, obj.item() does not return a native + # python float (in the first case - to avoid loss of precision, + # so we need to explicitly cast this down to a 64bit float) + obj = float(obj) + return obj + + +def _sanitize_numpy_keys( + d: Dict, + visited: Optional[Dict[int, Dict]] = None, +) -> Tuple[Dict, bool]: + """Returns a dictionary where all NumPy keys are converted. + + Args: + d: The dictionary to sanitize. + + Returns: + A sanitized dictionary, and a boolean indicating whether anything was + changed. + """ + out: Dict[Any, Any] = dict() + converted = False + + # Work with recursive dictionaries: if a dictionary has already been + # converted, reuse its converted value to retain the recursive structure + # of the input. + if visited is None: + visited = {id(d): out} + elif id(d) in visited: + return visited[id(d)], False + visited[id(d)] = out + + for key, value in d.items(): + if isinstance(value, dict): + value, converted_value = _sanitize_numpy_keys(value, visited) + converted |= converted_value + if isinstance(key, np.generic): + key = _numpy_generic_convert(key) + converted = True + out[key] = value + + return out, converted + + +def json_friendly( # noqa: C901 + obj: Any, +) -> Union[Tuple[Any, bool], Tuple[Union[None, str, float], bool]]: + """Convert an object into something that's more becoming of JSON.""" + converted = True + typename = get_full_typename(obj) + + if is_tf_eager_tensor_typename(typename): + obj = obj.numpy() + elif is_tf_tensor_typename(typename): + try: + obj = obj.eval() + except RuntimeError: + obj = obj.numpy() + elif is_pytorch_tensor_typename(typename) or is_fastai_tensor_typename(typename): + try: + if obj.requires_grad: + obj = obj.detach() + except AttributeError: + pass # before 0.4 is only present on variables + + try: + obj = obj.data + except RuntimeError: + pass # happens for Tensors before 0.4 + + if obj.size(): + obj = obj.cpu().detach().numpy() + else: + return obj.item(), True + elif is_jax_tensor_typename(typename): + obj = get_jax_tensor(obj) + + if is_numpy_array(obj): + if obj.size == 1: + obj = obj.flatten()[0] + elif obj.size <= 32: + obj = obj.tolist() + elif np and isinstance(obj, np.generic): + obj = _numpy_generic_convert(obj) + elif isinstance(obj, bytes): + obj = obj.decode("utf-8") + elif isinstance(obj, (datetime, date)): + obj = obj.isoformat() + elif callable(obj): + obj = ( + f"{obj.__module__}.{obj.__qualname__}" + if hasattr(obj, "__qualname__") and hasattr(obj, "__module__") + else str(obj) + ) + elif isinstance(obj, float) and math.isnan(obj): + obj = None + elif isinstance(obj, dict) and np: + obj, converted = _sanitize_numpy_keys(obj) + elif isinstance(obj, set): + # set is not json serializable, so we convert it to tuple + obj = tuple(obj) + elif isinstance(obj, enum.Enum): + obj = obj.name + else: + converted = False + if getsizeof(obj) > VALUE_BYTES_LIMIT: + wandb.termwarn( + f"Serializing object of type {type(obj).__name__} that is {getsizeof(obj)} bytes" + ) + return obj, converted + + +def json_friendly_val(val: Any) -> Any: + """Make any value (including dict, slice, sequence, dataclass) JSON friendly.""" + converted: Union[dict, list] + if isinstance(val, dict): + converted = {} + for key, value in val.items(): + converted[key] = json_friendly_val(value) + return converted + if isinstance(val, slice): + converted = dict( + slice_start=val.start, slice_step=val.step, slice_stop=val.stop + ) + return converted + val, _ = json_friendly(val) + if isinstance(val, Sequence) and not isinstance(val, str): + converted = [] + for value in val: + converted.append(json_friendly_val(value)) + return converted + if is_dataclass(val) and not isinstance(val, type): + converted = asdict(val) + return converted + else: + if val.__class__.__module__ not in ("builtins", "__builtin__"): + val = str(val) + return val + + +def alias_is_version_index(alias: str) -> bool: + return len(alias) >= 2 and alias[0] == "v" and alias[1:].isnumeric() + + +def convert_plots(obj: Any) -> Any: + if is_matplotlib_typename(get_full_typename(obj)): + tools = get_module( + "plotly.tools", + required=( + "plotly is required to log interactive plots, install with: " + "`pip install plotly` or convert the plot to an image with `wandb.Image(plt)`" + ), + ) + obj = tools.mpl_to_plotly(obj) + + if is_plotly_typename(get_full_typename(obj)): + return {"_type": "plotly", "plot": obj.to_plotly_json()} + else: + return obj + + +def maybe_compress_history(obj: Any) -> Tuple[Any, bool]: + if np and isinstance(obj, np.ndarray) and obj.size > 32: + return wandb.Histogram(obj, num_bins=32).to_json(), True + else: + return obj, False + + +def maybe_compress_summary(obj: Any, h5_typename: str) -> Tuple[Any, bool]: + if np and isinstance(obj, np.ndarray) and obj.size > 32: + return ( + { + "_type": h5_typename, # may not be ndarray + "var": np.var(obj).item(), + "mean": np.mean(obj).item(), + "min": np.amin(obj).item(), + "max": np.amax(obj).item(), + "10%": np.percentile(obj, 10), + "25%": np.percentile(obj, 25), + "75%": np.percentile(obj, 75), + "90%": np.percentile(obj, 90), + "size": obj.size, + }, + True, + ) + else: + return obj, False + + +def launch_browser(attempt_launch_browser: bool = True) -> bool: + """Decide if we should launch a browser.""" + _display_variables = ["DISPLAY", "WAYLAND_DISPLAY", "MIR_SOCKET"] + _webbrowser_names_blocklist = ["www-browser", "lynx", "links", "elinks", "w3m"] + + import webbrowser + + launch_browser = attempt_launch_browser + if launch_browser: + if "linux" in sys.platform and not any( + os.getenv(var) for var in _display_variables + ): + launch_browser = False + try: + browser = webbrowser.get() + if hasattr(browser, "name") and browser.name in _webbrowser_names_blocklist: + launch_browser = False + except webbrowser.Error: + launch_browser = False + + return launch_browser + + +def generate_id(length: int = 8) -> str: + # Do not use this; use wandb.sdk.lib.runid.generate_id instead. + # This is kept only for legacy code. + return runid.generate_id(length) + + +def parse_tfjob_config() -> Any: + """Attempt to parse TFJob config, returning False if it can't find it.""" + if os.getenv("TF_CONFIG"): + try: + return json.loads(os.environ["TF_CONFIG"]) + except ValueError: + return False + else: + return False + + +class WandBJSONEncoder(json.JSONEncoder): + """A JSON Encoder that handles some extra types.""" + + def default(self, obj: Any) -> Any: + if hasattr(obj, "json_encode"): + return obj.json_encode() + # if hasattr(obj, 'to_json'): + # return obj.to_json() + tmp_obj, converted = json_friendly(obj) + if converted: + return tmp_obj + return json.JSONEncoder.default(self, obj) + + +class WandBJSONEncoderOld(json.JSONEncoder): + """A JSON Encoder that handles some extra types.""" + + def default(self, obj: Any) -> Any: + tmp_obj, converted = json_friendly(obj) + tmp_obj, compressed = maybe_compress_summary(tmp_obj, get_h5_typename(obj)) + if converted: + return tmp_obj + return json.JSONEncoder.default(self, tmp_obj) + + +class WandBHistoryJSONEncoder(json.JSONEncoder): + """A JSON Encoder that handles some extra types. + + This encoder turns numpy like objects with a size > 32 into histograms. + """ + + def default(self, obj: Any) -> Any: + obj, converted = json_friendly(obj) + obj, compressed = maybe_compress_history(obj) + if converted: + return obj + return json.JSONEncoder.default(self, obj) + + +class JSONEncoderUncompressed(json.JSONEncoder): + """A JSON Encoder that handles some extra types. + + This encoder turns numpy like objects with a size > 32 into histograms. + """ + + def default(self, obj: Any) -> Any: + if is_numpy_array(obj): + return obj.tolist() + elif np and isinstance(obj, np.number): + return obj.item() + elif np and isinstance(obj, np.generic): + obj = obj.item() + return json.JSONEncoder.default(self, obj) + + +def json_dump_safer(obj: Any, fp: IO[str], **kwargs: Any) -> None: + """Convert obj to json, with some extra encodable types.""" + return dump(obj, fp, cls=WandBJSONEncoder, **kwargs) + + +def json_dumps_safer(obj: Any, **kwargs: Any) -> str: + """Convert obj to json, with some extra encodable types.""" + return dumps(obj, cls=WandBJSONEncoder, **kwargs) + + +# This is used for dumping raw json into files +def json_dump_uncompressed(obj: Any, fp: IO[str], **kwargs: Any) -> None: + """Convert obj to json, with some extra encodable types.""" + return dump(obj, fp, cls=JSONEncoderUncompressed, **kwargs) + + +def json_dumps_safer_history(obj: Any, **kwargs: Any) -> str: + """Convert obj to json, with some extra encodable types, including histograms.""" + return dumps(obj, cls=WandBHistoryJSONEncoder, **kwargs) + + +def make_json_if_not_number( + v: Union[int, float, str, Mapping, Sequence], +) -> Union[int, float, str]: + """If v is not a basic type convert it to json.""" + if isinstance(v, (float, int)): + return v + return json_dumps_safer(v) + + +def make_safe_for_json(obj: Any) -> Any: + """Replace invalid json floats with strings. Also converts to lists and dicts.""" + if isinstance(obj, Mapping): + return {k: make_safe_for_json(v) for k, v in obj.items()} + elif isinstance(obj, str): + # str's are Sequence, so we need to short-circuit + return obj + elif isinstance(obj, Sequence): + return [make_safe_for_json(v) for v in obj] + elif isinstance(obj, float): + # W&B backend and UI handle these strings + if obj != obj: # standard way to check for NaN + return "NaN" + elif obj == float("+inf"): + return "Infinity" + elif obj == float("-inf"): + return "-Infinity" + return obj + + +def no_retry_4xx(e: Exception) -> bool: + if not isinstance(e, requests.HTTPError): + return True + assert e.response is not None + if not (400 <= e.response.status_code < 500) or e.response.status_code == 429: + return True + body = json.loads(e.response.content) + raise UsageError(body["errors"][0]["message"]) + + +def parse_backend_error_messages(response: requests.Response) -> List[str]: + """Returns error messages stored in a backend response. + + If the response is not in an expected format, an empty list is returned. + + Args: + response: A response to an HTTP request to the W&B server. + """ + try: + data = response.json() + except requests.JSONDecodeError: + return [] + + if not isinstance(data, dict): + return [] + + # Backend error values are returned in one of two ways: + # - A string containing the error message + # - A JSON object with a "message" field that is a string + def get_message(error: Any) -> Optional[str]: + if isinstance(error, str): + return error + elif ( + isinstance(error, dict) + and (message := error.get("message")) + and isinstance(message, str) + ): + return message + else: + return None + + # The response can contain an "error" field with a single error + # or an "errors" field with a list of errors. + if error := data.get("error"): + message = get_message(error) + return [message] if message else [] + + elif (errors := data.get("errors")) and isinstance(errors, list): + messages: List[str] = [] + for error in errors: + message = get_message(error) + if message: + messages.append(message) + return messages + + else: + return [] + + +def no_retry_auth(e: Any) -> bool: + if hasattr(e, "exception"): + e = e.exception + if not isinstance(e, requests.HTTPError): + return True + if e.response is None: + return True + # Don't retry bad request errors; raise immediately + if e.response.status_code in (400, 409): + return False + # Retry all non-forbidden/unauthorized/not-found errors. + if e.response.status_code not in (401, 403, 404): + return True + + # Crash with more informational message on forbidden/unauthorized errors. + # UnauthorizedError + if e.response.status_code == 401: + raise AuthenticationError( + "The API key you provided is either invalid or missing. " + f"If the `{wandb.env.API_KEY}` environment variable is set, make sure it is correct. " + "Otherwise, to resolve this issue, you may try running the 'wandb login --relogin' command. " + "If you are using a local server, make sure that you're using the correct hostname. " + "If you're not sure, you can try logging in again using the 'wandb login --relogin --host [hostname]' command." + f"(Error {e.response.status_code}: {e.response.reason})" + ) + # ForbiddenError + if e.response.status_code == 403: + if wandb.run: + raise CommError(f"Permission denied to access {wandb.run.path}") + else: + raise CommError( + "It appears that you do not have permission to access the requested resource. " + "Please reach out to the project owner to grant you access. " + "If you have the correct permissions, verify that there are no issues with your networking setup." + f"(Error {e.response.status_code}: {e.response.reason})" + ) + + # NotFoundError + if e.response.status_code == 404: + # If error message is empty, raise a more generic NotFoundError message. + if parse_backend_error_messages(e.response): + return False + else: + raise LookupError( + f"Failed to find resource. Please make sure you have the correct resource path. " + f"(Error {e.response.status_code}: {e.response.reason})" + ) + return False + + +def check_retry_conflict(e: Any) -> Optional[bool]: + """Check if the exception is a conflict type so it can be retried. + + Returns: + True - Should retry this operation + False - Should not retry this operation + None - No decision, let someone else decide + """ + if hasattr(e, "exception"): + e = e.exception + if isinstance(e, requests.HTTPError) and e.response is not None: + if e.response.status_code == 409: + return True + return None + + +def check_retry_conflict_or_gone(e: Any) -> Optional[bool]: + """Check if the exception is a conflict or gone type, so it can be retried or not. + + Returns: + True - Should retry this operation + False - Should not retry this operation + None - No decision, let someone else decide + """ + if hasattr(e, "exception"): + e = e.exception + if isinstance(e, requests.HTTPError) and e.response is not None: + if e.response.status_code == 409: + return True + if e.response.status_code == 410: + return False + return None + + +def make_check_retry_fn( + fallback_retry_fn: CheckRetryFnType, + check_fn: Callable[[Exception], Optional[bool]], + check_timedelta: Optional[timedelta] = None, +) -> CheckRetryFnType: + """Return a check_retry_fn which can be used by lib.Retry(). + + Args: + fallback_fn: Use this function if check_fn didn't decide if a retry should happen. + check_fn: Function which returns bool if retry should happen or None if unsure. + check_timedelta: Optional retry timeout if we check_fn matches the exception + """ + + def check_retry_fn(e: Exception) -> Union[bool, timedelta]: + check = check_fn(e) + if check is None: + return fallback_retry_fn(e) + if check is False: + return False + if check_timedelta: + return check_timedelta + return True + + return check_retry_fn + + +def find_runner(program: str) -> Union[None, list, List[str]]: + """Return a command that will run program. + + Args: + program: The string name of the program to try to run. + + Returns: + commandline list of strings to run the program (eg. with subprocess.call()) or None + """ + if os.path.isfile(program) and not os.access(program, os.X_OK): + # program is a path to a non-executable file + try: + opened = open(program) + except OSError: # PermissionError doesn't exist in 2.7 + return None + first_line = opened.readline().strip() + if first_line.startswith("#!"): + return shlex.split(first_line[2:]) + if program.endswith(".py"): + return [sys.executable] + return None + + +def downsample(values: Sequence, target_length: int) -> list: + """Downsample 1d values to target_length, including start and end. + + Algorithm just rounds index down. + + Values can be any sequence, including a generator. + """ + if not target_length > 1: + raise UsageError("target_length must be > 1") + values = list(values) + if len(values) < target_length: + return values + ratio = float(len(values) - 1) / (target_length - 1) + result = [] + for i in range(target_length): + result.append(values[int(i * ratio)]) + return result + + +def has_num(dictionary: Mapping, key: Any) -> bool: + return key in dictionary and isinstance(dictionary[key], numbers.Number) + + +def docker_image_regex(image: str) -> Any: + """Regex match for valid docker image names.""" + if image: + return re.match( + r"^(?:(?=[^:\/]{1,253})(?!-)[a-zA-Z0-9-]{1,63}(? Optional[str]: + """Scan docker run args and attempt to find the most likely docker image argument. + + It excludes any arguments that start with a dash, and the argument after it if it + isn't a boolean switch. This can be improved, we currently fallback gracefully when + this fails. + """ + bool_args = [ + "-t", + "--tty", + "--rm", + "--privileged", + "--oom-kill-disable", + "--no-healthcheck", + "-i", + "--interactive", + "--init", + "--help", + "--detach", + "-d", + "--sig-proxy", + "-it", + "-itd", + ] + last_flag = -2 + last_arg = "" + possible_images = [] + if len(args) > 0 and args[0] == "run": + args.pop(0) + for i, arg in enumerate(args): + if arg.startswith("-"): + last_flag = i + last_arg = arg + elif "@sha256:" in arg: + # Because our regex doesn't match digests + possible_images.append(arg) + elif docker_image_regex(arg): + if last_flag == i - 2: + possible_images.append(arg) + elif "=" in last_arg: + possible_images.append(arg) + elif last_arg in bool_args and last_flag == i - 1: + possible_images.append(arg) + most_likely = None + for img in possible_images: + if ":" in img or "@" in img or "/" in img: + most_likely = img + break + if most_likely is None and len(possible_images) > 0: + most_likely = possible_images[0] + return most_likely + + +def load_yaml(file: Any) -> Any: + return yaml.safe_load(file) + + +def image_id_from_k8s() -> Optional[str]: + """Ping the k8s metadata service for the image id. + + Specify the KUBERNETES_NAMESPACE environment variable if your pods are not in the + default namespace: + + - name: KUBERNETES_NAMESPACE valueFrom: + fieldRef: + fieldPath: metadata.namespace + """ + token_path = "/var/run/secrets/kubernetes.io/serviceaccount/token" + + if not os.path.exists(token_path): + return None + + try: + with open(token_path) as token_file: + token = token_file.read() + except FileNotFoundError: + logger.warning(f"Token file not found at {token_path}.") + return None + except PermissionError as e: + current_uid = os.getuid() + warning = ( + f"Unable to read the token file at {token_path} due to permission error ({e})." + f"The current user id is {current_uid}. " + "Consider changing the securityContext to run the container as the current user." + ) + logger.warning(warning) + wandb.termwarn(warning) + return None + + if not token: + return None + + k8s_server = "https://{}:{}/api/v1/namespaces/{}/pods/{}".format( + os.getenv("KUBERNETES_SERVICE_HOST"), + os.getenv("KUBERNETES_PORT_443_TCP_PORT"), + os.getenv("KUBERNETES_NAMESPACE", "default"), + os.getenv("HOSTNAME"), + ) + try: + res = requests.get( + k8s_server, + verify="/var/run/secrets/kubernetes.io/serviceaccount/ca.crt", + timeout=3, + headers={"Authorization": f"Bearer {token}"}, + ) + res.raise_for_status() + except requests.RequestException: + return None + try: + return str( # noqa: B005 + res.json()["status"]["containerStatuses"][0]["imageID"] + ).strip("docker-pullable://") + except (ValueError, KeyError, IndexError): + logger.exception("Error checking kubernetes for image id") + return None + + +def async_call( + target: Callable, timeout: Optional[Union[int, float]] = None +) -> Callable: + """Wrap a method to run in the background with an optional timeout. + + Returns a new method that will call the original with any args, waiting for upto + timeout seconds. This new method blocks on the original and returns the result or + None if timeout was reached, along with the thread. You can check thread.is_alive() + to determine if a timeout was reached. If an exception is thrown in the thread, we + reraise it. + """ + q: queue.Queue = queue.Queue() + + def wrapped_target(q: "queue.Queue", *args: Any, **kwargs: Any) -> Any: + try: + q.put(target(*args, **kwargs)) + except Exception as e: + q.put(e) + + def wrapper( + *args: Any, **kwargs: Any + ) -> Union[Tuple[Exception, "threading.Thread"], Tuple[None, "threading.Thread"]]: + thread = threading.Thread( + target=wrapped_target, args=(q,) + args, kwargs=kwargs + ) + thread.daemon = True + thread.start() + + try: + result = q.get(True, timeout) + except queue.Empty: + return None, thread + + if isinstance(result, Exception): + raise result.with_traceback(sys.exc_info()[2]) + return result, thread + + return wrapper + + +def read_many_from_queue( + q: "queue.Queue", max_items: int, queue_timeout: Union[int, float] +) -> list: + try: + item = q.get(True, queue_timeout) + except queue.Empty: + return [] + items = [item] + for _ in range(max_items): + try: + item = q.get_nowait() + except queue.Empty: + return items + items.append(item) + return items + + +def stopwatch_now() -> float: + """Get a time value for interval comparisons. + + When possible it is a monotonic clock to prevent backwards time issues. + """ + return time.monotonic() + + +def class_colors(class_count: int) -> List[List[int]]: + # make class 0 black, and the rest equally spaced fully saturated hues + return [[0, 0, 0]] + [ + colorsys.hsv_to_rgb(i / (class_count - 1.0), 1.0, 1.0) # type: ignore + for i in range(class_count - 1) + ] + + +def _prompt_choice( + input_timeout: Union[int, float, None] = None, + jupyter: bool = False, +) -> str: + input_fn: Callable = input + prompt = term.LOG_STRING + if input_timeout is not None: + # delayed import to mitigate risk of timed_input complexity + from wandb.sdk.lib import timed_input + + input_fn = functools.partial(timed_input.timed_input, timeout=input_timeout) + # timed_input doesn't handle enhanced prompts + if platform.system() == "Windows": + prompt = "wandb" + + text = f"{prompt}: Enter your choice: " + if input_fn == input: + choice = input_fn(text) + else: + choice = input_fn(text, jupyter=jupyter) + return choice # type: ignore + + +def prompt_choices( + choices: Sequence[str], + input_timeout: Union[int, float, None] = None, + jupyter: bool = False, +) -> str: + """Allow a user to choose from a list of options.""" + for i, choice in enumerate(choices): + wandb.termlog(f"({i + 1}) {choice}") + + idx = -1 + while idx < 0 or idx > len(choices) - 1: + choice = _prompt_choice(input_timeout=input_timeout, jupyter=jupyter) + if not choice: + continue + idx = -1 + try: + idx = int(choice) - 1 + except ValueError: + pass + if idx < 0 or idx > len(choices) - 1: + wandb.termwarn("Invalid choice") + result = choices[idx] + wandb.termlog(f"You chose {result!r}") + return result + + +def guess_data_type(shape: Sequence[int], risky: bool = False) -> Optional[str]: + """Infer the type of data based on the shape of the tensors. + + Args: + shape (Sequence[int]): The shape of the data + risky(bool): some guesses are more likely to be wrong. + """ + # (samples,) or (samples,logits) + if len(shape) in (1, 2): + return "label" + # Assume image mask like fashion mnist: (no color channel) + # This is risky because RNNs often have 3 dim tensors: batch, time, channels + if risky and len(shape) == 3: + return "image" + if len(shape) == 4: + if shape[-1] in (1, 3, 4): + # (samples, height, width, Y \ RGB \ RGBA) + return "image" + else: + # (samples, height, width, logits) + return "segmentation_mask" + return None + + +def download_file_from_url( + dest_path: str, source_url: str, api_key: Optional[str] = None +) -> None: + auth = None + if not _thread_local_api_settings.cookies: + auth = ("api", api_key or "") + response = requests.get( + source_url, + auth=auth, + headers=_thread_local_api_settings.headers, + cookies=_thread_local_api_settings.cookies, + stream=True, + timeout=5, + ) + response.raise_for_status() + + if os.sep in dest_path: + filesystem.mkdir_exists_ok(os.path.dirname(dest_path)) + with fsync_open(dest_path, "wb") as file: + for data in response.iter_content(chunk_size=1024): + file.write(data) + + +def download_file_into_memory(source_url: str, api_key: Optional[str] = None) -> bytes: + auth = None + if not _thread_local_api_settings.cookies: + auth = ("api", api_key or "") + response = requests.get( + source_url, + auth=auth, + headers=_thread_local_api_settings.headers, + cookies=_thread_local_api_settings.cookies, + stream=True, + timeout=5, + ) + response.raise_for_status() + return response.content + + +def isatty(ob: IO) -> bool: + return hasattr(ob, "isatty") and ob.isatty() + + +def to_human_size(size: int, units: Optional[List[Tuple[str, Any]]] = None) -> str: + units = units or POW_10_BYTES + unit, value = units[0] + factor = round(float(size) / value, 1) + return ( + f"{factor}{unit}" + if factor < 1024 or len(units) == 1 + else to_human_size(size, units[1:]) + ) + + +def from_human_size(size: str, units: Optional[List[Tuple[str, Any]]] = None) -> int: + units = units or POW_10_BYTES + units_dict = {unit.upper(): value for (unit, value) in units} + regex = re.compile( + r"(\d+\.?\d*)\s*({})?".format("|".join(units_dict.keys())), re.IGNORECASE + ) + match = re.match(regex, size) + if not match: + raise ValueError("size must be of the form `10`, `10B` or `10 B`.") + factor, unit = ( + float(match.group(1)), + units_dict[match.group(2).upper()] if match.group(2) else 1, + ) + return int(factor * unit) + + +def auto_project_name(program: Optional[str]) -> str: + # if we're in git, set project name to git repo name + relative path within repo + from wandb.sdk.lib.gitlib import GitRepo + + root_dir = GitRepo().root_dir + if root_dir is None: + return "uncategorized" + # On windows, GitRepo returns paths in unix style, but os.path is windows + # style. Coerce here. + root_dir = to_native_slash_path(root_dir) + repo_name = os.path.basename(root_dir) + if program is None: + return str(repo_name) + if not os.path.isabs(program): + program = os.path.join(os.curdir, program) + prog_dir = os.path.dirname(os.path.abspath(program)) + if not prog_dir.startswith(root_dir): + return str(repo_name) + project = repo_name + sub_path = os.path.relpath(prog_dir, root_dir) + if sub_path != ".": + project += "-" + sub_path + return str(project.replace(os.sep, "_")) + + +def are_paths_on_same_drive(path1: str, path2: str) -> bool: + """Check if two paths are on the same drive. + + This check is only relevant on Windows, + since the concept of drives only exists on Windows. + """ + if platform.system() != "Windows": + return True + + try: + path1_drive = pathlib.Path(path1).resolve().drive + path2_drive = pathlib.Path(path2).resolve().drive + except OSError: + # If either path is not a valid Windows path, an OSError is raised. + return False + + return path1_drive == path2_drive + + +# TODO(hugh): Deprecate version here and use wandb/sdk/lib/paths.py +def to_forward_slash_path(path: str) -> str: + if platform.system() == "Windows": + path = path.replace("\\", "/") + return path + + +# TODO(hugh): Deprecate version here and use wandb/sdk/lib/paths.py +def to_native_slash_path(path: str) -> FilePathStr: + return FilePathStr(path.replace("/", os.sep)) + + +def check_and_warn_old(files: List[str]) -> bool: + if "wandb-metadata.json" in files: + wandb.termwarn("These runs were logged with a previous version of wandb.") + wandb.termwarn( + "Run pip install wandb<0.10.0 to get the old library and sync your runs." + ) + return True + return False + + +class ImportMetaHook: + def __init__(self) -> None: + self.modules: Dict[str, ModuleType] = dict() + self.on_import: Dict[str, list] = dict() + + def add(self, fullname: str, on_import: Callable) -> None: + self.on_import.setdefault(fullname, []).append(on_import) + + def install(self) -> None: + sys.meta_path.insert(0, self) # type: ignore + + def uninstall(self) -> None: + sys.meta_path.remove(self) # type: ignore + + def find_module( + self, fullname: str, path: Optional[str] = None + ) -> Optional["ImportMetaHook"]: + if fullname in self.on_import: + return self + return None + + def load_module(self, fullname: str) -> ModuleType: + self.uninstall() + mod = importlib.import_module(fullname) + self.install() + self.modules[fullname] = mod + on_imports = self.on_import.get(fullname) + if on_imports: + for f in on_imports: + f() + return mod + + def get_modules(self) -> Tuple[str, ...]: + return tuple(self.modules) + + def get_module(self, module: str) -> ModuleType: + return self.modules[module] + + +_import_hook: Optional[ImportMetaHook] = None + + +def add_import_hook(fullname: str, on_import: Callable) -> None: + global _import_hook + if _import_hook is None: + _import_hook = ImportMetaHook() + _import_hook.install() + _import_hook.add(fullname, on_import) + + +def host_from_path(path: Optional[str]) -> str: + """Return the host of the path.""" + url = urllib.parse.urlparse(path) + return str(url.netloc) + + +def uri_from_path(path: Optional[str]) -> str: + """Return the URI of the path.""" + url = urllib.parse.urlparse(path) + uri = url.path if url.path[0] != "/" else url.path[1:] + return str(uri) + + +def is_unicode_safe(stream: TextIO) -> bool: + """Return True if the stream supports UTF-8.""" + encoding = getattr(stream, "encoding", None) + return encoding.lower() in {"utf-8", "utf_8"} if encoding else False + + +def _has_internet() -> bool: + """Returns whether we have internet access. + + Checks for internet access by attempting to open a DNS connection to + Google's root servers. + """ + try: + s = socket.create_connection(("8.8.8.8", 53), 0.5) + s.close() + except OSError: + return False + + return True + + +def rand_alphanumeric( + length: int = 8, rand: Optional[Union[ModuleType, random.Random]] = None +) -> str: + wandb.termerror("rand_alphanumeric is deprecated, use 'secrets.token_hex'") + rand = rand or random + return "".join(rand.choice("0123456789ABCDEF") for _ in range(length)) + + +@contextlib.contextmanager +def fsync_open( + path: StrPath, mode: str = "w", encoding: Optional[str] = None +) -> Generator[IO[Any], None, None]: + """Open a path for I/O and guarantee that the file is flushed and synced.""" + with open(path, mode, encoding=encoding) as f: + yield f + + f.flush() + os.fsync(f.fileno()) + + +def _is_kaggle() -> bool: + return ( + os.getenv("KAGGLE_KERNEL_RUN_TYPE") is not None + or "kaggle_environments" in sys.modules + ) + + +def _is_likely_kaggle() -> bool: + # Telemetry to mark first runs from Kagglers. + return ( + _is_kaggle() + or os.path.exists( + os.path.expanduser(os.path.join("~", ".kaggle", "kaggle.json")) + ) + or "kaggle" in sys.modules + ) + + +def _is_databricks() -> bool: + # check if we are running inside a databricks notebook by + # inspecting sys.modules, searching for dbutils and verifying that + # it has the appropriate structure + + if "dbutils" in sys.modules: + dbutils = sys.modules["dbutils"] + if hasattr(dbutils, "shell"): + shell = dbutils.shell + if hasattr(shell, "sc"): + sc = shell.sc + if hasattr(sc, "appName"): + return bool(sc.appName == "Databricks Shell") + return False + + +def _is_py_requirements_or_dockerfile(path: str) -> bool: + file = os.path.basename(path) + return ( + file.endswith(".py") + or file.startswith("Dockerfile") + or file == "requirements.txt" + ) + + +def artifact_to_json(artifact: "Artifact") -> Dict[str, Any]: + return { + "_type": "artifactVersion", + "_version": "v0", + "id": artifact.id, + "version": artifact.source_version, + "sequenceName": artifact.source_name.split(":")[0], + "usedAs": artifact.use_as, + } + + +def check_dict_contains_nested_artifact(d: dict, nested: bool = False) -> bool: + for item in d.values(): + if isinstance(item, dict): + contains_artifacts = check_dict_contains_nested_artifact(item, True) + if contains_artifacts: + return True + elif (isinstance(item, wandb.Artifact) or _is_artifact_string(item)) and nested: + return True + return False + + +def load_json_yaml_dict(config: str) -> Any: + ext = os.path.splitext(config)[-1] + if ext == ".json": + with open(config) as f: + return json.load(f) + elif ext == ".yaml": + with open(config) as f: + return yaml.safe_load(f) + else: + try: + return json.loads(config) + except ValueError: + return None + + +def _parse_entity_project_item(path: str) -> tuple: + """Parse paths with the following formats: {item}, {project}/{item}, & {entity}/{project}/{item}. + + Args: + path: `str`, input path; must be between 0 and 3 in length. + + Returns: + tuple of length 3 - (item, project, entity) + + Example: + alias, project, entity = _parse_entity_project_item("myproj/mymodel:best") + + assert entity == "" + assert project == "myproj" + assert alias == "mymodel:best" + + """ + words = path.split("/") + if len(words) > 3: + raise ValueError( + "Invalid path: must be str the form {item}, {project}/{item}, or {entity}/{project}/{item}" + ) + padded_words = [""] * (3 - len(words)) + words + return tuple(reversed(padded_words)) + + +def _resolve_aliases(aliases: Optional[Union[str, Iterable[str]]]) -> List[str]: + """Add the 'latest' alias and ensure that all aliases are unique. + + Takes in `aliases` which can be None, str, or List[str] and returns List[str]. + Ensures that "latest" is always present in the returned list. + + Args: + aliases: `Optional[Union[str, List[str]]]` + + Returns: + List[str], with "latest" always present. + + Usage: + + ```python + aliases = _resolve_aliases(["best", "dev"]) + assert aliases == ["best", "dev", "latest"] + + aliases = _resolve_aliases("boom") + assert aliases == ["boom", "latest"] + ``` + """ + aliases = aliases or ["latest"] + + if isinstance(aliases, str): + aliases = [aliases] + + try: + return list(set(aliases) | {"latest"}) + except TypeError as exc: + raise ValueError("`aliases` must be Iterable or None") from exc + + +def _is_artifact_object(v: Any) -> "TypeGuard[wandb.Artifact]": + return isinstance(v, wandb.Artifact) + + +def _is_artifact_string(v: Any) -> "TypeGuard[str]": + return isinstance(v, str) and v.startswith("wandb-artifact://") + + +def _is_artifact_version_weave_dict(v: Any) -> "TypeGuard[dict]": + return isinstance(v, dict) and v.get("_type") == "artifactVersion" + + +def _is_artifact_representation(v: Any) -> bool: + return ( + _is_artifact_object(v) + or _is_artifact_string(v) + or _is_artifact_version_weave_dict(v) + ) + + +def parse_artifact_string(v: str) -> Tuple[str, Optional[str], bool]: + if not v.startswith("wandb-artifact://"): + raise ValueError(f"Invalid artifact string: {v}") + parsed_v = v[len("wandb-artifact://") :] + base_uri = None + url_info = urllib.parse.urlparse(parsed_v) + if url_info.scheme != "": + base_uri = f"{url_info.scheme}://{url_info.netloc}" + parts = url_info.path.split("/")[1:] + else: + parts = parsed_v.split("/") + if parts[0] == "_id": + # for now can't fetch paths but this will be supported in the future + # when we allow passing typed media objects, this can be extended + # to include paths + return parts[1], base_uri, True + + if len(parts) < 3: + raise ValueError(f"Invalid artifact string: {v}") + + # for now can't fetch paths but this will be supported in the future + # when we allow passing typed media objects, this can be extended + # to include paths + entity, project, name_and_alias_or_version = parts[:3] + return f"{entity}/{project}/{name_and_alias_or_version}", base_uri, False + + +def _get_max_cli_version() -> Union[str, None]: + max_cli_version = wandb.api.max_cli_version() + return str(max_cli_version) if max_cli_version is not None else None + + +def ensure_text( + string: Union[str, bytes], encoding: str = "utf-8", errors: str = "strict" +) -> str: + """Coerce s to str.""" + if isinstance(string, bytes): + return string.decode(encoding, errors) + elif isinstance(string, str): + return string + else: + raise TypeError(f"not expecting type {type(string)!r}") + + +def make_artifact_name_safe(name: str) -> str: + """Make an artifact name safe for use in artifacts.""" + # artifact names may only contain alphanumeric characters, dashes, underscores, and dots. + cleaned = re.sub(r"[^a-zA-Z0-9_\-.]", "_", name) + if len(cleaned) <= 128: + return cleaned + # truncate with dots in the middle using regex + return re.sub(r"(^.{63}).*(.{63}$)", r"\g<1>..\g<2>", cleaned) + + +def make_docker_image_name_safe(name: str) -> str: + """Make a docker image name safe for use in artifacts.""" + safe_chars = RE_DOCKER_IMAGE_NAME_CHARS.sub("__", name.lower()) + deduped = RE_DOCKER_IMAGE_NAME_SEPARATOR_REPEAT.sub("__", safe_chars) + trimmed_start = RE_DOCKER_IMAGE_NAME_SEPARATOR_START.sub("", deduped) + trimmed = RE_DOCKER_IMAGE_NAME_SEPARATOR_END.sub("", trimmed_start) + return trimmed if trimmed else "image" + + +def merge_dicts( + source: Dict[str, Any], + destination: Dict[str, Any], +) -> Dict[str, Any]: + """Recursively merge two dictionaries. + + This mutates the destination and its nested dictionaries and lists. + + Instances of `dict` are recursively merged and instances of `list` + are appended to the destination. If the destination type is not + `dict` or `list`, respectively, the key is overwritten with the + source value. + + For all other types, the source value overwrites the destination value. + """ + for key, value in source.items(): + if isinstance(value, dict): + node = destination.get(key) + if isinstance(node, dict): + merge_dicts(value, node) + else: + destination[key] = value + + elif isinstance(value, list): + dest_value = destination.get(key) + if isinstance(dest_value, list): + dest_value.extend(value) + else: + destination[key] = value + + else: + destination[key] = value + + return destination + + +def coalesce(*arg: Any) -> Any: + """Return the first non-none value in the list of arguments. + + Similar to ?? in C#. + """ + return next((a for a in arg if a is not None), None) + + +def recursive_cast_dictlike_to_dict(d: Dict[str, Any]) -> Dict[str, Any]: + for k, v in d.items(): + if isinstance(v, dict): + recursive_cast_dictlike_to_dict(v) + elif hasattr(v, "keys"): + d[k] = dict(v) + recursive_cast_dictlike_to_dict(d[k]) + return d + + +def remove_keys_with_none_values( + d: Union[Dict[str, Any], Any], +) -> Union[Dict[str, Any], Any]: + # otherwise iterrows will create a bunch of ugly charts + if not isinstance(d, dict): + return d + + if isinstance(d, dict): + new_dict = {} + for k, v in d.items(): + new_v = remove_keys_with_none_values(v) + if new_v is not None and not (isinstance(new_v, dict) and len(new_v) == 0): + new_dict[k] = new_v + return new_dict if new_dict else None + + +def batched(n: int, iterable: Iterable[T]) -> Generator[List[T], None, None]: + i = iter(iterable) + batch = list(itertools.islice(i, n)) + while batch: + yield batch + batch = list(itertools.islice(i, n)) + + +def random_string(length: int = 12) -> str: + """Generate a random string of a given length. + + :param length: Length of the string to generate. + :return: Random string. + """ + return "".join( + secrets.choice(string.ascii_lowercase + string.digits) for _ in range(length) + ) + + +def sample_with_exponential_decay_weights( + xs: Union[Iterable, Iterable[Iterable]], + ys: Iterable[Iterable], + keys: Optional[Iterable] = None, + sample_size: int = 1500, +) -> Tuple[List, List, Optional[List]]: + """Sample from a list of lists with weights that decay exponentially. + + May be used with the wandb.plot.line_series function. + """ + xs_array = np.array(xs) + ys_array = np.array(ys) + keys_array = np.array(keys) if keys else None + weights = np.exp(-np.arange(len(xs_array)) / len(xs_array)) + weights /= np.sum(weights) + sampled_indices = np.random.choice(len(xs_array), size=sample_size, p=weights) + sampled_xs = xs_array[sampled_indices].tolist() + sampled_ys = ys_array[sampled_indices].tolist() + sampled_keys = keys_array[sampled_indices].tolist() if keys_array else None + + return sampled_xs, sampled_ys, sampled_keys + + +@dataclasses.dataclass(frozen=True) +class InstalledDistribution: + """An installed distribution. + + Attributes: + key: The distribution name as it would be imported. + version: The distribution's version string. + """ + + key: str + version: str + + +def working_set() -> Iterable[InstalledDistribution]: + """Return the working set of installed distributions.""" + from importlib.metadata import distributions + + for d in distributions(): + try: + # In some distributions, the "Name" attribute may not be present, + # which can raise a KeyError. To handle this, we catch the exception + # and skip those distributions. + # For additional context, see: https://github.com/python/importlib_metadata/issues/371. + + # From Sentry events we observed that UnicodeDecodeError can occur when + # trying to decode the metadata of a distribution. To handle this, we catch + # the exception and skip those distributions. + yield InstalledDistribution(key=d.metadata["Name"], version=d.version) + except (KeyError, UnicodeDecodeError): + pass + + +def get_core_path() -> str: + """Returns the path to the wandb-core binary. + + The path can be set explicitly via the _WANDB_CORE_PATH environment + variable. Otherwise, the path to the binary in the current package + is returned. + + Returns: + str: The path to the wandb-core package. + + Raises: + WandbCoreNotAvailableError: If wandb-core was not built for the current system. + """ + # NOTE: Environment variable _WANDB_CORE_PATH is a temporary development feature + # to assist in running the core service from a live development directory. + path_from_env: str = os.environ.get("_WANDB_CORE_PATH", "") + if path_from_env: + wandb.termwarn( + f"Using wandb-core from path `_WANDB_CORE_PATH={path_from_env}`. " + "This is a development feature and may not work as expected." + ) + return path_from_env + + bin_path = pathlib.Path(__file__).parent / "bin" / "wandb-core" + if not bin_path.exists(): + raise WandbCoreNotAvailableError( + f"File not found: {bin_path}." + " Please contact support at support@wandb.com." + f" Your platform is: {platform.platform()}." + ) + + return str(bin_path) + + +class NonOctalStringDumper(yaml.Dumper): + """Prevents strings containing non-octal values like "008" and "009" from being converted to numbers in in the yaml string saved as the sweep config.""" + + def represent_scalar(self, tag, value, style=None): + if tag == "tag:yaml.org,2002:str" and value.startswith("0") and len(value) > 1: + return super().represent_scalar(tag, value, style="'") + return super().represent_scalar(tag, value, style) diff --git a/lib/python3.12/site-packages/wandb/wandb_agent.py b/lib/python3.12/site-packages/wandb/wandb_agent.py new file mode 100644 index 0000000000000000000000000000000000000000..fb14bd63ba21f006c403949ae986873e68bfd64a --- /dev/null +++ b/lib/python3.12/site-packages/wandb/wandb_agent.py @@ -0,0 +1,580 @@ +import logging +import multiprocessing +import os +import platform +import queue +import re +import signal +import socket +import subprocess +import sys +import time +import traceback +from typing import Any, Callable, Dict, List, Optional + +import yaml + +import wandb +from wandb import util, wandb_lib, wandb_sdk +from wandb.agents.pyagent import pyagent +from wandb.apis import InternalApi +from wandb.sdk.launch.sweeps import utils as sweep_utils +from wandb.sdk.lib import ipython + +logger = logging.getLogger(__name__) + + +class AgentError(Exception): + pass + + +class AgentProcess: + """Launch and manage a process.""" + + def __init__( + self, env=None, command=None, function=None, run_id=None, in_jupyter=None + ): + self._popen = None + self._proc = None + self._finished_q = multiprocessing.Queue() + self._proc_killed = False + + if command: + if platform.system() == "Windows": + kwargs = dict(creationflags=subprocess.CREATE_NEW_PROCESS_GROUP) + else: + kwargs = dict(preexec_fn=os.setpgrp) + if env.get(wandb.env.SERVICE): + env.pop(wandb.env.SERVICE) + self._popen = subprocess.Popen(command, env=env, **kwargs) + elif function: + self._proc = multiprocessing.Process( + target=self._start, + args=(self._finished_q, env, function, run_id, in_jupyter), + ) + self._proc.start() + else: + raise AgentError("Agent Process requires command or function") + + def _start(self, finished_q, env, function, run_id, in_jupyter): + if env: + for k, v in env.items(): + os.environ[k] = v + + # call user function + wandb.termlog(f"Agent Started Run: {run_id}") + if function: + function() + wandb.termlog(f"Agent Finished Run: {run_id}\n") + + # complete the run + run = wandb.run + if run: + wandb.join() + + # signal that the process is finished + finished_q.put(True) + + def poll(self): + if self._popen: + return self._popen.poll() + if self._proc_killed: + # we need to join process to prevent zombies + self._proc.join() + return True + try: + finished = self._finished_q.get(False, 0) + if finished: + return True + except queue.Empty: + pass + return + + def wait(self): + if self._popen: + # if on windows, wait() will block and we won't be able to interrupt + if platform.system() == "Windows": + while True: + p = self._popen.poll() + if p is not None: + return p + time.sleep(1) + return self._popen.wait() + return self._proc.join() + + def kill(self): + if self._popen: + return self._popen.kill() + pid = self._proc.pid + if pid: + ret = os.kill(pid, signal.SIGKILL) + self._proc_killed = True + return ret + return + + def terminate(self): + if self._popen: + # windows terminate is too strong, send Ctrl-C instead + if platform.system() == "Windows": + return self._popen.send_signal(signal.CTRL_C_EVENT) + return self._popen.terminate() + return self._proc.terminate() + + +class Agent: + POLL_INTERVAL = 5 + REPORT_INTERVAL = 0 + KILL_DELAY = 30 + FLAPPING_MAX_SECONDS = 60 + FLAPPING_MAX_FAILURES = 3 + MAX_INITIAL_FAILURES = 5 + DEFAULT_SWEEP_COMMAND: List[str] = [ + "${env}", + "${interpreter}", + "${program}", + "${args}", + ] + SWEEP_COMMAND_ENV_VAR_REGEX = re.compile(r"\$\{envvar\:([A-Z0-9_]*)\}") + + def __init__( + self, api, queue, sweep_id=None, function=None, in_jupyter=None, count=None + ): + self._api = api + self._queue = queue + self._run_processes = {} # keyed by run.id (GQL run name) + self._server_responses = [] + self._sweep_id = sweep_id + self._in_jupyter = in_jupyter + self._log = [] + self._running = True + self._last_report_time = None + self._function = function + self._report_interval = wandb.env.get_agent_report_interval( + self.REPORT_INTERVAL + ) + self._kill_delay = wandb.env.get_agent_kill_delay(self.KILL_DELAY) + self._finished = 0 + self._failed = 0 + self._count = count + self._sweep_command = [] + self._max_initial_failures = wandb.env.get_agent_max_initial_failures( + self.MAX_INITIAL_FAILURES + ) + if self._report_interval is None: + raise AgentError("Invalid agent report interval") + if self._kill_delay is None: + raise AgentError("Invalid agent kill delay") + # if the directory to log to is not set, set it + if os.environ.get("WANDB_DIR") is None: + os.environ["WANDB_DIR"] = os.path.abspath(os.getcwd()) + + def is_flapping(self): + """Determine if the process is flapping. + + Flapping occurs if the agents receives FLAPPING_MAX_FAILURES non-0 exit codes in + the first FLAPPING_MAX_SECONDS. + """ + if os.getenv(wandb.env.AGENT_DISABLE_FLAPPING) == "true": + return False + if time.time() < wandb.START_TIME + self.FLAPPING_MAX_SECONDS: + return self._failed >= self.FLAPPING_MAX_FAILURES + + def is_failing(self): + return ( + self._failed >= self._finished + and self._max_initial_failures <= self._failed + ) + + def run(self): # noqa: C901 + # TODO: catch exceptions, handle errors, show validation warnings, and make more generic + sweep_obj = self._api.sweep(self._sweep_id, "{}") + if sweep_obj: + sweep_yaml = sweep_obj.get("config") + if sweep_yaml: + sweep_config = yaml.safe_load(sweep_yaml) + if sweep_config: + sweep_command = sweep_config.get("command") + if sweep_command and isinstance(sweep_command, list): + self._sweep_command = sweep_command + + # TODO: include sweep ID + agent = self._api.register_agent(socket.gethostname(), sweep_id=self._sweep_id) + agent_id = agent["id"] + + try: + while self._running: + commands = util.read_many_from_queue( + self._queue, 100, self.POLL_INTERVAL + ) + for command in commands: + command["resp_queue"].put(self._process_command(command)) + + now = util.stopwatch_now() + if self._last_report_time is None or ( + self._report_interval != 0 + and now > self._last_report_time + self._report_interval + ): + logger.info("Running runs: %s", list(self._run_processes.keys())) + self._last_report_time = now + run_status = {} + for run_id, run_process in list(self._run_processes.items()): + poll_result = run_process.poll() + if poll_result is None: + run_status[run_id] = True + continue + elif ( + not isinstance(poll_result, bool) + and isinstance(poll_result, int) + and poll_result > 0 + ): + self._failed += 1 + if self.is_flapping(): + logger.error( + "Detected %i failed runs in the first %i seconds, shutting down.", + self.FLAPPING_MAX_FAILURES, + self.FLAPPING_MAX_SECONDS, + ) + logger.info( + "To disable this check set WANDB_AGENT_DISABLE_FLAPPING=true" + ) + self._running = False + break + if self.is_failing(): + logger.error( + "Detected %i failed runs in a row, shutting down.", + self._max_initial_failures, + ) + logger.info( + "To change this value set WANDB_AGENT_MAX_INITIAL_FAILURES=val" + ) + self._running = False + break + logger.info("Cleaning up finished run: %s", run_id) + + # wandb.teardown() was added with wandb service and is a hammer to make + # sure that active runs are finished before moving on to another agent run + # + # In the future, a lighter weight way to implement this could be to keep a + # service process open for all the agent instances and inform_finish when + # the run should be marked complete. This however could require + # inform_finish on every run created by this process. + if hasattr(wandb, "teardown"): + exit_code = 0 + if isinstance(poll_result, int): + exit_code = poll_result + elif isinstance(poll_result, bool): + exit_code = -1 + wandb.teardown(exit_code) + + del self._run_processes[run_id] + self._last_report_time = None + self._finished += 1 + + if self._count and self._finished >= self._count or not self._running: + self._running = False + continue + + commands = self._api.agent_heartbeat(agent_id, {}, run_status) + + # TODO: send _server_responses + self._server_responses = [] + for command in commands: + self._server_responses.append(self._process_command(command)) + + except KeyboardInterrupt: + try: + wandb.termlog( + "Ctrl-c pressed. Waiting for runs to end. Press ctrl-c again to terminate them." + ) + for _, run_process in self._run_processes.items(): + run_process.wait() + except KeyboardInterrupt: + pass + finally: + try: + if not self._in_jupyter: + wandb.termlog("Terminating and syncing runs. Press ctrl-c to kill.") + for _, run_process in self._run_processes.items(): + try: + run_process.terminate() + except OSError: + pass # if process is already dead + for _, run_process in self._run_processes.items(): + run_process.wait() + except KeyboardInterrupt: + wandb.termlog("Killing runs and quitting.") + for _, run_process in self._run_processes.items(): + try: + run_process.kill() + except OSError: + pass # if process is already dead + + def _process_command(self, command): + logger.info( + "Agent received command: %s" + % (command["type"] if "type" in command else "Unknown") + ) + response = { + "id": command.get("id"), + "result": None, + } + try: + command_type = command["type"] + if command_type == "run": + result = self._command_run(command) + elif command_type == "stop": + result = self._command_stop(command) + elif command_type == "exit": + result = self._command_exit(command) + elif command_type == "resume": + result = self._command_run(command) + else: + raise AgentError(f"No such command: {command_type}") # noqa: TRY301 + response["result"] = result + except Exception: + logger.exception("Exception while processing command: %s", command) + ex_type, ex, tb = sys.exc_info() + response["exception"] = f"{ex_type.__name__}: {str(ex)}" + response["traceback"] = traceback.format_tb(tb) + del tb + + self._log.append((command, response)) + + return response + + def _command_run(self, command): + logger.info( + "Agent starting run with config:\n" + + "\n".join( + ["\t{}: {}".format(k, v["value"]) for k, v in command["args"].items()] + ) + ) + if self._in_jupyter: + wandb.termlog( + f"Agent Starting Run: {command.get('run_id')} with config:\n" + + "\n".join( + [f"\t{k}: {v['value']}" for k, v in command["args"].items()] + ) + ) + + # Setup sweep command + sweep_command: List[str] = sweep_utils.create_sweep_command(self._sweep_command) + + run_id = command.get("run_id") + sweep_id = os.environ.get(wandb.env.SWEEP_ID) + # TODO(jhr): move into settings + config_file = os.path.join( + "wandb", "sweep-" + sweep_id, "config-" + run_id + ".yaml" + ) + json_file = os.path.join( + "wandb", "sweep-" + sweep_id, "config-" + run_id + ".json" + ) + + os.environ[wandb.env.RUN_ID] = run_id + + base_dir = os.environ.get(wandb.env.DIR, "") + sweep_param_path = os.path.join(base_dir, config_file) + os.environ[wandb.env.SWEEP_PARAM_PATH] = sweep_param_path + wandb_lib.config_util.save_config_file_from_dict( + sweep_param_path, command["args"] + ) + + env = dict(os.environ) + + sweep_vars: Dict[str, Any] = sweep_utils.create_sweep_command_args(command) + + if "${args_json_file}" in sweep_command: + with open(json_file, "w") as fp: + fp.write(sweep_vars["args_json"][0]) + + if self._function: + # make sure that each run regenerates setup singleton + wandb.teardown() + proc = AgentProcess( + function=self._function, + env=env, + run_id=run_id, + in_jupyter=self._in_jupyter, + ) + else: + sweep_vars["interpreter"] = ["python"] + sweep_vars["program"] = [command["program"]] + sweep_vars["args_json_file"] = [json_file] + if not platform.system() == "Windows": + sweep_vars["env"] = ["/usr/bin/env"] + command_list = [] + for c in sweep_command: + c = str(c) + if c.startswith("${") and c.endswith("}"): + replace_list = sweep_vars.get(c[2:-1]) + command_list += replace_list or [] + else: + command_list += [c] + logger.info( + "About to run command: {}".format( + " ".join(f'"{c}"' if " " in c else c for c in command_list) + ) + ) + proc = AgentProcess(command=command_list, env=env) + self._run_processes[run_id] = proc + + # we keep track of when we sent the sigterm to give processes a chance + # to handle the signal before sending sigkill every heartbeat + self._run_processes[run_id].last_sigterm_time = None + self._last_report_time = None + + def _command_stop(self, command): + run_id = command["run_id"] + if run_id in self._run_processes: + proc = self._run_processes[run_id] + now = util.stopwatch_now() + if proc.last_sigterm_time is None: + proc.last_sigterm_time = now + logger.info("Stop: %s", run_id) + try: + proc.terminate() + except OSError: # if process is already dead + pass + elif now > proc.last_sigterm_time + self._kill_delay: + logger.info("Kill: %s", run_id) + try: + proc.kill() + except OSError: # if process is already dead + pass + else: + logger.error("Run %s not running", run_id) + + def _command_exit(self, command): + logger.info("Received exit command. Killing runs and quitting.") + for _, proc in self._run_processes.items(): + try: + proc.kill() + except OSError: + # process is already dead + pass + self._running = False + + +class AgentApi: + def __init__(self, queue): + self._queue = queue + self._command_id = 0 + self._multiproc_manager = multiprocessing.Manager() + + def command(self, command): + command["origin"] = "local" + command["id"] = f"local-{self._command_id}" + self._command_id += 1 + resp_queue = self._multiproc_manager.Queue() + command["resp_queue"] = resp_queue + self._queue.put(command) + result = resp_queue.get() + print("result:", result) # noqa: T201 + if "exception" in result: + print("Exception occurred while running command") # noqa: T201 + for line in result["traceback"]: + print(line.strip()) # noqa: T201 + print(result["exception"]) # noqa: T201 + return result + + +def run_agent( + sweep_id, function=None, in_jupyter=None, entity=None, project=None, count=None +): + parts = dict(entity=entity, project=project, name=sweep_id) + err = sweep_utils.parse_sweep_id(parts) + if err: + wandb.termerror(err) + return + entity = parts.get("entity") or entity + project = parts.get("project") or project + sweep_id = parts.get("name") or sweep_id + + if entity: + wandb.env.set_entity(entity) + if project: + wandb.env.set_project(project) + if sweep_id: + # TODO(jhr): remove when jobspec is merged + os.environ[wandb.env.SWEEP_ID] = sweep_id + logger.setLevel(logging.DEBUG) + ch = logging.StreamHandler() + log_level = logging.DEBUG + if in_jupyter: + log_level = logging.ERROR + ch.setLevel(log_level) + formatter = logging.Formatter( + "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + ) + ch.setFormatter(formatter) + try: + logger.addHandler(ch) + + api = InternalApi() + queue = multiprocessing.Queue() + agent = Agent( + api, + queue, + sweep_id=sweep_id, + function=function, + in_jupyter=in_jupyter, + count=count, + ) + agent.run() + finally: + # make sure we remove the logging handler (important for jupyter notebooks) + logger.removeHandler(ch) + + +def agent( + sweep_id: str, + function: Optional[Callable] = None, + entity: Optional[str] = None, + project: Optional[str] = None, + count: Optional[int] = None, +) -> None: + """Start one or more sweep agents. + + The sweep agent uses the `sweep_id` to know which sweep it + is a part of, what function to execute, and (optionally) how + many agents to run. + + Args: + sweep_id: The unique identifier for a sweep. A sweep ID + is generated by W&B CLI or Python SDK. + function: A function to call instead of the "program" + specified in the sweep config. + entity: The username or team name where you want to send W&B + runs created by the sweep to. Ensure that the entity you + specify already exists. If you don't specify an entity, + the run will be sent to your default entity, + which is usually your username. + project: The name of the project where W&B runs created from + the sweep are sent to. If the project is not specified, the + run is sent to a project labeled "Uncategorized". + count: The number of sweep config trials to try. + """ + global _INSTANCES + _INSTANCES += 1 + try: + # make sure we are logged in + wandb_sdk.wandb_login._login(_silent=True) + if function: + return pyagent(sweep_id, function, entity, project, count) + return run_agent( + sweep_id, + function=function, + in_jupyter=ipython.in_jupyter(), + entity=entity, + project=project, + count=count, + ) + finally: + _INSTANCES -= 1 + + +_INSTANCES = 0 + + +def _is_running(): + return bool(_INSTANCES) diff --git a/lib/python3.12/site-packages/wandb/wandb_controller.py b/lib/python3.12/site-packages/wandb/wandb_controller.py new file mode 100644 index 0000000000000000000000000000000000000000..f819dc995e27192c7debdc87fa1042b3e77f1c98 --- /dev/null +++ b/lib/python3.12/site-packages/wandb/wandb_controller.py @@ -0,0 +1,719 @@ +"""Sweep controller. + +This module implements the sweep controller. + +On error an exception is raised: + ControllerError + +Example: + import wandb + + # + # create a sweep controller + # + # There are three different ways sweeps can be created: + # (1) create with sweep id from `wandb sweep` command + sweep_id = 'xyzxyz2' + tuner = wandb.controller(sweep_id) + # (2) create with sweep config + sweep_config = {} + tuner = wandb.controller() + tuner.configure(sweep_config) + tuner.create() + # (3) create by constructing programmatic sweep configuration + tuner = wandb.controller() + tuner.configure_search('random') + tuner.configure_program('train-dummy.py') + tuner.configure_parameter('param1', values=[1,2,3]) + tuner.configure_parameter('param2', values=[1,2,3]) + tuner.configure_controller(type="local") + tuner.create() + # + # run the sweep controller + # + # There are three different ways sweeps can be executed: + # (1) run to completion + tuner.run() + # (2) run in a simple loop + while not tuner.done(): + tuner.step() + tuner.print_status() + # (3) run in a more complex loop + while not tuner.done(): + params = tuner.search() + tuner.schedule(params) + runs = tuner.stopping() + if runs: + tuner.stop_runs(runs) +""" + +import json +import os +import random +import string +import time +from typing import Callable, Dict, List, Optional, Tuple, Union + +import yaml + +from wandb import env +from wandb.apis import InternalApi +from wandb.sdk import wandb_sweep +from wandb.sdk.launch.sweeps.utils import ( + handle_sweep_config_violations, + sweep_config_err_text_from_jsonschema_violations, +) +from wandb.util import get_module + +# TODO(jhr): Add metric status +# TODO(jhr): Add print_space +# TODO(jhr): Add print_summary + + +sweeps = get_module( + "sweeps", + required="wandb[sweeps] is required to use the local controller. " + "Please run `pip install wandb[sweeps]`.", +) + + +# This should be something like 'pending' (but we need to make sure everyone else is ok with that) +SWEEP_INITIAL_RUN_STATE = sweeps.RunState.pending + + +def _id_generator(size=10, chars=string.ascii_lowercase + string.digits): + return "".join(random.choice(chars) for _ in range(size)) + + +class ControllerError(Exception): + """Base class for sweep errors.""" + + +class _WandbController: + """Sweep controller class. + + Internal datastructures on the sweep object to coordinate local controller with + cloud controller. + + Data structures: + controller: { + schedule: [ + { id: SCHEDULE_ID + data: {param1: val1, param2: val2}}, + ] + earlystop: [RUN_ID, ...] + scheduler: + scheduled: [ + { id: SCHEDULE_ID + runid: RUN_ID}, + ] + + `controller` is only updated by the client + `scheduler` is only updated by the cloud backend + + Protocols: + Scheduling a run: + - client controller adds a schedule entry on the controller.schedule list + - cloud backend notices the new entry and creates a run with the parameters + - cloud backend adds a scheduled entry on the scheduler.scheduled list + - client controller notices that the run has been scheduled and removes it from + controller.schedule list + + Current implementation details: + - Runs are only schedule if there are no other runs scheduled. + + """ + + def __init__(self, sweep_id_or_config=None, entity=None, project=None): + # sweep id configured in constructor + self._sweep_id: Optional[str] = None + + # configured parameters + # Configuration to be created + self._create: Dict = {} + # Custom search + self._custom_search: Optional[ + Callable[ + [Union[dict, sweeps.SweepConfig], List[sweeps.SweepRun]], + Optional[sweeps.SweepRun], + ] + ] = None + # Custom stopping + self._custom_stopping: Optional[ + Callable[ + [Union[dict, sweeps.SweepConfig], List[sweeps.SweepRun]], + List[sweeps.SweepRun], + ] + ] = None + # Program function (used for future jupyter support) + self._program_function = None + + # The following are updated every sweep step + # raw sweep object (dict of strings) + self._sweep_obj = None + # parsed sweep config (dict) + self._sweep_config: Optional[Union[dict, sweeps.SweepConfig]] = None + # sweep metric used to optimize (str or None) + self._sweep_metric: Optional[str] = None + # list of _Run objects + self._sweep_runs: Optional[List[sweeps.SweepRun]] = None + # dictionary mapping name of run to run object + self._sweep_runs_map: Optional[Dict[str, sweeps.SweepRun]] = None + # scheduler dict (read only from controller) - used as feedback from the server + self._scheduler: Optional[Dict] = None + # controller dict (write only from controller) - used to send commands to server + self._controller: Optional[Dict] = None + # keep track of controller dict from previous step + self._controller_prev_step: Optional[Dict] = None + + # Internal + # Keep track of whether the sweep has been started + self._started: bool = False + # indicate whether there is more to schedule + self._done_scheduling: bool = False + # indicate whether the sweep needs to be created + self._defer_sweep_creation: bool = False + # count of logged lines since last status + self._logged: int = 0 + # last status line printed + self._laststatus: str = "" + # keep track of logged actions for print_actions() + self._log_actions: List[Tuple[str, str]] = [] + # keep track of logged debug for print_debug() + self._log_debug: List[str] = [] + + # all backend commands use internal api + environ = os.environ + if entity: + env.set_entity(entity, env=environ) + if project: + env.set_project(project, env=environ) + self._api = InternalApi(environ=environ) + + if isinstance(sweep_id_or_config, str): + self._sweep_id = sweep_id_or_config + elif isinstance(sweep_id_or_config, dict) or isinstance( + sweep_id_or_config, sweeps.SweepConfig + ): + self._create = sweeps.SweepConfig(sweep_id_or_config) + + # check for custom search and or stopping functions + for config_key, controller_attr in zip( + ["method", "early_terminate"], ["_custom_search", "_custom_stopping"] + ): + if callable(config_key in self._create and self._create[config_key]): + setattr(self, controller_attr, self._create[config_key]) + self._create[config_key] = "custom" + + self._sweep_id = self.create(from_dict=True) + elif sweep_id_or_config is None: + self._defer_sweep_creation = True + return + else: + raise ControllerError("Unhandled sweep controller type") + sweep_obj = self._sweep_object_read_from_backend() + if sweep_obj is None: + raise ControllerError("Can not find sweep") + self._sweep_obj = sweep_obj + + def configure_search( + self, + search: Union[ + str, + Callable[ + [Union[dict, sweeps.SweepConfig], List[sweeps.SweepRun]], + Optional[sweeps.SweepRun], + ], + ], + ): + self._configure_check() + if isinstance(search, str): + self._create["method"] = search + elif callable(search): + self._create["method"] = "custom" + self._custom_search = search + else: + raise ControllerError("Unhandled search type.") + + def configure_stopping( + self, + stopping: Union[ + str, + Callable[ + [Union[dict, sweeps.SweepConfig], List[sweeps.SweepRun]], + List[sweeps.SweepRun], + ], + ], + **kwargs, + ): + self._configure_check() + if isinstance(stopping, str): + self._create.setdefault("early_terminate", {}) + self._create["early_terminate"]["type"] = stopping + for k, v in kwargs.items(): + self._create["early_terminate"][k] = v + elif callable(stopping): + self._custom_stopping = stopping(kwargs) + self._create.setdefault("early_terminate", {}) + self._create["early_terminate"]["type"] = "custom" + else: + raise ControllerError("Unhandled stopping type.") + + def configure_metric(self, metric, goal=None): + self._configure_check() + self._create.setdefault("metric", {}) + self._create["metric"]["name"] = metric + if goal: + self._create["metric"]["goal"] = goal + + def configure_program(self, program): + self._configure_check() + if isinstance(program, str): + self._create["program"] = program + elif callable(program): + self._create["program"] = "__callable__" + self._program_function = program + raise ControllerError("Program functions are not supported yet") + else: + raise ControllerError("Unhandled sweep program type") + + def configure_name(self, name): + self._configure_check() + self._create["name"] = name + + def configure_description(self, description): + self._configure_check() + self._create["description"] = description + + def configure_parameter( + self, + name, + values=None, + value=None, + distribution=None, + min=None, + max=None, + mu=None, + sigma=None, + q=None, + a=None, + b=None, + ): + self._configure_check() + self._create.setdefault("parameters", {}).setdefault(name, {}) + if value is not None or ( + values is None and min is None and max is None and distribution is None + ): + self._create["parameters"][name]["value"] = value + if values is not None: + self._create["parameters"][name]["values"] = values + if distribution is not None: + self._create["parameters"][name]["distribution"] = distribution + if min is not None: + self._create["parameters"][name]["min"] = min + if max is not None: + self._create["parameters"][name]["max"] = max + if mu is not None: + self._create["parameters"][name]["mu"] = mu + if sigma is not None: + self._create["parameters"][name]["sigma"] = sigma + if q is not None: + self._create["parameters"][name]["q"] = q + if a is not None: + self._create["parameters"][name]["a"] = a + if b is not None: + self._create["parameters"][name]["b"] = b + + def configure_controller(self, type): + """Configure controller to local if type == 'local'.""" + self._configure_check() + self._create.setdefault("controller", {}) + self._create["controller"].setdefault("type", type) + + def configure(self, sweep_dict_or_config): + self._configure_check() + if self._create: + raise ControllerError("Already configured.") + if isinstance(sweep_dict_or_config, dict): + self._create = sweep_dict_or_config + elif isinstance(sweep_dict_or_config, str): + self._create = yaml.safe_load(sweep_dict_or_config) + else: + raise ControllerError("Unhandled sweep controller type") + + @property + def sweep_config(self) -> Union[dict, sweeps.SweepConfig]: + return self._sweep_config + + @property + def sweep_id(self) -> str: + return self._sweep_id + + def _log(self) -> None: + self._logged += 1 + + def _error(self, s: str) -> None: + print("ERROR:", s) # noqa: T201 + self._log() + + def _warn(self, s: str) -> None: + print("WARN:", s) # noqa: T201 + self._log() + + def _info(self, s: str) -> None: + print("INFO:", s) # noqa: T201 + self._log() + + def _debug(self, s: str) -> None: + print("DEBUG:", s) # noqa: T201 + self._log() + + def _configure_check(self) -> None: + if self._started: + raise ControllerError("Can not configure after sweep has been started.") + + def _validate(self, config: Dict) -> str: + violations = sweeps.schema_violations_from_proposed_config(config) + msg = ( + sweep_config_err_text_from_jsonschema_violations(violations) + if len(violations) > 0 + else "" + ) + return msg + + def create(self, from_dict: bool = False) -> str: + if self._started: + raise ControllerError("Can not create after sweep has been started.") + if not self._defer_sweep_creation and not from_dict: + raise ControllerError("Can not use create on already created sweep.") + if not self._create: + raise ControllerError("Must configure sweep before create.") + + # validate sweep config + self._create = sweeps.SweepConfig(self._create) + + # Create sweep + sweep_id, warnings = self._api.upsert_sweep(self._create) + handle_sweep_config_violations(warnings) + + print("Create sweep with ID:", sweep_id) # noqa: T201 + sweep_url = wandb_sweep._get_sweep_url(self._api, sweep_id) + if sweep_url: + print("Sweep URL:", sweep_url) # noqa: T201 + self._sweep_id = sweep_id + self._defer_sweep_creation = False + return sweep_id + + def run( + self, + verbose: bool = False, + print_status: bool = True, + print_actions: bool = False, + print_debug: bool = False, + ) -> None: + if verbose: + print_status = True + print_actions = True + print_debug = True + self._start_if_not_started() + while not self.done(): + if print_status: + self.print_status() + self.step() + if print_actions: + self.print_actions() + if print_debug: + self.print_debug() + time.sleep(5) + + def _sweep_object_read_from_backend(self) -> Optional[dict]: + specs_json = {} + if self._sweep_metric: + k = ["_step"] + k.append(self._sweep_metric) + specs_json = {"keys": k, "samples": 100000} + specs = json.dumps(specs_json) + # TODO(jhr): catch exceptions? + sweep_obj = self._api.sweep(self._sweep_id, specs) + if not sweep_obj: + return + self._sweep_obj = sweep_obj + self._sweep_config = yaml.safe_load(sweep_obj["config"]) + self._sweep_metric = self._sweep_config.get("metric", {}).get("name") + + _sweep_runs: List[sweeps.SweepRun] = [] + for r in sweep_obj["runs"]: + rr = r.copy() + if "summaryMetrics" in rr: + if rr["summaryMetrics"]: + rr["summaryMetrics"] = json.loads(rr["summaryMetrics"]) + if "config" not in rr: + raise ValueError("sweep object is missing config") + rr["config"] = json.loads(rr["config"]) + if "history" in rr: + if isinstance(rr["history"], list): + rr["history"] = [json.loads(d) for d in rr["history"]] + else: + raise ValueError( + "Invalid history value: expected list of json strings: {}".format( + rr["history"] + ) + ) + if "sampledHistory" in rr: + sampled_history = [] + for historyDictList in rr["sampledHistory"]: + sampled_history += historyDictList + rr["sampledHistory"] = sampled_history + _sweep_runs.append(sweeps.SweepRun(**rr)) + + self._sweep_runs = _sweep_runs + self._sweep_runs_map = {r.name: r for r in self._sweep_runs} + + self._controller = json.loads(sweep_obj.get("controller") or "{}") + self._scheduler = json.loads(sweep_obj.get("scheduler") or "{}") + self._controller_prev_step = self._controller.copy() + return sweep_obj + + def _sweep_object_sync_to_backend(self) -> None: + if self._controller == self._controller_prev_step: + return + sweep_obj_id = self._sweep_obj["id"] + controller = json.dumps(self._controller) + _, warnings = self._api.upsert_sweep( + self._sweep_config, controller=controller, obj_id=sweep_obj_id + ) + handle_sweep_config_violations(warnings) + self._controller_prev_step = self._controller.copy() + + def _start_if_not_started(self) -> None: + if self._started: + return + if self._defer_sweep_creation: + raise ControllerError( + "Must specify or create a sweep before running controller." + ) + obj = self._sweep_object_read_from_backend() + if not obj: + return + is_local = self._sweep_config.get("controller", {}).get("type") == "local" + if not is_local: + raise ControllerError( + "Only sweeps with a local controller are currently supported." + ) + self._started = True + # reset controller state, we might want to parse this and decide + # what we can continue and add a version key, but for now we can + # be safe and just reset things on start + self._controller = {} + self._sweep_object_sync_to_backend() + + def _parse_scheduled(self): + scheduled_list = self._scheduler.get("scheduled") or [] + started_ids = [] + stopped_runs = [] + done_runs = [] + for s in scheduled_list: + runid = s.get("runid") + objid = s.get("id") + r = self._sweep_runs_map.get(runid) + if not r: + continue + if r.stopped: + stopped_runs.append(runid) + summary = r.summary_metrics + if r.state == SWEEP_INITIAL_RUN_STATE and not summary: + continue + started_ids.append(objid) + if r.state != "running": + done_runs.append(runid) + return started_ids, stopped_runs, done_runs + + def _step(self) -> None: + self._start_if_not_started() + self._sweep_object_read_from_backend() + + started_ids, stopped_runs, done_runs = self._parse_scheduled() + + # Remove schedule entry from controller dict if already scheduled + schedule_list = self._controller.get("schedule", []) + new_schedule_list = [s for s in schedule_list if s.get("id") not in started_ids] + self._controller["schedule"] = new_schedule_list + + # Remove earlystop entry from controller if already stopped + earlystop_list = self._controller.get("earlystop", []) + new_earlystop_list = [ + r for r in earlystop_list if r not in stopped_runs and r not in done_runs + ] + self._controller["earlystop"] = new_earlystop_list + + # Clear out step logs + self._log_actions = [] + self._log_debug = [] + + def step(self) -> None: + self._step() + suggestion = self.search() + self.schedule(suggestion) + to_stop = self.stopping() + if len(to_stop) > 0: + self.stop_runs(to_stop) + + def done(self) -> bool: + self._start_if_not_started() + state = self._sweep_obj.get("state") + if state in [ + s.upper() + for s in ( + sweeps.RunState.preempting.value, + SWEEP_INITIAL_RUN_STATE.value, + sweeps.RunState.running.value, + ) + ]: + return False + return True + + def _search(self) -> Optional[sweeps.SweepRun]: + search = self._custom_search or sweeps.next_run + next_run = search(self._sweep_config, self._sweep_runs or []) + if next_run is None: + self._done_scheduling = True + return next_run + + def search(self) -> Optional[sweeps.SweepRun]: + self._start_if_not_started() + suggestion = self._search() + return suggestion + + def _stopping(self) -> List[sweeps.SweepRun]: + if "early_terminate" not in self.sweep_config: + return [] + stopper = self._custom_stopping or sweeps.stop_runs + stop_runs = stopper(self._sweep_config, self._sweep_runs or []) + + debug_lines = [ + " ".join([f"{k}={v}" for k, v in run.early_terminate_info.items()]) + for run in stop_runs + if run.early_terminate_info is not None + ] + if debug_lines: + self._log_debug += debug_lines + + return stop_runs + + def stopping(self) -> List[sweeps.SweepRun]: + self._start_if_not_started() + return self._stopping() + + def schedule(self, run: Optional[sweeps.SweepRun]) -> None: + self._start_if_not_started() + + # only schedule one run at a time (for now) + if self._controller and self._controller.get("schedule"): + return + + schedule_id = _id_generator() + + if run is None: + schedule_list = [{"id": schedule_id, "data": {"args": None}}] + else: + param_list = [ + "{}={}".format(k, v.get("value")) for k, v in sorted(run.config.items()) + ] + self._log_actions.append(("schedule", ",".join(param_list))) + + # schedule one run + schedule_list = [{"id": schedule_id, "data": {"args": run.config}}] + + self._controller["schedule"] = schedule_list + self._sweep_object_sync_to_backend() + + def stop_runs(self, runs: List[sweeps.SweepRun]) -> None: + earlystop_list = list({run.name for run in runs}) + self._log_actions.append(("stop", ",".join(earlystop_list))) + self._controller["earlystop"] = earlystop_list + self._sweep_object_sync_to_backend() + + def print_status(self) -> None: + status = _sweep_status(self._sweep_obj, self._sweep_config, self._sweep_runs) + if self._laststatus != status or self._logged: + print(status) # noqa: T201 + self._laststatus = status + self._logged = 0 + + def print_actions(self) -> None: + for action, line in self._log_actions: + self._info(f"{action.capitalize()} ({line})") + self._log_actions = [] + + def print_debug(self) -> None: + for line in self._log_debug: + self._debug(line) + self._log_debug = [] + + def print_space(self) -> None: + self._warn("Method not implemented yet.") + + def print_summary(self) -> None: + self._warn("Method not implemented yet.") + + +def _get_run_counts(runs: List[sweeps.SweepRun]) -> Dict[str, int]: + metrics = {} + categories = [name for name, _ in sweeps.RunState.__members__.items()] + ["unknown"] + for r in runs: + state = r.state + found = "unknown" + for c in categories: + if state == c: + found = c + break + metrics.setdefault(found, 0) + metrics[found] += 1 + return metrics + + +def _get_runs_status(metrics): + categories = [name for name, _ in sweeps.RunState.__members__.items()] + ["unknown"] + mlist = [] + for c in categories: + if not metrics.get(c): + continue + mlist.append(f"{c.capitalize()}: {metrics[c]}") + s = ", ".join(mlist) + return s + + +def _sweep_status( + sweep_obj: dict, + sweep_conf: Union[dict, sweeps.SweepConfig], + sweep_runs: List[sweeps.SweepRun], +) -> str: + sweep = sweep_obj["name"] + _ = sweep_obj["state"] + run_count = len(sweep_runs) + run_type_counts = _get_run_counts(sweep_runs) + stopped = len([r for r in sweep_runs if r.stopped]) + stopping = len([r for r in sweep_runs if r.should_stop]) + stopstr = "" + if stopped or stopping: + stopstr = f"Stopped: {stopped}" + if stopping: + stopstr += f" (Stopping: {stopping})" + runs_status = _get_runs_status(run_type_counts) + method = sweep_conf.get("method", "unknown") + stopping = sweep_conf.get("early_terminate", None) + sweep_options = [] + sweep_options.append(method) + if stopping: + sweep_options.append(stopping.get("type", "unknown")) + sweep_options = ",".join(sweep_options) + sections = [] + sections.append(f"Sweep: {sweep} ({sweep_options})") + if runs_status: + sections.append(f"Runs: {run_count} ({runs_status})") + else: + sections.append(f"Runs: {run_count}") + if stopstr: + sections.append(stopstr) + sections = " | ".join(sections) + return sections diff --git a/lib/python3.12/site-packages/wandb/wandb_run.py b/lib/python3.12/site-packages/wandb/wandb_run.py new file mode 100644 index 0000000000000000000000000000000000000000..05ce62facf71350e042e63e61a5e64043f3c37ef --- /dev/null +++ b/lib/python3.12/site-packages/wandb/wandb_run.py @@ -0,0 +1,9 @@ +"""Compatibility wandb_run module. + +In the future use: + from wandb.sdk.wandb_run import Run +""" + +from wandb.sdk.wandb_run import Run + +__all__ = ["Run"]