stikkireddy / dbtunnel

Proxy solution to run elegant Web UIs or interact with LLMs natively inside databricks notebooks.
https://stikkireddy.github.io/dbtunnel/
Apache License 2.0
18 stars 7 forks source link

Add support for Ray model serving #18

Closed natefleming closed 2 months ago

natefleming commented 2 months ago

This PR adds support for ray apps to dbtunnel.

This aims to cover use cases around use ray to serve REST endpoints.
The below example (included in examples/) demonstrates model serving.


ngrok_api_key: str = dbutils.secrets.get("fieldeng", "ngrok-api-key")
ngrok_tunnel_auth_token: str = dbutils.secrets.get("fieldeng", "ngrok-tunnel-auth-token")

from ray import serve
from ray.serve.deployment import Application

from transformers import pipeline

@serve.deployment(num_replicas=2, )
class SentimentAnalysis:

  def __init__(self):
    self._classifier = pipeline("sentiment-analysis")

  def __call__(self, request) -> str:
    input_text: str = request.query_params["input_text"]
    return self._classifier(input_text)[0]["label"]

app: Application = SentimentAnalysis.bind() #noqa: F821

from dbtunnel import dbtunnel

(
  dbtunnel.ray(app, port=8080)
    .inject_auth()
    .share_to_internet_via_ngrok(
        ngrok_api_token=ngrok_api_key, 
        ngrok_tunnel_auth_token=ngrok_tunnel_auth_token)
    .run()
)
stikkireddy commented 2 months ago

LGTM thanks for the feature!