diff --git a/benchmarks/README.md b/benchmarks/README.md index 9d1995573..9658f8678 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -96,16 +96,23 @@ 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, 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 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. 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 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. +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. + 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`. @@ -277,6 +284,8 @@ 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. @@ -303,16 +312,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 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: ```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 \ @@ -386,38 +400,101 @@ The manifest for this selection must be prepared with `--scale medium large xlar It is published as `large-1` in place of the previous example, so the worker and status commands below apply unchanged. The selection has 96 jobs; from the job durations of the 2026-10-11 run, they take about 14 hours of job time on `r7g.2xlarge` hosts, and the longest, xlarge nested pandas CSV, about 75 minutes. -`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 --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. -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, 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, 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-/`; 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 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. -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. +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 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. +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 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 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 + +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/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: + +```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/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 +``` + +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//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 @@ -433,6 +510,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. @@ -442,10 +520,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" \ @@ -460,29 +539,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..0afcdb315 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,22 @@ Resources: - Effect: Allow Action: [s3:GetObject, s3:PutObject, s3:DeleteObject, s3:AbortMultipartUpload, s3:ListMultipartUploadParts] Resource: !Sub '${Bucket.Arn}/*' + # Queue state, reports, and trial outputs go only below this stack's results prefix. + # Listing covers the whole bucket: without s3:ListBucket, reading a missing key + # returns 403 instead of 404, which breaks existence checks of fsspec and pyarrow, + # and ListBucketMultipartUploads does not support the s3:prefix condition key. + - !If + - HasResultsBucket + - Effect: Allow + Action: [s3:GetBucketLocation, s3:ListBucket, s3:ListBucketMultipartUploads] + Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket} + - !Ref AWS::NoValue + - !If + - HasResultsBucket + - Effect: Allow + Action: [s3:GetObject, s3:PutObject, s3:DeleteObject, s3:AbortMultipartUpload, s3:ListMultipartUploadParts] + Resource: !Sub arn:${AWS::Partition}:s3:::${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/* + - !Ref AWS::NoValue InstanceProfile: Type: AWS::IAM::InstanceProfile Properties: @@ -271,6 +326,8 @@ Resources: MetadataOptions: HttpTokens: required HttpEndpoint: enabled + # A worker reads the KeepHost tag from instance metadata. + InstanceMetadataTags: enabled ImageId: !If [IsArm64, !Ref Arm64ImageId, !Ref ImageId] InstanceType: !Ref InstanceType IamInstanceProfile: @@ -292,22 +349,48 @@ Resources: - Key: Name Value: !Sub '${AWS::StackName}-benchmark' UserData: - Fn::Base64: !Sub | - #!/bin/bash - set -euo pipefail - dnf install -y git aws-cfn-bootstrap - trap '/opt/aws/bin/cfn-signal -e $? --stack ${AWS::StackName} --resource Fleet --region ${AWS::Region}' EXIT - curl --proto '=https' --tlsv1.2 -LsSf https://astral.sh/uv/${UvVersion}/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh - git clone https://github.com/pyathena-dev/PyAthena.git /opt/pyathena - git -C /opt/pyathena fetch origin ${GitCommit} - git -C /opt/pyathena checkout --detach ${GitCommit} - export UV_PYTHON_INSTALL_DIR=/opt/uv-python - uv python install ${PythonVersion} - uv sync --directory /opt/pyathena/benchmarks --locked --no-dev --python ${PythonVersion} - chmod -R a+rX /opt/uv-python - chown -R ec2-user:ec2-user /opt/pyathena - printf '%s\n' '${AWS::StackId}' > /opt/pyathena/benchmarks/stack-id.txt - printf '%s\n' 'Setup complete. Run preflight and prepare manually.' + Fn::Base64: !Sub + - | + #!/bin/bash + set -euo pipefail + dnf install -y git aws-cfn-bootstrap + trap '/opt/aws/bin/cfn-signal -e $? --stack ${AWS::StackName} --resource Fleet --region ${AWS::Region}' EXIT + curl --proto '=https' --tlsv1.2 -LsSf https://astral.sh/uv/${UvVersion}/install.sh | env UV_INSTALL_DIR=/usr/local/bin sh + git clone https://github.com/pyathena-dev/PyAthena.git /opt/pyathena + git -C /opt/pyathena fetch origin ${GitCommit} + git -C /opt/pyathena checkout --detach ${GitCommit} + export UV_PYTHON_INSTALL_DIR=/opt/uv-python + uv python install ${PythonVersion} + uv sync --directory /opt/pyathena/benchmarks --locked --no-dev --python ${PythonVersion} + chmod -R a+rX /opt/uv-python + chown -R ec2-user:ec2-user /opt/pyathena + printf '%s\n' '${AWS::StackId}' > /opt/pyathena/benchmarks/stack-id.txt + printf '%s\n' '${Results}' > /opt/pyathena/benchmarks/results-location.txt + sudo -u ec2-user mkdir -p /opt/pyathena/benchmarks/results + if [ -n '${QueueName}' ]; then + cat > /etc/systemd/system/pyathena-bench-worker.service <<'UNIT' + [Unit] + Description=PyAthena benchmark fleet worker + Wants=network-online.target + After=network-online.target + [Service] + User=ec2-user + WorkingDirectory=/opt/pyathena/benchmarks + Environment=AWS_DEFAULT_REGION=${AWS::Region} + ExecStart=/opt/pyathena/.venv/bin/python -m pyathena_bench worker --results ${Results} --name ${QueueName} --api-slots ${ApiSlots} --wait --terminate-host + StandardOutput=append:/opt/pyathena/benchmarks/results/worker.log + StandardError=inherit + [Install] + WantedBy=multi-user.target + UNIT + systemctl daemon-reload + systemctl enable --now pyathena-bench-worker.service + fi + printf '%s\n' 'Setup complete. Run preflight and prepare manually.' + - Results: !If + - HasResultsBucket + - !Sub s3://${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/ + - !Sub s3://${Bucket}/ Fleet: Type: AWS::AutoScaling::AutoScalingGroup DependsOn: [Route, RouteAssociation, RouteAssociation2, RouteAssociation3] @@ -320,8 +403,17 @@ Resources: MinSize: 0 MaxSize: !Ref FleetSize DesiredCapacity: !Ref FleetSize - VPCZoneIdentifier: [!Ref Subnet, !Ref Subnet2, !Ref Subnet3] + VPCZoneIdentifier: !If + - UseFourZones + - [!Ref Subnet, !Ref Subnet2, !Ref Subnet3, !Ref Subnet4] + - [!Ref Subnet, !Ref Subnet2, !Ref Subnet3] HealthCheckType: EC2 + Tags: + # DependsOn cannot name the conditional fourth zone's route association; this + # reference makes hosts launch only after it exists, so bootstrap reaches the Internet. + - Key: pyathena-benchmark-route + Value: !If [UseFourZones, !Ref RouteAssociation4, none] + PropagateAtLaunch: false MixedInstancesPolicy: LaunchTemplate: LaunchTemplateSpecification: @@ -339,6 +431,11 @@ Outputs: Value: !Ref ScratchDatabase AutoScalingGroup: Value: !Ref Fleet + Results: + Value: !If + - HasResultsBucket + - !Sub s3://${ResultsBucket}/${ResultsPrefix}/${AWS::StackName}/ + - !Sub s3://${Bucket}/ SourceBucket: Value: !Ref SourceBucket SourcePrefix: diff --git a/benchmarks/pyathena_bench/__main__.py b/benchmarks/pyathena_bench/__main__.py index f4909acbe..2556d4db8 100644 --- a/benchmarks/pyathena_bench/__main__.py +++ b/benchmarks/pyathena_bench/__main__.py @@ -22,7 +22,9 @@ preflight, prepare, prepare_fixtures, + run_locations, session, + split_location, stack_resources, validate_manifest, ) @@ -32,15 +34,17 @@ API_FAMILIES, Filters, Queue, + SpotInterruptionError, expand_jobs, merge_jobs, retire_host, status, submit, + suspend_rebalance, work, ) from pyathena_bench.report import report -from pyathena_bench.runner import run +from pyathena_bench.runner import completed_trials, run, s3_checkpoint def parser() -> argparse.ArgumentParser: @@ -92,6 +96,18 @@ def parser() -> argparse.ArgumentParser: if name == "run": command.add_argument("--manifest", type=Path, required=True) command.add_argument("--out", type=Path, required=True) + command.add_argument( + "--resume", + type=Path, + nargs="+", + default=[], + help="Output directories of earlier attempts whose finished trials are skipped", + ) + command.add_argument( + "--checkpoint", + help="s3:// prefix below which each trial's start and record are written " + "between trials; set by fleet workers, whose job summary merges the attempts", + ) enqueue = commands.add_parser( "queue", help="Publish jobs and a prepared manifest for fleet workers" ) @@ -102,7 +118,11 @@ def parser() -> argparse.ArgumentParser: ) enqueue.add_argument("--manifest", type=Path, required=True) consume = commands.add_parser("worker", help="Run queued jobs on this host until none remain") - consume.add_argument("--stack", required=True) + location = consume.add_mutually_exclusive_group(required=True) + location.add_argument("--stack", help="Stack whose Results output holds the queue") + location.add_argument( + "--results", help="The stack's results location, which hosts know from their user data" + ) consume.add_argument("--name", required=True) consume.add_argument( "--api-slots", @@ -114,7 +134,19 @@ def parser() -> argparse.ArgumentParser: consume.add_argument( "--terminate-host", action="store_true", - help="Terminate this host through its Auto Scaling group after every job is claimed", + help="Suspend zone rebalancing at start, and terminate this host through its Auto " + "Scaling group when no job is left to claim", + ) + consume.add_argument( + "--wait", + action="store_true", + help="Wait for the queue to be published instead of failing", + ) + consume.add_argument( + "--lease-seconds", + type=float, + default=600.0, + help="Seconds after which a dead worker's claim and API slots return to the fleet", ) progress = commands.add_parser("status", help="Summarize a fleet queue") progress.add_argument("--stack", required=True) @@ -164,10 +196,13 @@ def main(argv: list[str] | None = None) -> int: if args.trials_only else f"b_{manifest['run_id']}_*" ), - "s3_prefix": f"s3://{resources['Bucket']}/runs/{manifest['run_id']}/" - + ("trials/" if args.trials_only else ""), + "s3_prefixes": [ + root + ("trials/" if args.trials_only else "") + for root in sorted( + set(run_locations(resources, manifest["run_id"]).values()) + ) + ], "preserve_inputs_and_fixtures": args.trials_only, - "query_output_prefix": f"s3://{resources['Bucket']}/runs/{manifest['run_id']}/", }, indent=2, ) @@ -177,20 +212,39 @@ def main(argv: list[str] | None = None) -> int: cleanup(settings, manifest, args.manifest, trials_only=args.trials_only) elif args.command in {"queue", "worker", "status"}: session_ = session(settings) - queue = Queue( - client(session_, "s3"), stack_resources(session_, args.stack)["Bucket"], args.name - ) + results = getattr(args, "results", None) or stack_resources(session_, args.stack)["Results"] + bucket, base = split_location(results) + queue = Queue(client(session_, "s3"), bucket, args.name, base) if args.command == "queue": validate_manifest(settings, read_json(args.manifest)) submit(queue, merge_jobs(read_json(p) for p in args.jobs), args.config, args.manifest) elif args.command == "worker": if args.api_slots < 1: raise ValueError("--api-slots must be positive") - ran = work(queue, args.workdir or Path("results/fleet") / args.name, args.api_slots) + if args.lease_seconds < 300: + raise ValueError("--lease-seconds must be at least 300") + autoscaling = client(session_, "autoscaling") + if args.terminate_host: + suspend_rebalance(autoscaling) + try: + ran = work( + queue, + args.workdir or Path("results/fleet") / args.name, + args.api_slots, + lease_seconds=args.lease_seconds, + wait=args.wait, + ) + except SpotInterruptionError as interrupted: + # The host is being reclaimed, and its group launches a replacement; after + # a rehearsed notice, the operator terminates the host the same way. + sys.stdout.write( + json.dumps({"ran": interrupted.ran, "interrupted": True}, indent=2) + "\n" + ) + return 0 sys.stdout.write(json.dumps({"ran": ran}, indent=2) + "\n") if args.terminate_host: # Only a normal return retires the host; an error keeps it for inspection. - retire_host(queue, client(session_, "autoscaling"), ran) + retire_host(queue, autoscaling, ran) else: sys.stdout.write(json.dumps(status(queue), indent=2) + "\n") else: @@ -248,8 +302,23 @@ def main(argv: list[str] | None = None) -> int: manifest = read_json(args.manifest) if args.suite == "init": prepare_fixtures(settings, manifest, args.manifest, scales, args.shape) - success = run(settings, manifest, cases, scales, args.out) - report(args.out) + checkpoint = ( + s3_checkpoint(client(session(settings), "s3"), args.checkpoint) + if args.checkpoint + else None + ) + success = run( + settings, + manifest, + cases, + scales, + args.out, + completed_trials(args.resume), + checkpoint, + ) + if not checkpoint: + # A fleet job's summary merges its attempts instead. + report(args.out) return 0 if success else 1 return 0 diff --git a/benchmarks/pyathena_bench/aws.py b/benchmarks/pyathena_bench/aws.py index c3f6008cc..9f54cabd0 100644 --- a/benchmarks/pyathena_bench/aws.py +++ b/benchmarks/pyathena_bench/aws.py @@ -37,6 +37,43 @@ 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 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", [])} @@ -52,6 +89,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 +99,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 +369,17 @@ 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 + 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. @@ -356,11 +400,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 +437,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 77b730429..8da889195 100644 --- a/benchmarks/pyathena_bench/fleet.py +++ b/benchmarks/pyathena_bench/fleet.py @@ -9,16 +9,23 @@ from __future__ import annotations +import contextlib +import http.client import json +import logging +import os +import signal import socket import subprocess import sys +import threading import time 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,14 +33,23 @@ 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 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. API_FAMILIES = {"cursor", "dict", "sqla_rest"} 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 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"} @dataclass(frozen=True) @@ -230,15 +246,44 @@ 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``, + ``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 ended. + """ + + 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 ended or expired before ``at``. + """ + return self.body["state"] != "running" or self.body["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. @@ -251,6 +296,30 @@ def key(self, *parts: str) -> str: """ return self.prefix + "/".join(parts) + def put(self, key: str, body: 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. @@ -264,18 +333,135 @@ 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 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 ended. + + 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. + """ + 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, + "attempt": 1, + "state": "running", + "renewed": at, + "expires": at + lease_seconds, + } + if current is None: + return self.put_lease(key, body, IfNoneMatch="*") + if not current.takeable(at): + return None + 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: + """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. + """ + 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 end_lease(self, lease: Lease, state: str) -> bool: + """End a lease that this worker still holds, so another worker can take it. + + 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: Lease to end. + state: ``released`` or ``interrupted``. + + Returns: + False if another worker took over the lease. + """ + 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=state, expires=0) def names(self, folder: str) -> set[str]: """List the object names directly below a queue folder. @@ -305,16 +491,22 @@ 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. + 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. 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: @@ -323,20 +515,29 @@ 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] = [] + 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): - key = self.key("slots", str(index)) - if self.put_if_absent(key, owner): - taken.append(key) + 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 key in taken: - self.release(key) + for lease in taken: + self.end_lease(lease, "released") raise - for key in taken: - self.release(key) + for lease in taken: + self.end_lease(lease, "released") return None def exists(self, *parts: str) -> bool: @@ -362,20 +563,45 @@ 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. + + 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).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 + def submit(queue: Queue, jobs: list[dict[str, Any]], config: Path, manifest: Path) -> None: """Publish jobs, configuration, and manifest for a fleet run. @@ -401,11 +627,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") @@ -415,12 +637,30 @@ 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, + checkpoint: str, +) -> 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. The + orchestrator runs in its own process group, which is killed when it + ends, so no trial process outlives it. 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. + checkpoint: The attempt's ``s3://`` prefix for per-trial checkpoints. Returns: Process exit code. @@ -428,18 +668,36 @@ 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), + "--checkpoint", checkpoint, *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, 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, ...]]: """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 @@ -447,18 +705,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)) @@ -507,42 +762,353 @@ 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 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 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, + claim: Lease, + slots: list[Lease], + lease_seconds: float, + notice: Callable[[], bool], + poll_seconds: float = 5.0, + ) -> None: + self.queue = queue + self.claim = claim + self.slots = list(slots) + self.lease_seconds = lease_seconds + self.notice = notice + self.renew_seconds = lease_seconds / 10 + self.poll_seconds = min(poll_seconds, self.renew_seconds) + self.interrupted = threading.Event() + 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, daemon=True), + threading.Thread(target=self._renew, daemon=True), + ] + + def __enter__(self) -> JobWatch: + for thread in self._threads: + thread.start() + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + traceback: TracebackType | None, + ) -> 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. + + 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() + for thread in self._threads: + if thread.is_alive(): + thread.join(timeout) + + def _watch(self) -> None: + while not self._stopped.wait(self.poll_seconds): + if check_notice(self.notice): + self.interrupted.set() + self.stop_event.set() + return + if time.time() >= self.claim.body["expires"]: + self._lose_claim() + return + + def _renew(self) -> None: + 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]]) -> 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 + ended, 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 + + +@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( + queue: Queue, + job: dict[str, Any], + claim: Lease, + slots: list[Lease], + workdir: Path, + runner: Callable[..., int], + settle: Callable[[Path], list[str]], + lease_seconds: float, + notice: Callable[[], bool], +) -> Outcome: + """Run one attempt of a claimed job and record its outcome. + + 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. + 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. + + 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. + notice: Function that returns whether an interruption notice exists. + + Returns: + 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" + + 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}" + for earlier in range(1, attempt) + if queue.download_tree( + job_dir / f"attempt-{earlier}", "results", job["id"], f"attempt-{earlier}" + ) + ] + errors = [error for earlier in resume for error in safe_settle(earlier)] + code = 1 + 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() + queue.end_lease(claim, "interrupted") + return Outcome("interrupted") + if code and not errors: + # A failed or stopped run can leave queries running. + errors = safe_settle(output) + upload_attempt() + + 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 abandon() + report_error = None + try: + report(job_dir) + except Exception as exc: + # 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 abandon() + 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 abandon() if recorded is None else Outcome("done", tuple(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, + notice: Callable[[], bool] | None = None, + 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: 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. 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; 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: - 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. + RuntimeError: If an attempt stopped by a lost claim could not be quiesced; + the worker's host is then kept. """ - 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) + 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()) @@ -558,59 +1124,60 @@ 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 check_notice(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(), + outcome = run_attempt( + queue, + job, + claim, + slots, + workdir, + runner, + settle, + lease_seconds, + notice, ) - completed.append(job["id"]) - if settle_errors: + 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: # Unconfirmed queries could overlap later trials on this host. return completed started = True break finally: for slot in slots: - queue.release(slot) + queue.end_lease(slot, "released") if not started: time.sleep(poll_seconds) @@ -639,55 +1206,126 @@ 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 | None = None +) -> 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, or None. + + Returns: + True if a notice exists; False also when instance metadata is unavailable. + """ + if simulated is not None and simulated.exists(): + return True + try: + (metadata or instance_metadata)("spot/instance-action") + except METADATA_ERRORS: + 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 METADATA_ERRORS: + 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() - 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(), + kept = keep_host(metadata) + 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 - 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 ) @@ -695,13 +1333,16 @@ 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 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. """ if not queue.exists("queue.json"): raise ValueError("Queue is missing or incompletely published") @@ -710,12 +1351,27 @@ 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() + + 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["state"] != "released": + leases[name] = { + **{k: lease.body.get(k) for k in ("host", "job", "attempt", "state")}, + "expired": lease.body["state"] == "running" and lease.takeable(at), + } + 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 255b2ed10..ba2b0fabb 100644 --- a/benchmarks/pyathena_bench/report.py +++ b/benchmarks/pyathena_bench/report.py @@ -75,6 +75,8 @@ def summarize(trials: list[dict[str, Any]]) -> list[dict[str, Any]]: reasons.add(f"Output files not listed: {trial['output_files_error']}") 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: @@ -90,7 +92,10 @@ 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 + ), "failed_warmups": sum( t.get("warmup", False) and t["status"] != "ok" for t in group ), @@ -155,14 +160,101 @@ def read_records(path: Path) -> list[dict[str, Any]]: return records -def report(directory: Path) -> None: +# 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") - 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"} - ) + 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 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. + + 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_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_trial_starts(source) + if event["trial"] not in recorded + ) + succeeded = {trial_key(t) for t in trials if t["status"] == "ok" and not t.get("warmup")} + merged = [] + counted: set[tuple[str, str, int | None]] = set() + for trial in trials: + if trial["status"] == "ok" and not trial.get("warmup"): + if len(sources) > 1 and trial_key(trial) in counted: + continue + 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 + + +def report(directory: Path) -> None: + 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 bfd774737..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,12 +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 CHECKPOINT, read_trials, trial_key + +_logger = logging.getLogger(__name__) class Observer: @@ -375,9 +382,80 @@ 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 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. + """ + 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( - 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, + checkpoint: Callable[[str, dict[str, Any]], 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 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. + """ + completed = completed or set() resources = validate_manifest(settings, manifest) validate_inputs(settings, manifest, scales) output.mkdir(parents=True, exist_ok=False) @@ -385,9 +463,12 @@ 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: + if (scale, case.id, None) in completed: + continue append_json( output / "trials.jsonl", { @@ -399,9 +480,16 @@ 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): + # 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 - prefix = f"runs/{manifest['run_id']}/trials/{trial}/" payload = { "settings": asdict(settings), "case": asdict(case), @@ -411,7 +499,7 @@ def run( "repetition": repetition, "warmup": repetition < 0, "database": resources["ScratchDatabase"], - "output": f"s3://{resources['Bucket']}/{prefix}", + "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 @@ -419,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" @@ -432,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 0b5ccb6a7..c5cfc16d0 100644 --- a/benchmarks/tests/test_aws.py +++ b/benchmarks/tests/test_aws.py @@ -10,7 +10,15 @@ 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, + run_locations, + split_location, + stack_resources, + validate_inputs, + validate_manifest, +) from pyathena_bench.config import Settings, read_json RUN = "a" * 32 @@ -24,6 +32,7 @@ "SourceTable": "pypi_file_downloads", "WorkGroup": "pyathena", "AutoScalingGroup": "fleet-test", + "Results": "s3://temporary/", } @@ -284,3 +293,85 @@ 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) + + +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", + "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) + 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") + + +@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 e2478ff8b..e8e7cd96e 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,61 @@ 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}/*" + ) + else: + # 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" + 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, 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 --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 93fd583a5..eed359b33 100644 --- a/benchmarks/tests/test_fleet.py +++ b/benchmarks/tests/test_fleet.py @@ -5,8 +5,14 @@ # # SPDX-License-Identifier: MIT +import csv +import http.client import io import json +import signal +import threading +import time +import urllib.error from dataclasses import replace from pathlib import Path from types import SimpleNamespace @@ -18,13 +24,20 @@ from pyathena_bench.config import Settings from pyathena_bench.fleet import ( Filters, + JobWatch, + Lease, Queue, + SpotInterruptionError, + check_notice, expand_jobs, + interruption_notice, merge_jobs, quiesce, retire_host, + run_job, status, submit, + suspend_rebalance, trial_evidence, work, ) @@ -33,24 +46,41 @@ 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 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"])] @@ -90,6 +120,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"])) @@ -275,11 +314,10 @@ 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, checkpoint): runs.append(item["id"]) - out = workdir / "results" / 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( @@ -288,7 +326,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( @@ -302,9 +340,10 @@ 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 not queue.names("slots") + assert [path.parent.name for path in settled] == ["heavy"] + 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 assert list(summary["failed"]) == ["heavy"] @@ -312,10 +351,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 @@ -344,6 +383,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.""" @@ -362,19 +413,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")]) @@ -386,12 +440,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 @@ -445,9 +497,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, checkpoint): 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" @@ -467,7 +518,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): @@ -476,7 +527,7 @@ def sleep_and_release(seconds): order = [] - def runner(item, workdir): + def runner(item, workdir, *args): order.append(item["id"]) return 0 @@ -496,7 +547,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): @@ -512,12 +565,12 @@ 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"}) - assert not queue.names("slots") + queue.acquire_slots(3, 2, {"host": "h"}, 600) + assert not live_slots(queue) def test_overlapping_submissions_cannot_share_a_name(self, tmp_path): s3, queue = queue_with([job("a")], tmp_path) @@ -603,5 +656,560 @@ 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, **changes): + 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, + **changes, + } + + +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_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") + 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" + + @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_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_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"] + + @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 + + 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 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 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, + settle=no_queries, + ) # fmt: skip + assert ran == ["heavy"] + assert not live_slots(queue) + + 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, 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: + 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, lease_seconds=0.5, notice=notice.is_set, + ) # fmt: skip + assert raised.value.ran == [] + 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 = [] + calls = [] + writes = [] + original_put = s3.put_object + 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=runner, + settle=lambda output: settled.append(output.name) or [], + ) # fmt: skip + assert ran == ["heavy", "next"] + 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") + 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" + 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 "Superseded: interrupted" in row["notes"] + + 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) + queue.put(queue.key("results", "a", "attempt-1", "events.jsonl"), {}) + 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, 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) + + 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_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), + ) + + def no_instance_action(path): + raise urllib.error.HTTPError(path, 404, "Not Found", None, None) + + 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_a_lost_claim_stops_the_attempt_and_quiesces_it(self, tmp_path): + _, queue = queue_with([job("a")], tmp_path) + stopped = [] + settled = [] + + def runner(item, workdir, out, resume, stop, checkpoint): + # Another worker took over the claim, for example after this host stalled. + 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 [], + lease_seconds=0.5) # fmt: skip + assert stopped == [True] + assert settled == ["attempt-1"] + assert ran == [] + assert not queue.names("done") + + def test_a_former_owner_records_nothing_after_a_takeover(self, tmp_path): + s3, queue = queue_with([job("a")], tmp_path) + + 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 + + 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): + _, queue = queue_with([job("heavy", heavy=True)], tmp_path) + + def runner_slot_only(item, workdir, out, resume, stop, checkpoint): + # Another worker took over the slot, for example after this host stalled. + 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_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 queue.read_lease(queue.key("slots", "0")).body["host"] == "other" + + def test_an_unconfirmed_stopped_attempt_keeps_the_host(self, tmp_path): + _, queue = queue_with([job("a")], tmp_path) + + 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"], lease_seconds=0.5) # fmt: skip + assert not queue.names("done") + + def test_cleanup_errors_survive_a_claim_lost_during_cleanup(self, tmp_path): + _, queue = queue_with([job("a"), job("b")], tmp_path) + calls = [] + + 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_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 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: + 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 = accepted_but_lost + 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): + config = tmp_path / "config.toml" + config.write_text("[scales]\nsmall = 10000\n") + queues = [] + + def stack_resources(*args): + pytest.fail("A worker with --results needs no stack outputs") + + 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 4a64b0ea1..39a2d66fc 100644 --- a/benchmarks/tests/test_report.py +++ b/benchmarks/tests/test_report.py @@ -137,3 +137,90 @@ def test_report_shows_the_most_output_files_and_the_largest_file(self, tmp_path) # A warmup is not in the columns, so its listing failure is not noted either. 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_counts_a_repetition_recorded_by_two_attempts_once(self, 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 09e015d79..69616f862 100644 --- a/benchmarks/tests/test_runner.py +++ b/benchmarks/tests/test_runner.py @@ -15,7 +15,7 @@ import pytest from pyathena_bench.cases import Case from pyathena_bench.config import Settings -from pyathena_bench.runner import Observer, output_files, run, supervise +from pyathena_bench.runner import Observer, completed_trials, output_files, run, supervise def successful_worker(pipe, payload): @@ -55,7 +55,7 @@ def run_stubs(monkeypatch): s3 = object() 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: {}) @@ -243,3 +243,94 @@ def summarize(s3, location): s3 is run_stubs and location.startswith(f"s3://out/runs/{'a' * 32}/trials/") for s3, location in calls ) + + +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) + + @pytest.mark.parametrize( + ("finished", "expected"), [([0], [-1, 1, 2]), ([0, 1, 2], [])], ids=["missing", "complete"] + ) + 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"), + ]