From 1c75926fb9d8f1d56fa1a3e48b6aa9786935eca5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 15:23:50 +0900 Subject: [PATCH 01/16] Resume benchmark fleet jobs after Spot interruptions Fleet claims and API slots become leases that the running worker renews with conditional writes; a dead worker's job and slots return to the fleet when its lease expires. A worker that sees a Spot interruption notice interrupts the running trial, uploads the attempt, releases its claim as interrupted, and exits without retiring its host; the group's replacement continues the job as the next attempt, running only the repetitions that earlier attempts did not record (run --resume), and the job's report merges the attempts. Queue state, reports, and trial outputs go to a results location that can be an existing bucket (ResultsBucket), such as the CI test bucket whose objects expire after a day, with write access only below the stack's prefix. Hosts start a worker for QueueName once the queue is published, suspend zone rebalancing at start, and are kept by a KeepHost instance tag. Spot is the default purchase option, and the group can use four Availability Zones. Refs #1136. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 139 +++-- benchmarks/cloudformation/benchmark.yaml | 103 +++- benchmarks/pyathena_bench/__main__.py | 66 ++- benchmarks/pyathena_bench/aws.py | 67 ++- benchmarks/pyathena_bench/fleet.py | 641 +++++++++++++++++++---- benchmarks/pyathena_bench/report.py | 44 +- benchmarks/pyathena_bench/runner.py | 50 +- benchmarks/tests/test_aws.py | 64 ++- benchmarks/tests/test_cloudformation.py | 75 ++- benchmarks/tests/test_fleet.py | 277 +++++++++- benchmarks/tests/test_report.py | 32 ++ benchmarks/tests/test_runner.py | 41 +- 12 files changed, 1397 insertions(+), 202 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 60383d6b8..1cfd09fa5 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -95,16 +95,22 @@ The defaults are Amazon Linux 2023 x86_64, `r7i.2xlarge` (8 vCPUs, 64 GiB), and Set `Architecture=arm64` with an `r7g` or `r8g` instance type to run on AWS Graviton; the template rejects an instance type that does not match the architecture. Measurements depend on the CPU architecture; `environment.json` records it, and runs on different architectures are not directly comparable. Instance size and disk size are parameters, so a later run can deliberately test another memory budget. -`Purchase=spot` launches Spot Instances of the same instance type instead of On-Demand Instances. -The Auto Scaling group spreads the hosts over three Availability Zones with the price-capacity-optimized allocation strategy and never mixes instance types, so every host stays identical. -The template uses the first three zones that `Fn::GetAZs` returns, so deployment fails when it returns fewer than three, which happens when the account's default VPC has default subnets in only one or two zones. -The group replaces an interrupted Spot Instance, but the job that ran on it keeps its claim; see the retry procedure in [Running on a fleet](#running-on-a-fleet). +The hosts are Spot Instances of one instance type by default; `Purchase=on-demand` launches On-Demand Instances of that type instead. +The Auto Scaling group spreads the hosts over `Zones` Availability Zones (4 by default, or 3) with the price-capacity-optimized allocation strategy and never mixes instance types, so every host stays identical. +The template uses the first `Zones` zones that `Fn::GetAZs` returns, so deployment fails when it returns fewer, which depends on the zones in which the account's default VPC has default subnets; `us-west-2` has four. +The group replaces an interrupted Spot Instance, and the replacement continues the interrupted job; see [Running on a fleet](#running-on-a-fleet). `environment.json` does not record the instance type or purchase option; a worker started with `--terminate-host` records both under `hosts/` in its queue, and the recovered `instances.json` lists the hosts that still exist. The AMI parameter resolves the current AL2023 image at deployment; record the resulting AMI when comparing environments. Source bucket access is read-only and restricted to the configured prefix. -The instances can write only to the scratch bucket and database, and can query the specified existing workgroup. +The instances can write only to the scratch bucket and database and to the results location below, and can query the specified existing workgroup. The supplied template assumes ordinary IAM access and SSE-S3 source objects; Lake Formation restrictions or a customer-managed KMS key require corresponding grants before execution. +The fleet queues, reports, and Athena trial outputs go to the stack's results location, the `Results` stack output. +With `ResultsBucket` set to an existing bucket, the results location is `s3://///`, and the instance role can write only below it; with `ResultsBucket` empty, the default, it is the scratch bucket. +Input snapshots, initialization fixtures, and scratch tables stay in the scratch bucket and database either way. +This repository's CI test bucket, which `AWS_ATHENA_S3_STAGING_DIR` names in CI, expires every object a day after it is written (`cloudformation/github_actions_oidc.yaml`). +As the results location, it removes leftovers by itself, but reports must be recovered within that day. + The deployment identity needs permission to create these resources and pass the instance role. Use an immutable, remotely accessible commit that contains this directory and the root `uv.lock`. Bootstrap checks out that commit, installs pinned uv and Python versions, and runs `uv sync --locked --no-dev` for the benchmark into `/opt/pyathena/.venv`. @@ -270,6 +276,7 @@ uv run --no-sync python -m pyathena_bench run \ ``` There is no automatic continuation of a partially recorded trial; rerun the selected comparison with the same manifest. +`run --resume` names the output directories of earlier attempts and skips the trials that they recorded as successful, which a fleet worker uses to continue an interrupted job. Large API row cases may exceed the default one-hour timeout: 10 million rows require at least 10,000 pages with arraysize 1000, or 100,000 pages with arraysize 100. Choose the timeout using a smaller-scale pilot before the full run. Do not run two orchestrators concurrently on the same dedicated host. @@ -295,16 +302,21 @@ With `--split-pages`, a job whose single trial needs at least that many pages be `queue --jobs` accepts several job files, rejects duplicate job IDs, and orders the combined jobs as `jobs` orders one selection. Deploy a new stack with the fleet size under its own name; changing `FleetSize` on an existing stack does not wait for the added hosts to finish bootstrap. +`QueueName` names the queue that every host consumes: after bootstrap, each host starts a worker that waits until a queue of that name is published, including the hosts that the group launches later to replace interrupted Spot Instances. +`ApiSlots` sets their `--api-slots`. +`ResultsBucket` selects the results location; the commands below use this repository's CI test bucket, whose objects expire after a day. Then open a session on one of its hosts: ```bash BENCHMARK_STACK=pyathena-benchmark-fleet +BENCHMARK_RESULTS_BUCKET="" uv run --env-file ../.env --locked aws cloudformation deploy \ --stack-name "$BENCHMARK_STACK" \ --template-file cloudformation/benchmark.yaml \ --capabilities CAPABILITY_IAM \ --tags Purpose=pyathena-benchmark \ - --parameter-overrides GitCommit="$BENCHMARK_COMMIT" FleetSize=20 + --parameter-overrides GitCommit="$BENCHMARK_COMMIT" FleetSize=12 Architecture=arm64 InstanceType=r7g.2xlarge \ + ResultsBucket="$BENCHMARK_RESULTS_BUCKET" QueueName=large-1 ApiSlots=10 BENCHMARK_GROUP=$(uv run --env-file ../.env --locked aws cloudformation describe-stacks --stack-name "$BENCHMARK_STACK" \ --query 'Stacks[0].Outputs[?OutputKey==`AutoScalingGroup`].OutputValue | [0]' --output text) BENCHMARK_INSTANCE=$(uv run --env-file ../.env --locked aws autoscaling describe-auto-scaling-groups \ @@ -366,38 +378,92 @@ uv run --no-sync python -m pyathena_bench --config results/fleet.toml queue \ The manifest for this selection must be prepared with `--scale medium large xlarge`. It is published as `large-1` in place of the previous example, so the worker and status commands below apply unchanged. -`queue` reserves the name with a conditional write, stores the configuration, manifest, and jobs under `fleet//` in the scratch bucket, and writes the `queue.json` marker last. -Workers refuse a queue without the marker. +`queue` reserves the name with a conditional write, stores the configuration, manifest, and jobs under `fleet//` in the results location, and writes the `queue.json` marker last. +Workers refuse a queue without the marker, or wait for it with `--wait`. A failed `queue` releases the reservation, so the same name can be published again; a `queue` process killed while publishing leaves `fleet//reservation.json`, which must be deleted before reusing the name. -Every worker runs its jobs with that stored configuration. -Start a worker on every host from the local machine: +Every worker runs its jobs with that stored configuration and manifest, so recovery and cleanup do not depend on the host where they were prepared. +The workers that the stack started begin as soon as `large-1` is published; follow them from the local machine: ```bash -uv run --env-file ../.env --locked aws ssm send-command --targets "Key=tag:aws:autoscaling:groupName,Values=$BENCHMARK_GROUP" \ - --document-name AWS-RunShellScript \ - --parameters 'commands=["sudo -iu ec2-user bash -lc \"cd /opt/pyathena/benchmarks && mkdir -p results && (nohup uv run --no-sync python -m pyathena_bench worker --stack $(cat /opt/pyathena/benchmarks/stack-id.txt) --name large-1 --terminate-host > results/worker-large-1.log 2>&1 &)\""]' uv run --env-file ../.env --locked python -m pyathena_bench status \ --stack "$BENCHMARK_STACK" --name large-1 ``` +For a queue with another name, start a worker on every host with SSM; a second worker on a host whose worker still runs would overlap trials: + +```bash +uv run --env-file ../.env --locked aws ssm send-command --targets "Key=tag:aws:autoscaling:groupName,Values=$BENCHMARK_GROUP" \ + --document-name AWS-RunShellScript \ + --parameters 'commands=["sudo -iu ec2-user bash -lc \"cd /opt/pyathena/benchmarks && mkdir -p results && (nohup uv run --no-sync python -m pyathena_bench worker --stack $(cat /opt/pyathena/benchmarks/stack-id.txt) --name large-2 --terminate-host > results/worker-large-2.log 2>&1 &)\""]' +``` + A worker claims a job by creating `claims/` with an S3 conditional write, so exactly one host runs each job. -After a job, the worker uploads its output directory to `results//` and its log to `logs/`, then writes `done/` with the exit code, host, and times. -A failed trial stops only its own job; `status` lists jobs with a non-zero exit code. +A claim, like each API slot under `slots/`, is a lease that its worker renews every tenth of `--lease-seconds` (600 by default) while the job runs. +When a worker dies, its claim and slots expire, and another worker takes them over with a conditional write, so the job and the slots return to the fleet without manual steps. +Each attempt of a job writes to `results//attempt-/`, and the worker uploads its trial records at every renewal, so a lost host loses at most the trials of the last minutes. +After a job, the worker uploads the attempt's output directory and its log to `logs/`, writes `results//summary.csv` and `summary.md` from every attempt, and writes `done/` with the exit code, attempt, host, and times. +A failed trial stops only its own job; `status` lists jobs with a non-zero exit code, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. After a failed job, the worker stops the job's reported queries and every queued or running query in the workgroup that writes under the job's trial output prefixes, which also covers queries whose IDs were never reported before a trial process died. It waits until they stop before claiming another job; this scans the workgroup's query history. If it cannot confirm that, it records the errors in the `done` marker and exits. -A worker that dies leaves a claim without a `done` marker, and it can also leave an API slot under `slots/`, which lowers the fleet's API-heavy capacity. -`status` lists slots in use; delete a slot object only after confirming that the host named in it no longer runs a job. -Workers exit when every job has been claimed. -With `--terminate-host`, a worker that returns normally records its jobs, instance type, Availability Zone, and purchase option under `hosts/.json` and terminates its host through the Auto Scaling group with a lower desired capacity, so no replacement starts and finished hosts stop costing money while others still run. -Before terminating, it suspends the group's `AZRebalance` process: hosts retiring unevenly across zones would otherwise let the group launch a host in one zone and terminate a busy host in another. -The process stays suspended until the stack is deleted. + +Each worker checks instance metadata for a Spot interruption notice (`spot/instance-action`), which EC2 publishes about two minutes before it reclaims an instance. +On a notice, the worker interrupts the running trial with SIGINT, which cancels the trial's reported queries, uploads the attempt, releases its claim as `interrupted` and its slots, and exits without terminating the host. +The group replaces the reclaimed instance, and the new host's worker takes over the interrupted claim as the next attempt: it downloads the earlier attempts, quiesces their trials as after a failed job, and runs only the repetitions that they did not record as successful, through `run --resume`. +The job's summary merges the attempts and leaves out a failed or incomplete trial whose repetition succeeded in a later attempt, so it reports the trials of an uninterrupted run. +Rebalance recommendations are not acted on; they often precede no interruption, and the notice leaves enough time to stop a trial. + +Workers exit when no job is left to claim: every job is done or claimed by a live lease. +With `--terminate-host`, a worker first suspends the group's `AZRebalance` process, which would otherwise let the group launch a host in one zone and terminate a busy host in another when Spot capacity or retired hosts leave the zones uneven; the process stays suspended until the stack is deleted. +A worker with `--terminate-host` that returns normally records its jobs, instance type, Availability Zone, and purchase option under `hosts/.json` and terminates its host through the group with a lower desired capacity, so no replacement starts and finished hosts stop costing money while others still run. This includes a worker that stopped after unconfirmed queries, whose errors are already in the `done` marker. -A worker that raises an error, for example for a missing queue, keeps its host. -To keep a host, create `/opt/pyathena/benchmarks/results/keep-host` before its worker finishes; keep the host where the manifests were prepared, so that its local `results/` can be recovered. -Start a worker on a replacement host for an interrupted Spot Instance by targeting that instance ID instead of the group; a second worker on a host that already runs one would overlap trials. -To retry jobs, wait until all workers have exited, run `cleanup --trials-only --execute` for the manifest, and publish the selected jobs under a new queue name with the same manifest. -Download `fleet//` with `aws s3 sync` before cleanup and stack deletion, as described below. +A worker that raises an error, for example for a job wider than `--api-slots`, keeps its host. +To keep a host after its worker finishes, tag it with `KeepHost=true` (`aws ec2 create-tags --resources --tags Key=KeepHost,Value=true`); the worker reads the tag from instance metadata. +A job whose last worker exited without finishing it, for example after a host error, stays claimed until its lease expires; then a worker started for the same queue takes it over. +To retry failed jobs, wait until all workers have exited, run `cleanup --trials-only --execute` for the manifest, and publish the selected jobs under a new queue name with the same manifest. + +### Pilot with a simulated interruption + +Before a fleet run, rehearse an interruption on one host. +Deploy with `FleetSize=1` and `QueueName=pilot-1`, prepare a small snapshot, and publish a job with several repetitions; the host's worker starts it: + +```bash +BENCHMARK_STACK=pyathena-benchmark-pilot +uv run --env-file ../.env --locked aws cloudformation deploy \ + --stack-name "$BENCHMARK_STACK" \ + --template-file cloudformation/benchmark.yaml \ + --capabilities CAPABILITY_IAM \ + --tags Purpose=pyathena-benchmark \ + --parameter-overrides GitCommit="$BENCHMARK_COMMIT" Architecture=arm64 InstanceType=r7g.2xlarge \ + ResultsBucket="$BENCHMARK_RESULTS_BUCKET" QueueName=pilot-1 +# In a session on the host, as ec2-user in /opt/pyathena/benchmarks: +uv run --no-sync python -m pyathena_bench prepare \ + --stack "$(cat stack-id.txt)" --manifest results/input-pilot.json --scale small +uv run --no-sync python -m pyathena_bench jobs --suite single --scale small --family pandas --api sync \ + --transport csv --output-kind native --warmups 0 --repetitions 6 > results/jobs-pilot.json +uv run --no-sync python -m pyathena_bench queue \ + --stack "$(cat stack-id.txt)" --name pilot-1 --jobs results/jobs-pilot.json --manifest results/input-pilot.json +``` + +When `status` shows the job claimed and `results//attempt-1/trials.jsonl` has a trial, create the file `/opt/pyathena/benchmarks/interruption-notice`, which a worker treats as an interruption notice. +Check that the claim is `interrupted`, then terminate the host without lowering the desired capacity, as EC2 does when it reclaims a Spot Instance: + +```bash +BENCHMARK_GROUP=$(uv run --env-file ../.env --locked aws cloudformation describe-stacks --stack-name "$BENCHMARK_STACK" \ + --query 'Stacks[0].Outputs[?OutputKey==`AutoScalingGroup`].OutputValue | [0]' --output text) +BENCHMARK_INSTANCE=$(uv run --env-file ../.env --locked aws autoscaling describe-auto-scaling-groups \ + --auto-scaling-group-names "$BENCHMARK_GROUP" --query 'AutoScalingGroups[0].Instances[0].InstanceId' --output text) +uv run --env-file ../.env --locked aws ssm send-command --instance-ids "$BENCHMARK_INSTANCE" \ + --document-name AWS-RunShellScript \ + --parameters 'commands=["sudo -u ec2-user touch /opt/pyathena/benchmarks/interruption-notice"]' +uv run --env-file ../.env --locked python -m pyathena_bench status --stack "$BENCHMARK_STACK" --name pilot-1 +uv run --env-file ../.env --locked aws autoscaling terminate-instance-in-auto-scaling-group \ + --instance-id "$BENCHMARK_INSTANCE" --no-should-decrement-desired-capacity +``` + +The group launches a replacement whose worker continues the job as `attempt-2` and then retires the host. +The pilot passes when `done/` records `attempt` 2 and exit code 0, `results//summary.csv` reports six successful trials and no failures, `hosts/` records the replacement with lifecycle `spot`, and the group's desired capacity drops to zero. +AWS Fault Injection Service (`aws:ec2:send-spot-instance-interruptions`) sends a real notice instead, but needs its own IAM role and experiment template. ## Reports and recovery @@ -412,6 +478,7 @@ uv run --no-sync python -m pyathena_bench report results/single-small ``` Report recovery ignores an incomplete final JSONL record with a warning and marks started trials without a final record as incomplete. +For a fleet job directory, `report results/recovered/fleet//results/` merges the trials of its `attempt-` subdirectories. Malformed records elsewhere remain errors. Small sample sizes make tail percentiles descriptive rather than reliable population estimates. @@ -421,10 +488,11 @@ Results are gitignored and are not uploaded by CI. ## Recovering reports and deleting the environment Stop the benchmark process before cleanup or stack deletion. -Recover all local manifests and reports first, including failed-run evidence. -Use the stack's generated bucket as a temporary transfer location; it is separate from `pyathena-benchmark`, which holds the input data. +Recover the fleet queues and reports first, including failed-run evidence; with an expiring results bucket, do so within a day of the run. +Each queue holds its manifest, configuration, and jobs, so no host needs to be kept for recovery. +Local files on a remaining host, such as a manifest that was never queued, can go through the stack's generated bucket; it is separate from `pyathena-benchmark`, which holds the input data. -On EC2: +On EC2, only for such local files: ```bash BENCHMARK_BUCKET=$(aws cloudformation describe-stacks --stack-name "$BENCHMARK_STACK_ID" \ @@ -439,29 +507,32 @@ BENCHMARK_BUCKET=$(uv run --env-file ../.env --locked aws cloudformation describ --query 'Stacks[0].Outputs[?OutputKey==`Bucket`].OutputValue | [0]' --output text) mkdir -p results/recovered uv run --env-file ../.env --locked aws s3 sync "s3://$BENCHMARK_BUCKET/reports/" results/recovered/ +BENCHMARK_RESULTS=$(uv run --env-file ../.env --locked aws cloudformation describe-stacks --stack-name "$BENCHMARK_STACK" \ + --query 'Stacks[0].Outputs[?OutputKey==`Results`].OutputValue | [0]' --output text) uv run --env-file ../.env --locked aws cloudformation describe-stacks --stack-name "$BENCHMARK_STACK" > results/recovered/stack.json -uv run --env-file ../.env --locked aws s3 sync "s3://$BENCHMARK_BUCKET/fleet/" results/recovered/fleet/ +uv run --env-file ../.env --locked aws s3 sync "${BENCHMARK_RESULTS}fleet/" results/recovered/fleet/ uv run --env-file ../.env --locked aws ec2 describe-instances \ --filters "Name=tag:aws:autoscaling:groupName,Values=$BENCHMARK_GROUP" > results/recovered/instances.json ``` Verify the downloaded files before proceeding. -For every preparation manifest, preview and then clean that run using its original configuration: +For every preparation manifest, preview and then clean that run using its original configuration; a queued manifest is `results/recovered/fleet//manifest.json`: ```bash uv run --env-file ../.env --locked python -m pyathena_bench cleanup \ - --manifest results/recovered/input.json + --manifest results/recovered/fleet/large-1/manifest.json uv run --env-file ../.env --locked python -m pyathena_bench cleanup \ - --manifest results/recovered/input.json --execute + --manifest results/recovered/fleet/large-1/manifest.json --execute ``` -Cleanup checks the live stack identity, cancels and waits for matching active queries, removes run-prefixed scratch tables, aborts matching multipart uploads, and removes the run's S3 prefix. +Cleanup checks the live stack identity, cancels and waits for matching active queries, removes run-prefixed scratch tables, aborts matching multipart uploads, and removes the run's prefix `runs//` in the scratch bucket and in the results location. History discovery uses [batches of up to 50 query IDs](https://docs.aws.amazon.com/athena/latest/APIReference/API_BatchGetQueryExecution.html); completed queries need no cancellation calls. Unavailable recorded query metadata produces a warning and falls back to scanning live history; errors inspecting that history stop cleanup before deletion. The transfer copy under `reports/` remains until the final stack teardown. For multiple manifests, repeat cleanup for each; never clean while another process is still submitting queries. Confirm the scratch database is empty and the only remaining S3 objects are the recovered reports and fleet queues. +In a separate results bucket, the fleet queues under `fleet/` remain until they expire, or until they are removed with `aws s3 rm "${BENCHMARK_RESULTS}fleet/" --recursive`. An unexpected table or object is a reason to inspect the corresponding run before deleting anything further. ```bash diff --git a/benchmarks/cloudformation/benchmark.yaml b/benchmarks/cloudformation/benchmark.yaml index ace157ce3..b94c3330f 100644 --- a/benchmarks/cloudformation/benchmark.yaml +++ b/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//; 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,28 @@ 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. + - !If + - HasResultsBucket + - Effect: Allow + Action: s3:GetBucketLocation + Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket} + - !Ref AWS::NoValue + - !If + - HasResultsBucket + - Effect: Allow + Action: [s3:ListBucket, s3:ListBucketMultipartUploads] + Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket} + Condition: + StringLike: + s3:prefix: !Sub '${ResultsPrefix}/${AWS::StackName}/*' + - !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 +332,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: @@ -307,6 +370,26 @@ Resources: 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 + 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 --stack ${AWS::StackId} --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.' Fleet: Type: AWS::AutoScaling::AutoScalingGroup @@ -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: diff --git a/benchmarks/pyathena_bench/__main__.py b/benchmarks/pyathena_bench/__main__.py index 40e1883c1..a802725ae 100644 --- a/benchmarks/pyathena_bench/__main__.py +++ b/benchmarks/pyathena_bench/__main__.py @@ -23,6 +23,7 @@ prepare, prepare_fixtures, session, + split_location, stack_resources, validate_manifest, ) @@ -31,15 +32,18 @@ from pyathena_bench.fleet import ( Filters, Queue, + SpotInterruptionError, expand_jobs, + interruption_notice, 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 def parser() -> argparse.ArgumentParser: @@ -85,6 +89,13 @@ 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", + ) enqueue = commands.add_parser( "queue", help="Publish jobs and a prepared manifest for fleet workers" ) @@ -107,7 +118,17 @@ 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) @@ -157,10 +178,12 @@ 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": [ + f"{root}runs/{manifest['run_id']}/" + + ("trials/" if args.trials_only else "") + for root in sorted({f"s3://{resources['Bucket']}/", resources["Results"]}) + ], "preserve_inputs_and_fixtures": args.trials_only, - "query_output_prefix": f"s3://{resources['Bucket']}/runs/{manifest['run_id']}/", }, indent=2, ) @@ -170,20 +193,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 - ) + bucket, base = split_location(stack_resources(session_, args.stack)["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 < 120: + raise ValueError("--lease-seconds must be at least 120") + 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, + renew_seconds=args.lease_seconds / 10, + notice=interruption_notice, + wait=args.wait, + ) + except SpotInterruptionError as interrupted: + # The host is being reclaimed; its group launches a replacement. + 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") elif args.command == "jobs": @@ -248,7 +290,9 @@ 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) + success = run( + settings, manifest, cases, scales, args.out, completed_trials(args.resume) + ) report(args.out) return 0 if success else 1 return 0 diff --git a/benchmarks/pyathena_bench/aws.py b/benchmarks/pyathena_bench/aws.py index c3f6008cc..516d40e4e 100644 --- a/benchmarks/pyathena_bench/aws.py +++ b/benchmarks/pyathena_bench/aws.py @@ -37,6 +37,24 @@ def client(session_: Any, service: str) -> Any: ) +def split_location(location: str) -> tuple[str, str]: + """Split an ``s3://bucket/prefix/`` location. + + Args: + location: S3 location ending with a slash. + + Returns: + The bucket and the key prefix, which is empty or ends with a slash. + + Raises: + ValueError: If the location is not an S3 location ending with a slash. + """ + if not re.fullmatch(r"s3://[a-z0-9][a-z0-9.-]+[a-z0-9]/([^/]+/)*", location): + raise ValueError(f"Invalid S3 location: {location!r}") + bucket, _, prefix = location[len("s3://") :].partition("/") + return bucket, prefix + + def stack_resources(session_: Any, stack: str) -> dict[str, str]: data = client(session_, "cloudformation").describe_stacks(StackName=stack)["Stacks"][0] tags = {item["Key"]: item["Value"] for item in data.get("Tags", [])} @@ -52,6 +70,7 @@ def stack_resources(session_: Any, stack: str) -> dict[str, str]: "SourceTable", "WorkGroup", "AutoScalingGroup", + "Results", } if not required <= result.keys(): raise ValueError("Stack outputs are incomplete") @@ -61,6 +80,8 @@ def stack_resources(session_: Any, stack: str) -> dict[str, str]: ): raise ValueError("Scratch resources must differ from source resources") identifier(result["ScratchDatabase"]) + if split_location(result["Results"])[0] == result["SourceBucket"]: + raise ValueError("Results must not be written to the source bucket") result["StackId"] = data["StackId"] return result @@ -329,13 +350,23 @@ def active_queries(athena: Any, workgroup: str, prefixes: tuple[str, ...]) -> se def cleanup( settings: Settings, manifest: dict[str, Any], path: Path, trials_only: bool = False ) -> None: - """Remove this run's data after checking the current stack identity.""" + """Remove this run's data after checking the current stack identity. + + The run's objects are under ``runs//`` in the scratch bucket + (inputs and fixtures) and in the stack's results location (trial outputs); + both are cleaned. + """ resources = validate_manifest(settings, manifest) session_ = session(settings) - bucket = resources["Bucket"] run_id = manifest["run_id"] - root = f"runs/{run_id}/" - prefix = root + "trials/" if trials_only else root + results_bucket, results_prefix = split_location(resources["Results"]) + locations = sorted( + { + (resources["Bucket"], f"runs/{run_id}/"), + (results_bucket, f"{results_prefix}runs/{run_id}/"), + } + ) + roots = tuple(f"s3://{bucket}/{root}" for bucket, root in locations) athena = client(session_, "athena") # Discover queries whose worker died before returning its query ID. # The query prefix is checked against live Athena output metadata. @@ -356,11 +387,11 @@ def cleanup( continue if query.get("WorkGroup") != settings.workgroup or not query.get( "ResultConfiguration", {} - ).get("OutputLocation", "").startswith(f"s3://{bucket}/{root}"): + ).get("OutputLocation", "").startswith(roots): raise ValueError("Manifest contains a query outside this run") if query["Status"]["State"] in {"RUNNING", "QUEUED"}: ids.add(entry["query_id"]) - ids |= active_queries(athena, settings.workgroup, (f"s3://{bucket}/{root}",)) + ids |= active_queries(athena, settings.workgroup, roots) errors = cancel_queries(settings, ids) if errors: raise RuntimeError("Could not cancel queries: " + "; ".join(errors)) @@ -393,15 +424,21 @@ def cleanup( for table in tables: glue.delete_table(DatabaseName=resources["ScratchDatabase"], Name=table) s3 = client(session_, "s3") - for page in s3.get_paginator("list_multipart_uploads").paginate(Bucket=bucket, Prefix=prefix): - for upload in page.get("Uploads", []): - s3.abort_multipart_upload(Bucket=bucket, Key=upload["Key"], UploadId=upload["UploadId"]) - for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix): - objects = [{"Key": obj["Key"]} for obj in page.get("Contents", [])] - if objects: - response = s3.delete_objects(Bucket=bucket, Delete={"Objects": objects}) - if response.get("Errors"): - raise RuntimeError(str(response["Errors"])) + for bucket, root in locations: + prefix = root + "trials/" if trials_only else root + for page in s3.get_paginator("list_multipart_uploads").paginate( + Bucket=bucket, Prefix=prefix + ): + for upload in page.get("Uploads", []): + s3.abort_multipart_upload( + Bucket=bucket, Key=upload["Key"], UploadId=upload["UploadId"] + ) + for page in s3.get_paginator("list_objects_v2").paginate(Bucket=bucket, Prefix=prefix): + objects = [{"Key": obj["Key"]} for obj in page.get("Contents", [])] + if objects: + response = s3.delete_objects(Bucket=bucket, Delete={"Objects": objects}) + if response.get("Errors"): + raise RuntimeError(str(response["Errors"])) if not trials_only: manifest["cleaned"] = True write_json(path, manifest) diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index b14fafad1..b1489a72c 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -10,15 +10,19 @@ from __future__ import annotations import json +import signal import socket import subprocess import sys +import threading import time +import urllib.error import urllib.request from collections.abc import Callable, Iterable from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path +from types import TracebackType from typing import Any from botocore.exceptions import ClientError @@ -26,13 +30,18 @@ from pyathena_bench.aws import active_queries, client, session from pyathena_bench.cases import Case, matrix from pyathena_bench.config import Settings +from pyathena_bench.report import report # Row cursors page through GetQueryResults, which has an account-wide request-rate quota. API_FAMILIES = {"cursor", "dict"} QUEUE_ROOT = "fleet" -# A host with this file keeps running after its worker exits, for inspection. -KEEP_HOST = Path("results/keep-host") +# An instance tag with the value "true" keeps a host running after its worker exits. +KEEP_HOST_TAG = "KeepHost" +# A file that a worker treats like a Spot interruption notice, to rehearse interruptions. +SIMULATED_NOTICE = Path("interruption-notice") METADATA_URL = "http://169.254.169.254/latest" +# S3 errors of a conditional write whose condition failed. +PRECONDITION_ERRORS = {"PreconditionFailed", "ConditionalRequestConflict"} @dataclass(frozen=True) @@ -219,15 +228,45 @@ def now() -> str: return datetime.now(UTC).isoformat(timespec="seconds") +@dataclass +class Lease: + """A claim or API slot object held by one worker until it expires. + + The body records the owner, the attempt number, the state (``running`` or + ``interrupted``), and the expiry as a Unix time. A worker renews its leases + while it runs a job; another worker may take over a lease that has expired + or was released as interrupted. Objects without an expiry, written by + earlier versions, never expire. + """ + + key: str + body: dict[str, Any] + etag: str + + def takeable(self, at: float) -> bool: + """Return whether another worker may take over the lease. + + Args: + at: Current Unix time. + + Returns: + True if the lease was interrupted or expired before ``at``. + """ + expires = self.body.get("expires") + return self.body.get("state") == "interrupted" or (expires is not None and expires < at) + + class Queue: - """S3 layout of one named fleet run under ``fleet//`` in the scratch bucket.""" + """S3 layout of one named fleet run under ``fleet//`` in a bucket.""" - def __init__(self, s3: Any, bucket: str, name: str) -> None: + def __init__(self, s3: Any, bucket: str, name: str, base: str = "") -> None: if not name or "/" in name: raise ValueError("Queue name must be a non-empty path segment") + if base and not base.endswith("/"): + raise ValueError("Queue base prefix must end with a slash") self.s3 = s3 self.bucket = bucket - self.prefix = f"{QUEUE_ROOT}/{name}/" + self.prefix = f"{base}{QUEUE_ROOT}/{name}/" def key(self, *parts: str) -> str: """Return the object key for a queue path. @@ -240,6 +279,30 @@ def key(self, *parts: str) -> str: """ return self.prefix + "/".join(parts) + def put(self, key: str, body: dict[str, Any], **condition: str) -> str | None: + """Write a JSON object, optionally only if a condition holds. + + Args: + key: Object key. + body: JSON-serializable content. + **condition: ``IfNoneMatch="*"`` or ``IfMatch=``. + + Returns: + The new ETag, or None if the condition failed. + + Raises: + ClientError: For S3 errors other than a failed precondition. + """ + try: + response = self.s3.put_object( + Bucket=self.bucket, Key=key, Body=json.dumps(body).encode(), **condition + ) + except ClientError as error: + if error.response["Error"]["Code"] in PRECONDITION_ERRORS: + return None + raise + return str(response["ETag"]) + def put_if_absent(self, key: str, body: dict[str, Any]) -> bool: """Create an object only if no object exists at the key. @@ -253,18 +316,68 @@ def put_if_absent(self, key: str, body: dict[str, Any]) -> bool: Raises: ClientError: For S3 errors other than a failed precondition. """ + return self.put(key, body, IfNoneMatch="*") is not None + + def read_lease(self, key: str) -> Lease | None: + """Read a claim or slot object. + + Args: + key: Object key. + + Returns: + The lease, or None if the object does not exist. + """ try: - self.s3.put_object( - Bucket=self.bucket, Key=key, Body=json.dumps(body).encode(), IfNoneMatch="*" - ) - return True + response = self.s3.get_object(Bucket=self.bucket, Key=key) except ClientError as error: - if error.response["Error"]["Code"] in { - "PreconditionFailed", - "ConditionalRequestConflict", - }: - return False + if error.response["Error"]["Code"] in {"NoSuchKey", "404"}: + return None raise + return Lease(key, json.loads(response["Body"].read()), str(response["ETag"])) + + def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Lease | None: + """Take a lease on an object that is absent, expired, or interrupted. + + Creation and takeover are conditional writes, so exactly one worker + obtains each lease. A takeover increments the attempt number. + + Args: + key: Object key of the claim or slot. + owner: Description of the worker and job. + lease_seconds: Seconds until the lease expires unless renewed. + + Returns: + The lease, or None if another worker holds it. + """ + body = {**owner, "attempt": 1, "state": "running", "expires": time.time() + lease_seconds} + etag = self.put(key, body, IfNoneMatch="*") + if etag is None: + current = self.read_lease(key) + if current is None or not current.takeable(time.time()): + return None + body["attempt"] = current.body.get("attempt", 1) + 1 + etag = self.put(key, body, IfMatch=current.etag) + if etag is None: + return None + return Lease(key, body, etag) + + def renew(self, lease: Lease, lease_seconds: float, **changes: Any) -> bool: + """Extend or update a lease if this worker still holds it. + + Args: + lease: Lease obtained from :meth:`acquire`; updated in place. + lease_seconds: Seconds from now until the lease expires. + **changes: Other fields to change, such as ``state``. + + Returns: + False if another worker took over the lease. + """ + body = {**lease.body, "expires": time.time() + lease_seconds, **changes} + etag = self.put(lease.key, body, IfMatch=lease.etag) + if etag is None: + return False + lease.body, lease.etag = body, etag + return True def names(self, folder: str) -> set[str]: """List the object names directly below a queue folder. @@ -294,16 +407,21 @@ def read(self, *parts: str) -> bytes: """ return self.s3.get_object(Bucket=self.bucket, Key=self.key(*parts))["Body"].read() - def acquire_slots(self, slots: int, count: int, owner: dict[str, Any]) -> list[str] | None: + def acquire_slots( + self, slots: int, count: int, owner: dict[str, Any], lease_seconds: float + ) -> list[Lease] | None: """Take API slots that bound simultaneous GetQueryResults paging. + A slot whose lease expired, because its worker died, is taken over. + Args: slots: Number of API slots in the fleet. count: Slots needed, one per simultaneous paging query. owner: Description stored in each slot object. + lease_seconds: Seconds until each slot lease expires unless renewed. Returns: - Slot object keys, or None if not enough slots are free. Slots taken + Slot leases, or None if not enough slots are free. Slots taken before a shortage or an error are released. Raises: @@ -312,20 +430,20 @@ def acquire_slots(self, slots: int, count: int, owner: dict[str, Any]) -> list[s needed = max(count, 1) if needed > slots: raise ValueError(f"A job needs {needed} API slots, but only {slots} exist") - taken: list[str] = [] + taken: list[Lease] = [] try: for index in range(slots): - key = self.key("slots", str(index)) - if self.put_if_absent(key, owner): - taken.append(key) + lease = self.acquire(self.key("slots", str(index)), owner, lease_seconds) + if lease is not None: + taken.append(lease) if len(taken) == needed: return taken except BaseException: - for key in taken: - self.release(key) + for lease in taken: + self.release(lease.key) raise - for key in taken: - self.release(key) + for lease in taken: + self.release(lease.key) return None def exists(self, *parts: str) -> bool: @@ -351,20 +469,41 @@ def release(self, key: str) -> None: """ self.s3.delete_object(Bucket=self.bucket, Key=key) - def upload_tree(self, directory: Path, *parts: str) -> None: + def upload_tree(self, directory: Path, *parts: str, pattern: str = "*") -> None: """Upload every file below a directory to a queue path. Args: directory: Local directory; a missing directory uploads nothing. *parts: Destination path segments below the queue prefix. + pattern: Glob pattern of the files to upload, relative to ``directory``. """ if not directory.exists(): return - for path in sorted(p for p in directory.rglob("*") if p.is_file()): + for path in sorted(p for p in directory.rglob(pattern) if p.is_file()): self.s3.upload_file( str(path), self.bucket, self.key(*parts, path.relative_to(directory).as_posix()) ) + def download_tree(self, directory: Path, *parts: str) -> list[str]: + """Download every object below a queue path into a local directory. + + Args: + directory: Local destination directory. + *parts: Source path segments below the queue prefix. + + Returns: + Relative paths of the downloaded files. + """ + source = self.key(*parts) + "/" + names = sorted(self.names("/".join(parts))) + for name in names: + path = directory / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes( + self.s3.get_object(Bucket=self.bucket, Key=source + name)["Body"].read() + ) + return names + def submit(queue: Queue, jobs: list[dict[str, Any]], config: Path, manifest: Path) -> None: """Publish jobs, configuration, and manifest for a fleet run. @@ -404,12 +543,26 @@ def submit(queue: Queue, jobs: list[dict[str, Any]], config: Path, manifest: Pat raise -def run_job(job: dict[str, Any], workdir: Path) -> int: - """Run one job as a separate orchestrator process. +def run_job( + job: dict[str, Any], + workdir: Path, + output: Path, + resume: list[Path], + stop: threading.Event, +) -> int: + """Run one job attempt as a separate orchestrator process. + + When ``stop`` is set, the orchestrator receives SIGINT, which stops the + running trial, cancels its reported queries, and records the trial as + interrupted; it is killed if it has not exited a minute later. Args: job: Job dictionary. workdir: Directory with the queue's ``config.toml`` and ``manifest.json``. + output: New output directory for this attempt. + resume: Output directories of earlier attempts whose completed trials + are not repeated. + stop: Event that interrupts the attempt. Returns: Process exit code. @@ -417,11 +570,21 @@ def run_job(job: dict[str, Any], workdir: Path) -> int: (workdir / "logs").mkdir(parents=True, exist_ok=True) command = [ sys.executable, "-m", "pyathena_bench", "--config", str(workdir / "config.toml"), - "run", "--manifest", str(workdir / "manifest.json"), - "--out", str(workdir / "results" / job["id"]), *job["args"], + "run", "--manifest", str(workdir / "manifest.json"), "--out", str(output), *job["args"], ] # fmt: skip - with (workdir / "logs" / f"{job['id']}.log").open("w") as stream: - return subprocess.call(command, stdout=stream, stderr=subprocess.STDOUT) + if resume: + command += ["--resume", *map(str, resume)] + with (workdir / "logs" / f"{job['id']}-{output.name}.log").open("w") as stream: + process = subprocess.Popen(command, stdout=stream, stderr=subprocess.STDOUT) + while process.poll() is None: + if stop.wait(1): + process.send_signal(signal.SIGINT) + try: + return process.wait(timeout=60) + except subprocess.TimeoutExpired: + process.kill() + return process.wait() + return process.returncode def trial_evidence(output: Path) -> tuple[set[str], tuple[str, ...]]: @@ -496,41 +659,267 @@ def quiesce(settings: Settings, output: Path) -> list[str]: return errors +class SpotInterruptionError(Exception): + """A Spot interruption notice stopped the worker before every job was claimed.""" + + def __init__(self, ran: list[str]) -> None: + super().__init__("Stopped by a Spot interruption notice") + self.ran = ran + + +class JobWatch: + """Renew a running job's leases, upload its progress, and watch for interruption. + + A background thread checks for an interruption notice every + ``poll_seconds`` and, every ``renew_seconds``, renews the leases and calls + ``progress``. A lease that another worker took over is no longer renewed and + is listed in ``lost``. Transient errors are retried at the next renewal. + """ + + def __init__( + self, + queue: Queue, + leases: list[Lease], + lease_seconds: float, + renew_seconds: float, + progress: Callable[[], None], + notice: Callable[[], bool], + poll_seconds: float = 5.0, + ) -> None: + self.queue = queue + self.leases = list(leases) + self.lease_seconds = lease_seconds + self.renew_seconds = renew_seconds + self.progress = progress + self.notice = notice + self.poll_seconds = min(poll_seconds, renew_seconds) + self.interrupted = threading.Event() + self.lost: list[str] = [] + self._stopped = threading.Event() + self._thread = threading.Thread(target=self._run, daemon=True) + + def __enter__(self) -> JobWatch: + self._thread.start() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.stop() + + def stop(self) -> None: + """Stop the thread and wait for it, so the caller may update the leases.""" + self._stopped.set() + if self._thread.is_alive(): + self._thread.join() + + def _run(self) -> None: + renew_at = time.monotonic() + self.renew_seconds + while not self._stopped.wait(self.poll_seconds): + if not self.interrupted.is_set() and self.notice(): + self.interrupted.set() + if time.monotonic() < renew_at: + continue + renew_at = time.monotonic() + self.renew_seconds + for lease in list(self.leases): + try: + if not self.queue.renew(lease, self.lease_seconds): + self.leases.remove(lease) + self.lost.append(lease.key) + except Exception: + # The lease has time left for a retry. + continue + try: + self.progress() + except Exception: + # The next renewal uploads again. + continue + + +def pending_jobs(queue: Queue, jobs: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Return the jobs that are neither done nor held by a live claim. + + Args: + queue: Queue to inspect. + jobs: Jobs in queue order. + + Returns: + Jobs without a ``done`` marker whose claim is absent, expired, or + interrupted, in queue order. + """ + done = queue.names("done") + claimed = queue.names("claims") + at = time.time() + pending = [] + for job in jobs: + if job["id"] in done: + continue + if job["id"] in claimed: + lease = queue.read_lease(queue.key("claims", job["id"])) + if lease is not None and not lease.takeable(at): + continue + pending.append(job) + return pending + + +def run_attempt( + queue: Queue, + job: dict[str, Any], + claim: Lease, + slots: list[Lease], + workdir: Path, + runner: Callable[..., int], + settle: Callable[[Path], list[str]], + lease_seconds: float, + renew_seconds: float, + notice: Callable[[], bool], +) -> list[str] | None: + """Run one attempt of a claimed job and record its outcome. + + The attempt writes to ``results//attempt-/``. A takeover (attempt + 2 or later) first downloads the earlier attempts, quiesces their trials, + and passes them to the runner so that completed trials are not repeated. + Trial records are uploaded at every lease renewal. After the attempt, the + job-level report merges every attempt, and the ``done`` marker records the + exit code. + + Args: + queue: Queue of the job. + job: Job dictionary. + claim: Claim lease of the job. + slots: API slot leases held for the job. + workdir: Local working directory. + runner: Function that runs the attempt; see :func:`run_job`. + settle: Function that quiesces an output directory and returns errors. + lease_seconds: Seconds a lease lasts without renewal. + renew_seconds: Seconds between renewals. + notice: Function that returns whether an interruption notice exists. + + Returns: + Errors that leave queries unconfirmed, or None if an interruption + notice stopped the attempt, whose claim is then released as + interrupted. + """ + job_dir = workdir / "results" / job["id"] + attempt = claim.body["attempt"] + output = job_dir / f"attempt-{attempt}" + remote = ("results", job["id"], output.name) + + def progress() -> None: + queue.upload_tree(output, *remote, pattern="*.jsonl") + + with JobWatch(queue, [claim, *slots], lease_seconds, renew_seconds, progress, notice) as watch: + start = now() + resume = [ + job_dir / f"attempt-{earlier}" + for earlier in range(1, attempt) + if queue.download_tree( + job_dir / f"attempt-{earlier}", "results", job["id"], f"attempt-{earlier}" + ) + ] + errors: list[str] = [] + for earlier_output in resume: + try: + errors += settle(earlier_output) + except Exception as exc: + errors.append( + f"Could not quiesce {earlier_output.name}: {type(exc).__name__}: {exc}" + ) + code = 1 + if not errors: + code = runner(job, workdir, output, resume, watch.interrupted) + if watch.interrupted.is_set(): + watch.stop() + queue.upload_tree(output, *remote) + queue.renew(claim, 0, state="interrupted") + return None + try: + errors = settle(output) if code else [] + except Exception as exc: + errors = [f"Could not quiesce the job: {type(exc).__name__}: {exc}"] + queue.upload_tree(output, *remote) + log = workdir / "logs" / f"{job['id']}-{output.name}.log" + if log.exists(): + queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) + try: + report(job_dir) + for name in ("summary.csv", "summary.md"): + queue.s3.upload_file( + str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) + ) + except (ValueError, OSError): + pass # No trial was recorded; the attempt's log explains why. + queue.s3.put_object( + Bucket=queue.bucket, + Key=queue.key("done", job["id"]), + Body=json.dumps( + { + "exit_code": code, + "host": socket.gethostname(), + "attempt": attempt, + "start": start, + "end": now(), + "quiesce_errors": errors, + "lost_leases": watch.lost, + } + ).encode(), + ) + return errors + + def work( queue: Queue, workdir: Path, api_slots: int, poll_seconds: float = 30.0, - runner: Callable[[dict[str, Any], Path], int] = run_job, + runner: Callable[..., int] = run_job, settle: Callable[[Path], list[str]] | None = None, + lease_seconds: float = 600.0, + renew_seconds: float = 60.0, + notice: Callable[[], bool] = lambda: False, + wait: bool = False, ) -> list[str]: - """Claim and run queued jobs until every job has been claimed. + """Claim and run queued jobs until no job is left to claim. - Each host runs one worker, so jobs on a host never overlap. An API-heavy - job starts only after the worker takes one shared slot per simultaneous - paging query; while it waits, no later API-heavy job takes slots. After a - failed job, the worker cancels the job's observed queries and waits until - they stop before claiming another job. + Each host runs one worker, so jobs on a host never overlap. A claim and + each API slot are leases that the worker renews while the job runs; a job + or slot whose lease expired, because its worker died, returns to the + fleet. An API-heavy job starts only after the worker takes one shared slot + per simultaneous paging query; while it waits, no later API-heavy job + takes slots. After a failed job, the worker cancels the job's observed + queries and waits until they stop before claiming another job. Args: queue: Queue to consume. workdir: Local directory for queue files, results, and logs. api_slots: Fleet-wide number of simultaneous GetQueryResults paging queries. - poll_seconds: Wait before rescanning when only slot-limited jobs remain. - runner: Function that runs one job and returns its exit code. + poll_seconds: Wait before rescanning when only slot-limited jobs remain, + and between checks for an unpublished queue. + runner: Function that runs one attempt; see :func:`run_job`. settle: Function that quiesces a failed job, given its output directory, and returns errors; defaults to :func:`quiesce` with the queued configuration. + lease_seconds: Seconds a claim or slot lasts without renewal. + renew_seconds: Seconds between renewals. + notice: Function that returns whether a Spot interruption notice exists. + wait: Whether to wait for the queue to be published instead of failing. Returns: - IDs of the jobs this worker ran. + IDs of the jobs this worker ran to completion. Raises: - ValueError: If the queue has not been completely published, or a job - needs more API slots than ``api_slots``. + ValueError: If the queue has not been completely published and + ``wait`` is false, or a job needs more API slots than ``api_slots``. + SpotInterruptionError: After an interruption notice, once the running attempt's + records are uploaded and its claim is released as interrupted. """ - if not queue.exists("queue.json"): - raise ValueError("Queue is missing or incompletely published") + while not queue.exists("queue.json"): + if not wait: + raise ValueError("Queue is missing or incompletely published") + time.sleep(poll_seconds) workdir.mkdir(parents=True, exist_ok=True) for name in ("config.toml", "manifest.json", "jobs.json"): (workdir / name).write_bytes(queue.read(name)) @@ -547,59 +936,54 @@ def settle(output: Path) -> list[str]: host = socket.gethostname() completed: list[str] = [] while True: - claimed = queue.names("claims") - pending = [j for j in jobs if j["id"] not in claimed] + if notice(): + raise SpotInterruptionError(completed) + pending = pending_jobs(queue, jobs) if not pending: return completed started = False reserved = False for job in pending: owner = {"host": host, "job": job["id"], "at": now()} - slots: list[str] = [] + slots: list[Lease] = [] if job["api_heavy"]: if reserved: continue - acquired = queue.acquire_slots(api_slots, job.get("api_weight", 1), owner) + acquired = queue.acquire_slots( + api_slots, job.get("api_weight", 1), owner, lease_seconds + ) if acquired is None: # Let slots drain for this job instead of starving it with smaller ones. reserved = True continue slots = acquired try: - if not queue.put_if_absent(queue.key("claims", job["id"]), owner): + claim = queue.acquire(queue.key("claims", job["id"]), owner, lease_seconds) + if claim is None: continue - start = now() - code = runner(job, workdir) - try: - settle_errors = settle(workdir / "results" / job["id"]) if code else [] - except Exception as exc: - settle_errors = [f"Could not quiesce the job: {type(exc).__name__}: {exc}"] - queue.upload_tree(workdir / "results" / job["id"], "results", job["id"]) - log = workdir / "logs" / f"{job['id']}.log" - if log.exists(): - queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) - queue.s3.put_object( - Bucket=queue.bucket, - Key=queue.key("done", job["id"]), - Body=json.dumps( - { - "exit_code": code, - "host": host, - "start": start, - "end": now(), - "quiesce_errors": settle_errors, - } - ).encode(), + errors = run_attempt( + queue, + job, + claim, + slots, + workdir, + runner, + settle, + lease_seconds, + renew_seconds, + notice, ) + if errors is None: + raise SpotInterruptionError(completed) completed.append(job["id"]) - if settle_errors: + if errors: # Unconfirmed queries could overlap later trials on this host. return completed started = True break finally: for slot in slots: - queue.release(slot) + queue.release(slot.key) if not started: time.sleep(poll_seconds) @@ -628,34 +1012,96 @@ def instance_metadata(path: str, timeout: float = 2.0) -> str: return response.read().decode() +def interruption_notice( + metadata: Callable[[str], str] | None = None, simulated: Path = SIMULATED_NOTICE +) -> bool: + """Return whether this host has a Spot interruption notice. + + EC2 publishes ``spot/instance-action`` in instance metadata about two + minutes before it stops or terminates a Spot Instance; the path does not + exist otherwise. A file at ``simulated`` counts as a notice, so that a + pilot can rehearse an interruption. + + Args: + metadata: Function that reads an instance metadata path; defaults to + :func:`instance_metadata`. + simulated: File that counts as a notice. + + Returns: + True if a notice exists; False also when instance metadata is unavailable. + """ + if simulated.exists(): + return True + try: + (metadata or instance_metadata)("spot/instance-action") + except (urllib.error.URLError, OSError): + return False + return True + + +def keep_host(metadata: Callable[[str], str]) -> bool: + """Return whether the host's ``KeepHost`` instance tag is ``true``. + + The launch template exposes instance tags in instance metadata. + + Args: + metadata: Function that reads an instance metadata path. + + Returns: + True if the tag exists with the value ``true``. + """ + try: + return metadata(f"tags/instance/{KEEP_HOST_TAG}").strip().lower() == "true" + except (urllib.error.URLError, OSError): + return False + + +def suspend_rebalance(autoscaling: Any, metadata: Callable[[str], str] = instance_metadata) -> str: + """Suspend zone rebalancing in the host's Auto Scaling group. + + With ``AZRebalance`` active, the group may launch a host in one zone and + terminate a busy host in another when Spot capacity or retired hosts leave + the zones uneven. The process stays suspended until the stack is deleted. + + Args: + autoscaling: Auto Scaling client. + metadata: Function that reads an instance metadata path. + + Returns: + Name of the Auto Scaling group. + """ + host = metadata("instance-id") + group = autoscaling.describe_auto_scaling_instances(InstanceIds=[host])["AutoScalingInstances"][ + 0 + ]["AutoScalingGroupName"] + autoscaling.suspend_processes(AutoScalingGroupName=group, ScalingProcesses=["AZRebalance"]) + return str(group) + + def retire_host( queue: Queue, autoscaling: Any, ran: list[str], - keep: Path = KEEP_HOST, metadata: Callable[[str], str] = instance_metadata, ) -> bool: """Record a finished worker and terminate its host unless the host is kept. The record under ``hosts/`` in the queue names the host's instance type, - Availability Zone, and purchase option. Before terminating, the group's - ``AZRebalance`` process is suspended: hosts that retire unevenly across - zones would otherwise let the group replace busy hosts in other zones. - Terminating with a decremented desired capacity keeps the group from - launching a replacement. + Availability Zone, and purchase option. Terminating with a decremented + desired capacity keeps the group from launching a replacement. A host + whose ``KeepHost`` tag is ``true`` keeps running. Args: queue: Queue the worker consumed. autoscaling: Auto Scaling client. ran: IDs of the jobs the worker ran. - keep: File whose presence keeps the host running. metadata: Function that reads an instance metadata path. Returns: True if termination was requested. """ host = metadata("instance-id") - kept = keep.exists() + kept = keep_host(metadata) queue.s3.put_object( Bucket=queue.bucket, Key=queue.key("hosts", f"{host}.json"), @@ -673,10 +1119,6 @@ def retire_host( ) if kept: return False - group = autoscaling.describe_auto_scaling_instances(InstanceIds=[host])["AutoScalingInstances"][ - 0 - ]["AutoScalingGroupName"] - autoscaling.suspend_processes(AutoScalingGroupName=group, ScalingProcesses=["AZRebalance"]) autoscaling.terminate_instance_in_auto_scaling_group( InstanceId=host, ShouldDecrementDesiredCapacity=True ) @@ -684,13 +1126,15 @@ def retire_host( def status(queue: Queue) -> dict[str, Any]: - """Summarize queue progress and failed jobs. + """Summarize queue progress, failed jobs, and leases. Args: queue: Queue to inspect. Returns: - Counts of jobs, claims, and completions, with failed job details. + Counts of jobs, claims, and completions, with failed job details, the + claims of unfinished jobs, and the API slots in use, each lease with + its host, attempt, state, and whether it has expired. """ if not queue.exists("queue.json"): raise ValueError("Queue is missing or incompletely published") @@ -701,10 +1145,25 @@ def status(queue: Queue) -> dict[str, Any]: record = json.loads(queue.read("done", name)) if record["exit_code"] != 0: failed[name] = record + at = time.time() + + def describe(folder: str, names: Iterable[str]) -> dict[str, Any]: + leases = {} + for name in sorted(names): + lease = queue.read_lease(queue.key(folder, name)) + if lease is not None: + leases[name] = { + **{k: lease.body.get(k) for k in ("host", "job", "attempt", "state")}, + "expired": lease.takeable(at) and lease.body.get("state") != "interrupted", + } + return leases + + claims = queue.names("claims") return { "jobs": len(jobs), - "claimed": len(queue.names("claims")), + "claimed": len(claims), "done": len(done), - "slots_in_use": sorted(queue.names("slots")), + "running": describe("claims", claims - done), + "slots_in_use": describe("slots", queue.names("slots")), "failed": failed, } diff --git a/benchmarks/pyathena_bench/report.py b/benchmarks/pyathena_bench/report.py index 311488587..433a0165c 100644 --- a/benchmarks/pyathena_bench/report.py +++ b/benchmarks/pyathena_bench/report.py @@ -144,14 +144,44 @@ def read_records(path: Path) -> list[dict[str, Any]]: return records +def collect_trials(directory: Path) -> list[dict[str, Any]]: + """Read the trials of an output directory and of its attempt subdirectories. + + A fleet job keeps each attempt in an ``attempt-`` subdirectory; their + trials are merged. + A started trial without a final record is reported as incomplete. A failed + or incomplete trial is dropped when another attempt recorded the same + repetition as successful, so a resumed job reports the trials of an + uninterrupted run. + + Args: + directory: Output directory of a run, or a job directory with attempts. + + Returns: + Trial records. + """ + sources = [directory, *sorted(p for p in directory.glob("attempt-*") if p.is_dir())] + trials = [t for source in sources for t in read_records(source / "trials.jsonl")] + recorded = {trial.get("trial") for trial in trials} + trials.extend( + {**event, "status": "incomplete", "error": "Trial started without a final record"} + for source in sources + for event in read_records(source / "events.jsonl") + if event["event"] == "trial_start" and event["trial"] not in recorded + ) + succeeded = { + (t["scale"], t["case_id"], t.get("repetition")) for t in trials if t["status"] == "ok" + } + return [ + t + for t in trials + if t["status"] in {"ok", "unsupported"} + or (t["scale"], t["case_id"], t.get("repetition")) not in succeeded + ] + + def report(directory: Path) -> None: - trials = read_records(directory / "trials.jsonl") - completed = {trial.get("trial") for trial in trials} - for event in read_records(directory / "events.jsonl"): - if event["event"] == "trial_start" and event["trial"] not in completed: - trials.append( - {**event, "status": "incomplete", "error": "Trial started without a final record"} - ) + trials = collect_trials(directory) rows = summarize(trials) if not rows: raise ValueError("No trials to report") diff --git a/benchmarks/pyathena_bench/runner.py b/benchmarks/pyathena_bench/runner.py index a55ac8721..89e3914b7 100644 --- a/benchmarks/pyathena_bench/runner.py +++ b/benchmarks/pyathena_bench/runner.py @@ -35,6 +35,7 @@ from pyathena_bench.aws import cancel_queries, session, validate_inputs, validate_manifest from pyathena_bench.cases import Case from pyathena_bench.config import Settings, select_sql, write_json +from pyathena_bench.report import read_records class Observer: @@ -325,9 +326,49 @@ def environment(settings: Settings) -> dict[str, Any]: } +def completed_trials(directories: list[Path]) -> set[tuple[str, str, int | None]]: + """Collect the trials that earlier attempts of a job recorded as finished. + + Args: + directories: Output directories of earlier attempts. + + Returns: + ``(scale, case_id, repetition)`` for successful trials, and + ``(scale, case_id, None)`` for cases recorded as unsupported. + """ + completed: set[tuple[str, str, int | None]] = set() + for directory in directories: + for trial in read_records(directory / "trials.jsonl"): + if trial["status"] == "ok": + completed.add((trial["scale"], trial["case_id"], trial["repetition"])) + elif trial["status"] == "unsupported": + completed.add((trial["scale"], trial["case_id"], None)) + return completed + + def run( - settings: Settings, manifest: dict[str, Any], cases: list[Case], scales: list[str], output: Path + settings: Settings, + manifest: dict[str, Any], + cases: list[Case], + scales: list[str], + output: Path, + completed: set[tuple[str, str, int | None]] | None = None, ) -> bool: + """Run every trial of the selected cases into a new output directory. + + Args: + settings: Configuration. + manifest: Prepared manifest. + cases: Cases to run. + scales: Scales to run. + output: New output directory. + completed: Trials from :func:`completed_trials` to skip, so that a + resumed job runs only the missing repetitions. + + Returns: + True if every trial succeeded. + """ + completed = completed or set() resources = validate_manifest(settings, manifest) validate_inputs(settings, manifest, scales) output.mkdir(parents=True, exist_ok=False) @@ -337,6 +378,8 @@ def run( for scale in scales: for case in cases: if case.unsupported: + if (scale, case.id, None) in completed: + continue append_json( output / "trials.jsonl", { @@ -349,8 +392,9 @@ def run( ) continue for repetition in range(-settings.warmups, settings.repetitions): + if (scale, case.id, repetition) in completed: + continue trial = uuid.uuid4().hex - prefix = f"runs/{manifest['run_id']}/trials/{trial}/" payload = { "settings": asdict(settings), "case": asdict(case), @@ -360,7 +404,7 @@ def run( "repetition": repetition, "warmup": repetition < 0, "database": resources["ScratchDatabase"], - "output": f"s3://{resources['Bucket']}/{prefix}", + "output": f"{resources['Results']}runs/{manifest['run_id']}/trials/{trial}/", "temp_table": f"b_{manifest['run_id']}_{trial}", "sql": select_sql( resources["ScratchDatabase"], manifest["scales"][scale]["table"], case.shape diff --git a/benchmarks/tests/test_aws.py b/benchmarks/tests/test_aws.py index 0b5ccb6a7..53286d9e9 100644 --- a/benchmarks/tests/test_aws.py +++ b/benchmarks/tests/test_aws.py @@ -10,7 +10,13 @@ import pytest from botocore.exceptions import ClientError -from pyathena_bench.aws import cleanup, prepare, validate_inputs, validate_manifest +from pyathena_bench.aws import ( + cleanup, + prepare, + split_location, + validate_inputs, + validate_manifest, +) from pyathena_bench.config import Settings, read_json RUN = "a" * 32 @@ -24,6 +30,7 @@ "SourceTable": "pypi_file_downloads", "WorkGroup": "pyathena", "AutoScalingGroup": "fleet-test", + "Results": "s3://temporary/", } @@ -284,3 +291,58 @@ def paginator(operation): else: with pytest.raises(ClientError, match="AccessDeniedException"): cleanup(Settings(), data, tmp_path / "manifest.json") + + +def test_cleanup_covers_a_separate_results_location(aws_metadata, monkeypatch, tmp_path): + resources = {**RESOURCES, "Results": "s3://results/bench/stack/"} + monkeypatch.setattr( + "pyathena_bench.aws.stack_resources", lambda session, stack: dict(resources) + ) + actions = [] + scanned = {} + + class Fake: + def get_paginator(self, operation): + def pages(**kwargs): + actions.append((operation, kwargs)) + return { + "get_tables": [{"TableList": []}], + "list_multipart_uploads": [{}], + "list_objects_v2": [{}], + }[operation] + + return SimpleNamespace(paginate=pages) + + def active(athena, workgroup, prefixes): + scanned["prefixes"] = prefixes + return set() + + monkeypatch.setattr("pyathena_bench.aws.client", lambda *args: Fake()) + monkeypatch.setattr("pyathena_bench.aws.active_queries", active) + monkeypatch.setattr("pyathena_bench.aws.cancel_queries", lambda settings, ids: []) + data = manifest() + data["resources"] = dict(resources) + cleanup(Settings(), data, tmp_path / "manifest.json", trials_only=True) + assert scanned["prefixes"] == ( + f"s3://results/bench/stack/runs/{RUN}/", + f"s3://temporary/runs/{RUN}/", + ) + listed = {(k["Bucket"], k["Prefix"]) for op, k in actions if op == "list_objects_v2"} + assert listed == { + ("results", f"bench/stack/runs/{RUN}/trials/"), + ("temporary", f"runs/{RUN}/trials/"), + } + + +@pytest.mark.parametrize( + ("location", "expected"), + [("s3://bucket/", ("bucket", "")), ("s3://bucket/a/b/", ("bucket", "a/b/"))], +) +def test_split_location(location, expected): + assert split_location(location) == expected + + +@pytest.mark.parametrize("location", ["s3://bucket", "s3://bucket/a", "bucket/a/", "s3://B/"]) +def test_split_location_rejects_invalid_locations(location): + with pytest.raises(ValueError, match="Invalid S3 location"): + split_location(location) diff --git a/benchmarks/tests/test_cloudformation.py b/benchmarks/tests/test_cloudformation.py index e2478ff8b..49fdd538a 100644 --- a/benchmarks/tests/test_cloudformation.py +++ b/benchmarks/tests/test_cloudformation.py @@ -23,7 +23,8 @@ def test_template_has_no_retained_storage_or_inbound_ports(): assert launch["MetadataOptions"]["HttpTokens"] == "required" assert "VersioningConfiguration" not in resources["Bucket"]["Properties"] statements = resources["Role"]["Properties"]["Policies"][0]["PolicyDocument"]["Statement"] - for statement in statements: + # Conditional results-bucket statements are checked separately. + for statement in (s for s in statements if "Fn::If" not in s): if "SourceBucket" in str(statement["Resource"]): assert statement["Action"] in {"s3:ListBucket", "s3:GetBucketLocation", "s3:GetObject"} if "SourceDatabase" in str(statement["Resource"]): @@ -72,23 +73,34 @@ def test_architecture_selects_the_image_and_rules_match_instance_families(): assert all(t.startswith(("r7g.", "r8g.")) for t in families["arm64"]) -def test_purchase_option_keeps_one_instance_type_across_three_zones(): +def test_spot_by_default_keeps_one_instance_type_across_three_or_four_zones(): template, errors = decode( str(Path(__file__).resolve().parents[1] / "cloudformation/benchmark.yaml") ) assert not errors resources = template["Resources"] - assert template["Parameters"]["Purchase"]["Default"] == "on-demand" + assert template["Parameters"]["Purchase"]["Default"] == "spot" + assert template["Parameters"]["Zones"]["Default"] == 4 policy = resources["Fleet"]["Properties"]["MixedInstancesPolicy"] assert "Overrides" not in policy["LaunchTemplate"] distribution = policy["InstancesDistribution"] assert distribution["OnDemandBaseCapacity"] == 0 assert distribution["OnDemandPercentageAboveBaseCapacity"] == {"Fn::If": ["IsSpot", 0, 100]} assert distribution["SpotAllocationStrategy"] == "price-capacity-optimized" - subnets = resources["Fleet"]["Properties"]["VPCZoneIdentifier"] - names = [s["Ref"] for s in subnets] - zones = [resources[n]["Properties"]["AvailabilityZone"]["Fn::Select"][0] for n in names] - assert sorted(zones) == [0, 1, 2] + condition, four, three = resources["Fleet"]["Properties"]["VPCZoneIdentifier"]["Fn::If"] + assert condition == "UseFourZones" + for subnets, expected in ((four, [0, 1, 2, 3]), (three, [0, 1, 2])): + names = [s["Ref"] for s in subnets] + zones = [resources[n]["Properties"]["AvailabilityZone"]["Fn::Select"][0] for n in names] + assert sorted(zones) == expected + assert resources["Subnet4"]["Condition"] == "UseFourZones" + assert resources["RouteAssociation4"]["Condition"] == "UseFourZones" + # Hosts launch after every zone's route exists, including the conditional fourth one. + assert {"RouteAssociation", "RouteAssociation2", "RouteAssociation3"} <= set( + resources["Fleet"]["DependsOn"] + ) + [tag] = resources["Fleet"]["Properties"]["Tags"] + assert tag["Value"]["Fn::If"][:2] == ["UseFourZones", {"Ref": "RouteAssociation4"}] interface = resources["LaunchTemplate"]["Properties"]["LaunchTemplateData"][ "NetworkInterfaces" ][0] @@ -103,6 +115,7 @@ def test_hosts_can_terminate_instances_only_in_their_own_stack(): statements = template["Resources"]["Role"]["Properties"]["Policies"][0]["PolicyDocument"][ "Statement" ] + statements = [s for s in statements if "Fn::If" not in s] scoped = [s for s in statements if "autoscaling" in str(s["Action"]) and "Condition" in s] unscoped = [s for s in statements if "autoscaling" in str(s["Action"]) and "Condition" not in s] # Only the read-only lookup of the host's group is unconditional. @@ -117,3 +130,51 @@ def test_hosts_can_terminate_instances_only_in_their_own_stack(): "autoscaling:ResourceTag/aws:cloudformation:stack-id": {"Ref": "AWS::StackId"} } } + + +def test_results_location_is_scoped_to_the_stack_prefix(): + template, errors = decode( + str(Path(__file__).resolve().parents[1] / "cloudformation/benchmark.yaml") + ) + assert not errors + statements = template["Resources"]["Role"]["Properties"]["Policies"][0]["PolicyDocument"][ + "Statement" + ] + results = [s["Fn::If"] for s in statements if "Fn::If" in s] + assert results + assert all(r[0] == "HasResultsBucket" for r in results) + assert all(r[2] == {"Ref": "AWS::NoValue"} for r in results) + for _, statement, _ in results: + if "s3:PutObject" in statement["Action"]: + assert statement["Resource"]["Fn::Sub"].endswith( + ":s3:::${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/*" + ) + elif "s3:ListBucket" in statement["Action"]: + assert statement["Condition"] == { + "StringLike": {"s3:prefix": {"Fn::Sub": "${ResultsPrefix}/${AWS::StackName}/*"}} + } + else: + assert statement["Action"] == "s3:GetBucketLocation" + condition, results_bucket, scratch = template["Outputs"]["Results"]["Value"]["Fn::If"] + assert condition == "HasResultsBucket" + assert results_bucket == { + "Fn::Sub": "s3://${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/" + } + assert scratch == {"Fn::Sub": "s3://${Bucket}/"} + + +def test_hosts_start_the_queue_worker_and_expose_their_tags(): + template, errors = decode( + str(Path(__file__).resolve().parents[1] / "cloudformation/benchmark.yaml") + ) + assert not errors + launch = template["Resources"]["LaunchTemplate"]["Properties"]["LaunchTemplateData"] + assert launch["MetadataOptions"]["InstanceMetadataTags"] == "enabled" + user_data = launch["UserData"]["Fn::Base64"]["Fn::Sub"] + assert "if [ -n '${QueueName}' ]; then" in user_data + assert ( + "ExecStart=/opt/pyathena/.venv/bin/python -m pyathena_bench worker --stack ${AWS::StackId}" + " --name ${QueueName} --api-slots ${ApiSlots} --wait --terminate-host" + ) in user_data + assert "systemctl enable --now pyathena-bench-worker.service" in user_data + assert template["Parameters"]["QueueName"]["Default"] == "" diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 6078ac0c3..41a5188ef 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -5,8 +5,12 @@ # # SPDX-License-Identifier: MIT +import csv import io import json +import threading +import time +import urllib.error from dataclasses import replace from pathlib import Path from types import SimpleNamespace @@ -17,13 +21,17 @@ from pyathena_bench.config import Settings from pyathena_bench.fleet import ( Filters, + JobWatch, Queue, + SpotInterruptionError, expand_jobs, + interruption_notice, merge_jobs, quiesce, retire_host, status, submit, + suspend_rebalance, trial_evidence, work, ) @@ -32,24 +40,38 @@ class FakeS3: - """In-memory S3 with conditional creation, as used by the fleet queue.""" + """In-memory S3 with conditional writes, as used by the fleet queue.""" def __init__(self): self.objects = {} + self.etags = {} + self.writes = 0 def put_object(self, **kwargs): - if kwargs.get("IfNoneMatch") == "*" and kwargs["Key"] in self.objects: + key = kwargs["Key"] + if kwargs.get("IfNoneMatch") == "*" and key in self.objects: raise ClientError({"Error": {"Code": "PreconditionFailed"}}, "PutObject") - self.objects[kwargs["Key"]] = kwargs["Body"] + if "IfMatch" in kwargs and self.etags.get(key) != kwargs["IfMatch"]: + raise ClientError({"Error": {"Code": "PreconditionFailed"}}, "PutObject") + self.writes += 1 + self.objects[key] = kwargs["Body"] + self.etags[key] = f'"{self.writes}"' + return {"ETag": self.etags[key]} def get_object(self, **kwargs): - return {"Body": io.BytesIO(self.objects[kwargs["Key"]])} + if kwargs["Key"] not in self.objects: + raise ClientError({"Error": {"Code": "NoSuchKey"}}, "GetObject") + return { + "Body": io.BytesIO(self.objects[kwargs["Key"]]), + "ETag": self.etags.get(kwargs["Key"], '"0"'), + } def delete_object(self, **kwargs): self.objects.pop(kwargs["Key"], None) + self.etags.pop(kwargs["Key"], None) def upload_file(self, filename, bucket, key): - self.objects[key] = Path(filename).read_bytes() + self.put_object(Key=key, Body=Path(filename).read_bytes()) def list_objects_v2(self, **kwargs): keys = [k for k in self.objects if k.startswith(kwargs["Prefix"])] @@ -212,9 +234,8 @@ def test_workers_run_each_job_once_and_record_results(self, tmp_path): runs = [] settled = [] - def runner(item, workdir): + def runner(item, workdir, out, resume, stop): runs.append(item["id"]) - out = workdir / "results" / item["id"] out.mkdir(parents=True) (out / "summary.csv").write_text("x") return 0 if item["id"] == "light" else 1 @@ -225,7 +246,7 @@ def runner(item, workdir): api_slots=1, poll_seconds=0, runner=runner, - settle=settled.append, + settle=lambda output: settled.append(output) or [], ) assert ( work( @@ -239,8 +260,8 @@ def runner(item, workdir): == [] ) assert sorted(runs) == ["heavy", "light"] - assert [path.name for path in settled] == ["heavy"] - assert queue.key("results", "light", "summary.csv") in s3.objects + assert [path.parent.name for path in settled] == ["heavy"] + assert queue.key("results", "light", "attempt-1", "summary.csv") in s3.objects assert not queue.names("slots") summary = status(queue) assert summary["done"] == 2 @@ -249,10 +270,10 @@ def runner(item, workdir): def test_api_heavy_jobs_wait_for_a_free_slot(self, tmp_path): s3, queue = queue_with([job("heavy", heavy=True), job("light")], tmp_path) other_host = queue.key("slots", "0") - s3.put_object(Bucket="bucket", Key=other_host, Body=b"{}") + assert queue.acquire(other_host, {"host": "other"}, 600) order = [] - def runner(item, workdir): + def runner(item, workdir, *args): order.append(item["id"]) s3.delete_object(Bucket="bucket", Key=other_host) return 0 @@ -281,6 +302,18 @@ def test_duplicate_ids_across_lists_are_rejected(self): } +def metadata_with(**tags): + def read(path): + if path.startswith("tags/instance/"): + name = path.removeprefix("tags/instance/") + if name not in tags: + raise urllib.error.HTTPError(path, 404, "Not Found", None, None) + return tags[name] + return METADATA[path] + + return read + + class FakeAutoScaling: """Auto Scaling client that records calls in order.""" @@ -299,19 +332,22 @@ def terminate_instance_in_auto_scaling_group(self, **kwargs): class TestRetireHost: - def test_finished_worker_records_its_host_and_terminates_it_after_suspending_rebalance( - self, tmp_path - ): - s3, queue = queue_with([job("a")], tmp_path) + def test_worker_start_suspends_zone_rebalancing(self): autoscaling = FakeAutoScaling() - keep = tmp_path / "keep-host" - assert retire_host(queue, autoscaling, ["a"], keep=keep, metadata=METADATA.__getitem__) + assert suspend_rebalance(autoscaling, metadata_with()) == "fleet-group" assert autoscaling.calls == [ ("describe", {"InstanceIds": ["i-123"]}), ( "suspend", {"AutoScalingGroupName": "fleet-group", "ScalingProcesses": ["AZRebalance"]}, ), + ] + + def test_finished_worker_records_its_host_and_terminates_it(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + autoscaling = FakeAutoScaling() + assert retire_host(queue, autoscaling, ["a"], metadata=metadata_with(KeepHost="false")) + assert autoscaling.calls == [ ("terminate", {"InstanceId": "i-123", "ShouldDecrementDesiredCapacity": True}), ] record = json.loads(s3.objects[queue.key("hosts", "i-123.json")]) @@ -323,12 +359,10 @@ def test_finished_worker_records_its_host_and_terminates_it_after_suspending_reb "spot", ) - def test_keep_file_leaves_the_host_and_its_group_unchanged(self, tmp_path): + def test_keep_tag_leaves_the_host_and_its_group_unchanged(self, tmp_path): s3, queue = queue_with([job("a")], tmp_path) - keep = tmp_path / "keep-host" - keep.touch() autoscaling = FakeAutoScaling() - assert not retire_host(queue, autoscaling, [], keep=keep, metadata=METADATA.__getitem__) + assert not retire_host(queue, autoscaling, [], metadata=metadata_with(KeepHost="true")) assert autoscaling.calls == [] assert json.loads(s3.objects[queue.key("hosts", "i-123.json")])["kept"] is True @@ -382,9 +416,8 @@ def test_failed_jobs_are_quiesced_and_unconfirmed_queries_stop_the_worker(self, _, queue = queue_with([job("first"), job("second")], tmp_path) runs = [] - def runner(item, workdir): + def runner(item, workdir, out, resume, stop): runs.append(item["id"]) - out = workdir / "results" / item["id"] out.mkdir(parents=True) (out / "events.jsonl").write_text( json.dumps({"event": "query", "query_id": "q1", "trial": "t"}) + "\n" @@ -404,7 +437,7 @@ def test_waiting_heavy_job_reserves_slots_from_later_heavy_jobs(self, tmp_path, [job("wide", heavy=True, api_weight=2), job("narrow", heavy=True)], tmp_path ) other_host = queue.key("slots", "0") - s3.put_object(Bucket="bucket", Key=other_host, Body=b"{}") + assert queue.acquire(other_host, {"host": "other"}, 600) attempts = [] def sleep_and_release(seconds): @@ -413,7 +446,7 @@ def sleep_and_release(seconds): order = [] - def runner(item, workdir): + def runner(item, workdir, *args): order.append(item["id"]) return 0 @@ -449,11 +482,11 @@ def test_slot_acquisition_errors_release_taken_slots(self, tmp_path): def flaky(**kwargs): if kwargs["Key"].endswith("slots/1"): raise ClientError({"Error": {"Code": "InternalError"}}, "PutObject") - original(**kwargs) + return original(**kwargs) s3.put_object = flaky with pytest.raises(ClientError): - queue.acquire_slots(3, 2, {"host": "h"}) + queue.acquire_slots(3, 2, {"host": "h"}, 600) assert not queue.names("slots") def test_overlapping_submissions_cannot_share_a_name(self, tmp_path): @@ -542,3 +575,191 @@ def settle(output): assert errors == [ "Could not quiesce the job: RuntimeError: Could not inspect query history" ] + + +def ok_trial(repetition): + return { + "trial": f"t{repetition}", + "case_id": "case", + "scale": "small", + "repetition": repetition, + "warmup": False, + "status": "ok", + "queries": [{"total_seconds": 1}], + "successful_queries_per_second": 1, + "rss_peak_bytes": 1, + "max_threads": 1, + } + + +class TestLeases: + def test_only_one_worker_takes_over_an_expired_lease(self, tmp_path): + _, queue = queue_with([], tmp_path) + key = queue.key("claims", "a") + assert queue.acquire(key, {"host": "dead"}, -1).body["attempt"] == 1 + stale = queue.read_lease(key) + first = queue.acquire(key, {"host": "one"}, 600) + assert first.body["attempt"] == 2 + assert queue.acquire(key, {"host": "two"}, 600) is None + # A worker that read the expired lease before the takeover cannot overwrite it. + assert queue.put(key, {"host": "late"}, IfMatch=stale.etag) is None + assert queue.renew(first, 600) + assert not queue.renew(stale, 600) + assert queue.read_lease(key).body["host"] == "one" + + def test_live_claims_are_left_and_expired_ones_return_to_the_fleet(self, tmp_path): + _, queue = queue_with([job("live"), job("expired")], tmp_path) + queue.acquire(queue.key("claims", "live"), {"host": "busy"}, 600) + queue.acquire(queue.key("claims", "expired"), {"host": "dead"}, -1) + runs = [] + + def runner(item, workdir, out, resume, stop): + runs.append((item["id"], out.name, resume)) + return 0 + + ran = work( + queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner, settle=no_queries + ) + assert ran == ["expired"] + assert runs == [("expired", "attempt-2", [])] + running = status(queue)["running"] + assert running["live"] == { + "host": "busy", + "job": None, + "attempt": 1, + "state": "running", + "expired": False, + } + + def test_an_expired_slot_returns_to_the_fleet(self, tmp_path): + _, queue = queue_with([job("heavy", heavy=True)], tmp_path) + queue.acquire(queue.key("slots", "0"), {"host": "dead"}, -1) + assert status(queue)["slots_in_use"]["0"]["expired"] + ran = work( + queue, + tmp_path / "w", + api_slots=1, + poll_seconds=0, + runner=lambda *a: 0, + settle=no_queries, + ) + assert ran == ["heavy"] + assert not queue.names("slots") + + def test_watch_renews_leases_and_uploads_progress(self, tmp_path): + _, queue = queue_with([], tmp_path) + lease = queue.acquire(queue.key("claims", "a"), {"host": "h"}, 1) + expires = lease.body["expires"] + uploads = [] + with JobWatch(queue, [lease], 600, 0.01, lambda: uploads.append(1), lambda: False): + deadline = time.monotonic() + 5 + while not uploads and time.monotonic() < deadline: + time.sleep(0.01) + assert uploads + assert queue.read_lease(lease.key).body["expires"] > expires + 500 + + def test_watch_stops_renewing_a_lease_taken_over_by_another_worker(self, tmp_path): + _, queue = queue_with([], tmp_path) + lease = queue.acquire(queue.key("claims", "a"), {"host": "h"}, -1) + assert queue.acquire(lease.key, {"host": "other"}, 600) + with JobWatch(queue, [lease], 600, 0.01, lambda: None, lambda: False) as watch: + deadline = time.monotonic() + 5 + while not watch.lost and time.monotonic() < deadline: + time.sleep(0.01) + assert watch.lost == [lease.key] + assert queue.read_lease(lease.key).body["host"] == "other" + + +class TestInterruption: + def test_notice_comes_from_instance_metadata_or_a_simulation_file(self, tmp_path): + def missing(path): + raise urllib.error.HTTPError(path, 404, "Not Found", None, None) + + absent = tmp_path / "notice" + assert not interruption_notice(missing, absent) + assert interruption_notice(lambda path: '{"action": "terminate"}', absent) + absent.touch() + assert interruption_notice(missing, absent) + + def test_interrupted_job_is_resumed_by_another_worker_with_the_same_trials(self, tmp_path): + s3, queue = queue_with([job("heavy", heavy=True), job("next")], tmp_path) + notice = threading.Event() + + def interrupted_runner(item, workdir, out, resume, stop): + out.mkdir(parents=True) + (out / "trials.jsonl").write_text(json.dumps(ok_trial(0)) + "\n") + notice.set() + assert stop.wait(5) + with (out / "trials.jsonl").open("a") as stream: + stream.write(json.dumps({**ok_trial(1), "status": "interrupted"}) + "\n") + return 1 + + with pytest.raises(SpotInterruptionError) as raised: + work( + queue, + tmp_path / "w1", + api_slots=1, + poll_seconds=0, + runner=interrupted_runner, + settle=no_queries, + renew_seconds=0.01, + notice=notice.is_set, + ) + assert raised.value.ran == [] + claim = queue.read_lease(queue.key("claims", "heavy")) + assert claim.body["state"] == "interrupted" + assert not queue.names("slots") + assert queue.key("results", "heavy", "attempt-1", "trials.jsonl") in s3.objects + assert not queue.names("done") + + settled = [] + runs = [] + + def resumed_runner(item, workdir, out, resume, stop): + runs.append((item["id"], out.name, [p.name for p in resume])) + out.mkdir(parents=True) + if item["id"] == "heavy": + (out / "trials.jsonl").write_text( + "".join(json.dumps(ok_trial(r)) + "\n" for r in (1, 2)) + ) + return 0 + + ran = work( + queue, + tmp_path / "w2", + api_slots=1, + poll_seconds=0, + runner=resumed_runner, + settle=lambda output: settled.append(output.name) or [], + ) + assert ran == ["heavy", "next"] + assert runs == [("heavy", "attempt-2", ["attempt-1"]), ("next", "attempt-1", [])] + # The interrupted attempt's queries are quiesced before the job continues. + assert settled == ["attempt-1"] + summary = s3.objects[queue.key("results", "heavy", "summary.csv")].decode() + row = next(csv.DictReader(io.StringIO(summary))) + assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") + assert json.loads(queue.read("done", "heavy"))["attempt"] == 2 + + def test_a_notice_between_jobs_stops_the_worker_without_claiming(self, tmp_path): + _, queue = queue_with([job("a")], tmp_path) + with pytest.raises(SpotInterruptionError): + work(queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, notice=lambda: True) + assert not queue.names("claims") + + def test_worker_waits_for_publication(self, tmp_path, monkeypatch): + s3 = FakeS3() + queue = Queue(s3, "bucket", "run1", "results/stack/") + config = tmp_path / "config.toml" + config.write_text("[scales]\nsmall = 10000\n") + manifest = tmp_path / "manifest.json" + manifest.write_text("{}") + monkeypatch.setattr( + "pyathena_bench.fleet.time.sleep", + lambda seconds: submit(queue, [job("a")], config, manifest), + ) + ran = work( + queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, settle=no_queries, wait=True + ) + assert ran == ["a"] + assert all(key.startswith("results/stack/fleet/run1/") for key in s3.objects) diff --git a/benchmarks/tests/test_report.py b/benchmarks/tests/test_report.py index 7e59e598d..16cdea7fc 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -109,3 +109,35 @@ def test_report_exposes_client_init_and_memory_increase(self, tmp_path): assert "Client result (s)" in markdown assert "Init (s)" in markdown assert "RSS increase (B)" in markdown + + +def test_report_merges_attempts_and_drops_superseded_failures(tmp_path): + def trial(repetition, status="ok"): + return { + "trial": f"t{repetition}-{status}", + "case_id": "case", + "scale": "small", + "repetition": repetition, + "warmup": False, + "status": status, + "queries": [{"total_seconds": 1}], + "successful_queries_per_second": 1, + "rss_peak_bytes": 1, + "max_threads": 1, + } + + first = tmp_path / "attempt-1" + second = tmp_path / "attempt-2" + first.mkdir() + second.mkdir() + (first / "trials.jsonl").write_text( + json.dumps(trial(0)) + "\n" + json.dumps(trial(1, "interrupted")) + "\n" + ) + # The interrupted attempt also started a trial that it never recorded. + start = {"event": "trial_start", "trial": "lost", "case_id": "case", "scale": "small"} + (first / "events.jsonl").write_text(json.dumps({**start, "repetition": 2}) + "\n") + (second / "trials.jsonl").write_text(json.dumps(trial(1)) + "\n" + json.dumps(trial(2)) + "\n") + report(tmp_path) + with (tmp_path / "summary.csv").open() as stream: + row = next(csv.DictReader(stream)) + assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") diff --git a/benchmarks/tests/test_runner.py b/benchmarks/tests/test_runner.py index 07edbf1b9..5c22b1b38 100644 --- a/benchmarks/tests/test_runner.py +++ b/benchmarks/tests/test_runner.py @@ -14,7 +14,7 @@ import pytest from pyathena_bench.cases import Case from pyathena_bench.config import Settings -from pyathena_bench.runner import Observer, run, supervise +from pyathena_bench.runner import Observer, completed_trials, run, supervise def successful_worker(pipe, payload): @@ -110,7 +110,7 @@ def test_failed_trial_stops_suite_before_any_more_queries(self, monkeypatch, tmp calls = [] monkeypatch.setattr( "pyathena_bench.runner.validate_manifest", - lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out"}, + lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out", "Results": "s3://out/"}, ) monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) @@ -132,3 +132,40 @@ def fail(payload, events, settings): assert calls[0]["warmup"] result = json.loads((tmp_path / "run/trials.jsonl").read_text()) assert result["status"] == status + + +def test_resumed_run_skips_finished_trials_and_writes_trial_outputs_to_results( + monkeypatch, tmp_path +): + monkeypatch.setattr( + "pyathena_bench.runner.validate_manifest", + lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out", "Results": "s3://res/p/"}, + ) + monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) + monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) + payloads = [] + + def succeed(payload, events, settings): + payloads.append(payload) + return {"status": "ok"} + + monkeypatch.setattr("pyathena_bench.runner.supervise", succeed) + manifest = { + "run_id": "a" * 32, + "scales": {"small": {"table": "input", "rows": 10}}, + "fixtures": {}, + } + case = Case("cursor") + earlier = tmp_path / "attempt-1" + earlier.mkdir() + records = [ + {"scale": "small", "case_id": case.id, "repetition": 0, "status": "ok"}, + {"scale": "small", "case_id": case.id, "repetition": 1, "status": "interrupted"}, + ] + (earlier / "trials.jsonl").write_text("".join(json.dumps(r) + "\n" for r in records)) + settings = replace(Settings(), warmups=0, repetitions=3) + assert run( + settings, manifest, [case], ["small"], tmp_path / "attempt-2", completed_trials([earlier]) + ) + assert [p["repetition"] for p in payloads] == [1, 2] + assert all(p["output"].startswith(f"s3://res/p/runs/{'a' * 32}/trials/") for p in payloads) From 1dfcab197e5c9ffc822fe2ee126affbb1b56508d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 15:26:10 +0900 Subject: [PATCH 02/16] Skip the run after an early notice and keep done markers on report errors A notice that arrived while earlier attempts were being quiesced still started the orchestrator; the attempt is now released as interrupted before it runs, and its log is uploaded with the attempt. A report failure no longer skips the done marker, which would let another worker rerun a finished job; the error is recorded in the marker. The README notes that replacements continue jobs only with QueueName set. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 2 +- benchmarks/pyathena_bench/fleet.py | 32 ++++++++----- benchmarks/tests/test_fleet.py | 72 ++++++++++++++++++++++++++++++ 3 files changed, 93 insertions(+), 13 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 1cfd09fa5..f5a2a7cf2 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -98,7 +98,7 @@ Instance size and disk size are parameters, so a later run can deliberately test The hosts are Spot Instances of one instance type by default; `Purchase=on-demand` launches On-Demand Instances of that type instead. The Auto Scaling group spreads the hosts over `Zones` Availability Zones (4 by default, or 3) with the price-capacity-optimized allocation strategy and never mixes instance types, so every host stays identical. The template uses the first `Zones` zones that `Fn::GetAZs` returns, so deployment fails when it returns fewer, which depends on the zones in which the account's default VPC has default subnets; `us-west-2` has four. -The group replaces an interrupted Spot Instance, and the replacement continues the interrupted job; see [Running on a fleet](#running-on-a-fleet). +The group replaces an interrupted Spot Instance; with `QueueName` set, the replacement's worker continues the interrupted job, see [Running on a fleet](#running-on-a-fleet). `environment.json` does not record the instance type or purchase option; a worker started with `--terminate-host` records both under `hosts/` in its queue, and the recovered `instances.json` lists the hosts that still exist. The AMI parameter resolves the current AL2023 image at deployment; record the resulting AMI when comparing environments. Source bucket access is read-only and restricted to the configured prefix. diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index b1489a72c..73f6a7cd9 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -828,30 +828,37 @@ def progress() -> None: errors.append( f"Could not quiesce {earlier_output.name}: {type(exc).__name__}: {exc}" ) + log = workdir / "logs" / f"{job['id']}-{output.name}.log" + + def upload_attempt() -> None: + queue.upload_tree(output, *remote) + if log.exists(): + queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) + code = 1 - if not errors: + if not errors and not watch.interrupted.is_set(): code = runner(job, workdir, output, resume, watch.interrupted) - if watch.interrupted.is_set(): - watch.stop() - queue.upload_tree(output, *remote) - queue.renew(claim, 0, state="interrupted") - return None + if watch.interrupted.is_set(): + watch.stop() + upload_attempt() + queue.renew(claim, 0, state="interrupted") + return None + if not errors: try: errors = settle(output) if code else [] except Exception as exc: errors = [f"Could not quiesce the job: {type(exc).__name__}: {exc}"] - queue.upload_tree(output, *remote) - log = workdir / "logs" / f"{job['id']}-{output.name}.log" - if log.exists(): - queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) + upload_attempt() + report_error = None try: report(job_dir) for name in ("summary.csv", "summary.md"): queue.s3.upload_file( str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) ) - except (ValueError, OSError): - pass # No trial was recorded; the attempt's log explains why. + except Exception as exc: + # A job without trials, or with unreadable records, still gets its done marker. + report_error = f"{type(exc).__name__}: {exc}" queue.s3.put_object( Bucket=queue.bucket, Key=queue.key("done", job["id"]), @@ -864,6 +871,7 @@ def progress() -> None: "end": now(), "quiesce_errors": errors, "lost_leases": watch.lost, + "report_error": report_error, } ).encode(), ) diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 41a5188ef..3fec519c0 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -8,6 +8,7 @@ import csv import io import json +import signal import threading import time import urllib.error @@ -29,6 +30,7 @@ merge_jobs, quiesce, retire_host, + run_job, status, submit, suspend_rebalance, @@ -741,6 +743,34 @@ def resumed_runner(item, workdir, out, resume, stop): assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") assert json.loads(queue.read("done", "heavy"))["attempt"] == 2 + def test_a_notice_while_quiescing_an_earlier_attempt_skips_the_run(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + queue.acquire(queue.key("claims", "a"), {"host": "dead"}, -1) + s3.put_object(Key=queue.key("results", "a", "attempt-1", "events.jsonl"), Body=b"") + notice = threading.Event() + + def settle(output): + notice.set() + time.sleep(0.2) + return [] + + def runner(*args): + pytest.fail("A job must not start after an interruption notice") + + with pytest.raises(SpotInterruptionError): + work( + queue, + tmp_path / "w", + api_slots=1, + poll_seconds=0, + runner=runner, + settle=settle, + renew_seconds=0.01, + notice=notice.is_set, + ) + lease = queue.read_lease(queue.key("claims", "a")) + assert (lease.body["state"], lease.body["attempt"]) == ("interrupted", 2) + def test_a_notice_between_jobs_stops_the_worker_without_claiming(self, tmp_path): _, queue = queue_with([job("a")], tmp_path) with pytest.raises(SpotInterruptionError): @@ -763,3 +793,45 @@ def test_worker_waits_for_publication(self, tmp_path, monkeypatch): ) assert ran == ["a"] assert all(key.startswith("results/stack/fleet/run1/") for key in s3.objects) + + +class FakeProcess: + """A child process that runs until it receives SIGINT.""" + + def __init__(self, command, stdout, stderr): + self.command = command + self.signals = [] + + def poll(self): + return 130 if self.signals else None + + def send_signal(self, signum): + self.signals.append(signum) + + def wait(self, timeout=None): + return 130 + + +def test_run_job_resumes_earlier_attempts_and_interrupts_with_sigint(tmp_path, monkeypatch): + processes = [] + + def popen(command, stdout, stderr): + processes.append(FakeProcess(command, stdout, stderr)) + return processes[-1] + + monkeypatch.setattr("pyathena_bench.fleet.subprocess.Popen", popen) + stop = threading.Event() + stop.set() + earlier = tmp_path / "results" / "a" / "attempt-1" + code = run_job( + {"id": "a", "args": ["--suite", "single"]}, + tmp_path, + tmp_path / "results" / "a" / "attempt-2", + [earlier], + stop, + ) + assert code == 130 + [process] = processes + assert process.signals == [signal.SIGINT] + assert process.command[-4:] == ["--suite", "single", "--resume", str(earlier)] + assert (tmp_path / "logs" / "a-attempt-2.log").exists() From af75752b77b99a7a9abf1ed680a7642d08e81253 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 15:27:27 +0900 Subject: [PATCH 03/16] Keep a lease whose conditional write succeeded before a retry When the response to a successful conditional write is lost, the SDK retries the write and the retry fails its condition. A renewal then gave up a running job's lease, and a creation or takeover skipped a claim that it held. A failed condition now re-reads the object and keeps the lease when it holds exactly the written body. README wording on rebalance recommendations, the bucket's expiry, lost records, and the pilot's pass criterion is corrected. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 8 ++--- benchmarks/pyathena_bench/fleet.py | 47 +++++++++++++++++++++--------- benchmarks/tests/test_fleet.py | 23 +++++++++++++++ 3 files changed, 61 insertions(+), 17 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index f5a2a7cf2..e4118894e 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -108,7 +108,7 @@ The supplied template assumes ordinary IAM access and SSE-S3 source objects; Lak The fleet queues, reports, and Athena trial outputs go to the stack's results location, the `Results` stack output. With `ResultsBucket` set to an existing bucket, the results location is `s3://///`, and the instance role can write only below it; with `ResultsBucket` empty, the default, it is the scratch bucket. Input snapshots, initialization fixtures, and scratch tables stay in the scratch bucket and database either way. -This repository's CI test bucket, which `AWS_ATHENA_S3_STAGING_DIR` names in CI, expires every object a day after it is written (`cloudformation/github_actions_oidc.yaml`). +This repository's CI test bucket, which `AWS_ATHENA_S3_STAGING_DIR` names in CI, expires every object one day after it is written, at the following midnight UTC (`cloudformation/github_actions_oidc.yaml`). As the results location, it removes leftovers by itself, but reports must be recovered within that day. The deployment identity needs permission to create these resources and pass the instance role. @@ -400,7 +400,7 @@ uv run --env-file ../.env --locked aws ssm send-command --targets "Key=tag:aws:a A worker claims a job by creating `claims/` with an S3 conditional write, so exactly one host runs each job. A claim, like each API slot under `slots/`, is a lease that its worker renews every tenth of `--lease-seconds` (600 by default) while the job runs. When a worker dies, its claim and slots expire, and another worker takes them over with a conditional write, so the job and the slots return to the fleet without manual steps. -Each attempt of a job writes to `results//attempt-/`, and the worker uploads its trial records at every renewal, so a lost host loses at most the trials of the last minutes. +Each attempt of a job writes to `results//attempt-/`, and the worker uploads its trial records at every renewal, so a lost host loses its running trial and at most the records written since the last renewal. After a job, the worker uploads the attempt's output directory and its log to `logs/`, writes `results//summary.csv` and `summary.md` from every attempt, and writes `done/` with the exit code, attempt, host, and times. A failed trial stops only its own job; `status` lists jobs with a non-zero exit code, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. After a failed job, the worker stops the job's reported queries and every queued or running query in the workgroup that writes under the job's trial output prefixes, which also covers queries whose IDs were never reported before a trial process died. @@ -411,7 +411,7 @@ Each worker checks instance metadata for a Spot interruption notice (`spot/insta On a notice, the worker interrupts the running trial with SIGINT, which cancels the trial's reported queries, uploads the attempt, releases its claim as `interrupted` and its slots, and exits without terminating the host. The group replaces the reclaimed instance, and the new host's worker takes over the interrupted claim as the next attempt: it downloads the earlier attempts, quiesces their trials as after a failed job, and runs only the repetitions that they did not record as successful, through `run --resume`. The job's summary merges the attempts and leaves out a failed or incomplete trial whose repetition succeeded in a later attempt, so it reports the trials of an uninterrupted run. -Rebalance recommendations are not acted on; they often precede no interruption, and the notice leaves enough time to stop a trial. +Rebalance recommendations are not acted on; an interruption does not always follow them, and the notice leaves enough time to stop a trial. Workers exit when no job is left to claim: every job is done or claimed by a live lease. With `--terminate-host`, a worker first suspends the group's `AZRebalance` process, which would otherwise let the group launch a host in one zone and terminate a busy host in another when Spot capacity or retired hosts leave the zones uneven; the process stays suspended until the stack is deleted. @@ -462,7 +462,7 @@ uv run --env-file ../.env --locked aws autoscaling terminate-instance-in-auto-sc ``` The group launches a replacement whose worker continues the job as `attempt-2` and then retires the host. -The pilot passes when `done/` records `attempt` 2 and exit code 0, `results//summary.csv` reports six successful trials and no failures, `hosts/` records the replacement with lifecycle `spot`, and the group's desired capacity drops to zero. +The pilot passes when `done/` records `attempt` 2 and exit code 0, every case in `results//summary.csv` reports six successful trials and no failures, `hosts/` records the replacement with lifecycle `spot`, and the group's desired capacity drops to zero. AWS Fault Injection Service (`aws:ec2:send-spot-instance-interruptions`) sends a real notice instead, but needs its own IAM role and experiment template. ## Reports and recovery diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index 73f6a7cd9..efb569a2f 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -335,6 +335,29 @@ def read_lease(self, key: str) -> Lease | None: raise return Lease(key, json.loads(response["Body"].read()), str(response["ETag"])) + def put_lease(self, key: str, body: dict[str, Any], **condition: str) -> Lease | None: + """Write a lease conditionally, recognizing a write that already succeeded. + + When the response to a successful conditional write is lost, the SDK + retries it, and the retry fails its condition. An object that holds + exactly ``body`` is therefore this worker's own write. + + Args: + key: Object key. + body: Lease content, unique through its expiry time. + **condition: ``IfNoneMatch="*"`` or ``IfMatch=``. + + Returns: + The lease, or None if the condition failed for another writer's object. + """ + etag = self.put(key, body, **condition) + if etag is not None: + return Lease(key, body, etag) + current = self.read_lease(key) + if current is not None and current.body == body: + return current + return None + def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Lease | None: """Take a lease on an object that is absent, expired, or interrupted. @@ -350,16 +373,14 @@ def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Leas The lease, or None if another worker holds it. """ body = {**owner, "attempt": 1, "state": "running", "expires": time.time() + lease_seconds} - etag = self.put(key, body, IfNoneMatch="*") - if etag is None: - current = self.read_lease(key) - if current is None or not current.takeable(time.time()): - return None - body["attempt"] = current.body.get("attempt", 1) + 1 - etag = self.put(key, body, IfMatch=current.etag) - if etag is None: - return None - return Lease(key, body, etag) + lease = self.put_lease(key, body, IfNoneMatch="*") + if lease is not None: + return lease + current = self.read_lease(key) + if current is None or not current.takeable(time.time()): + return None + body["attempt"] = current.body.get("attempt", 1) + 1 + return self.put_lease(key, body, IfMatch=current.etag) def renew(self, lease: Lease, lease_seconds: float, **changes: Any) -> bool: """Extend or update a lease if this worker still holds it. @@ -373,10 +394,10 @@ def renew(self, lease: Lease, lease_seconds: float, **changes: Any) -> bool: False if another worker took over the lease. """ body = {**lease.body, "expires": time.time() + lease_seconds, **changes} - etag = self.put(lease.key, body, IfMatch=lease.etag) - if etag is None: + written = self.put_lease(lease.key, body, IfMatch=lease.etag) + if written is None: return False - lease.body, lease.etag = body, etag + lease.body, lease.etag = written.body, written.etag return True def names(self, folder: str) -> set[str]: diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 3fec519c0..bd24398d8 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -609,6 +609,29 @@ def test_only_one_worker_takes_over_an_expired_lease(self, tmp_path): assert not queue.renew(stale, 600) assert queue.read_lease(key).body["host"] == "one" + @pytest.mark.parametrize("operation", ["create", "takeover", "renew"]) + def test_a_retried_write_that_already_succeeded_keeps_the_lease(self, tmp_path, operation): + s3, queue = queue_with([], tmp_path) + key = queue.key("claims", "a") + if operation != "create": + lease = queue.acquire(key, {"host": "h"}, -1 if operation == "takeover" else 600) + original = s3.put_object + + def lost_response(**kwargs): + # The first write succeeds, but its response is lost and the SDK retries. + original(**kwargs) + return original(**kwargs) + + s3.put_object = lost_response + if operation == "renew": + assert queue.renew(lease, 600) + assert queue.read_lease(key).etag == lease.etag + else: + taken = queue.acquire(key, {"host": "me"}, 600) + assert taken.body["host"] == "me" + assert taken.body["attempt"] == (1 if operation == "create" else 2) + assert queue.read_lease(key).etag == taken.etag + def test_live_claims_are_left_and_expired_ones_return_to_the_fleet(self, tmp_path): _, queue = queue_with([job("live"), job("expired")], tmp_path) queue.acquire(queue.key("claims", "live"), {"host": "busy"}, 600) From 4b06a6c3098aaed93d85b0652ff64d6078b115b7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 15:53:41 +0900 Subject: [PATCH 04/16] Stop stale lease owners and wait for dead workers' claims Independent review findings: - Workers started during stack creation now wait for the stack's outputs (worker --wait) instead of exiting. - A worker whose lease was taken over, or whose renewals failed until it expired, stops the attempt and publishes no summary or done marker; a final conditional renewal precedes the done marker. - Slots are released by overwriting the held lease, not by deleting the key, so another worker's takeover is kept. - Workers do not exit while an unfinished job's claim has missed renewals; they wait to take it over when it expires. - A separate thread watches for interruption notices; on a notice the claim is released as interrupted before the attempt is uploaded, even after a concurrent renewal. - Checkpoint and final uploads are serialized, and changed trial records are checkpointed every few seconds. - A resumed run with no trials left exits successfully. - ListBucketMultipartUploads on the results bucket is granted without the s3:prefix condition, which it does not support. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 11 +- benchmarks/cloudformation/benchmark.yaml | 5 +- benchmarks/pyathena_bench/__main__.py | 34 ++- benchmarks/pyathena_bench/aws.py | 6 +- benchmarks/pyathena_bench/fleet.py | 318 ++++++++++++++++------- benchmarks/tests/test_cloudformation.py | 4 +- benchmarks/tests/test_fleet.py | 132 +++++++++- 7 files changed, 399 insertions(+), 111 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index e4118894e..cd06c7055 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -302,7 +302,7 @@ With `--split-pages`, a job whose single trial needs at least that many pages be `queue --jobs` accepts several job files, rejects duplicate job IDs, and orders the combined jobs as `jobs` orders one selection. Deploy a new stack with the fleet size under its own name; changing `FleetSize` on an existing stack does not wait for the added hosts to finish bootstrap. -`QueueName` names the queue that every host consumes: after bootstrap, each host starts a worker that waits until a queue of that name is published, including the hosts that the group launches later to replace interrupted Spot Instances. +`QueueName` names the queue that every host consumes: after bootstrap, each host starts a worker that waits for the stack's outputs and then until a queue of that name is published, including the hosts that the group launches later to replace interrupted Spot Instances. `ApiSlots` sets their `--api-slots`. `ResultsBucket` selects the results location; the commands below use this repository's CI test bucket, whose objects expire after a day. Then open a session on one of its hosts: @@ -400,7 +400,10 @@ uv run --env-file ../.env --locked aws ssm send-command --targets "Key=tag:aws:a A worker claims a job by creating `claims/` with an S3 conditional write, so exactly one host runs each job. A claim, like each API slot under `slots/`, is a lease that its worker renews every tenth of `--lease-seconds` (600 by default) while the job runs. When a worker dies, its claim and slots expire, and another worker takes them over with a conditional write, so the job and the slots return to the fleet without manual steps. -Each attempt of a job writes to `results//attempt-/`, and the worker uploads its trial records at every renewal, so a lost host loses its running trial and at most the records written since the last renewal. +A worker releases a slot by overwriting its lease as released, never by deleting it, so a slot that another worker took over stays held. +A worker whose lease was taken over, or whose renewals failed until the lease expired, stops the attempt and records nothing for the job; the new owner continues it. +Each attempt of a job writes to `results//attempt-/`, and the worker uploads changed trial records every five seconds, so a lost host loses its running trial and at most the records of the last seconds. +A trial whose start was not uploaded before its host died is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. After a job, the worker uploads the attempt's output directory and its log to `logs/`, writes `results//summary.csv` and `summary.md` from every attempt, and writes `done/` with the exit code, attempt, host, and times. A failed trial stops only its own job; `status` lists jobs with a non-zero exit code, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. After a failed job, the worker stops the job's reported queries and every queued or running query in the workgroup that writes under the job's trial output prefixes, which also covers queries whose IDs were never reported before a trial process died. @@ -408,12 +411,14 @@ It waits until they stop before claiming another job; this scans the workgroup's If it cannot confirm that, it records the errors in the `done` marker and exits. Each worker checks instance metadata for a Spot interruption notice (`spot/instance-action`), which EC2 publishes about two minutes before it reclaims an instance. -On a notice, the worker interrupts the running trial with SIGINT, which cancels the trial's reported queries, uploads the attempt, releases its claim as `interrupted` and its slots, and exits without terminating the host. +On a notice, the worker interrupts the running trial with SIGINT, which cancels the trial's reported queries, releases its claim as `interrupted`, uploads the attempt, releases its slots, and exits without terminating the host. +A separate thread watches for the notice, so slow S3 requests do not delay it. The group replaces the reclaimed instance, and the new host's worker takes over the interrupted claim as the next attempt: it downloads the earlier attempts, quiesces their trials as after a failed job, and runs only the repetitions that they did not record as successful, through `run --resume`. The job's summary merges the attempts and leaves out a failed or incomplete trial whose repetition succeeded in a later attempt, so it reports the trials of an uninterrupted run. Rebalance recommendations are not acted on; an interruption does not always follow them, and the notice leaves enough time to stop a trial. Workers exit when no job is left to claim: every job is done or claimed by a live lease. +A claim that has missed renewals belongs to a worker that has probably died, so workers keep waiting until it expires and one of them takes it over. With `--terminate-host`, a worker first suspends the group's `AZRebalance` process, which would otherwise let the group launch a host in one zone and terminate a busy host in another when Spot capacity or retired hosts leave the zones uneven; the process stays suspended until the stack is deleted. A worker with `--terminate-host` that returns normally records its jobs, instance type, Availability Zone, and purchase option under `hosts/.json` and terminates its host through the group with a lower desired capacity, so no replacement starts and finished hosts stop costing money while others still run. This includes a worker that stopped after unconfirmed queries, whose errors are already in the `done` marker. diff --git a/benchmarks/cloudformation/benchmark.yaml b/benchmarks/cloudformation/benchmark.yaml index b94c3330f..bbaf4646f 100644 --- a/benchmarks/cloudformation/benchmark.yaml +++ b/benchmarks/cloudformation/benchmark.yaml @@ -300,16 +300,17 @@ Resources: 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. + # ListBucketMultipartUploads does not support the s3:prefix condition key. - !If - HasResultsBucket - Effect: Allow - Action: s3:GetBucketLocation + Action: [s3:GetBucketLocation, s3:ListBucketMultipartUploads] Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket} - !Ref AWS::NoValue - !If - HasResultsBucket - Effect: Allow - Action: [s3:ListBucket, s3:ListBucketMultipartUploads] + Action: s3:ListBucket Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket} Condition: StringLike: diff --git a/benchmarks/pyathena_bench/__main__.py b/benchmarks/pyathena_bench/__main__.py index a802725ae..06a2c1089 100644 --- a/benchmarks/pyathena_bench/__main__.py +++ b/benchmarks/pyathena_bench/__main__.py @@ -13,10 +13,13 @@ import json import logging import sys +import time from dataclasses import asdict, replace from pathlib import Path +from typing import Any from pyathena_bench.aws import ( + IncompleteStackError, cleanup, client, preflight, @@ -150,6 +153,31 @@ def parser() -> argparse.ArgumentParser: return root +def wait_for_stack( + session_: Any, stack: str, wait: bool, poll_seconds: float = 30.0 +) -> dict[str, str]: + """Read the stack's resources, waiting for its outputs if requested. + + A host's worker starts during stack creation, before the outputs exist. + + Args: + session_: Boto3 session. + stack: Stack name or ID. + wait: Whether to wait while the outputs are incomplete. + poll_seconds: Seconds between attempts. + + Returns: + The stack resources. + """ + while True: + try: + return stack_resources(session_, stack) + except IncompleteStackError: + if not wait: + raise + time.sleep(poll_seconds) + + def main(argv: list[str] | None = None) -> int: args = parser().parse_args(argv) if args.command == "report": @@ -193,7 +221,8 @@ 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) - bucket, base = split_location(stack_resources(session_, args.stack)["Results"]) + resources = wait_for_stack(session_, args.stack, getattr(args, "wait", False)) + bucket, base = split_location(resources["Results"]) queue = Queue(client(session_, "s3"), bucket, args.name, base) if args.command == "queue": validate_manifest(settings, read_json(args.manifest)) @@ -293,6 +322,9 @@ def main(argv: list[str] | None = None) -> int: success = run( settings, manifest, cases, scales, args.out, completed_trials(args.resume) ) + if args.resume and not (args.out / "trials.jsonl").exists(): + # Earlier attempts recorded every trial; the job's report merges them. + return 0 if success else 1 report(args.out) return 0 if success else 1 return 0 diff --git a/benchmarks/pyathena_bench/aws.py b/benchmarks/pyathena_bench/aws.py index 516d40e4e..b7b190b20 100644 --- a/benchmarks/pyathena_bench/aws.py +++ b/benchmarks/pyathena_bench/aws.py @@ -37,6 +37,10 @@ def client(session_: Any, service: str) -> Any: ) +class IncompleteStackError(ValueError): + """The stack has no complete outputs yet, for example while it is being created.""" + + def split_location(location: str) -> tuple[str, str]: """Split an ``s3://bucket/prefix/`` location. @@ -73,7 +77,7 @@ def stack_resources(session_: Any, stack: str) -> dict[str, str]: "Results", } if not required <= result.keys(): - raise ValueError("Stack outputs are incomplete") + raise IncompleteStackError("Stack outputs are incomplete") if ( result["SourceBucket"] == result["Bucket"] or result["ScratchDatabase"] == result["SourceDatabase"] diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index efb569a2f..c73f4042b 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -232,11 +232,12 @@ def now() -> str: class Lease: """A claim or API slot object held by one worker until it expires. - The body records the owner, the attempt number, the state (``running`` or - ``interrupted``), and the expiry as a Unix time. A worker renews its leases + The body records the owner, the attempt number, the state (``running``, + ``interrupted``, or ``released``), and the times of the last write + (``renewed``) and of the expiry, as Unix times. A worker renews its leases while it runs a job; another worker may take over a lease that has expired - or was released as interrupted. Objects without an expiry, written by - earlier versions, never expire. + or was released. Objects without an expiry, written by earlier versions, + never expire. """ key: str @@ -253,7 +254,28 @@ def takeable(self, at: float) -> bool: True if the lease was interrupted or expired before ``at``. """ expires = self.body.get("expires") - return self.body.get("state") == "interrupted" or (expires is not None and expires < at) + return self.body.get("state") in {"interrupted", "released"} or ( + expires is not None and expires < at + ) + + def stale(self, at: float) -> bool: + """Return whether a running lease missed renewals but has not expired yet. + + A live worker renews a lease ten times per lease period, so a lease + that has gone a fifth of its period without renewal belongs to a + worker that has probably died. + + Args: + at: Current Unix time. + + Returns: + True if the lease is running, not takeable, and overdue for renewal. + """ + renewed = self.body.get("renewed") + expires = self.body.get("expires") + if renewed is None or expires is None or self.takeable(at): + return False + return at - renewed > (expires - renewed) / 5 class Queue: @@ -372,7 +394,14 @@ def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Leas Returns: The lease, or None if another worker holds it. """ - body = {**owner, "attempt": 1, "state": "running", "expires": time.time() + lease_seconds} + at = time.time() + body = { + **owner, + "attempt": 1, + "state": "running", + "renewed": at, + "expires": at + lease_seconds, + } lease = self.put_lease(key, body, IfNoneMatch="*") if lease is not None: return lease @@ -393,13 +422,49 @@ def renew(self, lease: Lease, lease_seconds: float, **changes: Any) -> bool: Returns: False if another worker took over the lease. """ - body = {**lease.body, "expires": time.time() + lease_seconds, **changes} + at = time.time() + body = {**lease.body, "renewed": at, "expires": at + lease_seconds, **changes} written = self.put_lease(lease.key, body, IfMatch=lease.etag) if written is None: return False lease.body, lease.etag = written.body, written.etag return True + def release_lease(self, lease: Lease) -> bool: + """Release a lease that this worker still holds, so another worker can take it. + + The object is overwritten as released rather than deleted, so that a + lease another worker took over in the meantime stays intact. + + Args: + lease: Lease to release. + + Returns: + False if another worker took over the lease. + """ + return self.renew(lease, 0, state="released", expires=0) + + def mark_interrupted(self, lease: Lease) -> bool: + """Release a claim as interrupted, so another worker continues the job at once. + + If the lease was renewed concurrently with a newer ETag, the current + object is still released when it belongs to the same owner and attempt. + + Args: + lease: Claim lease of the interrupted attempt. + + Returns: + False if another worker took over the claim. + """ + if self.renew(lease, 0, state="interrupted", expires=0): + return True + current = self.read_lease(lease.key) + if current is None or any( + current.body.get(k) != lease.body.get(k) for k in ("host", "job", "attempt") + ): + return False + return self.renew(current, 0, state="interrupted", expires=0) + def names(self, folder: str) -> set[str]: """List the object names directly below a queue folder. @@ -461,10 +526,10 @@ def acquire_slots( return taken except BaseException: for lease in taken: - self.release(lease.key) + self.release_lease(lease) raise for lease in taken: - self.release(lease.key) + self.release_lease(lease) return None def exists(self, *parts: str) -> bool: @@ -691,10 +756,13 @@ def __init__(self, ran: list[str]) -> None: class JobWatch: """Renew a running job's leases, upload its progress, and watch for interruption. - A background thread checks for an interruption notice every - ``poll_seconds`` and, every ``renew_seconds``, renews the leases and calls - ``progress``. A lease that another worker took over is no longer renewed and - is listed in ``lost``. Transient errors are retried at the next renewal. + One thread checks for an interruption notice every ``poll_seconds``. + Another calls ``progress`` every ``poll_seconds`` and renews the leases + every ``renew_seconds``, so a slow S3 request does not delay the notice. + ``stop_requested`` is set on a notice (with ``interrupted``) and when a + lease is lost: another worker took it over, or renewals failed until it + expired. A lost lease is listed in ``lost`` and no longer renewed. + ``upload_lock`` serializes uploads of the attempt's files. """ def __init__( @@ -715,12 +783,18 @@ def __init__( self.notice = notice self.poll_seconds = min(poll_seconds, renew_seconds) self.interrupted = threading.Event() + self.stop_requested = threading.Event() self.lost: list[str] = [] + self.upload_lock = threading.Lock() self._stopped = threading.Event() - self._thread = threading.Thread(target=self._run, daemon=True) + self._threads = [ + threading.Thread(target=self._watch_notice, daemon=True), + threading.Thread(target=self._renew, daemon=True), + ] def __enter__(self) -> JobWatch: - self._thread.start() + for thread in self._threads: + thread.start() return self def __exit__( @@ -731,36 +805,51 @@ def __exit__( ) -> None: self.stop() - def stop(self) -> None: - """Stop the thread and wait for it, so the caller may update the leases.""" + def stop(self, timeout: float | None = None) -> None: + """Stop the threads and wait for them. + + Args: + timeout: Seconds to wait for each thread, or None to wait until it + ends, after which the caller may update the leases. + """ self._stopped.set() - if self._thread.is_alive(): - self._thread.join() + for thread in self._threads: + if thread.is_alive(): + thread.join(timeout) - def _run(self) -> None: - renew_at = time.monotonic() + self.renew_seconds + def _watch_notice(self) -> None: while not self._stopped.wait(self.poll_seconds): - if not self.interrupted.is_set() and self.notice(): + if self.notice(): self.interrupted.set() + self.stop_requested.set() + return + + def _renew(self) -> None: + renew_at = time.monotonic() + self.renew_seconds + while not self._stopped.wait(self.poll_seconds): + try: + with self.upload_lock: + self.progress() + except Exception: + # The next poll uploads again. + pass if time.monotonic() < renew_at: continue renew_at = time.monotonic() + self.renew_seconds for lease in list(self.leases): try: - if not self.queue.renew(lease, self.lease_seconds): - self.leases.remove(lease) - self.lost.append(lease.key) + if self.queue.renew(lease, self.lease_seconds): + continue except Exception: - # The lease has time left for a retry. - continue - try: - self.progress() - except Exception: - # The next renewal uploads again. - continue + if time.time() < lease.body["expires"]: + # The lease has time left for a retry. + continue + self.leases.remove(lease) + self.lost.append(lease.key) + self.stop_requested.set() -def pending_jobs(queue: Queue, jobs: list[dict[str, Any]]) -> list[dict[str, Any]]: +def pending_jobs(queue: Queue, jobs: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], bool]: """Return the jobs that are neither done nor held by a live claim. Args: @@ -769,21 +858,39 @@ def pending_jobs(queue: Queue, jobs: list[dict[str, Any]]) -> list[dict[str, Any Returns: Jobs without a ``done`` marker whose claim is absent, expired, or - interrupted, in queue order. + released, in queue order, and whether another unfinished job's claim + is stale: its worker has probably died, and the job returns to the + fleet when the claim expires. """ done = queue.names("done") claimed = queue.names("claims") at = time.time() pending = [] + stale = False for job in jobs: if job["id"] in done: continue if job["id"] in claimed: lease = queue.read_lease(queue.key("claims", job["id"])) if lease is not None and not lease.takeable(at): + stale = stale or lease.stale(at) continue pending.append(job) - return pending + return pending, stale + + +@dataclass(frozen=True) +class Outcome: + """Result of one job attempt. + + ``state`` is ``done`` when the attempt recorded the job's outcome, + ``interrupted`` after a Spot interruption notice, and ``lost`` when another + worker took over the job's claim or a slot. ``errors`` lists queries that + could not be confirmed as stopped. + """ + + state: str + errors: tuple[str, ...] = () def run_attempt( @@ -797,15 +904,17 @@ def run_attempt( lease_seconds: float, renew_seconds: float, notice: Callable[[], bool], -) -> list[str] | None: +) -> Outcome: """Run one attempt of a claimed job and record its outcome. The attempt writes to ``results//attempt-/``. A takeover (attempt 2 or later) first downloads the earlier attempts, quiesces their trials, and passes them to the runner so that completed trials are not repeated. - Trial records are uploaded at every lease renewal. After the attempt, the - job-level report merges every attempt, and the ``done`` marker records the - exit code. + Changed trial records are uploaded every few seconds. An interruption + notice or a lost lease stops the runner; after a notice, the claim is + released as interrupted before the attempt is uploaded. Otherwise, after a + final check that the claim is still held, the job-level report merges + every attempt, and the ``done`` marker records the exit code. Args: queue: Queue of the job. @@ -820,19 +929,30 @@ def run_attempt( notice: Function that returns whether an interruption notice exists. Returns: - Errors that leave queries unconfirmed, or None if an interruption - notice stopped the attempt, whose claim is then released as - interrupted. + The outcome of the attempt. """ job_dir = workdir / "results" / job["id"] attempt = claim.body["attempt"] output = job_dir / f"attempt-{attempt}" remote = ("results", job["id"], output.name) + log = workdir / "logs" / f"{job['id']}-{output.name}.log" + uploaded: dict[str, int] = {} def progress() -> None: - queue.upload_tree(output, *remote, pattern="*.jsonl") + for path in sorted(output.glob("*.jsonl")): + size = path.stat().st_size + if uploaded.get(path.name) != size: + queue.s3.upload_file(str(path), queue.bucket, queue.key(*remote, path.name)) + uploaded[path.name] = size with JobWatch(queue, [claim, *slots], lease_seconds, renew_seconds, progress, notice) as watch: + + def upload_attempt() -> None: + with watch.upload_lock: + queue.upload_tree(output, *remote) + if log.exists(): + queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) + start = now() resume = [ job_dir / f"attempt-{earlier}" @@ -849,54 +969,52 @@ def progress() -> None: errors.append( f"Could not quiesce {earlier_output.name}: {type(exc).__name__}: {exc}" ) - log = workdir / "logs" / f"{job['id']}-{output.name}.log" - - def upload_attempt() -> None: - queue.upload_tree(output, *remote) - if log.exists(): - queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) - code = 1 - if not errors and not watch.interrupted.is_set(): - code = runner(job, workdir, output, resume, watch.interrupted) - if watch.interrupted.is_set(): - watch.stop() + if not errors and not watch.stop_requested.is_set(): + code = runner(job, workdir, output, resume, watch.stop_requested) + if watch.stop_requested.is_set(): + # Within the two minutes of a notice, release the claim before uploading. + watch.stop(timeout=10) + interrupted = watch.interrupted.is_set() + if interrupted: + queue.mark_interrupted(claim) upload_attempt() - queue.renew(claim, 0, state="interrupted") - return None + return Outcome("interrupted" if interrupted else "lost") if not errors: try: errors = settle(output) if code else [] except Exception as exc: errors = [f"Could not quiesce the job: {type(exc).__name__}: {exc}"] upload_attempt() - report_error = None - try: - report(job_dir) - for name in ("summary.csv", "summary.md"): - queue.s3.upload_file( - str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) - ) - except Exception as exc: - # A job without trials, or with unreadable records, still gets its done marker. - report_error = f"{type(exc).__name__}: {exc}" - queue.s3.put_object( - Bucket=queue.bucket, - Key=queue.key("done", job["id"]), - Body=json.dumps( - { - "exit_code": code, - "host": socket.gethostname(), - "attempt": attempt, - "start": start, - "end": now(), - "quiesce_errors": errors, - "lost_leases": watch.lost, - "report_error": report_error, - } - ).encode(), - ) - return errors + if watch.lost or not queue.renew(claim, lease_seconds): + # Another worker continues the job and records its outcome. + return Outcome("lost") + report_error = None + try: + report(job_dir) + for name in ("summary.csv", "summary.md"): + queue.s3.upload_file( + str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) + ) + except Exception as exc: + # A job without trials, or with unreadable records, still gets its done marker. + report_error = f"{type(exc).__name__}: {exc}" + queue.s3.put_object( + Bucket=queue.bucket, + Key=queue.key("done", job["id"]), + Body=json.dumps( + { + "exit_code": code, + "host": socket.gethostname(), + "attempt": attempt, + "start": start, + "end": now(), + "quiesce_errors": errors, + "report_error": report_error, + } + ).encode(), + ) + return Outcome("done", tuple(errors)) def work( @@ -916,10 +1034,12 @@ def work( Each host runs one worker, so jobs on a host never overlap. A claim and each API slot are leases that the worker renews while the job runs; a job or slot whose lease expired, because its worker died, returns to the - fleet. An API-heavy job starts only after the worker takes one shared slot - per simultaneous paging query; while it waits, no later API-heavy job - takes slots. After a failed job, the worker cancels the job's observed - queries and waits until they stop before claiming another job. + fleet. While another unfinished job's claim is stale, the worker waits to + take it over instead of exiting. An API-heavy job starts only after the + worker takes one shared slot per simultaneous paging query; while it + waits, no later API-heavy job takes slots. After a failed job, the worker + cancels the job's observed queries and waits until they stop before + claiming another job. Args: queue: Queue to consume. @@ -967,8 +1087,8 @@ def settle(output: Path) -> list[str]: while True: if notice(): raise SpotInterruptionError(completed) - pending = pending_jobs(queue, jobs) - if not pending: + pending, stale = pending_jobs(queue, jobs) + if not pending and not stale: return completed started = False reserved = False @@ -990,7 +1110,7 @@ def settle(output: Path) -> list[str]: claim = queue.acquire(queue.key("claims", job["id"]), owner, lease_seconds) if claim is None: continue - errors = run_attempt( + outcome = run_attempt( queue, job, claim, @@ -1002,17 +1122,18 @@ def settle(output: Path) -> list[str]: renew_seconds, notice, ) - if errors is None: + if outcome.state == "interrupted": raise SpotInterruptionError(completed) - completed.append(job["id"]) - if errors: - # Unconfirmed queries could overlap later trials on this host. - return completed + if outcome.state == "done": + completed.append(job["id"]) + if outcome.errors: + # Unconfirmed queries could overlap later trials on this host. + return completed started = True break finally: for slot in slots: - queue.release(slot.key) + queue.release_lease(slot) if not started: time.sleep(poll_seconds) @@ -1180,10 +1301,11 @@ def describe(folder: str, names: Iterable[str]) -> dict[str, Any]: leases = {} for name in sorted(names): lease = queue.read_lease(queue.key(folder, name)) - if lease is not None: + if lease is not None and lease.body.get("state") != "released": leases[name] = { **{k: lease.body.get(k) for k in ("host", "job", "attempt", "state")}, - "expired": lease.takeable(at) and lease.body.get("state") != "interrupted", + "expired": lease.takeable(at) and lease.body.get("state") == "running", + "stale": lease.stale(at), } return leases diff --git a/benchmarks/tests/test_cloudformation.py b/benchmarks/tests/test_cloudformation.py index 49fdd538a..74b8c3ff5 100644 --- a/benchmarks/tests/test_cloudformation.py +++ b/benchmarks/tests/test_cloudformation.py @@ -154,7 +154,9 @@ def test_results_location_is_scoped_to_the_stack_prefix(): "StringLike": {"s3:prefix": {"Fn::Sub": "${ResultsPrefix}/${AWS::StackName}/*"}} } else: - assert statement["Action"] == "s3:GetBucketLocation" + # Neither action supports the s3:prefix condition key. + assert statement["Action"] == ["s3:GetBucketLocation", "s3:ListBucketMultipartUploads"] + assert "Condition" not in statement condition, results_bucket, scratch = template["Outputs"]["Results"]["Value"]["Fn::If"] assert condition == "HasResultsBucket" assert results_bucket == { diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index bd24398d8..d931a9d78 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -18,11 +18,13 @@ import pytest from botocore.exceptions import ClientError -from pyathena_bench.__main__ import main +from pyathena_bench.__main__ import main, wait_for_stack +from pyathena_bench.aws import IncompleteStackError from pyathena_bench.config import Settings from pyathena_bench.fleet import ( Filters, JobWatch, + Lease, Queue, SpotInterruptionError, expand_jobs, @@ -113,6 +115,15 @@ def no_queries(ids): return [] +def live_slots(queue): + """Return the API slots that a worker holds; released slots stay as objects.""" + return [ + name + for name in queue.names("slots") + if not queue.read_lease(queue.key("slots", name)).takeable(time.time()) + ] + + class TestExpandJobs: def test_row_cursor_jobs_are_split_by_arraysize_and_marked_api_heavy(self): jobs = expand_jobs(SETTINGS, ["single"], ["small"], ["flat"], Filters(family=["cursor"])) @@ -264,7 +275,7 @@ def runner(item, workdir, out, resume, stop): assert sorted(runs) == ["heavy", "light"] assert [path.parent.name for path in settled] == ["heavy"] assert queue.key("results", "light", "attempt-1", "summary.csv") in s3.objects - assert not queue.names("slots") + assert not live_slots(queue) summary = status(queue) assert summary["done"] == 2 assert list(summary["failed"]) == ["heavy"] @@ -489,7 +500,7 @@ def flaky(**kwargs): s3.put_object = flaky with pytest.raises(ClientError): queue.acquire_slots(3, 2, {"host": "h"}, 600) - assert not queue.names("slots") + assert not live_slots(queue) def test_overlapping_submissions_cannot_share_a_name(self, tmp_path): s3, queue = queue_with([job("a")], tmp_path) @@ -654,6 +665,7 @@ def runner(item, workdir, out, resume, stop): "attempt": 1, "state": "running", "expired": False, + "stale": False, } def test_an_expired_slot_returns_to_the_fleet(self, tmp_path): @@ -669,7 +681,7 @@ def test_an_expired_slot_returns_to_the_fleet(self, tmp_path): settle=no_queries, ) assert ran == ["heavy"] - assert not queue.names("slots") + assert not live_slots(queue) def test_watch_renews_leases_and_uploads_progress(self, tmp_path): _, queue = queue_with([], tmp_path) @@ -733,7 +745,7 @@ def interrupted_runner(item, workdir, out, resume, stop): assert raised.value.ran == [] claim = queue.read_lease(queue.key("claims", "heavy")) assert claim.body["state"] == "interrupted" - assert not queue.names("slots") + assert not live_slots(queue) assert queue.key("results", "heavy", "attempt-1", "trials.jsonl") in s3.objects assert not queue.names("done") @@ -858,3 +870,113 @@ def popen(command, stdout, stderr): assert process.signals == [signal.SIGINT] assert process.command[-4:] == ["--suite", "single", "--resume", str(earlier)] assert (tmp_path / "logs" / "a-attempt-2.log").exists() + + +class TestLeaseOwnership: + def test_releasing_an_expired_slot_keeps_another_workers_takeover(self, tmp_path): + _, queue = queue_with([], tmp_path) + key = queue.key("slots", "0") + mine = queue.acquire(key, {"host": "slow"}, -1) + theirs = queue.acquire(key, {"host": "other"}, 600) + assert not queue.release_lease(mine) + assert queue.read_lease(key).body == theirs.body + assert live_slots(queue) == ["0"] + + def test_interruption_is_marked_after_a_concurrent_renewal(self, tmp_path): + _, queue = queue_with([], tmp_path) + key = queue.key("claims", "a") + lease = queue.acquire(key, {"host": "h", "job": "a"}, 600) + stale = Lease(key, dict(lease.body), lease.etag) + assert queue.renew(lease, 600) + assert queue.mark_interrupted(stale) + assert queue.read_lease(key).body["state"] == "interrupted" + + def test_a_lost_claim_stops_the_attempt_without_recording_the_job(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + stopped = [] + + def runner(item, workdir, out, resume, stop): + # Another worker took over the claim, for example after this host stalled. + s3.put_object( + Key=queue.key("claims", "a"), + Body=json.dumps( + {"host": "other", "attempt": 2, "state": "running", "renewed": time.time(), + "expires": time.time() + 600} + ).encode(), + ) # fmt: skip + stopped.append(stop.wait(5)) + return 130 + + ran = work( + queue, + tmp_path / "w", + api_slots=1, + poll_seconds=0, + runner=runner, + settle=no_queries, + renew_seconds=0.01, + ) + assert stopped == [True] + assert ran == [] + assert not queue.names("done") + assert queue.read_lease(queue.key("claims", "a")).body["host"] == "other" + + def test_a_stale_claim_keeps_the_worker_until_it_can_take_it_over(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + at = time.time() + # The worker renewed its 200-second lease long ago and has probably died. + s3.put_object( + Key=queue.key("claims", "a"), + Body=json.dumps( + {"host": "dead", "attempt": 1, "state": "running", "renewed": at - 199.5, + "expires": at + 0.5} + ).encode(), + ) # fmt: skip + assert status(queue)["running"]["a"]["stale"] + runs = [] + ran = work( + queue, + tmp_path / "w", + api_slots=1, + poll_seconds=0.2, + runner=lambda item, workdir, out, *a: runs.append(out.name) or 0, + settle=no_queries, + ) + assert ran == ["a"] + assert runs == ["attempt-2"] + + +def test_worker_waits_for_stack_outputs(monkeypatch): + calls = [] + + def resources(session_, stack): + calls.append(stack) + if len(calls) < 3: + raise IncompleteStackError("Stack outputs are incomplete") + return {"Results": "s3://bucket/"} + + monkeypatch.setattr("pyathena_bench.__main__.stack_resources", resources) + monkeypatch.setattr("pyathena_bench.__main__.time.sleep", lambda seconds: None) + assert wait_for_stack(None, "stack", wait=True) == {"Results": "s3://bucket/"} + assert len(calls) == 3 + calls.clear() + with pytest.raises(IncompleteStackError): + wait_for_stack(None, "stack", wait=False) + + +def test_a_resumed_run_without_remaining_trials_succeeds(tmp_path, monkeypatch): + config = tmp_path / "config.toml" + config.write_text("[scales]\nsmall = 10000\n") + manifest = tmp_path / "manifest.json" + manifest.write_text("{}") + monkeypatch.setattr("pyathena_bench.__main__.run", lambda *args: True) + + def report(directory): + raise ValueError("No trials to report") + + monkeypatch.setattr("pyathena_bench.__main__.report", report) + argv = ["--config", str(config), "run", "--suite", "single", "--scale", "small"] + paths = ["--manifest", str(manifest), "--out", str(tmp_path / "attempt-2")] + assert main([*argv, *paths, "--resume", str(tmp_path / "attempt-1")]) == 0 + with pytest.raises(ValueError, match="No trials"): + main([*argv, "--manifest", str(manifest), "--out", str(tmp_path / "attempt-3")]) From e033f4547ea0795975da1cc79ec5c3489efad69e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 15:54:17 +0900 Subject: [PATCH 05/16] Bound the watch shutdown after a notice or a lost lease Co-Authored-By: Claude Opus 5.5 --- benchmarks/pyathena_bench/fleet.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index c73f4042b..9e7cb1590 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -803,7 +803,9 @@ def __exit__( exc: BaseException | None, traceback: TracebackType | None, ) -> None: - self.stop() + # After a notice or a lost lease, a renewal stuck in S3 must not delay the + # release of the slots; a late renewal fails against the changed ETags. + self.stop(10 if self.stop_requested.is_set() else None) def stop(self, timeout: float | None = None) -> None: """Stop the threads and wait for them. From 6531dbf05f81eaf969ad0146bf25ce76b5586093 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 16:28:21 +0900 Subject: [PATCH 06/16] Confirm live claims before exiting and keep the first done marker Independent follow-up findings: - The first done marker of a job is created with a conditional write and never replaced; the watch keeps renewing through publication, and only the worker that recorded the marker uploads the job summary. - A worker exits only after every other unfinished claim has been renewed since it found nothing to claim; a claim whose worker died right after a renewal is taken over when it expires. - The notice thread also stops the job when the claim's local expiry passes, so stalled renewals cannot hide a lost claim. - Losing an API slot no longer abandons a job that still holds its claim; the done marker lists lost slots. - Releasing a lease reconciles with this worker's own concurrent renewal, like marking a claim interrupted. - On a notice, the attempt is uploaded before the claim is released, with bounded waits for the upload lock, and the report counts a repetition recorded by two attempts once. - --lease-seconds must be at least 300. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 11 +- benchmarks/pyathena_bench/__main__.py | 4 +- benchmarks/pyathena_bench/fleet.py | 154 +++++++++++++++----------- benchmarks/pyathena_bench/report.py | 22 ++-- benchmarks/tests/test_fleet.py | 98 +++++++++++++++- benchmarks/tests/test_report.py | 23 ++++ 6 files changed, 231 insertions(+), 81 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index cd06c7055..5cb2361ef 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -398,10 +398,11 @@ uv run --env-file ../.env --locked aws ssm send-command --targets "Key=tag:aws:a ``` A worker claims a job by creating `claims/` with an S3 conditional write, so exactly one host runs each job. -A claim, like each API slot under `slots/`, is a lease that its worker renews every tenth of `--lease-seconds` (600 by default) while the job runs. +A claim, like each API slot under `slots/`, is a lease that its worker renews every tenth of `--lease-seconds` (600 by default, at least 300) while the job runs. When a worker dies, its claim and slots expire, and another worker takes them over with a conditional write, so the job and the slots return to the fleet without manual steps. A worker releases a slot by overwriting its lease as released, never by deleting it, so a slot that another worker took over stays held. -A worker whose lease was taken over, or whose renewals failed until the lease expired, stops the attempt and records nothing for the job; the new owner continues it. +A worker whose claim was taken over, or whose renewals failed until the claim expired, stops the attempt and records nothing for the job; the new owner continues it, and the first `done` marker of a job is never replaced. +A lost API slot does not stop the job; the `done` marker lists it under `lost_slots`. Each attempt of a job writes to `results//attempt-/`, and the worker uploads changed trial records every five seconds, so a lost host loses its running trial and at most the records of the last seconds. A trial whose start was not uploaded before its host died is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. After a job, the worker uploads the attempt's output directory and its log to `logs/`, writes `results//summary.csv` and `summary.md` from every attempt, and writes `done/` with the exit code, attempt, host, and times. @@ -411,14 +412,14 @@ It waits until they stop before claiming another job; this scans the workgroup's If it cannot confirm that, it records the errors in the `done` marker and exits. Each worker checks instance metadata for a Spot interruption notice (`spot/instance-action`), which EC2 publishes about two minutes before it reclaims an instance. -On a notice, the worker interrupts the running trial with SIGINT, which cancels the trial's reported queries, releases its claim as `interrupted`, uploads the attempt, releases its slots, and exits without terminating the host. +On a notice, the worker interrupts the running trial with SIGINT, which cancels the trial's reported queries, uploads the attempt, releases its claim as `interrupted` and its slots, and exits without terminating the host. A separate thread watches for the notice, so slow S3 requests do not delay it. The group replaces the reclaimed instance, and the new host's worker takes over the interrupted claim as the next attempt: it downloads the earlier attempts, quiesces their trials as after a failed job, and runs only the repetitions that they did not record as successful, through `run --resume`. The job's summary merges the attempts and leaves out a failed or incomplete trial whose repetition succeeded in a later attempt, so it reports the trials of an uninterrupted run. Rebalance recommendations are not acted on; an interruption does not always follow them, and the notice leaves enough time to stop a trial. -Workers exit when no job is left to claim: every job is done or claimed by a live lease. -A claim that has missed renewals belongs to a worker that has probably died, so workers keep waiting until it expires and one of them takes it over. +Workers exit when no job is left to claim and every other unfinished job's claim has been renewed since they found nothing to claim, which takes up to a renewal interval. +A claim whose worker died is never renewed, so workers keep waiting until it expires and one of them takes it over; `status` marks a claim that has missed renewals as `stale`. With `--terminate-host`, a worker first suspends the group's `AZRebalance` process, which would otherwise let the group launch a host in one zone and terminate a busy host in another when Spot capacity or retired hosts leave the zones uneven; the process stays suspended until the stack is deleted. A worker with `--terminate-host` that returns normally records its jobs, instance type, Availability Zone, and purchase option under `hosts/.json` and terminates its host through the group with a lower desired capacity, so no replacement starts and finished hosts stop costing money while others still run. This includes a worker that stopped after unconfirmed queries, whose errors are already in the `done` marker. diff --git a/benchmarks/pyathena_bench/__main__.py b/benchmarks/pyathena_bench/__main__.py index 06a2c1089..ba31c8e73 100644 --- a/benchmarks/pyathena_bench/__main__.py +++ b/benchmarks/pyathena_bench/__main__.py @@ -230,8 +230,8 @@ def main(argv: list[str] | None = None) -> int: elif args.command == "worker": if args.api_slots < 1: raise ValueError("--api-slots must be positive") - if args.lease_seconds < 120: - raise ValueError("--lease-seconds must be at least 120") + 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) diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index 9e7cb1590..7ef9976d5 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -431,10 +431,7 @@ def renew(self, lease: Lease, lease_seconds: float, **changes: Any) -> bool: return True def release_lease(self, lease: Lease) -> bool: - """Release a lease that this worker still holds, so another worker can take it. - - The object is overwritten as released rather than deleted, so that a - lease another worker took over in the meantime stays intact. + """Release a lease that this worker still holds; see :meth:`end_lease`. Args: lease: Lease to release. @@ -442,28 +439,31 @@ def release_lease(self, lease: Lease) -> bool: Returns: False if another worker took over the lease. """ - return self.renew(lease, 0, state="released", expires=0) + return self.end_lease(lease, "released") - def mark_interrupted(self, lease: Lease) -> bool: - """Release a claim as interrupted, so another worker continues the job at once. + def end_lease(self, lease: Lease, state: str) -> bool: + """End a lease that this worker still holds, so another worker can take it. - If the lease was renewed concurrently with a newer ETag, the current - object is still released when it belongs to the same owner and attempt. + The object is overwritten with ``state`` rather than deleted, so that a + lease another worker took over in the meantime stays intact. If a + concurrent renewal by this worker gave the object a newer ETag, the + current object is ended when it has the same owner and attempt. Args: - lease: Claim lease of the interrupted attempt. + lease: Lease to end. + state: ``released`` or ``interrupted``. Returns: - False if another worker took over the claim. + False if another worker took over the lease. """ - if self.renew(lease, 0, state="interrupted", expires=0): + if self.renew(lease, 0, state=state, expires=0): return True current = self.read_lease(lease.key) if current is None or any( current.body.get(k) != lease.body.get(k) for k in ("host", "job", "attempt") ): return False - return self.renew(current, 0, state="interrupted", expires=0) + return self.renew(current, 0, state=state, expires=0) def names(self, folder: str) -> set[str]: """List the object names directly below a queue folder. @@ -756,12 +756,13 @@ def __init__(self, ran: list[str]) -> None: class JobWatch: """Renew a running job's leases, upload its progress, and watch for interruption. - One thread checks for an interruption notice every ``poll_seconds``. - Another calls ``progress`` every ``poll_seconds`` and renews the leases - every ``renew_seconds``, so a slow S3 request does not delay the notice. - ``stop_requested`` is set on a notice (with ``interrupted``) and when a - lease is lost: another worker took it over, or renewals failed until it - expired. A lost lease is listed in ``lost`` and no longer renewed. + One thread checks every ``poll_seconds`` for an interruption notice and + for the claim's expiry, without S3 requests. Another calls ``progress`` + every ``poll_seconds`` and renews the leases every ``renew_seconds``. + ``stop_requested`` is set on a notice (with ``interrupted``) and when the + claim, the first lease, is lost: another worker took it over, or + renewals failed or stalled until it expired. A lost lease is listed in + ``lost`` and no longer renewed; losing an API slot does not stop the job. ``upload_lock`` serializes uploads of the attempt's files. """ @@ -776,6 +777,7 @@ def __init__( poll_seconds: float = 5.0, ) -> None: self.queue = queue + self.claim = leases[0] self.leases = list(leases) self.lease_seconds = lease_seconds self.renew_seconds = renew_seconds @@ -825,6 +827,10 @@ def _watch_notice(self) -> None: self.interrupted.set() self.stop_requested.set() return + if time.time() >= self.claim.body["expires"]: + # Renewals failed or stalled; another worker may own the job now. + self.stop_requested.set() + return def _renew(self) -> None: renew_at = time.monotonic() + self.renew_seconds @@ -848,10 +854,13 @@ def _renew(self) -> None: continue self.leases.remove(lease) self.lost.append(lease.key) - self.stop_requested.set() + if lease is self.claim: + self.stop_requested.set() -def pending_jobs(queue: Queue, jobs: list[dict[str, Any]]) -> tuple[list[dict[str, Any]], bool]: +def pending_jobs( + queue: Queue, jobs: list[dict[str, Any]] +) -> tuple[list[dict[str, Any]], dict[str, float | None]]: """Return the jobs that are neither done nor held by a live claim. Args: @@ -860,25 +869,24 @@ def pending_jobs(queue: Queue, jobs: list[dict[str, Any]]) -> tuple[list[dict[st Returns: Jobs without a ``done`` marker whose claim is absent, expired, or - released, in queue order, and whether another unfinished job's claim - is stale: its worker has probably died, and the job returns to the - fleet when the claim expires. + released, in queue order, and the last renewal time of every other + unfinished job's claim. """ done = queue.names("done") claimed = queue.names("claims") at = time.time() pending = [] - stale = False + running: dict[str, float | None] = {} for job in jobs: if job["id"] in done: continue if job["id"] in claimed: lease = queue.read_lease(queue.key("claims", job["id"])) if lease is not None and not lease.takeable(at): - stale = stale or lease.stale(at) + running[job["id"]] = lease.body.get("renewed") continue pending.append(job) - return pending, stale + return pending, running @dataclass(frozen=True) @@ -949,11 +957,16 @@ def progress() -> None: with JobWatch(queue, [claim, *slots], lease_seconds, renew_seconds, progress, notice) as watch: - def upload_attempt() -> None: - with watch.upload_lock: + def upload_attempt(timeout: float = -1) -> None: + # A checkpoint stuck in S3 holds the lock; it uploads the same files. + if not watch.upload_lock.acquire(timeout=timeout): + return + try: queue.upload_tree(output, *remote) if log.exists(): queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) + finally: + watch.upload_lock.release() start = now() resume = [ @@ -975,36 +988,33 @@ def upload_attempt() -> None: if not errors and not watch.stop_requested.is_set(): code = runner(job, workdir, output, resume, watch.stop_requested) if watch.stop_requested.is_set(): - # Within the two minutes of a notice, release the claim before uploading. watch.stop(timeout=10) - interrupted = watch.interrupted.is_set() - if interrupted: - queue.mark_interrupted(claim) - upload_attempt() - return Outcome("interrupted" if interrupted else "lost") + # Upload before releasing the claim, so that the next attempt resumes + # from every recorded trial; both are small within a notice's two minutes. + upload_attempt(timeout=10) + if not watch.interrupted.is_set(): + # The claim was lost; another worker continues the job. + return Outcome("lost") + queue.end_lease(claim, "interrupted") + return Outcome("interrupted") if not errors: try: errors = settle(output) if code else [] except Exception as exc: errors = [f"Could not quiesce the job: {type(exc).__name__}: {exc}"] upload_attempt() - if watch.lost or not queue.renew(claim, lease_seconds): - # Another worker continues the job and records its outcome. - return Outcome("lost") - report_error = None - try: - report(job_dir) - for name in ("summary.csv", "summary.md"): - queue.s3.upload_file( - str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) - ) - except Exception as exc: - # A job without trials, or with unreadable records, still gets its done marker. - report_error = f"{type(exc).__name__}: {exc}" - queue.s3.put_object( - Bucket=queue.bucket, - Key=queue.key("done", job["id"]), - Body=json.dumps( + report_error = None + try: + report(job_dir) + except Exception as exc: + # A job without trials, or with unreadable records, still gets its done marker. + report_error = f"{type(exc).__name__}: {exc}" + if claim.key in watch.lost: + return Outcome("lost") + # The first done marker wins, so a worker that lost the claim unnoticed + # cannot replace the outcome that the claim's new owner recorded. + recorded = queue.put_lease( + queue.key("done", job["id"]), { "exit_code": code, "host": socket.gethostname(), @@ -1013,9 +1023,17 @@ def upload_attempt() -> None: "end": now(), "quiesce_errors": errors, "report_error": report_error, - } - ).encode(), - ) + "lost_slots": [key for key in watch.lost if key != claim.key], + }, + IfNoneMatch="*", + ) + if recorded is None: + return Outcome("lost") + if report_error is None: + for name in ("summary.csv", "summary.md"): + queue.s3.upload_file( + str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) + ) return Outcome("done", tuple(errors)) @@ -1036,10 +1054,11 @@ def work( Each host runs one worker, so jobs on a host never overlap. A claim and each API slot are leases that the worker renews while the job runs; a job or slot whose lease expired, because its worker died, returns to the - fleet. While another unfinished job's claim is stale, the worker waits to - take it over instead of exiting. An API-heavy job starts only after the - worker takes one shared slot per simultaneous paging query; while it - waits, no later API-heavy job takes slots. After a failed job, the worker + fleet. Before exiting, the worker waits until every other unfinished + job's claim has been renewed, so a job whose worker died is taken over. + An API-heavy job starts only after the worker takes one shared slot per + simultaneous paging query; while it waits, no later API-heavy job takes + slots. After a failed job, the worker cancels the job's observed queries and waits until they stop before claiming another job. @@ -1086,12 +1105,23 @@ def settle(output: Path) -> list[str]: host = socket.gethostname() completed: list[str] = [] + # Renewal times of other workers' claims when this worker found nothing to claim. + waiting: dict[str, float | None] = {} while True: if notice(): raise SpotInterruptionError(completed) - pending, stale = pending_jobs(queue, jobs) - if not pending and not stale: - return completed + pending, running = pending_jobs(queue, jobs) + if pending: + waiting = {} + else: + for job_id, renewed in running.items(): + waiting.setdefault(job_id, renewed) + # Exit once every other claim has been renewed since; a claim whose + # worker died is never renewed, and is taken over when it expires. + if all( + renewed is None or renewed != waiting[job_id] for job_id, renewed in running.items() + ): + return completed started = False reserved = False for job in pending: diff --git a/benchmarks/pyathena_bench/report.py b/benchmarks/pyathena_bench/report.py index 433a0165c..72790112f 100644 --- a/benchmarks/pyathena_bench/report.py +++ b/benchmarks/pyathena_bench/report.py @@ -152,7 +152,7 @@ def collect_trials(directory: Path) -> list[dict[str, Any]]: A started trial without a final record is reported as incomplete. A failed or incomplete trial is dropped when another attempt recorded the same repetition as successful, so a resumed job reports the trials of an - uninterrupted run. + uninterrupted run; so is a second successful record of a repetition. Args: directory: Output directory of a run, or a job directory with attempts. @@ -172,12 +172,20 @@ def collect_trials(directory: Path) -> list[dict[str, Any]]: succeeded = { (t["scale"], t["case_id"], t.get("repetition")) for t in trials if t["status"] == "ok" } - return [ - t - for t in trials - if t["status"] in {"ok", "unsupported"} - or (t["scale"], t["case_id"], t.get("repetition")) not in succeeded - ] + merged = [] + kept: set[tuple[Any, ...]] = set() + for trial in trials: + key = (trial["scale"], trial["case_id"], trial.get("repetition")) + if trial["status"] == "ok": + # Two attempts can both record a repetition when a takeover resumed + # from a checkpoint older than the interrupted attempt's last upload. + if len(sources) > 1 and key in kept: + continue + kept.add(key) + elif trial["status"] != "unsupported" and key in succeeded: + continue + merged.append(trial) + return merged def report(directory: Path) -> None: diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index d931a9d78..3fd1a0da3 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -643,19 +643,27 @@ def lost_response(**kwargs): assert taken.body["attempt"] == (1 if operation == "create" else 2) assert queue.read_lease(key).etag == taken.etag - def test_live_claims_are_left_and_expired_ones_return_to_the_fleet(self, tmp_path): + def test_live_claims_are_left_and_expired_ones_return_to_the_fleet(self, tmp_path, monkeypatch): _, queue = queue_with([job("live"), job("expired")], tmp_path) - queue.acquire(queue.key("claims", "live"), {"host": "busy"}, 600) + busy = queue.acquire(queue.key("claims", "live"), {"host": "busy"}, 600) queue.acquire(queue.key("claims", "expired"), {"host": "dead"}, -1) runs = [] + waits = [] def runner(item, workdir, out, resume, stop): runs.append((item["id"], out.name, resume)) return 0 + def sleep(seconds): + # The busy worker renews its claim while this worker waits to exit. + waits.append(seconds) + assert queue.renew(busy, 600) + + monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) ran = work( queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner, settle=no_queries ) + assert len(waits) == 1 assert ran == ["expired"] assert runs == [("expired", "attempt-2", [])] running = status(queue)["running"] @@ -888,13 +896,26 @@ def test_interruption_is_marked_after_a_concurrent_renewal(self, tmp_path): lease = queue.acquire(key, {"host": "h", "job": "a"}, 600) stale = Lease(key, dict(lease.body), lease.etag) assert queue.renew(lease, 600) - assert queue.mark_interrupted(stale) + assert queue.end_lease(stale, "interrupted") assert queue.read_lease(key).body["state"] == "interrupted" - - def test_a_lost_claim_stops_the_attempt_without_recording_the_job(self, tmp_path): + # A stale copy of a slot is released the same way. + slot = queue.acquire(queue.key("slots", "0"), {"host": "h", "job": "a"}, 600) + copy = Lease(slot.key, dict(slot.body), slot.etag) + assert queue.renew(slot, 600) + assert queue.release_lease(copy) + assert live_slots(queue) == [] + + def test_a_lost_claim_stops_the_attempt_without_recording_the_job(self, tmp_path, monkeypatch): s3, queue = queue_with([job("a")], tmp_path) stopped = [] + def sleep(seconds): + # The new owner renews the claim while this worker waits to exit. + other = queue.read_lease(queue.key("claims", "a")) + assert queue.renew(other, 600) + + monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) + def runner(item, workdir, out, resume, stop): # Another worker took over the claim, for example after this host stalled. s3.put_object( @@ -980,3 +1001,70 @@ def report(directory): assert main([*argv, *paths, "--resume", str(tmp_path / "attempt-1")]) == 0 with pytest.raises(ValueError, match="No trials"): main([*argv, "--manifest", str(manifest), "--out", str(tmp_path / "attempt-3")]) + + +class TestLeaseExpiry: + def test_watch_stops_the_job_when_renewals_fail_until_the_claim_expires(self, tmp_path): + _, queue = queue_with([], tmp_path) + claim = queue.acquire(queue.key("claims", "a"), {"host": "h"}, 0.2) + + def unavailable(lease, seconds, **changes): + raise ClientError({"Error": {"Code": "SlowDown"}}, "PutObject") + + queue.renew = unavailable + with JobWatch(queue, [claim], 600, 0.01, lambda: None, lambda: False) as watch: + assert watch.stop_requested.wait(5) + assert not watch.interrupted.is_set() + + def test_losing_a_slot_does_not_stop_the_job(self, tmp_path): + s3, queue = queue_with([job("heavy", heavy=True)], tmp_path) + + def runner(item, workdir, out, resume, stop): + # Another worker took over the slot, for example after this host stalled. + s3.put_object( + Key=queue.key("slots", "0"), + Body=json.dumps( + {"host": "other", "state": "running", "expires": time.time() + 600} + ).encode(), + ) + time.sleep(0.2) + assert not stop.is_set() + return 0 + + assert work( + queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner, settle=no_queries, + renew_seconds=0.01, + ) == ["heavy"] # fmt: skip + assert json.loads(queue.read("done", "heavy"))["lost_slots"] == [queue.key("slots", "0")] + # The other worker's slot is kept. + assert json.loads(s3.objects[queue.key("slots", "0")])["host"] == "other" + + def test_a_fresh_claim_that_is_never_renewed_is_taken_over_after_it_expires(self, tmp_path): + _, queue = queue_with([job("a")], tmp_path) + # Its worker died right after renewing it. + queue.acquire(queue.key("claims", "a"), {"host": "dead"}, 0.5) + runs = [] + ran = work( + queue, + tmp_path / "w", + api_slots=1, + poll_seconds=0.2, + runner=lambda item, workdir, out, *a: runs.append(out.name) or 0, + settle=no_queries, + ) + assert ran == ["a"] + assert runs == ["attempt-2"] + + def test_an_existing_done_marker_is_not_replaced(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + + def runner(item, workdir, out, resume, stop): + # The claim's new owner finished the job while this worker stalled. + queue.put_if_absent(queue.key("done", "a"), {"exit_code": 0, "host": "other"}) + out.mkdir(parents=True) + (out / "trials.jsonl").write_text(json.dumps(ok_trial(0)) + "\n") + return 0 + + assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries) == [] + assert json.loads(queue.read("done", "a"))["host"] == "other" + assert queue.key("results", "a", "summary.csv") not in s3.objects diff --git a/benchmarks/tests/test_report.py b/benchmarks/tests/test_report.py index 16cdea7fc..553782182 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -141,3 +141,26 @@ def trial(repetition, status="ok"): with (tmp_path / "summary.csv").open() as stream: row = next(csv.DictReader(stream)) assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") + + +def test_report_counts_a_repetition_recorded_by_two_attempts_once(tmp_path): + record = { + "case_id": "case", + "scale": "small", + "repetition": 0, + "warmup": False, + "status": "ok", + "queries": [{"total_seconds": 1}], + "successful_queries_per_second": 1, + "rss_peak_bytes": 1, + "max_threads": 1, + } + for attempt in (1, 2): + directory = tmp_path / f"attempt-{attempt}" + directory.mkdir() + (directory / "trials.jsonl").write_text( + json.dumps({**record, "trial": f"t{attempt}"}) + "\n" + ) + report(tmp_path) + with (tmp_path / "summary.csv").open() as stream: + assert next(csv.DictReader(stream))["successful_trials"] == "1" From 9e349def875e880902a19826fe9162c2a7846497 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 16:29:00 +0900 Subject: [PATCH 07/16] Treat a claim that expired after the run as lost before recording the job Co-Authored-By: Claude Opus 5.5 --- benchmarks/pyathena_bench/fleet.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index 7ef9976d5..d6c471958 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -1009,7 +1009,10 @@ def upload_attempt(timeout: float = -1) -> None: except Exception as exc: # A job without trials, or with unreadable records, still gets its done marker. report_error = f"{type(exc).__name__}: {exc}" - if claim.key in watch.lost: + if claim.key in watch.lost or ( + watch.stop_requested.is_set() and not watch.interrupted.is_set() + ): + # The claim was taken over or expired while the job was settled or reported. return Outcome("lost") # The first done marker wins, so a worker that lost the claim unnoticed # cannot replace the outcome that the claim's new owner recorded. From 0f85514d371120f296bde8f3d14be1567837a7ac Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 16:54:14 +0900 Subject: [PATCH 08/16] Confirm claim ownership synchronously before recording a job Independent follow-up findings: - After the watch stops, a conditional renewal of the claim must succeed before the job is recorded, so a former owner that stalled through a takeover cannot record the job first; the renewal leaves at least --lease-seconds for publication. - Summaries are uploaded before the done marker is created. - A worker exits only when every other unfinished claim was renewed since it started waiting and is still fresh. - The expiry test now blocks a renewal and requires the stop before unblocking it. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 4 +- benchmarks/pyathena_bench/fleet.py | 48 +++++++++++--------- benchmarks/tests/test_fleet.py | 71 ++++++++++++++++++++++++++---- 3 files changed, 92 insertions(+), 31 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 5cb2361ef..b542e44a6 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -405,7 +405,7 @@ A worker whose claim was taken over, or whose renewals failed until the claim ex A lost API slot does not stop the job; the `done` marker lists it under `lost_slots`. Each attempt of a job writes to `results//attempt-/`, and the worker uploads changed trial records every five seconds, so a lost host loses its running trial and at most the records of the last seconds. A trial whose start was not uploaded before its host died is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. -After a job, the worker uploads the attempt's output directory and its log to `logs/`, writes `results//summary.csv` and `summary.md` from every attempt, and writes `done/` with the exit code, attempt, host, and times. +After a job, the worker uploads the attempt's output directory and its log to `logs/`, renews its claim conditionally to confirm that it still owns the job, writes `results//summary.csv` and `summary.md` from every attempt, and then creates `done/` with the exit code, attempt, host, and times. A failed trial stops only its own job; `status` lists jobs with a non-zero exit code, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. After a failed job, the worker stops the job's reported queries and every queued or running query in the workgroup that writes under the job's trial output prefixes, which also covers queries whose IDs were never reported before a trial process died. It waits until they stop before claiming another job; this scans the workgroup's query history. @@ -418,7 +418,7 @@ The group replaces the reclaimed instance, and the new host's worker takes over The job's summary merges the attempts and leaves out a failed or incomplete trial whose repetition succeeded in a later attempt, so it reports the trials of an uninterrupted run. Rebalance recommendations are not acted on; an interruption does not always follow them, and the notice leaves enough time to stop a trial. -Workers exit when no job is left to claim and every other unfinished job's claim has been renewed since they found nothing to claim, which takes up to a renewal interval. +Workers exit when no job is left to claim and every other unfinished job's claim has been renewed since they found nothing to claim and is still fresh, which takes up to a renewal interval. A claim whose worker died is never renewed, so workers keep waiting until it expires and one of them takes it over; `status` marks a claim that has missed renewals as `stale`. With `--terminate-host`, a worker first suspends the group's `AZRebalance` process, which would otherwise let the group launch a host in one zone and terminate a busy host in another when Spot capacity or retired hosts leave the zones uneven; the process stays suspended until the stack is deleted. A worker with `--terminate-host` that returns normally records its jobs, instance type, Availability Zone, and purchase option under `hosts/.json` and terminates its host through the group with a lower desired capacity, so no replacement starts and finished hosts stop costing money while others still run. diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index d6c471958..a265466cd 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -860,7 +860,7 @@ def _renew(self) -> None: def pending_jobs( queue: Queue, jobs: list[dict[str, Any]] -) -> tuple[list[dict[str, Any]], dict[str, float | None]]: +) -> tuple[list[dict[str, Any]], dict[str, Lease]]: """Return the jobs that are neither done nor held by a live claim. Args: @@ -869,21 +869,20 @@ def pending_jobs( Returns: Jobs without a ``done`` marker whose claim is absent, expired, or - released, in queue order, and the last renewal time of every other - unfinished job's claim. + released, in queue order, and the claims of the other unfinished jobs. """ done = queue.names("done") claimed = queue.names("claims") at = time.time() pending = [] - running: dict[str, float | None] = {} + running: dict[str, Lease] = {} for job in jobs: if job["id"] in done: continue if job["id"] in claimed: lease = queue.read_lease(queue.key("claims", job["id"])) if lease is not None and not lease.takeable(at): - running[job["id"]] = lease.body.get("renewed") + running[job["id"]] = lease continue pending.append(job) return pending, running @@ -1009,13 +1008,23 @@ def upload_attempt(timeout: float = -1) -> None: except Exception as exc: # A job without trials, or with unreadable records, still gets its done marker. report_error = f"{type(exc).__name__}: {exc}" - if claim.key in watch.lost or ( - watch.stop_requested.is_set() and not watch.interrupted.is_set() + watch.stop() + # Renewing the claim conditionally proves that this worker still owns the + # job and leaves at least --lease-seconds for the few seconds of publication. + if ( + claim.key in watch.lost + or (watch.stop_requested.is_set() and not watch.interrupted.is_set()) + or not queue.renew(claim, lease_seconds) ): - # The claim was taken over or expired while the job was settled or reported. + # The claim was taken over or expired; its new owner records the job. return Outcome("lost") - # The first done marker wins, so a worker that lost the claim unnoticed - # cannot replace the outcome that the claim's new owner recorded. + if report_error is None: + # Summaries go first, so a recorded job always has them. + for name in ("summary.csv", "summary.md"): + queue.s3.upload_file( + str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) + ) + # The first done marker of a job is never replaced. recorded = queue.put_lease( queue.key("done", job["id"]), { @@ -1032,11 +1041,6 @@ def upload_attempt(timeout: float = -1) -> None: ) if recorded is None: return Outcome("lost") - if report_error is None: - for name in ("summary.csv", "summary.md"): - queue.s3.upload_file( - str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) - ) return Outcome("done", tuple(errors)) @@ -1117,12 +1121,16 @@ def settle(output: Path) -> list[str]: if pending: waiting = {} else: - for job_id, renewed in running.items(): - waiting.setdefault(job_id, renewed) - # Exit once every other claim has been renewed since; a claim whose - # worker died is never renewed, and is taken over when it expires. + at = time.time() + for job_id, lease in running.items(): + waiting.setdefault(job_id, lease.body.get("renewed")) + # Exit once every other claim has been renewed since and is still + # fresh; a claim whose worker died stops being renewed and is taken + # over when it expires. if all( - renewed is None or renewed != waiting[job_id] for job_id, renewed in running.items() + lease.body.get("renewed") is None + or (lease.body["renewed"] != waiting[job_id] and not lease.stale(at)) + for job_id, lease in running.items() ): return completed started = False diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 3fd1a0da3..55bf90bf2 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -1004,17 +1004,24 @@ def report(directory): class TestLeaseExpiry: - def test_watch_stops_the_job_when_renewals_fail_until_the_claim_expires(self, tmp_path): + def test_watch_stops_the_job_when_a_stalled_renewal_outlasts_the_claim(self, tmp_path): _, queue = queue_with([], tmp_path) claim = queue.acquire(queue.key("claims", "a"), {"host": "h"}, 0.2) + stalled = threading.Event() + released = threading.Event() - def unavailable(lease, seconds, **changes): - raise ClientError({"Error": {"Code": "SlowDown"}}, "PutObject") + def stall(lease, seconds, **changes): + stalled.set() + assert released.wait(5) + return False - queue.renew = unavailable + queue.renew = stall with JobWatch(queue, [claim], 600, 0.01, lambda: None, lambda: False) as watch: + assert stalled.wait(5) + # The claim expires while its renewal is still blocked in S3. assert watch.stop_requested.wait(5) - assert not watch.interrupted.is_set() + assert not watch.interrupted.is_set() + released.set() def test_losing_a_slot_does_not_stop_the_job(self, tmp_path): s3, queue = queue_with([job("heavy", heavy=True)], tmp_path) @@ -1055,16 +1062,62 @@ def test_a_fresh_claim_that_is_never_renewed_is_taken_over_after_it_expires(self assert ran == ["a"] assert runs == ["attempt-2"] - def test_an_existing_done_marker_is_not_replaced(self, tmp_path): + def test_a_former_owner_records_nothing_after_a_takeover(self, tmp_path, monkeypatch): s3, queue = queue_with([job("a")], tmp_path) def runner(item, workdir, out, resume, stop): - # The claim's new owner finished the job while this worker stalled. - queue.put_if_absent(queue.key("done", "a"), {"exit_code": 0, "host": "other"}) + # This host stalled; the claim expired and another worker took it over, + # before this worker's watch noticed. + s3.put_object( + Key=queue.key("claims", "a"), + Body=json.dumps( + {"host": "other", "job": "a", "attempt": 2, "state": "running", + "renewed": time.time(), "expires": time.time() + 600} + ).encode(), + ) # fmt: skip out.mkdir(parents=True) (out / "trials.jsonl").write_text(json.dumps(ok_trial(0)) + "\n") return 0 + def sleep(seconds): + other = queue.read_lease(queue.key("claims", "a")) + assert queue.renew(other, 600) + + monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries) == [] - assert json.loads(queue.read("done", "a"))["host"] == "other" + assert not queue.names("done") assert queue.key("results", "a", "summary.csv") not in s3.objects + + +def test_a_claim_renewed_once_then_stale_does_not_let_the_worker_exit(tmp_path, monkeypatch): + s3, queue = queue_with([job("a")], tmp_path) + queue.acquire(queue.key("claims", "a"), {"host": "x", "job": "a"}, 600) + sleeps = [] + real_sleep = time.sleep + + def sleep(seconds): + sleeps.append(seconds) + if len(sleeps) == 1: + # Its worker renewed the claim once and then died long ago. + at = time.time() + s3.put_object( + Key=queue.key("claims", "a"), + Body=json.dumps( + {"host": "x", "job": "a", "attempt": 1, "state": "running", + "renewed": at - 199.7, "expires": at + 0.3} + ).encode(), + ) # fmt: skip + real_sleep(0.2) + + monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) + runs = [] + ran = work( + queue, + tmp_path / "w", + api_slots=1, + poll_seconds=0, + runner=lambda item, workdir, out, *a: runs.append(out.name) or 0, + settle=no_queries, + ) + assert ran == ["a"] + assert runs == ["attempt-2"] From f0dc49f53085e8afbe39e2ad70c483a477ba0d8e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 17:05:43 +0900 Subject: [PATCH 09/16] Recheck the claim before recording a job and before exiting Independent follow-up findings: - The done marker is not created when publication outlasted the renewed claim. - A worker does not exit while another unfinished claim has expired since the scan. - The stalled-renewal test keeps the renewal blocked until cleanup, and the resume test checks that the summary precedes the done marker. Co-Authored-By: Claude Opus 5.5 --- benchmarks/pyathena_bench/fleet.py | 9 ++++++++- benchmarks/tests/test_fleet.py | 29 +++++++++++++++++++++++------ 2 files changed, 31 insertions(+), 7 deletions(-) diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index a265466cd..e02378da4 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -1024,6 +1024,9 @@ def upload_attempt(timeout: float = -1) -> None: queue.s3.upload_file( str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) ) + if time.time() >= claim.body["expires"]: + # Publication outlasted the renewed claim, which another worker may own now. + return Outcome("lost") # The first done marker of a job is never replaced. recorded = queue.put_lease( queue.key("done", job["id"]), @@ -1129,7 +1132,11 @@ def settle(output: Path) -> list[str]: # over when it expires. if all( lease.body.get("renewed") is None - or (lease.body["renewed"] != waiting[job_id] and not lease.stale(at)) + or ( + lease.body["renewed"] != waiting[job_id] + and not lease.stale(at) + and not lease.takeable(at) + ) for job_id, lease in running.items() ): return completed diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 55bf90bf2..51d310362 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -759,6 +759,14 @@ def interrupted_runner(item, workdir, out, resume, stop): settled = [] runs = [] + writes = [] + original_put = s3.put_object + + def recording_put(**kwargs): + writes.append(kwargs["Key"]) + return original_put(**kwargs) + + s3.put_object = recording_put def resumed_runner(item, workdir, out, resume, stop): runs.append((item["id"], out.name, [p.name for p in resume])) @@ -781,6 +789,9 @@ def resumed_runner(item, workdir, out, resume, stop): assert runs == [("heavy", "attempt-2", ["attempt-1"]), ("next", "attempt-1", [])] # The interrupted attempt's queries are quiesced before the job continues. assert settled == ["attempt-1"] + order = [key for key in writes if key.endswith(("heavy/summary.csv", "done/heavy"))] + # A recorded job always has its summary. + assert order == [queue.key("results", "heavy", "summary.csv"), queue.key("done", "heavy")] summary = s3.objects[queue.key("results", "heavy", "summary.csv")].decode() row = next(csv.DictReader(io.StringIO(summary))) assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") @@ -1012,15 +1023,21 @@ def test_watch_stops_the_job_when_a_stalled_renewal_outlasts_the_claim(self, tmp def stall(lease, seconds, **changes): stalled.set() - assert released.wait(5) + released.wait() return False queue.renew = stall - with JobWatch(queue, [claim], 600, 0.01, lambda: None, lambda: False) as watch: - assert stalled.wait(5) - # The claim expires while its renewal is still blocked in S3. - assert watch.stop_requested.wait(5) - assert not watch.interrupted.is_set() + watch = JobWatch(queue, [claim], 600, 0.01, lambda: None, lambda: False) + try: + with watch: + assert stalled.wait(5) + # The claim expires while its renewal is still blocked in S3. + assert watch.stop_requested.wait(5) + assert not watch.interrupted.is_set() + assert stalled.is_set() + assert not watch.lost + released.set() + finally: released.set() def test_losing_a_slot_does_not_stop_the_job(self, tmp_path): From 26748f7fae29e72d17ff514216f5d6a6b2166dd7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 17:17:37 +0900 Subject: [PATCH 10/16] Keep each job summary below the attempt that the done marker names A former owner whose upload stalled past its claim could overwrite the summary that the claim's new owner published under the job's shared keys. The merged summary now goes to job-summary.* in the publishing attempt's directory, and the first done marker names that attempt. The stalled-renewal test unblocks the renewal before the watch is joined. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 5 +++-- benchmarks/pyathena_bench/fleet.py | 16 ++++++++++------ benchmarks/tests/test_fleet.py | 21 +++++++++++---------- 3 files changed, 24 insertions(+), 18 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index b542e44a6..1af025d3e 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -405,7 +405,7 @@ A worker whose claim was taken over, or whose renewals failed until the claim ex A lost API slot does not stop the job; the `done` marker lists it under `lost_slots`. Each attempt of a job writes to `results//attempt-/`, and the worker uploads changed trial records every five seconds, so a lost host loses its running trial and at most the records of the last seconds. A trial whose start was not uploaded before its host died is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. -After a job, the worker uploads the attempt's output directory and its log to `logs/`, renews its claim conditionally to confirm that it still owns the job, writes `results//summary.csv` and `summary.md` from every attempt, and then creates `done/` with the exit code, attempt, host, and times. +After a job, the worker uploads the attempt's output directory and its log to `logs/`, renews its claim conditionally to confirm that it still owns the job, writes `job-summary.csv` and `job-summary.md` from every attempt into its attempt directory, and then creates `done/` with the exit code, attempt, host, times, and the summary's attempt directory under `summary`. A failed trial stops only its own job; `status` lists jobs with a non-zero exit code, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. After a failed job, the worker stops the job's reported queries and every queued or running query in the workgroup that writes under the job's trial output prefixes, which also covers queries whose IDs were never reported before a trial process died. It waits until they stop before claiming another job; this scans the workgroup's query history. @@ -416,6 +416,7 @@ On a notice, the worker interrupts the running trial with SIGINT, which cancels A separate thread watches for the notice, so slow S3 requests do not delay it. The group replaces the reclaimed instance, and the new host's worker takes over the interrupted claim as the next attempt: it downloads the earlier attempts, quiesces their trials as after a failed job, and runs only the repetitions that they did not record as successful, through `run --resume`. The job's summary merges the attempts and leaves out a failed or incomplete trial whose repetition succeeded in a later attempt, so it reports the trials of an uninterrupted run. +`report results/recovered/fleet//results/` produces the same merged summary from the recovered attempts. Rebalance recommendations are not acted on; an interruption does not always follow them, and the notice leaves enough time to stop a trial. Workers exit when no job is left to claim and every other unfinished job's claim has been renewed since they found nothing to claim and is still fresh, which takes up to a renewal interval. @@ -468,7 +469,7 @@ uv run --env-file ../.env --locked aws autoscaling terminate-instance-in-auto-sc ``` The group launches a replacement whose worker continues the job as `attempt-2` and then retires the host. -The pilot passes when `done/` records `attempt` 2 and exit code 0, every case in `results//summary.csv` reports six successful trials and no failures, `hosts/` records the replacement with lifecycle `spot`, and the group's desired capacity drops to zero. +The pilot passes when `done/` records `attempt` 2 and exit code 0, every case in `results//attempt-2/job-summary.csv` reports six successful trials and no failures, `hosts/` records the replacement with lifecycle `spot`, and the group's desired capacity drops to zero. AWS Fault Injection Service (`aws:ec2:send-spot-instance-interruptions`) sends a real notice instead, but needs its own IAM role and experiment template. ## Reports and recovery diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index e02378da4..57abcc814 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -920,10 +920,12 @@ def run_attempt( 2 or later) first downloads the earlier attempts, quiesces their trials, and passes them to the runner so that completed trials are not repeated. Changed trial records are uploaded every few seconds. An interruption - notice or a lost lease stops the runner; after a notice, the claim is - released as interrupted before the attempt is uploaded. Otherwise, after a - final check that the claim is still held, the job-level report merges - every attempt, and the ``done`` marker records the exit code. + notice or a lost claim stops the runner; after a notice, the attempt is + uploaded and the claim released as interrupted. Otherwise, after a + conditional renewal confirms that the claim is still held, the job's + report merges every attempt into ``job-summary.csv`` and + ``job-summary.md`` of this attempt, and the first ``done`` marker of the + job records the exit code and names that attempt. Args: queue: Queue of the job. @@ -1019,10 +1021,11 @@ def upload_attempt(timeout: float = -1) -> None: # The claim was taken over or expired; its new owner records the job. return Outcome("lost") if report_error is None: - # Summaries go first, so a recorded job always has them. + # The job's summary goes below this attempt, which the done marker names, + # so a former owner's late upload cannot replace the recorded summary. for name in ("summary.csv", "summary.md"): queue.s3.upload_file( - str(job_dir / name), queue.bucket, queue.key("results", job["id"], name) + str(job_dir / name), queue.bucket, queue.key(*remote, f"job-{name}") ) if time.time() >= claim.body["expires"]: # Publication outlasted the renewed claim, which another worker may own now. @@ -1038,6 +1041,7 @@ def upload_attempt(timeout: float = -1) -> None: "end": now(), "quiesce_errors": errors, "report_error": report_error, + "summary": None if report_error else "/".join(remote[1:]), "lost_slots": [key for key in watch.lost if key != claim.key], }, IfNoneMatch="*", diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 51d310362..9899d0ac9 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -789,10 +789,12 @@ def resumed_runner(item, workdir, out, resume, stop): assert runs == [("heavy", "attempt-2", ["attempt-1"]), ("next", "attempt-1", [])] # The interrupted attempt's queries are quiesced before the job continues. assert settled == ["attempt-1"] - order = [key for key in writes if key.endswith(("heavy/summary.csv", "done/heavy"))] - # A recorded job always has its summary. - assert order == [queue.key("results", "heavy", "summary.csv"), queue.key("done", "heavy")] - summary = s3.objects[queue.key("results", "heavy", "summary.csv")].decode() + summary_key = queue.key("results", "heavy", "attempt-2", "job-summary.csv") + order = [key for key in writes if key in {summary_key, queue.key("done", "heavy")}] + # A recorded job always has its summary, below the attempt that the marker names. + assert order == [summary_key, queue.key("done", "heavy")] + assert json.loads(queue.read("done", "heavy"))["summary"] == "heavy/attempt-2" + summary = s3.objects[summary_key].decode() row = next(csv.DictReader(io.StringIO(summary))) assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") assert json.loads(queue.read("done", "heavy"))["attempt"] == 2 @@ -1028,17 +1030,16 @@ def stall(lease, seconds, **changes): queue.renew = stall watch = JobWatch(queue, [claim], 600, 0.01, lambda: None, lambda: False) - try: - with watch: + with watch: + try: assert stalled.wait(5) # The claim expires while its renewal is still blocked in S3. assert watch.stop_requested.wait(5) assert not watch.interrupted.is_set() - assert stalled.is_set() assert not watch.lost + finally: + # Unblock the renewal before leaving the block joins the thread. released.set() - finally: - released.set() def test_losing_a_slot_does_not_stop_the_job(self, tmp_path): s3, queue = queue_with([job("heavy", heavy=True)], tmp_path) @@ -1103,7 +1104,7 @@ def sleep(seconds): monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries) == [] assert not queue.names("done") - assert queue.key("results", "a", "summary.csv") not in s3.objects + assert not [key for key in s3.objects if key.endswith("job-summary.csv")] def test_a_claim_renewed_once_then_stale_does_not_let_the_worker_exit(tmp_path, monkeypatch): From 3eefc13d39e4f87aaad95a99ae3c48c2d0ddd2f8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 17:38:50 +0900 Subject: [PATCH 11/16] Address code review findings on resume, leases, and the results grant - A resumed attempt runs its own warmups before the missing measured repetitions, and skips a case whose repetitions are all recorded. - An attempt stopped by a lost claim quiesces its queries before the worker claims another job, and the orchestrator runs in its own process group, which is killed when it ends. - With ResultsBucket set, the role can list the whole bucket, so a read of a missing key returns 404 instead of 403; object access stays below the stack's prefix. - Unexpected errors from notice checks are logged and treated as no notice instead of ending the watch thread or the worker. - A stack without the Results output (earlier template) uses the scratch bucket; --wait waits only while a stack operation is in progress. - Superseded failures are counted as superseded_failures and noted instead of dropped. - Trial records are uploaded at trial starts and ends, and renewals come before uploads. - acquire() reads first and writes only when the lease is absent or takeable. - A worker removes a rehearsal notice file when it starts; README explains Zones=3 and the operator's termination after a rehearsal. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 19 +-- benchmarks/cloudformation/benchmark.yaml | 15 +-- benchmarks/pyathena_bench/__main__.py | 10 +- benchmarks/pyathena_bench/aws.py | 9 +- benchmarks/pyathena_bench/fleet.py | 142 +++++++++++++++-------- benchmarks/pyathena_bench/report.py | 23 ++-- benchmarks/pyathena_bench/runner.py | 18 ++- benchmarks/tests/test_aws.py | 33 ++++++ benchmarks/tests/test_cloudformation.py | 15 ++- benchmarks/tests/test_fleet.py | 103 +++++++++++++++- benchmarks/tests/test_report.py | 5 +- benchmarks/tests/test_runner.py | 35 ++++++ 12 files changed, 332 insertions(+), 95 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 1af025d3e..31d80cd4c 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -97,7 +97,8 @@ Measurements depend on the CPU architecture; `environment.json` records it, and Instance size and disk size are parameters, so a later run can deliberately test another memory budget. The hosts are Spot Instances of one instance type by default; `Purchase=on-demand` launches On-Demand Instances of that type instead. The Auto Scaling group spreads the hosts over `Zones` Availability Zones (4 by default, or 3) with the price-capacity-optimized allocation strategy and never mixes instance types, so every host stays identical. -The template uses the first `Zones` zones that `Fn::GetAZs` returns, so deployment fails when it returns fewer, which depends on the zones in which the account's default VPC has default subnets; `us-west-2` has four. +The template uses the first `Zones` zones that `Fn::GetAZs` returns, which depends on the zones in which the account's default VPC has default subnets; `us-west-2`, the default region, has four, and the fourth widens Spot capacity, which was missing in two of three zones in the 4.0.0 run. +Where `Fn::GetAZs` returns three zones, deploy with `Zones=3`; with the default, the deployment fails with a template error before it creates any resource. The group replaces an interrupted Spot Instance; with `QueueName` set, the replacement's worker continues the interrupted job, see [Running on a fleet](#running-on-a-fleet). `environment.json` does not record the instance type or purchase option; a worker started with `--terminate-host` records both under `hosts/` in its queue, and the recovered `instances.json` lists the hosts that still exist. The AMI parameter resolves the current AL2023 image at deployment; record the resulting AMI when comparing environments. @@ -106,7 +107,8 @@ The instances can write only to the scratch bucket and database and to the resul The supplied template assumes ordinary IAM access and SSE-S3 source objects; Lake Formation restrictions or a customer-managed KMS key require corresponding grants before execution. The fleet queues, reports, and Athena trial outputs go to the stack's results location, the `Results` stack output. -With `ResultsBucket` set to an existing bucket, the results location is `s3://///`, and the instance role can write only below it; with `ResultsBucket` empty, the default, it is the scratch bucket. +With `ResultsBucket` set to an existing bucket, the results location is `s3://///`, and the instance role can read and write objects only below it, while it can list the whole bucket, because S3 answers a read of a missing key with 403 instead of 404 without that permission; with `ResultsBucket` empty, the default, it is the scratch bucket. +A stack created from an earlier template has no `Results` output and keeps everything in the scratch bucket. Input snapshots, initialization fixtures, and scratch tables stay in the scratch bucket and database either way. This repository's CI test bucket, which `AWS_ATHENA_S3_STAGING_DIR` names in CI, expires every object one day after it is written, at the following midnight UTC (`cloudformation/github_actions_oidc.yaml`). As the results location, it removes leftovers by itself, but reports must be recovered within that day. @@ -276,7 +278,7 @@ uv run --no-sync python -m pyathena_bench run \ ``` There is no automatic continuation of a partially recorded trial; rerun the selected comparison with the same manifest. -`run --resume` names the output directories of earlier attempts and skips the trials that they recorded as successful, which a fleet worker uses to continue an interrupted job. +`run --resume` names the output directories of earlier attempts and skips the measured trials that they recorded as successful, after running its own warmups, which a fleet worker uses to continue an interrupted job. Large API row cases may exceed the default one-hour timeout: 10 million rows require at least 10,000 pages with arraysize 1000, or 100,000 pages with arraysize 100. Choose the timeout using a smaller-scale pilot before the full run. Do not run two orchestrators concurrently on the same dedicated host. @@ -401,9 +403,9 @@ A worker claims a job by creating `claims/` with an S3 conditional write, s A claim, like each API slot under `slots/`, is a lease that its worker renews every tenth of `--lease-seconds` (600 by default, at least 300) while the job runs. When a worker dies, its claim and slots expire, and another worker takes them over with a conditional write, so the job and the slots return to the fleet without manual steps. A worker releases a slot by overwriting its lease as released, never by deleting it, so a slot that another worker took over stays held. -A worker whose claim was taken over, or whose renewals failed until the claim expired, stops the attempt and records nothing for the job; the new owner continues it, and the first `done` marker of a job is never replaced. +A worker whose claim was taken over, or whose renewals failed until the claim expired, stops the attempt, quiesces its queries as after a failed job, and records nothing for the job; the new owner continues it, and the first `done` marker of a job is never replaced. A lost API slot does not stop the job; the `done` marker lists it under `lost_slots`. -Each attempt of a job writes to `results//attempt-/`, and the worker uploads changed trial records every five seconds, so a lost host loses its running trial and at most the records of the last seconds. +Each attempt of a job writes to `results//attempt-/`, and the worker uploads its trial records within seconds of each trial's start and end, so a lost host loses only its running trial. A trial whose start was not uploaded before its host died is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. After a job, the worker uploads the attempt's output directory and its log to `logs/`, renews its claim conditionally to confirm that it still owns the job, writes `job-summary.csv` and `job-summary.md` from every attempt into its attempt directory, and then creates `done/` with the exit code, attempt, host, times, and the summary's attempt directory under `summary`. A failed trial stops only its own job; `status` lists jobs with a non-zero exit code, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. @@ -414,8 +416,8 @@ If it cannot confirm that, it records the errors in the `done` marker and exits. Each worker checks instance metadata for a Spot interruption notice (`spot/instance-action`), which EC2 publishes about two minutes before it reclaims an instance. On a notice, the worker interrupts the running trial with SIGINT, which cancels the trial's reported queries, uploads the attempt, releases its claim as `interrupted` and its slots, and exits without terminating the host. A separate thread watches for the notice, so slow S3 requests do not delay it. -The group replaces the reclaimed instance, and the new host's worker takes over the interrupted claim as the next attempt: it downloads the earlier attempts, quiesces their trials as after a failed job, and runs only the repetitions that they did not record as successful, through `run --resume`. -The job's summary merges the attempts and leaves out a failed or incomplete trial whose repetition succeeded in a later attempt, so it reports the trials of an uninterrupted run. +The group replaces the reclaimed instance, and the new host's worker takes over the interrupted claim as the next attempt: it downloads the earlier attempts, quiesces their trials as after a failed job, and runs its own warmups and then only the measured repetitions that they did not record as successful, through `run --resume`. +The job's summary merges the attempts and counts a failed or incomplete trial whose repetition succeeded in a later attempt as `superseded_failures` instead of `failed_trials`, with its reason in the notes, so it reports the trials of an uninterrupted run without hiding failures. `report results/recovered/fleet//results/` produces the same merged summary from the recovered attempts. Rebalance recommendations are not acted on; an interruption does not always follow them, and the notice leaves enough time to stop a trial. @@ -452,7 +454,8 @@ uv run --no-sync python -m pyathena_bench queue \ --stack "$(cat stack-id.txt)" --name pilot-1 --jobs results/jobs-pilot.json --manifest results/input-pilot.json ``` -When `status` shows the job claimed and `results//attempt-1/trials.jsonl` has a trial, create the file `/opt/pyathena/benchmarks/interruption-notice`, which a worker treats as an interruption notice. +When `status` shows the job claimed and `results//attempt-1/trials.jsonl` has a trial, create the file `/opt/pyathena/benchmarks/interruption-notice`, which a worker treats as an interruption notice and removes when it starts. +As after a real notice, the worker exits and leaves its host running, so terminate the host as below. Check that the claim is `interrupted`, then terminate the host without lowering the desired capacity, as EC2 does when it reclaims a Spot Instance: ```bash diff --git a/benchmarks/cloudformation/benchmark.yaml b/benchmarks/cloudformation/benchmark.yaml index bbaf4646f..3a9b8e75a 100644 --- a/benchmarks/cloudformation/benchmark.yaml +++ b/benchmarks/cloudformation/benchmark.yaml @@ -300,22 +300,15 @@ Resources: 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. - # ListBucketMultipartUploads does not support the s3:prefix condition key. + # 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:ListBucketMultipartUploads] + Action: [s3:GetBucketLocation, s3:ListBucket, s3:ListBucketMultipartUploads] Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket} - !Ref AWS::NoValue - - !If - - HasResultsBucket - - Effect: Allow - Action: s3:ListBucket - Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket} - Condition: - StringLike: - s3:prefix: !Sub '${ResultsPrefix}/${AWS::StackName}/*' - - !Ref AWS::NoValue - !If - HasResultsBucket - Effect: Allow diff --git a/benchmarks/pyathena_bench/__main__.py b/benchmarks/pyathena_bench/__main__.py index ba31c8e73..f50a9bc56 100644 --- a/benchmarks/pyathena_bench/__main__.py +++ b/benchmarks/pyathena_bench/__main__.py @@ -33,6 +33,7 @@ from pyathena_bench.cases import FAMILIES, matrix from pyathena_bench.config import Settings, read_json from pyathena_bench.fleet import ( + SIMULATED_NOTICE, Filters, Queue, SpotInterruptionError, @@ -125,7 +126,9 @@ def parser() -> argparse.ArgumentParser: "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" + "--wait", + action="store_true", + help="Wait for the stack's outputs and for the queue to be published instead of failing", ) consume.add_argument( "--lease-seconds", @@ -233,6 +236,8 @@ def main(argv: list[str] | None = None) -> int: if args.lease_seconds < 300: raise ValueError("--lease-seconds must be at least 300") autoscaling = client(session_, "autoscaling") + # A rehearsed notice applies to one worker; a restarted worker starts clean. + SIMULATED_NOTICE.unlink(missing_ok=True) if args.terminate_host: suspend_rebalance(autoscaling) try: @@ -246,7 +251,8 @@ def main(argv: list[str] | None = None) -> int: wait=args.wait, ) except SpotInterruptionError as interrupted: - # The host is being reclaimed; its group launches a replacement. + # 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" ) diff --git a/benchmarks/pyathena_bench/aws.py b/benchmarks/pyathena_bench/aws.py index b7b190b20..af33eab2e 100644 --- a/benchmarks/pyathena_bench/aws.py +++ b/benchmarks/pyathena_bench/aws.py @@ -38,7 +38,7 @@ def client(session_: Any, service: str) -> Any: class IncompleteStackError(ValueError): - """The stack has no complete outputs yet, for example while it is being created.""" + """The stack's outputs are incomplete while an operation on it is in progress.""" def split_location(location: str) -> tuple[str, str]: @@ -74,10 +74,13 @@ def stack_resources(session_: Any, stack: str) -> dict[str, str]: "SourceTable", "WorkGroup", "AutoScalingGroup", - "Results", } if not required <= result.keys(): - raise IncompleteStackError("Stack outputs are incomplete") + if data.get("StackStatus", "").endswith("_IN_PROGRESS"): + raise IncompleteStackError("Stack outputs are incomplete") + raise ValueError("Stack outputs are incomplete") + # A stack from a template before the Results output keeps everything in the scratch bucket. + result.setdefault("Results", f"s3://{result['Bucket']}/") if ( result["SourceBucket"] == result["Bucket"] or result["ScratchDatabase"] == result["SourceDatabase"] diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index 57abcc814..36bf373f5 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -9,7 +9,11 @@ from __future__ import annotations +import contextlib +import http.client import json +import logging +import os import signal import socket import subprocess @@ -40,6 +44,7 @@ # A file that a worker treats like a Spot interruption notice, to rehearse interruptions. SIMULATED_NOTICE = Path("interruption-notice") METADATA_URL = "http://169.254.169.254/latest" +_logger = logging.getLogger(__name__) # S3 errors of a conditional write whose condition failed. PRECONDITION_ERRORS = {"PreconditionFailed", "ConditionalRequestConflict"} @@ -394,6 +399,7 @@ def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Leas Returns: The lease, or None if another worker holds it. """ + current = self.read_lease(key) at = time.time() body = { **owner, @@ -402,11 +408,10 @@ def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Leas "renewed": at, "expires": at + lease_seconds, } - lease = self.put_lease(key, body, IfNoneMatch="*") - if lease is not None: - return lease - current = self.read_lease(key) - if current is None or not current.takeable(time.time()): + if current is None: + # Exactly one worker creates the object; another one leaves it. + return self.put_lease(key, body, IfNoneMatch="*") + if not current.takeable(at): return None body["attempt"] = current.body.get("attempt", 1) + 1 return self.put_lease(key, body, IfMatch=current.etag) @@ -640,7 +645,9 @@ def run_job( When ``stop`` is set, the orchestrator receives SIGINT, which stops the running trial, cancels its reported queries, and records the trial as - interrupted; it is killed if it has not exited a minute later. + interrupted; it is killed if it has not exited a minute later. The + orchestrator runs in its own process group, which is killed when it + ends, so no trial process outlives it. Args: job: Job dictionary. @@ -661,16 +668,22 @@ def run_job( if resume: command += ["--resume", *map(str, resume)] with (workdir / "logs" / f"{job['id']}-{output.name}.log").open("w") as stream: - process = subprocess.Popen(command, stdout=stream, stderr=subprocess.STDOUT) - while process.poll() is None: - if stop.wait(1): - process.send_signal(signal.SIGINT) - try: - return process.wait(timeout=60) - except subprocess.TimeoutExpired: - process.kill() - return process.wait() - return process.returncode + process = subprocess.Popen( + command, stdout=stream, stderr=subprocess.STDOUT, start_new_session=True + ) + try: + while process.poll() is None: + if stop.wait(1): + process.send_signal(signal.SIGINT) + try: + return process.wait(timeout=60) + except subprocess.TimeoutExpired: + process.kill() + return process.wait() + return process.returncode + finally: + with contextlib.suppress(ProcessLookupError, PermissionError): + os.killpg(process.pid, signal.SIGKILL) def trial_evidence(output: Path) -> tuple[set[str], tuple[str, ...]]: @@ -823,7 +836,7 @@ def stop(self, timeout: float | None = None) -> None: def _watch_notice(self) -> None: while not self._stopped.wait(self.poll_seconds): - if self.notice(): + if check_notice(self.notice): self.interrupted.set() self.stop_requested.set() return @@ -835,27 +848,30 @@ def _watch_notice(self) -> None: def _renew(self) -> None: renew_at = time.monotonic() + self.renew_seconds while not self._stopped.wait(self.poll_seconds): + # Renewals come first, so that a slow upload does not delay them. + if time.monotonic() >= renew_at: + renew_at = time.monotonic() + self.renew_seconds + self._renew_leases() try: with self.upload_lock: self.progress() except Exception: # The next poll uploads again. - pass - if time.monotonic() < renew_at: - continue - renew_at = time.monotonic() + self.renew_seconds - for lease in list(self.leases): - try: - if self.queue.renew(lease, self.lease_seconds): - continue - except Exception: - if time.time() < lease.body["expires"]: - # The lease has time left for a retry. - continue - self.leases.remove(lease) - self.lost.append(lease.key) - if lease is self.claim: - self.stop_requested.set() + _logger.warning("Could not upload the attempt's progress", exc_info=True) + + def _renew_leases(self) -> None: + for lease in list(self.leases): + try: + if self.queue.renew(lease, self.lease_seconds): + continue + except Exception: + if time.time() < lease.body["expires"]: + # The lease has time left for a retry. + continue + self.leases.remove(lease) + self.lost.append(lease.key) + if lease is self.claim: + self.stop_requested.set() def pending_jobs( @@ -919,7 +935,7 @@ def run_attempt( The attempt writes to ``results//attempt-/``. A takeover (attempt 2 or later) first downloads the earlier attempts, quiesces their trials, and passes them to the runner so that completed trials are not repeated. - Changed trial records are uploaded every few seconds. An interruption + Trial records are uploaded within seconds of each trial's start and end. An interruption notice or a lost claim stops the runner; after a notice, the attempt is uploaded and the claim released as interrupted. Otherwise, after a conditional renewal confirms that the claim is still held, the job's @@ -947,14 +963,25 @@ def run_attempt( output = job_dir / f"attempt-{attempt}" remote = ("results", job["id"], output.name) log = workdir / "logs" / f"{job['id']}-{output.name}.log" - uploaded: dict[str, int] = {} + checked: dict[str, int] = {} def progress() -> None: - for path in sorted(output.glob("*.jsonl")): - size = path.stat().st_size - if uploaded.get(path.name) != size: - queue.s3.upload_file(str(path), queue.bucket, queue.key(*remote, path.name)) - uploaded[path.name] = size + # Upload at trial boundaries: when a trial starts (its output prefix is + # needed to quiesce it) and when it is recorded (needed to resume). + events, trials = output / "events.jsonl", output / "trials.jsonl" + boundary = False + if events.exists(): + with events.open("rb") as stream: + stream.seek(checked.get("events", 0)) + boundary = b'"trial_start"' in stream.read() + checked["events"] = stream.tell() + if trials.exists() and trials.stat().st_size != checked.get("trials"): + checked["trials"] = trials.stat().st_size + boundary = True + if boundary: + for path in (events, trials): + if path.exists(): + queue.s3.upload_file(str(path), queue.bucket, queue.key(*remote, path.name)) with JobWatch(queue, [claim, *slots], lease_seconds, renew_seconds, progress, notice) as watch: @@ -994,8 +1021,13 @@ def upload_attempt(timeout: float = -1) -> None: # from every recorded trial; both are small within a notice's two minutes. upload_attempt(timeout=10) if not watch.interrupted.is_set(): - # The claim was lost; another worker continues the job. - return Outcome("lost") + # The claim was lost; another worker continues the job. This + # attempt's queries are stopped before this worker claims another. + try: + errors += settle(output) + except Exception as exc: + errors.append(f"Could not quiesce the job: {type(exc).__name__}: {exc}") + return Outcome("lost", tuple(errors)) queue.end_lease(claim, "interrupted") return Outcome("interrupted") if not errors: @@ -1122,7 +1154,7 @@ def settle(output: Path) -> list[str]: # Renewal times of other workers' claims when this worker found nothing to claim. waiting: dict[str, float | None] = {} while True: - if notice(): + if check_notice(notice): raise SpotInterruptionError(completed) pending, running = pending_jobs(queue, jobs) if pending: @@ -1180,9 +1212,9 @@ def settle(output: Path) -> list[str]: raise SpotInterruptionError(completed) if outcome.state == "done": completed.append(job["id"]) - if outcome.errors: - # Unconfirmed queries could overlap later trials on this host. - return completed + if outcome.errors: + # Unconfirmed queries could overlap later trials on this host. + return completed started = True break finally: @@ -1216,6 +1248,22 @@ def instance_metadata(path: str, timeout: float = 2.0) -> str: return response.read().decode() +def check_notice(notice: Callable[[], bool]) -> bool: + """Call a notice check, treating an unexpected error as no notice. + + Args: + notice: Function that returns whether an interruption notice exists. + + Returns: + The result, or False if the check raised. + """ + try: + return notice() + except Exception: + _logger.warning("Could not check for an interruption notice", exc_info=True) + return False + + def interruption_notice( metadata: Callable[[str], str] | None = None, simulated: Path = SIMULATED_NOTICE ) -> bool: @@ -1238,7 +1286,7 @@ def interruption_notice( return True try: (metadata or instance_metadata)("spot/instance-action") - except (urllib.error.URLError, OSError): + except (urllib.error.URLError, OSError, http.client.HTTPException): return False return True diff --git a/benchmarks/pyathena_bench/report.py b/benchmarks/pyathena_bench/report.py index 72790112f..78308aa1a 100644 --- a/benchmarks/pyathena_bench/report.py +++ b/benchmarks/pyathena_bench/report.py @@ -72,6 +72,8 @@ def summarize(trials: list[dict[str, Any]]) -> list[dict[str, Any]]: if trial["status"] == "ok": continue prefix = "Warmup: " if trial.get("warmup") else "" + if trial.get("superseded"): + prefix = "Superseded: " failures = [q for q in trial.get("queries", []) if q["status"] != "ok"] if failures: for query in failures: @@ -87,7 +89,11 @@ def summarize(trials: list[dict[str, Any]]) -> list[dict[str, Any]]: "scale": scale, "case": case_id, "successful_trials": len(good), - "failed_trials": sum(t["status"] not in {"ok", "unsupported"} for t in measured), + "failed_trials": sum( + t["status"] not in {"ok", "unsupported"} and not t.get("superseded") + for t in measured + ), + "superseded_failures": sum(bool(t.get("superseded")) for t in measured), "failed_warmups": sum( t.get("warmup", False) and t["status"] != "ok" for t in group ), @@ -150,9 +156,10 @@ def collect_trials(directory: Path) -> list[dict[str, Any]]: A fleet job keeps each attempt in an ``attempt-`` subdirectory; their trials are merged. A started trial without a final record is reported as incomplete. A failed - or incomplete trial is dropped when another attempt recorded the same - repetition as successful, so a resumed job reports the trials of an - uninterrupted run; so is a second successful record of a repetition. + or incomplete trial is marked ``superseded`` when another attempt recorded + the same measured repetition as successful, so a resumed job counts the + trials of an uninterrupted run while its failures stay visible; a second + successful record of a repetition is dropped. Args: directory: Output directory of a run, or a job directory with attempts. @@ -170,7 +177,9 @@ def collect_trials(directory: Path) -> list[dict[str, Any]]: if event["event"] == "trial_start" and event["trial"] not in recorded ) succeeded = { - (t["scale"], t["case_id"], t.get("repetition")) for t in trials if t["status"] == "ok" + (t["scale"], t["case_id"], t.get("repetition")) + for t in trials + if t["status"] == "ok" and not t.get("warmup") } merged = [] kept: set[tuple[Any, ...]] = set() @@ -182,8 +191,8 @@ def collect_trials(directory: Path) -> list[dict[str, Any]]: if len(sources) > 1 and key in kept: continue kept.add(key) - elif trial["status"] != "unsupported" and key in succeeded: - continue + elif trial["status"] != "unsupported" and key in succeeded and not trial.get("warmup"): + trial = {**trial, "superseded": True} merged.append(trial) return merged diff --git a/benchmarks/pyathena_bench/runner.py b/benchmarks/pyathena_bench/runner.py index 89e3914b7..4e0d0d45d 100644 --- a/benchmarks/pyathena_bench/runner.py +++ b/benchmarks/pyathena_bench/runner.py @@ -333,13 +333,14 @@ def completed_trials(directories: list[Path]) -> set[tuple[str, str, int | None] directories: Output directories of earlier attempts. Returns: - ``(scale, case_id, repetition)`` for successful trials, and - ``(scale, case_id, None)`` for cases recorded as unsupported. + ``(scale, case_id, repetition)`` for successful measured trials, and + ``(scale, case_id, None)`` for cases recorded as unsupported. Warmups + are not included: each attempt runs its own before its measured trials. """ completed: set[tuple[str, str, int | None]] = set() for directory in directories: for trial in read_records(directory / "trials.jsonl"): - if trial["status"] == "ok": + if trial["status"] == "ok" and not trial.get("warmup"): completed.add((trial["scale"], trial["case_id"], trial["repetition"])) elif trial["status"] == "unsupported": completed.add((trial["scale"], trial["case_id"], None)) @@ -363,7 +364,8 @@ def run( scales: Scales to run. output: New output directory. completed: Trials from :func:`completed_trials` to skip, so that a - resumed job runs only the missing repetitions. + resumed job runs only the missing measured repetitions, after its + own warmups; a case without missing repetitions is skipped. Returns: True if every trial succeeded. @@ -391,8 +393,14 @@ def run( }, ) continue + if all( + (scale, case.id, repetition) in completed + for repetition in range(settings.repetitions) + ): + continue for repetition in range(-settings.warmups, settings.repetitions): - if (scale, case.id, repetition) in completed: + # Warmups run again on this host; measured trials are not repeated. + if repetition >= 0 and (scale, case.id, repetition) in completed: continue trial = uuid.uuid4().hex payload = { diff --git a/benchmarks/tests/test_aws.py b/benchmarks/tests/test_aws.py index 53286d9e9..948012210 100644 --- a/benchmarks/tests/test_aws.py +++ b/benchmarks/tests/test_aws.py @@ -11,9 +11,11 @@ import pytest from botocore.exceptions import ClientError from pyathena_bench.aws import ( + IncompleteStackError, cleanup, prepare, split_location, + stack_resources, validate_inputs, validate_manifest, ) @@ -346,3 +348,34 @@ def test_split_location(location, expected): def test_split_location_rejects_invalid_locations(location): with pytest.raises(ValueError, match="Invalid S3 location"): split_location(location) + + +@pytest.mark.parametrize( + ("status", "outputs", "error"), + [ + ("CREATE_IN_PROGRESS", {}, IncompleteStackError), + ("CREATE_COMPLETE", {}, ValueError), + ( + "CREATE_COMPLETE", + {k: v for k, v in RESOURCES.items() if k not in {"StackId", "Results"}}, + None, + ), + ], + ids=["creating", "complete_without_outputs", "earlier_template"], +) +def test_stack_resources_waits_only_while_in_progress(monkeypatch, status, outputs, error): + stack = { + "StackId": "arn:stack", + "StackStatus": status, + "Tags": [{"Key": "Purpose", "Value": "pyathena-benchmark"}], + "Outputs": [{"OutputKey": k, "OutputValue": v} for k, v in outputs.items()], + } + fake = SimpleNamespace(describe_stacks=lambda **kwargs: {"Stacks": [stack]}) + monkeypatch.setattr("pyathena_bench.aws.client", lambda *args: fake) + if error is None: + # A stack from a template before the Results output keeps results in the scratch bucket. + assert stack_resources(None, "stack")["Results"] == "s3://temporary/" + return + with pytest.raises(error) as raised: + stack_resources(None, "stack") + assert (type(raised.value) is IncompleteStackError) is (error is IncompleteStackError) diff --git a/benchmarks/tests/test_cloudformation.py b/benchmarks/tests/test_cloudformation.py index 74b8c3ff5..fb99d1a28 100644 --- a/benchmarks/tests/test_cloudformation.py +++ b/benchmarks/tests/test_cloudformation.py @@ -149,13 +149,16 @@ def test_results_location_is_scoped_to_the_stack_prefix(): assert statement["Resource"]["Fn::Sub"].endswith( ":s3:::${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/*" ) - elif "s3:ListBucket" in statement["Action"]: - assert statement["Condition"] == { - "StringLike": {"s3:prefix": {"Fn::Sub": "${ResultsPrefix}/${AWS::StackName}/*"}} - } else: - # Neither action supports the s3:prefix condition key. - assert statement["Action"] == ["s3:GetBucketLocation", "s3:ListBucketMultipartUploads"] + # Listing and bucket metadata only; object access stays below the prefix. + assert statement["Action"] == [ + "s3:GetBucketLocation", + "s3:ListBucket", + "s3:ListBucketMultipartUploads", + ] + assert statement["Resource"] == { + "Fn::Sub": "arn:${AWS::Partition}:s3:::${ResultsBucket}" + } assert "Condition" not in statement condition, results_bucket, scratch = template["Outputs"]["Results"]["Value"]["Fn::If"] assert condition == "HasResultsBucket" diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 9899d0ac9..327b615f6 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import csv +import http.client import io import json import signal @@ -27,6 +28,7 @@ Lease, Queue, SpotInterruptionError, + check_notice, expand_jobs, interruption_notice, merge_jobs, @@ -736,7 +738,9 @@ def interrupted_runner(item, workdir, out, resume, stop): notice.set() assert stop.wait(5) with (out / "trials.jsonl").open("a") as stream: - stream.write(json.dumps({**ok_trial(1), "status": "interrupted"}) + "\n") + interrupted = {**ok_trial(1), "status": "interrupted"} + del interrupted["queries"] + stream.write(json.dumps(interrupted) + "\n") return 1 with pytest.raises(SpotInterruptionError) as raised: @@ -797,6 +801,9 @@ def resumed_runner(item, workdir, out, resume, stop): summary = s3.objects[summary_key].decode() row = next(csv.DictReader(io.StringIO(summary))) assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") + # The interrupted trial stays visible. + assert row["superseded_failures"] == "1" + assert "Superseded: interrupted" in row["notes"] assert json.loads(queue.read("done", "heavy"))["attempt"] == 2 def test_a_notice_while_quiescing_an_earlier_attempt_skips_the_run(self, tmp_path): @@ -854,8 +861,11 @@ def test_worker_waits_for_publication(self, tmp_path, monkeypatch): class FakeProcess: """A child process that runs until it receives SIGINT.""" - def __init__(self, command, stdout, stderr): + pid = 4321 + + def __init__(self, command, stdout, stderr, start_new_session): self.command = command + self.start_new_session = start_new_session self.signals = [] def poll(self): @@ -871,11 +881,15 @@ def wait(self, timeout=None): def test_run_job_resumes_earlier_attempts_and_interrupts_with_sigint(tmp_path, monkeypatch): processes = [] - def popen(command, stdout, stderr): - processes.append(FakeProcess(command, stdout, stderr)) + def popen(command, stdout, stderr, start_new_session): + processes.append(FakeProcess(command, stdout, stderr, start_new_session)) return processes[-1] + killed = [] monkeypatch.setattr("pyathena_bench.fleet.subprocess.Popen", popen) + monkeypatch.setattr( + "pyathena_bench.fleet.os.killpg", lambda pid, signum: killed.append((pid, signum)) + ) stop = threading.Event() stop.set() earlier = tmp_path / "results" / "a" / "attempt-1" @@ -889,6 +903,9 @@ def popen(command, stdout, stderr): assert code == 130 [process] = processes assert process.signals == [signal.SIGINT] + # The orchestrator's own process group, including any trial process, is killed. + assert process.start_new_session + assert killed == [(4321, signal.SIGKILL)] assert process.command[-4:] == ["--suite", "single", "--resume", str(earlier)] assert (tmp_path / "logs" / "a-attempt-2.log").exists() @@ -921,6 +938,7 @@ def test_interruption_is_marked_after_a_concurrent_renewal(self, tmp_path): def test_a_lost_claim_stops_the_attempt_without_recording_the_job(self, tmp_path, monkeypatch): s3, queue = queue_with([job("a")], tmp_path) stopped = [] + settled = [] def sleep(seconds): # The new owner renews the claim while this worker waits to exit. @@ -947,9 +965,11 @@ def runner(item, workdir, out, resume, stop): api_slots=1, poll_seconds=0, runner=runner, - settle=no_queries, + settle=lambda output: settled.append(output.name) or [], renew_seconds=0.01, ) + # The stopped attempt's queries are quiesced before the worker goes on. + assert settled == ["attempt-1"] assert stopped == [True] assert ran == [] assert not queue.names("done") @@ -1139,3 +1159,76 @@ def sleep(seconds): ) assert ran == ["a"] assert runs == ["attempt-2"] + + +class TestReviewRepairs: + def test_acquire_does_not_write_a_live_lease(self, tmp_path): + s3, queue = queue_with([], tmp_path) + key = queue.key("claims", "a") + queue.acquire(key, {"host": "busy"}, 600) + attempts = [] + original = s3.put_object + s3.put_object = lambda **kwargs: attempts.append(kwargs["Key"]) or original(**kwargs) + assert queue.acquire(key, {"host": "me"}, 600) is None + assert attempts == [] + + def test_notice_errors_do_not_stop_the_watch_or_the_worker(self, tmp_path): + def incomplete(path): + raise http.client.IncompleteRead(b"") + + assert not interruption_notice(incomplete, tmp_path / "notice") + + def broken(): + raise RuntimeError("unexpected") + + assert not check_notice(broken) + _, queue = queue_with([job("a")], tmp_path) + assert work(queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, settle=no_queries, + notice=broken) == ["a"] # fmt: skip + + def test_progress_uploads_at_trial_boundaries(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + events_key = queue.key("results", "a", "attempt-1", "events.jsonl") + uploads = [] + original = s3.upload_file + + def recording(filename, bucket, key): + uploads.append(key) + return original(filename, bucket, key) + + s3.upload_file = recording + + def runner(item, workdir, out, resume, stop): + out.mkdir(parents=True) + events = out / "events.jsonl" + events.write_text(json.dumps({"event": "trial_start", "trial": "t"}) + "\n") + deadline = time.monotonic() + 5 + while events_key not in uploads and time.monotonic() < deadline: + time.sleep(0.01) + started = uploads.count(events_key) + with events.open("a") as stream: + for _ in range(20): + stream.write(json.dumps({"event": "query", "query_id": "q"}) + "\n") + stream.flush() + time.sleep(0.01) + # Query events between trial boundaries are not uploaded. + assert uploads.count(events_key) == started == 1 + return 0 + + assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries, + renew_seconds=0.01) == ["a"] # fmt: skip + + def test_worker_removes_a_rehearsal_notice_when_it_starts(self, tmp_path, monkeypatch): + config = tmp_path / "config.toml" + config.write_text("[scales]\nsmall = 10000\n") + monkeypatch.chdir(tmp_path) + Path("interruption-notice").touch() + monkeypatch.setattr( + "pyathena_bench.__main__.wait_for_stack", lambda *a: {"Results": "s3://bucket/"} + ) + monkeypatch.setattr("pyathena_bench.__main__.session", lambda settings: None) + monkeypatch.setattr("pyathena_bench.__main__.client", lambda *a: FakeS3()) + monkeypatch.setattr("pyathena_bench.__main__.work", lambda *a, **k: []) + argv = ["--config", str(config), "worker", "--stack", "s", "--name", "q"] + assert main(argv) == 0 + assert not Path("interruption-notice").exists() diff --git a/benchmarks/tests/test_report.py b/benchmarks/tests/test_report.py index 553782182..506657648 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -131,7 +131,7 @@ def trial(repetition, status="ok"): first.mkdir() second.mkdir() (first / "trials.jsonl").write_text( - json.dumps(trial(0)) + "\n" + json.dumps(trial(1, "interrupted")) + "\n" + json.dumps(trial(0)) + "\n" + json.dumps({**trial(1, "interrupted"), "queries": []}) + "\n" ) # The interrupted attempt also started a trial that it never recorded. start = {"event": "trial_start", "trial": "lost", "case_id": "case", "scale": "small"} @@ -141,6 +141,9 @@ def trial(repetition, status="ok"): with (tmp_path / "summary.csv").open() as stream: row = next(csv.DictReader(stream)) assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") + # The superseded failures stay visible: the interrupted and the unrecorded trial. + assert row["superseded_failures"] == "2" + assert "Superseded: interrupted" in row["notes"] def test_report_counts_a_repetition_recorded_by_two_attempts_once(tmp_path): diff --git a/benchmarks/tests/test_runner.py b/benchmarks/tests/test_runner.py index 5c22b1b38..4867b0656 100644 --- a/benchmarks/tests/test_runner.py +++ b/benchmarks/tests/test_runner.py @@ -169,3 +169,38 @@ def succeed(payload, events, settings): ) assert [p["repetition"] for p in payloads] == [1, 2] assert all(p["output"].startswith(f"s3://res/p/runs/{'a' * 32}/trials/") for p in payloads) + + +@pytest.mark.parametrize( + ("finished", "expected"), [([0], [-1, 1, 2]), ([0, 1, 2], [])], ids=["missing", "complete"] +) +def test_resumed_run_repeats_its_warmups_before_missing_repetitions( + monkeypatch, tmp_path, finished, expected +): + monkeypatch.setattr( + "pyathena_bench.runner.validate_manifest", + lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out", "Results": "s3://out/"}, + ) + monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) + monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) + payloads = [] + monkeypatch.setattr( + "pyathena_bench.runner.supervise", + lambda payload, events, settings: payloads.append(payload) or {"status": "ok"}, + ) + case = Case("cursor") + earlier = tmp_path / "attempt-1" + earlier.mkdir() + records = [ + {"scale": "small", "case_id": case.id, "repetition": -1, "warmup": True, "status": "ok"} + ] + records += [ + {"scale": "small", "case_id": case.id, "repetition": r, "status": "ok"} for r in finished + ] + (earlier / "trials.jsonl").write_text("".join(json.dumps(r) + "\n" for r in records)) + manifest = {"run_id": "a" * 32, "scales": {"small": {"table": "t", "rows": 10}}, "fixtures": {}} + settings = replace(Settings(), warmups=1, repetitions=3) + assert run( + settings, manifest, [case], ["small"], tmp_path / "attempt-2", completed_trials([earlier]) + ) + assert [p["repetition"] for p in payloads] == expected From ae77dab737c27daaba4032c2de604cc45fa4fa54 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 17:50:14 +0900 Subject: [PATCH 12/16] Keep the host after an unconfirmed stopped attempt and retry uploads - A worker that cannot confirm that a stopped attempt's queries ended raises instead of retiring its host, since the unfinished job may have no other worker. - A trial boundary stays pending until its upload succeeds. - README: a three-zone region with the default Zones fails with a stack error. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 3 +- benchmarks/pyathena_bench/fleet.py | 20 ++++++++--- benchmarks/tests/test_fleet.py | 55 ++++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 6 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 31d80cd4c..9bd7d4100 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -98,7 +98,7 @@ Instance size and disk size are parameters, so a later run can deliberately test The hosts are Spot Instances of one instance type by default; `Purchase=on-demand` launches On-Demand Instances of that type instead. The Auto Scaling group spreads the hosts over `Zones` Availability Zones (4 by default, or 3) with the price-capacity-optimized allocation strategy and never mixes instance types, so every host stays identical. The template uses the first `Zones` zones that `Fn::GetAZs` returns, which depends on the zones in which the account's default VPC has default subnets; `us-west-2`, the default region, has four, and the fourth widens Spot capacity, which was missing in two of three zones in the 4.0.0 run. -Where `Fn::GetAZs` returns three zones, deploy with `Zones=3`; with the default, the deployment fails with a template error before it creates any resource. +Where `Fn::GetAZs` returns three zones, deploy with `Zones=3`; with the default, the deployment fails with a stack error. The group replaces an interrupted Spot Instance; with `QueueName` set, the replacement's worker continues the interrupted job, see [Running on a fleet](#running-on-a-fleet). `environment.json` does not record the instance type or purchase option; a worker started with `--terminate-host` records both under `hosts/` in its queue, and the recovered `instances.json` lists the hosts that still exist. The AMI parameter resolves the current AL2023 image at deployment; record the resulting AMI when comparing environments. @@ -404,6 +404,7 @@ A claim, like each API slot under `slots/`, is a lease that its worker renews ev When a worker dies, its claim and slots expire, and another worker takes them over with a conditional write, so the job and the slots return to the fleet without manual steps. A worker releases a slot by overwriting its lease as released, never by deleting it, so a slot that another worker took over stays held. A worker whose claim was taken over, or whose renewals failed until the claim expired, stops the attempt, quiesces its queries as after a failed job, and records nothing for the job; the new owner continues it, and the first `done` marker of a job is never replaced. +If that worker cannot confirm that the queries stopped, it raises an error and keeps its host, because the job may have no other worker. A lost API slot does not stop the job; the `done` marker lists it under `lost_slots`. Each attempt of a job writes to `results//attempt-/`, and the worker uploads its trial records within seconds of each trial's start and end, so a lost host loses only its running trial. A trial whose start was not uploaded before its host died is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index 36bf373f5..e8b1f2d1d 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -967,21 +967,23 @@ def run_attempt( def progress() -> None: # Upload at trial boundaries: when a trial starts (its output prefix is - # needed to quiesce it) and when it is recorded (needed to resume). + # needed to quiesce it) and when it is recorded (needed to resume). A + # boundary stays pending until its upload succeeds. events, trials = output / "events.jsonl", output / "trials.jsonl" - boundary = False if events.exists(): with events.open("rb") as stream: stream.seek(checked.get("events", 0)) - boundary = b'"trial_start"' in stream.read() + if b'"trial_start"' in stream.read(): + checked["pending"] = 1 checked["events"] = stream.tell() if trials.exists() and trials.stat().st_size != checked.get("trials"): checked["trials"] = trials.stat().st_size - boundary = True - if boundary: + checked["pending"] = 1 + if checked.get("pending"): for path in (events, trials): if path.exists(): queue.s3.upload_file(str(path), queue.bucket, queue.key(*remote, path.name)) + checked["pending"] = 0 with JobWatch(queue, [claim, *slots], lease_seconds, renew_seconds, progress, notice) as watch: @@ -1131,6 +1133,8 @@ def work( ``wait`` is false, or a job needs more API slots than ``api_slots``. SpotInterruptionError: After an interruption notice, once the running attempt's records are uploaded and its claim is released as interrupted. + RuntimeError: If an attempt stopped by a lost claim could not be quiesced; + the worker's host is then kept. """ while not queue.exists("queue.json"): if not wait: @@ -1210,6 +1214,12 @@ def settle(output: Path) -> list[str]: ) if outcome.state == "interrupted": raise SpotInterruptionError(completed) + if outcome.state == "lost" and outcome.errors: + # The job is unfinished, so the host is kept rather than retired. + raise RuntimeError( + f"Could not quiesce a stopped attempt of {job['id']}: " + + "; ".join(outcome.errors) + ) if outcome.state == "done": completed.append(job["id"]) if outcome.errors: diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 327b615f6..60ac7875c 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -1232,3 +1232,58 @@ def test_worker_removes_a_rehearsal_notice_when_it_starts(self, tmp_path, monkey argv = ["--config", str(config), "worker", "--stack", "s", "--name", "q"] assert main(argv) == 0 assert not Path("interruption-notice").exists() + + +class TestReviewRepairsFollowUp: + def test_an_unconfirmed_stopped_attempt_keeps_the_host(self, tmp_path, monkeypatch): + s3, queue = queue_with([job("a")], tmp_path) + + def runner(item, workdir, out, resume, stop): + s3.put_object( + Key=queue.key("claims", "a"), + Body=json.dumps( + {"host": "other", "job": "a", "attempt": 2, "state": "running", + "renewed": time.time(), "expires": time.time() + 600} + ).encode(), + ) # fmt: skip + assert stop.wait(5) + return 130 + + with pytest.raises(RuntimeError, match="Could not quiesce a stopped attempt of a"): + work( + queue, + tmp_path / "w", + api_slots=1, + runner=runner, + settle=lambda output: ["q1: still RUNNING"], + renew_seconds=0.01, + ) + assert not queue.names("done") + + def test_a_failed_trial_start_upload_is_retried(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + events_key = queue.key("results", "a", "attempt-1", "events.jsonl") + original = s3.upload_file + failures = [] + + def flaky(filename, bucket, key): + if key == events_key and not failures: + failures.append(key) + raise ClientError({"Error": {"Code": "SlowDown"}}, "PutObject") + return original(filename, bucket, key) + + s3.upload_file = flaky + + def runner(item, workdir, out, resume, stop): + out.mkdir(parents=True) + (out / "events.jsonl").write_text(json.dumps({"event": "trial_start"}) + "\n") + deadline = time.monotonic() + 5 + while events_key not in s3.objects and time.monotonic() < deadline: + time.sleep(0.01) + # The trial is still running; its start was uploaded after one failure. + assert failures == [events_key] + assert events_key in s3.objects + return 0 + + assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries, + renew_seconds=0.01) == ["a"] # fmt: skip From b5795b5daeeff377ff47cc59c83d5019df9eea6f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 17:55:24 +0900 Subject: [PATCH 13/16] Keep cleanup errors when the claim is lost after the run An attempt whose claim was lost during cleanup or publication returned without the errors of quiescing its failed run, so the worker went on to further trials while a query could still run. Co-Authored-By: Claude Opus 5.5 --- benchmarks/pyathena_bench/fleet.py | 6 +++--- benchmarks/tests/test_fleet.py | 26 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index e8b1f2d1d..045ee1b3b 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -1053,7 +1053,7 @@ def upload_attempt(timeout: float = -1) -> None: or not queue.renew(claim, lease_seconds) ): # The claim was taken over or expired; its new owner records the job. - return Outcome("lost") + return Outcome("lost", tuple(errors)) if report_error is None: # The job's summary goes below this attempt, which the done marker names, # so a former owner's late upload cannot replace the recorded summary. @@ -1063,7 +1063,7 @@ def upload_attempt(timeout: float = -1) -> None: ) if time.time() >= claim.body["expires"]: # Publication outlasted the renewed claim, which another worker may own now. - return Outcome("lost") + return Outcome("lost", tuple(errors)) # The first done marker of a job is never replaced. recorded = queue.put_lease( queue.key("done", job["id"]), @@ -1081,7 +1081,7 @@ def upload_attempt(timeout: float = -1) -> None: IfNoneMatch="*", ) if recorded is None: - return Outcome("lost") + return Outcome("lost", tuple(errors)) return Outcome("done", tuple(errors)) diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 60ac7875c..716b97c23 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -1287,3 +1287,29 @@ def runner(item, workdir, out, resume, stop): assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries, renew_seconds=0.01) == ["a"] # fmt: skip + + +def test_cleanup_errors_survive_a_claim_lost_during_cleanup(tmp_path): + s3, queue = queue_with([job("a"), job("b")], tmp_path) + ran = [] + + def settle(output): + # The failed job's queries do not stop, and the claim is lost meanwhile. + s3.put_object( + Key=queue.key("claims", "a"), + Body=json.dumps( + {"host": "other", "job": "a", "attempt": 2, "state": "running", + "renewed": time.time(), "expires": time.time() + 600} + ).encode(), + ) # fmt: skip + return ["q1: still RUNNING"] + + def runner(item, workdir, out, resume, stop): + ran.append(item["id"]) + return 1 + + with pytest.raises(RuntimeError, match="q1: still RUNNING"): + work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=settle) + # No further trials start while the query may still run. + assert ran == ["a"] + assert not queue.names("done") From 7bbf1e14836abe9e556a7418608a582d1cc95cdd Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 18:18:55 +0900 Subject: [PATCH 14/16] Simplify the fleet's resume, leases, and worker start /simplify on #1145: - The run writes each trial's start and record to the attempt's checkpoint/ in S3 between trials (run --checkpoint), instead of the worker re-uploading growing JSONL files during trials; the worker uploads the attempt's files once when it ends, and a run under a fleet job writes no summary of its own. - Hosts get the results location from user data (worker --results, results-location.txt), so a worker needs no stack outputs; the wait for them and the fallback for stacks without the Results output are removed. - Workers exit when nothing is pending; an expired claim is taken over by any running worker. - Leases always have attempt, state, renewed, and expires; one end_lease() for released and interrupted; the renew interval is a tenth of the lease period; JobWatch sets one claim_lost event. - report.trial_key() defines a trial across attempts; superseded failures are noted instead of counted in a separate column. - aws.run_locations() gives the run's S3 prefixes for cleanup and trial outputs; METADATA_ERRORS, safe_settle(), download_file, Queue.put for JSON writes, and slots written only when enough are free. - The rehearsal notice file lives in the worker's directory, removed by the worker when it starts. - Tests grouped by the object under test with shared helpers. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 22 +- benchmarks/cloudformation/benchmark.yaml | 78 +- benchmarks/pyathena_bench/__main__.py | 81 +-- benchmarks/pyathena_bench/aws.py | 38 +- benchmarks/pyathena_bench/fleet.py | 518 ++++++------- benchmarks/pyathena_bench/report.py | 93 ++- benchmarks/pyathena_bench/runner.py | 85 ++- benchmarks/tests/test_aws.py | 40 +- benchmarks/tests/test_cloudformation.py | 9 +- benchmarks/tests/test_fleet.py | 886 +++++++++-------------- benchmarks/tests/test_report.py | 123 ++-- benchmarks/tests/test_runner.py | 155 ++-- 12 files changed, 977 insertions(+), 1151 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index ffcd12aa6..29e8ec579 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -109,7 +109,6 @@ The supplied template assumes ordinary IAM access and SSE-S3 source objects; Lak The fleet queues, reports, and Athena trial outputs go to the stack's results location, the `Results` stack output. With `ResultsBucket` set to an existing bucket, the results location is `s3://///`, and the instance role can read and write objects only below it, while it can list the whole bucket, because S3 answers a read of a missing key with 403 instead of 404 without that permission; with `ResultsBucket` empty, the default, it is the scratch bucket. -A stack created from an earlier template has no `Results` output and keeps everything in the scratch bucket. Input snapshots, initialization fixtures, and scratch tables stay in the scratch bucket and database either way. This repository's CI test bucket, which `AWS_ATHENA_S3_STAGING_DIR` names in CI, expires every object one day after it is written, at the following midnight UTC (`cloudformation/github_actions_oidc.yaml`). As the results location, it removes leftovers by itself, but reports must be recovered within that day. @@ -286,6 +285,7 @@ uv run --no-sync python -m pyathena_bench run \ There is no automatic continuation of a partially recorded trial; rerun the selected comparison with the same manifest. `run --resume` names the output directories of earlier attempts and skips the measured trials that they recorded as successful, after running its own warmups, which a fleet worker uses to continue an interrupted job. +A fleet worker also passes `--checkpoint`, with which the run writes each trial's start and record to S3 between trials and writes no summary of its own; the job's summary merges its attempts. Large API row cases may exceed the default one-hour timeout: 10 million rows require at least 10,000 pages with arraysize 1000, or 100,000 pages with arraysize 100. Choose the timeout using a smaller-scale pilot before the full run. Do not run two orchestrators concurrently on the same dedicated host. @@ -312,7 +312,7 @@ With `--split-pages`, a job whose single trial needs at least that many pages be `queue --jobs` accepts several job files, rejects duplicate job IDs, and orders the combined jobs as `jobs` orders one selection. Deploy a new stack with the fleet size under its own name; changing `FleetSize` on an existing stack does not wait for the added hosts to finish bootstrap. -`QueueName` names the queue that every host consumes: after bootstrap, each host starts a worker that waits for the stack's outputs and then until a queue of that name is published, including the hosts that the group launches later to replace interrupted Spot Instances. +`QueueName` names the queue that every host consumes: after bootstrap, each host starts a worker with the stack's results location from its user data, also written to `results-location.txt` next to `stack-id.txt`, and the worker waits until a queue of that name is published; this includes the hosts that the group launches later to replace interrupted Spot Instances. `ApiSlots` sets their `--api-slots`. `ResultsBucket` selects the results location; the commands below use this repository's CI test bucket, whose objects expire after a day. Then open a session on one of its hosts: @@ -416,7 +416,7 @@ For a queue with another name, start a worker on every host with SSM; a second w ```bash uv run --env-file ../.env --locked aws ssm send-command --targets "Key=tag:aws:autoscaling:groupName,Values=$BENCHMARK_GROUP" \ --document-name AWS-RunShellScript \ - --parameters 'commands=["sudo -iu ec2-user bash -lc \"cd /opt/pyathena/benchmarks && mkdir -p results && (nohup uv run --no-sync python -m pyathena_bench worker --stack $(cat /opt/pyathena/benchmarks/stack-id.txt) --name large-2 --terminate-host > results/worker-large-2.log 2>&1 &)\""]' + --parameters 'commands=["sudo -iu ec2-user bash -lc \"cd /opt/pyathena/benchmarks && mkdir -p results && (nohup uv run --no-sync python -m pyathena_bench worker --results $(cat /opt/pyathena/benchmarks/results-location.txt) --name large-2 --terminate-host > results/worker-large-2.log 2>&1 &)\""]' ``` A worker claims a job by creating `claims/` with an S3 conditional write, so exactly one host runs each job. @@ -426,8 +426,8 @@ A worker releases a slot by overwriting its lease as released, never by deleting A worker whose claim was taken over, or whose renewals failed until the claim expired, stops the attempt, quiesces its queries as after a failed job, and records nothing for the job; the new owner continues it, and the first `done` marker of a job is never replaced. If that worker cannot confirm that the queries stopped, it raises an error and keeps its host, because the job may have no other worker. A lost API slot does not stop the job; the `done` marker lists it under `lost_slots`. -Each attempt of a job writes to `results//attempt-/`, and the worker uploads its trial records within seconds of each trial's start and end, so a lost host loses only its running trial. -A trial whose start was not uploaded before its host died is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. +Each attempt of a job writes to `results//attempt-/`; between trials, its run writes each trial's start and record below `checkpoint/` there, so a lost host loses only its running trial, and the worker uploads the attempt's files when it ends. +A trial whose start checkpoint could not be written, which the run logs, is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. After a job, the worker uploads the attempt's output directory and its log to `logs/`, renews its claim conditionally to confirm that it still owns the job, writes `job-summary.csv` and `job-summary.md` from every attempt into its attempt directory, and then creates `done/` with the exit code, attempt, host, times, and the summary's attempt directory under `summary`. A failed trial stops only its own job; `status` lists jobs with a non-zero exit code, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. After a failed job, the worker stops the job's reported queries and every queued or running query in the workgroup that writes under the job's trial output prefixes, which also covers queries whose IDs were never reported before a trial process died. @@ -438,18 +438,18 @@ Each worker checks instance metadata for a Spot interruption notice (`spot/insta On a notice, the worker interrupts the running trial with SIGINT, which cancels the trial's reported queries, uploads the attempt, releases its claim as `interrupted` and its slots, and exits without terminating the host. A separate thread watches for the notice, so slow S3 requests do not delay it. The group replaces the reclaimed instance, and the new host's worker takes over the interrupted claim as the next attempt: it downloads the earlier attempts, quiesces their trials as after a failed job, and runs its own warmups and then only the measured repetitions that they did not record as successful, through `run --resume`. -The job's summary merges the attempts and counts a failed or incomplete trial whose repetition succeeded in a later attempt as `superseded_failures` instead of `failed_trials`, with its reason in the notes, so it reports the trials of an uninterrupted run without hiding failures. +The job's summary merges the attempts and leaves a failed or incomplete trial whose repetition succeeded in a later attempt out of `failed_trials`, noting its reason as `Superseded: ...`, so it reports the trials of an uninterrupted run without hiding failures. `report results/recovered/fleet//results/` produces the same merged summary from the recovered attempts. Rebalance recommendations are not acted on; an interruption does not always follow them, and the notice leaves enough time to stop a trial. -Workers exit when no job is left to claim and every other unfinished job's claim has been renewed since they found nothing to claim and is still fresh, which takes up to a renewal interval. -A claim whose worker died is never renewed, so workers keep waiting until it expires and one of them takes it over; `status` marks a claim that has missed renewals as `stale`. +Workers exit when no job is left to claim. +A claim whose worker died expires and is taken over by a worker that is still running; `status` marks it as expired. With `--terminate-host`, a worker first suspends the group's `AZRebalance` process, which would otherwise let the group launch a host in one zone and terminate a busy host in another when Spot capacity or retired hosts leave the zones uneven; the process stays suspended until the stack is deleted. A worker with `--terminate-host` that returns normally records its jobs, instance type, Availability Zone, and purchase option under `hosts/.json` and terminates its host through the group with a lower desired capacity, so no replacement starts and finished hosts stop costing money while others still run. This includes a worker that stopped after unconfirmed queries, whose errors are already in the `done` marker. A worker that raises an error, for example for a job wider than `--api-slots`, keeps its host. To keep a host after its worker finishes, tag it with `KeepHost=true` (`aws ec2 create-tags --resources --tags Key=KeepHost,Value=true`); the worker reads the tag from instance metadata. -A job whose last worker exited without finishing it, for example after a host error, stays claimed until its lease expires; then a worker started for the same queue takes it over. +A job whose worker died after the other workers exited, or after a host error, stays claimed until its lease expires; then a worker started for the same queue, as below, takes it over. To retry failed jobs, wait until all workers have exited, run `cleanup --trials-only --execute` for the manifest, and publish the selected jobs under a new queue name with the same manifest. ### Pilot with a simulated interruption @@ -475,7 +475,7 @@ uv run --no-sync python -m pyathena_bench queue \ --stack "$(cat stack-id.txt)" --name pilot-1 --jobs results/jobs-pilot.json --manifest results/input-pilot.json ``` -When `status` shows the job claimed and `results//attempt-1/trials.jsonl` has a trial, create the file `/opt/pyathena/benchmarks/interruption-notice`, which a worker treats as an interruption notice and removes when it starts. +When `status` shows the job claimed and `results//attempt-1/checkpoint/` has a trial record, create the file `interruption-notice` in the worker's directory, `/opt/pyathena/benchmarks/results/fleet/pilot-1/`, which the worker treats as an interruption notice and removes when it starts. As after a real notice, the worker exits and leaves its host running, so terminate the host as below. Check that the claim is `interrupted`, then terminate the host without lowering the desired capacity, as EC2 does when it reclaims a Spot Instance: @@ -486,7 +486,7 @@ BENCHMARK_INSTANCE=$(uv run --env-file ../.env --locked aws autoscaling describe --auto-scaling-group-names "$BENCHMARK_GROUP" --query 'AutoScalingGroups[0].Instances[0].InstanceId' --output text) uv run --env-file ../.env --locked aws ssm send-command --instance-ids "$BENCHMARK_INSTANCE" \ --document-name AWS-RunShellScript \ - --parameters 'commands=["sudo -u ec2-user touch /opt/pyathena/benchmarks/interruption-notice"]' + --parameters 'commands=["sudo -u ec2-user touch /opt/pyathena/benchmarks/results/fleet/pilot-1/interruption-notice"]' uv run --env-file ../.env --locked python -m pyathena_bench status --stack "$BENCHMARK_STACK" --name pilot-1 uv run --env-file ../.env --locked aws autoscaling terminate-instance-in-auto-scaling-group \ --instance-id "$BENCHMARK_INSTANCE" --no-should-decrement-desired-capacity diff --git a/benchmarks/cloudformation/benchmark.yaml b/benchmarks/cloudformation/benchmark.yaml index 3a9b8e75a..0afcdb315 100644 --- a/benchmarks/cloudformation/benchmark.yaml +++ b/benchmarks/cloudformation/benchmark.yaml @@ -349,42 +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 - 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 --stack ${AWS::StackId} --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.' + 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] diff --git a/benchmarks/pyathena_bench/__main__.py b/benchmarks/pyathena_bench/__main__.py index 0ff88e895..2556d4db8 100644 --- a/benchmarks/pyathena_bench/__main__.py +++ b/benchmarks/pyathena_bench/__main__.py @@ -13,18 +13,16 @@ import json import logging import sys -import time from dataclasses import asdict, replace from pathlib import Path -from typing import Any from pyathena_bench.aws import ( - IncompleteStackError, cleanup, client, preflight, prepare, prepare_fixtures, + run_locations, session, split_location, stack_resources, @@ -34,12 +32,10 @@ from pyathena_bench.config import Settings, read_json from pyathena_bench.fleet import ( API_FAMILIES, - SIMULATED_NOTICE, Filters, Queue, SpotInterruptionError, expand_jobs, - interruption_notice, merge_jobs, retire_host, status, @@ -48,7 +44,7 @@ work, ) from pyathena_bench.report import report -from pyathena_bench.runner import completed_trials, run +from pyathena_bench.runner import completed_trials, run, s3_checkpoint def parser() -> argparse.ArgumentParser: @@ -107,6 +103,11 @@ def parser() -> argparse.ArgumentParser: 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" ) @@ -117,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", @@ -135,7 +140,7 @@ def parser() -> argparse.ArgumentParser: consume.add_argument( "--wait", action="store_true", - help="Wait for the stack's outputs and for the queue to be published instead of failing", + help="Wait for the queue to be published instead of failing", ) consume.add_argument( "--lease-seconds", @@ -163,31 +168,6 @@ def parser() -> argparse.ArgumentParser: return root -def wait_for_stack( - session_: Any, stack: str, wait: bool, poll_seconds: float = 30.0 -) -> dict[str, str]: - """Read the stack's resources, waiting for its outputs if requested. - - A host's worker starts during stack creation, before the outputs exist. - - Args: - session_: Boto3 session. - stack: Stack name or ID. - wait: Whether to wait while the outputs are incomplete. - poll_seconds: Seconds between attempts. - - Returns: - The stack resources. - """ - while True: - try: - return stack_resources(session_, stack) - except IncompleteStackError: - if not wait: - raise - time.sleep(poll_seconds) - - def main(argv: list[str] | None = None) -> int: args = parser().parse_args(argv) if args.command == "report": @@ -217,9 +197,10 @@ def main(argv: list[str] | None = None) -> int: else f"b_{manifest['run_id']}_*" ), "s3_prefixes": [ - f"{root}runs/{manifest['run_id']}/" - + ("trials/" if args.trials_only else "") - for root in sorted({f"s3://{resources['Bucket']}/", resources["Results"]}) + 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, }, @@ -231,8 +212,8 @@ 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) - resources = wait_for_stack(session_, args.stack, getattr(args, "wait", False)) - bucket, base = split_location(resources["Results"]) + 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)) @@ -243,8 +224,6 @@ def main(argv: list[str] | None = None) -> int: if args.lease_seconds < 300: raise ValueError("--lease-seconds must be at least 300") autoscaling = client(session_, "autoscaling") - # A rehearsed notice applies to one worker; a restarted worker starts clean. - SIMULATED_NOTICE.unlink(missing_ok=True) if args.terminate_host: suspend_rebalance(autoscaling) try: @@ -253,8 +232,6 @@ def main(argv: list[str] | None = None) -> int: args.workdir or Path("results/fleet") / args.name, args.api_slots, lease_seconds=args.lease_seconds, - renew_seconds=args.lease_seconds / 10, - notice=interruption_notice, wait=args.wait, ) except SpotInterruptionError as interrupted: @@ -325,13 +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) + 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) + settings, + manifest, + cases, + scales, + args.out, + completed_trials(args.resume), + checkpoint, ) - if args.resume and not (args.out / "trials.jsonl").exists(): - # Earlier attempts recorded every trial; the job's report merges them. - return 0 if success else 1 - report(args.out) + if not checkpoint: + # A fleet job's summary merges its attempts instead. + report(args.out) return 0 if success else 1 return 0 diff --git a/benchmarks/pyathena_bench/aws.py b/benchmarks/pyathena_bench/aws.py index af33eab2e..9f54cabd0 100644 --- a/benchmarks/pyathena_bench/aws.py +++ b/benchmarks/pyathena_bench/aws.py @@ -37,10 +37,6 @@ def client(session_: Any, service: str) -> Any: ) -class IncompleteStackError(ValueError): - """The stack's outputs are incomplete while an operation on it is in progress.""" - - def split_location(location: str) -> tuple[str, str]: """Split an ``s3://bucket/prefix/`` location. @@ -59,6 +55,25 @@ def split_location(location: str) -> tuple[str, str]: return bucket, prefix +def run_locations(resources: dict[str, str], run_id: str) -> dict[str, str]: + """Return the S3 prefixes of a prepared run. + + Args: + resources: Stack resources from :func:`stack_resources`. + run_id: Run ID of the manifest. + + Returns: + ``scratch``, below which inputs and fixtures are, and ``results``, below + which trial outputs are (``trials//``); both ``s3://`` prefixes + ending with a slash, the same when the results location is the scratch + bucket. + """ + return { + "scratch": f"s3://{resources['Bucket']}/runs/{run_id}/", + "results": f"{resources['Results']}runs/{run_id}/", + } + + def stack_resources(session_: Any, stack: str) -> dict[str, str]: data = client(session_, "cloudformation").describe_stacks(StackName=stack)["Stacks"][0] tags = {item["Key"]: item["Value"] for item in data.get("Tags", [])} @@ -74,13 +89,10 @@ def stack_resources(session_: Any, stack: str) -> dict[str, str]: "SourceTable", "WorkGroup", "AutoScalingGroup", + "Results", } if not required <= result.keys(): - if data.get("StackStatus", "").endswith("_IN_PROGRESS"): - raise IncompleteStackError("Stack outputs are incomplete") raise ValueError("Stack outputs are incomplete") - # A stack from a template before the Results output keeps everything in the scratch bucket. - result.setdefault("Results", f"s3://{result['Bucket']}/") if ( result["SourceBucket"] == result["Bucket"] or result["ScratchDatabase"] == result["SourceDatabase"] @@ -366,14 +378,8 @@ def cleanup( resources = validate_manifest(settings, manifest) session_ = session(settings) run_id = manifest["run_id"] - results_bucket, results_prefix = split_location(resources["Results"]) - locations = sorted( - { - (resources["Bucket"], f"runs/{run_id}/"), - (results_bucket, f"{results_prefix}runs/{run_id}/"), - } - ) - roots = tuple(f"s3://{bucket}/{root}" for bucket, root in locations) + roots = tuple(sorted(set(run_locations(resources, run_id).values()))) + locations = [split_location(root) for root in roots] athena = client(session_, "athena") # Discover queries whose worker died before returning its query ID. # The query prefix is checked against live Athena output metadata. diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index e3963e72a..5ea522991 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -20,7 +20,6 @@ import sys import threading import time -import urllib.error import urllib.request from collections.abc import Callable, Iterable from dataclasses import dataclass @@ -34,7 +33,7 @@ from pyathena_bench.aws import active_queries, client, session from pyathena_bench.cases import Case, matrix from pyathena_bench.config import Settings -from pyathena_bench.report import report +from pyathena_bench.report import read_trial_starts, report # Row cursors page through GetQueryResults, which has an account-wide request-rate quota; # so does read_sql_query() through the REST driver. @@ -42,10 +41,13 @@ QUEUE_ROOT = "fleet" # An instance tag with the value "true" keeps a host running after its worker exits. KEEP_HOST_TAG = "KeepHost" -# A file that a worker treats like a Spot interruption notice, to rehearse interruptions. -SIMULATED_NOTICE = Path("interruption-notice") +# A file in a worker's directory that it treats like a Spot interruption notice, to +# rehearse interruptions. +SIMULATED_NOTICE = "interruption-notice" METADATA_URL = "http://169.254.169.254/latest" _logger = logging.getLogger(__name__) +# Errors of an instance metadata request, including a missing path (HTTP 404). +METADATA_ERRORS = (OSError, http.client.HTTPException) # S3 errors of a conditional write whose condition failed. PRECONDITION_ERRORS = {"PreconditionFailed", "ConditionalRequestConflict"} @@ -252,8 +254,7 @@ class Lease: ``interrupted``, or ``released``), and the times of the last write (``renewed``) and of the expiry, as Unix times. A worker renews its leases while it runs a job; another worker may take over a lease that has expired - or was released. Objects without an expiry, written by earlier versions, - never expire. + or was ended. """ key: str @@ -267,31 +268,9 @@ def takeable(self, at: float) -> bool: at: Current Unix time. Returns: - True if the lease was interrupted or expired before ``at``. + True if the lease was ended or expired before ``at``. """ - expires = self.body.get("expires") - return self.body.get("state") in {"interrupted", "released"} or ( - expires is not None and expires < at - ) - - def stale(self, at: float) -> bool: - """Return whether a running lease missed renewals but has not expired yet. - - A live worker renews a lease ten times per lease period, so a lease - that has gone a fifth of its period without renewal belongs to a - worker that has probably died. - - Args: - at: Current Unix time. - - Returns: - True if the lease is running, not takeable, and overdue for renewal. - """ - renewed = self.body.get("renewed") - expires = self.body.get("expires") - if renewed is None or expires is None or self.takeable(at): - return False - return at - renewed > (expires - renewed) / 5 + return self.body["state"] != "running" or self.body["expires"] < at class Queue: @@ -317,7 +296,7 @@ def key(self, *parts: str) -> str: """ return self.prefix + "/".join(parts) - def put(self, key: str, body: dict[str, Any], **condition: str) -> str | None: + def put(self, key: str, body: Any, **condition: str) -> str | None: """Write a JSON object, optionally only if a condition holds. Args: @@ -397,7 +376,7 @@ def put_lease(self, key: str, body: dict[str, Any], **condition: str) -> Lease | return None def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Lease | None: - """Take a lease on an object that is absent, expired, or interrupted. + """Take a lease on an object that is absent, expired, or ended. Creation and takeover are conditional writes, so exactly one worker obtains each lease. A takeover increments the attempt number. @@ -410,7 +389,22 @@ def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Leas Returns: The lease, or None if another worker holds it. """ - current = self.read_lease(key) + return self.take(key, self.read_lease(key), owner, lease_seconds) + + def take( + self, key: str, current: Lease | None, owner: dict[str, Any], lease_seconds: float + ) -> Lease | None: + """Take a lease given the object that :meth:`read_lease` returned for it. + + Args: + key: Object key of the claim or slot. + current: The object as read, or None if it did not exist. + owner: Description of the worker and job. + lease_seconds: Seconds until the lease expires unless renewed. + + Returns: + The lease, or None if another worker holds it or wrote it since. + """ at = time.time() body = { **owner, @@ -420,11 +414,10 @@ def acquire(self, key: str, owner: dict[str, Any], lease_seconds: float) -> Leas "expires": at + lease_seconds, } if current is None: - # Exactly one worker creates the object; another one leaves it. return self.put_lease(key, body, IfNoneMatch="*") if not current.takeable(at): return None - body["attempt"] = current.body.get("attempt", 1) + 1 + body["attempt"] = current.body["attempt"] + 1 return self.put_lease(key, body, IfMatch=current.etag) def renew(self, lease: Lease, lease_seconds: float, **changes: Any) -> bool: @@ -446,17 +439,6 @@ def renew(self, lease: Lease, lease_seconds: float, **changes: Any) -> bool: lease.body, lease.etag = written.body, written.etag return True - def release_lease(self, lease: Lease) -> bool: - """Release a lease that this worker still holds; see :meth:`end_lease`. - - Args: - lease: Lease to release. - - Returns: - False if another worker took over the lease. - """ - return self.end_lease(lease, "released") - def end_lease(self, lease: Lease, state: str) -> bool: """End a lease that this worker still holds, so another worker can take it. @@ -514,7 +496,8 @@ def acquire_slots( ) -> list[Lease] | None: """Take API slots that bound simultaneous GetQueryResults paging. - A slot whose lease expired, because its worker died, is taken over. + Every slot is read first, and slots are written only when enough are + free; a slot whose lease expired, because its worker died, is free. Args: slots: Number of API slots in the fleet. @@ -532,20 +515,29 @@ def acquire_slots( needed = max(count, 1) if needed > slots: raise ValueError(f"A job needs {needed} API slots, but only {slots} exist") + at = time.time() + free = [] + for index in range(slots): + key = self.key("slots", str(index)) + current = self.read_lease(key) + if current is None or current.takeable(at): + free.append((key, current)) + if len(free) < needed: + return None taken: list[Lease] = [] try: - for index in range(slots): - lease = self.acquire(self.key("slots", str(index)), owner, lease_seconds) + for key, current in free: + lease = self.take(key, current, owner, lease_seconds) if lease is not None: taken.append(lease) if len(taken) == needed: return taken except BaseException: for lease in taken: - self.release_lease(lease) + self.end_lease(lease, "released") raise for lease in taken: - self.release_lease(lease) + self.end_lease(lease, "released") return None def exists(self, *parts: str) -> bool: @@ -601,9 +593,7 @@ def download_tree(self, directory: Path, *parts: str) -> list[str]: for name in names: path = directory / name path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes( - self.s3.get_object(Bucket=self.bucket, Key=source + name)["Body"].read() - ) + self.s3.download_file(self.bucket, source + name, str(path)) return names @@ -631,11 +621,7 @@ def submit(queue: Queue, jobs: list[dict[str, Any]], config: Path, manifest: Pat raise ValueError("Queue already exists; choose a new name") for name, path in (("config.toml", config), ("manifest.json", manifest)): queue.s3.upload_file(str(path), queue.bucket, queue.key(name)) - queue.s3.put_object( - Bucket=queue.bucket, - Key=queue.key("jobs.json"), - Body=json.dumps(jobs, indent=1).encode(), - ) + queue.put(queue.key("jobs.json"), jobs) # Workers start only after this marker. if not queue.put_if_absent(queue.key("queue.json"), {"created": now(), "jobs": len(jobs)}): raise ValueError("Queue already exists; choose a new name") @@ -651,6 +637,7 @@ def run_job( output: Path, resume: list[Path], stop: threading.Event, + checkpoint: str, ) -> int: """Run one job attempt as a separate orchestrator process. @@ -667,6 +654,7 @@ def run_job( resume: Output directories of earlier attempts whose completed trials are not repeated. stop: Event that interrupts the attempt. + checkpoint: The attempt's ``s3://`` prefix for per-trial checkpoints. Returns: Process exit code. @@ -674,7 +662,8 @@ def run_job( (workdir / "logs").mkdir(parents=True, exist_ok=True) command = [ sys.executable, "-m", "pyathena_bench", "--config", str(workdir / "config.toml"), - "run", "--manifest", str(workdir / "manifest.json"), "--out", str(output), *job["args"], + "run", "--manifest", str(workdir / "manifest.json"), "--out", str(output), + "--checkpoint", checkpoint, *job["args"], ] # fmt: skip if resume: command += ["--resume", *map(str, resume)] @@ -701,7 +690,8 @@ def trial_evidence(output: Path) -> tuple[set[str], tuple[str, ...]]: """Collect the query IDs and output prefixes that a job's trials used. Args: - output: Job output directory with ``events.jsonl``. + output: Job output directory with ``events.jsonl`` or the checkpoint + of a fleet attempt. Returns: Reported query IDs and the S3 output prefixes of started trials; both @@ -709,18 +699,15 @@ def trial_evidence(output: Path) -> tuple[set[str], tuple[str, ...]]: """ events = output / "events.jsonl" ids: set[str] = set() - prefixes: set[str] = set() - if not events.exists(): - return ids, () - for line in events.read_text().splitlines(): - try: - event = json.loads(line) - except json.JSONDecodeError: - continue - if event.get("event") == "query" and event.get("query_id"): - ids.add(event["query_id"]) - elif event.get("event") == "trial_start" and event.get("output"): - prefixes.add(event["output"]) + if events.exists(): + for line in events.read_text().splitlines(): + try: + event = json.loads(line) + except json.JSONDecodeError: + continue + if event.get("event") == "query" and event.get("query_id"): + ids.add(event["query_id"]) + prefixes = {e["output"] for e in read_trial_starts(output) if e.get("output")} return ids, tuple(sorted(prefixes)) @@ -778,43 +765,40 @@ def __init__(self, ran: list[str]) -> None: class JobWatch: - """Renew a running job's leases, upload its progress, and watch for interruption. + """Renew a running job's leases and watch for an interruption notice. One thread checks every ``poll_seconds`` for an interruption notice and - for the claim's expiry, without S3 requests. Another calls ``progress`` - every ``poll_seconds`` and renews the leases every ``renew_seconds``. - ``stop_requested`` is set on a notice (with ``interrupted``) and when the - claim, the first lease, is lost: another worker took it over, or - renewals failed or stalled until it expired. A lost lease is listed in - ``lost`` and no longer renewed; losing an API slot does not stop the job. - ``upload_lock`` serializes uploads of the attempt's files. + for the claim's expiry, without S3 requests; another renews the claim and + the API slots ten times per lease period. ``interrupted`` is set on a + notice, and ``claim_lost`` when another worker took over the claim or + renewals failed or stalled until it expired; ``stop`` is set with either + and stops the runner. A lost slot is listed in ``lost_slots`` and does not + stop the job. """ def __init__( self, queue: Queue, - leases: list[Lease], + claim: Lease, + slots: list[Lease], lease_seconds: float, - renew_seconds: float, - progress: Callable[[], None], notice: Callable[[], bool], poll_seconds: float = 5.0, ) -> None: self.queue = queue - self.claim = leases[0] - self.leases = list(leases) + self.claim = claim + self.slots = list(slots) self.lease_seconds = lease_seconds - self.renew_seconds = renew_seconds - self.progress = progress self.notice = notice - self.poll_seconds = min(poll_seconds, renew_seconds) + self.renew_seconds = lease_seconds / 10 + self.poll_seconds = min(poll_seconds, self.renew_seconds) self.interrupted = threading.Event() - self.stop_requested = threading.Event() - self.lost: list[str] = [] - self.upload_lock = threading.Lock() + self.claim_lost = threading.Event() + self.stop_event = threading.Event() + self.lost_slots: list[str] = [] self._stopped = threading.Event() self._threads = [ - threading.Thread(target=self._watch_notice, daemon=True), + threading.Thread(target=self._watch, daemon=True), threading.Thread(target=self._renew, daemon=True), ] @@ -829,9 +813,9 @@ def __exit__( exc: BaseException | None, traceback: TracebackType | None, ) -> None: - # After a notice or a lost lease, a renewal stuck in S3 must not delay the - # release of the slots; a late renewal fails against the changed ETags. - self.stop(10 if self.stop_requested.is_set() else None) + # After a stop, a renewal stuck in S3 must not delay the release of the + # slots; a late renewal fails against the changed ETags. + self.stop(10 if self.stop_event.is_set() else None) def stop(self, timeout: float | None = None) -> None: """Stop the threads and wait for them. @@ -845,49 +829,38 @@ def stop(self, timeout: float | None = None) -> None: if thread.is_alive(): thread.join(timeout) - def _watch_notice(self) -> None: + def _watch(self) -> None: while not self._stopped.wait(self.poll_seconds): if check_notice(self.notice): self.interrupted.set() - self.stop_requested.set() + self.stop_event.set() return if time.time() >= self.claim.body["expires"]: - # Renewals failed or stalled; another worker may own the job now. - self.stop_requested.set() + self._lose_claim() return def _renew(self) -> None: - renew_at = time.monotonic() + self.renew_seconds - while not self._stopped.wait(self.poll_seconds): - # Renewals come first, so that a slow upload does not delay them. - if time.monotonic() >= renew_at: - renew_at = time.monotonic() + self.renew_seconds - self._renew_leases() - try: - with self.upload_lock: - self.progress() - except Exception: - # The next poll uploads again. - _logger.warning("Could not upload the attempt's progress", exc_info=True) - - def _renew_leases(self) -> None: - for lease in list(self.leases): - try: - if self.queue.renew(lease, self.lease_seconds): - continue - except Exception: - if time.time() < lease.body["expires"]: - # The lease has time left for a retry. - continue - self.leases.remove(lease) - self.lost.append(lease.key) - if lease is self.claim: - self.stop_requested.set() + while not self._stopped.wait(self.renew_seconds): + if not self.claim_lost.is_set() and not self._renew_lease(self.claim): + self._lose_claim() + for slot in list(self.slots): + if not self._renew_lease(slot): + self.slots.remove(slot) + self.lost_slots.append(slot.key) + + def _renew_lease(self, lease: Lease) -> bool: + try: + return self.queue.renew(lease, self.lease_seconds) + except Exception: + # The lease has time left for a retry until it expires. + return time.time() < lease.body["expires"] + + def _lose_claim(self) -> None: + self.claim_lost.set() + self.stop_event.set() -def pending_jobs( - queue: Queue, jobs: list[dict[str, Any]] -) -> tuple[list[dict[str, Any]], dict[str, Lease]]: +def pending_jobs(queue: Queue, jobs: list[dict[str, Any]]) -> list[dict[str, Any]]: """Return the jobs that are neither done nor held by a live claim. Args: @@ -896,23 +869,21 @@ def pending_jobs( Returns: Jobs without a ``done`` marker whose claim is absent, expired, or - released, in queue order, and the claims of the other unfinished jobs. + ended, in queue order. """ done = queue.names("done") claimed = queue.names("claims") at = time.time() pending = [] - running: dict[str, Lease] = {} for job in jobs: if job["id"] in done: continue if job["id"] in claimed: lease = queue.read_lease(queue.key("claims", job["id"])) if lease is not None and not lease.takeable(at): - running[job["id"]] = lease continue pending.append(job) - return pending, running + return pending @dataclass(frozen=True) @@ -938,19 +909,19 @@ def run_attempt( runner: Callable[..., int], settle: Callable[[Path], list[str]], lease_seconds: float, - renew_seconds: float, notice: Callable[[], bool], ) -> Outcome: """Run one attempt of a claimed job and record its outcome. - The attempt writes to ``results//attempt-/``. A takeover (attempt - 2 or later) first downloads the earlier attempts, quiesces their trials, + The attempt writes to ``results//attempt-/``, and its orchestrator + writes per-trial checkpoints there between trials. A takeover (attempt 2 + or later) first downloads the earlier attempts, quiesces their trials, and passes them to the runner so that completed trials are not repeated. - Trial records are uploaded within seconds of each trial's start and end. An interruption - notice or a lost claim stops the runner; after a notice, the attempt is - uploaded and the claim released as interrupted. Otherwise, after a - conditional renewal confirms that the claim is still held, the job's - report merges every attempt into ``job-summary.csv`` and + An interruption notice or a lost claim stops the runner: after a notice, + the attempt is uploaded and the claim released as interrupted; after a + lost claim, the attempt's queries are quiesced and nothing is recorded. + Otherwise, after a conditional renewal confirms that the claim is still + held, the job's report merges every attempt into ``job-summary.csv`` and ``job-summary.md`` of this attempt, and the first ``done`` marker of the job records the exit code and names that attempt. @@ -963,7 +934,6 @@ def run_attempt( runner: Function that runs the attempt; see :func:`run_job`. settle: Function that quiesces an output directory and returns errors. lease_seconds: Seconds a lease lasts without renewal. - renew_seconds: Seconds between renewals. notice: Function that returns whether an interruption notice exists. Returns: @@ -974,41 +944,19 @@ def run_attempt( output = job_dir / f"attempt-{attempt}" remote = ("results", job["id"], output.name) log = workdir / "logs" / f"{job['id']}-{output.name}.log" - checked: dict[str, int] = {} - - def progress() -> None: - # Upload at trial boundaries: when a trial starts (its output prefix is - # needed to quiesce it) and when it is recorded (needed to resume). A - # boundary stays pending until its upload succeeds. - events, trials = output / "events.jsonl", output / "trials.jsonl" - if events.exists(): - with events.open("rb") as stream: - stream.seek(checked.get("events", 0)) - if b'"trial_start"' in stream.read(): - checked["pending"] = 1 - checked["events"] = stream.tell() - if trials.exists() and trials.stat().st_size != checked.get("trials"): - checked["trials"] = trials.stat().st_size - checked["pending"] = 1 - if checked.get("pending"): - for path in (events, trials): - if path.exists(): - queue.s3.upload_file(str(path), queue.bucket, queue.key(*remote, path.name)) - checked["pending"] = 0 - - with JobWatch(queue, [claim, *slots], lease_seconds, renew_seconds, progress, notice) as watch: - - def upload_attempt(timeout: float = -1) -> None: - # A checkpoint stuck in S3 holds the lock; it uploads the same files. - if not watch.upload_lock.acquire(timeout=timeout): - return - try: - queue.upload_tree(output, *remote) - if log.exists(): - queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) - finally: - watch.upload_lock.release() + def upload_attempt() -> None: + queue.upload_tree(output, *remote) + if log.exists(): + queue.s3.upload_file(str(log), queue.bucket, queue.key("logs", log.name)) + + def safe_settle(path: Path) -> list[str]: + try: + return settle(path) + except Exception as exc: + return [f"Could not quiesce {path.name}: {type(exc).__name__}: {exc}"] + + with JobWatch(queue, claim, slots, lease_seconds, notice) as watch: start = now() resume = [ job_dir / f"attempt-{earlier}" @@ -1017,83 +965,65 @@ def upload_attempt(timeout: float = -1) -> None: job_dir / f"attempt-{earlier}", "results", job["id"], f"attempt-{earlier}" ) ] - errors: list[str] = [] - for earlier_output in resume: - try: - errors += settle(earlier_output) - except Exception as exc: - errors.append( - f"Could not quiesce {earlier_output.name}: {type(exc).__name__}: {exc}" - ) + errors = [error for earlier in resume for error in safe_settle(earlier)] code = 1 - if not errors and not watch.stop_requested.is_set(): - code = runner(job, workdir, output, resume, watch.stop_requested) - if watch.stop_requested.is_set(): + if not errors and not watch.stop_event.is_set(): + code = runner( + job, + workdir, + output, + resume, + watch.stop_event, + f"s3://{queue.bucket}/{queue.key(*remote)}/", + ) + if watch.interrupted.is_set(): watch.stop(timeout=10) # Upload before releasing the claim, so that the next attempt resumes # from every recorded trial; both are small within a notice's two minutes. - upload_attempt(timeout=10) - if not watch.interrupted.is_set(): - # The claim was lost; another worker continues the job. This - # attempt's queries are stopped before this worker claims another. - try: - errors += settle(output) - except Exception as exc: - errors.append(f"Could not quiesce the job: {type(exc).__name__}: {exc}") - return Outcome("lost", tuple(errors)) + upload_attempt() queue.end_lease(claim, "interrupted") return Outcome("interrupted") - if not errors: - try: - errors = settle(output) if code else [] - except Exception as exc: - errors = [f"Could not quiesce the job: {type(exc).__name__}: {exc}"] + if code and not errors: + # A failed or stopped run can leave queries running. + errors = safe_settle(output) upload_attempt() - report_error = None - try: - report(job_dir) - except Exception as exc: - # A job without trials, or with unreadable records, still gets its done marker. - report_error = f"{type(exc).__name__}: {exc}" - watch.stop() - # Renewing the claim conditionally proves that this worker still owns the - # job and leaves at least --lease-seconds for the few seconds of publication. - if ( - claim.key in watch.lost - or (watch.stop_requested.is_set() and not watch.interrupted.is_set()) - or not queue.renew(claim, lease_seconds) - ): - # The claim was taken over or expired; its new owner records the job. - return Outcome("lost", tuple(errors)) - if report_error is None: - # The job's summary goes below this attempt, which the done marker names, - # so a former owner's late upload cannot replace the recorded summary. - for name in ("summary.csv", "summary.md"): - queue.s3.upload_file( - str(job_dir / name), queue.bucket, queue.key(*remote, f"job-{name}") - ) - if time.time() >= claim.body["expires"]: - # Publication outlasted the renewed claim, which another worker may own now. - return Outcome("lost", tuple(errors)) - # The first done marker of a job is never replaced. - recorded = queue.put_lease( - queue.key("done", job["id"]), - { - "exit_code": code, - "host": socket.gethostname(), - "attempt": attempt, - "start": start, - "end": now(), - "quiesce_errors": errors, - "report_error": report_error, - "summary": None if report_error else "/".join(remote[1:]), - "lost_slots": [key for key in watch.lost if key != claim.key], - }, - IfNoneMatch="*", - ) - if recorded is None: - return Outcome("lost", tuple(errors)) - return Outcome("done", tuple(errors)) + lost = Outcome("lost", tuple(errors)) + # The conditional renewal proves that this worker still owns the job, so a + # former owner never records it, and leaves at least --lease-seconds for + # the few seconds of publication. + if watch.claim_lost.is_set() or not queue.renew(claim, lease_seconds): + return lost + report_error = None + try: + report(job_dir) + # The job's summary goes below this attempt, which the done marker names, + # so a former owner's late upload cannot replace the recorded summary. + for name in ("summary.csv", "summary.md"): + queue.s3.upload_file( + str(job_dir / name), queue.bucket, queue.key(*remote, f"job-{name}") + ) + except Exception as exc: + # A job without trials, or with unreadable records, still gets its done marker. + report_error = f"{type(exc).__name__}: {exc}" + if time.time() >= claim.body["expires"]: + # Publication outlasted the renewed claim, which another worker may own now. + return lost + recorded = queue.put_lease( + queue.key("done", job["id"]), + { + "exit_code": code, + "host": socket.gethostname(), + "attempt": attempt, + "start": start, + "end": now(), + "quiesce_errors": errors, + "report_error": report_error, + "summary": None if report_error else "/".join(remote[1:]), + "lost_slots": watch.lost_slots, + }, + IfNoneMatch="*", + ) + return lost if recorded is None else Outcome("done", tuple(errors)) def work( @@ -1104,8 +1034,7 @@ def work( runner: Callable[..., int] = run_job, settle: Callable[[Path], list[str]] | None = None, lease_seconds: float = 600.0, - renew_seconds: float = 60.0, - notice: Callable[[], bool] = lambda: False, + notice: Callable[[], bool] | None = None, wait: bool = False, ) -> list[str]: """Claim and run queued jobs until no job is left to claim. @@ -1113,11 +1042,9 @@ def work( Each host runs one worker, so jobs on a host never overlap. A claim and each API slot are leases that the worker renews while the job runs; a job or slot whose lease expired, because its worker died, returns to the - fleet. Before exiting, the worker waits until every other unfinished - job's claim has been renewed, so a job whose worker died is taken over. - An API-heavy job starts only after the worker takes one shared slot per - simultaneous paging query; while it waits, no later API-heavy job takes - slots. After a failed job, the worker + fleet: a running worker takes it over. An API-heavy job starts only after + the worker takes one shared slot per simultaneous paging query; while it + waits, no later API-heavy job takes slots. After a failed job, the worker cancels the job's observed queries and waits until they stop before claiming another job. @@ -1131,9 +1058,12 @@ def work( settle: Function that quiesces a failed job, given its output directory, and returns errors; defaults to :func:`quiesce` with the queued configuration. - lease_seconds: Seconds a claim or slot lasts without renewal. - renew_seconds: Seconds between renewals. - notice: Function that returns whether a Spot interruption notice exists. + lease_seconds: Seconds a claim or slot lasts without renewal; the + worker renews its leases ten times per period. + notice: Function that returns whether a Spot interruption notice + exists; defaults to :func:`interruption_notice` with the + rehearsal file ``interruption-notice`` in ``workdir``, which the + worker removes when it starts. wait: Whether to wait for the queue to be published instead of failing. Returns: @@ -1152,6 +1082,14 @@ def work( raise ValueError("Queue is missing or incompletely published") time.sleep(poll_seconds) workdir.mkdir(parents=True, exist_ok=True) + if notice is None: + rehearsal = workdir / SIMULATED_NOTICE + # A rehearsed notice stops one worker; a restarted worker starts clean. + rehearsal.unlink(missing_ok=True) + + def notice() -> bool: + return interruption_notice(simulated=rehearsal) + for name in ("config.toml", "manifest.json", "jobs.json"): (workdir / name).write_bytes(queue.read(name)) jobs = json.loads((workdir / "jobs.json").read_text()) @@ -1166,31 +1104,12 @@ def settle(output: Path) -> list[str]: host = socket.gethostname() completed: list[str] = [] - # Renewal times of other workers' claims when this worker found nothing to claim. - waiting: dict[str, float | None] = {} while True: if check_notice(notice): raise SpotInterruptionError(completed) - pending, running = pending_jobs(queue, jobs) - if pending: - waiting = {} - else: - at = time.time() - for job_id, lease in running.items(): - waiting.setdefault(job_id, lease.body.get("renewed")) - # Exit once every other claim has been renewed since and is still - # fresh; a claim whose worker died stops being renewed and is taken - # over when it expires. - if all( - lease.body.get("renewed") is None - or ( - lease.body["renewed"] != waiting[job_id] - and not lease.stale(at) - and not lease.takeable(at) - ) - for job_id, lease in running.items() - ): - return completed + pending = pending_jobs(queue, jobs) + if not pending: + return completed started = False reserved = False for job in pending: @@ -1220,7 +1139,6 @@ def settle(output: Path) -> list[str]: runner, settle, lease_seconds, - renew_seconds, notice, ) if outcome.state == "interrupted": @@ -1240,7 +1158,7 @@ def settle(output: Path) -> list[str]: break finally: for slot in slots: - queue.release_lease(slot) + queue.end_lease(slot, "released") if not started: time.sleep(poll_seconds) @@ -1286,7 +1204,7 @@ def check_notice(notice: Callable[[], bool]) -> bool: def interruption_notice( - metadata: Callable[[str], str] | None = None, simulated: Path = SIMULATED_NOTICE + metadata: Callable[[str], str] | None = None, simulated: Path | None = None ) -> bool: """Return whether this host has a Spot interruption notice. @@ -1298,16 +1216,16 @@ def interruption_notice( Args: metadata: Function that reads an instance metadata path; defaults to :func:`instance_metadata`. - simulated: File that counts as a notice. + simulated: File that counts as a notice, or None. Returns: True if a notice exists; False also when instance metadata is unavailable. """ - if simulated.exists(): + if simulated is not None and simulated.exists(): return True try: (metadata or instance_metadata)("spot/instance-action") - except (urllib.error.URLError, OSError, http.client.HTTPException): + except METADATA_ERRORS: return False return True @@ -1325,7 +1243,7 @@ def keep_host(metadata: Callable[[str], str]) -> bool: """ try: return metadata(f"tags/instance/{KEEP_HOST_TAG}").strip().lower() == "true" - except (urllib.error.URLError, OSError): + except METADATA_ERRORS: return False @@ -1375,20 +1293,17 @@ def retire_host( """ host = metadata("instance-id") kept = keep_host(metadata) - queue.s3.put_object( - Bucket=queue.bucket, - Key=queue.key("hosts", f"{host}.json"), - Body=json.dumps( - { - "host": socket.gethostname(), - "instance_type": metadata("instance-type"), - "availability_zone": metadata("placement/availability-zone"), - "lifecycle": metadata("instance-life-cycle"), - "ran": ran, - "end": now(), - "kept": kept, - } - ).encode(), + queue.put( + queue.key("hosts", f"{host}.json"), + { + "host": socket.gethostname(), + "instance_type": metadata("instance-type"), + "availability_zone": metadata("placement/availability-zone"), + "lifecycle": metadata("instance-life-cycle"), + "ran": ran, + "end": now(), + "kept": kept, + }, ) if kept: return False @@ -1424,11 +1339,10 @@ def describe(folder: str, names: Iterable[str]) -> dict[str, Any]: leases = {} for name in sorted(names): lease = queue.read_lease(queue.key(folder, name)) - if lease is not None and lease.body.get("state") != "released": + if lease is not None and lease.body["state"] != "released": leases[name] = { **{k: lease.body.get(k) for k in ("host", "job", "attempt", "state")}, - "expired": lease.takeable(at) and lease.body.get("state") == "running", - "stale": lease.stale(at), + "expired": lease.body["state"] == "running" and lease.takeable(at), } return leases diff --git a/benchmarks/pyathena_bench/report.py b/benchmarks/pyathena_bench/report.py index d26cb56fc..ba2b0fabb 100644 --- a/benchmarks/pyathena_bench/report.py +++ b/benchmarks/pyathena_bench/report.py @@ -96,7 +96,6 @@ def summarize(trials: list[dict[str, Any]]) -> list[dict[str, Any]]: t["status"] not in {"ok", "unsupported"} and not t.get("superseded") for t in measured ), - "superseded_failures": sum(bool(t.get("superseded")) for t in measured), "failed_warmups": sum( t.get("warmup", False) and t["status"] != "ok" for t in group ), @@ -161,16 +160,68 @@ def read_records(path: Path) -> list[dict[str, Any]]: return records +# A fleet attempt's per-trial records in S3, written between trials; see runner.run(). +CHECKPOINT = "checkpoint" + + +def trial_key(trial: dict[str, Any]) -> tuple[str, str, int | None]: + """Return the scale, case, and repetition that identify a trial across attempts. + + Args: + trial: Trial record or ``trial_start`` event. + + Returns: + The key; the repetition is None for a case recorded as unsupported. + """ + return (trial["scale"], trial["case_id"], trial.get("repetition")) + + +def read_trials(directory: Path) -> list[dict[str, Any]]: + """Read the trial records of one output directory, including its checkpoint. + + Args: + directory: Output directory of a run or of a fleet attempt. + + Returns: + Records from ``trials.jsonl`` and checkpoint records missing from it. + """ + trials = read_records(directory / "trials.jsonl") + recorded = {trial.get("trial") for trial in trials} + for path in sorted((directory / CHECKPOINT).glob("*.json")): + if not path.name.endswith("-start.json") and path.stem not in recorded: + trials.append(json.loads(path.read_text())) + return trials + + +def read_trial_starts(directory: Path) -> list[dict[str, Any]]: + """Read the ``trial_start`` events of one output directory, including its checkpoint. + + Args: + directory: Output directory of a run or of a fleet attempt. + + Returns: + Events from ``events.jsonl`` and checkpoint events missing from it. + """ + starts = [e for e in read_records(directory / "events.jsonl") if e["event"] == "trial_start"] + started = {event.get("trial") for event in starts} + starts.extend( + json.loads(path.read_text()) + for path in sorted((directory / CHECKPOINT).glob("*-start.json")) + if path.name.removesuffix("-start.json") not in started + ) + return starts + + def collect_trials(directory: Path) -> list[dict[str, Any]]: """Read the trials of an output directory and of its attempt subdirectories. A fleet job keeps each attempt in an ``attempt-`` subdirectory; their - trials are merged. - A started trial without a final record is reported as incomplete. A failed - or incomplete trial is marked ``superseded`` when another attempt recorded - the same measured repetition as successful, so a resumed job counts the - trials of an uninterrupted run while its failures stay visible; a second - successful record of a repetition is dropped. + trials are merged. A started trial without a final record is reported as + incomplete. A failed or incomplete trial whose measured repetition another + attempt recorded as successful is marked ``superseded``: it is noted but + not counted as failed, so a resumed job counts the trials of an + uninterrupted run. A repetition that two attempts recorded as successful + counts once. Args: directory: Output directory of a run, or a job directory with attempts. @@ -179,31 +230,25 @@ def collect_trials(directory: Path) -> list[dict[str, Any]]: Trial records. """ sources = [directory, *sorted(p for p in directory.glob("attempt-*") if p.is_dir())] - trials = [t for source in sources for t in read_records(source / "trials.jsonl")] + trials = [t for source in sources for t in read_trials(source)] recorded = {trial.get("trial") for trial in trials} trials.extend( {**event, "status": "incomplete", "error": "Trial started without a final record"} for source in sources - for event in read_records(source / "events.jsonl") - if event["event"] == "trial_start" and event["trial"] not in recorded + for event in read_trial_starts(source) + if event["trial"] not in recorded ) - succeeded = { - (t["scale"], t["case_id"], t.get("repetition")) - for t in trials - if t["status"] == "ok" and not t.get("warmup") - } + succeeded = {trial_key(t) for t in trials if t["status"] == "ok" and not t.get("warmup")} merged = [] - kept: set[tuple[Any, ...]] = set() + counted: set[tuple[str, str, int | None]] = set() for trial in trials: - key = (trial["scale"], trial["case_id"], trial.get("repetition")) - if trial["status"] == "ok": - # Two attempts can both record a repetition when a takeover resumed - # from a checkpoint older than the interrupted attempt's last upload. - if len(sources) > 1 and key in kept: + if trial["status"] == "ok" and not trial.get("warmup"): + if len(sources) > 1 and trial_key(trial) in counted: continue - kept.add(key) - elif trial["status"] != "unsupported" and key in succeeded and not trial.get("warmup"): - trial = {**trial, "superseded": True} + counted.add(trial_key(trial)) + elif trial["status"] != "unsupported" and not trial.get("warmup"): + if trial_key(trial) in succeeded: + trial = {**trial, "superseded": True} merged.append(trial) return merged diff --git a/benchmarks/pyathena_bench/runner.py b/benchmarks/pyathena_bench/runner.py index d28f497ad..7aaf186b3 100644 --- a/benchmarks/pyathena_bench/runner.py +++ b/benchmarks/pyathena_bench/runner.py @@ -13,6 +13,7 @@ import hashlib import importlib.metadata import json +import logging import multiprocessing import os import platform @@ -24,6 +25,7 @@ import threading import time import uuid +from collections.abc import Callable from dataclasses import asdict from multiprocessing.connection import Connection from pathlib import Path @@ -36,13 +38,17 @@ from pyathena_bench.aws import ( cancel_queries, client, + run_locations, session, + split_location, validate_inputs, validate_manifest, ) from pyathena_bench.cases import Case from pyathena_bench.config import Settings, select_sql, write_json -from pyathena_bench.report import read_records +from pyathena_bench.report import CHECKPOINT, read_trials, trial_key + +_logger = logging.getLogger(__name__) class Observer: @@ -387,14 +393,38 @@ def completed_trials(directories: list[Path]) -> set[tuple[str, str, int | None] ``(scale, case_id, None)`` for cases recorded as unsupported. Warmups are not included: each attempt runs its own before its measured trials. """ - completed: set[tuple[str, str, int | None]] = set() - for directory in directories: - for trial in read_records(directory / "trials.jsonl"): - if trial["status"] == "ok" and not trial.get("warmup"): - completed.add((trial["scale"], trial["case_id"], trial["repetition"])) - elif trial["status"] == "unsupported": - completed.add((trial["scale"], trial["case_id"], None)) - return completed + return { + trial_key(trial) + for directory in directories + for trial in read_trials(directory) + if (trial["status"] == "ok" and not trial.get("warmup")) or trial["status"] == "unsupported" + } + + +def s3_checkpoint(s3: Any, location: str) -> Callable[[str, dict[str, Any]], None]: + """Return a function that writes a fleet attempt's per-trial records to S3. + + Args: + s3: S3 client. + location: The attempt's ``s3://`` prefix; records go below ``checkpoint/``. + + Returns: + A function that writes a record under a name, logging a failure: a + missing checkpoint only makes a resumed job repeat that trial. + """ + bucket, prefix = split_location(location) + + def write(name: str, record: dict[str, Any]) -> None: + try: + s3.put_object( + Bucket=bucket, + Key=f"{prefix}{CHECKPOINT}/{name}", + Body=json.dumps(record, sort_keys=True).encode(), + ) + except Exception: + _logger.warning("Could not write the checkpoint %s", name, exc_info=True) + + return write def run( @@ -404,6 +434,7 @@ def run( scales: list[str], output: Path, completed: set[tuple[str, str, int | None]] | None = None, + checkpoint: Callable[[str, dict[str, Any]], None] | None = None, ) -> bool: """Run every trial of the selected cases into a new output directory. @@ -416,6 +447,10 @@ def run( completed: Trials from :func:`completed_trials` to skip, so that a resumed job runs only the missing measured repetitions, after its own warmups; a case without missing repetitions is skipped. + checkpoint: Function from :func:`s3_checkpoint` that receives each + trial's ``trial_start`` event before the trial and its record + after it, between trials, so that a fleet worker can resume the + job if the host is lost. Returns: True if every trial succeeded. @@ -428,6 +463,7 @@ def run( write_json(output / "manifest.json", manifest) success = True s3 = client(session(settings), "s3") + results_root = run_locations(resources, manifest["run_id"])["results"] for scale in scales: for case in cases: if case.unsupported: @@ -463,7 +499,7 @@ def run( "repetition": repetition, "warmup": repetition < 0, "database": resources["ScratchDatabase"], - "output": f"{resources['Results']}runs/{manifest['run_id']}/trials/{trial}/", + "output": f"{results_root}trials/{trial}/", "temp_table": f"b_{manifest['run_id']}_{trial}", "sql": select_sql( resources["ScratchDatabase"], manifest["scales"][scale]["table"], case.shape @@ -471,7 +507,10 @@ def run( "expected_rows": manifest["scales"][scale]["rows"], "fixture": manifest["fixtures"].get(f"{scale}-{case.shape}-{case.transport}"), } - append_json(output / "events.jsonl", {"event": "trial_start", **payload}) + start = {"event": "trial_start", **payload} + append_json(output / "events.jsonl", start) + if checkpoint: + checkpoint(f"{trial}-start.json", start) result = supervise(payload, output / "events.jsonl", settings) if ( result["status"] == "ok" @@ -484,18 +523,18 @@ def run( except Exception as exc: # The measurement stands; only the file summary is missing. result["output_files_error"] = f"{type(exc).__name__}: {exc}" - append_json( - output / "trials.jsonl", - { - "trial": trial, - "case_id": case.id, - "case": asdict(case), - "scale": scale, - "repetition": repetition, - "warmup": repetition < 0, - **result, - }, - ) + record = { + "trial": trial, + "case_id": case.id, + "case": asdict(case), + "scale": scale, + "repetition": repetition, + "warmup": repetition < 0, + **result, + } + append_json(output / "trials.jsonl", record) + if checkpoint: + checkpoint(f"{trial}.json", record) success = success and result["status"] == "ok" # A failed worker can leave an unobserved query running in Athena. # Do not let it contaminate subsequent trials before cleanup. diff --git a/benchmarks/tests/test_aws.py b/benchmarks/tests/test_aws.py index 948012210..c5cfc16d0 100644 --- a/benchmarks/tests/test_aws.py +++ b/benchmarks/tests/test_aws.py @@ -11,9 +11,9 @@ import pytest from botocore.exceptions import ClientError from pyathena_bench.aws import ( - IncompleteStackError, cleanup, prepare, + run_locations, split_location, stack_resources, validate_inputs, @@ -350,32 +350,28 @@ def test_split_location_rejects_invalid_locations(location): split_location(location) -@pytest.mark.parametrize( - ("status", "outputs", "error"), - [ - ("CREATE_IN_PROGRESS", {}, IncompleteStackError), - ("CREATE_COMPLETE", {}, ValueError), - ( - "CREATE_COMPLETE", - {k: v for k, v in RESOURCES.items() if k not in {"StackId", "Results"}}, - None, - ), - ], - ids=["creating", "complete_without_outputs", "earlier_template"], -) -def test_stack_resources_waits_only_while_in_progress(monkeypatch, status, outputs, error): +def test_stack_resources_require_the_results_output(monkeypatch): + outputs = {k: v for k, v in RESOURCES.items() if k != "StackId"} stack = { "StackId": "arn:stack", - "StackStatus": status, "Tags": [{"Key": "Purpose", "Value": "pyathena-benchmark"}], "Outputs": [{"OutputKey": k, "OutputValue": v} for k, v in outputs.items()], } fake = SimpleNamespace(describe_stacks=lambda **kwargs: {"Stacks": [stack]}) monkeypatch.setattr("pyathena_bench.aws.client", lambda *args: fake) - if error is None: - # A stack from a template before the Results output keeps results in the scratch bucket. - assert stack_resources(None, "stack")["Results"] == "s3://temporary/" - return - with pytest.raises(error) as raised: + assert stack_resources(None, "stack")["Results"] == "s3://temporary/" + stack["Outputs"] = [o for o in stack["Outputs"] if o["OutputKey"] != "Results"] + with pytest.raises(ValueError, match="Stack outputs are incomplete"): stack_resources(None, "stack") - assert (type(raised.value) is IncompleteStackError) is (error is IncompleteStackError) + + +@pytest.mark.parametrize( + ("results", "expected"), + [ + ("s3://temporary/", f"s3://temporary/runs/{RUN}/"), + ("s3://results/bench/stack/", f"s3://results/bench/stack/runs/{RUN}/"), + ], +) +def test_run_locations(results, expected): + locations = run_locations({**RESOURCES, "Results": results}, RUN) + assert locations == {"scratch": f"s3://temporary/runs/{RUN}/", "results": expected} diff --git a/benchmarks/tests/test_cloudformation.py b/benchmarks/tests/test_cloudformation.py index fb99d1a28..e8e7cd96e 100644 --- a/benchmarks/tests/test_cloudformation.py +++ b/benchmarks/tests/test_cloudformation.py @@ -175,11 +175,16 @@ def test_hosts_start_the_queue_worker_and_expose_their_tags(): assert not errors launch = template["Resources"]["LaunchTemplate"]["Properties"]["LaunchTemplateData"] assert launch["MetadataOptions"]["InstanceMetadataTags"] == "enabled" - user_data = launch["UserData"]["Fn::Base64"]["Fn::Sub"] + user_data, variables = launch["UserData"]["Fn::Base64"]["Fn::Sub"] assert "if [ -n '${QueueName}' ]; then" in user_data + # The worker reads the results location from user data, before the stack's outputs exist. assert ( - "ExecStart=/opt/pyathena/.venv/bin/python -m pyathena_bench worker --stack ${AWS::StackId}" + "ExecStart=/opt/pyathena/.venv/bin/python -m pyathena_bench worker --results ${Results}" " --name ${QueueName} --api-slots ${ApiSlots} --wait --terminate-host" ) in user_data + assert ( + "printf '%s\\n' '${Results}' > /opt/pyathena/benchmarks/results-location.txt" in user_data + ) + assert variables["Results"] == template["Outputs"]["Results"]["Value"] assert "systemctl enable --now pyathena-bench-worker.service" in user_data assert template["Parameters"]["QueueName"]["Default"] == "" diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index d96881cc5..5a29330b3 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -19,8 +19,7 @@ import pytest from botocore.exceptions import ClientError -from pyathena_bench.__main__ import main, wait_for_stack -from pyathena_bench.aws import IncompleteStackError +from pyathena_bench.__main__ import main from pyathena_bench.cases import Case from pyathena_bench.config import Settings from pyathena_bench.fleet import ( @@ -80,6 +79,9 @@ def delete_object(self, **kwargs): def upload_file(self, filename, bucket, key): self.put_object(Key=key, Body=Path(filename).read_bytes()) + def download_file(self, bucket, key, filename): + Path(filename).write_bytes(self.objects[key]) + def list_objects_v2(self, **kwargs): keys = [k for k in self.objects if k.startswith(kwargs["Prefix"])] return {"KeyCount": len(keys[: kwargs.get("MaxKeys", 1000)])} @@ -312,7 +314,7 @@ def test_workers_run_each_job_once_and_record_results(self, tmp_path): runs = [] settled = [] - def runner(item, workdir, out, resume, stop): + def runner(item, workdir, out, resume, stop, checkpoint): runs.append(item["id"]) out.mkdir(parents=True) (out / "summary.csv").write_text("x") @@ -494,7 +496,7 @@ def test_failed_jobs_are_quiesced_and_unconfirmed_queries_stop_the_worker(self, _, queue = queue_with([job("first"), job("second")], tmp_path) runs = [] - def runner(item, workdir, out, resume, stop): + def runner(item, workdir, out, resume, stop, checkpoint): runs.append(item["id"]) out.mkdir(parents=True) (out / "events.jsonl").write_text( @@ -544,7 +546,9 @@ def test_trial_evidence_reads_reported_queries_and_trial_prefixes(self, tmp_path ] ) ) - assert trial_evidence(tmp_path) == ({"a"}, ("s3://b/runs/r/trials/t1/",)) + # A trial process killed while writing leaves an incomplete final record. + with pytest.warns(UserWarning, match="incomplete final record"): + assert trial_evidence(tmp_path) == ({"a"}, ("s3://b/runs/r/trials/t1/",)) assert trial_evidence(tmp_path / "missing") == (set(), ()) def test_jobs_wider_than_the_api_slots_stop_the_worker_before_claiming(self, tmp_path): @@ -651,11 +655,11 @@ def settle(output): assert ran == ["first"] errors = status(queue)["failed"]["first"]["quiesce_errors"] assert errors == [ - "Could not quiesce the job: RuntimeError: Could not inspect query history" + "Could not quiesce attempt-1: RuntimeError: Could not inspect query history" ] -def ok_trial(repetition): +def ok_trial(repetition, **changes): return { "trial": f"t{repetition}", "case_id": "case", @@ -667,10 +671,34 @@ def ok_trial(repetition): "successful_queries_per_second": 1, "rss_peak_bytes": 1, "max_threads": 1, + **changes, } -class TestLeases: +def take_over(queue, job_id, host="other", attempt=2): + """Write the claim that another worker holds after taking the job over.""" + at = time.time() + queue.put( + queue.key("claims", job_id), + {"host": host, "job": job_id, "attempt": attempt, "state": "running", + "renewed": at, "expires": at + 600}, + ) # fmt: skip + + +def recording_runner(calls, code=0, trials=None): + """Return a runner that records its calls and writes the given trial records.""" + + def runner(item, workdir, out, resume, stop, checkpoint): + calls.append((item["id"], out.name, [p.name for p in resume])) + out.mkdir(parents=True) + records = (trials or {}).get(item["id"], []) + (out / "trials.jsonl").write_text("".join(json.dumps(t) + "\n" for t in records)) + return code + + return runner + + +class TestLease: def test_only_one_worker_takes_over_an_expired_lease(self, tmp_path): _, queue = queue_with([], tmp_path) key = queue.key("claims", "a") @@ -708,152 +736,238 @@ def lost_response(**kwargs): assert taken.body["attempt"] == (1 if operation == "create" else 2) assert queue.read_lease(key).etag == taken.etag - def test_live_claims_are_left_and_expired_ones_return_to_the_fleet(self, tmp_path, monkeypatch): - _, queue = queue_with([job("live"), job("expired")], tmp_path) - busy = queue.acquire(queue.key("claims", "live"), {"host": "busy"}, 600) - queue.acquire(queue.key("claims", "expired"), {"host": "dead"}, -1) - runs = [] - waits = [] + def test_acquire_does_not_write_a_live_lease(self, tmp_path): + s3, queue = queue_with([], tmp_path) + key = queue.key("claims", "a") + queue.acquire(key, {"host": "busy"}, 600) + attempts = [] + original = s3.put_object + s3.put_object = lambda **kwargs: attempts.append(kwargs["Key"]) or original(**kwargs) + assert queue.acquire(key, {"host": "me"}, 600) is None + assert attempts == [] - def runner(item, workdir, out, resume, stop): - runs.append((item["id"], out.name, resume)) - return 0 + def test_ending_an_expired_lease_keeps_another_workers_takeover(self, tmp_path): + _, queue = queue_with([], tmp_path) + key = queue.key("slots", "0") + mine = queue.acquire(key, {"host": "slow"}, -1) + theirs = queue.acquire(key, {"host": "other"}, 600) + assert not queue.end_lease(mine, "released") + assert queue.read_lease(key).body == theirs.body + assert live_slots(queue) == ["0"] - def sleep(seconds): - # The busy worker renews its claim while this worker waits to exit. - waits.append(seconds) - assert queue.renew(busy, 600) + @pytest.mark.parametrize("state", ["interrupted", "released"]) + def test_ending_a_lease_after_its_own_concurrent_renewal(self, tmp_path, state): + _, queue = queue_with([], tmp_path) + key = queue.key("claims", "a") + lease = queue.acquire(key, {"host": "h", "job": "a"}, 600) + copy = Lease(key, dict(lease.body), lease.etag) + assert queue.renew(lease, 600) + assert queue.end_lease(copy, state) + assert queue.read_lease(key).body["state"] == state - monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) - ran = work( - queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner, settle=no_queries + def test_slots_are_written_only_when_enough_are_free(self, tmp_path): + s3, queue = queue_with([], tmp_path) + queue.acquire(queue.key("slots", "0"), {"host": "busy"}, 600) + attempts = [] + original = s3.put_object + s3.put_object = lambda **kwargs: attempts.append(kwargs["Key"]) or original(**kwargs) + assert queue.acquire_slots(2, 2, {"host": "me"}, 600) is None + assert attempts == [] + + +class TestJobWatch: + def test_renews_the_claim_and_slots(self, tmp_path): + _, queue = queue_with([], tmp_path) + claim = queue.acquire(queue.key("claims", "a"), {"host": "h"}, 1) + slot = queue.acquire(queue.key("slots", "0"), {"host": "h"}, 1) + renewed = claim.body["renewed"], slot.body["renewed"] + # The renew interval is a tenth of the lease period. + with JobWatch(queue, claim, [slot], 1, lambda: False) as watch: + time.sleep(0.5) + assert not watch.stop_event.is_set() + assert queue.read_lease(claim.key).body["renewed"] > renewed[0] + assert queue.read_lease(slot.key).body["renewed"] > renewed[1] + + def test_a_claim_taken_over_by_another_worker_is_lost(self, tmp_path): + _, queue = queue_with([], tmp_path) + claim = queue.acquire(queue.key("claims", "a"), {"host": "h"}, 0.1) + take_over(queue, "a") + with JobWatch(queue, claim, [], 0.1, lambda: False) as watch: + assert watch.stop_event.wait(5) + assert watch.claim_lost.is_set() + assert not watch.interrupted.is_set() + assert queue.read_lease(claim.key).body["host"] == "other" + + def test_a_stalled_renewal_that_outlasts_the_claim_loses_it(self, tmp_path): + _, queue = queue_with([], tmp_path) + claim = queue.acquire(queue.key("claims", "a"), {"host": "h"}, 0.2) + stalled = threading.Event() + released = threading.Event() + + def stall(lease, seconds, **changes): + stalled.set() + released.wait() + return False + + queue.renew = stall + with JobWatch(queue, claim, [], 0.1, lambda: False) as watch: + try: + assert stalled.wait(5) + # The claim expires while its renewal is still blocked in S3. + assert watch.claim_lost.wait(5) + assert not watch.interrupted.is_set() + finally: + # Unblock the renewal before leaving the block joins the thread. + released.set() + + +class TestInterruptionNotice: + def test_notice_comes_from_instance_metadata_or_a_rehearsal_file(self, tmp_path): + def missing(path): + raise urllib.error.HTTPError(path, 404, "Not Found", None, None) + + rehearsal = tmp_path / "notice" + assert not interruption_notice(missing, rehearsal) + assert interruption_notice(lambda path: '{"action": "terminate"}', rehearsal) + rehearsal.touch() + assert interruption_notice(missing, rehearsal) + + def test_notice_errors_count_as_no_notice(self, tmp_path): + def incomplete(path): + raise http.client.IncompleteRead(b"") + + assert not interruption_notice(incomplete) + + def broken(): + raise RuntimeError("unexpected") + + assert not check_notice(broken) + _, queue = queue_with([job("a")], tmp_path) + assert work(queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, settle=no_queries, + notice=broken) == ["a"] # fmt: skip + + +class FakeProcess: + """A child process that runs until it receives SIGINT.""" + + pid = 4321 + + def __init__(self, command, stdout, stderr, start_new_session): + self.command = command + self.start_new_session = start_new_session + self.signals = [] + + def poll(self): + return 130 if self.signals else None + + def send_signal(self, signum): + self.signals.append(signum) + + def wait(self, timeout=None): + return 130 + + +class TestRunJob: + def test_resumes_writes_checkpoints_and_interrupts_with_sigint(self, tmp_path, monkeypatch): + processes = [] + + def popen(command, stdout, stderr, start_new_session): + processes.append(FakeProcess(command, stdout, stderr, start_new_session)) + return processes[-1] + + killed = [] + monkeypatch.setattr("pyathena_bench.fleet.subprocess.Popen", popen) + monkeypatch.setattr( + "pyathena_bench.fleet.os.killpg", lambda pid, signum: killed.append((pid, signum)) + ) + stop = threading.Event() + stop.set() + earlier = tmp_path / "results" / "a" / "attempt-1" + code = run_job( + {"id": "a", "args": ["--suite", "single"]}, + tmp_path, + tmp_path / "results" / "a" / "attempt-2", + [earlier], + stop, + "s3://bucket/fleet/q/results/a/attempt-2/", ) - assert len(waits) == 1 + assert code == 130 + [process] = processes + assert process.signals == [signal.SIGINT] + assert process.command[-6:] == [ + "--checkpoint", "s3://bucket/fleet/q/results/a/attempt-2/", + "--suite", "single", "--resume", str(earlier), + ] # fmt: skip + # The orchestrator's own process group, including any trial process, is killed. + assert process.start_new_session + assert killed == [(4321, signal.SIGKILL)] + assert (tmp_path / "logs" / "a-attempt-2.log").exists() + + +class TestWork: + def test_expired_claims_are_taken_over_and_live_ones_left(self, tmp_path): + _, queue = queue_with([job("live"), job("expired")], tmp_path) + queue.acquire(queue.key("claims", "live"), {"host": "busy", "job": "live"}, 600) + queue.acquire(queue.key("claims", "expired"), {"host": "dead"}, -1) + calls = [] + ran = work( + queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=recording_runner(calls), + settle=no_queries, + ) # fmt: skip + # Nothing else is pending, so the worker exits while the live claim runs. assert ran == ["expired"] - assert runs == [("expired", "attempt-2", [])] - running = status(queue)["running"] - assert running["live"] == { - "host": "busy", - "job": None, - "attempt": 1, - "state": "running", - "expired": False, - "stale": False, - } + assert calls == [("expired", "attempt-2", [])] + assert status(queue)["running"]["live"] == { + "host": "busy", "job": "live", "attempt": 1, "state": "running", "expired": False, + } # fmt: skip def test_an_expired_slot_returns_to_the_fleet(self, tmp_path): _, queue = queue_with([job("heavy", heavy=True)], tmp_path) queue.acquire(queue.key("slots", "0"), {"host": "dead"}, -1) assert status(queue)["slots_in_use"]["0"]["expired"] ran = work( - queue, - tmp_path / "w", - api_slots=1, - poll_seconds=0, - runner=lambda *a: 0, + queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=lambda *a: 0, settle=no_queries, - ) + ) # fmt: skip assert ran == ["heavy"] assert not live_slots(queue) - def test_watch_renews_leases_and_uploads_progress(self, tmp_path): - _, queue = queue_with([], tmp_path) - lease = queue.acquire(queue.key("claims", "a"), {"host": "h"}, 1) - expires = lease.body["expires"] - uploads = [] - with JobWatch(queue, [lease], 600, 0.01, lambda: uploads.append(1), lambda: False): - deadline = time.monotonic() + 5 - while not uploads and time.monotonic() < deadline: - time.sleep(0.01) - assert uploads - assert queue.read_lease(lease.key).body["expires"] > expires + 500 - - def test_watch_stops_renewing_a_lease_taken_over_by_another_worker(self, tmp_path): - _, queue = queue_with([], tmp_path) - lease = queue.acquire(queue.key("claims", "a"), {"host": "h"}, -1) - assert queue.acquire(lease.key, {"host": "other"}, 600) - with JobWatch(queue, [lease], 600, 0.01, lambda: None, lambda: False) as watch: - deadline = time.monotonic() + 5 - while not watch.lost and time.monotonic() < deadline: - time.sleep(0.01) - assert watch.lost == [lease.key] - assert queue.read_lease(lease.key).body["host"] == "other" - - -class TestInterruption: - def test_notice_comes_from_instance_metadata_or_a_simulation_file(self, tmp_path): - def missing(path): - raise urllib.error.HTTPError(path, 404, "Not Found", None, None) - - absent = tmp_path / "notice" - assert not interruption_notice(missing, absent) - assert interruption_notice(lambda path: '{"action": "terminate"}', absent) - absent.touch() - assert interruption_notice(missing, absent) - - def test_interrupted_job_is_resumed_by_another_worker_with_the_same_trials(self, tmp_path): + def test_an_interrupted_job_is_resumed_by_another_worker_with_the_same_trials(self, tmp_path): s3, queue = queue_with([job("heavy", heavy=True), job("next")], tmp_path) notice = threading.Event() - def interrupted_runner(item, workdir, out, resume, stop): + def interrupted_runner(item, workdir, out, resume, stop, checkpoint): out.mkdir(parents=True) (out / "trials.jsonl").write_text(json.dumps(ok_trial(0)) + "\n") notice.set() assert stop.wait(5) with (out / "trials.jsonl").open("a") as stream: - interrupted = {**ok_trial(1), "status": "interrupted"} - del interrupted["queries"] - stream.write(json.dumps(interrupted) + "\n") + stream.write(json.dumps(ok_trial(1, status="interrupted", queries=[])) + "\n") return 1 with pytest.raises(SpotInterruptionError) as raised: work( - queue, - tmp_path / "w1", - api_slots=1, - poll_seconds=0, - runner=interrupted_runner, - settle=no_queries, - renew_seconds=0.01, - notice=notice.is_set, - ) + queue, tmp_path / "w1", api_slots=1, poll_seconds=0, runner=interrupted_runner, + settle=no_queries, lease_seconds=0.5, notice=notice.is_set, + ) # fmt: skip assert raised.value.ran == [] - claim = queue.read_lease(queue.key("claims", "heavy")) - assert claim.body["state"] == "interrupted" + assert queue.read_lease(queue.key("claims", "heavy")).body["state"] == "interrupted" assert not live_slots(queue) assert queue.key("results", "heavy", "attempt-1", "trials.jsonl") in s3.objects assert not queue.names("done") settled = [] - runs = [] + calls = [] writes = [] original_put = s3.put_object - - def recording_put(**kwargs): - writes.append(kwargs["Key"]) - return original_put(**kwargs) - - s3.put_object = recording_put - - def resumed_runner(item, workdir, out, resume, stop): - runs.append((item["id"], out.name, [p.name for p in resume])) - out.mkdir(parents=True) - if item["id"] == "heavy": - (out / "trials.jsonl").write_text( - "".join(json.dumps(ok_trial(r)) + "\n" for r in (1, 2)) - ) - return 0 - + s3.put_object = lambda **kwargs: writes.append(kwargs["Key"]) or original_put(**kwargs) + runner = recording_runner(calls, trials={"heavy": [ok_trial(1), ok_trial(2)]}) ran = work( - queue, - tmp_path / "w2", - api_slots=1, - poll_seconds=0, - runner=resumed_runner, + queue, tmp_path / "w2", api_slots=1, poll_seconds=0, runner=runner, settle=lambda output: settled.append(output.name) or [], - ) + ) # fmt: skip assert ran == ["heavy", "next"] - assert runs == [("heavy", "attempt-2", ["attempt-1"]), ("next", "attempt-1", [])] + assert calls == [("heavy", "attempt-2", ["attempt-1"]), ("next", "attempt-1", [])] # The interrupted attempt's queries are quiesced before the job continues. assert settled == ["attempt-1"] summary_key = queue.key("results", "heavy", "attempt-2", "job-summary.csv") @@ -861,18 +975,40 @@ def resumed_runner(item, workdir, out, resume, stop): # A recorded job always has its summary, below the attempt that the marker names. assert order == [summary_key, queue.key("done", "heavy")] assert json.loads(queue.read("done", "heavy"))["summary"] == "heavy/attempt-2" - summary = s3.objects[summary_key].decode() - row = next(csv.DictReader(io.StringIO(summary))) + row = next(csv.DictReader(io.StringIO(s3.objects[summary_key].decode()))) assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") # The interrupted trial stays visible. - assert row["superseded_failures"] == "1" assert "Superseded: interrupted" in row["notes"] - assert json.loads(queue.read("done", "heavy"))["attempt"] == 2 - def test_a_notice_while_quiescing_an_earlier_attempt_skips_the_run(self, tmp_path): + def test_a_job_resumes_from_the_checkpoint_of_a_lost_host(self, tmp_path): s3, queue = queue_with([job("a")], tmp_path) + # The host died after one trial; only its checkpoint is in S3. + queue.acquire(queue.key("claims", "a"), {"host": "dead", "job": "a"}, -1) + checkpoint = ("results", "a", "attempt-1", "checkpoint") + start = {"event": "trial_start", "trial": "t1", "output": "s3://r/runs/x/trials/t1/", + "case_id": "case", "scale": "small", "repetition": 1} # fmt: skip + queue.put(queue.key(*checkpoint, "t0.json"), ok_trial(0)) + queue.put(queue.key(*checkpoint, "t1-start.json"), start) + prefixes = [] + calls = [] + + def settle(output): + prefixes.extend(trial_evidence(output)[1]) + return [] + + runner = recording_runner(calls, trials={"a": [ok_trial(1)]}) + assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=settle) == ["a"] + assert calls == [("a", "attempt-2", ["attempt-1"])] + # The started trial is quiesced, and the finished one is not repeated. + assert prefixes == ["s3://r/runs/x/trials/t1/"] + summary = queue.key("results", "a", "attempt-2", "job-summary.csv") + row = next(csv.DictReader(io.StringIO(s3.objects[summary].decode()))) + assert row["successful_trials"] == "2" + + def test_a_notice_while_quiescing_an_earlier_attempt_skips_the_run(self, tmp_path): + _, queue = queue_with([job("a")], tmp_path) queue.acquire(queue.key("claims", "a"), {"host": "dead"}, -1) - s3.put_object(Key=queue.key("results", "a", "attempt-1", "events.jsonl"), Body=b"") + queue.put(queue.key("results", "a", "attempt-1", "events.jsonl"), {}) notice = threading.Event() def settle(output): @@ -884,16 +1020,8 @@ def runner(*args): pytest.fail("A job must not start after an interruption notice") with pytest.raises(SpotInterruptionError): - work( - queue, - tmp_path / "w", - api_slots=1, - poll_seconds=0, - runner=runner, - settle=settle, - renew_seconds=0.01, - notice=notice.is_set, - ) + work(queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner, + settle=settle, lease_seconds=0.5, notice=notice.is_set) # fmt: skip lease = queue.read_lease(queue.key("claims", "a")) assert (lease.body["state"], lease.body["attempt"]) == ("interrupted", 2) @@ -903,476 +1031,128 @@ def test_a_notice_between_jobs_stops_the_worker_without_claiming(self, tmp_path) work(queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, notice=lambda: True) assert not queue.names("claims") - def test_worker_waits_for_publication(self, tmp_path, monkeypatch): + def test_waits_for_publication_and_removes_a_rehearsal_notice(self, tmp_path, monkeypatch): s3 = FakeS3() queue = Queue(s3, "bucket", "run1", "results/stack/") config = tmp_path / "config.toml" config.write_text("[scales]\nsmall = 10000\n") manifest = tmp_path / "manifest.json" manifest.write_text("{}") + rehearsal = tmp_path / "w" / "interruption-notice" + rehearsal.parent.mkdir() + rehearsal.touch() monkeypatch.setattr( "pyathena_bench.fleet.time.sleep", lambda seconds: submit(queue, [job("a")], config, manifest), ) - ran = work( - queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, settle=no_queries, wait=True - ) - assert ran == ["a"] - assert all(key.startswith("results/stack/fleet/run1/") for key in s3.objects) - - -class FakeProcess: - """A child process that runs until it receives SIGINT.""" - - pid = 4321 - - def __init__(self, command, stdout, stderr, start_new_session): - self.command = command - self.start_new_session = start_new_session - self.signals = [] - - def poll(self): - return 130 if self.signals else None - - def send_signal(self, signum): - self.signals.append(signum) - - def wait(self, timeout=None): - return 130 + def no_instance_action(path): + raise urllib.error.HTTPError(path, 404, "Not Found", None, None) -def test_run_job_resumes_earlier_attempts_and_interrupts_with_sigint(tmp_path, monkeypatch): - processes = [] - - def popen(command, stdout, stderr, start_new_session): - processes.append(FakeProcess(command, stdout, stderr, start_new_session)) - return processes[-1] - - killed = [] - monkeypatch.setattr("pyathena_bench.fleet.subprocess.Popen", popen) - monkeypatch.setattr( - "pyathena_bench.fleet.os.killpg", lambda pid, signum: killed.append((pid, signum)) - ) - stop = threading.Event() - stop.set() - earlier = tmp_path / "results" / "a" / "attempt-1" - code = run_job( - {"id": "a", "args": ["--suite", "single"]}, - tmp_path, - tmp_path / "results" / "a" / "attempt-2", - [earlier], - stop, - ) - assert code == 130 - [process] = processes - assert process.signals == [signal.SIGINT] - # The orchestrator's own process group, including any trial process, is killed. - assert process.start_new_session - assert killed == [(4321, signal.SIGKILL)] - assert process.command[-4:] == ["--suite", "single", "--resume", str(earlier)] - assert (tmp_path / "logs" / "a-attempt-2.log").exists() - - -class TestLeaseOwnership: - def test_releasing_an_expired_slot_keeps_another_workers_takeover(self, tmp_path): - _, queue = queue_with([], tmp_path) - key = queue.key("slots", "0") - mine = queue.acquire(key, {"host": "slow"}, -1) - theirs = queue.acquire(key, {"host": "other"}, 600) - assert not queue.release_lease(mine) - assert queue.read_lease(key).body == theirs.body - assert live_slots(queue) == ["0"] + monkeypatch.setattr("pyathena_bench.fleet.instance_metadata", no_instance_action) + ran = work(queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, settle=no_queries, + wait=True) # fmt: skip + # The default notice reads the rehearsal file, which the worker removed when it started. + assert not rehearsal.exists() + assert ran == ["a"] - def test_interruption_is_marked_after_a_concurrent_renewal(self, tmp_path): - _, queue = queue_with([], tmp_path) - key = queue.key("claims", "a") - lease = queue.acquire(key, {"host": "h", "job": "a"}, 600) - stale = Lease(key, dict(lease.body), lease.etag) - assert queue.renew(lease, 600) - assert queue.end_lease(stale, "interrupted") - assert queue.read_lease(key).body["state"] == "interrupted" - # A stale copy of a slot is released the same way. - slot = queue.acquire(queue.key("slots", "0"), {"host": "h", "job": "a"}, 600) - copy = Lease(slot.key, dict(slot.body), slot.etag) - assert queue.renew(slot, 600) - assert queue.release_lease(copy) - assert live_slots(queue) == [] - - def test_a_lost_claim_stops_the_attempt_without_recording_the_job(self, tmp_path, monkeypatch): - s3, queue = queue_with([job("a")], tmp_path) + def test_a_lost_claim_stops_the_attempt_and_quiesces_it(self, tmp_path): + _, queue = queue_with([job("a")], tmp_path) stopped = [] settled = [] - def sleep(seconds): - # The new owner renews the claim while this worker waits to exit. - other = queue.read_lease(queue.key("claims", "a")) - assert queue.renew(other, 600) - - monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) - - def runner(item, workdir, out, resume, stop): + def runner(item, workdir, out, resume, stop, checkpoint): # Another worker took over the claim, for example after this host stalled. - s3.put_object( - Key=queue.key("claims", "a"), - Body=json.dumps( - {"host": "other", "attempt": 2, "state": "running", "renewed": time.time(), - "expires": time.time() + 600} - ).encode(), - ) # fmt: skip + take_over(queue, "a") stopped.append(stop.wait(5)) return 130 - ran = work( - queue, - tmp_path / "w", - api_slots=1, - poll_seconds=0, - runner=runner, - settle=lambda output: settled.append(output.name) or [], - renew_seconds=0.01, - ) - # The stopped attempt's queries are quiesced before the worker goes on. - assert settled == ["attempt-1"] + ran = work(queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner, + settle=lambda output: settled.append(output.name) or [], + lease_seconds=0.5) # fmt: skip assert stopped == [True] + assert settled == ["attempt-1"] assert ran == [] assert not queue.names("done") - assert queue.read_lease(queue.key("claims", "a")).body["host"] == "other" - def test_a_stale_claim_keeps_the_worker_until_it_can_take_it_over(self, tmp_path): + def test_a_former_owner_records_nothing_after_a_takeover(self, tmp_path): s3, queue = queue_with([job("a")], tmp_path) - at = time.time() - # The worker renewed its 200-second lease long ago and has probably died. - s3.put_object( - Key=queue.key("claims", "a"), - Body=json.dumps( - {"host": "dead", "attempt": 1, "state": "running", "renewed": at - 199.5, - "expires": at + 0.5} - ).encode(), - ) # fmt: skip - assert status(queue)["running"]["a"]["stale"] - runs = [] - ran = work( - queue, - tmp_path / "w", - api_slots=1, - poll_seconds=0.2, - runner=lambda item, workdir, out, *a: runs.append(out.name) or 0, - settle=no_queries, - ) - assert ran == ["a"] - assert runs == ["attempt-2"] - - -def test_worker_waits_for_stack_outputs(monkeypatch): - calls = [] - - def resources(session_, stack): - calls.append(stack) - if len(calls) < 3: - raise IncompleteStackError("Stack outputs are incomplete") - return {"Results": "s3://bucket/"} - - monkeypatch.setattr("pyathena_bench.__main__.stack_resources", resources) - monkeypatch.setattr("pyathena_bench.__main__.time.sleep", lambda seconds: None) - assert wait_for_stack(None, "stack", wait=True) == {"Results": "s3://bucket/"} - assert len(calls) == 3 - calls.clear() - with pytest.raises(IncompleteStackError): - wait_for_stack(None, "stack", wait=False) - - -def test_a_resumed_run_without_remaining_trials_succeeds(tmp_path, monkeypatch): - config = tmp_path / "config.toml" - config.write_text("[scales]\nsmall = 10000\n") - manifest = tmp_path / "manifest.json" - manifest.write_text("{}") - monkeypatch.setattr("pyathena_bench.__main__.run", lambda *args: True) - def report(directory): - raise ValueError("No trials to report") - - monkeypatch.setattr("pyathena_bench.__main__.report", report) - argv = ["--config", str(config), "run", "--suite", "single", "--scale", "small"] - paths = ["--manifest", str(manifest), "--out", str(tmp_path / "attempt-2")] - assert main([*argv, *paths, "--resume", str(tmp_path / "attempt-1")]) == 0 - with pytest.raises(ValueError, match="No trials"): - main([*argv, "--manifest", str(manifest), "--out", str(tmp_path / "attempt-3")]) - - -class TestLeaseExpiry: - def test_watch_stops_the_job_when_a_stalled_renewal_outlasts_the_claim(self, tmp_path): - _, queue = queue_with([], tmp_path) - claim = queue.acquire(queue.key("claims", "a"), {"host": "h"}, 0.2) - stalled = threading.Event() - released = threading.Event() - - def stall(lease, seconds, **changes): - stalled.set() - released.wait() - return False + def runner(item, workdir, out, resume, stop, checkpoint): + # This host stalled; another worker took the claim over before the + # watch noticed. + take_over(queue, "a") + out.mkdir(parents=True) + (out / "trials.jsonl").write_text(json.dumps(ok_trial(0)) + "\n") + return 0 - queue.renew = stall - watch = JobWatch(queue, [claim], 600, 0.01, lambda: None, lambda: False) - with watch: - try: - assert stalled.wait(5) - # The claim expires while its renewal is still blocked in S3. - assert watch.stop_requested.wait(5) - assert not watch.interrupted.is_set() - assert not watch.lost - finally: - # Unblock the renewal before leaving the block joins the thread. - released.set() + assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries) == [] + assert not queue.names("done") + assert not [key for key in s3.objects if key.endswith("job-summary.csv")] def test_losing_a_slot_does_not_stop_the_job(self, tmp_path): - s3, queue = queue_with([job("heavy", heavy=True)], tmp_path) + _, queue = queue_with([job("heavy", heavy=True)], tmp_path) - def runner(item, workdir, out, resume, stop): + def runner_slot_only(item, workdir, out, resume, stop, checkpoint): # Another worker took over the slot, for example after this host stalled. - s3.put_object( - Key=queue.key("slots", "0"), - Body=json.dumps( - {"host": "other", "state": "running", "expires": time.time() + 600} - ).encode(), - ) + queue.put(queue.key("slots", "0"), {"host": "other", "state": "running", + "expires": time.time() + 600}) # fmt: skip time.sleep(0.2) assert not stop.is_set() return 0 - assert work( - queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner, settle=no_queries, - renew_seconds=0.01, - ) == ["heavy"] # fmt: skip + assert work(queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner_slot_only, + settle=no_queries, lease_seconds=0.5) == ["heavy"] # fmt: skip assert json.loads(queue.read("done", "heavy"))["lost_slots"] == [queue.key("slots", "0")] # The other worker's slot is kept. - assert json.loads(s3.objects[queue.key("slots", "0")])["host"] == "other" - - def test_a_fresh_claim_that_is_never_renewed_is_taken_over_after_it_expires(self, tmp_path): - _, queue = queue_with([job("a")], tmp_path) - # Its worker died right after renewing it. - queue.acquire(queue.key("claims", "a"), {"host": "dead"}, 0.5) - runs = [] - ran = work( - queue, - tmp_path / "w", - api_slots=1, - poll_seconds=0.2, - runner=lambda item, workdir, out, *a: runs.append(out.name) or 0, - settle=no_queries, - ) - assert ran == ["a"] - assert runs == ["attempt-2"] - - def test_a_former_owner_records_nothing_after_a_takeover(self, tmp_path, monkeypatch): - s3, queue = queue_with([job("a")], tmp_path) - - def runner(item, workdir, out, resume, stop): - # This host stalled; the claim expired and another worker took it over, - # before this worker's watch noticed. - s3.put_object( - Key=queue.key("claims", "a"), - Body=json.dumps( - {"host": "other", "job": "a", "attempt": 2, "state": "running", - "renewed": time.time(), "expires": time.time() + 600} - ).encode(), - ) # fmt: skip - out.mkdir(parents=True) - (out / "trials.jsonl").write_text(json.dumps(ok_trial(0)) + "\n") - return 0 - - def sleep(seconds): - other = queue.read_lease(queue.key("claims", "a")) - assert queue.renew(other, 600) - - monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) - assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries) == [] - assert not queue.names("done") - assert not [key for key in s3.objects if key.endswith("job-summary.csv")] - + assert queue.read_lease(queue.key("slots", "0")).body["host"] == "other" -def test_a_claim_renewed_once_then_stale_does_not_let_the_worker_exit(tmp_path, monkeypatch): - s3, queue = queue_with([job("a")], tmp_path) - queue.acquire(queue.key("claims", "a"), {"host": "x", "job": "a"}, 600) - sleeps = [] - real_sleep = time.sleep - - def sleep(seconds): - sleeps.append(seconds) - if len(sleeps) == 1: - # Its worker renewed the claim once and then died long ago. - at = time.time() - s3.put_object( - Key=queue.key("claims", "a"), - Body=json.dumps( - {"host": "x", "job": "a", "attempt": 1, "state": "running", - "renewed": at - 199.7, "expires": at + 0.3} - ).encode(), - ) # fmt: skip - real_sleep(0.2) - - monkeypatch.setattr("pyathena_bench.fleet.time.sleep", sleep) - runs = [] - ran = work( - queue, - tmp_path / "w", - api_slots=1, - poll_seconds=0, - runner=lambda item, workdir, out, *a: runs.append(out.name) or 0, - settle=no_queries, - ) - assert ran == ["a"] - assert runs == ["attempt-2"] - - -class TestReviewRepairs: - def test_acquire_does_not_write_a_live_lease(self, tmp_path): - s3, queue = queue_with([], tmp_path) - key = queue.key("claims", "a") - queue.acquire(key, {"host": "busy"}, 600) - attempts = [] - original = s3.put_object - s3.put_object = lambda **kwargs: attempts.append(kwargs["Key"]) or original(**kwargs) - assert queue.acquire(key, {"host": "me"}, 600) is None - assert attempts == [] - - def test_notice_errors_do_not_stop_the_watch_or_the_worker(self, tmp_path): - def incomplete(path): - raise http.client.IncompleteRead(b"") - - assert not interruption_notice(incomplete, tmp_path / "notice") - - def broken(): - raise RuntimeError("unexpected") - - assert not check_notice(broken) + def test_an_unconfirmed_stopped_attempt_keeps_the_host(self, tmp_path): _, queue = queue_with([job("a")], tmp_path) - assert work(queue, tmp_path / "w", api_slots=1, runner=lambda *a: 0, settle=no_queries, - notice=broken) == ["a"] # fmt: skip - - def test_progress_uploads_at_trial_boundaries(self, tmp_path): - s3, queue = queue_with([job("a")], tmp_path) - events_key = queue.key("results", "a", "attempt-1", "events.jsonl") - uploads = [] - original = s3.upload_file - def recording(filename, bucket, key): - uploads.append(key) - return original(filename, bucket, key) - - s3.upload_file = recording - - def runner(item, workdir, out, resume, stop): - out.mkdir(parents=True) - events = out / "events.jsonl" - events.write_text(json.dumps({"event": "trial_start", "trial": "t"}) + "\n") - deadline = time.monotonic() + 5 - while events_key not in uploads and time.monotonic() < deadline: - time.sleep(0.01) - started = uploads.count(events_key) - with events.open("a") as stream: - for _ in range(20): - stream.write(json.dumps({"event": "query", "query_id": "q"}) + "\n") - stream.flush() - time.sleep(0.01) - # Query events between trial boundaries are not uploaded. - assert uploads.count(events_key) == started == 1 - return 0 - - assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries, - renew_seconds=0.01) == ["a"] # fmt: skip - - def test_worker_removes_a_rehearsal_notice_when_it_starts(self, tmp_path, monkeypatch): - config = tmp_path / "config.toml" - config.write_text("[scales]\nsmall = 10000\n") - monkeypatch.chdir(tmp_path) - Path("interruption-notice").touch() - monkeypatch.setattr( - "pyathena_bench.__main__.wait_for_stack", lambda *a: {"Results": "s3://bucket/"} - ) - monkeypatch.setattr("pyathena_bench.__main__.session", lambda settings: None) - monkeypatch.setattr("pyathena_bench.__main__.client", lambda *a: FakeS3()) - monkeypatch.setattr("pyathena_bench.__main__.work", lambda *a, **k: []) - argv = ["--config", str(config), "worker", "--stack", "s", "--name", "q"] - assert main(argv) == 0 - assert not Path("interruption-notice").exists() - - -class TestReviewRepairsFollowUp: - def test_an_unconfirmed_stopped_attempt_keeps_the_host(self, tmp_path, monkeypatch): - s3, queue = queue_with([job("a")], tmp_path) - - def runner(item, workdir, out, resume, stop): - s3.put_object( - Key=queue.key("claims", "a"), - Body=json.dumps( - {"host": "other", "job": "a", "attempt": 2, "state": "running", - "renewed": time.time(), "expires": time.time() + 600} - ).encode(), - ) # fmt: skip + def runner(item, workdir, out, resume, stop, checkpoint): + take_over(queue, "a") assert stop.wait(5) return 130 with pytest.raises(RuntimeError, match="Could not quiesce a stopped attempt of a"): - work( - queue, - tmp_path / "w", - api_slots=1, - runner=runner, - settle=lambda output: ["q1: still RUNNING"], - renew_seconds=0.01, - ) + work(queue, tmp_path / "w", api_slots=1, runner=runner, + settle=lambda output: ["q1: still RUNNING"], lease_seconds=0.5) # fmt: skip assert not queue.names("done") - def test_a_failed_trial_start_upload_is_retried(self, tmp_path): - s3, queue = queue_with([job("a")], tmp_path) - events_key = queue.key("results", "a", "attempt-1", "events.jsonl") - original = s3.upload_file - failures = [] - - def flaky(filename, bucket, key): - if key == events_key and not failures: - failures.append(key) - raise ClientError({"Error": {"Code": "SlowDown"}}, "PutObject") - return original(filename, bucket, key) - - s3.upload_file = flaky - - def runner(item, workdir, out, resume, stop): - out.mkdir(parents=True) - (out / "events.jsonl").write_text(json.dumps({"event": "trial_start"}) + "\n") - deadline = time.monotonic() + 5 - while events_key not in s3.objects and time.monotonic() < deadline: - time.sleep(0.01) - # The trial is still running; its start was uploaded after one failure. - assert failures == [events_key] - assert events_key in s3.objects - return 0 + def test_cleanup_errors_survive_a_claim_lost_during_cleanup(self, tmp_path): + _, queue = queue_with([job("a"), job("b")], tmp_path) + calls = [] - assert work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries, - renew_seconds=0.01) == ["a"] # fmt: skip + def settle(output): + # The failed job's queries do not stop, and the claim is lost meanwhile. + take_over(queue, "a") + return ["q1: still RUNNING"] + with pytest.raises(RuntimeError, match="q1: still RUNNING"): + work(queue, tmp_path / "w", api_slots=1, runner=recording_runner(calls, code=1), + settle=settle) # fmt: skip + # No further trials start while the query may still run. + assert [call[0] for call in calls] == ["a"] + assert not queue.names("done") -def test_cleanup_errors_survive_a_claim_lost_during_cleanup(tmp_path): - s3, queue = queue_with([job("a"), job("b")], tmp_path) - ran = [] - def settle(output): - # The failed job's queries do not stop, and the claim is lost meanwhile. - s3.put_object( - Key=queue.key("claims", "a"), - Body=json.dumps( - {"host": "other", "job": "a", "attempt": 2, "state": "running", - "renewed": time.time(), "expires": time.time() + 600} - ).encode(), - ) # fmt: skip - return ["q1: still RUNNING"] +class TestWorkerCommand: + def test_worker_uses_the_results_location_without_the_stack(self, tmp_path, monkeypatch): + config = tmp_path / "config.toml" + config.write_text("[scales]\nsmall = 10000\n") + queues = [] - def runner(item, workdir, out, resume, stop): - ran.append(item["id"]) - return 1 + def stack_resources(*args): + pytest.fail("A worker with --results needs no stack outputs") - with pytest.raises(RuntimeError, match="q1: still RUNNING"): - work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=settle) - # No further trials start while the query may still run. - assert ran == ["a"] - assert not queue.names("done") + monkeypatch.setattr("pyathena_bench.__main__.stack_resources", stack_resources) + monkeypatch.setattr("pyathena_bench.__main__.session", lambda settings: None) + monkeypatch.setattr("pyathena_bench.__main__.client", lambda *a: FakeS3()) + monkeypatch.setattr( + "pyathena_bench.__main__.work", lambda queue, *a, **k: queues.append(queue) or [] + ) + argv = ["--config", str(config), "worker", "--results", "s3://bucket/p/stack/"] + assert main([*argv, "--name", "q", "--wait"]) == 0 + [queue] = queues + assert (queue.bucket, queue.prefix) == ("bucket", "p/stack/fleet/q/") diff --git a/benchmarks/tests/test_report.py b/benchmarks/tests/test_report.py index f88da6cd5..39a2d66fc 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -138,60 +138,89 @@ def test_report_shows_the_most_output_files_and_the_largest_file(self, tmp_path) row = summarize([trials[0], {**failed_listing, "warmup": True}])[0] assert row["notes"] == "" + def test_report_merges_attempts_and_notes_superseded_failures(self, tmp_path): + def trial(repetition, status="ok"): + return { + "trial": f"t{repetition}-{status}", + "case_id": "case", + "scale": "small", + "repetition": repetition, + "warmup": False, + "status": status, + "queries": [{"total_seconds": 1}], + "successful_queries_per_second": 1, + "rss_peak_bytes": 1, + "max_threads": 1, + } + + first = tmp_path / "attempt-1" + second = tmp_path / "attempt-2" + first.mkdir() + second.mkdir() + (first / "trials.jsonl").write_text( + json.dumps(trial(0)) + + "\n" + + json.dumps({**trial(1, "interrupted"), "queries": []}) + + "\n" + ) + # The interrupted attempt also started a trial that it never recorded. + start = {"event": "trial_start", "trial": "lost", "case_id": "case", "scale": "small"} + (first / "events.jsonl").write_text(json.dumps({**start, "repetition": 2}) + "\n") + (second / "trials.jsonl").write_text( + json.dumps(trial(1)) + "\n" + json.dumps(trial(2)) + "\n" + ) + report(tmp_path) + with (tmp_path / "summary.csv").open() as stream: + row = next(csv.DictReader(stream)) + assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") + # The superseded failures stay visible in the notes: the interrupted and the + # unrecorded trial. + assert "Superseded: interrupted" in row["notes"] + assert "Superseded: Trial started without a final record" in row["notes"] -def test_report_merges_attempts_and_drops_superseded_failures(tmp_path): - def trial(repetition, status="ok"): - return { - "trial": f"t{repetition}-{status}", + def test_report_counts_a_repetition_recorded_by_two_attempts_once(self, tmp_path): + record = { "case_id": "case", "scale": "small", - "repetition": repetition, + "repetition": 0, "warmup": False, - "status": status, + "status": "ok", "queries": [{"total_seconds": 1}], "successful_queries_per_second": 1, "rss_peak_bytes": 1, "max_threads": 1, } + for attempt in (1, 2): + directory = tmp_path / f"attempt-{attempt}" + directory.mkdir() + (directory / "trials.jsonl").write_text( + json.dumps({**record, "trial": f"t{attempt}"}) + "\n" + ) + report(tmp_path) + with (tmp_path / "summary.csv").open() as stream: + assert next(csv.DictReader(stream))["successful_trials"] == "1" - first = tmp_path / "attempt-1" - second = tmp_path / "attempt-2" - first.mkdir() - second.mkdir() - (first / "trials.jsonl").write_text( - json.dumps(trial(0)) + "\n" + json.dumps({**trial(1, "interrupted"), "queries": []}) + "\n" - ) - # The interrupted attempt also started a trial that it never recorded. - start = {"event": "trial_start", "trial": "lost", "case_id": "case", "scale": "small"} - (first / "events.jsonl").write_text(json.dumps({**start, "repetition": 2}) + "\n") - (second / "trials.jsonl").write_text(json.dumps(trial(1)) + "\n" + json.dumps(trial(2)) + "\n") - report(tmp_path) - with (tmp_path / "summary.csv").open() as stream: - row = next(csv.DictReader(stream)) - assert (row["successful_trials"], row["failed_trials"]) == ("3", "0") - # The superseded failures stay visible: the interrupted and the unrecorded trial. - assert row["superseded_failures"] == "2" - assert "Superseded: interrupted" in row["notes"] - - -def test_report_counts_a_repetition_recorded_by_two_attempts_once(tmp_path): - record = { - "case_id": "case", - "scale": "small", - "repetition": 0, - "warmup": False, - "status": "ok", - "queries": [{"total_seconds": 1}], - "successful_queries_per_second": 1, - "rss_peak_bytes": 1, - "max_threads": 1, - } - for attempt in (1, 2): - directory = tmp_path / f"attempt-{attempt}" - directory.mkdir() - (directory / "trials.jsonl").write_text( - json.dumps({**record, "trial": f"t{attempt}"}) + "\n" - ) - report(tmp_path) - with (tmp_path / "summary.csv").open() as stream: - assert next(csv.DictReader(stream))["successful_trials"] == "1" + def test_report_reads_the_checkpoint_of_an_attempt(self, tmp_path): + record = { + "trial": "t0", + "case_id": "case", + "scale": "small", + "repetition": 0, + "warmup": False, + "status": "ok", + "queries": [{"total_seconds": 1}], + "successful_queries_per_second": 1, + "rss_peak_bytes": 1, + "max_threads": 1, + } + checkpoint = tmp_path / "attempt-1" / "checkpoint" + checkpoint.mkdir(parents=True) + # The host died after recording one trial and starting another. + (checkpoint / "t0.json").write_text(json.dumps(record)) + start = {"event": "trial_start", "trial": "t1", "case_id": "case", "scale": "small"} + (checkpoint / "t1-start.json").write_text(json.dumps({**start, "repetition": 1})) + (checkpoint / "t0-start.json").write_text(json.dumps({**start, "trial": "t0"})) + report(tmp_path) + with (tmp_path / "summary.csv").open() as stream: + row = next(csv.DictReader(stream)) + assert (row["successful_trials"], row["failed_trials"]) == ("1", "1") diff --git a/benchmarks/tests/test_runner.py b/benchmarks/tests/test_runner.py index 70e6fbc23..69616f862 100644 --- a/benchmarks/tests/test_runner.py +++ b/benchmarks/tests/test_runner.py @@ -245,73 +245,92 @@ def summarize(s3, location): ) -def test_resumed_run_skips_finished_trials_and_writes_trial_outputs_to_results( - monkeypatch, tmp_path -): - monkeypatch.setattr( - "pyathena_bench.runner.validate_manifest", - lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out", "Results": "s3://res/p/"}, - ) - monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) - monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) - payloads = [] - - def succeed(payload, events, settings): - payloads.append(payload) - return {"status": "ok"} +class TestResume: + def test_finished_trials_are_skipped_and_outputs_go_to_the_results_location( + self, monkeypatch, tmp_path, run_stubs + ): + monkeypatch.setattr( + "pyathena_bench.runner.validate_manifest", + lambda *args: {"ScratchDatabase": "s", "Bucket": "out", "Results": "s3://res/p/"}, + ) + payloads = [] + monkeypatch.setattr( + "pyathena_bench.runner.supervise", + lambda payload, events, settings: payloads.append(payload) or {"status": "ok"}, + ) + case = Case("cursor") + earlier = tmp_path / "attempt-1" + earlier.mkdir() + records = [ + {"scale": "small", "case_id": case.id, "repetition": 0, "status": "ok"}, + {"scale": "small", "case_id": case.id, "repetition": 1, "status": "interrupted"}, + ] + (earlier / "trials.jsonl").write_text("".join(json.dumps(r) + "\n" for r in records)) + settings = replace(Settings(), warmups=0, repetitions=3) + assert run( + settings, RUN_MANIFEST, [case], ["small"], tmp_path / "a2", completed_trials([earlier]) + ) + assert [p["repetition"] for p in payloads] == [1, 2] + prefix = f"s3://res/p/runs/{RUN_MANIFEST['run_id']}/trials/" + assert all(p["output"].startswith(prefix) for p in payloads) - monkeypatch.setattr("pyathena_bench.runner.supervise", succeed) - manifest = { - "run_id": "a" * 32, - "scales": {"small": {"table": "input", "rows": 10}}, - "fixtures": {}, - } - case = Case("cursor") - earlier = tmp_path / "attempt-1" - earlier.mkdir() - records = [ - {"scale": "small", "case_id": case.id, "repetition": 0, "status": "ok"}, - {"scale": "small", "case_id": case.id, "repetition": 1, "status": "interrupted"}, - ] - (earlier / "trials.jsonl").write_text("".join(json.dumps(r) + "\n" for r in records)) - settings = replace(Settings(), warmups=0, repetitions=3) - assert run( - settings, manifest, [case], ["small"], tmp_path / "attempt-2", completed_trials([earlier]) + @pytest.mark.parametrize( + ("finished", "expected"), [([0], [-1, 1, 2]), ([0, 1, 2], [])], ids=["missing", "complete"] ) - assert [p["repetition"] for p in payloads] == [1, 2] - assert all(p["output"].startswith(f"s3://res/p/runs/{'a' * 32}/trials/") for p in payloads) - - -@pytest.mark.parametrize( - ("finished", "expected"), [([0], [-1, 1, 2]), ([0, 1, 2], [])], ids=["missing", "complete"] -) -def test_resumed_run_repeats_its_warmups_before_missing_repetitions( - monkeypatch, tmp_path, finished, expected -): - monkeypatch.setattr( - "pyathena_bench.runner.validate_manifest", - lambda *args: {"ScratchDatabase": "scratch", "Bucket": "out", "Results": "s3://out/"}, - ) - monkeypatch.setattr("pyathena_bench.runner.validate_inputs", lambda *args: None) - monkeypatch.setattr("pyathena_bench.runner.environment", lambda *args: {}) - payloads = [] - monkeypatch.setattr( - "pyathena_bench.runner.supervise", - lambda payload, events, settings: payloads.append(payload) or {"status": "ok"}, - ) - case = Case("cursor") - earlier = tmp_path / "attempt-1" - earlier.mkdir() - records = [ - {"scale": "small", "case_id": case.id, "repetition": -1, "warmup": True, "status": "ok"} - ] - records += [ - {"scale": "small", "case_id": case.id, "repetition": r, "status": "ok"} for r in finished - ] - (earlier / "trials.jsonl").write_text("".join(json.dumps(r) + "\n" for r in records)) - manifest = {"run_id": "a" * 32, "scales": {"small": {"table": "t", "rows": 10}}, "fixtures": {}} - settings = replace(Settings(), warmups=1, repetitions=3) - assert run( - settings, manifest, [case], ["small"], tmp_path / "attempt-2", completed_trials([earlier]) - ) - assert [p["repetition"] for p in payloads] == expected + def test_each_attempt_runs_its_own_warmups( + self, monkeypatch, tmp_path, run_stubs, finished, expected + ): + payloads = [] + monkeypatch.setattr( + "pyathena_bench.runner.supervise", + lambda payload, events, settings: payloads.append(payload) or {"status": "ok"}, + ) + case = Case("cursor") + checkpoint = tmp_path / "attempt-1" / "checkpoint" + checkpoint.mkdir(parents=True) + warmup = {"scale": "small", "case_id": case.id, "repetition": -1, "warmup": True} + records = [{**warmup, "trial": "w", "status": "ok"}] + records += [ + {"trial": f"t{r}", "scale": "small", "case_id": case.id, "repetition": r, + "status": "ok"} + for r in finished + ] # fmt: skip + # Earlier records can come from the checkpoint of a lost host. + for record in records: + (checkpoint / f"{record['trial']}.json").write_text(json.dumps(record)) + settings = replace(Settings(), warmups=1, repetitions=3) + assert run( + settings, + RUN_MANIFEST, + [case], + ["small"], + tmp_path / "attempt-2", + completed_trials([tmp_path / "attempt-1"]), + ) + assert [p["repetition"] for p in payloads] == expected + + def test_checkpoints_are_written_between_trials(self, monkeypatch, tmp_path, run_stubs): + order = [] + monkeypatch.setattr( + "pyathena_bench.runner.supervise", + lambda payload, events, settings: order.append(("trial", payload["trial"])) + or {"status": "ok"}, + ) + settings = replace(Settings(), warmups=0, repetitions=2) + assert run( + settings, + RUN_MANIFEST, + [Case("cursor")], + ["small"], + tmp_path / "run", + checkpoint=lambda name, record: order.append(("checkpoint", name)), + ) + trials = [t for kind, t in order if kind == "trial"] + assert order == [ + ("checkpoint", f"{trials[0]}-start.json"), + ("trial", trials[0]), + ("checkpoint", f"{trials[0]}.json"), + ("checkpoint", f"{trials[1]}-start.json"), + ("trial", trials[1]), + ("checkpoint", f"{trials[1]}.json"), + ] From cd932b278b8e32355b37da3b336a01afb96c7c1b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 18:34:06 +0900 Subject: [PATCH 15/16] Release locally lost claims, contain downloads, and keep publication retryable Independent review findings on the simplification: - A claim lost only by its local expiry (a renewal that S3 accepted stalled) is released conditionally, which works only while it is still this worker's, so the job returns to the fleet. - Downloaded attempt files must resolve below the attempt directory. - A failed summary upload releases the claim and propagates instead of recording the job; status lists jobs with a report error as failed. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 2 +- benchmarks/pyathena_bench/fleet.py | 48 +++++++++++++++++-------- benchmarks/tests/test_fleet.py | 58 ++++++++++++++++++++++++++++-- 3 files changed, 91 insertions(+), 17 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 29e8ec579..9658f8678 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -429,7 +429,7 @@ A lost API slot does not stop the job; the `done` marker lists it under `lost_sl Each attempt of a job writes to `results//attempt-/`; between trials, its run writes each trial's start and record below `checkpoint/` there, so a lost host loses only its running trial, and the worker uploads the attempt's files when it ends. A trial whose start checkpoint could not be written, which the run logs, is not quiesced by the next attempt; its queries write to their own trial prefix and are cancelled by `cleanup`. After a job, the worker uploads the attempt's output directory and its log to `logs/`, renews its claim conditionally to confirm that it still owns the job, writes `job-summary.csv` and `job-summary.md` from every attempt into its attempt directory, and then creates `done/` with the exit code, attempt, host, times, and the summary's attempt directory under `summary`. -A failed trial stops only its own job; `status` lists jobs with a non-zero exit code, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. +A failed trial stops only its own job; `status` lists jobs with a non-zero exit code or a report error, the claims of running jobs, and the slots in use, each with its host, attempt, and whether its lease has expired. After a failed job, the worker stops the job's reported queries and every queued or running query in the workgroup that writes under the job's trial output prefixes, which also covers queries whose IDs were never reported before a trial process died. It waits until they stop before claiming another job; this scans the workgroup's query history. If it cannot confirm that, it records the errors in the `done` marker and exits. diff --git a/benchmarks/pyathena_bench/fleet.py b/benchmarks/pyathena_bench/fleet.py index 5ea522991..8da889195 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -587,11 +587,17 @@ def download_tree(self, directory: Path, *parts: str) -> list[str]: Returns: Relative paths of the downloaded files. + + Raises: + ValueError: If an object name would be written outside ``directory``. """ source = self.key(*parts) + "/" names = sorted(self.names("/".join(parts))) + root = directory.resolve() for name in names: - path = directory / name + path = (directory / name).resolve() + if not path.is_relative_to(root): + raise ValueError(f"Object {source + name} is outside {directory}") path.parent.mkdir(parents=True, exist_ok=True) self.s3.download_file(self.bucket, source + name, str(path)) return names @@ -987,27 +993,40 @@ def safe_settle(path: Path) -> list[str]: # A failed or stopped run can leave queries running. errors = safe_settle(output) upload_attempt() - lost = Outcome("lost", tuple(errors)) + + def abandon() -> Outcome: + # A claim lost only by its local expiry, for example while a renewal that + # S3 accepted stalled, may still be this worker's; releasing it, which + # works only then, returns the job to the fleet. + queue.end_lease(claim, "released") + return Outcome("lost", tuple(errors)) + # The conditional renewal proves that this worker still owns the job, so a # former owner never records it, and leaves at least --lease-seconds for # the few seconds of publication. if watch.claim_lost.is_set() or not queue.renew(claim, lease_seconds): - return lost + return abandon() report_error = None try: report(job_dir) - # The job's summary goes below this attempt, which the done marker names, - # so a former owner's late upload cannot replace the recorded summary. - for name in ("summary.csv", "summary.md"): - queue.s3.upload_file( - str(job_dir / name), queue.bucket, queue.key(*remote, f"job-{name}") - ) except Exception as exc: - # A job without trials, or with unreadable records, still gets its done marker. + # A job without trials, or with unreadable records, gets a failed done marker. report_error = f"{type(exc).__name__}: {exc}" + if report_error is None: + try: + # The job's summary goes below this attempt, which the done marker names, + # so a former owner's late upload cannot replace the recorded summary. + for name in ("summary.csv", "summary.md"): + queue.s3.upload_file( + str(job_dir / name), queue.bucket, queue.key(*remote, f"job-{name}") + ) + except Exception: + # Another worker publishes the job again from the uploaded attempts. + queue.end_lease(claim, "released") + raise if time.time() >= claim.body["expires"]: # Publication outlasted the renewed claim, which another worker may own now. - return lost + return abandon() recorded = queue.put_lease( queue.key("done", job["id"]), { @@ -1023,7 +1042,7 @@ def safe_settle(path: Path) -> list[str]: }, IfNoneMatch="*", ) - return lost if recorded is None else Outcome("done", tuple(errors)) + return abandon() if recorded is None else Outcome("done", tuple(errors)) def work( @@ -1320,7 +1339,8 @@ def status(queue: Queue) -> dict[str, Any]: queue: Queue to inspect. Returns: - Counts of jobs, claims, and completions, with failed job details, the + Counts of jobs, claims, and completions, with the details of failed + jobs (a non-zero exit code or a report error), the claims of unfinished jobs, and the API slots in use, each lease with its host, attempt, state, and whether it has expired. """ @@ -1331,7 +1351,7 @@ def status(queue: Queue) -> dict[str, Any]: failed = {} for name in sorted(done): record = json.loads(queue.read("done", name)) - if record["exit_code"] != 0: + if record["exit_code"] != 0 or record.get("report_error"): failed[name] = record at = time.time() diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 5a29330b3..62a6c3c4f 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -317,7 +317,7 @@ def test_workers_run_each_job_once_and_record_results(self, tmp_path): def runner(item, workdir, out, resume, stop, checkpoint): runs.append(item["id"]) out.mkdir(parents=True) - (out / "summary.csv").write_text("x") + (out / "trials.jsonl").write_text(json.dumps(ok_trial(0)) + "\n") return 0 if item["id"] == "light" else 1 assert work( @@ -341,7 +341,8 @@ def runner(item, workdir, out, resume, stop, checkpoint): ) assert sorted(runs) == ["heavy", "light"] assert [path.parent.name for path in settled] == ["heavy"] - assert queue.key("results", "light", "attempt-1", "summary.csv") in s3.objects + assert queue.key("results", "light", "attempt-1", "trials.jsonl") in s3.objects + assert queue.key("results", "light", "attempt-1", "job-summary.csv") in s3.objects assert not live_slots(queue) summary = status(queue) assert summary["done"] == 2 @@ -699,6 +700,13 @@ def runner(item, workdir, out, resume, stop, checkpoint): class TestLease: + def test_downloads_stay_below_the_destination(self, tmp_path): + _, queue = queue_with([], tmp_path) + queue.put(queue.key("results", "a", "attempt-1", "..", "..", "escape.json"), {}) + with pytest.raises(ValueError, match="outside"): + queue.download_tree(tmp_path / "w" / "attempt-1", "results", "a", "attempt-1") + assert not (tmp_path / "w" / "escape.json").exists() + def test_only_one_worker_takes_over_an_expired_lease(self, tmp_path): _, queue = queue_with([], tmp_path) key = queue.key("claims", "a") @@ -1136,6 +1144,52 @@ def settle(output): assert [call[0] for call in calls] == ["a"] assert not queue.names("done") + def test_a_claim_lost_only_locally_returns_to_the_fleet(self, tmp_path): + _, queue = queue_with([job("a")], tmp_path) + renew = queue.renew + failing = threading.Event() + failing.set() + + def unavailable(lease, seconds, **changes): + # S3 is unreachable during the first attempt, until its claim expires locally. + if failing.is_set() and changes.get("state") is None: + raise ClientError({"Error": {"Code": "SlowDown"}}, "PutObject") + return renew(lease, seconds, **changes) + + queue.renew = unavailable + calls = [] + + def runner(item, workdir, out, resume, stop, checkpoint): + calls.append(out.name) + if len(calls) == 1: + assert stop.wait(5) + failing.clear() + return 130 + return 0 + + assert work(queue, tmp_path / "w", api_slots=1, poll_seconds=0, runner=runner, + settle=no_queries, lease_seconds=0.3) == ["a"] # fmt: skip + # The first attempt released its claim, so the job ran again instead of + # staying claimed by a worker that had given it up. + assert calls == ["attempt-1", "attempt-2"] + assert json.loads(queue.read("done", "a"))["attempt"] == 2 + + def test_a_failed_summary_upload_releases_the_claim(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + original = s3.upload_file + + def failing(filename, bucket, key): + if key.endswith("job-summary.csv"): + raise ClientError({"Error": {"Code": "InternalError"}}, "PutObject") + return original(filename, bucket, key) + + s3.upload_file = failing + runner = recording_runner([], trials={"a": [ok_trial(0)]}) + with pytest.raises(ClientError): + work(queue, tmp_path / "w", api_slots=1, runner=runner, settle=no_queries) + assert not queue.names("done") + assert queue.read_lease(queue.key("claims", "a")).body["state"] == "released" + class TestWorkerCommand: def test_worker_uses_the_results_location_without_the_stack(self, tmp_path, monkeypatch): From 0e5bbb7d3f982b77edfa710467a5c88cd93c7d2b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 12 Oct 2026 18:39:47 +0900 Subject: [PATCH 16/16] Test a claim that stays live in S3 after it expires locally Co-Authored-By: Claude Opus 5.5 --- benchmarks/tests/test_fleet.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/benchmarks/tests/test_fleet.py b/benchmarks/tests/test_fleet.py index 62a6c3c4f..eed359b33 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -1150,13 +1150,16 @@ def test_a_claim_lost_only_locally_returns_to_the_fleet(self, tmp_path): failing = threading.Event() failing.set() - def unavailable(lease, seconds, **changes): - # S3 is unreachable during the first attempt, until its claim expires locally. + def accepted_but_lost(lease, seconds, **changes): + # During the first attempt, S3 applies each renewal, but its response + # is lost, so the claim stays live in S3 while it expires locally. if failing.is_set() and changes.get("state") is None: - raise ClientError({"Error": {"Code": "SlowDown"}}, "PutObject") + at = time.time() + queue.put(lease.key, {**lease.body, "renewed": at, "expires": at + 600}) + raise ClientError({"Error": {"Code": "RequestTimeout"}}, "PutObject") return renew(lease, seconds, **changes) - queue.renew = unavailable + queue.renew = accepted_but_lost calls = [] def runner(item, workdir, out, resume, stop, checkpoint):