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

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

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

## 仕組み

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

https://jupyter-server-proxy.readthedocs.io/en/latest/

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

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

```shell:Shell
python -m http.server 8080
```

:::note
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サーバーにアクセスできます。

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/d6f90002-7af6-d6ab-d75d-273907e53dee.png)

## Streamlitの場合

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

```shell: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」にアクセスしてみます。

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

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/5d483720-5816-fdf4-9810-62f9cc562ba1.png)

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/5876ca96-6a80-779b-6dd5-d26207cfd185.png)

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」を追加する

```shell:起動コマンド
streamlit run app.py --server.baseUrlPath /jupyterlab/default/proxy/absolute/8501
```

```text:アクセスURL
https://********.studio.****.sagemaker.aws/jupyterlab/default/proxy/absolute/8501/
```

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

![](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/41574/9d9333df-1698-cd32-74e5-134e40d3b444.png)

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


## おまけ

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

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

https://developers-jp.googleblog.com/2024/10/gemma-2-for-japan.html

```shell: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:Shell
streamlit run app.py --server.baseUrlPath /jupyterlab/default/proxy/absolute/8501
```

<iframe width="640" height="480" src="https://www.youtube.com/embed/zxLPBRiQNbA?si=pArZof6L7HEqngUH&amp;controls=0" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture; web-share" referrerpolicy="strict-origin-when-cross-origin" allowfullscreen></iframe>

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

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