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

前提

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

必要知識

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

今回のゴール

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

Spaceで利用するGPUの選択

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

Dockerfileの編集

FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04

ARG DEBIAN_FRONTEND=noninteractive

RUN ln -sf /usr/share/zoneinfo/Asia/Tokyo /etc/localtime

RUN apt-get update && \
    apt-get install --no-install-recommends -y \
    build-essential \
    curl \
    python3.9 \
    python3-pip \
    git \
    ffmpeg

# Set the working directory to /code
WORKDIR /code

COPY ./requirements.txt /code/requirements.txt

# Set up a new user named "user" with user ID 1000
RUN useradd -m -u 1000 user

# Switch to the "user" user
USER user

# Set home to the user's home directory
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:$PATH

RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt

# Set the working directory to the user's home directory
WORKDIR $HOME/app

# Copy the current directory contents into the container at $HOME/app setting the owner to the user
COPY --chown=user . $HOME/app

# Start the FastAPI app on port 7860, the default port expected by Spaces
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

前回から大きく変わった点はベースイメージにnvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04を指定している点である。
このイメージは、NVIDIAが提供する公式のCUDAコンテナで、GPUを活用するアプリケーションの開発環境として利用される。

app.pyの編集

from fastapi import FastAPI
import subprocess
import torch
import os
import huggingface_hub
import subprocess
import sys
import requests

# Create a new FastAPI app instance
app = FastAPI()

@app.get("/")
def home():
    device_name = torch.cuda.get_device_name(0) # GPUデバイス名の取得

    torch_version = str(torch.__version__) # torchのバージョンの取得

    info = {
        device_name,
        torch_version
    }

    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を活用した様々な機械学習モデルの推論や処理に対応できるサーバーの開発が期待できる。