diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index 344e2f6..b9e96a3 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -1,20 +1,33 @@ -name: Python Package +# This workflow will upload a Python Package using Twine when a release is created +# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries + +name: Upload Python Package on: release: - types: [ created ] + types: [published] + workflow_dispatch: + +permissions: + contents: read jobs: deploy: runs-on: ubuntu-latest steps: - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v3 + with: + python-version: '3.12' - name: Install dependencies - run: pip install setuptools wheel twine - - name: Build and publish - env: - TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} - TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} run: | - python setup.py sdist bdist_wheel - twine upload dist/* + python -m pip install -U pip + pip install build + - name: Build package + run: python -m build + - name: Publish package + uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29 + with: + user: __token__ + password: ${{ secrets.PYPI_API_TOKEN }} diff --git a/.gitignore b/.gitignore index 5901bfb..2fecae4 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,4 @@ build dist *.egg-info .mypy_cache +/venv diff --git a/LICENSE b/LICENSE index 0e259d4..62102f0 100644 --- a/LICENSE +++ b/LICENSE @@ -1,121 +1,21 @@ -Creative Commons Legal Code - -CC0 1.0 Universal - - CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE - LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN - ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS - INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES - REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS - PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM - THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED - HEREUNDER. - -Statement of Purpose - -The laws of most jurisdictions throughout the world automatically confer -exclusive Copyright and Related Rights (defined below) upon the creator -and subsequent owner(s) (each and all, an "owner") of an original work of -authorship and/or a database (each, a "Work"). - -Certain owners wish to permanently relinquish those rights to a Work for -the purpose of contributing to a commons of creative, cultural and -scientific works ("Commons") that the public can reliably and without fear -of later claims of infringement build upon, modify, incorporate in other -works, reuse and redistribute as freely as possible in any form whatsoever -and for any purposes, including without limitation commercial purposes. -These owners may contribute to the Commons to promote the ideal of a free -culture and the further production of creative, cultural and scientific -works, or to gain reputation or greater distribution for their Work in -part through the use and efforts of others. - -For these and/or other purposes and motivations, and without any -expectation of additional consideration or compensation, the person -associating CC0 with a Work (the "Affirmer"), to the extent that he or she -is an owner of Copyright and Related Rights in the Work, voluntarily -elects to apply CC0 to the Work and publicly distribute the Work under its -terms, with knowledge of his or her Copyright and Related Rights in the -Work and the meaning and intended legal effect of CC0 on those rights. - -1. Copyright and Related Rights. A Work made available under CC0 may be -protected by copyright and related or neighboring rights ("Copyright and -Related Rights"). Copyright and Related Rights include, but are not -limited to, the following: - - i. the right to reproduce, adapt, distribute, perform, display, - communicate, and translate a Work; - ii. moral rights retained by the original author(s) and/or performer(s); -iii. publicity and privacy rights pertaining to a person's image or - likeness depicted in a Work; - iv. rights protecting against unfair competition in regards to a Work, - subject to the limitations in paragraph 4(a), below; - v. rights protecting the extraction, dissemination, use and reuse of data - in a Work; - vi. database rights (such as those arising under Directive 96/9/EC of the - European Parliament and of the Council of 11 March 1996 on the legal - protection of databases, and under any national implementation - thereof, including any amended or successor version of such - directive); and -vii. other similar, equivalent or corresponding rights throughout the - world based on applicable law or treaty, and any national - implementations thereof. - -2. Waiver. To the greatest extent permitted by, but not in contravention -of, applicable law, Affirmer hereby overtly, fully, permanently, -irrevocably and unconditionally waives, abandons, and surrenders all of -Affirmer's Copyright and Related Rights and associated claims and causes -of action, whether now known or unknown (including existing as well as -future claims and causes of action), in the Work (i) in all territories -worldwide, (ii) for the maximum duration provided by applicable law or -treaty (including future time extensions), (iii) in any current or future -medium and for any number of copies, and (iv) for any purpose whatsoever, -including without limitation commercial, advertising or promotional -purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each -member of the public at large and to the detriment of Affirmer's heirs and -successors, fully intending that such Waiver shall not be subject to -revocation, rescission, cancellation, termination, or any other legal or -equitable action to disrupt the quiet enjoyment of the Work by the public -as contemplated by Affirmer's express Statement of Purpose. - -3. Public License Fallback. Should any part of the Waiver for any reason -be judged legally invalid or ineffective under applicable law, then the -Waiver shall be preserved to the maximum extent permitted taking into -account Affirmer's express Statement of Purpose. In addition, to the -extent the Waiver is so judged Affirmer hereby grants to each affected -person a royalty-free, non transferable, non sublicensable, non exclusive, -irrevocable and unconditional license to exercise Affirmer's Copyright and -Related Rights in the Work (i) in all territories worldwide, (ii) for the -maximum duration provided by applicable law or treaty (including future -time extensions), (iii) in any current or future medium and for any number -of copies, and (iv) for any purpose whatsoever, including without -limitation commercial, advertising or promotional purposes (the -"License"). The License shall be deemed effective as of the date CC0 was -applied by Affirmer to the Work. Should any part of the License for any -reason be judged legally invalid or ineffective under applicable law, such -partial invalidity or ineffectiveness shall not invalidate the remainder -of the License, and in such case Affirmer hereby affirms that he or she -will not (i) exercise any of his or her remaining Copyright and Related -Rights in the Work or (ii) assert any associated claims and causes of -action with respect to the Work, in either case contrary to Affirmer's -express Statement of Purpose. - -4. Limitations and Disclaimers. - - a. No trademark or patent rights held by Affirmer are waived, abandoned, - surrendered, licensed or otherwise affected by this document. - b. Affirmer offers the Work as-is and makes no representations or - warranties of any kind concerning the Work, express, implied, - statutory or otherwise, including without limitation warranties of - title, merchantability, fitness for a particular purpose, non - infringement, or the absence of latent or other defects, accuracy, or - the present or absence of errors, whether or not discoverable, all to - the greatest extent permissible under applicable law. - c. Affirmer disclaims responsibility for clearing rights of other persons - that may apply to the Work or any use thereof, including without - limitation any person's Copyright and Related Rights in the Work. - Further, Affirmer disclaims responsibility for obtaining any necessary - consents, permissions or other rights required for any use of the - Work. - d. Affirmer understands and acknowledges that Creative Commons is not a - party to this document and has no duty or obligation with respect to - this CC0 or use of the Work. +MIT License + +Copyright (c) 2024 Ayaka + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/MANIFEST.in b/MANIFEST.in deleted file mode 100644 index 1aba38f..0000000 --- a/MANIFEST.in +++ /dev/null @@ -1 +0,0 @@ -include LICENSE diff --git a/README.md b/README.md index 81c0b78..adc68ef 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,8 @@ jax-smi ## Approach +Update: Since v2.0.0, `jax-smi` calls `tpu-info` directly on TPU platforms. + Save the memory profile to `/dev/shm/memory.prof` in a separate thread every 1 second using [`jax.profiler.save_device_memory_profile()`](https://jax.readthedocs.io/en/latest/_autosummary/jax.profiler.save_device_memory_profile.html). Inspect the memory profile with `go tool pprof -tags /dev/shm/memory.prof`. diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..c221beb --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,43 @@ +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +where = ["src"] + +[project] +name = "jax-smi" +authors = [ + { name="Ayaka", email="ayaka@mail.shn.hk" }, +] +dependencies = [ + "jax>=0.2.16", + "fire", + "tpu-info @ git+https://github.com/google/cloud-accelerator-diagnostics.git@6439522892666f9a50ce3bfad7f9a12d4c7e136b#egg=tpu-info&subdirectory=tpu_info", +] +dynamic = ["version"] +description = "JAX Synergistic Memory Inspector" +readme = "README.md" +requires-python = ">=3.8" +classifiers = [ + "Development Status :: 5 - Production/Stable", + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Topic :: Scientific/Engineering :: Mathematics", + "Topic :: Scientific/Engineering :: Artificial Intelligence", + "Topic :: Software Development :: Libraries :: Python Modules", + "Topic :: System :: Distributed Computing", + "Programming Language :: Python :: 3", +] + +[project.urls] +Homepage = "https://github.com/ayaka14732/jax-smi" +Issues = "https://github.com/ayaka14732/jax-smi/issues" + +[project.scripts] +jax-smi = "jax_smi.cli_tool:main" + +[tool.setuptools.dynamic] +version = {attr = "jax_smi.__version__"} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..15b9ece --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +-e git+https://github.com/google/cloud-accelerator-diagnostics.git@6439522892666f9a50ce3bfad7f9a12d4c7e136b#egg=tpu-info&subdirectory=tpu_info diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index 8183238..0000000 --- a/setup.cfg +++ /dev/null @@ -1,2 +0,0 @@ -[metadata] -license_files = LICENSE diff --git a/setup.py b/setup.py deleted file mode 100644 index b7bfea7..0000000 --- a/setup.py +++ /dev/null @@ -1,48 +0,0 @@ -from setuptools import setup, find_packages -from os import path - -here = path.abspath(path.dirname(__file__)) - -with open(path.join(here, 'README.md')) as f: - long_description = f.read() - -setup( - name='jax-smi', - version='1.0.3', - description='JAX Synergistic Memory Inspector', - long_description=long_description, - long_description_content_type='text/markdown', - url='https://github.com/ayaka14732/jax-smi', - author='Ayaka Mikazuki', - author_email='ayaka@mail.shn.hk', - classifiers=[ - 'Development Status :: 5 - Production/Stable', - 'Intended Audience :: Developers', - 'Intended Audience :: Education', - 'Intended Audience :: Science/Research', - 'License :: CC0 1.0 Universal (CC0 1.0) Public Domain Dedication', - 'Topic :: Scientific/Engineering :: Mathematics', - 'Topic :: Scientific/Engineering :: Artificial Intelligence', - 'Topic :: Software Development :: Libraries :: Python Modules', - 'Topic :: System :: Distributed Computing', - 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3.8', - 'Programming Language :: Python :: 3.9', - 'Programming Language :: Python :: 3.10', - 'Programming Language :: Python :: 3.11', - ], - keywords='jax machine-learning', - packages=find_packages('src'), - package_dir={'': 'src'}, - include_package_data=True, - python_requires='>=3.8, <4', - install_requires=['jax>=0.2.16', 'fire'], - entry_points = { - 'console_scripts': ['jax-smi=jax_smi.cli_tool:main'], - }, - project_urls={ - 'Bug Reports': 'https://github.com/ayaka14732/jax-smi/issues', - 'Source': 'https://github.com/ayaka14732/jax-smi', - }, - zip_safe=False, -) diff --git a/src/jax_smi/__init__.py b/src/jax_smi/__init__.py index 9edaacd..e78fc61 100644 --- a/src/jax_smi/__init__.py +++ b/src/jax_smi/__init__.py @@ -1 +1,3 @@ from .initialise_tracking import initialise_tracking + +__version__ = '2.0.0' diff --git a/src/jax_smi/cli_tool.py b/src/jax_smi/cli_tool.py index 2f9eb37..31d7aa7 100644 --- a/src/jax_smi/cli_tool.py +++ b/src/jax_smi/cli_tool.py @@ -1,9 +1,29 @@ -import curses -import fire -import time +from datetime import datetime import subprocess +import time + +import fire + +from .common_utils import ON_TPU + +def print_info_tpu(interval: float=1.) -> None: + import rich.console + from tpu_info.cli import print_chip_info + + console = rich.console.Console() + + try: + while True: + console.clear() + print(datetime.now().strftime('%Y-%m-%d %H:%M:%S')) + print_chip_info() + time.sleep(interval) + except KeyboardInterrupt: + pass + +def print_info_non_tpu(interval: float=1., dir_prefix: str='/dev/shm') -> None: + import curses -def run(interval: float=1., dir_prefix: str='/dev/shm'): stdscr = curses.initscr() try: while True: @@ -19,5 +39,11 @@ def run(interval: float=1., dir_prefix: str='/dev/shm'): except KeyboardInterrupt: curses.endwin() +def run(interval: float=1., dir_prefix: str='/dev/shm'): + if ON_TPU: + print_info_tpu() + else: + print_info_non_tpu(interval=interval, dir_prefix=dir_prefix) + def main(): fire.Fire(run) diff --git a/src/jax_smi/common_utils.py b/src/jax_smi/common_utils.py new file mode 100644 index 0000000..90a97df --- /dev/null +++ b/src/jax_smi/common_utils.py @@ -0,0 +1,4 @@ +import glob + +# https://github.com/yixiaoer/tpux/blob/1c158a568c6164dff93047f8f13b0bf1e8838725/src/tpux/cli.py#L215 +ON_TPU = len(glob.glob('/dev/accel*')) > 0 diff --git a/src/jax_smi/initialise_tracking.py b/src/jax_smi/initialise_tracking.py index ddd153e..096133d 100644 --- a/src/jax_smi/initialise_tracking.py +++ b/src/jax_smi/initialise_tracking.py @@ -1,4 +1,10 @@ +from .common_utils import ON_TPU + def initialise_tracking(interval: float=1., dir_prefix: str='/dev/shm') -> None: + # on TPU, we use `tpu-info` directly, so nothing needs to be done here + if ON_TPU: + return + import jax import threading