Add two-worker smoke test
This commit is contained in:
Emil
Executable
+165
@@ -0,0 +1,165 @@
|
||||
#!/usr/bin/env bash
|
||||
#
|
||||
# End-to-end check: two Python workers process separate coordinator shards.
|
||||
#
|
||||
# Requires Docker, curl, python3, and an installed scimesh-worker (normally
|
||||
# from this repository's .venv). The test uses its own Compose project, ports,
|
||||
# volumes, and temporary worker directories, leaving a developer stack alone.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
|
||||
COORDINATOR_DIR="$ROOT_DIR/coordinator"
|
||||
COMPOSE_PROJECT=${COMPOSE_PROJECT:-scimesh-two-worker-smoke}
|
||||
COORDINATOR_PORT=${COORDINATOR_PORT:-18081}
|
||||
POSTGRES_PORT=${POSTGRES_PORT:-55434}
|
||||
HOST="http://127.0.0.1:${COORDINATOR_PORT}"
|
||||
TOKEN=${SCIMESH_SMOKE_TOKEN:-two-worker-smoke-token}
|
||||
WORK_DIR=$(mktemp -d "${TMPDIR:-/tmp}/scimesh-two-worker-smoke.XXXXXX")
|
||||
WORKER_ONE_PID=""
|
||||
WORKER_TWO_PID=""
|
||||
WORKER_PYTHON=${SCIMESH_WORKER_PYTHON:-"$ROOT_DIR/.venv/bin/python"}
|
||||
|
||||
cleanup() {
|
||||
local exit_code=$?
|
||||
if [[ "$exit_code" -ne 0 ]]; then
|
||||
printf '\nTwo-worker smoke failed; worker logs follow.\n' >&2
|
||||
sed -n '1,200p' "$WORK_DIR/worker-a.log" >&2 || true
|
||||
sed -n '1,200p' "$WORK_DIR/worker-b.log" >&2 || true
|
||||
(
|
||||
cd "$COORDINATOR_DIR"
|
||||
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
|
||||
docker compose -p "$COMPOSE_PROJECT" logs coordinator >&2 || true
|
||||
)
|
||||
fi
|
||||
if [[ -n "$WORKER_ONE_PID" ]]; then kill "$WORKER_ONE_PID" 2>/dev/null || true; fi
|
||||
if [[ -n "$WORKER_TWO_PID" ]]; then kill "$WORKER_TWO_PID" 2>/dev/null || true; fi
|
||||
if [[ -n "$WORKER_ONE_PID" ]]; then wait "$WORKER_ONE_PID" 2>/dev/null || true; fi
|
||||
if [[ -n "$WORKER_TWO_PID" ]]; then wait "$WORKER_TWO_PID" 2>/dev/null || true; fi
|
||||
(
|
||||
cd "$COORDINATOR_DIR"
|
||||
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
|
||||
docker compose -p "$COMPOSE_PROJECT" down -v --remove-orphans >/dev/null 2>&1 || true
|
||||
)
|
||||
rm -rf "$WORK_DIR"
|
||||
exit "$exit_code"
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
require() {
|
||||
command -v "$1" >/dev/null || {
|
||||
printf 'missing required command: %s\n' "$1" >&2
|
||||
exit 2
|
||||
}
|
||||
}
|
||||
|
||||
for command in docker curl python3; do require "$command"; done
|
||||
[[ -x "$WORKER_PYTHON" ]] || {
|
||||
printf 'worker Python is not executable: %s\n' "$WORKER_PYTHON" >&2
|
||||
printf 'Set SCIMESH_WORKER_PYTHON to a Python environment with SciMesh and RDKit.\n' >&2
|
||||
exit 2
|
||||
}
|
||||
|
||||
printf 'Starting isolated coordinator on %s (project %s)\n' "$HOST" "$COMPOSE_PROJECT"
|
||||
(
|
||||
cd "$COORDINATOR_DIR"
|
||||
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
|
||||
WORKER_AUTH_TOKEN="$TOKEN" UI_AUTH_TOKEN= \
|
||||
docker compose -p "$COMPOSE_PROJECT" up -d --build
|
||||
)
|
||||
|
||||
for _ in $(seq 1 45); do
|
||||
if curl -fsS "$HOST/health" >/dev/null; then break; fi
|
||||
sleep 1
|
||||
done
|
||||
curl -fsS "$HOST/health" >/dev/null || {
|
||||
printf 'coordinator did not become healthy\n' >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
start_worker() {
|
||||
local worker_name=$1
|
||||
local worker_dir=$2
|
||||
SCIMESH_COORDINATOR_URL="$HOST" \
|
||||
SCIMESH_BEARER_TOKEN="$TOKEN" \
|
||||
SCIMESH_WORKER_NAME="$worker_name" \
|
||||
SCIMESH_POLL_INTERVAL=0.2 \
|
||||
"$WORKER_PYTHON" -m scimesh.worker.cli --work-dir "$worker_dir" >"$worker_dir.log" 2>&1 &
|
||||
STARTED_WORKER_PID=$!
|
||||
}
|
||||
|
||||
start_worker two-worker-smoke-a "$WORK_DIR/worker-a"
|
||||
WORKER_ONE_PID=$STARTED_WORKER_PID
|
||||
start_worker two-worker-smoke-b "$WORK_DIR/worker-b"
|
||||
WORKER_TWO_PID=$STARTED_WORKER_PID
|
||||
|
||||
for _ in $(seq 1 30); do
|
||||
registered=$(docker compose -p "$COMPOSE_PROJECT" -f "$COORDINATOR_DIR/docker-compose.yml" \
|
||||
exec -T postgres psql -U scimesh -d scimesh -Atc "SELECT count(*) FROM workers")
|
||||
[[ "$registered" == "2" ]] && break
|
||||
sleep 1
|
||||
done
|
||||
[[ "${registered:-0}" == "2" ]] || {
|
||||
printf 'workers did not register; logs follow\n' >&2
|
||||
sed -n '1,160p' "$WORK_DIR/worker-a.log" >&2 || true
|
||||
sed -n '1,160p' "$WORK_DIR/worker-b.log" >&2 || true
|
||||
exit 1
|
||||
}
|
||||
|
||||
DATASET="$WORK_DIR/fixture.tsv"
|
||||
printf '%s\n' \
|
||||
$'chembl_id\tcanonical_smiles' \
|
||||
$'TEST001\tCC' $'TEST002\tCCC' $'TEST003\tCCCC' $'TEST004\tCCCO' $'TEST005\tCCN' \
|
||||
$'TEST006\tCCCl' $'TEST007\tCCBr' $'TEST008\tCCF' $'TEST009\tCC=O' $'TEST010\tCC#N' \
|
||||
$'TEST011\tCO' $'TEST012\tCOC' $'TEST013\tCOCC' $'TEST014\tCN' $'TEST015\tCNC' \
|
||||
$'TEST016\tO=C=O' $'TEST017\tC1CC1' $'TEST018\tc1ccccc1' $'TEST019\tCC(C)O' $'TEST020\tCC(C)N' \
|
||||
>"$DATASET"
|
||||
|
||||
response=$(curl -fsS -H "Authorization: Bearer $TOKEN" -X POST "$HOST/jobs/upload" \
|
||||
-F 'workload=similarity-search' \
|
||||
-F 'parameters={"query_smiles":"CCO","top_k":5,"progress_every":0}' \
|
||||
-F 'chunk_rows=5' \
|
||||
-F 'max_rows=20' \
|
||||
-F "file=@${DATASET};type=text/tab-separated-values")
|
||||
job_id=$(printf '%s' "$response" | python3 -c 'import json,sys; print(json.load(sys.stdin)["job_id"])')
|
||||
task_count=$(printf '%s' "$response" | python3 -c 'import json,sys; print(json.load(sys.stdin)["task_count"])')
|
||||
[[ "$task_count" == "4" ]] || {
|
||||
printf 'expected four shards, got response: %s\n' "$response" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
printf 'Submitted job %s with four shards\n' "$job_id"
|
||||
for _ in $(seq 1 90); do
|
||||
job=$(curl -fsS -H "Authorization: Bearer $TOKEN" "$HOST/jobs/$job_id")
|
||||
status=$(printf '%s' "$job" | python3 -c 'import json,sys; print(json.load(sys.stdin)["status"])')
|
||||
[[ "$status" == "completed" || "$status" == "failed" || "$status" == "cancelled" ]] && break
|
||||
sleep 1
|
||||
done
|
||||
|
||||
printf '%s' "$job" | python3 -c '
|
||||
import json, sys
|
||||
job = json.load(sys.stdin)
|
||||
assert job["status"] == "completed", job
|
||||
assert job["total"] == 4, job
|
||||
assert job["completed"] == 4, job
|
||||
assert job["failed"] == 0, job
|
||||
'
|
||||
|
||||
task_check=$(docker compose -p "$COMPOSE_PROJECT" -f "$COORDINATOR_DIR/docker-compose.yml" \
|
||||
exec -T postgres psql -U scimesh -d scimesh -Atc \
|
||||
"SELECT count(*) FROM tasks WHERE job_id = '$job_id'::uuid AND status = 'completed' AND attempt = 1 AND result_artifact_id IS NOT NULL")
|
||||
[[ "$task_check" == "4" ]] || {
|
||||
printf 'expected four first-attempt tasks with coordinator artifacts, got %s\n' "$task_check" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
worker_one_results=$(find "$WORK_DIR/worker-a" -name result.csv -type f | wc -l | tr -d ' ')
|
||||
worker_two_results=$(find "$WORK_DIR/worker-b" -name result.csv -type f | wc -l | tr -d ' ')
|
||||
[[ "$worker_one_results" -ge 1 && "$worker_two_results" -ge 1 ]] || {
|
||||
printf 'both workers must process at least one shard (a=%s, b=%s)\n' \
|
||||
"$worker_one_results" "$worker_two_results" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
printf 'PASS: 4/4 shards completed; worker-a=%s, worker-b=%s\n' \
|
||||
"$worker_one_results" "$worker_two_results"
|
||||
Reference in New Issue
Block a user