前提

この記事では、HuggingFace Space を app としてではなく API サーバーとして使う方法を扱う。もともとインターンで HuggingFace Inference Endpoints を使ってモデルの API 提供を行っていたのだが、エンドポイントを起動するたびに０からビルドしなくてはならず時間がかかっていたためビルドキャッシュが残る HuggingFace Space の DockerSDK を使った API サーバーの構築に着手した。 https://huggingface.co/docs/hub/spaces-sdks-docker 前回の記事では API サーバーの作成までを行ったが、GPU の利用までは触れなかったため今回の記事ではここを扱う。今回の記事は前回の記事と重複する内容は記述しないので適宜参照よろしくお願いします。

↓ 前回の記事

必要知識

HuggingFace Space の利用
Docker のセットアップ

今回のゴール

GPU を利用できる API サーバーを Huggingface Space 上で作成する

Space で利用する GPU の選択

Space の Settings タブにある Space Hardware から利用したい GPU を選択する。(デフォルトでは CPU basic になっている)

今回はA10G smallを選択。

Dockerfile の編集

1FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
2
3ARG DEBIAN_FRONTEND=noninteractive
4
5RUN ln -sf /usr/share/zoneinfo/Asia/Tokyo /etc/localtime
6
7RUN apt-get update && \
8    apt-get install --no-install-recommends -y \
9    build-essential \
10    curl \
11    python3.9 \
12    python3-pip \
13    git \
14    ffmpeg
15
16# Set the working directory to /code
17WORKDIR /code
18
19COPY ./requirements.txt /code/requirements.txt
20
21# Set up a new user named "user" with user ID 1000
22RUN useradd -m -u 1000 user
23
24# Switch to the "user" user
25USER user
26
27# Set home to the user's home directory
28ENV HOME=/home/user \
29    PATH=/home/user/.local/bin:$PATH
30
31RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
32
33# Set the working directory to the user's home directory
34WORKDIR $HOME/app
35
36# Copy the current directory contents into the container at $HOME/app setting the owner to the user
37COPY --chown=user . $HOME/app
38
39# Start the FastAPI app on port 7860, the default port expected by Spaces
40CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

前回から大きく変わった点はベースイメージにnvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04を指定している点である。

このイメージは、NVIDIA が提供する公式の CUDA コンテナで、GPU を活用するアプリケーションの開発環境として利用される。

app.py の編集

1from fastapi import FastAPI
2import subprocess
3import torch
4import os
5import huggingface_hub
6import subprocess
7import sys
8import requests
9
10# Create a new FastAPI app instance
11app = FastAPI()
12
13@app.get("/")
14def home():
15    device_name = torch.cuda.get_device_name(0) # GPUデバイス名の取得
16
17    torch_version = str(torch.__version__) # torchのバージョンの取得
18
19    info = {
20        device_name,
21        torch_version
22    }
23
24    return info

今回は GPU がプログラム上から見えていることを確認するためtorch.cuda.get_device_name(0)で cuda のデバイス名を表示、また現在の torch のバージョンをルートページに表示するようにした。

Space の起動

コードの変更をコミットすると自動でビルドが始まるが、GPU を選択すると画像のようにBuilding on {GPU名}となる。ビルドが終わるとRunning on {GPU名}となる。 Space が起動すると、ちゃんとルートページに使用している GPU 名と現在の torch のバージョンが表示された！

まとめ

本記事では、前回の記事の続きとして HuggingFace Space 上で GPU を利用できる API サーバーの作成を行った。

GPU 活用のための環境構築: NVIDIA が提供する公式 CUDA イメージ（nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04）をベースに採用することで、最新の GPU 計算環境と cuDNN ライブラリを手軽に利用可能となった。

今回の実装で、HuggingFace Space を活用した API サーバー構築の基盤が整った。GPU を活用した様々な機械学習モデルの推論や処理に対応できるサーバーの開発が期待できる。

HuggingFace SpaceをAPIサーバーとして使う その2[GPUの利用]

前提