復元力のある長時間実行 GPU ジョブを設計する

このチュートリアルでは、最大 12 時間実行される、GPU 対応の復元性に優れた Cloud Run ジョブを設計してデプロイする方法について説明します。

Cloud Run は拡張タスク タイムアウトをサポートしていますが、中断のない実行の SLA は提供していません。GPU 対応タスクはライブ マイグレーションされないため、毎週の定期メンテナンスで実行中に終了します。そのため、拡張タスクは少なくとも 1 回は中断されることが予想されます。コンピューティング リソースを浪費することなくジョブを正常に完了するには、1 回の中断のない実行に依存するのではなく、チェックポイントを設定して再開するようにアプリケーションを設計します。

このチュートリアルでは、チェックポイントの実装、プラットフォーム SIGTERM シグナルの処理、最大再試行回数の構成方法について説明します。

目標

このチュートリアルでは、次のことを行います。

  • チェックポイントを実装し、SIGTERM シグナルを処理する Python アプリケーションを準備します。
  • チェックポイントを永続化する Cloud Storage バケットを作成します。
  • GPU を有効にし、タスク タイムアウトを 12 時間に設定し、再試行ポリシーを構成して、Cloud Run ジョブをビルドしてデプロイします。
  • ジョブを実行し、ログでプリエンプション イベントをモニタリングする方法を理解する。

費用

このドキュメントでは、課金対象である次の Google Cloudコンポーネントを使用します。

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。

新規の Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

  1. Google Cloud CLI をインストールします。

  2. フェデレーション ID(連携 ID)を使用するように gcloud CLI を構成します。

    詳細については、 フェデレーション ID を使用して gcloud CLI にログインするをご覧ください。

  3. gcloud CLI を初期化するには、次のコマンドを実行します。

    gcloud init
  4. Google Cloud プロジェクトを作成または選択します。

    プロジェクトの選択または作成に必要なロール

    • プロジェクトを選択する: プロジェクトを選択する際に、特定の IAM ロールは必要ありません。ロールが付与されているプロジェクトを選択できます。
    • プロジェクトを作成する: プロジェクトを作成するには、resourcemanager.projects.create 権限を含むプロジェクト作成者のロール(roles/resourcemanager.projectCreator)が必要です。詳しくは、ロールを付与する方法をご覧ください。
    • Google Cloud プロジェクトを作成します。

      gcloud projects create PROJECT_ID

      PROJECT_ID は、作成する Google Cloud プロジェクトの名前に置き換えます。

    • 作成した Google Cloud プロジェクトを選択します。

      gcloud config set project PROJECT_ID

      PROJECT_ID は、 Google Cloud プロジェクト名に置き換えます。

  5. Google Cloud プロジェクトに対して課金が有効になっていることを確認します。

  6. Cloud Run、Artifact Registry、Cloud Build、Cloud Storage API を有効にします(まだ有効になっていない場合)。

    API を有効にするために必要なロール

    API を有効にするには、serviceusage.services.enable 権限が必要です。プロジェクトを作成した場合は、オーナーロール(roles/owner)を介してこの権限がすでに付与されている可能性があります。それ以外の場合は、Service Usage 管理者ロール(roles/serviceusage.serviceUsageAdmin)を介してこの権限を取得できます。ロールを付与する方法を確認する。

    gcloud services enable run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com storage.googleapis.com
  7. リージョンを環境変数として設定します。
    export REGION=REGION
  8. ジョブのサービス アカウントを作成します。
    gcloud iam service-accounts create gpu-job-sa
  9. 選択したリージョンに [十分な GPU 割り当て](/run/docs/configuring/jobs/gpu#request-quota)(Nvidia L4 など)があることを確認します。

アプリケーションを準備する

長時間実行ジョブの復元性を確保するには、アプリケーションで次の処理を行う必要があります。

  1. 進行状況(チェックポイント)を、マウントされた Cloud Storage バケットなどの永続的な場所に定期的に保存します。

    SIGTERM の後、タスクが強制終了されるまでに約 10 秒かかります。この時間では、大規模なモデルのチェックポイントを書き込むには不十分です。SIGTERM ハンドラは、小さな状態のみをベスト エフォートでフラッシュするものとして扱い、定期的なチェックポイントを主な保護手段として利用します。チェックポイントの間隔は、失っても問題のない作業量に基づいて選択します。チェックポイントの頻度が高いほど、有用な作業に費やされない時間が長くなります。目標は、ユースケースに適したバランスを実現することです。

  2. 起動時に最後のチェックポイントから状態を復元する。

  3. Cloud Run がメンテナンス プロセスの 10 秒前に送信する SIGTERM シグナルをリッスンし、終了前に最終チェックポイントを保存しようとします。

次の例は、これらの機能を実装する方法を示しています。

アプリケーション コードを作成する

  1. プロジェクトのディレクトリを作成します。

    mkdir gc-run-gpu-resilience
    cd gc-run-gpu-resilience
    
  2. task.py という名前のファイルを作成し、次のスクリプトを貼り付けて、チェックポイント処理とシグナル処理をサポートするトレーニングまたはバッチ処理をシミュレートします。

    import os
    import sys
    import time
    import signal
    
    # The checkpoint directory will be mounted to this path
    CHECKPOINT_DIR = "/checkpoints"
    # Cloud Run sets these variables automatically
    JOB_NAME = os.environ.get("CLOUD_RUN_JOB", "local-job")
    TASK_INDEX = os.environ.get("CLOUD_RUN_TASK_INDEX", "0")
    CHECKPOINT_FILE = os.path.join(
        CHECKPOINT_DIR, f"{JOB_NAME}-task-{TASK_INDEX}-checkpoint.txt")
    TOTAL_STEPS = 100
    current_step = 0
    is_exiting = False
    
    def save_checkpoint(step):
        print(f"Saving checkpoint at step {step} to {CHECKPOINT_FILE}...")
        try:
            # Write to a temp file and rename to ensure atomicity
            temp_file = CHECKPOINT_FILE + ".tmp"
            with open(temp_file, "w") as f:
                f.write(str(step))
            os.rename(temp_file, CHECKPOINT_FILE)
            print("Checkpoint saved successfully.")
        except Exception as e:
            print(f"Failed to save checkpoint: {e}")
    
    def load_checkpoint():
        print(f"Looking for a checkpoint at {CHECKPOINT_FILE}...")
        global current_step
        if os.path.exists(CHECKPOINT_FILE):
            try:
                with open(CHECKPOINT_FILE, "r") as f:
                    current_step = int(f.read().strip())
                print(f"Resuming from checkpointed step {current_step}.")
            except Exception as e:
                print(f"Failed to load checkpoint: {e}. Starting from scratch.")
        else:
            print("No checkpoint found. Starting from scratch.")
    
    def sigterm_handler(signum, frame):
        global is_exiting
        print("Received SIGTERM signal from platform. Initiating graceful shutdown...")
        is_exiting = True
    
    # Register the signal handler for SIGTERM
    signal.signal(signal.SIGTERM, sigterm_handler)
    
    # Load existing checkpoint if any to resume work
    load_checkpoint()
    
    # Simulate work loop
    while current_step < TOTAL_STEPS and not is_exiting:
        print(f"Processing step {current_step}/{TOTAL_STEPS}...")
        # Simulate heavy GPU computation
        time.sleep(1)
        current_step += 1
        # Periodically save checkpoint (e.g., every 5 steps)
        if current_step % 5 == 0:
            save_checkpoint(current_step)
    
    if is_exiting:
        # Save final state before exiting due to preemption
        # Best-effort only: with ~10s before SIGKILL this will not complete for
        # large state. Periodic checkpointing (above) is what actually protects
        # you.
        save_checkpoint(current_step)
        # Note that save_checkpoint already writes-then-renames for atomicity,
        # which is good, but the reason it matters here is that a truncated
        # checkpoint from an interrupted final save must not clobber the last
        # good one.
        print("Graceful shutdown complete. Exiting with error code to trigger retry.")
        # Exit with a non-zero code to signal to Cloud Run that the task
        # was interrupted and needs to be retried from the checkpoint.
        sys.exit(143)
        # Note that this statement deliberately consumes of the configured retry
        # attempts.
    
    print("Task completed successfully.")
    sys.exit(0)
    
  3. アプリケーションをコンテナ化する Dockerfile を作成します。

より詳細で現実的な実装については、こちらの例をご覧ください。

チェックポイント用の Cloud Storage バケットを作成する

チェックポイント ファイルを保存する Cloud Storage バケットを作成し、ジョブのサービス アカウントにアクセス権を付与します。

  1. バケットを作成します。

    gcloud storage buckets create gs://gpu-checkpoints-PROJECT_ID \
        --location REGION
    
  2. ジョブのサービス アカウントに Storage オブジェクト管理者ロールを付与します。

    gcloud storage buckets add-iam-policy-binding gs://gpu-checkpoints-PROJECT_ID \
        --member="serviceAccount:gpu-job-sa@PROJECT_ID." \
        --role="roles/storage.objectAdmin"
    

ジョブをデプロイする

ジョブを Cloud Run にデプロイします。

gcloud beta run jobs deploy gpu-resilient-job \
    --source=. \
    --region REGION \
    --gpu 1 \
    --gpu-type nvidia-l4 \
    --no-gpu-zonal-redundancy \
    --task-timeout 168h \
    --max-retries 3 \
    --service-account gpu-job-sa@PROJECT_ID. \
    --add-volume=name=checkpoint-volume,type=cloud-storage,bucket=gpu-checkpoints-PROJECT_ID \
    --add-volume-mount=volume=checkpoint-volume,mount-path=/checkpoints

次の重要なフラグに注意してください。

  • --task-timeout 168h: タスクを最大 12 時間実行できます。
  • --max-retries 3: タイムアウトが 1 時間を超える GPU ジョブの場合は、デフォルト構成の 3 回の再試行を設定することをおすすめします。プラットフォームのメンテナンス中にサイレント エラーが発生しないように、再試行回数を 0 に設定しないことをおすすめします。最後に、メンテナンス関連の障害に対して少なくとも 1 回の再試行、12 時間の実行期間中のアプリケーション障害に対して追加の再試行を予算に組み込むことをおすすめします。
  • --add-volume と --add-volume-mount: Cloud Storage バケットを /checkpoints にマウントして、アプリケーションがタスクの再試行間で状態を保持できるようにします。

ジョブを実行してモニタリングする

  1. ジョブを実行します。

    gcloud run jobs execute gpu-resilient-job --region REGION
    

    このコマンドを実行すると、次のようなメッセージが返されます。

    Execution [gpu-resilient-job-abcde] has been successfully started.
    
    View details about this execution by running `gcloud run jobs executions
    describe gpu-resilient-job-abcde`.
    
    See logs for this execution at: https://console.cloud.google.com/logs/viewer?project=YOUR_PROJECT_ID&resource=cloud_run_job/job_name/my-job/location/my-region/execution_id/gpu-resilient-job-abcde
    

    文字列 gpu-resilient-job-abcde は、この特定の実行試行に対して自動生成された一意の ID です。この ID を使用して、この特定の実行をモニタリングまたはキャンセルできます。出力には、 Google Cloud コンソールの [ログ エクスプローラ] ページへの直接リンクも含まれています。このページでは、ジョブタスクの出力やエラーを確認できます。

  2. ログをモニタリングします。メンテナンス イベントが発生すると、Cloud Run は警告をログに記録します。次のように表示されます。

    • アプリケーション ログ: Received SIGTERM signal from platform. Initiating graceful shutdown... の後に Saving checkpoint...。
    • プリエンプションを示すシステムログ: Container instance preemption initiated by the platform due to underlying hardware maintenance.(理由: INSTANCE_PREEMPTED)。
    • Cloud Run がタスクを自動的に再起動する。
    • 再起動時のアプリケーション ログ: Resuming from checkpointed step X.。

クリーンアップ

Google Cloud アカウントで追加料金が発生しないようにするには、このクイックスタートでデプロイしたすべてのリソースを削除します。

リポジトリを削除する

Cloud Run では、ジョブの実行時間に対してのみ料金が発生します。ただし、コンテナ イメージを Artifact Registry に保存した場合にも料金が発生する可能性があります。Artifact Registry リポジトリを削除するには、Artifact Registry ドキュメントのリポジトリを削除するの手順を行います。

ジョブを削除する

Cloud Run ジョブでは、ジョブタスクの実行時にのみ料金が発生します。Cloud Run ジョブを削除するには、次のいずれかの操作を行います。

コンソール

ジョブを削除するには:

  1. Google Cloud コンソールで Cloud Run に移動します。

    Cloud Run に移動

  2. 削除するジョブをジョブリストで探し、そのチェックボックスをクリックして選択します。

  3. [削除] をクリックします。これにより、進行中のすべてのジョブ実行と実行中のすべてのコンテナ インスタンスが終了します。

gcloud

ジョブを削除するには、次のコマンドを実行します。

gcloud run jobs delete JOB_NAME

JOB_NAME は、ジョブ名に置き換えます。

テスト プロジェクトを削除する

Google Cloud プロジェクトを削除すると、そのプロジェクト内のすべてのリソースに対する課金が停止します。プロジェクト内のすべての Google Cloud リソースを解放する手順は次のとおりです。

    Google Cloud プロジェクトを削除する:

    gcloud projects delete PROJECT_ID

Cloud Storage バケットを削除します。

gcloud storage buckets delete gs://gpu-checkpoints-PROJECT_ID

次のステップ