How to run a cuPy python program on NVIDIA NGC?

Question

I have a python program that utilizes NVIDIA GPU using cuPy package. The program runs fine with my local laptop's GPU and on a ubuntu GPU cluster. Now for scalability reasons, I want to make my program able to run on NVIDIA GPU Cloud (NGC).

To run the program on NGC, what I understood so far is, I need to do the following:

• Containerized my program
• install docker
• run the program using docker

with my very limited experience with Docker, I am not able to understand the following:

• how do I install cuPy docker image? Should I use install Chainer image instead?
• my program contains some with cupy.cuda.Device(device=device_id) statements to select different GPUs, can I keep the statements? or do I need to remove these statements?

I want to start with trying to run the following simple program on NGC, but not sure the exact steps I need to follow.

MWE:

import cupy as cp

# Create a random matrix on the GPU
a_gpu = cp.random.rand(1000, 1000)

# Perform a matrix multiplication on the GPU
result_gpu = cp.dot(a_gpu, a_gpu)

# Transfer the result back to the CPU
result_cpu = cp.asnumpy(result_gpu)

# Print the result
print(result_cpu)