Amazon ECS の GPU 共有でコンテナごとのメモリ・演算の上限が守られるか試してみた

AWSTemplateFormatVersion: "2010-09-09"
Description: Amazon ECS GPU sharing (NVIDIA MPS) benchmark on g5.xlarge

Parameters:
  VpcId:
    Type: AWS::EC2::VPC::Id
  SubnetIds:
    Type: List<AWS::EC2::Subnet::Id>
  DesiredCapacity:
    Type: Number
    AllowedValues: [0, 1]
    Default: 0
  BenchImage:
    Type: String
    Default: pytorch/pytorch:2.14.1-cuda12.6-cudnn9-runtime
  GpuAmiId:
    Type: AWS::SSM::Parameter::Value<AWS::EC2::Image::Id>
    Default: /aws/service/ecs/optimized-ami/amazon-linux-2023/gpu/recommended/image_id

Resources:
  Cluster:
    Type: AWS::ECS::Cluster
    Properties:
      ClusterName: mps-bench
      ClusterSettings:
        - Name: containerInsights
          Value: enhanced

  LogGroup:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: /ecs/mps-bench
      RetentionInDays: 14

  BenchScript:
    Type: AWS::SSM::Parameter
    Properties:
      Name: /mps-bench/bench.py
      Type: String
      Value: |
        import json, os, time, urllib.request

        MODE = os.environ.get("MODE", "info")
        KEYS = ["CUDA_MPS_PINNED_DEVICE_MEM_LIMIT", "CUDA_MPS_ACTIVE_THREAD_PERCENTAGE",
                "CUDA_MPS_PIPE_DIRECTORY", "NVIDIA_VISIBLE_DEVICES"]
        ORIG = {k: os.environ.get(k) for k in KEYS}


        def emit(**kw):
            kw.update(run=os.environ.get("RUN_ID", ""), mode=MODE, ts=round(time.time(), 3))
            print(json.dumps(kw), flush=True)


        # 上限の強制のされ方を確かめるため、CUDA の初期化前に環境変数を書き換える
        BYPASS = os.environ.get("BYPASS", "")
        if BYPASS == "unset":
            for k in KEYS[:2]:
                os.environ.pop(k, None)
        elif BYPASS == "raise":
            os.environ["CUDA_MPS_PINNED_DEVICE_MEM_LIMIT"] = "0=20G"
            os.environ["CUDA_MPS_ACTIVE_THREAD_PERCENTAGE"] = "100"

        try:
            url = os.environ["ECS_CONTAINER_METADATA_URI_V4"] + "/task"
            TASK = json.load(urllib.request.urlopen(url, timeout=2))["TaskARN"].split("/")[-1]
        except Exception:
            TASK = None

        # 全員が揃ってから計測を始めるため、指定時刻まで待つ
        start_at = float(os.environ.get("START_AT", "0"))
        if start_at > time.time():
            time.sleep(min(start_at - time.time(), 300))

        import torch

        p = torch.cuda.get_device_properties(0)
        free, total = torch.cuda.mem_get_info()
        emit(event="info", task=TASK, bypass=BYPASS, env=ORIG,
             effective={k: os.environ.get(k) for k in KEYS[:2]},
             gpu=p.name, uuid=str(getattr(p, "uuid", "")), sms=p.multi_processor_count,
             total_mib=total >> 20, free_mib=free >> 20,
             torch=torch.__version__, cuda=torch.version.cuda)

        if MODE == "mem":
            chunk = int(os.environ.get("CHUNK_MIB", "256"))
            cap = min(int(os.environ.get("MAX_MIB", "24000")), 24000)
            bufs, err = [], None
            while len(bufs) * chunk < cap:
                try:
                    bufs.append(torch.empty(chunk << 20, dtype=torch.uint8, device="cuda"))
                except Exception as e:
                    err = f"{type(e).__name__}: {str(e)[:400]}"
                    break
            torch.cuda.synchronize()
            emit(event="mem", allocated_mib=len(bufs) * chunk, error=err)
            time.sleep(min(float(os.environ.get("HOLD_SEC", "0")), 300))

        elif MODE == "compute":
            n = int(os.environ.get("N", "8192"))
            dur = min(float(os.environ.get("DURATION", "60")), 600)
            win = float(os.environ.get("WINDOW", "5"))
            a = torch.randn(n, n, device="cuda", dtype=torch.float16)
            b = torch.randn_like(a)
            for _ in range(3):
                a @ b
            torch.cuda.synchronize()
            t_start, total_it = time.time(), 0
            while time.time() - t_start < dur and total_it < 200000:
                w0, it = time.time(), 0
                while time.time() - w0 < win:
                    for _ in range(10):
                        a @ b
                    torch.cuda.synchronize()
                    it += 10
                w1 = time.time()
                total_it += it
                emit(event="window", t0=round(w0, 3), t1=round(w1, 3),
                     tflops=round(2 * n**3 * it / (w1 - w0) / 1e12, 2))
            emit(event="compute_done", iters=total_it,
                 tflops=round(2 * n**3 * total_it / (time.time() - t_start) / 1e12, 2))

        elif MODE == "crash":
            # 範囲外のインデックスで GPU 側のアサーションを発生させる
            x = torch.zeros(8, device="cuda")
            try:
                x[torch.tensor([1 << 20], device="cuda")].sum().item()
                torch.cuda.synchronize()
                emit(event="crash", error=None)
            except Exception as e:
                emit(event="crash", error=f"{type(e).__name__}: {str(e)[:400]}")

  InstanceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal: { Service: ec2.amazonaws.com }
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role
        - arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore

  InstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Roles: [!Ref InstanceRole]

  ExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal: { Service: ecs-tasks.amazonaws.com }
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonECSTaskExecutionRolePolicy
      Policies:
        - PolicyName: read-bench-script
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: ssm:GetParameters
                Resource: !Sub arn:aws:ssm:${AWS::Region}:${AWS::AccountId}:parameter/mps-bench/*

  InstanceSecurityGroup:
    Type: AWS::EC2::SecurityGroup
    Properties:
      GroupDescription: mps-bench instances (no inbound)
      VpcId: !Ref VpcId

  LaunchTemplate:
    Type: AWS::EC2::LaunchTemplate
    Properties:
      LaunchTemplateData:
        ImageId: !Ref GpuAmiId
        InstanceType: g5.xlarge
        IamInstanceProfile: { Arn: !GetAtt InstanceProfile.Arn }
        SecurityGroupIds: [!Ref InstanceSecurityGroup]
        MetadataOptions: { HttpTokens: required, HttpPutResponseHopLimit: 2 }
        BlockDeviceMappings:
          - DeviceName: /dev/xvda
            Ebs: { VolumeSize: 80, VolumeType: gp3, DeleteOnTermination: true }
        UserData:
          Fn::Base64: !Sub |
            #!/bin/bash
            cat >> /etc/ecs/ecs.config <<EOF
            ECS_CLUSTER=${Cluster}
            ECS_ENABLE_GPU_SUPPORT=true
            ECS_IMAGE_PULL_BEHAVIOR=prefer-cached
            EOF
            systemctl enable --now amazon-ssm-agent || (dnf install -y amazon-ssm-agent && systemctl enable --now amazon-ssm-agent)

  AutoScalingGroup:
    Type: AWS::AutoScaling::AutoScalingGroup
    Properties:
      MinSize: "0"
      MaxSize: "1"
      DesiredCapacity: !Ref DesiredCapacity
      VPCZoneIdentifier: !Ref SubnetIds
      LaunchTemplate:
        LaunchTemplateId: !Ref LaunchTemplate
        Version: !GetAtt LaunchTemplate.LatestVersionNumber
      Tags:
        - Key: Name
          Value: mps-bench
          PropagateAtLaunch: true

  TdDedicated:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: mps-bench-dedicated
      RequiresCompatibilities: [EC2]
      NetworkMode: bridge
      ExecutionRoleArn: !GetAtt ExecutionRole.Arn
      ContainerDefinitions:
        - Name: bench
          Image: !Ref BenchImage
          Essential: true
          Cpu: 512
          Memory: 3072
          Command: ["sh", "-c", "printf '%s' \"$BENCH_SCRIPT\" > /tmp/bench.py && exec python /tmp/bench.py"]
          Secrets:
            - Name: BENCH_SCRIPT
              ValueFrom: !Ref BenchScript
          ResourceRequirements:
            - Type: GPU
              Value: "1"
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: !Ref LogGroup
              awslogs-region: !Ref AWS::Region
              awslogs-stream-prefix: mps-bench-dedicated

  TdMps25:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: mps-bench-m4096-c25
      RequiresCompatibilities: [EC2]
      NetworkMode: bridge
      ExecutionRoleArn: !GetAtt ExecutionRole.Arn
      ContainerDefinitions:
        - Name: bench
          Image: !Ref BenchImage
          Essential: true
          Cpu: 512
          Memory: 3072
          Command: ["sh", "-c", "printf '%s' \"$BENCH_SCRIPT\" > /tmp/bench.py && exec python /tmp/bench.py"]
          Secrets:
            - Name: BENCH_SCRIPT
              ValueFrom: !Ref BenchScript
          ResourceRequirements:
            - Type: GPU
              SharingStrategy:
                NvidiaMps:
                  Memory: 4096
                  MaxComputePercent: 25
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: !Ref LogGroup
              awslogs-region: !Ref AWS::Region
              awslogs-stream-prefix: mps-bench-m4096-c25

  TdMps50:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: mps-bench-m4096-c50
      RequiresCompatibilities: [EC2]
      NetworkMode: bridge
      ExecutionRoleArn: !GetAtt ExecutionRole.Arn
      ContainerDefinitions:
        - Name: bench
          Image: !Ref BenchImage
          Essential: true
          Cpu: 512
          Memory: 3072
          Command: ["sh", "-c", "printf '%s' \"$BENCH_SCRIPT\" > /tmp/bench.py && exec python /tmp/bench.py"]
          Secrets:
            - Name: BENCH_SCRIPT
              ValueFrom: !Ref BenchScript
          ResourceRequirements:
            - Type: GPU
              SharingStrategy:
                NvidiaMps:
                  Memory: 4096
                  MaxComputePercent: 50
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: !Ref LogGroup
              awslogs-region: !Ref AWS::Region
              awslogs-stream-prefix: mps-bench-m4096-c50

  TdMps75:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: mps-bench-m4096-c75
      RequiresCompatibilities: [EC2]
      NetworkMode: bridge
      ExecutionRoleArn: !GetAtt ExecutionRole.Arn
      ContainerDefinitions:
        - Name: bench
          Image: !Ref BenchImage
          Essential: true
          Cpu: 512
          Memory: 3072
          Command: ["sh", "-c", "printf '%s' \"$BENCH_SCRIPT\" > /tmp/bench.py && exec python /tmp/bench.py"]
          Secrets:
            - Name: BENCH_SCRIPT
              ValueFrom: !Ref BenchScript
          ResourceRequirements:
            - Type: GPU
              SharingStrategy:
                NvidiaMps:
                  Memory: 4096
                  MaxComputePercent: 75
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: !Ref LogGroup
              awslogs-region: !Ref AWS::Region
              awslogs-stream-prefix: mps-bench-m4096-c75

  TdMps100:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: mps-bench-m4096
      RequiresCompatibilities: [EC2]
      NetworkMode: bridge
      ExecutionRoleArn: !GetAtt ExecutionRole.Arn
      ContainerDefinitions:
        - Name: bench
          Image: !Ref BenchImage
          Essential: true
          Cpu: 512
          Memory: 3072
          Command: ["sh", "-c", "printf '%s' \"$BENCH_SCRIPT\" > /tmp/bench.py && exec python /tmp/bench.py"]
          Secrets:
            - Name: BENCH_SCRIPT
              ValueFrom: !Ref BenchScript
          ResourceRequirements:
            - Type: GPU
              SharingStrategy:
                NvidiaMps:
                  Memory: 4096
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: !Ref LogGroup
              awslogs-region: !Ref AWS::Region
              awslogs-stream-prefix: mps-bench-m4096

  TdMpsBig:
    Type: AWS::ECS::TaskDefinition
    Properties:
      Family: mps-bench-m20480
      RequiresCompatibilities: [EC2]
      NetworkMode: bridge
      ExecutionRoleArn: !GetAtt ExecutionRole.Arn
      ContainerDefinitions:
        - Name: bench
          Image: !Ref BenchImage
          Essential: true
          Cpu: 512
          Memory: 3072
          Command: ["sh", "-c", "printf '%s' \"$BENCH_SCRIPT\" > /tmp/bench.py && exec python /tmp/bench.py"]
          Secrets:
            - Name: BENCH_SCRIPT
              ValueFrom: !Ref BenchScript
          ResourceRequirements:
            - Type: GPU
              SharingStrategy:
                NvidiaMps:
                  Memory: 20480
          LogConfiguration:
            LogDriver: awslogs
            Options:
              awslogs-group: !Ref LogGroup
              awslogs-region: !Ref AWS::Region
              awslogs-stream-prefix: mps-bench-m20480

Outputs:
  ClusterName:
    Value: !Ref Cluster
  AutoScalingGroupName:
    Value: !Ref AutoScalingGroup

#!/usr/bin/env bash
# mps-bench のタスクを起動し、終了後にログの JSON 行を集めて保存する。
# 使い方: run.sh <ラベル> '<ファミリー>|K=V|K=V' ['<ファミリー>|K=V' ...] [--delay 秒] [--no-wait]
# 同じ呼び出しで指定したタスクには同じ START_AT を渡し、計測の開始を揃える。
set -eo pipefail

CLUSTER=mps-bench
LOG_GROUP=/ecs/mps-bench
OUT="$(cd "$(dirname "$0")" && pwd)/results"

label=$1; shift
delay=45 wait=true specs=()
while (($#)); do
  case $1 in
    --delay) delay=$2; shift 2 ;;
    --no-wait) wait=false; shift ;;
    *) specs+=("$1"); shift ;;
  esac
done
start_at=$(($(date +%s) + delay))
mkdir -p "$OUT"
save() { printf '%s\n' "$rec" > "$OUT/$label.json"; }

tasks='[]' failures='[]'
for i in "${!specs[@]}"; do
  IFS='|' read -ra parts <<<"${specs[$i]}"
  family=${parts[0]}
  env=$(jq -n --arg run "$label-$i" --arg at "$start_at" '{RUN_ID: $run, START_AT: $at}')
  for kv in "${parts[@]:1}"; do
    env=$(jq --arg k "${kv%%=*}" --arg v "${kv#*=}" '.[$k] = $v' <<<"$env")
  done
  overrides=$(jq -c '{containerOverrides: [{name: "bench",
    environment: (to_entries | map({name: .key, value: .value}))}]}' <<<"$env")
  r=$(aws ecs run-task --cluster "$CLUSTER" --task-definition "$family" \
    --launch-type EC2 --count 1 --overrides "$overrides" --output json)
  failures=$(jq -c --argjson r "$r" '. + $r.failures' <<<"$failures")
  tasks=$(jq -c --argjson r "$r" --arg f "$family" --argjson env "$env" \
    '. + [$r.tasks[] | {arn: .taskArn, family: $f, env: $env}]' <<<"$tasks")
done
rec=$(jq -n --arg label "$label" --argjson at "$start_at" \
  --argjson tasks "$tasks" --argjson failures "$failures" \
  '{label: $label, start_at: $at, tasks: $tasks, failures: $failures}')
jq -c '{launched: (.tasks | length), failures}' <<<"$rec"

arns=$(jq -r '.tasks[].arn' <<<"$rec")
if [[ -z $arns ]]; then save; exit; fi
sleep $((delay < 20 ? delay : 20))
gpu='[]'
ci=$(aws ecs list-container-instances --cluster "$CLUSTER" --query containerInstanceArns --output text)
if [[ -n $ci && $ci != None ]]; then
  gpu=$(aws ecs describe-container-instances --cluster "$CLUSTER" --container-instances $ci \
    --query 'containerInstances[].gpuDevices[]' --output json)
fi
running=$(aws ecs describe-tasks --cluster "$CLUSTER" --tasks $arns --query tasks --output json)
rec=$(jq --argjson g "$gpu" --argjson d "$running" \
  '.gpu_devices_running = $g | .describe_tasks_running = $d' <<<"$rec")
if ! $wait; then save; exit; fi

# 待機は 1 回あたり最大 10 分。2 回で打ち切る
for _ in 1 2; do
  aws ecs wait tasks-stopped --cluster "$CLUSTER" --tasks $arns && break
done
stopped=$(aws ecs describe-tasks --cluster "$CLUSTER" --tasks $arns --query tasks --output json)
n=$(jq '.tasks | length' <<<"$rec")
for ((i = 0; i < n; i++)); do
  arn=$(jq -r ".tasks[$i].arn" <<<"$rec")
  family=$(jq -r ".tasks[$i].family" <<<"$rec")
  lines=$(aws logs filter-log-events --log-group-name "$LOG_GROUP" \
    --log-stream-names "$family/bench/${arn##*/}" \
    --query 'events[].message' --output json 2>/dev/null || echo '[]')
  rec=$(jq --argjson i "$i" --arg arn "$arn" --argjson d "$stopped" --argjson lines "$lines" '
    ($d[] | select(.taskArn == $arn)) as $t | $t.containers[0] as $c
    | [$lines[] | fromjson? | objects] as $json
    | .tasks[$i] += {
        stopped_reason: $t.stoppedReason, exit_code: $c.exitCode, reason: $c.reason,
        gpu_ids: $c.gpuIds, gpu_sharing: $c.gpuSharingStrategy,
        json: $json, other_lines: (($lines | length) - ($json | length)),
        tail: ([$lines[] | select(startswith("{") | not)] | .[-5:])
      }' <<<"$rec")
done
save
jq -c '.tasks[] | {family, exit: .exit_code, gpu: .gpu_ids,
  events: [.json[] | select(.event | IN("info", "mem", "compute_done", "crash"))],
  other_lines}' <<<"$rec"