Back to Serve

Using Jina on Colab with GPU/TPU support

docs/Using_Jina_on_Colab.ipynb

3.34.04.5 KB
Original Source

Using Jina on Colab with GPU/TPU support

In this tutorial, we will cover two ways of using Colab as illustrated below:

1. Change runtime type

Go to menu Runtime -> Change run time type -> GPU/TPU

2. Install Jina

You will be asked to restart the kernel after installation. Go ahead and click "Restart Runtime"

python
!pip install jina

3. Import

python
from jina import Flow, Document, Executor, requests, DocumentArray
import jina

print(jina.__version__)

4. Build a GPU Executor

python
import torch

class GPUExec(Executor):

  @requests
  def foo(self, docs: DocumentArray, **kwargs):
    docs[0].tags['cuda'] = torch.cuda.is_available()

5. Add GPU Executor to a Flow

python
f = Flow().add(uses=GPUExec)

f.start()

6. Send request to the running Flow

The result should contain a single Document which has .tags['cuda'] field set to True via GPUExec.

python
r = f.post('/', Document())
print(r[0].tags)

Now we see that one can directly connect to a running Flow inside Google Colab. The trick here is to use f.start() to start the Flow and use f.close() to close the Flow. This usage is different than our classic with context manager.

python
f.close()

Connecting from local to a running Flow on Colab

A more interesting use case is to connect to Colab from a local machine. In this case, we need to install ngrox.

python
!pip install pyngrok

1. Expose via HTTP

Easiest and most compatible, no registration needed on ngrox. However, HTTP protocol is not performant on large data and streaming.

python
f = Flow(protocol='http', port=54321).add(uses=GPUExec)

f.start()
python
!ngrok http 54321 --log "stdout"

Notice the last line, you should see something like the following:

t=2022-06-11T20:06:49+0000 lvl=info msg="started tunnel" obj=tunnels name=command_line addr=http://localhost:54321 url=https://7ee9-35-197-36-214.ngrok.io

Now use the string after url= to replace the following curl command. In this example, my url=https://7ee9-35-197-36-214.ngrok.io. Yours would be different!

bash
curl -X POST https://7ee9-35-197-36-214.ngrok.io/post \
     -H 'Content-Type: application/json' \
     -d '{"data":[{}]}'

Running this curl command in your local terminal. You will get:

json
{"header":{"requestId":"008c15ac9770412b9ea8fcd3b0944520","status":null,"execEndpoint":"/","targetExecutor":""},"parameters":null,"routes":[{"executor":"gateway","startTime":"2022-06-11T20:09:30.189579+00:00","endTime":"2022-06-11T20:09:30.227964+00:00","status":null},{"executor":"executor0","startTime":"2022-06-11T20:09:30.190086+00:00","endTime":"2022-06-11T20:09:30.227796+00:00","status":null}],"data":[{"id":"402bddfd49311dab4128dcfae5cec9ba","parent_id":null,"granularity":null,"adjacency":null,"blob":null,"tensor":null,"mime_type":null,"text":null,"weight":null,"uri":null,"tags":{"cuda":true},"offset":null,"location":null,"embedding":null,"modality":null,"evaluations":null,"scores":null,"chunks":null,"matches":null}]}

Where one can see "tags":{"cuda":true} is successfully returned.

To stop the server, click stop the last cell. Then close the Flow.

python
f.close()

2. Expose via gRPC

Using ngrox to forward gRPC is also possible. But you will need to first sign up at https://dashboard.ngrok.com/signup

After signing up, you can get a token. Then simply add your token via (replacing YOUR_TOKEN_HERE)

python
!ngrok authtoken 2ARf0Y7QrG3E2TsFaX7W3KWvfGD_6R972KKvbqeucpCrCuEHv

Now let's start the Flow with gRPC gateway again.

python
f = Flow(port=54321).add(uses=GPUExec)

f.start()

Let's get its public address via ngrok

python
!ngrok tcp 54321 --log "stdout"

At the last line, you should see something like:

t=2022-06-11T20:29:11+0000 lvl=info msg="started tunnel" obj=tunnels name=command_line addr=//localhost:54321 url=tcp://6.tcp.ngrok.io:18096

Grab the text after url=tcp:// in my case it is 6.tcp.ngrok.io:18096.

Now build a client using this address from your local laptop/Python environment.

Copy paste the code below to your local Python, remmeber to change your address.

python
from jina import Client, Document


c = Client(host='grpc://6.tcp.ngrok.io:18096')
r = c.post('/', Document())
print(r[0].tags['cuda'])

And you will get

text
True

Showing the connection is success!

Now enjoy the free GPU/TPU to build your awesome Jina applications!