SageMaker StudioのJupyterLab環境でStreamlitを使用する（Gemma 2 JPNと会話するおまけ付き）

SageMaker StudioのJupyterLab環境で、Proxy機能が使えることがわかりました！

この機能を使えば、StreamlitなどのWebアプリを動作させることができます。

仕組み

SageMaker StudioのJupyterLab環境には、「Jupyter Server Proxy」という拡張機能がインストールされています。

この拡張機能が、jupyterLab環境で起動しているWebサーバーのポートにプロキシしてくれます。

例えば、以下のコマンドでWebサーバーを起動します。

Shell

python -m http.server 8080

Pythonに標準で含まれるHTTPサーバー機能です。
https://docs.python.org/ja/3/library/http.server.html#module-http.server

SageMaker StudioのJupyterLab環境にアクセスしているURLは以下のような形式です。

https://**********.studio.us-east-1.sagemaker.aws/jupyterlab/default/lab

このURLの「/lab」以下を「/proxy/{ポート番号}」に変えると、そのポートにProxyしてくれます。

今回の8080ポートの場合は、以下のURLへアクセスします。

https://**********.studio.us-east-1.sagemaker.aws/jupyterlab/default/proxy/8080

こうすることで、8080番ポートで起動しているWebサーバーにアクセスできます。

Streamlitの場合

Streamlitの場合はもうひと手間必要ですが、まずは、先程紹介した方法でアクセスしてみましょう。
以下のコマンドででもアプリを起動します。

Shell

streamlit hello

Collecting usage statistics. To deactivate, set browser.gatherUsageStats to false.


  Welcome to Streamlit. Check out our demo in your browser.

  Local URL: http://localhost:8501
  Network URL: http://169.255.255.2:8501
  External URL: http://52.4.240.77:8501

  Ready to create your own Python apps super quickly?
  Head over to https://docs.streamlit.io

  May you create awesome apps!

ポート番号未指定の場合はデフォルトの8501番ポートで起動します。
「/jupyterlab/default/proxy/8501」にアクセスしてみます。

結果は真っ白な画面が表示されました。

Chromeの開発者モードで確認すると、404エラーになっているものがありました。

HTMLの取得に成功しているものの、CSSやJavaScriptの取得に失敗しています。

一番下のファイルは「/jupyterlab/default/proxy/static/js/main.33ba0d1c.js」というURLでアクセスしていますが、「/jupyterlab/default/proxy /8501 /static/js/main.33ba0d1c.js」にアクセスする必要があります。

間にProxyが入っていることで相対パスが上手く処理できていません。

以下の2点を対応することで、正しく動作するようになります。

• URLに /absolute を追加する
• Streamlitの起動パラメーターで、絶対パスを指定する

具体的には、

• URLを「/jupyterlab/default/proxy/absolute/8501/」に変更する
• Streamlitの起動パラメーターに「--server.baseUrlPath /jupyterlab/default/proxy/absolute/8501」を追加する

これで正しく表示されました😀

こちらのドキュメントに記載がありました
https://jupyter-server-proxy.readthedocs.io/en/latest/arbitrary-ports-hosts.html

おまけ

JupyterLab環境でStreamlitが起動できたので、先日登場したGemma 2 JPNとチャットをしてみましょう。

使用したインスタンスタイプはml.g4dn.2xlargeです。

Shell

conda install -y accelerate pytorch transformers huggingface_hub pytorch-cuda=12.4 -c pytorch -c nvidia

app.py

import streamlit as st
import torch
from transformers import pipeline

if "pipe" not in st.session_state:
    st.session_state.pipe = pipeline(
        "text-generation",
        model="google/gemma-2-2b-jpn-it",
        model_kwargs={"torch_dtype": torch.bfloat16},
        device="cuda",
    )
pipe = st.session_state.pipe

st.title("Chat with Gemma 2 JPN on SageMaker")
st.subheader("ml.g4dn.2xlarge")

if "messages" not in st.session_state:
    st.session_state.messages = []
messages = st.session_state.messages

for m in messages:
    with st.chat_message(m["role"]):
        st.write(m["content"])

if prompt := st.chat_input():

    messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.write(prompt)

    outputs = pipe(messages, return_full_text=False, max_new_tokens=1024)
    assistant_response = outputs[0]["generated_text"].strip()

    messages.append({"role": "assistant", "content": assistant_response})
    with st.chat_message("assistant"):
        st.write(assistant_response)

パスを指定して起動します

Shell

streamlit run app.py --server.baseUrlPath /jupyterlab/default/proxy/absolute/8501

爆速ではありませんが、実用的な速さに感じました。量子化などもしていないので精度的にも最も良い状態と思われます。

ml.g4dn.2xlargeは時間あたり$0.94なので、用途によっては検討の余地があると思いましたが、いかがでしょうか？？