Resume benchmark fleet jobs after Spot interruptions

benchmarks/cloudformation/benchmark.yaml

@@ -26,9 +26,14 @@ Parameters:
     Description: CPU architecture of the benchmark hosts; it must match the instance type.
   Purchase:
     Type: String
-    Default: on-demand
+    Default: spot
     AllowedValues: [on-demand, spot]
     Description: Purchase option of the benchmark hosts; every host uses the same instance type either way.
+  Zones:
+    Type: Number
+    Default: 4
+    AllowedValues: [3, 4]
+    Description: Availability Zones to spread the hosts over, the first ones that Fn::GetAZs returns; more zones widen Spot capacity.
   InstanceType:
     Type: String
     Default: r7i.2xlarge
@@ -74,6 +79,25 @@ Parameters:
     Type: String
     Default: pypi_file_downloads
     AllowedPattern: '[a-z][a-z0-9_]*'
+  ResultsBucket:
+    Type: String
+    Default: ''
+    AllowedPattern: '([a-z0-9][a-z0-9.-]+[a-z0-9])?'
+    Description: Existing bucket for queue state, reports, and trial outputs under ResultsPrefix/<stack name>/; empty uses the scratch bucket.
+  ResultsPrefix:
+    Type: String
+    Default: pyathena-benchmark
+    AllowedPattern: '[a-zA-Z0-9_-]+(/[a-zA-Z0-9_-]+)*'
+  QueueName:
+    Type: String
+    Default: ''
+    AllowedPattern: '[a-zA-Z0-9._-]*'
+    Description: Fleet queue that every host's worker consumes once it is published; empty starts no worker.
+  ApiSlots:
+    Type: Number
+    Default: 10
+    MinValue: 1
+    Description: Fleet-wide limit of simultaneous GetQueryResults paging queries for the started workers.
 
 Rules:
   X86InstanceType:
@@ -90,6 +114,8 @@ Rules:
 Conditions:
   IsArm64: !Equals [!Ref Architecture, arm64]
   IsSpot: !Equals [!Ref Purchase, spot]
+  UseFourZones: !Equals [!Ref Zones, 4]
+  HasResultsBucket: !Not [!Equals [!Ref ResultsBucket, '']]
 
 Resources:
   Vpc:
@@ -105,7 +131,7 @@ Resources:
     Properties:
       VpcId: !Ref Vpc
       InternetGatewayId: !Ref Gateway
-  # Three Availability Zones give Spot capacity more than one pool.
+  # Several Availability Zones give Spot capacity more than one pool.
   Subnet:
     Type: AWS::EC2::Subnet
     Properties:
@@ -124,6 +150,13 @@ Resources:
       VpcId: !Ref Vpc
       CidrBlock: 10.93.2.0/24
       AvailabilityZone: !Select [2, !GetAZs '']
+  Subnet4:
+    Type: AWS::EC2::Subnet
+    Condition: UseFourZones
+    Properties:
+      VpcId: !Ref Vpc
+      CidrBlock: 10.93.3.0/24
+      AvailabilityZone: !Select [3, !GetAZs '']
   RouteTable:
     Type: AWS::EC2::RouteTable
     Properties:
@@ -150,6 +183,12 @@ Resources:
     Properties:
       SubnetId: !Ref Subnet3
       RouteTableId: !Ref RouteTable
+  RouteAssociation4:
+    Type: AWS::EC2::SubnetRouteTableAssociation
+    Condition: UseFourZones
+    Properties:
+      SubnetId: !Ref Subnet4
+      RouteTableId: !Ref RouteTable
   SecurityGroup:
     Type: AWS::EC2::SecurityGroup
     Properties:
@@ -260,6 +299,22 @@ Resources:
               - Effect: Allow
                 Action: [s3:GetObject, s3:PutObject, s3:DeleteObject, s3:AbortMultipartUpload, s3:ListMultipartUploadParts]
                 Resource: !Sub '${Bucket.Arn}/*'
+              # Queue state, reports, and trial outputs go only below this stack's results prefix.
+              # Listing covers the whole bucket: without s3:ListBucket, reading a missing key
+              # returns 403 instead of 404, which breaks existence checks of fsspec and pyarrow,
+              # and ListBucketMultipartUploads does not support the s3:prefix condition key.
+              - !If
+                - HasResultsBucket
+                - Effect: Allow
+                  Action: [s3:GetBucketLocation, s3:ListBucket, s3:ListBucketMultipartUploads]
+                  Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket}
+                - !Ref AWS::NoValue
+              - !If
+                - HasResultsBucket
+                - Effect: Allow
+                  Action: [s3:GetObject, s3:PutObject, s3:DeleteObject, s3:AbortMultipartUpload, s3:ListMultipartUploadParts]
+                  Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/*
+                - !Ref AWS::NoValue
   InstanceProfile:
     Type: AWS::IAM::InstanceProfile
     Properties:
@@ -271,6 +326,8 @@ Resources:
         MetadataOptions:
           HttpTokens: required
           HttpEndpoint: enabled
+          # A worker reads the KeepHost tag from instance metadata.
+          InstanceMetadataTags: enabled
         ImageId: !If [IsArm64, !Ref Arm64ImageId, !Ref ImageId]
         InstanceType: !Ref InstanceType
         IamInstanceProfile:
@@ -292,22 +349,48 @@ Resources:
               - Key: Name
                 Value: !Sub '${AWS::StackName}-benchmark'
         UserData:
-          Fn::Base64: !Sub |
-            #!/bin/bash
-            set -euo pipefail
-            dnf install -y git aws-cfn-bootstrap
-            trap '/opt/aws/bin/cfn-signal -e $? --stack ${AWS::StackName} --resource Fleet --region ${AWS::Region}' EXIT
-            curl --proto '=https' --tlsv1.2 -LsSf https://astral.sh/uv/${UvVersion}/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh
-            git clone https://github.com/pyathena-dev/PyAthena.git /opt/pyathena
-            git -C /opt/pyathena fetch origin ${GitCommit}
-            git -C /opt/pyathena checkout --detach ${GitCommit}
-            export UV_PYTHON_INSTALL_DIR=/opt/uv-python
-            uv python install ${PythonVersion}
-            uv sync --directory /opt/pyathena/benchmarks --locked --no-dev --python ${PythonVersion}
-            chmod -R a+rX /opt/uv-python
-            chown -R ec2-user:ec2-user /opt/pyathena
-            printf '%s\n' '${AWS::StackId}' > /opt/pyathena/benchmarks/stack-id.txt
-            printf '%s\n' 'Setup complete. Run preflight and prepare manually.'
+          Fn::Base64: !Sub
+            - |
+              #!/bin/bash
+              set -euo pipefail
+              dnf install -y git aws-cfn-bootstrap
+              trap '/opt/aws/bin/cfn-signal -e $? --stack ${AWS::StackName} --resource Fleet --region ${AWS::Region}' EXIT
+              curl --proto '=https' --tlsv1.2 -LsSf https://astral.sh/uv/${UvVersion}/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh
+              git clone https://github.com/pyathena-dev/PyAthena.git /opt/pyathena
+              git -C /opt/pyathena fetch origin ${GitCommit}
+              git -C /opt/pyathena checkout --detach ${GitCommit}
+              export UV_PYTHON_INSTALL_DIR=/opt/uv-python
+              uv python install ${PythonVersion}
+              uv sync --directory /opt/pyathena/benchmarks --locked --no-dev --python ${PythonVersion}
+              chmod -R a+rX /opt/uv-python
+              chown -R ec2-user:ec2-user /opt/pyathena
+              printf '%s\n' '${AWS::StackId}' > /opt/pyathena/benchmarks/stack-id.txt
+              printf '%s\n' '${Results}' > /opt/pyathena/benchmarks/results-location.txt
+              sudo -u ec2-user mkdir -p /opt/pyathena/benchmarks/results
+              if [ -n '${QueueName}' ]; then
+              cat > /etc/systemd/system/pyathena-bench-worker.service <<'UNIT'
+              [Unit]
+              Description=PyAthena benchmark fleet worker
+              Wants=network-online.target
+              After=network-online.target
+              [Service]
+              User=ec2-user
+              WorkingDirectory=/opt/pyathena/benchmarks
+              Environment=AWS_DEFAULT_REGION=${AWS::Region}
+              ExecStart=/opt/pyathena/.venv/bin/python -m pyathena_bench worker --results ${Results} --name ${QueueName} --api-slots ${ApiSlots} --wait --terminate-host
+              StandardOutput=append:/opt/pyathena/benchmarks/results/worker.log
+              StandardError=inherit
+              [Install]
+              WantedBy=multi-user.target
+              UNIT
+              systemctl daemon-reload
+              systemctl enable --now pyathena-bench-worker.service
+              fi
+              printf '%s\n' 'Setup complete. Run preflight and prepare manually.'
+            - Results: !If
+                - HasResultsBucket
+                - !Sub s3://${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/
+                - !Sub s3://${Bucket}/
   Fleet:
     Type: AWS::AutoScaling::AutoScalingGroup
     DependsOn: [Route, RouteAssociation, RouteAssociation2, RouteAssociation3]
@@ -320,8 +403,17 @@ Resources:
       MinSize: 0
       MaxSize: !Ref FleetSize
       DesiredCapacity: !Ref FleetSize
-      VPCZoneIdentifier: [!Ref Subnet, !Ref Subnet2, !Ref Subnet3]
+      VPCZoneIdentifier: !If
+        - UseFourZones
+        - [!Ref Subnet, !Ref Subnet2, !Ref Subnet3, !Ref Subnet4]
+        - [!Ref Subnet, !Ref Subnet2, !Ref Subnet3]
       HealthCheckType: EC2
+      Tags:
+        # DependsOn cannot name the conditional fourth zone's route association; this
+        # reference makes hosts launch only after it exists, so bootstrap reaches the Internet.
+        - Key: pyathena-benchmark-route
+          Value: !If [UseFourZones, !Ref RouteAssociation4, none]
+          PropagateAtLaunch: false
       MixedInstancesPolicy:
         LaunchTemplate:
           LaunchTemplateSpecification:
@@ -339,6 +431,11 @@ Outputs:
     Value: !Ref ScratchDatabase
   AutoScalingGroup:
     Value: !Ref Fleet
+  Results:
+    Value: !If
+      - HasResultsBucket
+      - !Sub s3://${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/
+      - !Sub s3://${Bucket}/
   SourceBucket:
     Value: !Ref SourceBucket
   SourcePrefix:

benchmarks/pyathena_bench/__main__.py

@@ -22,7 +22,9 @@
     preflight,
     prepare,
     prepare_fixtures,
+    run_locations,
     session,
+    split_location,
     stack_resources,
     validate_manifest,
 )
@@ -32,15 +34,17 @@
     API_FAMILIES,
     Filters,
     Queue,
+    SpotInterruptionError,
     expand_jobs,
     merge_jobs,
     retire_host,
     status,
     submit,
+    suspend_rebalance,
     work,
 )
 from pyathena_bench.report import report
-from pyathena_bench.runner import run
+from pyathena_bench.runner import completed_trials, run, s3_checkpoint
 
 
 def parser() -> argparse.ArgumentParser:
@@ -92,6 +96,18 @@ def parser() -> argparse.ArgumentParser:
         if name == "run":
             command.add_argument("--manifest", type=Path, required=True)
             command.add_argument("--out", type=Path, required=True)
+            command.add_argument(
+                "--resume",
+                type=Path,
+                nargs="+",
+                default=[],
+                help="Output directories of earlier attempts whose finished trials are skipped",
+            )
+            command.add_argument(
+                "--checkpoint",
+                help="s3:// prefix below which each trial's start and record are written "
+                "between trials; set by fleet workers, whose job summary merges the attempts",
+            )
     enqueue = commands.add_parser(
         "queue", help="Publish jobs and a prepared manifest for fleet workers"
     )
@@ -102,7 +118,11 @@ def parser() -> argparse.ArgumentParser:
     )
     enqueue.add_argument("--manifest", type=Path, required=True)
     consume = commands.add_parser("worker", help="Run queued jobs on this host until none remain")
-    consume.add_argument("--stack", required=True)
+    location = consume.add_mutually_exclusive_group(required=True)
+    location.add_argument("--stack", help="Stack whose Results output holds the queue")
+    location.add_argument(
+        "--results", help="The stack's results location, which hosts know from their user data"
+    )
     consume.add_argument("--name", required=True)
     consume.add_argument(
         "--api-slots",
@@ -114,7 +134,19 @@ def parser() -> argparse.ArgumentParser:
     consume.add_argument(
         "--terminate-host",
         action="store_true",
-        help="Terminate this host through its Auto Scaling group after every job is claimed",
+        help="Suspend zone rebalancing at start, and terminate this host through its Auto "
+        "Scaling group when no job is left to claim",
+    )
+    consume.add_argument(
+        "--wait",
+        action="store_true",
+        help="Wait for the queue to be published instead of failing",
+    )
+    consume.add_argument(
+        "--lease-seconds",
+        type=float,
+        default=600.0,
+        help="Seconds after which a dead worker's claim and API slots return to the fleet",
     )
     progress = commands.add_parser("status", help="Summarize a fleet queue")
     progress.add_argument("--stack", required=True)
@@ -164,10 +196,13 @@ def main(argv: list[str] | None = None) -> int:
                         if args.trials_only
                         else f"b_{manifest['run_id']}_*"
                     ),
-                    "s3_prefix": f"s3://{resources['Bucket']}/runs/{manifest['run_id']}/"
-                    + ("trials/" if args.trials_only else ""),
+                    "s3_prefixes": [
+                        root + ("trials/" if args.trials_only else "")
+                        for root in sorted(
+                            set(run_locations(resources, manifest["run_id"]).values())
+                        )
+                    ],
                     "preserve_inputs_and_fixtures": args.trials_only,
-                    "query_output_prefix": f"s3://{resources['Bucket']}/runs/{manifest['run_id']}/",
                 },
                 indent=2,
             )
@@ -177,20 +212,39 @@ def main(argv: list[str] | None = None) -> int:
             cleanup(settings, manifest, args.manifest, trials_only=args.trials_only)
     elif args.command in {"queue", "worker", "status"}:
         session_ = session(settings)
-        queue = Queue(
-            client(session_, "s3"), stack_resources(session_, args.stack)["Bucket"], args.name
-        )
+        results = getattr(args, "results", None) or stack_resources(session_, args.stack)["Results"]
+        bucket, base = split_location(results)
+        queue = Queue(client(session_, "s3"), bucket, args.name, base)
         if args.command == "queue":
             validate_manifest(settings, read_json(args.manifest))
             submit(queue, merge_jobs(read_json(p) for p in args.jobs), args.config, args.manifest)
         elif args.command == "worker":
             if args.api_slots < 1:
                 raise ValueError("--api-slots must be positive")
-            ran = work(queue, args.workdir or Path("results/fleet") / args.name, args.api_slots)
+            if args.lease_seconds < 300:
+                raise ValueError("--lease-seconds must be at least 300")
+            autoscaling = client(session_, "autoscaling")
+            if args.terminate_host:
+                suspend_rebalance(autoscaling)
+            try:
+                ran = work(
+                    queue,
+                    args.workdir or Path("results/fleet") / args.name,
+                    args.api_slots,
+                    lease_seconds=args.lease_seconds,
+                    wait=args.wait,
+                )
+            except SpotInterruptionError as interrupted:
+                # The host is being reclaimed, and its group launches a replacement; after
+                # a rehearsed notice, the operator terminates the host the same way.
+                sys.stdout.write(
+                    json.dumps({"ran": interrupted.ran, "interrupted": True}, indent=2) + "\n"
+                )
+                return 0
             sys.stdout.write(json.dumps({"ran": ran}, indent=2) + "\n")
             if args.terminate_host:
                 # Only a normal return retires the host; an error keeps it for inspection.
-                retire_host(queue, client(session_, "autoscaling"), ran)
+                retire_host(queue, autoscaling, ran)
         else:
             sys.stdout.write(json.dumps(status(queue), indent=2) + "\n")
     else:
@@ -248,8 +302,23 @@ def main(argv: list[str] | None = None) -> int:
             manifest = read_json(args.manifest)
             if args.suite == "init":
                 prepare_fixtures(settings, manifest, args.manifest, scales, args.shape)
-            success = run(settings, manifest, cases, scales, args.out)
-            report(args.out)
+            checkpoint = (
+                s3_checkpoint(client(session(settings), "s3"), args.checkpoint)
+                if args.checkpoint
+                else None
+            )
+            success = run(
+                settings,
+                manifest,
+                cases,
+                scales,
+                args.out,
+                completed_trials(args.resume),
+                checkpoint,
+            )
+            if not checkpoint:
+                # A fleet job's summary merges its attempts instead.
+                report(args.out)
             return 0 if success else 1
     return 0