├── .eslintignore ├── .eslintrc.js ├── .github └── workflows │ └── build.yml ├── .gitignore ├── .prettierignore ├── .prettierrc ├── .stylelintrc ├── CHANGELOG.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── Test.ipynb ├── install.json ├── jupyter-config ├── nb-config │ └── jupyterlab_genv.json └── server-config │ └── jupyterlab_genv.json ├── jupyterlab_genv ├── __init__.py ├── __main__.py ├── _version.py ├── genv │ ├── __init__.py │ ├── control.py │ ├── devices.py │ └── envs.py ├── genv_provisioner.py └── handlers.py ├── package-lock.json ├── package.json ├── pyproject.toml ├── resources └── readme │ ├── activate.gif │ ├── attach.gif │ ├── commands.gif │ └── overview.gif ├── setup.py ├── src ├── dialogs.tsx ├── handler.ts └── index.tsx ├── style ├── base.css ├── index.css └── index.js ├── tsconfig.json └── yarn.lock /.eslintignore: -------------------------------------------------------------------------------- 1 | node_modules 2 | dist 3 | coverage 4 | **/*.d.ts 5 | tests 6 | 7 | **/__tests__ 8 | ui-tests 9 | -------------------------------------------------------------------------------- /.eslintrc.js: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | extends: [ 3 | 'eslint:recommended', 4 | 'plugin:@typescript-eslint/eslint-recommended', 5 | 'plugin:@typescript-eslint/recommended', 6 | 'plugin:prettier/recommended' 7 | ], 8 | parser: '@typescript-eslint/parser', 9 | parserOptions: { 10 | project: 'tsconfig.json', 11 | sourceType: 'module' 12 | }, 13 | plugins: ['@typescript-eslint'], 14 | rules: { 15 | '@typescript-eslint/naming-convention': [ 16 | 'error', 17 | { 18 | selector: 'interface', 19 | format: ['PascalCase'], 20 | custom: { 21 | regex: '^I[A-Z]', 22 | match: true 23 | } 24 | } 25 | ], 26 | '@typescript-eslint/no-unused-vars': ['warn', { args: 'none' }], 27 | '@typescript-eslint/no-explicit-any': 'off', 28 | '@typescript-eslint/no-namespace': 'off', 29 | '@typescript-eslint/no-use-before-define': 'off', 30 | '@typescript-eslint/quotes': [ 31 | 'error', 32 | 'single', 33 | { avoidEscape: true, allowTemplateLiterals: false } 34 | ], 35 | curly: ['error', 'all'], 36 | eqeqeq: 'error', 37 | 'prefer-arrow-callback': 'error' 38 | } 39 | }; 40 | -------------------------------------------------------------------------------- /.github/workflows/build.yml: -------------------------------------------------------------------------------- 1 | name: Build 2 | 3 | on: 4 | push: 5 | branches: main 6 | pull_request: 7 | branches: '*' 8 | 9 | jobs: 10 | build: 11 | runs-on: ubuntu-latest 12 | 13 | steps: 14 | - name: Checkout 15 | uses: actions/checkout@v2 16 | 17 | - name: Base Setup 18 | uses: jupyterlab/maintainer-tools/.github/actions/base-setup@v1 19 | 20 | - name: Install dependencies 21 | run: python -m pip install -U jupyterlab~=3.1 22 | 23 | - name: Lint the extension 24 | run: | 25 | set -eux 26 | jlpm 27 | jlpm run lint:check 28 | 29 | - name: Build the extension 30 | run: | 31 | set -eux 32 | python -m pip install .[test] 33 | 34 | jupyter server extension list 35 | jupyter server extension list 2>&1 | grep -ie "jupyterlab_genv.*OK" 36 | 37 | jupyter labextension list 38 | jupyter labextension list 2>&1 | grep -ie "jupyterlab_genv.*OK" 39 | python -m jupyterlab.browser_check 40 | 41 | - name: Package the extension 42 | run: | 43 | set -eux 44 | 45 | pip install build 46 | python -m build 47 | pip uninstall -y "jupyterlab_genv" jupyterlab 48 | 49 | - name: Upload extension packages 50 | uses: actions/upload-artifact@v2 51 | with: 52 | name: extension-artifacts 53 | path: dist/jupyterlab_genv* 54 | if-no-files-found: error 55 | 56 | test_isolated: 57 | needs: build 58 | runs-on: ubuntu-latest 59 | 60 | steps: 61 | - name: Checkout 62 | uses: actions/checkout@v2 63 | - name: Install Python 64 | uses: actions/setup-python@v2 65 | with: 66 | python-version: '3.9' 67 | architecture: 'x64' 68 | - uses: actions/download-artifact@v2 69 | with: 70 | name: extension-artifacts 71 | - name: Install and Test 72 | run: | 73 | set -eux 74 | # Remove NodeJS, twice to take care of system and locally installed node versions. 75 | sudo rm -rf $(which node) 76 | sudo rm -rf $(which node) 77 | 78 | pip install "jupyterlab~=3.1" jupyterlab_genv*.whl 79 | 80 | 81 | jupyter server extension list 82 | jupyter server extension list 2>&1 | grep -ie "jupyterlab_genv.*OK" 83 | 84 | jupyter labextension list 85 | jupyter labextension list 2>&1 | grep -ie "jupyterlab_genv.*OK" 86 | python -m jupyterlab.browser_check --no-chrome-test 87 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.bundle.* 2 | lib/ 3 | node_modules/ 4 | .eslintcache 5 | .stylelintcache 6 | *.egg-info/ 7 | .ipynb_checkpoints 8 | *.tsbuildinfo 9 | jupyterlab_genv/labextension 10 | 11 | # Integration tests 12 | ui-tests/test-results/ 13 | ui-tests/playwright-report/ 14 | 15 | # Created by https://www.gitignore.io/api/python 16 | # Edit at https://www.gitignore.io/?templates=python 17 | 18 | ### Python ### 19 | # Byte-compiled / optimized / DLL files 20 | __pycache__/ 21 | *.py[cod] 22 | *$py.class 23 | 24 | # C extensions 25 | *.so 26 | 27 | # Distribution / packaging 28 | .Python 29 | build/ 30 | develop-eggs/ 31 | dist/ 32 | downloads/ 33 | eggs/ 34 | .eggs/ 35 | lib/ 36 | lib64/ 37 | parts/ 38 | sdist/ 39 | var/ 40 | wheels/ 41 | pip-wheel-metadata/ 42 | share/python-wheels/ 43 | .installed.cfg 44 | *.egg 45 | MANIFEST 46 | 47 | # PyInstaller 48 | # Usually these files are written by a python script from a template 49 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 50 | *.manifest 51 | *.spec 52 | 53 | # Installer logs 54 | pip-log.txt 55 | pip-delete-this-directory.txt 56 | 57 | # Unit test / coverage reports 58 | htmlcov/ 59 | .tox/ 60 | .nox/ 61 | .coverage 62 | .coverage.* 63 | .cache 64 | nosetests.xml 65 | coverage.xml 66 | *.cover 67 | .hypothesis/ 68 | .pytest_cache/ 69 | 70 | # Translations 71 | *.mo 72 | *.pot 73 | 74 | # Scrapy stuff: 75 | .scrapy 76 | 77 | # Sphinx documentation 78 | docs/_build/ 79 | 80 | # PyBuilder 81 | target/ 82 | 83 | # pyenv 84 | .python-version 85 | 86 | # celery beat schedule file 87 | celerybeat-schedule 88 | 89 | # SageMath parsed files 90 | *.sage.py 91 | 92 | # Spyder project settings 93 | .spyderproject 94 | .spyproject 95 | 96 | # Rope project settings 97 | .ropeproject 98 | 99 | # Mr Developer 100 | .mr.developer.cfg 101 | .project 102 | .pydevproject 103 | 104 | # mkdocs documentation 105 | /site 106 | 107 | # mypy 108 | .mypy_cache/ 109 | .dmypy.json 110 | dmypy.json 111 | 112 | # Pyre type checker 113 | .pyre/ 114 | 115 | # End of https://www.gitignore.io/api/python 116 | 117 | # OSX files 118 | .DS_Store 119 | -------------------------------------------------------------------------------- /.prettierignore: -------------------------------------------------------------------------------- 1 | node_modules 2 | **/node_modules 3 | **/lib 4 | **/package.json 5 | jupyterlab_genv 6 | -------------------------------------------------------------------------------- /.prettierrc: -------------------------------------------------------------------------------- 1 | { 2 | "singleQuote": true, 3 | "trailingComma": "none", 4 | "arrowParens": "avoid", 5 | "endOfLine": "auto" 6 | } 7 | -------------------------------------------------------------------------------- /.stylelintrc: -------------------------------------------------------------------------------- 1 | { 2 | "extends": [ 3 | "stylelint-config-recommended", 4 | "stylelint-config-standard", 5 | "stylelint-prettier/recommended" 6 | ], 7 | "rules": { 8 | "property-no-vendor-prefix": null, 9 | "selector-no-vendor-prefix": null, 10 | "value-no-vendor-prefix": null 11 | } 12 | } 13 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Changelog 2 | 3 | All notable changes to this project will be documented in this file. 4 | 5 | The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 6 | and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). 7 | 8 | ## [0.4.0] - 2024-02-19 9 | 10 | ### Added 11 | 12 | - Supporting `jupyter-server` versions 2.x. 13 | 14 | ## [0.3.0] - 2023-06-12 15 | 16 | ### Changed 17 | 18 | - Using Genv 1.0.0 APIs. 19 | 20 | ## [0.2.0] - 2023-04-30 21 | 22 | ### Changed 23 | 24 | - Using renamed command `genv-devices find`. 25 | 26 | ## [0.1.2] - 2022-12-11 27 | 28 | ### Changed 29 | 30 | - Initializing spawned terminals using `eval "$(genv init -)"`. 31 | 32 | ## [0.1.1] - 2022-09-24 33 | 34 | ### Added 35 | 36 | - Respect environment variable `GENV_ROOT` to support [Conda](https://docs.conda.io/en/latest/) installation of `genv`. 37 | 38 | ### Changed 39 | 40 | - Start using changelog. 41 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | BSD 3-Clause License 2 | 3 | Copyright (c) 2022, Raz Rotenberg 4 | All rights reserved. 5 | 6 | Redistribution and use in source and binary forms, with or without 7 | modification, are permitted provided that the following conditions are met: 8 | 9 | 1. Redistributions of source code must retain the above copyright notice, this 10 | list of conditions and the following disclaimer. 11 | 12 | 2. Redistributions in binary form must reproduce the above copyright notice, 13 | this list of conditions and the following disclaimer in the documentation 14 | and/or other materials provided with the distribution. 15 | 16 | 3. Neither the name of the copyright holder nor the names of its 17 | contributors may be used to endorse or promote products derived from 18 | this software without specific prior written permission. 19 | 20 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 21 | AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 22 | IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE 23 | DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE 24 | FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL 25 | DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR 26 | SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER 27 | CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, 28 | OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 29 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 30 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include LICENSE 2 | include *.md 3 | include pyproject.toml 4 | recursive-include jupyter-config *.json 5 | 6 | include package.json 7 | include install.json 8 | include ts*.json 9 | include yarn.lock 10 | 11 | graft jupyterlab_genv/labextension 12 | 13 | # Javascript files 14 | graft src 15 | graft style 16 | prune **/node_modules 17 | prune lib 18 | prune binder 19 | 20 | # Patterns to exclude from any directory 21 | global-exclude *~ 22 | global-exclude *.pyc 23 | global-exclude *.pyo 24 | global-exclude .git 25 | global-exclude .ipynb_checkpoints 26 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # GPU Environment Management for JupyterLab [![Join the community at https://join.slack.com/t/genvcommunity/shared_invite/zt-1i70tphdc-DmFgK5yr3HFI8Txx1yFXBw](https://img.shields.io/badge/Slack-genv-ff007f?logo=slack)](https://join.slack.com/t/genvcommunity/shared_invite/zt-1i70tphdc-DmFgK5yr3HFI8Txx1yFXBw) [![Join the chat at https://gitter.im/run-ai-genv/community](https://badges.gitter.im/run-ai-genv/community.svg)](https://gitter.im/run-ai-genv/community?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge) 2 | 3 | [![Github Actions Status](https://github.com/run-ai/jupyterlab_genv/workflows/Build/badge.svg)](https://github.com/run-ai/jupyterlab_genv/actions/workflows/build.yml) 4 | 5 | A JupyterLab extension for managing GPU environments using [genv](https://github.com/run-ai/genv). 6 | 7 | The [_genv_](https://github.com/run-ai/genv) extension lets you interactively control, configure and monitor the GPU resources that your Jupyter Notebooks are using. 8 | 9 | ![Overview](/resources/readme/overview.gif) 10 | 11 | ## 🏃🏻 Be an early runner in the genv community! 12 | 13 | [](https://join.slack.com/t/genvcommunity/shared_invite/zt-1i70tphdc-DmFgK5yr3HFI8Txx1yFXBw) 14 | 15 | Join our Slack channel with the creators of _genv_ and start building your models faster! 16 | 17 | - Installation and setup support as well as best practice tips and tricks directly for your use-case 18 | - Discuss possible features 19 | - Monthly coffee breaks to get to know the rest of the community 20 | 21 | Looking forward to seeing you as a part of the community! 22 | 23 | ## Table of Contents 24 | 25 | - [Getting Started](#getting-started) 26 | - [Installation](#installation) 27 | - [Conda](#conda) 28 | - [Pip](#pip) 29 | - [Install _genv_ Kernels](#install-genv-kernels) 30 | - [Usage](#usage) 31 | - [Activate Your Environment](#activate-your-environment) 32 | - [Attach GPUs to Your Environment](#attach-gpus-to-your-environment) 33 | - [See Devices and Environments](#see-devices-and-environments) 34 | - [Development](#development) 35 | - [Publish](#publish) 36 | - [PyPI](#pypi) 37 | - [Conda](#conda-1) 38 | 39 | ## Getting Started 40 | 41 | Read the _genv_ [reference](https://github.com/run-ai/genv#usage) to get started. 42 | 43 | ## Installation 44 | 45 | ### Requirements 46 | 47 | JupyterLab >= 3.0 48 | 49 | ### Conda 50 | 51 | If you are using [Conda](https://docs.conda.io/en/latest/), it is best to install the `jupyterlab_genv` [package](https://anaconda.org/conda-forge/jupyterlab_genv) from the channel [conda-forge](https://conda-forge.org/): 52 | 53 | ```bash 54 | conda install -c conda-forge jupyterlab_genv 55 | ``` 56 | 57 | ### Pip 58 | 59 | Alternatively, you can install `jupyterlab_genv` from [PyPI](https://pypi.org/project/jupyterlab-genv/) using `pip`: 60 | 61 | ```bash 62 | pip install jupyterlab_genv 63 | ``` 64 | 65 | ### Install _genv_ Kernels 66 | 67 | After installing `jupyterlab_genv`, you will need to install _genv_ Jupyter kernels using: 68 | 69 | ```bash 70 | python -m jupyterlab_genv install 71 | ``` 72 | 73 | ## Usage 74 | 75 | ### Activate Your Environment 76 | 77 | To activate your environment, you will have to select a _genv_ [kernel](#install-genv-kernels). 78 | 79 | Then, click the `GPUs` button on the Jupyter Notebook toolbar. 80 | A dialog should pop up where you can choose either to create a new environment for your Jupyter Notebook, or to use an existing one. 81 | 82 | Then, you can open a terminal activated in your environment. 83 | From there you will be able to configure the environment and attach devices. 84 | 85 | ![Activate](/resources/readme/activate.gif) 86 | 87 | ### Attach GPUs to Your Environment 88 | 89 | Configuring the environment and attaching devices is done from the _genv_ terminal. 90 | 91 | Make sure to restart your kernel after running the command in the terminal for it to take effect. 92 | 93 | ![Attach](/resources/readme/attach.gif) 94 | 95 | ### See Devices and Environments 96 | 97 | You can open the devices and environments widgets to see information. 98 | 99 | Open the command palette (`Command/Ctrl Shift C`) and type `GPUs`. 100 | 101 | ![Commands](/resources/readme/commands.gif) 102 | 103 | ## Development 104 | 105 | ### Setup 106 | 107 | You will need to create a virtual environment once using the command: 108 | 109 | ```bash 110 | conda create -n jupyterlab_genv --override-channels --strict-channel-priority -c conda-forge -c nodefaults jupyterlab=3 cookiecutter nodejs jupyter-packaging git 111 | ``` 112 | 113 | Then, activate the virtual environment when you want to work on the project: 114 | 115 | ```bash 116 | conda activate jupyterlab_genv 117 | ``` 118 | 119 | ### Install 120 | 121 | Use the following commands to install the Python package and enable it in JupyterLab: 122 | 123 | ```bash 124 | # Install package in development mode 125 | pip install -e . 126 | # Link your development version of the extension with JupyterLab 127 | jupyter labextension develop . --overwrite 128 | # Server extension must be manually installed in develop mode 129 | jupyter server extension enable jupyterlab_genv 130 | ``` 131 | 132 | If you make any changes you will need to rebuild the extension Typescript source using: 133 | 134 | ``` 135 | jlpm build 136 | ``` 137 | 138 | Alternatively, you can watch the source directory using: 139 | 140 | ``` 141 | jlpm watch 142 | ``` 143 | 144 | With the `jlpm watch` command running, every saved change will immediately be built locally and available in your running JupyterLab. Refresh JupyterLab to load the change in your browser (you may need to wait several seconds for the extension to be rebuilt). 145 | 146 | ### Run 147 | 148 | Run JupyterLab using the command: 149 | 150 | ```bash 151 | jupyter lab 152 | ``` 153 | 154 | > Running `SHELL=bash jupyter lab --no-browser` is even better 155 | 156 | ### Uninstall 157 | 158 | ```bash 159 | # Server extension must be manually disabled in develop mode 160 | jupyter server extension disable jupyterlab_genv 161 | pip uninstall jupyterlab_genv 162 | ``` 163 | 164 | In development mode, you will also need to remove the symlink created by `jupyter labextension develop` command. 165 | To find its location, you can run `jupyter labextension list` to figure out where the `labextensions` folder is located. 166 | Then you can remove the symlink named `jupyterlab_genv` within that folder. 167 | 168 | ### Reference 169 | 170 | #### List all kernel provisioners 171 | 172 | ```bash 173 | jupyter kernelspec provisioners 174 | ``` 175 | 176 | #### Install a kernel provisioner 177 | 178 | To add a kernel provisioner to a kernel spec, edit its `kernel.json` file. 179 | For example, to install a kernel provisioner for the `python3` kernel spec, run: 180 | 181 | ```bash 182 | vim $CONDA_PREFIX/share/jupyter/kernels/python3/kernel.json 183 | ``` 184 | 185 | And add: 186 | 187 | ``` 188 | "metadata": { 189 | "kernel_provisioner": { 190 | "provisioner_name": "genv-provisioner" 191 | } 192 | } 193 | ``` 194 | 195 | #### List all available kernel specs 196 | 197 | ```bash 198 | ls -la $CONDA_PREFIX/share/jupyter/kernels/ 199 | ``` 200 | 201 | #### List all running kernels 202 | 203 | ```bash 204 | ls -la $(jupyter --runtime-dir)/kernel-*.json 205 | ``` 206 | 207 | #### List Jupyter server extensions 208 | 209 | ```bash 210 | jupyter server extension list 211 | ``` 212 | 213 | #### List JupyterLab extensions 214 | 215 | ```bash 216 | jupyter labextension list 217 | ``` 218 | 219 | ## Publish 220 | 221 | The `jupyterlab_genv` package is manually published to both [PyPI](https://pypi.org/project/jupyterlab-genv/) and [conda-forge](https://anaconda.org/conda-forge/jupyterlab_genv). 222 | 223 | We do not publish the frontend part as an npm package because the Python package is a prebuilt server extension, and the frontend part alone is useless. 224 | 225 | Also make sure to update the [changelog](./CHANGELOG.md) ([here's](https://keepachangelog.com/en/1.0.0/#how) how) and lint the project by running `npm run lint`. 226 | 227 | ### Bump Version 228 | 229 | The [cookiecutter template](https://github.com/jupyterlab/extension-cookiecutter-ts) uses `tbump` for bumping the version. 230 | However, for some reason this does not work at the moment, and we bump the version manually. 231 | 232 | Search for the current version in the project files and replace the relevant instances. 233 | Here is a list of files that you should update: 234 | 235 | - [package.json](package.json#L3) 236 | - [package-lock.json](package-lock.json#L3) 237 | - [pyproject.toml](pyproject.toml#L7) (also [here](pyproject.toml#L84) for future `tbump` support) 238 | - [jupyterlab_genv/\_version.py](jupyterlab_genv/_version.py#L6) 239 | 240 | After pushing these changes, create a release on [GitHub](https://github.com/run-ai/jupyterlab_genv/releases). 241 | 242 | ### PyPI 243 | 244 | #### Prerequisites 245 | 246 | ```bash 247 | pip install build twine tbump 248 | ``` 249 | 250 | #### Create a Python Package 251 | 252 | Create a Python source package (`.tar.gz`) and the binary package (`.whl`) in the `dist/` directory using: 253 | 254 | ```bash 255 | python -m build 256 | ``` 257 | 258 | > `python setup.py sdist bdist_wheel` is deprecated and will not work for this package. 259 | 260 | Then, upload the package to [PyPI](https://pypi.org/project/jupyterlab-genv/) using: 261 | 262 | ```bash 263 | twine upload dist/* 264 | ``` 265 | 266 | > We upload to PyPI with the organizational user [runai](https://pypi.org/user/runai/) 267 | 268 | ### Conda 269 | 270 | The Conda package is managed using its [feedstock](https://github.com/conda-forge/jupyterlab_genv-feedstock). 271 | 272 | After publishing to [PyPI](#pypi), update the [version](https://github.com/conda-forge/jupyterlab_genv-feedstock/blob/main/recipe/meta.yaml#L2) and [sha256](https://github.com/conda-forge/jupyterlab_genv-feedstock/blob/main/recipe/meta.yaml#L10) fields in the recipe `meta.yaml` file. 273 | 274 | A few minutes after pushing these changes, you should be able to see that the Conda [package](https://anaconda.org/conda-forge/jupyterlab_genv) version was updated. 275 | 276 | > You can get the SHA256 hash from [PyPI](https://pypi.org/project/jupyterlab-genv/#files) 277 | -------------------------------------------------------------------------------- /Test.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "id": "f65bbbc5-c713-461d-82e3-dd8129221c02", 7 | "metadata": {}, 8 | "outputs": [], 9 | "source": [ 10 | "import os\n", 11 | "\n", 12 | "if 'CUDA_VISIBLE_DEVICES' not in os.environ:\n", 13 | " print('Not running in a GPU environment')\n", 14 | "else:\n", 15 | " print('Running in a GPU environment')\n", 16 | " indices = [int(index) for index in os.environ['CUDA_VISIBLE_DEVICES'].split(',') if len(index)]\n", 17 | " if len(indices) == 0:\n", 18 | " print('Detached from devices')\n", 19 | " else:\n", 20 | " print(f'Using devices {indices}')" 21 | ] 22 | } 23 | ], 24 | "metadata": { 25 | "kernelspec": { 26 | "display_name": "Python 3 (ipykernel)", 27 | "language": "python", 28 | "name": "python3" 29 | }, 30 | "language_info": { 31 | "codemirror_mode": { 32 | "name": "ipython", 33 | "version": 3 34 | }, 35 | "file_extension": ".py", 36 | "mimetype": "text/x-python", 37 | "name": "python", 38 | "nbconvert_exporter": "python", 39 | "pygments_lexer": "ipython3", 40 | "version": "3.10.6" 41 | } 42 | }, 43 | "nbformat": 4, 44 | "nbformat_minor": 5 45 | } 46 | -------------------------------------------------------------------------------- /install.json: -------------------------------------------------------------------------------- 1 | { 2 | "packageManager": "python", 3 | "packageName": "jupyterlab_genv", 4 | "uninstallInstructions": "Use your Python package manager (pip, conda, etc.) to uninstall the package jupyterlab_genv" 5 | } 6 | -------------------------------------------------------------------------------- /jupyter-config/nb-config/jupyterlab_genv.json: -------------------------------------------------------------------------------- 1 | { 2 | "NotebookApp": { 3 | "nbserver_extensions": { 4 | "jupyterlab_genv": true 5 | } 6 | } 7 | } 8 | -------------------------------------------------------------------------------- /jupyter-config/server-config/jupyterlab_genv.json: -------------------------------------------------------------------------------- 1 | { 2 | "ServerApp": { 3 | "jpserver_extensions": { 4 | "jupyterlab_genv": true 5 | } 6 | } 7 | } 8 | -------------------------------------------------------------------------------- /jupyterlab_genv/__init__.py: -------------------------------------------------------------------------------- 1 | import json 2 | from pathlib import Path 3 | 4 | from ._version import __version__ 5 | from .handlers import setup_handlers 6 | 7 | 8 | 9 | HERE = Path(__file__).parent.resolve() 10 | 11 | 12 | with (HERE / "labextension" / "package.json").open() as fid: 13 | data = json.load(fid) 14 | 15 | 16 | def _jupyter_labextension_paths(): 17 | return [{ 18 | "src": "labextension", 19 | "dest": data["name"] 20 | }] 21 | 22 | 23 | 24 | def _jupyter_server_extension_points(): 25 | return [{ 26 | "module": "jupyterlab_genv" 27 | }] 28 | 29 | 30 | def _load_jupyter_server_extension(server_app): 31 | """Registers the API handler to receive HTTP requests from the frontend extension. 32 | 33 | Parameters 34 | ---------- 35 | server_app: jupyterlab.labapp.LabApp 36 | JupyterLab application instance 37 | """ 38 | setup_handlers(server_app.web_app) 39 | server_app.log.info("Registered {name} server extension".format(**data)) 40 | 41 | 42 | # For backward compatibility with notebook server - useful for Binder/JupyterHub 43 | load_jupyter_server_extension = _load_jupyter_server_extension 44 | 45 | -------------------------------------------------------------------------------- /jupyterlab_genv/__main__.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import json 3 | import os 4 | import shutil 5 | 6 | import jupyter_client.kernelspec 7 | 8 | def do_install() -> None: 9 | kernel_specs = jupyter_client.kernelspec.find_kernel_specs() 10 | 11 | for kernel_spec_name, kernel_spec_dir in kernel_specs.items(): 12 | if kernel_spec_name.endswith('-genv') or f'{kernel_spec_name}-genv' in kernel_specs: 13 | continue 14 | 15 | dst = f'{kernel_spec_dir}-genv' 16 | 17 | print(f'Installing genv wrapper for kernel spec "{kernel_spec_name}" at {dst}') 18 | 19 | shutil.copytree(kernel_spec_dir, dst) 20 | 21 | kernel_json_path = f'{os.path.join(dst, "kernel.json")}' 22 | 23 | with open(kernel_json_path, 'r') as f: 24 | kernel_json = json.load(f) 25 | 26 | kernel_json['display_name'] = f"{kernel_json['display_name']} (genv)" 27 | 28 | if 'metadata' not in kernel_json: 29 | kernel_json['metadata'] = {} 30 | 31 | kernel_json['metadata']['kernel_provisioner'] = { 'provisioner_name': 'genv-provisioner' } 32 | 33 | with open(kernel_json_path, 'w') as f: 34 | json.dump(kernel_json, f, indent=1) 35 | 36 | if __name__ == "__main__": 37 | parser = argparse.ArgumentParser(description=f'JupyterLab genv extension') 38 | subparsers = parser.add_subparsers(dest='command', required=True) 39 | subparsers.add_parser('install', help='Install genv kernel specs') 40 | args = parser.parse_args() 41 | 42 | if args.command == 'install': 43 | do_install() 44 | -------------------------------------------------------------------------------- /jupyterlab_genv/_version.py: -------------------------------------------------------------------------------- 1 | import json 2 | from pathlib import Path 3 | 4 | __all__ = ["__version__"] 5 | 6 | __version__ = "0.4.0" 7 | -------------------------------------------------------------------------------- /jupyterlab_genv/genv/__init__.py: -------------------------------------------------------------------------------- 1 | from . import devices 2 | from . import envs 3 | -------------------------------------------------------------------------------- /jupyterlab_genv/genv/control.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | 3 | async def exec(command: str) -> str: 4 | proc = await asyncio.create_subprocess_shell(f'genv {command}', 5 | stdout=asyncio.subprocess.PIPE, 6 | stderr=asyncio.subprocess.PIPE) 7 | 8 | stdout, stderr = await proc.communicate() 9 | 10 | assert proc.returncode == 0 11 | 12 | return stdout.decode('utf-8').strip() 13 | -------------------------------------------------------------------------------- /jupyterlab_genv/genv/devices.py: -------------------------------------------------------------------------------- 1 | from typing import Dict, List 2 | 3 | from . import control 4 | 5 | async def _exec(command: str) -> str: 6 | return await control.exec(f'devices {command}') 7 | 8 | async def ps() -> Dict: 9 | stdout = await _exec("ps --format csv --no-header --timestamp") 10 | lines = [line for line in stdout.splitlines() if len(line)] 11 | 12 | infos = [] 13 | 14 | for line in lines: 15 | id, eid, env, attached = line.split(',') 16 | id = int(id) 17 | 18 | infos.append({ 19 | "eid": eid, 20 | "env": env, 21 | }) 22 | 23 | return infos 24 | 25 | async def find(eid: str) -> List[int]: 26 | return [int(index) for index in (await _exec(f'find --eid {eid}')).split(',') if len(index) > 0] 27 | -------------------------------------------------------------------------------- /jupyterlab_genv/genv/envs.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Dict, List 3 | 4 | from . import control 5 | 6 | async def _exec(command: str) -> str: 7 | return await control.exec(f'envs {command}') 8 | 9 | async def activate(eid: str, kernel_id: str) -> None: 10 | await _exec(f'activate --eid {eid} --uid {os.getuid()} --kernel-id {kernel_id}') 11 | 12 | async def find(kernel_id: str) -> str: 13 | return await _exec(f'find --kernel-id {kernel_id}') 14 | 15 | async def ps() -> List[Dict]: 16 | stdout = await _exec("ps --format csv --no-header --timestamp") 17 | lines = [line for line in stdout.splitlines() if len(line)] 18 | 19 | infos = [] 20 | 21 | for line in lines: 22 | eid, user, name, created, pids = line.split(',') 23 | 24 | infos.append({ 25 | "eid": eid, 26 | "user": user, 27 | "name": name, 28 | "pids": [int(pid) for pid in pids.split(' ') if len(pid)], 29 | }) 30 | 31 | return infos 32 | -------------------------------------------------------------------------------- /jupyterlab_genv/genv_provisioner.py: -------------------------------------------------------------------------------- 1 | import jupyter_client 2 | import os 3 | from typing import Any, Dict 4 | 5 | from . import genv 6 | 7 | class GenvProvisioner(jupyter_client.LocalProvisioner): 8 | async def pre_launch(self, **kwargs: Any) -> Dict[str, Any]: 9 | eid = await genv.envs.find(self.kernel_id) or self.kernel_id 10 | indices = await genv.devices.find(eid) 11 | 12 | env = kwargs.pop('env', os.environ).copy() 13 | env.update({ 'CUDA_VISIBLE_DEVICES': ','.join(str(index) for index in indices) }) 14 | # TODO(raz): add 'shims' to PATH so that nvidia-smi shim would run 15 | kwargs['env'] = env 16 | 17 | return await super().pre_launch(**kwargs) 18 | -------------------------------------------------------------------------------- /jupyterlab_genv/handlers.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | from jupyter_server.base.handlers import APIHandler 4 | from jupyter_server.utils import url_path_join 5 | import tornado 6 | 7 | from . import genv 8 | 9 | class DevicesHandler(APIHandler): 10 | @tornado.web.authenticated 11 | async def get(self): 12 | self.finish(json.dumps(await genv.devices.ps())) 13 | 14 | class EnvsHandler(APIHandler): 15 | @tornado.web.authenticated 16 | async def get(self): 17 | self.finish(json.dumps(await genv.envs.ps())) 18 | 19 | class ActivateHandler(APIHandler): 20 | @tornado.web.authenticated 21 | async def post(self): 22 | body = self.get_json_body() 23 | await genv.envs.activate(body['eid'], body['kernel_id']) 24 | self.finish() 25 | 26 | class FindHandler(APIHandler): 27 | @tornado.web.authenticated 28 | async def get(self): 29 | kernel_id = self.get_query_argument('kernel_id') 30 | self.finish(json.dumps(await genv.envs.find(kernel_id))) 31 | 32 | def setup_handlers(web_app): 33 | host_pattern = ".*$" 34 | 35 | base_url = web_app.settings["base_url"] 36 | url = lambda uri: url_path_join(base_url, "jupyterlab-genv", uri) 37 | 38 | handlers = [ 39 | (url("devices"), DevicesHandler), 40 | (url("envs"), EnvsHandler), 41 | (url("activate"), ActivateHandler), 42 | (url("find"), FindHandler), 43 | ] 44 | 45 | web_app.add_handlers(host_pattern, handlers) 46 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "jupyterlab_genv", 3 | "version": "0.4.0", 4 | "description": "A JupyterLab extension.", 5 | "keywords": [ 6 | "jupyter", 7 | "jupyterlab", 8 | "jupyterlab-extension" 9 | ], 10 | "homepage": "https://github.com/run-ai/jupyterlab_genv", 11 | "bugs": { 12 | "url": "https://github.com/run-ai/jupyterlab_genv/issues" 13 | }, 14 | "license": "BSD-3-Clause", 15 | "author": { 16 | "name": "Raz Rotenberg", 17 | "email": "raz.rotenberg@gmail.com" 18 | }, 19 | "files": [ 20 | "lib/**/*.{d.ts,eot,gif,html,jpg,js,js.map,json,png,svg,woff2,ttf}", 21 | "style/**/*.{css,js,eot,gif,html,jpg,json,png,svg,woff2,ttf}" 22 | ], 23 | "main": "lib/index.js", 24 | "types": "lib/index.d.ts", 25 | "style": "style/index.css", 26 | "repository": { 27 | "type": "git", 28 | "url": "https://github.com/run-ai/jupyterlab_genv.git" 29 | }, 30 | "scripts": { 31 | "build": "jlpm build:lib && jlpm build:labextension:dev", 32 | "build:prod": "jlpm clean && jlpm build:lib && jlpm build:labextension", 33 | "build:labextension": "jupyter labextension build .", 34 | "build:labextension:dev": "jupyter labextension build --development True .", 35 | "build:lib": "tsc", 36 | "clean": "jlpm clean:lib", 37 | "clean:lib": "rimraf lib tsconfig.tsbuildinfo", 38 | "clean:lintcache": "rimraf .eslintcache .stylelintcache", 39 | "clean:labextension": "rimraf jupyterlab_genv/labextension", 40 | "clean:all": "jlpm clean:lib && jlpm clean:labextension && jlpm clean:lintcache", 41 | "eslint": "jlpm eslint:check --fix", 42 | "eslint:check": "eslint . --cache --ext .ts,.tsx", 43 | "install:extension": "jlpm build", 44 | "lint": "jlpm stylelint && jlpm prettier && jlpm eslint", 45 | "lint:check": "jlpm stylelint:check && jlpm prettier:check && jlpm eslint:check", 46 | "prettier": "jlpm prettier:base --write --list-different", 47 | "prettier:base": "prettier \"**/*{.ts,.tsx,.js,.jsx,.css,.json,.md}\"", 48 | "prettier:check": "jlpm prettier:base --check", 49 | "stylelint": "jlpm stylelint:check --fix", 50 | "stylelint:check": "stylelint --cache \"style/**/*.css\"", 51 | "watch": "run-p watch:src watch:labextension", 52 | "watch:src": "tsc -w", 53 | "watch:labextension": "jupyter labextension watch ." 54 | }, 55 | "dependencies": { 56 | "@jupyterlab/application": "^3.4.7", 57 | "@jupyterlab/apputils": "^3.4.7", 58 | "@jupyterlab/coreutils": "^5.1.0", 59 | "@jupyterlab/notebook": "^3.4.7", 60 | "@jupyterlab/services": "^6.1.0", 61 | "@jupyterlab/terminal": "^3.4.7", 62 | "@lumino/messaging": "^1.10.2", 63 | "@lumino/widgets": "^1.34.0" 64 | }, 65 | "devDependencies": { 66 | "@jupyterlab/builder": "^3.6.5", 67 | "@typescript-eslint/eslint-plugin": "^4.8.1", 68 | "@typescript-eslint/parser": "^4.8.1", 69 | "eslint": "^7.14.0", 70 | "eslint-config-prettier": "^6.15.0", 71 | "eslint-plugin-prettier": "^3.1.4", 72 | "mkdirp": "^1.0.3", 73 | "npm-run-all": "^4.1.5", 74 | "prettier": "^2.1.1", 75 | "rimraf": "^3.0.2", 76 | "stylelint": "^15.10.1", 77 | "stylelint-config-prettier": "^9.0.3", 78 | "stylelint-config-recommended": "^6.0.0", 79 | "stylelint-config-standard": "~24.0.0", 80 | "stylelint-prettier": "^2.0.0", 81 | "typescript": "~4.1.3" 82 | }, 83 | "sideEffects": [ 84 | "style/*.css", 85 | "style/index.js" 86 | ], 87 | "styleModule": "style/index.js", 88 | "publishConfig": { 89 | "access": "public" 90 | }, 91 | "jupyterlab": { 92 | "discovery": { 93 | "server": { 94 | "managers": [ 95 | "pip" 96 | ], 97 | "base": { 98 | "name": "jupyterlab_genv" 99 | } 100 | } 101 | }, 102 | "extension": true, 103 | "outputDir": "jupyterlab_genv/labextension" 104 | }, 105 | "jupyter-releaser": { 106 | "hooks": { 107 | "before-build-npm": [ 108 | "python -m pip install jupyterlab~=3.1", 109 | "jlpm" 110 | ], 111 | "before-build-python": [ 112 | "jlpm clean:all" 113 | ] 114 | } 115 | } 116 | } 117 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["hatchling>=1.3.1", "jupyterlab~=3.1"] 3 | build-backend = "hatchling.build" 4 | 5 | [project] 6 | name = "jupyterlab_genv" 7 | version = "0.4.0" 8 | description = "A JupyterLab extension for managing GPU environments using genv." 9 | readme = "README.md" 10 | license = { file = "LICENSE" } 11 | requires-python = ">=3.7" 12 | authors = [ 13 | { name = "Raz Rotenberg", email = "raz.rotenberg@gmail.com" }, 14 | ] 15 | keywords = ["Jupyter", "JupyterLab", "JupyterLab3"] 16 | classifiers = [ 17 | "Framework :: Jupyter", 18 | "Framework :: Jupyter :: JupyterLab", 19 | "Framework :: Jupyter :: JupyterLab :: 3", 20 | "Framework :: Jupyter :: JupyterLab :: Extensions", 21 | "Framework :: Jupyter :: JupyterLab :: Extensions :: Prebuilt", 22 | "License :: OSI Approved :: BSD License", 23 | "Programming Language :: Python", 24 | "Programming Language :: Python :: 3", 25 | "Programming Language :: Python :: 3.7", 26 | "Programming Language :: Python :: 3.8", 27 | "Programming Language :: Python :: 3.9", 28 | "Programming Language :: Python :: 3.10", 29 | ] 30 | dependencies = [ 31 | "genv", 32 | "jupyter_server>=1.6,<3" 33 | ] 34 | 35 | [project.optional-dependencies] 36 | test = [ 37 | ] 38 | 39 | [project.urls] 40 | Homepage = "https://github.com/run-ai/jupyterlab_genv" 41 | 42 | [tool.hatch.build] 43 | artifacts = ["jupyterlab_genv/labextension"] 44 | 45 | [tool.hatch.build.targets.wheel.shared-data] 46 | "jupyterlab_genv/labextension" = "share/jupyter/labextensions/jupyterlab_genv" 47 | "install.json" = "share/jupyter/labextensions/jupyterlab_genv/install.json" 48 | "jupyter-config/server-config" = "etc/jupyter/jupyter_server_config.d" 49 | "jupyter-config/nb-config" = "etc/jupyter/jupyter_notebook_config.d" 50 | 51 | [tool.hatch.build.targets.sdist] 52 | exclude = [".github"] 53 | 54 | [tool.hatch.build.hooks.jupyter-builder] 55 | dependencies = ["hatch-jupyter-builder>=0.5"] 56 | build-function = "hatch_jupyter_builder.npm_builder" 57 | ensured-targets = [ 58 | "jupyterlab_genv/labextension/static/style.js", 59 | "jupyterlab_genv//labextension/package.json", 60 | ] 61 | skip-if-exists = ["jupyterlab_genv/labextension/static/style.js"] 62 | 63 | [tool.hatch.build.hooks.jupyter-builder.build-kwargs] 64 | build_cmd = "build:prod" 65 | npm = ["jlpm"] 66 | 67 | [tool.hatch.build.hooks.jupyter-builder.editable-build-kwargs] 68 | build_cmd = "install:extension" 69 | npm = ["jlpm"] 70 | source_dir = "src" 71 | build_dir = "jupyterlab_genv/labextension" 72 | 73 | [tool.tbump] 74 | field = [ 75 | { name = "channel", default = "" }, 76 | { name = "release", default = "" }, 77 | ] 78 | file = [ 79 | { src = "pyproject.toml", version_template = "version = \"{major}.{minor}.{patch}{channel}{release}\"" }, 80 | { src = "jupyterlab_genv/_version.py" }, 81 | { src = "package.json" }, 82 | ] 83 | 84 | [tool.tbump.version] 85 | current = "0.4.0" 86 | regex = "(?P\\d+)\\.(?P\\d+)\\.(?P\\d+)((?Pa|b|rc|.dev)(?P\\d+))?" 87 | 88 | [tool.tbump.git] 89 | message_template = "Bump to {new_version}" 90 | tag_template = "v{new_version}" 91 | 92 | [project.entry-points."jupyter_client.kernel_provisioners"] 93 | genv-provisioner = "jupyterlab_genv.genv_provisioner:GenvProvisioner" 94 | -------------------------------------------------------------------------------- /resources/readme/activate.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/run-ai/jupyterlab_genv/4b29d360fd198ceb750ab0c206b14159cc9c34d5/resources/readme/activate.gif -------------------------------------------------------------------------------- /resources/readme/attach.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/run-ai/jupyterlab_genv/4b29d360fd198ceb750ab0c206b14159cc9c34d5/resources/readme/attach.gif -------------------------------------------------------------------------------- /resources/readme/commands.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/run-ai/jupyterlab_genv/4b29d360fd198ceb750ab0c206b14159cc9c34d5/resources/readme/commands.gif -------------------------------------------------------------------------------- /resources/readme/overview.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/run-ai/jupyterlab_genv/4b29d360fd198ceb750ab0c206b14159cc9c34d5/resources/readme/overview.gif -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | __import__('setuptools').setup() 2 | -------------------------------------------------------------------------------- /src/dialogs.tsx: -------------------------------------------------------------------------------- 1 | import { 2 | Dialog, 3 | InputDialog, 4 | showDialog, 5 | ReactWidget 6 | } from '@jupyterlab/apputils'; 7 | 8 | import React from 'react'; 9 | 10 | export namespace Dialogs { 11 | export async function noKernel(): Promise { 12 | const { button } = await showDialog({ 13 | title: 'No Kernel', 14 | body: 'You need a kernel in order to run in a GPU environment.', 15 | buttons: [ 16 | Dialog.cancelButton({ label: 'Later' }), 17 | Dialog.warnButton({ label: 'Select kernel', accept: true }) 18 | ] 19 | }); 20 | 21 | return button.accept; 22 | } 23 | 24 | export async function notSupportedKernel(): Promise { 25 | const { button } = await showDialog({ 26 | title: 'Not a genv Kernel', 27 | body: ReactWidget.create( 28 | <> 29 | Please select a genv kernel. 30 |
31 | If you don't have any, run the following command: 32 |
33 |
34 | python -m jupyterlab_genv install 35 | 36 | ), 37 | buttons: [ 38 | Dialog.cancelButton({ label: 'Later' }), 39 | Dialog.warnButton({ label: 'Select kernel', accept: true }) 40 | ] 41 | }); 42 | 43 | return button.accept; 44 | } 45 | 46 | export async function activate( 47 | envs: { eid: string; name: string }[], 48 | kernel_id: string 49 | ): Promise { 50 | const placeholder = 'Create a new environment'; 51 | 52 | function desc(env: { eid: string; name: string }): string { 53 | return env.name ? `${env.name} (${env.eid})` : env.eid; 54 | } 55 | 56 | const values = new Map([ 57 | [placeholder, kernel_id], 58 | ...(envs.map(env => [desc(env), env.eid]) as [string, string][]) 59 | ]); 60 | 61 | let { value } = await InputDialog.getItem({ 62 | title: 'Activate GPU Environment', 63 | items: [...values.keys()], 64 | okLabel: 'Next' 65 | }); 66 | 67 | if (value) { 68 | value = values.get(value) || value; 69 | } 70 | 71 | return value; 72 | } 73 | 74 | export async function configure(eid: string): Promise { 75 | const { button } = await showDialog({ 76 | title: 'Configure GPU Environment', 77 | body: ReactWidget.create( 78 | <> 79 | Open a terminal and run the following command: 80 |
81 |
82 | genv activate --id {eid} 83 |
84 | Then, configure the environment with normal genv commands. 85 |
86 |
87 | If you are not familiar with how to configure genv environments, check 88 | out the genv reference. 89 |
90 | You can find it at https://github.com/run-ai/genv. 91 |
92 |
93 | IMPORTANT 94 | You will need to restart the kernel for changes form the terminal to 95 | effect. 96 | 97 | ), 98 | buttons: [ 99 | Dialog.cancelButton({ label: 'Later' }), 100 | Dialog.okButton({ label: 'Open a terminal' }) 101 | ] 102 | }); 103 | 104 | return button.accept; 105 | } 106 | } 107 | -------------------------------------------------------------------------------- /src/handler.ts: -------------------------------------------------------------------------------- 1 | import { URLExt } from '@jupyterlab/coreutils'; 2 | 3 | import { ServerConnection } from '@jupyterlab/services'; 4 | 5 | /** 6 | * Call the API extension 7 | * 8 | * @param endPoint API REST end point for the extension 9 | * @param init Initial values for the request 10 | * @returns The response body interpreted as JSON 11 | */ 12 | async function requestAPI( 13 | endPoint = '', 14 | init: RequestInit = {} 15 | ): Promise { 16 | // Make request to Jupyter API 17 | const settings = ServerConnection.makeSettings(); 18 | const requestUrl = URLExt.join( 19 | settings.baseUrl, 20 | 'jupyterlab-genv', // API Namespace 21 | endPoint 22 | ); 23 | 24 | let response: Response; 25 | try { 26 | response = await ServerConnection.makeRequest(requestUrl, init, settings); 27 | } catch (error) { 28 | throw new ServerConnection.NetworkError(error); 29 | } 30 | 31 | let data: any = await response.text(); 32 | 33 | if (data.length > 0) { 34 | try { 35 | data = JSON.parse(data); 36 | } catch (error) { 37 | console.log('Not a JSON response body.', response); 38 | } 39 | } 40 | 41 | if (!response.ok) { 42 | throw new ServerConnection.ResponseError(response, data.message || data); 43 | } 44 | 45 | return data; 46 | } 47 | 48 | export namespace Handler { 49 | export async function activate( 50 | kernel_id: string, 51 | eid: string 52 | ): Promise { 53 | const body = JSON.stringify({ 54 | eid: eid, 55 | kernel_id: kernel_id 56 | }); 57 | 58 | await requestAPI('activate', { 59 | body: body, 60 | method: 'POST' 61 | }); 62 | } 63 | 64 | export async function devices(): Promise<{ eid: string }[]> { 65 | return await requestAPI('devices'); 66 | } 67 | 68 | export async function envs(): Promise< 69 | { eid: string; name: string; user: string }[] 70 | > { 71 | return await requestAPI('envs'); 72 | } 73 | 74 | export async function find(kernel_id: string): Promise { 75 | return (await requestAPI(`find?kernel_id=${kernel_id}`)) || null; 76 | } 77 | } 78 | -------------------------------------------------------------------------------- /src/index.tsx: -------------------------------------------------------------------------------- 1 | import { 2 | JupyterFrontEnd, 3 | JupyterFrontEndPlugin 4 | } from '@jupyterlab/application'; 5 | 6 | import { 7 | ICommandPalette, 8 | MainAreaWidget, 9 | ToolbarButton 10 | } from '@jupyterlab/apputils'; 11 | 12 | import { refreshIcon } from '@jupyterlab/ui-components'; 13 | 14 | import { Kernel } from '@jupyterlab/services'; 15 | import { ITerminal } from '@jupyterlab/terminal'; 16 | 17 | import { DocumentRegistry } from '@jupyterlab/docregistry'; 18 | import { INotebookModel, NotebookPanel } from '@jupyterlab/notebook'; 19 | 20 | import { IDisposable } from '@lumino/disposable'; 21 | import { Message } from '@lumino/messaging'; 22 | import { Widget } from '@lumino/widgets'; 23 | 24 | import { Handler } from './handler'; 25 | 26 | import { Dialogs } from './dialogs'; 27 | 28 | async function openTerminal(eid: string, app: JupyterFrontEnd): Promise { 29 | // NOTE(raz): the terminal is returned only when it's created in the first time. 30 | // this means that we can't send commands to the terminal if it's already running. 31 | // we should consider either creating a terminal per kernel or fixing this. 32 | // we tried opening a terminal per kernel but it seems like terminal names can't 33 | // be long enough to contain a kernel identifier. 34 | // here's a reference: 35 | // https://github.com/jupyterlab/jupyterlab/blob/v3.4.7/packages/terminal-extension/src/index.ts#L323 36 | const terminal: MainAreaWidget | undefined = 37 | await app.commands.execute('terminal:open', { name: 'genv' }); 38 | 39 | if (terminal) { 40 | app.shell.add(terminal, 'main', { mode: 'split-bottom' }); 41 | 42 | terminal.content.session.send({ 43 | type: 'stdin', 44 | content: [ 45 | [ 46 | '# this is a terminal for configuring your genv environment.', 47 | '# it will be activated in your environment.', 48 | '# you can configure your environment and attach devices from here.', 49 | '# ', 50 | '# you can start with running the following command:', 51 | '# ', 52 | '# genv attach --help', 53 | '# ', 54 | '# for more information check out the reference at https://github.com/run-ai/genv', 55 | '# ', 56 | '# IMPORTANT: you will need to restart your Jupyter kernel after configuring the environment from the terminal.', 57 | '', 58 | 'eval "$(genv init -)"', 59 | `genv activate --id ${eid}` 60 | ] 61 | .map(line => line + '\n') 62 | .join('') 63 | ] 64 | }); 65 | } 66 | } 67 | 68 | async function handleClick( 69 | kernel: Kernel.IKernelConnection | null | undefined, 70 | app: JupyterFrontEnd 71 | ) { 72 | if (kernel) { 73 | const spec = await kernel.spec; 74 | 75 | if (spec?.name.endsWith('-genv')) { 76 | let eid: string | null = await Handler.find(kernel.id); 77 | 78 | if (!eid) { 79 | const envs = await Handler.envs(); 80 | 81 | eid = await Dialogs.activate(envs, kernel.id); 82 | 83 | if (eid) { 84 | await Handler.activate(kernel.id, eid); 85 | } 86 | } 87 | 88 | if (eid) { 89 | if (await Dialogs.configure(eid)) { 90 | await openTerminal(eid, app); 91 | } 92 | } 93 | } else { 94 | if (await Dialogs.notSupportedKernel()) { 95 | await app.commands.execute('notebook:change-kernel'); 96 | } 97 | } 98 | } else { 99 | if (await Dialogs.noKernel()) { 100 | await app.commands.execute('notebook:change-kernel'); 101 | } 102 | } 103 | } 104 | 105 | export class ButtonExtension 106 | implements DocumentRegistry.IWidgetExtension 107 | { 108 | constructor(app: JupyterFrontEnd) { 109 | this._app = app; 110 | } 111 | 112 | createNew( 113 | panel: NotebookPanel, 114 | _context: DocumentRegistry.IContext 115 | ): IDisposable { 116 | // Create the toolbar button 117 | const mybutton = new ToolbarButton({ 118 | label: 'GPUs', 119 | tooltip: 'Configure the GPU environment', 120 | onClick: async () => { 121 | await handleClick(panel.sessionContext.session?.kernel, this._app); 122 | } 123 | }); 124 | 125 | // Add the toolbar button to the notebook toolbar 126 | panel.toolbar.insertItem(10, 'mybutton', mybutton); 127 | 128 | // The ToolbarButton class implements `IDisposable`, so the 129 | // button *is* the extension for the purposes of this method. 130 | return mybutton; 131 | } 132 | 133 | private _app; 134 | } 135 | 136 | class DevicesWidget extends Widget { 137 | private div?: HTMLDivElement; 138 | 139 | async onUpdateRequest(_msg: Message): Promise { 140 | const devices = await Handler.devices(); 141 | 142 | if (this.div) { 143 | this.node.removeChild(this.div); 144 | } 145 | 146 | this.div = document.createElement('div'); 147 | this.node.appendChild(this.div); 148 | 149 | for (const index in devices) { 150 | const device = devices[index]; 151 | const div = document.createElement('div'); 152 | 153 | if (device.eid) { 154 | div.innerText = `GPU ${index}: used by environment ${device.eid}`; 155 | } else { 156 | div.innerText = `GPU ${index}: available`; 157 | } 158 | 159 | this.div.appendChild(div); 160 | } 161 | } 162 | } 163 | 164 | class EnvsWidget extends Widget { 165 | private div?: HTMLDivElement; 166 | 167 | async onUpdateRequest(_msg: Message): Promise { 168 | const envs = await Handler.envs(); 169 | 170 | if (this.div) { 171 | this.node.removeChild(this.div); 172 | } 173 | 174 | this.div = document.createElement('div'); 175 | this.node.appendChild(this.div); 176 | 177 | for (const env of envs) { 178 | const div = document.createElement('div'); 179 | 180 | div.innerText = `${env.eid} ${env.user}`; 181 | 182 | if (env.name) { 183 | div.innerText += ` ${env.name}`; 184 | } 185 | 186 | this.div.appendChild(div); 187 | } 188 | } 189 | } 190 | 191 | const plugin: JupyterFrontEndPlugin = { 192 | id: 'jupyterlab_genv:plugin', 193 | autoStart: true, 194 | requires: [ICommandPalette], 195 | activate: async (app: JupyterFrontEnd, palette: ICommandPalette) => { 196 | app.docRegistry.addWidgetExtension('Notebook', new ButtonExtension(app)); 197 | 198 | const devicesContent = new DevicesWidget(); 199 | const devicesWidget = new MainAreaWidget({ content: devicesContent }); 200 | 201 | devicesWidget.id = 'jupyterlab_genv.devices'; 202 | devicesWidget.title.label = 'GPUs: Devices'; 203 | devicesWidget.title.closable = true; 204 | devicesWidget.toolbar.insertItem( 205 | 0, 206 | 'refresh', 207 | new ToolbarButton({ 208 | icon: refreshIcon, 209 | tooltip: 'Refresh', 210 | onClick: () => { 211 | devicesContent.update(); 212 | } 213 | }) 214 | ); 215 | 216 | const devicesCommand = 'jupyterlab_genv.devices.open'; 217 | 218 | app.commands.addCommand(devicesCommand, { 219 | label: 'GPUs: Show Devices', 220 | execute: () => { 221 | if (!devicesWidget.isAttached) { 222 | app.shell.add(devicesWidget, 'main'); 223 | } 224 | 225 | app.shell.activateById(devicesWidget.id); 226 | } 227 | }); 228 | 229 | palette.addItem({ command: devicesCommand, category: 'GPUs' }); 230 | 231 | const envsContent = new EnvsWidget(); 232 | const envsWidget = new MainAreaWidget({ content: envsContent }); 233 | 234 | envsWidget.id = 'jupyterlab_genv.envs'; 235 | envsWidget.title.label = 'GPUs: Environments'; 236 | envsWidget.title.closable = true; 237 | envsWidget.toolbar.insertItem( 238 | 0, 239 | 'refresh', 240 | new ToolbarButton({ 241 | icon: refreshIcon, 242 | tooltip: 'Refresh', 243 | onClick: () => { 244 | envsContent.update(); 245 | } 246 | }) 247 | ); 248 | 249 | const envsCommand = 'jupyterlab_genv.envs.open'; 250 | 251 | app.commands.addCommand(envsCommand, { 252 | label: 'GPUs: Show Environments', 253 | execute: () => { 254 | if (!envsWidget.isAttached) { 255 | app.shell.add(envsWidget, 'main'); 256 | } 257 | 258 | app.shell.activateById(envsWidget.id); 259 | } 260 | }); 261 | 262 | palette.addItem({ command: envsCommand, category: 'GPUs' }); 263 | } 264 | }; 265 | 266 | export default plugin; 267 | -------------------------------------------------------------------------------- /style/base.css: -------------------------------------------------------------------------------- 1 | /* 2 | See the JupyterLab Developer Guide for useful CSS Patterns: 3 | 4 | https://jupyterlab.readthedocs.io/en/stable/developer/css.html 5 | */ 6 | -------------------------------------------------------------------------------- /style/index.css: -------------------------------------------------------------------------------- 1 | @import url('base.css'); 2 | -------------------------------------------------------------------------------- /style/index.js: -------------------------------------------------------------------------------- 1 | import './base.css'; 2 | -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "allowSyntheticDefaultImports": true, 4 | "composite": true, 5 | "declaration": true, 6 | "esModuleInterop": true, 7 | "incremental": true, 8 | "jsx": "react", 9 | "module": "esnext", 10 | "moduleResolution": "node", 11 | "noEmitOnError": true, 12 | "noImplicitAny": true, 13 | "noUnusedLocals": true, 14 | "preserveWatchOutput": true, 15 | "resolveJsonModule": true, 16 | "outDir": "lib", 17 | "rootDir": "src", 18 | "strict": true, 19 | "strictNullChecks": true, 20 | "target": "es2017", 21 | "types": [] 22 | }, 23 | "include": ["src/*"] 24 | } 25 | --------------------------------------------------------------------------------