Make組ブログ

Python、Webアプリや製品・サービス開発についてhirokikyが書きます。

Dockerで稼働しているCeleryが停止時にタスクを完了しない問題と修正方法

DockerでCeleryを稼働しているとき、設定によってはタスクを安全に完了(Graceful Shutdown)できていない場合があります。それぞれの登場人物の仕様を理解して間違いないようにしましょう。

コンテナランタイムの仕様

AWS Fargateではコンテナを停止する際、最初に SIGTERM が送信され、それでも30秒以内にプロセスが停止しない場合は SIGKILL が送信されます。この SIGTERM が適切にCeleryのプロセスに届いていないと安全なシャットダウンができません。

aws.amazon.com

CMDでシェル形式で書くとSIGTERMが届かない

Dockerでプロセスを起動する際、 CMD コマンドを指定することがあります。

CMD celery -A myapp

CMD をこのように指定する場合、Dockerは /bin/bash -c を先頭につけてこのコマンドを実行します。この場合 SIGTERMCMD に指定しているコマンドに届かず、 /bin/bash で止まってしまいます。

その結果、Dockerコンテナにコンテナランタイムが停止するよう命令しても、Celeryのようなタスクが完了せずに中断されてしまいます。

修正方法1:exec形式で書く

シェル形式でなくexec形式で書くと /bin/bash でラップされずに起動します。

CMD ["/usr/local/bin/celery", "-A", "myapp"]

Dockerfile リファレンス — Docker-docs-ja 24.0 ドキュメント

@naoina さんご指摘いただきありがとうございました!

修正方法2:ENTRYPOINTに書く

Celeryの起動コマンドをENTRYPOINTに設定し、CMDにはオプションパラメータを渡すように変更しても良いです。

ENTRYPOINT ["/usr/local/bin/celery"]
CMD ["-A", "myapp", "worker", "-l", "ERROR"]

こちらでもDockerコンテナが直接Celeryを起動するので、コンテナランタイムからのSIGTERMを受け取れます。

ENTRYPOINTとCMDとは何か

ENTRYPOINTには実行されるべきコマンドを書き、CMDには docker run で起動する際のデフォルトのオプションを書きます。デフォルトの実行コマンドを書き換えられたくない場合などにENTRYPOINTを使うことで、コンテナの用途を明確にできます。

www.docker.com

CMDだけの場合は2種類の書き方があるので、気をつけましょう。

とはいえSIGTERMの話とENTRYPOINT or CMDの話は直接的に関係がないので、よくDocker、bash、コンテナランタイム(Fargate)の挙動を知っておく必要があります。

おわりに

今回はとくにタスクの扱いが繊細なCeleryで書きましたが、他のアプリケーションでも同じことです。たとえばGunicornやFastAPIのBackgroundTaskの挙動も関係します。

現在、Shodoではアドベントカレンダー応援クーポンを配布しております。80%オフでShodoを最長3ヶ月間使えるクーポンです。以下のクーポンコードをご購入時に入力して、このアドベントカレンダーの季節にShodoのAI校正をお役立てください。

XMAS2024

shodo.ink

執筆:@hirokiky
Shodoで執筆されました

Optimum Neuronでresume_from_checkpointをするとTypeError None // int

先日の記事のように Optimun Neuronを使うとHuggingFace Trainerを簡単にAWS Neuron上で扱えます。

ですがまだあまり安定していないというのが正直なところです。今回の問題は0.0.23から0.0.25の範囲で確認済みです。記事に書いてしまっていますが、もし良ければちゃんと調査してIssueにしてください(すいません、時間が取れればやりたいと思ってます)。

Optimun Neuronでチェックポイントからの再開時にTypeError

Optimun Neuronでチェックポイントを残しつつ学習し、実際に resume_from_checkpoint=True をつけて再開すると以下の行で TypeError が発生します。

        train_bs_state = trainer_state.train_batch_size // max(1, training_args.n_gpu)

https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L1874

これは trainer_state.train_batch_size にNoneが設定されており、intとの比較でエラーとなっています。呼び出し元はこちら。

            self.compare_trainer_and_checkpoint_args(self.args, self.state)

https://github.com/huggingface/optimum-neuron/blob/v0.0.25/optimum/neuron/trainers.py#L900

Optimun Neuronの0.0.25ではNeuronTrainerの実装を見ても設定する処理がないようです。
https://github.com/huggingface/optimum-neuron/blob/v0.0.25/optimum/neuron/trainers.py

HuggingFace Trainerでは設定しています。
https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L2249

問題のとりあえずの回避方法

かなり雑な回避方法ですが、Optimun NeuronのNeuron Trainerでチェックポイントを今すぐ使いたい人は以下を参考にしてください(ちゃんと本家を修正しましょう。すいません)。この方法では読み込み時に、 training_state の値ではなく指定されたバッチサイズを使っています。

from optimum.neuron import NeuronTrainer

# Avoid a bug that NeuronTrainer does not store batch size to checkpoint
# And it cause None // int TypeError in compare_trainer_and_checkpoint_args.
class Trainer(NeuronTrainer):
    def compare_trainer_and_checkpoint_args(self, training_args, training_state):
        if training_state.train_batch_size is None:
            training_state.train_batch_size = self._train_batch_size
        super().compare_trainer_and_checkpoint_args(training_args, training_state)

バッチサイズの指定が別にならないようにだけは注意してください。

training_state.train_batch_size をちゃんと保存するようにしたかったのですが、良いところで処理を入れられずこう回避しています。

PRを出す人に向けて

個人的なメモでもありますが、PRを出す際はテストを修正して証明するのが手っ取り早いと思います。

ですが現状、Optimun Neuronで resume_from_checkpoint のテストはスキップされているようです。こちらを直してテストが通らないことを説明するのがPR提出のスタートラインになるかと思います。

https://github.com/huggingface/optimum-neuron/blob/v0.0.25/tests/test_trainers.py#L278

おわりに

Optimun Neuronはとても便利ですが、少しまだ不安定な様相も見せています。今後の開発に期待したいのと、同時に僕たちも協力していきたいものです。今回はこの記事に一旦まとめるということでご容赦ください(時間がとれればちゃんと検証してPRします)。

現在、Shodoではアドベントカレンダー応援クーポンを配布しております。80%オフでShodoを最長3ヶ月間使えるクーポンです。以下のクーポンコードをご購入時に入力して、このアドベントカレンダーの季節にShodoのAI校正をブログの執筆にお役立てください。

XMAS2024

shodo.ink

執筆:@hirokiky
Shodoで執筆されました

Optimun NeuronでHuggingFace TrainerをAWS Neuron上で簡単に動かす

Optimun Neuron を使うと、HuggingFace TrainerをAWS Neuron上で気軽に扱えるようになります。具体的には以下のように、TrainerとTrainingArgumentsを置き換えるだけです。

from optimum.neuron import NeuronTrainer as Trainer
from optimum.neuron import NeuronTrainingArguments as TrainingArguments

こうするとTrn1上で、今までと同じコードで学習が可能です。AWS Neuronを使う場合はモデルのコンパイルが必要なのですが、これも自動で行われます。学習の間に動的にコンパイルを実行し、学習をしてくれます。この動的なコンパイルは長くて30分ほどかかりますが、学習自体が速いため総合的な時間は短縮されるでしょう。

作成されるモデルはコンパイル済みのものではありません。AWS Neuronというとそれ専用になるイメージがありますが、大丈夫です。

スクリプトの実行時は torchrun というコマンドを使います。

torchrun --nproc_per_node=2 train_model.py

公式のサンプルノートブックがおすすめ

こちらに公式で公開されているノートブックが用意されています。やりたいことに合わせて参考にするのがおすすめです!

github.com

ただNotebookであれこれ試すとなるとNeuronのコンパイルが頻繁に走って面倒な場合があります。個人的には手元のGPUなどでしっかりチューニングしてから、AWS Neuronで学習させるのが良いと思います。データセットの準備なども手元で済ませるほうが良いです(これについては別途記事に)。

SageMakerでは勝手にtorchrunをしてくれる

SageMakerでTrn1を使う場合は、自動で train_model.pytorchrun で実行してくれます。ですのでやることといえばインポートを切り替えるくらいなのでとても便利です。

モデルのコンパイル時間にどれくらいかかるか?などもモデルによりけりなので、ぜひ一度試してみると良いでしょう。インスタンスタイプを変えるだけで、勝手にSageMakerがよしなにやってくれます。

他にも知っておくと良いかも、ということ

AWS Neuronでの学習について詳しく知りたいとき

Optimun Neuronを使うと良い感じにラップされすぎていて、逆に何をしているのか分かりにくいときがあります。
pytorh-xlaを使ってNeuron上でTrainerを使った学習のサンプルがAWSの公式であるのでこちらもおすすめです。

aws.amazon.com

この記事、以前のインタビューでもご一緒した常世さんが書いてらっしゃいます。

trn1.32xlargeでマルチノードを使う場合は設定が必要

Trn1でマルチノードで学習させる際は、EFAの設定が必要になります。SecurityGroupで所属するマシン同士ですべてのInbound、Outboundを許可したり、専用のネットワークインタフェースを作成する必要があります。EC2の場合、手動で準備する必要がありかなり手間は大きいです。以下のドキュメントでコマンドによる作成方法などが書かれていますのでこちらを参考にしてください。

awsdocs-neuron.readthedocs-hosted.com

設定がうまくいっているかチェックするコマンドがとても役立ちます。ネットワークのPing Pongも確認してくれるので、SecurityGroupの設定が正しいかなども確認できます(実際にこのコマンドのおかげで間違いに気づいたこともありました)。

チェックポイントから再開させるとバグ

Optimun Neuronのバージョンによってはチェックポイントから再開する際にバグが発生する可能性があります。細かいことは次回の記事で記載いたします。

おわりに

AWS NeuronはTrnもInfも情報が少ないですが、公式できちんと整備されたドキュメントやサンプルはたくさんあります。ぜひお手元のモデルでの学習に活用してみてください。ShodoではAWS Neuronも最大限活用し、AIを使ったサービス運営を行っています。

現在、Shodoではアドベントカレンダー応援クーポンを配布しております。80%オフでShodoを最長3ヶ月間使えるクーポンです。以下のクーポンコードをご購入時に入力して、このアドベントカレンダーの季節にShodoのAI校正をブログの執筆にお役立てください。

XMAS2024

shodo.ink

執筆:@hirokiky
Shodoで執筆されました

SageMakerでのトレーニング時は事前にデータをエンコードしておく

SageMakerでトレーニングする際、強めのマシンを使う場合は事前にデータをエンコードしておくと良いです。データ量が多いと dataset.map(encoder) をするのも案外時間がかかります。そうするとGPUやTrnを有効活用していない時間も課金されてしまうので、事前にエンコードしたものを使いましょう。とくにSpotInstanceを使って繰り返しSageMakerのサーバーを起動する場合はやっておくことをおすすめします。

事前にデータをエンコードする

事前にエンコードするには、以下のように普段通りのエンコードし、 dataset.save_to_disk() を呼び出します。

dataset = load_dataset(...)
encoder = Encoder()
dataset = dataset.map(encoder)
dataset = dataset.remove_columns([...])
dataset.save_to_disk("data/mydataset_enc/")

読み込む際は load_from_disk() を使います。

from datasets import load_from_disk

load_from_disk("data/mydataset_enc/")

このフォルダーをS3にアップロードしておきましょう。

SageMakerでエンコードしたデータセットを使う

SageMakerでの学習を開始する際に、先ほどアップロードしたS3のパスを指定します。

from sagemaker.huggingface import HuggingFace

estimator = HuggingFace(
    entry_point="train_model.py",
)
estimator.fit({
    "train": "s3://my-dataset/mydataset_enc/",
})

このように指定するとSageMakerが実行前に自動でこのフォルダーをダウンロードしてくれます。
SageMakerの内部で実行する train_model.py では以下のようにファイルを読み込みましょう。

dataset = load_from_disk(os.environ["SM_CHANNEL_TRAIN"])

以上です!

これでSageMakerを起動するたびに大規模なデータセットエンコードを待つ必要がありません。

おわりに

Shodoでは独自のAIモデルを使ってAI校正のサービスを提供しています。

現在、Shodoではアドベントカレンダー応援クーポンを配布しております。80%オフでShodoを最長3ヶ月間使えるクーポンです。以下のクーポンコードをご購入時に入力して、このアドベントカレンダーの季節にShodoのAI校正をブログの執筆にお役立てください。

XMAS2024

shodo.ink

執筆:@hirokiky
Shodoで執筆されました

HuggingFace Trainerで学習とハイパーパラメーターチューニングを両立するプログラムを書く

HuggingFace Trainer を使っていて、トレーニング用のプログラムとハイパーパラメーターチューニング用のプログラムを両立させる方法を考えてみました。

HuggingFace Trainerでは trainer.hyperparameter_search(...) メソッドを使ってハイパーパラメーターのチューニングが可能です。たとえばOptunaを使う場合も backend="optuna" と指定するだけで済み、Pruneの設定なども簡単に行えます。

レーニングもハイパーパラメーターチューニングも Trainer(...) を使いますので、ある程度処理を共通化したいところです。ですがいくつか違う点があるのでその差分をうまく吸収するコードを書いていきましょう。

たとえばこんな点があります:

  • .hyperparameter_search() 時はモデルを model= で渡すのではなく、 model_init= に関数を渡す
  • ハイパーパラメーターチューニング時に保存やevalの実行頻度を変える
    • save_strategy="no" を指定してモデルを保存しない(学習時は save_strategy="epoch"
    • eval_strategy="epoch" を指定(学習時は eval_strategy="steps"

レーニング用とチューニング用で再利用する

以下のように、 Trainer を作る関数を用意します。この関数の結果に trainer.train_model() するか trainer.hyperparameter_search() を呼ぶことになります。

def model_init():
    return AutoModel.from_pretrained(...)


def trainer_builder(
    epochs,
    output_dir,
    learning_rate,
    model_init=False,
    **trainer_arguments,
):
    # データセットの用意
    dataset = load_dataset(...)

    train_data = ...
    val_data = ...

    if model_init:
        model_kwargs = {"model_init": model_init, "model": False}
    else:
        model_kwargs = {"model": model_init()}

    args = TrainerArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        learning_rate=learning_rate,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        **trainer_kwargs,
    )

    return Trainer(
        **model_kwargs,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
    )

レーニング用のプログラムでは以下のように trainer_builder() を呼びだします。
コマンドの引数をうまく取ったり、そのまま受け渡したりしてハイパーパラメーターチューニング時との差分を埋めています。

@click.command()
@click.option("--epoch", type=int, default=32)
@click.option("--output_dir", default="output")
@click.option("--learning_rate", type=float, default=1e-05)
@click.option("--resume_ckpt", default=False, is_flag=True)
def main(resume_ckpt, **cmd_kwargs):
    trainer = trainer_builder(
        save_strategy="epoch",
        eval_strategy="steps",
        eval_steps=1000,
        **cmd_kwargs,
    )
    trainer.train(resume_from_checkpoint=resume_ckpt)

他にもコマンドの引数として受け取りたいパラメーターがあれば随時追加すればOKです。

ハイパーパラメーターチューニングは以下のようにします。

def hp(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-05, 1e-04),
    }


def main():
    trainer = trainer_builder(
        epochs=10,
        output_dir="output",
        learning_rate=1e-05,  # どのみち上書きされる
        batch_size=32,
        model_init=True,
        save_strategy="no",
        eval_strategy="epoch",
        report_to=[],
    )
    result = trainer.hyperparameter_search(
        hp,
        direction="minimize",
        n_traials=10,
        backend="optuna",
        pruner=optuna.pruners.MedianPruner(),
    )
    pprint(result)

これでトレーニング用とチューニング用でコードをある程度共通化できます。

注釈

この書き方が万全という自信はなく、まだ実験中なので他に良い書き方があれば教えてください。たとえばTrainerを共通化するのではなく、データセットの準備を共通化すれば済むなどの考え方は他にもありそうです。また、トレーニング用のコードでは model_init の代わりに model 引数を使っていますが、 前者で済むのであれば切り替えは不要です(.hyperparameter_search() の場合は model_init が必須です)。

おわりに

Shodoではモデルのファインチューニングを行っており、自社でAIを運用しております。モデルのハイパーパラメーターチューニングや、継続的な開発をしやすいプログラムが重要になりますので、こういった工夫が重要になります。ぜひほかにも「こうしたほうが良いよ!」等のアドバイスがあればぜひ教えていただけると嬉しいです。

また現在、Shodoではアドベントカレンダー応援クーポンを配布しております。80%オフでShodoを最長3ヶ月間使えるクーポンです。以下のクーポンコードをご購入時に入力して、このアドベントカレンダーの季節にShodoのAI校正をブログの執筆にお役立てください。

XMAS2024

shodo.ink

執筆:@hirokiky
Shodoで執筆されました

SageMakerでのトレーニング実行前に FileNotFoundError: [Errno 2] No such file or directory: 'train'

SageMakerでトレーニング中に No such file or directory: 'train' というエラーが発生しました。

Traceback (most recent call last):
  File "/usr/local/bin/dockerd-entrypoint.py", line 28, in <module>
    subprocess.check_call(shlex.split(" ".join(sys.argv[1:])))
  File "/opt/conda/lib/python3.10/subprocess.py", line 364, in check_call
    retcode = call(*popenargs, **kwargs)
  File "/opt/conda/lib/python3.10/subprocess.py", line 345, in call
    with Popen(*popenargs, **kwargs) as p:
  File "/opt/conda/lib/python3.10/subprocess.py", line 971, in __init__
    self._execute_child(args, executable, preexec_fn, close_fds,
  File "/opt/conda/lib/python3.10/subprocess.py", line 1863, in _execute_child
    raise child_exception_type(errno_num, err_msg, err_filename)

2024-11-01T14:50:11.898Z
FileNotFoundError: [Errno 2] No such file or directory: 'train'

見たところ私自身のプログラムが原因ではないようでした。
エラーが発生したときの実行ログはこちらです。

Starting
Preparing the instances for training
Downloading
Downloading the training image
Training
Training image download completed. Training in progress.
Uploading
Uploading generated training model
Failed
Training job failed

問題の原因と解決方法

SageMakerでトレーニング時に指定するDockerイメージに誤りがありました。
training 用のコンテナーを指定すべきところ、inference 用にしてしまっていました。
このトレーニングではDockerイメージを直接指定しています。

763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference-neuronx:2.1.2-neuronx-py310-sdk2.20.1-ubuntu20.04

763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training-neuronx:2.1.2-neuronx-py310-sdk2.20.1-ubuntu20.04

単なるうっかり間違いでしたが、エラーの特定には時間がかかりました。

問題発見の糸口

エラーの表示から自身のプログラムが起動されていないことは分かったので、インポート時のエラーかDockerの起動周りで問題があることは想像できました。そこでDockerのイメージから内容を探りました。

SageMakerで使われるDockerイメージは deep-learning-containersというリポジトリーで管理されています。検索するとエラーの根本になっていた dockerd-entrypoint.pyが見つかったので、さらにそれの呼び出し元を探したところ以下にありました。

Dockerfileのエントリーポイントとして設定されています。

ENTRYPOINT ["python", "/usr/local/bin/dockerd-entrypoint.py"]

https://github.com/aws/deep-learning-containers/blob/f97413f4b6833cea3e68ad3c7b9b67adcf72bc4a/huggingface/pytorch/inference/docker/1.13/py3/sdk2.14.1/Dockerfile.neuronx#L187

結果、CMDに渡される引数を確認する必要があると分かり、 sagemaker-training-toolkit というSageMaker側で動いているプログラムを検索しました。

docker run で検索すると以下のテストがあり、sagemaker-training-toolkit からDockerがこう呼ばれ、エラーが発生したと予想できました。

        command = (
            "docker run --name sagemaker-training-toolkit-test "
            "sagemaker-training-toolkit-test:dummy train"
        )

https://github.com/aws/sagemaker-training-toolkit/blob/54721512597e77530a68e14c35a857be2f5c5687/test/integration/local/test_dummy.py#L27

そこでDocker Imageで同じ動作(docker run ... image-name train)を再現しました。

docker pull 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference-neuronx:2.1.2-neuronx-py310-sdk2.20.1-ubuntu20.04
docker run --rm 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-inference-neuronx:2.1.2-neuronx-py310-sdk2.20.1-ubuntu20.04 train

すると同じエラーが発生しました。

Traceback (most recent call last):
  File "/usr/local/bin/dockerd-entrypoint.py", line 28, in <module>
    subprocess.check_call(shlex.split(" ".join(sys.argv[1:])))
  File "/opt/conda/lib/python3.10/subprocess.py", line 364, in check_call
    retcode = call(*popenargs, **kwargs)
  File "/opt/conda/lib/python3.10/subprocess.py", line 345, in call
    with Popen(*popenargs, **kwargs) as p:
  File "/opt/conda/lib/python3.10/subprocess.py", line 971, in __init__
    self._execute_child(args, executable, preexec_fn, close_fds,
  File "/opt/conda/lib/python3.10/subprocess.py", line 1863, in _execute_child
    raise child_exception_type(errno_num, err_msg, err_filename)
FileNotFoundError: [Errno 2] No such file or directory: 'train'

これでDockerイメージの相性に間違いがあると分かったので、イメージを再確認したところ training 用ではなかったと分かりました。

おわりに

今回は単にうっかりミスでしたが、深いところまで調べる必要がありました。さすがのChatGPTもこの問題については解決できないようでしたし、少し勝った気持ちになり嬉しかったです。

SageMakerをフル活用しているとAWS提供のDeepLearningコンテナーやツールキットも知っていく必要があります。
かなり色々なところをケアしてくれる良いサービスですし拡張性もありますが、拡張やカスタムコンテナーなどの領域に進んでいく際は各リポジトリーについても知っておくことをおすすめします。

SageMakerとNeuron(inf1/inf2/trn1)については語り尽くせないのですが、ぜひ少しでも参考になれば幸いです。 昨今色々と騒がれていますが、これが本当のAIだよと伝えたいです。

執筆:@hirokiky
Shodoで執筆されました