设计弹性长时间运行的 GPU 作业

本教程介绍了如何设计和部署运行时间长达 12 小时、支持 GPU 的弹性 Cloud Run 作业。

虽然 Cloud Run 支持延长任务超时时间,但它不提供不间断执行的 SLA。由于启用了 GPU 的任务不会实时迁移,因此每周例行维护会在任务运行期间终止这些任务。因此,任何长时间运行的任务都应做好至少中断一次的准备。为确保作业成功完成而不浪费计算资源,请将应用设计为可设置检查点并恢复,而不是依赖于单个不间断运行。

本教程将指导您如何实现检查点、处理平台 SIGTERM 信号以及配置最大重试次数。

目标

在本教程中,您将执行以下操作:

  • 准备一个实现检查点并处理 SIGTERM 信号的 Python 应用。
  • 创建一个 Cloud Storage 存储桶来持久保存检查点。
  • 构建并部署一个 Cloud Run 作业,该作业启用了 GPU、任务超时时间为 12 小时,并配置了重试政策。
  • 运行作业,并了解如何在日志中监控抢占事件。

费用

在本文档中,您将使用 Google Cloud的以下收费组件:

如需根据您的预计使用量来估算费用,请使用价格计算器。

新 Google Cloud 用户可能有资格申请免费试用。

准备工作

  1. 登录您的 Google Cloud 账号。如果您是 Google Cloud新手,请 创建一个账号来评估我们的产品在实际场景中的表现。新客户还可获享 $300 赠金,用于运行、测试和部署工作负载。
  2. 安装 Google Cloud CLI。

  3. 如果您使用的是外部身份提供方 (IdP),则必须先 使用联合身份登录 gcloud CLI。

  4. 如需初始化 gcloud CLI,请运行以下命令:

    gcloud init
  5. 创建或选择 Google Cloud 项目。

    选择或创建项目所需的角色

    • 选择项目:选择项目不需要特定的 IAM 角色,您可以选择已获授角色的任何项目。
    • 创建项目:如需创建项目,您需要拥有 Project Creator 角色 (roles/resourcemanager.projectCreator),该角色包含 resourcemanager.projects.create 权限。了解如何授予角色。
    • 创建 Google Cloud 项目:

      gcloud projects create PROJECT_ID

      将 PROJECT_ID 替换为您要创建的 Google Cloud 项目的名称。

    • 选择您创建的 Google Cloud 项目:

      gcloud config set project PROJECT_ID

      将 PROJECT_ID 替换为您的 Google Cloud 项目名称。

  6. 验证是否已为您的 Google Cloud 项目启用结算功能。

  7. 安装 Google Cloud CLI。

  8. 如果您使用的是外部身份提供方 (IdP),则必须先 使用联合身份登录 gcloud CLI。

  9. 如需初始化 gcloud CLI,请运行以下命令:

    gcloud init
  10. 创建或选择 Google Cloud 项目。

    选择或创建项目所需的角色

    • 选择项目:选择项目不需要特定的 IAM 角色,您可以选择已获授角色的任何项目。
    • 创建项目:如需创建项目,您需要拥有 Project Creator 角色 (roles/resourcemanager.projectCreator),该角色包含 resourcemanager.projects.create 权限。了解如何授予角色。
    • 创建 Google Cloud 项目:

      gcloud projects create PROJECT_ID

      将 PROJECT_ID 替换为您要创建的 Google Cloud 项目的名称。

    • 选择您创建的 Google Cloud 项目:

      gcloud config set project PROJECT_ID

      将 PROJECT_ID 替换为您的 Google Cloud 项目名称。

  11. 验证是否已为您的 Google Cloud 项目启用结算功能。

  12. 启用 Cloud Run、Artifact Registry、Cloud Build 和 Cloud Storage API(如果尚未启用):

    启用 API 所需的角色

    如需启用 API,您需要拥有 serviceusage.services.enable 权限。如果您创建了项目,则可能已经通过 Owner 角色 (roles/owner) 拥有此权限。否则,您可以通过 Service Usage Admin 角色 (roles/serviceusage.serviceUsageAdmin) 获得此权限。了解如何授予角色。

    gcloud services enable run.googleapis.com cloudbuild.googleapis.com artifactregistry.googleapis.com storage.googleapis.com
  13. 将区域设置为环境变量:
    export REGION=REGION
  14. 为作业创建服务账号:
    gcloud iam service-accounts create gpu-job-sa
  15. 确保您在所选区域中有 [足够的 GPU 配额](/run/docs/configuring/jobs/gpu#request-quota)(例如 Nvidia L4)。

准备申请材料

为了使长时间运行的作业具有弹性,您的应用必须:

  1. 定期将进度(检查点)保存到永久性位置,例如装载的 Cloud Storage 存储桶。

    在 SIGTERM 之后,您有大约 10 秒的时间来终止任务。这段时间不足以写入大型模型检查点。将 SIGTERM 处理程序视为仅尽力而为地刷新少量状态,并依靠定期检查点作为主要保护。根据您能承受的工作损失量来选择检查点间隔,请注意,检查点越频繁,意味着不执行有用工作的时间越多。目标是针对您的应用场景实现适当的平衡。

  2. 在启动时从上一个检查点恢复状态。

  3. 监听 SIGTERM 信号(Cloud Run 会在维护流程开始前 10 秒发送该信号),并尝试在退出前保存最终检查点。

以下示例说明了如何实现这些功能。

创建应用代码

  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 Object Admin 角色:

    gcloud storage buckets add-iam-policy-binding gs://gpu-checkpoints-PROJECT_ID \
        --member="serviceAccount:gpu-job-sa@PROJECT_ID.iam.gserviceaccount.com" \
        --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.iam.gserviceaccount.com \
    --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。最后,我们建议您在 12 小时的执行期内,为与维护相关的故障至少安排 1 次重试,并为应用故障安排额外的重试。
  • --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。您可以使用它来监控或取消此特定运行。输出还包含一个直接链接,可前往 Google Cloud 控制台中的 Logs Explorer 页面,以查看作业任务的输出或错误。

  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

后续步骤