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 ROCm + JAX

There are two ways to access and use ROCm + JAX libraries: either by SSHing and accessing the underlying container or using JupyterLab from a browser. SSHing allows more flexibility and is suitable if you have a pure-Python project to deploy. On the other hand, JupyterLab is more accessible and comes with pre-installed tutorials.

Approach A: SSH

First of all, connect to Droplet 1-Click with SSH. Use the droplet’s public IP address to log in as root:

ssh root@your_droplet_public_IP

In case of a problem, make sure to check our SSH guide.

After successfully connecting to the droplet, now it’s time to connect to the underlying container with ROCm + JAX libraries:

docker exec -it rocm /bin/bash

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

Approach B: JupyterLab

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, visualizations, and narrative text.

Quickly Get Started With Your Jupyter Droplet

1. Access the Droplet Console:

   • Navigate to the Droplets page.
   • Locate your newly created Jupyter Droplet and click on its name.
   • Under the “Access” tab, select Console. This will open an in-browser terminal session connected to your droplet.
   • Log in as the root user using the password you set during droplet creation.

2. Login via SSH (Method 1: Using SSH Key):

• If you selected an SSH key during droplet creation, follow these steps:- Open your preferred SSH client (e.g., PuTTY, Terminal).

  • Use the droplet’s public IP address to log in as root:

  ssh root@your_droplet_public_IP

  • Ensure your SSH key is added to the SSH agent, or specify the key file directly:

  ssh -i /path/to/your/private_key root@your_droplet_public_IP

  • Once connected, you will be logged in as the root user without needing a password.

2. Check the Message of the Day (MOTD) for Jupyter Access Token:

   • Upon successful login via console or SSH, the Message of the Day (MOTD) will be displayed.
   • This message includes important information such as the Jupyter Lab access token. Take note of this token as you’ll need it to log in to Jupyter Lab.

5. Access the Jupyter Lab IDE:

   • Open a web browser and navigate to the URL for your Jupyter Lab instance, which will be http://your_droplet_public_IP.
   • This will bring you to the login page for Jupyter Lab.

6. Enter the Jupyter Lab Access Token:

   • On the login page, enter the access token that you noted from the MOTD.
   • This token is required for the first login to ensure secure access to your Jupyter Lab environment.

7. Start Using Jupyter Lab:

   • 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.