JAX

Run a ready-to-use Docker image with JAX and JupyterLab IDE on top of your ROCm-enabled system! All software is preinstalled under containers, so you can jump straight into your workflow—get started with either SSH or Jupyter.

Software Included

Creating an App using the Control Panel

Click the Deploy to DigitalOcean button to create a Droplet based on this 1-Click App. If you aren’t logged in, this link will prompt you to log in with your DigitalOcean account.

Creating an App using the API

In addition to creating a Droplet from the JAX 1-Click App using the control panel, you can also use the DigitalOcean API. As an example, to create a 4GB JAX Droplet in the SFO2 region, you can use the following curl command. You need to either save your API access token) to an environment variable or substitute it in the command below.

curl -X POST -H 'Content-Type: application/json' \
         -H 'Authorization: Bearer '$TOKEN'' -d \
        '{"name":"choose_a_name","region":"sfo2","size":"s-2vcpu-4gb","image": "digitaloceanai-jax"}' \
        "https://api.digitalocean.com/v2/droplets"

Getting Started After Deploying JAX

Accessing JAX + ROCm resources and libraries

There are two ways to utilise JAX + ROCm resources and libraries: directly from JupyterLab, or via a Docker container.

• JupyterLab is preferable if you have notebook projects or just want to explore and experiment.
• Accessing Docker containers is preferable if you have a Python project that requires deployment.

The JupyterLab option is recommended for first-time users.

Either way will require you to SSH into the droplet first.

SSHing into your AMD 1-Click

The simplest way to SSH is to use the DigitalOcean web console. After your Droplet is created, you can find the “Web Console” button in the top right corner of the Droplet page:

After clicking the button, a separate browser page with Web Console will open.

Please note that the web console may be inaccessible for the first minute or two after the droplet is created.

In this case, refresh the page and try again.

Exploring via Jupyter

JupyterLab is an interactive development environment for working with notebooks, code, and data. It offers a flexible and powerful interface for data science and scientific computing, allowing you to create and share documents that contain live code, equations, visualisations, and narrative text.

After the web console establishes the SSH connection, you will see your Droplet terminal with a welcome message.

Inside the welcome message, locate the Token field and copy it:

Now go back to the Droplet page and copy the IPv4 address of your 1-Click; you can locate it under “Connection Details”:

Paste IPv4 into your browser and press Enter to access the JupyterLab environment:

If you see a page that says “Your JupyterLab is being configured”, just wait a minute until JupyterLab starts, and the page will automatically refresh itself.

After accessing JupyterLab, you will be greeted by an authentication screen. Here you can use the token from earlier to log in to your JupyterLab:

After entering the access token, you’ll be logged into Jupyter Lab’s web-based IDE.

You can now start creating, editing, and running Jupyter notebooks directly within your browser!

Exploring via Docker

After SSHing into the Droplet, run this command to connect to Docker container hosting all the ROCm software:

docker exec -it rocm /bin/bash

Done! You can now freely implement any project that requires JAX + ROCm libraries.