166 lines
6.5 KiB
Bash
Executable File
166 lines
6.5 KiB
Bash
Executable File
#!/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" --max-tasks 2 >"$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" 2>/dev/null || printf '0')
|
|
[[ "$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"
|