Zenn
🧨

poetry を用いて アプリケーションでは torch-cpu, バッチでは torch(GPU) を使う

2023/11/17に公開
poetry
tech

zenn で書くのは初めてです。よろしくお願いします。
初めてなのでちょっとした覚え書きを書いてみます。

poetry 環境下で torch を使っているとき、2023/11/17 時点では pyproject.toml で以下のようにすることで torch-cpu, torch(GPU) を切り替えられることは皆さんご存知かなと思います。

torch = [
    {version = "^2.0.1", markers="sys_platform == 'darwin'", source = "default"},
    {version = "^2.0.1", markers="sys_platform != 'darwin'", source = "torch-cpu"},
]

[[tool.poetry.source]]
name = "default"
url = "https://pypi.python.org/simple/"
priority = "default"

[[tool.poetry.source]]
name = "torch-cpu"
url = "https://download.pytorch.org/whl/cpu"
priority = "explicit"

さて、モノレポというほどではなくても同じリポジトリ内でモデル学習のバッチと、それをホストして運用される API とを同時に扱っているケースは多いのではないか、と思います。そんなとき、学習バッチでは速度の観点 GPU を使うけど、API では費用などの面から CPU で推論したいんじゃないでしょうか。
torch(GPU) を使うと、nvidia runtime のパッケージが大量に入ってくるので、イメージサイズを減らしたいという観点からもこうしたモチベーションは結構普通なんじゃないかなと思います。

こんな時に役立つのが poetry dependencies group です。以下のような形で app, batch を切り替えてモジュールを利用することができます。
sys_platform の部分は macos, linux 上を判定するためのものです。

[tool.poetry.group.app]
optional = true

[tool.poetry.group.app.dependencies]
torch = [
    {version = "^2.0.1", markers="sys_platform == 'darwin'", source = "default"},
    {version = "^2.0.1", markers="sys_platform != 'darwin'", source = "torch-cpu"},
]

[tool.poetry.group.batch]
optional = true

[tool.poetry.group.batch.dependencies]
torch = {version = "^2.0.1", source = "default"}

こんな風にした場合、 poetry install --with app --without batch --sync のようなコマンドでインストールすることで、app 定義だけを lock file に反映させつつインストールできます。
sync オプションをつけずにやっていてしばらくハマりました...

Dockerfile は次のようになります。

docker/app/Dockerfile

FROM public.ecr.aws/docker/library/python:3.10-slim as builder
WORKDIR /app
RUN apt-get update \
    && apt-get install -y build-essential git \
    && pip install --no-cache-dir poetry

COPY poetry.lock pyproject.toml ./

ENV POETRY_NO_INTERACTION=1 \
  POETRY_VIRTUALENVS_IN_PROJECT=1 \
  POETRY_VIRTUALENVS_CREATE=1 \
  POETRY_CACHE_DIR=/tmp/poetry_cache
RUN poetry install --no-root --with app --without batch,dev --sync && rm -rf /tmp/poetry_cache

FROM public.ecr.aws/docker/library/python:3.10-slim
WORKDIR /app
RUN apt-get update \
    && apt-get install -y build-essential \
    && pip install --no-cache-dir poetry
COPY --from=builder /app/.venv /app/.venv

COPY src ./src
COPY Makefile poetry.lock pyproject.toml REVISION ./
RUN mkdir model

CMD ["make", "run-api"]

docker/batch/Dockerfile

FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

RUN apt-get update -y \
    && apt-get -y install build-essential python3-pip \
    && rm -rf /var/lib/apt/lists

ENV LC_ALL=C.UTF-8
ENV LANG=C.UTF-8

WORKDIR /work

COPY poetry.lock pyproject.toml ./
RUN python3 -m pip install --no-cache-dir --upgrade pip && \
    python3 -m pip install --no-cache-dir poetry awscli
RUN poetry install --no-root --with batch --without app,dev --sync

COPY src ./src
COPY scripts ./scripts
COPY tests ./tests
COPY sql ./sql
COPY configs ./configs
COPY Makefile setup.cfg REVISION ./
RUN mkdir model output data

CMD ["/bin/bash"]

ということで、poetry 上で dependencies group を使いつつ torch-cpu, torch(GPU) を切り替える方法の紹介でした。

Discussion