Compare commits
84
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
8c5350abdf | ||
|
|
9fe81bc531 | ||
|
|
b281d5811a | ||
|
|
771952e22e | ||
|
|
565d4466e4 | ||
|
|
59e7fb0155 | ||
|
|
320f52615e | ||
|
|
f2977a990e | ||
|
|
5665e7df98 | ||
|
|
361eb2e344 | ||
|
|
749396da05 | ||
|
|
700a96a259 | ||
|
|
5d738e0a14 | ||
|
|
a18b8b8ae4 | ||
|
|
706bc85e17 | ||
|
|
9a8221163a | ||
|
|
644c287002 | ||
|
|
f20cc7fe00 | ||
|
|
284aef5d6f | ||
|
|
f059ac626c | ||
|
|
5c5a2af0a1 | ||
|
|
bc76f386e5 | ||
|
|
19fbb8e926 | ||
|
|
96169086f0 | ||
|
|
c43af32495 | ||
|
|
11e9333033 | ||
|
|
0a759a3f01 | ||
|
|
b9a975b0ea | ||
|
|
a201dd5ef9 | ||
|
|
dc75411907 | ||
|
|
6cdc115d60 | ||
|
|
fa76133efc | ||
|
|
7d8998408c | ||
|
|
e2a57175a0 | ||
|
|
2991ed202b | ||
|
|
9b235282fc | ||
|
|
3a1461315f | ||
|
|
172ff76fb8 | ||
|
|
6f14eeb32e | ||
|
|
87a483c2fb | ||
|
|
18d58cce84 | ||
|
|
dcabfcd0c3 | ||
|
|
e9cf6f0842 | ||
|
|
779ff8c10e | ||
|
|
e584cfc481 | ||
|
|
4ac19999a9 | ||
|
|
df4bdc9de9 | ||
|
|
9a458ec4ef | ||
|
|
49eb662798 | ||
|
|
5a9a10c681 | ||
|
|
f5ead0a450 | ||
|
|
c8c6455caf | ||
|
|
33f629f387 | ||
|
|
a7e949a0a7 | ||
|
|
163cbe14bf | ||
|
|
80ff72a0fe | ||
|
|
c6a66747eb | ||
|
|
0c1f5f06d4 | ||
|
|
67407220c3 | ||
|
|
1b1b971378 | ||
|
|
73196579e8 | ||
|
|
a3db1a1e67 | ||
|
|
16db1e41f7 | ||
|
|
ad9cc8f95c | ||
|
|
746958884e | ||
|
|
7ad28b939d | ||
|
|
1012f5d95a | ||
|
|
f5baec507c | ||
|
|
5a1414bee9 | ||
|
|
8b738efd5d | ||
|
|
d0aeb7fc95 | ||
|
|
a055473706 | ||
|
|
6e67daa9eb | ||
|
|
0f3a2d92d8 | ||
|
|
0bef7604fd | ||
|
|
6ef92908a1 | ||
|
|
f953112cfd | ||
|
|
19cbf7f113 | ||
|
|
9ec8f50313 | ||
|
|
08f5478a66 | ||
|
|
bde6cdb4ba | ||
|
|
43ceec1f77 | ||
|
|
f5b16b057f | ||
|
|
f8de0b2b9d |
@@ -60,7 +60,7 @@ jobs:
|
||||
run: go install -tags 'postgres' github.com/golang-migrate/migrate/v4/cmd/migrate@v4.17.1
|
||||
|
||||
- name: apply migrations
|
||||
run: migrate -path migrations -database "$TEST_DATABASE_URL" up
|
||||
run: migrate -path internal/storage/postgres/migrations -database "$TEST_DATABASE_URL" up
|
||||
|
||||
- name: integration tests
|
||||
run: go test -tags=integration ./internal/storage/postgres/ -v
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
name: python
|
||||
|
||||
on:
|
||||
push:
|
||||
paths:
|
||||
- "scimesh/**"
|
||||
- "tests/**"
|
||||
- "pyproject.toml"
|
||||
- ".github/workflows/python.yml"
|
||||
pull_request:
|
||||
paths:
|
||||
- "scimesh/**"
|
||||
- "tests/**"
|
||||
- "pyproject.toml"
|
||||
- ".github/workflows/python.yml"
|
||||
|
||||
jobs:
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.11"
|
||||
cache: pip
|
||||
- run: python -m pip install --upgrade pip
|
||||
- run: python -m pip install -e '.[dev]'
|
||||
- run: pytest -q
|
||||
@@ -0,0 +1,107 @@
|
||||
name: release
|
||||
|
||||
# Builds static coordinator and worker-agent binaries for every major
|
||||
# platform on a v* tag push, attaches them (plus SHA-256 checksums) to the
|
||||
# GitHub Release, and pushes the coordinator image to GHCR.
|
||||
#
|
||||
# git tag v1.0.0 && git push origin v1.0.0
|
||||
on:
|
||||
push:
|
||||
tags: ["v*"]
|
||||
|
||||
permissions:
|
||||
contents: write
|
||||
packages: write
|
||||
|
||||
jobs:
|
||||
binaries:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
fail-fast: false
|
||||
matrix:
|
||||
os: [linux, darwin, windows]
|
||||
arch: [amd64, arm64]
|
||||
defaults:
|
||||
run:
|
||||
working-directory: coordinator
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version-file: coordinator/go.mod
|
||||
cache-dependency-path: coordinator/go.sum
|
||||
|
||||
- name: vet
|
||||
run: go vet ./...
|
||||
|
||||
- name: build coordinator and worker-agent
|
||||
env:
|
||||
VERSION: ${{ github.ref_name }}
|
||||
run: |
|
||||
mkdir -p dist
|
||||
for cmd in coordinator worker-agent; do
|
||||
CGO_ENABLED=0 GOOS=${{ matrix.os }} GOARCH=${{ matrix.arch }} \
|
||||
go build -trimpath \
|
||||
-ldflags="-s -w -X main.version=${VERSION#v}" \
|
||||
-o "dist/${cmd}-${{ matrix.os }}-${{ matrix.arch }}${{ matrix.os == 'windows' && '.exe' || '' }}" \
|
||||
"./cmd/${cmd}"
|
||||
done
|
||||
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: binaries-${{ matrix.os }}-${{ matrix.arch }}
|
||||
path: coordinator/dist/*
|
||||
if-no-files-found: error
|
||||
|
||||
release:
|
||||
needs: binaries
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
path: artifacts
|
||||
# Only the binary artifacts: the image job also uploads a buildkit
|
||||
# cache artifact (*.dockerbuild) that download-artifact cannot fetch.
|
||||
pattern: binaries-*
|
||||
merge-multiple: true
|
||||
|
||||
- name: checksums
|
||||
working-directory: artifacts
|
||||
run: sha256sum * > SHA256SUMS.txt
|
||||
|
||||
- uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
files: artifacts/*
|
||||
generate_release_notes: true
|
||||
|
||||
image:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: docker/setup-buildx-action@v3
|
||||
|
||||
- uses: docker/login-action@v3
|
||||
with:
|
||||
registry: ghcr.io
|
||||
username: ${{ github.actor }}
|
||||
password: ${{ secrets.GITHUB_TOKEN }}
|
||||
|
||||
- uses: docker/metadata-action@v5
|
||||
id: meta
|
||||
with:
|
||||
images: ghcr.io/${{ github.repository }}/coordinator
|
||||
tags: |
|
||||
type=semver,pattern={{version}}
|
||||
type=semver,pattern={{major}}.{{minor}}
|
||||
type=raw,value=latest
|
||||
|
||||
- uses: docker/build-push-action@v6
|
||||
with:
|
||||
context: coordinator
|
||||
push: true
|
||||
tags: ${{ steps.meta.outputs.tags }}
|
||||
labels: ${{ steps.meta.outputs.labels }}
|
||||
build-args: |
|
||||
VERSION=${{ github.ref_name }}
|
||||
@@ -0,0 +1,66 @@
|
||||
name: users
|
||||
|
||||
on:
|
||||
push:
|
||||
paths:
|
||||
- "users/**"
|
||||
- ".github/workflows/users.yml"
|
||||
pull_request:
|
||||
paths:
|
||||
- "users/**"
|
||||
- ".github/workflows/users.yml"
|
||||
|
||||
defaults:
|
||||
run:
|
||||
working-directory: users
|
||||
|
||||
jobs:
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
|
||||
services:
|
||||
postgres:
|
||||
image: postgres:16-alpine
|
||||
env:
|
||||
POSTGRES_USER: scimesh
|
||||
POSTGRES_PASSWORD: scimesh
|
||||
POSTGRES_DB: scimesh_users
|
||||
ports:
|
||||
- 5432:5432
|
||||
options: >-
|
||||
--health-cmd "pg_isready -U scimesh"
|
||||
--health-interval 5s
|
||||
--health-timeout 3s
|
||||
--health-retries 10
|
||||
|
||||
env:
|
||||
TEST_DATABASE_URL: postgres://scimesh:scimesh@localhost:5432/scimesh_users?sslmode=disable
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- uses: actions/setup-go@v5
|
||||
with:
|
||||
go-version-file: users/go.mod
|
||||
cache-dependency-path: users/go.sum
|
||||
|
||||
- name: go vet
|
||||
run: go vet ./...
|
||||
|
||||
- name: gofmt
|
||||
run: test -z "$(gofmt -l .)" || (gofmt -l . && exit 1)
|
||||
|
||||
- name: unit tests (race)
|
||||
run: go test -race ./...
|
||||
|
||||
- name: lint
|
||||
run: go run github.com/golangci/golangci-lint/v2/cmd/golangci-lint@v2.12.2 run --build-tags=integration ./...
|
||||
|
||||
- name: install migrate CLI
|
||||
run: go install -tags 'postgres' github.com/golang-migrate/migrate/v4/cmd/migrate@v4.17.1
|
||||
|
||||
- name: apply migrations
|
||||
run: migrate -path migrations -database "$TEST_DATABASE_URL" up
|
||||
|
||||
- name: integration tests
|
||||
run: go test -tags=integration ./internal/storage/postgres/ -v
|
||||
@@ -11,3 +11,10 @@ results/
|
||||
*_similarities.csv
|
||||
test_results.csv
|
||||
test_structures/
|
||||
|
||||
# Local coordinator-worker execution state
|
||||
worker-data*/
|
||||
scimesh-worker-data/
|
||||
coordinator/.demo/
|
||||
site/
|
||||
coordinator/bin/
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
COMPLETED
|
||||
# Session Goal
|
||||
|
||||
давай теперь почистим проект от линего кода
|
||||
|
||||
## Plan
|
||||
|
||||
1. Аудит: ruff/pyflakes — неиспользуемые импорты по scimesh/ и tests/; grep — неиспользуемые функции/модули (после рефакторингов могли остаться мёртвые экспорты, например в descriptors/search/graph core и sdk/_validation).
|
||||
2. Удаление мёртвого кода: неиспользуемые импорты, функции, дубли (например write_descriptor_shards/concatenate_descriptor_shards, если вытеснены дефолтами batch), устаревшие файлы-обёртки.
|
||||
3. Проверка, что ничего публичного/API не сломано: pyright 0 ошибок, pytest зелёный, go test/vet, mkdocs build.
|
||||
4. Финал: полный прогон, COMPLETED.
|
||||
|
||||
## Progress
|
||||
|
||||
Эта сессия (доп. задача): Go-агент доведён до паритета, Python-демон удалён.
|
||||
- [x] Go-агент: token provider (static + worker-key exchange + 401 refresh на API/download/upload), CLEANUP_AFTER_SECONDS (очистка attempt-директорий), тесты auth (exchange/cache/reject/select/401-retry).
|
||||
- [x] Python: удалены daemon.py, cli.py, config.py, coordinator.py, artifacts.py, auth.py, transport.py; `scimesh/worker/` = только task.py + runners.py (SDK-мост, allowlist из env) + models.py (ClaimedTask/RunResult); консольный скрипт scimesh-worker убран из pyproject.
|
||||
- [x] Тесты: удалены test_worker_daemon.py, test_worker_auth.py; 208 pytest зелёные.
|
||||
- [x] Демо/смок переведены на Go-агент (demo-ui.sh: build_agent + env; two-worker-smoke.sh: AGENT_BIN + TASK_RUNNER_JSON). `make smoke-two-worker` PASS: 4/4 шардов, worker-a=2, worker-b=2.
|
||||
- [x] Документация: README, AGENTS.md, STATUS, handoff, mkdocs worker-integration/cli.
|
||||
- [x] Верификация: ruff clean; 208 pytest; pyright 0 (scimesh+tests); go test 11 пакетов + vet; mkdocs 0 warnings.
|
||||
@@ -5,8 +5,10 @@
|
||||
SciMesh is a Python package for molecular-similarity workloads. Source lives in
|
||||
`scimesh/`: `chemistry/` reads data and makes fingerprints, `workloads/`
|
||||
contains commands, and `core/` provides the workload protocol and registry.
|
||||
The worker daemon in `scimesh/worker/` is a coordinator client, not a database
|
||||
client. Tests are in `tests/`; specifications in `docs/`; roadmap: `PLAN.md`.
|
||||
The Go worker agent (`coordinator/internal/agent/`) is a coordinator client,
|
||||
not a database client; the Python side of a claimed task lives in
|
||||
`scimesh/worker/` (the per-task SDK execution entry). Tests are in `tests/`;
|
||||
specifications in `docs/`; roadmap: `PLAN.md`.
|
||||
|
||||
For distributed work, read `.agents/`, `docs/api-contract.md`,
|
||||
and `STATUS.md`. Use one CTX task per pull request; local workloads are the
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
.DEFAULT_GOAL := help
|
||||
|
||||
.PHONY: help agent coordinator demo-ui demo-down demo-logs smoke-two-worker docs docs-serve
|
||||
|
||||
help:
|
||||
@printf '%s\n' \
|
||||
'SciMesh developer commands:' \
|
||||
' make agent Build the Go worker agent (coordinator/bin/worker-agent).' \
|
||||
' make coordinator Build the coordinator server as a static binary' \
|
||||
' (coordinator/bin/coordinator).' \
|
||||
' make demo-ui Start the local UI pipeline demo with 2 Go worker agents.' \
|
||||
' make demo-ui WORKERS=3 Start the demo with 3 workers.' \
|
||||
' make demo-logs Follow coordinator logs for the demo.' \
|
||||
' make demo-down Stop demo containers and workers.' \
|
||||
' make smoke-two-worker E2E: two Go agents process 4 shards and the' \
|
||||
' result must match the local CLI reference.' \
|
||||
' make docs Build the MkDocs site into site/.' \
|
||||
' make docs-serve Serve the MkDocs site at http://localhost:8000.' \
|
||||
'' \
|
||||
'After make demo-ui: open http://localhost:18080/ui (operator / demo-ui-secret).'
|
||||
|
||||
# Convenient entry points from the repository root. Extra settings are passed
|
||||
# through, for example: make demo-ui WORKERS=3
|
||||
agent:
|
||||
$(MAKE) -C coordinator agent
|
||||
|
||||
coordinator:
|
||||
$(MAKE) -C coordinator coordinator
|
||||
|
||||
demo-ui:
|
||||
$(MAKE) -C coordinator demo-ui
|
||||
|
||||
demo-down:
|
||||
$(MAKE) -C coordinator demo-down
|
||||
|
||||
demo-logs:
|
||||
$(MAKE) -C coordinator demo-logs
|
||||
|
||||
smoke-two-worker:
|
||||
./scripts/two-worker-smoke.sh
|
||||
|
||||
docs:
|
||||
.venv/bin/mkdocs build
|
||||
|
||||
docs-serve:
|
||||
.venv/bin/mkdocs serve
|
||||
@@ -5,11 +5,13 @@
|
||||
> platform. It is intentionally detailed enough to split into independent task
|
||||
> briefs for developers or coding agents.
|
||||
>
|
||||
> **Planning baseline.** This branch starts from `Workers`: the Python package
|
||||
> has local `similarity-search` and `similarity-graph` workloads plus a Worker
|
||||
> Daemon client. The coordinator and PostgreSQL implementation do not yet
|
||||
> exist. The Worker contract and the Go/PostgreSQL design briefs in `docs/` are
|
||||
> part of this plan.
|
||||
> **Current planning baseline (2026-08-01).** The Go/PostgreSQL coordinator,
|
||||
> Python Worker Agent, versioned distributed-workload protocol, artifact-backed
|
||||
> task lifecycle, reducer orchestration, operator UI, User Service, and
|
||||
> distributed `similarity-search` are implemented on `main`. The evidence-based
|
||||
> completion tracker is [`STATUS.md`](STATUS.md); this document defines the
|
||||
> remaining direction and dependencies. Earlier descriptions of a missing
|
||||
> coordinator are historical context, not current work.
|
||||
|
||||
---
|
||||
|
||||
@@ -51,7 +53,7 @@ Coordinator reducer -> final artifact -> download/status API
|
||||
|
||||
### 2.1 In scope
|
||||
|
||||
- Go 1.22+ coordinator service with PostgreSQL 15+;
|
||||
- Go 1.25+ coordinator service with PostgreSQL 15+;
|
||||
- Python Worker Daemon running existing SciMesh workloads locally;
|
||||
- durable job, task, worker, and artifact metadata;
|
||||
- local coordinator-managed artifact storage for the first deployment;
|
||||
@@ -66,7 +68,8 @@ Coordinator reducer -> final artifact -> download/status API
|
||||
|
||||
- cloud object storage, Kubernetes, autoscaling, and multi-region operation;
|
||||
- arbitrary shell commands sent by coordinator to workers;
|
||||
- user accounts, multi-tenancy, billing, or sophisticated authorization;
|
||||
- billing and sophisticated multi-tenant administration beyond the implemented
|
||||
User Service and owner scoping;
|
||||
- GPU scheduling and multiprocessing inside a worker;
|
||||
- Docker as a required runtime dependency;
|
||||
- video/CV processing implementation;
|
||||
@@ -490,14 +493,84 @@ brute-force graph for both `greater` and `less` threshold directions.
|
||||
|
||||
### 7.3 Future workload policy
|
||||
|
||||
A new workload is accepted only when it supplies:
|
||||
A workload is more than a runner. It must define validation, planner and
|
||||
reducer behavior, worker allowlist/capabilities, input and output artifacts,
|
||||
UI/API parameters, reproducible execution environment, result verification,
|
||||
golden cross-worker fixtures, and applicable resource limits. The future public
|
||||
contract is described in [`docs/scimesh-sdk-roadmap.md`](docs/scimesh-sdk-roadmap.md).
|
||||
Its normative future interfaces and execution semantics are defined in the
|
||||
design-draft [`docs/scimesh-sdk-contract.md`](docs/scimesh-sdk-contract.md).
|
||||
|
||||
- an input/parameter validator;
|
||||
- an explicit sharding strategy;
|
||||
- bounded-memory task execution;
|
||||
- deterministic reduction semantics;
|
||||
- fixture-based local and distributed correctness tests;
|
||||
- a `describe()` payload for UI/API discovery.
|
||||
Every workload declaration must classify its task decomposition and input/output
|
||||
artifact shapes, determinism, reduction semantics, verifier mode, supported
|
||||
trust modes, CPU/memory/accelerator needs, and maximum output growth. The
|
||||
initial profiles are:
|
||||
|
||||
| Profile | Current acceptance policy |
|
||||
| --- | --- |
|
||||
| Byte-exact deterministic | Supported for untrusted quorum when whole artifacts have identical SHA-256. |
|
||||
| Canonical-exact deterministic | Deferred until the parser, schema, ordering, encoding, and serializer are versioned. |
|
||||
| Numeric deterministic with tolerance | Deferred until structured numeric comparison exists. |
|
||||
| Stochastic/search-based | Requires domain-specific evidence, repeated runs, or trusted execution. |
|
||||
| Trusted-only or domain-verified | May be planned only with an explicit trust policy and verifier. |
|
||||
|
||||
The current untrusted quorum records one vote per owner and accepts a task only
|
||||
when distinct owners upload artifacts with the same complete-file SHA-256. It
|
||||
therefore supports only the byte-exact profile (or a workload that first makes
|
||||
its output byte-identical through a specified canonicalization step). Reducers
|
||||
must fail safely rather than silently merge inconsistent partial results.
|
||||
|
||||
Before a workload is admitted to untrusted execution it needs a reproducibility
|
||||
gate: pinned environment/container digest and dependency versions; fixed locale,
|
||||
timezone, UTF-8/newline/CSV settings; explicit invalid-row and algorithm
|
||||
options; canonical representation and ordering; deterministic filenames/archive
|
||||
metadata; golden fixtures from two independently provisioned workers; local vs
|
||||
distributed parity; and retry/out-of-order completion tests. A loose dependency
|
||||
constraint is insufficient for byte-exact quorum.
|
||||
|
||||
Near-term critical path:
|
||||
|
||||
```text
|
||||
distributed similarity-graph
|
||||
-> reliability, security, and cross-language CI
|
||||
-> stable first release
|
||||
-> SDK foundation and descriptor-batch
|
||||
-> additional deterministic workloads
|
||||
```
|
||||
|
||||
Initial deterministic-workload backlog: `descriptor-batch` (the first SDK
|
||||
reference workload), molecule standardization, SMARTS screening, fingerprint
|
||||
export, fixed-template SMIRKS enumeration with strict caps, and reaction
|
||||
validation/descriptors. `similarity-graph` remains ahead of this backlog.
|
||||
Bounded combinatorial libraries and seeded conformers need specialized controls.
|
||||
ML, retrosynthesis, docking, QM, molecular dynamics, and GPU workloads are
|
||||
deferred until verifier/trust and reproducibility requirements are met.
|
||||
|
||||
### 7.4 Future verification, concurrency, and accelerators
|
||||
|
||||
Verification is a future versioned workload capability, not permanent
|
||||
whole-file-SHA logic. Planned modes are `ExactArtifactVerifier`,
|
||||
`CanonicalRecordVerifier`, `NumericToleranceVerifier`, `DomainSpecificVerifier`,
|
||||
and `TrustedWorkerPolicy`. Exact SHA-256 remains the first and safest mode;
|
||||
canonical and numeric modes must compare bounded structured data and publish
|
||||
sanitized evidence and failure reasons.
|
||||
|
||||
Worker concurrency remains **1** until implemented and tested. Its target model
|
||||
is one physical machine running one Worker Agent with `N` execution slots and
|
||||
one isolated subprocess per active Task, rather than one registered worker per
|
||||
CPU core. `max_concurrency` must be separate from `cpu_count`; each task keeps
|
||||
its own heartbeat, attempt directory, lease lifecycle, resource request, and
|
||||
graceful-drain behavior. CPU-bound scientific code should use processes and
|
||||
avoid nested oversubscription.
|
||||
|
||||
Accelerator support is also deferred. The coordinator matches generic resource
|
||||
requirements; the Worker Agent discovers and isolates devices (including
|
||||
`CUDA_VISIBLE_DEVICES`) and owns process/accounting lifecycle; the Python
|
||||
workload owns batching, memory strategy, deterministic output, and scientific
|
||||
validation; reducers/verifiers define CPU/GPU-independent semantics. CUDA and
|
||||
scientific kernels do not belong in the Go coordinator. GPU work follows stable
|
||||
CPU slices, generic resource requirements, pinned worker images, and tested
|
||||
CPU/GPU or domain-valid equivalence.
|
||||
|
||||
---
|
||||
|
||||
@@ -820,6 +893,138 @@ CTX-09 enables final result downloads.
|
||||
- failure/retry scenarios have automated coverage;
|
||||
- README contains architecture diagram, security caveat, and troubleshooting.
|
||||
|
||||
### CTX-13 — In-worker CPU parallelism
|
||||
|
||||
**Goal:** Allow a worker to use a bounded, configured number of CPU threads or
|
||||
processes while preserving the existing one-task-per-lease coordinator model.
|
||||
|
||||
**Depends on:** CTX-12.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- worker concurrency is an explicit configuration value with a safe default of
|
||||
one;
|
||||
- a task's internal parallel execution has bounded memory and does not build a
|
||||
dense N×N similarity matrix;
|
||||
- CPU-parallel `similarity-search` and `similarity-graph` outputs match the
|
||||
single-threaded local reference byte-for-byte where ordering is observable;
|
||||
- result ordering is deterministic across worker counts and block sizes;
|
||||
- cancellation, lease loss, and worker failure stop child work safely and do
|
||||
not report a successful result;
|
||||
- benchmarks and tests cover one-worker and multi-worker configurations.
|
||||
|
||||
### CTX-14 — GPU-accelerated workload execution
|
||||
|
||||
**Goal:** Add an optional GPU execution backend for supported molecular
|
||||
workloads, while retaining the validated CPU implementation as the reference
|
||||
and fallback.
|
||||
|
||||
**Depends on:** CTX-13.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- GPU capability and backend version are advertised explicitly by a worker;
|
||||
- the coordinator schedules GPU work only to compatible workers and CPU-only
|
||||
workers continue to claim CPU tasks;
|
||||
- unsupported hardware, unavailable drivers, and GPU execution errors produce
|
||||
sanitized failures or a documented CPU fallback;
|
||||
- GPU results match the CPU reference within a documented, tested numerical
|
||||
tolerance and preserve deterministic output ordering;
|
||||
- GPU memory use is bounded and no dense N×N similarity matrix is created;
|
||||
- CPU-only CI verifies backend selection and contract behavior, with GPU
|
||||
integration tests documented for compatible runners.
|
||||
|
||||
### CTX-15 — User Service and access control
|
||||
|
||||
**Goal:** Introduce a dedicated User Service for user identity and access
|
||||
control, without coupling workers to user credentials or moving scientific
|
||||
workload logic into the service.
|
||||
|
||||
**Depends on:** CTX-12.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- the service has a versioned, documented API in
|
||||
[`docs/user-service-api-contract.md`](docs/user-service-api-contract.md) and
|
||||
owns user identity data;
|
||||
- credentials and authentication tokens are stored and handled securely; they
|
||||
are never logged or exposed to workers;
|
||||
- authenticated identity is propagated to coordinator requests through an
|
||||
explicit, validated boundary;
|
||||
- authorization restricts access to jobs and artifacts to the intended user or
|
||||
project;
|
||||
- unauthenticated, expired-token, and cross-user access attempts have
|
||||
automated failure tests;
|
||||
- the existing single-operator demo remains usable through a documented local
|
||||
development configuration.
|
||||
|
||||
### CTX-16 — Workload SDK foundation
|
||||
|
||||
**Goal:** Provide a strict Python authoring SDK for installed, allowlisted
|
||||
scientific workloads while retaining the CTX-07 distributed protocol as a
|
||||
compatible wire profile.
|
||||
|
||||
**Depends on:** CTX-07 and CTX-08. Coordinator-backed generalized scheduling
|
||||
also depends on CTX-10 through CTX-14, but the Python contract and local
|
||||
conformance runtime can land independently and must fail closed for unavailable
|
||||
features.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- public manifest, workflow, task, artifact, resource, execution, provenance,
|
||||
and verifier value objects are immutable, typed, JSON-safe, versioned, and
|
||||
strict about unknown fields;
|
||||
- installed workload discovery requires an administrator allowlist plus exact
|
||||
workload version and package digest; job parameters cannot select code;
|
||||
- compatibility negotiation covers SDK/protocol/profile/feature/environment
|
||||
versions and occurs before planner invocation;
|
||||
- plans/tasks pin package and manifest digests plus selected trust mode, and
|
||||
quorum candidates carry coordinator-owned candidate/owner and scientific
|
||||
binding identities;
|
||||
- `core-batch-v1` has a trusted local conformance executor with atomic resource
|
||||
reservation, sealed-output/provenance validation, declared verifier
|
||||
invocation, and golden scientific parity;
|
||||
- exact, canonical-record, and structured numeric-tolerance verifier
|
||||
primitives return bounded sanitized decisions;
|
||||
- the existing distributed `similarity-search` is available through an adapter
|
||||
without changing its wire schema, worker alias boundary, or scientific
|
||||
result, and parity is tested;
|
||||
- advanced dynamic, stream, accelerator, gang, and side-effect profiles are
|
||||
rejected unless an enforcing runtime advertises their required features;
|
||||
- an author guide documents package entry points, security boundaries,
|
||||
conformance tests, and current coordinator/Worker limitations.
|
||||
|
||||
---
|
||||
|
||||
### CTX-17 — Self-provisioning coordinator and setup wizard
|
||||
|
||||
**Goal:** A downloaded coordinator binary should bring up a working platform
|
||||
with as little manual configuration as possible: it provisions its own
|
||||
schema, and a `setup` command walks the operator through the remaining
|
||||
environment (database creation, secrets, admin account).
|
||||
|
||||
**Depends on:** the Go coordinator and the release build pipeline. Step 1
|
||||
(embedded migrations, `AUTO_MIGRATE`) and step 2 (the `coordinator setup`
|
||||
wizard: database reachability and creation, schema migration, `.env` with a
|
||||
generated `JWT_SECRET`, readiness summary) are implemented; step 3 — a fully
|
||||
embedded userservice — is the remaining work.
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- the binary embeds the migrations and applies pending ones on startup by
|
||||
default (`AUTO_MIGRATE=false` opts out for managed databases); applying is
|
||||
idempotent and safe under concurrent starts;
|
||||
- `coordinator setup` (interactive, then non-interactive with `--yes`) checks
|
||||
database reachability, offers to create the role/database when credentials
|
||||
allow it, applies migrations, writes a `.env` with a generated `JWT_SECRET`
|
||||
and storage path, and prints exact next steps;
|
||||
- `coordinator --version` and the setup output agree on the release build;
|
||||
- the wizard explains what it cannot do itself: running PostgreSQL and the
|
||||
userservice, with concrete commands (docker compose, systemd) to finish;
|
||||
- the Docker image keeps working without the separate migrate step, and the
|
||||
release workflow publishes the binaries that support `setup`;
|
||||
- setup fails closed on non-interactive input and never logs secrets.
|
||||
|
||||
---
|
||||
|
||||
## 10. Suggested assignment bundles
|
||||
@@ -835,6 +1040,7 @@ parallel unless one engineer owns integration.
|
||||
| Distributed computation | CTX-07, CTX-08, CTX-10 | Scientific Python engineer |
|
||||
| Product surface | CTX-09, CTX-11 | Full-stack/backend engineer |
|
||||
| Quality gate | CTX-12 | DevOps/QA engineer |
|
||||
| Workload SDK | CTX-16 | Scientific Python/platform engineer |
|
||||
|
||||
Suggested order for a small team:
|
||||
|
||||
@@ -948,15 +1154,24 @@ Before merging a task, reviewer checks:
|
||||
Do not start these before CTX-12 is accepted.
|
||||
|
||||
- Replace local artifact storage with S3/MinIO behind an `ArtifactStore` API.
|
||||
- Add worker labels/capacity-aware scheduling and concurrency > 1.
|
||||
- Add worker labels and capacity-aware scheduling.
|
||||
- Implement CTX-13 for bounded in-worker CPU parallelism.
|
||||
- Implement CTX-14 for optional GPU-accelerated workload execution.
|
||||
- Implement CTX-15 for the User Service and authenticated user/project access.
|
||||
- Add cancellation propagation to workers.
|
||||
- Add image outputs and final PDF reporting to job artifacts.
|
||||
- Add CV/video workloads using the same planner/runner/reducer contract.
|
||||
- Add observability export (Prometheus/OpenTelemetry).
|
||||
- Add per-user/project authorization and signed artifact URLs.
|
||||
- Add signed artifact URLs.
|
||||
- Add shard caching and content-addressed input deduplication.
|
||||
- Add job priority and fair scheduling.
|
||||
- Add a CLI for submitting and monitoring remote jobs.
|
||||
- Implement CTX-17 step 3: a fully embedded userservice
|
||||
(`coordinator userservice` subcommand) so one binary can serve the whole
|
||||
platform without containers.
|
||||
- Replace PostgreSQL with an embedded SQLite backend for fully self-contained
|
||||
single-binary deployments (large storage-layer change; postgres row locks,
|
||||
transactions, and integration tests must be re-derived).
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
# SciMesh
|
||||
|
||||
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
|
||||
currently runs exact similarity search and sparse similarity-graph construction
|
||||
locally in one Python process; it creates no dense similarity matrix. A Python
|
||||
Worker client and the planned Go/PostgreSQL coordinator contract are tracked in
|
||||
the repository, but distributed execution is not available yet; see
|
||||
runs exact similarity search and sparse similarity-graph construction locally in
|
||||
one Python process; it creates no dense similarity matrix. The Go/PostgreSQL
|
||||
coordinator and Go worker agents (which execute SDK workloads in a Python
|
||||
subprocess) can run a shard-based `similarity-search`
|
||||
pipeline locally. After every shard succeeds, the coordinator deterministically
|
||||
merges its candidates into one final global top-k CSV. See
|
||||
[`STATUS.md`](STATUS.md).
|
||||
|
||||
The ChEMBL TSV database is intentionally not included in this repository. Download it separately and pass its path to the commands below. The expected columns are `chembl_id` and `canonical_smiles`.
|
||||
@@ -43,6 +45,39 @@ scimesh similarity-search --help
|
||||
scimesh similarity-graph --help
|
||||
```
|
||||
|
||||
## Manual pipeline demo
|
||||
|
||||
To inspect the coordinator, Web UI, and distributed `similarity-search`
|
||||
pipeline by hand, install development dependencies once and start the isolated
|
||||
demo from the repository root:
|
||||
|
||||
```bash
|
||||
python3 -m venv .venv
|
||||
.venv/bin/pip install -e '.[dev]'
|
||||
make demo-ui
|
||||
```
|
||||
|
||||
The MkDocs documentation site is served inside the UI at `/ui/docs/`
|
||||
(`make docs` builds it from `mkdocs/`; the demo mounts `site/`
|
||||
automatically, or set `SCIMESH_DOCS_DIR` for a manual coordinator). The site
|
||||
covers the complete Workload SDK: guides (`mkdocs/sdk/`), the full
|
||||
auto-generated API reference for `scimesh.sdk` (`mkdocs/api/`), and the
|
||||
documentation rules the site is written by (`mkdocs/approach.md`).
|
||||
|
||||
Open `http://localhost:18080/ui` and sign in with username `operator` and
|
||||
password `demo-ui-secret`. The command starts PostgreSQL, the coordinator, and
|
||||
two Go worker agents (built by `make agent`; each executes the SDK workload
|
||||
in a Python subprocess). Upload a small ChEMBL TSV, then use the job page
|
||||
to follow shard progress, inspect bounded **Preview CSV** results, and see a
|
||||
live processing-speed chart in shards per minute. The **Workloads** page shows
|
||||
the installed SDK workload library (descriptions, parameters, and artifact
|
||||
schemas) from the embedded catalog; regenerate it with
|
||||
`make workloads-export` (or `scimesh workload export`) whenever workloads
|
||||
change. To change the worker count, run `make demo-ui WORKERS=3`; stop
|
||||
everything with `make demo-down`.
|
||||
|
||||
Run `make help` to display these commands in the terminal.
|
||||
|
||||
## Similarity search
|
||||
|
||||
`similarity-search` finds the top-k molecules most similar to a query. The query is supplied either by ChEMBL ID or by SMILES. It uses Morgan fingerprints with `radius=2` and `fpSize=2048`, Tanimoto similarity, streaming TSV reads, and a bounded heap. Invalid SMILES and the query molecule are skipped.
|
||||
@@ -110,3 +145,40 @@ pytest
|
||||
```
|
||||
|
||||
The package separates common dataset parsing and fingerprints from independent workloads. Add future workloads through the workload registry without changing the main CLI.
|
||||
|
||||
## Workload SDK
|
||||
|
||||
`scimesh.sdk` is the framework only: strict and immutable workload manifests,
|
||||
typed artifact ports, static map/reduce plans, resource eligibility and local
|
||||
reservations, exact/canonical/numeric verifier primitives, installed-package
|
||||
allowlisting, and a local conformance executor. It contains no scientific
|
||||
workload code. Workloads are user scripts built on the SDK: the built-in
|
||||
`similarity-search`, `similarity-graph`, and `descriptor-batch` live in
|
||||
`scimesh/workloads/` (each a small package with `core.py` + `definition.py`),
|
||||
composed by `scimesh/workloads/library.py` and registered through
|
||||
`scimesh.workloads` entry points. The Worker Agent executes those SDK-built
|
||||
workloads directly (see `scimesh/worker/runners.py`), so the same scientific
|
||||
handlers run locally, in conformance, and on claimed coordinator tasks.
|
||||
`scimesh workload list` and `scimesh workload run` run any SDK workload from
|
||||
the command line. See the
|
||||
[SDK author guide](docs/workload-sdk.md), [contract](docs/scimesh-sdk-contract.md),
|
||||
and [delivery roadmap](docs/scimesh-sdk-roadmap.md).
|
||||
|
||||
Dynamic workflows, real Worker concurrency, coordinator-backed GPU allocation,
|
||||
streaming, and gang execution remain fail-closed until their versioned runtime
|
||||
features are implemented; declaring those profiles does not silently enable
|
||||
them.
|
||||
|
||||
The included `LocalCoreBatchExecutor` is a trusted, single-threaded in-process
|
||||
conformance harness. It validates scientific parity, sealed outputs, provenance,
|
||||
and limits, but intentionally refuses profiles that claim network/process
|
||||
isolation, secrets, accelerators, gangs, checkpoints, or retries; those require
|
||||
the future enforcing Agent runtime.
|
||||
|
||||
## Team
|
||||
|
||||
- [Emil](https://github.com/emil28092005) — Project Lead
|
||||
- [Kristina](https://github.com/kristtma) — Tech Lead
|
||||
- [Veniamin](https://t.me/Veniamin_Kt) — Scientific Lead
|
||||
- [Arkhip](https://github.com/hIpa-ussr) — Programmer
|
||||
- [Reranchik](https://github.com/RERAN4K) — Programmer
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# SciMesh Status
|
||||
|
||||
**Updated:** 2026-07-23
|
||||
**Branch baseline:** `main` at `b4a89dd` (coordinator merge)
|
||||
**Updated:** 2026-08-01
|
||||
**Branch baseline:** `main`; this revision adds the Workload SDK foundation.
|
||||
|
||||
## Current state
|
||||
|
||||
@@ -18,8 +18,18 @@ the reference behaviour for future distributed execution:
|
||||
The Go coordinator and its PostgreSQL-backed task lifecycle are implemented:
|
||||
registration, atomic claiming, lease renewal, artifact storage, dataset
|
||||
chunking, result/failure reporting, and job progress. The Python worker now
|
||||
uses the live coordinator contract; its HTTP path was exercised against a real
|
||||
Docker PostgreSQL stack on 2026-07-23.
|
||||
uses the live coordinator contract. Completed similarity-search shard results
|
||||
are reduced once into a checksum-protected final CSV, which is downloadable
|
||||
through the coordinator. The full Go checks (including a fresh migration and
|
||||
real PostgreSQL smoke test) passed on 2026-07-24.
|
||||
|
||||
The User Service is merged into `main`. It owns user accounts, authentication,
|
||||
roles, and verified-contributor status; the coordinator scopes user jobs and
|
||||
worker operations to the authenticated owner. Its documented v1 contract is in
|
||||
[`docs/user-service-api-contract.md`](docs/user-service-api-contract.md).
|
||||
Users can create and revoke worker keys for self-service Worker Agent
|
||||
enrollment. Untrusted workers require quorum agreement from distinct owners on
|
||||
the complete result-artifact SHA-256 before a task is accepted.
|
||||
|
||||
## Milestone tracker
|
||||
|
||||
@@ -31,28 +41,49 @@ Docker PostgreSQL stack on 2026-07-23.
|
||||
| CTX-03 Transactional queue | Implemented | Real-PostgreSQL integration tests cover atomic claims and concurrency. |
|
||||
| CTX-04 Worker registry and HTTP API | Implemented | Registration, claim, heartbeat, result, failure, and status endpoints. |
|
||||
| CTX-05 Artifact storage | Implemented | Coordinator-owned inputs/results, checksum verification, and upload flow. |
|
||||
| CTX-06 Python Worker live-contract alignment | Implemented | Worker completed a real uploaded shard via HTTP on 2026-07-23. |
|
||||
| CTX-07 Distributed workload protocol | Not started | Depends on artifact and Worker contracts. |
|
||||
| CTX-08 Distributed similarity-search | Not started | Local reference exists. |
|
||||
| CTX-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. |
|
||||
| CTX-10 Distributed similarity-graph | Not started | Local reference exists. |
|
||||
| CTX-11 Dashboard/operator view | In progress | `feat/web-interface` adds a protected local view: job/task/worker status, dataset upload, diagnostic partial-artifact download, and polling. Final-result reduction remains CTX-09. |
|
||||
| CTX-06 Python Worker live-contract alignment | Superseded | The Python worker daemon was removed; the Go worker agent (`coordinator/internal/agent/` + `cmd/worker-agent`) now implements the lifecycle (register/claim/heartbeat/download/upload/submit/fail, token refresh, cleanup) and executes SDK workloads via the Python task entry `scimesh/worker/task.py`. E2E: `make smoke-two-worker` passes 4/4 shards with two agents. |
|
||||
| CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. |
|
||||
| CTX-08 Distributed similarity-search | Implemented | Python planner resolves `query_id` once, creates deterministic shard plans, worker adapter emits exact partial top-k CSVs/metrics, and reducer matches the local reference. |
|
||||
| CTX-09 Reducer and final-result API | Implemented | Atomic `reducing` claim, deterministic coordinator-side top-k reducer, sanitized reducer failure, final artifact persistence, `result_uri`, and final CSV download. |
|
||||
| CTX-10 Distributed similarity-graph | Not started | Local reference exists; the SDK-built local graph workload already enforces the pair-coverage invariant. |
|
||||
| CTX-11 Dashboard/operator view | Implemented | Protected live control room: MkDocs documentation served at `/ui/docs/` (SCIMESH_DOCS_DIR; the demo mounts `site/` automatically), recent-run/worker overview, real pipeline-stage visualization, shard attempts and safe failures, validated similarity-search upload, coordinator artifacts, final-result download, bounded polling, a Workload library page rendering the embedded catalog from `scimesh workload export` (`/ui/workloads`, regenerated via `make workloads-export`). |
|
||||
| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
|
||||
| CTX-15 User Service and access control | Implemented | User/owner scoping, verified contributors, worker keys, self-service enrollment, and quorum-backed untrusted workers are merged; local Go/Python and Docker/PostgreSQL checks passed. |
|
||||
| MkDocs documentation site | Implemented | A standalone documentation site (`mkdocs/`, `docs_dir: mkdocs`) covering the complete Workload SDK: guides (overview, authoring workloads, CLI, worker integration), the full auto-generated API reference for all 15 `scimesh.sdk` modules (mkdocstrings), and the writing rules (`mkdocs/approach.md`). Built with `make docs`, served inside the UI at `/ui/docs/`; the project's internal `docs/` directory is not part of the site. |
|
||||
| CTX-16 Workload SDK foundation | Implemented | `scimesh.sdk` provides strict immutable manifests/plans/artifacts, digest/trust-pinned tasks, typed DAGs, compatibility negotiation, verifier primitives with owner/binding-safe quorum inputs, resource eligibility/local allocation, measured package discovery, a trusted local core-batch conformance harness, and strict package discovery. Enforcing coordinator/Worker profiles remain fail-closed. |
|
||||
| SDK roadmap step 3: `descriptor-batch` | Implemented | The first SDK-built reference workload (`scimesh/workloads/descriptors/`): pinned 81-name RDKit 2D descriptor set, canonical one-row-per-input CSV, deterministic row-bounded shards, shard-index concatenation with one header, byte-identical local/distributed output, and a two-worker `untrusted_quorum` verifier test. |
|
||||
| SDK-built `similarity-search` and `similarity-graph` | Implemented | Both workloads are SDK-built packages (`scimesh/workloads/search/`, `scimesh/workloads/graph/`) built on the `MapReduceWorkload` authoring scaffold (`scimesh/sdk/batch.py`); they reuse the local scientific cores and are byte-identical to the single-process references (search; graph for both threshold directions and any block size). The graph reducer enforces the CTX-10 pair-coverage invariant. `scimesh/workloads/library.py` composes the built-in registry/runtime. |
|
||||
| SDK-built `molwt-filter` | Implemented | The minimal authoring example (`scimesh/workloads/molwt_filter/`): filters molecules by exact RDKit molecular weight with only one scientific hook, using the scaffold's new default sharding and concatenation hooks. Registered in the built-in library and as a `scimesh.workloads` entry point. |
|
||||
| SDK authoring scaffold | Implemented | `MapReduceWorkload` (exported from `scimesh.sdk`) assembles manifest, map/reduce stages, workflow, and digest-pinned handlers from three scientific hooks (partition/compute/merge), with overridable hooks for domain validation, plan-time resolution, custom task planning, and partial-key policy. The generic `scimesh workload list|run` CLI and the worker's allowlist-driven loading (`SCIMESH_WORKLOAD_ALLOWLIST`, `SCIMESH_CAPABILITIES`) let new workloads run without touching other code. |
|
||||
|
||||
## Next recommended assignment
|
||||
|
||||
Assign **CTX-07** to the workload role: define distributed job planning and
|
||||
reduction boundaries before implementing distributed search or graph execution.
|
||||
Assign **CTX-10** to the distributed-science role: implement deterministic
|
||||
block-pair planning and reduction for `similarity-graph`.
|
||||
|
||||
## Known constraints
|
||||
|
||||
- Planner/reducer semantics are not implemented; the operator UI labels
|
||||
`partial_result` files as diagnostic and cannot present them as final output.
|
||||
Use the local `scimesh` CLI for complete workload results.
|
||||
- The worker/coordinator flow currently accepts both underscore API workload
|
||||
names and hyphenated CLI names while the contract is consolidated.
|
||||
- A real-stack worker test uses a small `query_smiles` shard. Resolving a
|
||||
`query_id` once and sharing it across shards belongs to CTX-07.
|
||||
- The worker/coordinator flow accepts both underscore API workload names and
|
||||
hyphenated names at the runner boundary; the runner normalizes them.
|
||||
- The worker executes SDK-built workloads through `scimesh/worker/runners.py`
|
||||
(a workload-generic v1-wire bridge over `TaskSpec`/`LocalTaskContext`);
|
||||
`query_id` resolution and parameter validation live in the workload itself.
|
||||
`max_rows` is a plan-time option and is rejected per task by the stage
|
||||
projection.
|
||||
- A real-stack worker test uses a small `query_smiles` shard. The Python
|
||||
planner resolves `query_id` once and shares `query_smiles`; the upload UI
|
||||
currently accepts `query_smiles` only.
|
||||
- The coordinator accepts uploaded distributed jobs only for
|
||||
`similarity-search` with `query_smiles`. It rejects `similarity-graph` until
|
||||
CTX-10 supplies cross-shard pair planning.
|
||||
- The SDK can execute `core-batch-v1` locally, but the protocol-v1 coordinator
|
||||
still has flat single-input/single-result tasks and no package/resource
|
||||
leases. General DAG, concurrent-Agent, GPU, stream, and gang execution needs
|
||||
a versioned coordinator/Worker rollout; unsupported features fail before
|
||||
planner invocation.
|
||||
- The local SDK executor is intentionally trusted and in-process. It does not
|
||||
enforce process/network/timeout/credential isolation and rejects declarations
|
||||
that would require those guarantees.
|
||||
|
||||
## Update rule
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
# build fails with "the --mount option requires BuildKit".
|
||||
|
||||
# --- build stage ----------------------------------------------------------
|
||||
FROM golang:1.24-alpine AS build
|
||||
FROM golang:1.25-alpine AS build
|
||||
|
||||
WORKDIR /src
|
||||
|
||||
@@ -16,6 +16,9 @@ RUN --mount=type=cache,target=/go/pkg/mod go mod download
|
||||
|
||||
COPY . .
|
||||
|
||||
# Release builds inject the tag via build-arg; local builds stay "dev".
|
||||
ARG VERSION=dev
|
||||
|
||||
# The cache mounts persist the module cache and the compiler's build cache
|
||||
# *across* builds, so a rebuild after a code edit recompiles only what changed
|
||||
# instead of the whole dependency tree.
|
||||
@@ -25,7 +28,7 @@ COPY . .
|
||||
RUN --mount=type=cache,target=/go/pkg/mod \
|
||||
--mount=type=cache,target=/root/.cache/go-build \
|
||||
CGO_ENABLED=0 GOOS=linux go build \
|
||||
-trimpath -ldflags="-s -w" \
|
||||
-trimpath -ldflags="-s -w -X main.version=${VERSION#v}" \
|
||||
-o /out/coordinator ./cmd/coordinator
|
||||
|
||||
# --- runtime stage --------------------------------------------------------
|
||||
|
||||
+104
-3
@@ -1,4 +1,6 @@
|
||||
.PHONY: build run test test-integration vet lint tidy check migrate-up migrate-down up down down-clean logs ps rebuild psql smoke
|
||||
.DEFAULT_GOAL := help
|
||||
|
||||
.PHONY: help build run test test-integration vet lint tidy check migrate-up migrate-down up down down-clean logs ps rebuild psql smoke agent coordinator setup workloads-export demo-ui demo-down demo-reset demo-logs
|
||||
|
||||
# `check` deliberately uses its own Compose project and host ports. This keeps
|
||||
# it from connecting to or replacing a developer's local PostgreSQL instance.
|
||||
@@ -10,6 +12,105 @@ CHECK_TOKEN ?= dev-token
|
||||
CHECK_DATABASE_URL ?= postgres://scimesh:scimesh@localhost:$(CHECK_POSTGRES_PORT)/scimesh?sslmode=disable
|
||||
CHECK_COMPOSE = POSTGRES_PORT=$(CHECK_POSTGRES_PORT) COORDINATOR_PORT=$(CHECK_COORDINATOR_PORT) docker compose -p $(CHECK_PROJECT)
|
||||
|
||||
# --- local manual demo ---------------------------------------------------
|
||||
# A separate project and ports mean this demo cannot collide with the normal
|
||||
# `make up` stack or a developer's local PostgreSQL on 5432.
|
||||
DEMO_PROJECT ?= scimesh-demo
|
||||
DEMO_POSTGRES_PORT ?= 55432
|
||||
DEMO_COORDINATOR_PORT ?= 18080
|
||||
DEMO_UI_TOKEN ?= demo-ui-secret
|
||||
DEMO_WORKER_TOKEN ?= demo-worker-token
|
||||
DEMO_WORKERS ?= 2
|
||||
# Short public knob for `make demo-ui WORKERS=3`; DEMO_WORKERS remains useful
|
||||
# for scripts and backwards-compatible documentation.
|
||||
WORKERS ?= $(DEMO_WORKERS)
|
||||
DEMO_DIR ?= .demo
|
||||
|
||||
# workloads.json is the UI workload catalog, generated from the Python SDK
|
||||
# workload library. It is checked in so the binary embeds it; regenerate it
|
||||
# whenever workloads or their manifests change (requires the Python venv).
|
||||
WORKLOADS_JSON := internal/workloads/workloads.json
|
||||
|
||||
# Version injected into the binaries via -ldflags; falls back to "dev".
|
||||
VERSION ?= $(shell git describe --tags --always --dirty 2>/dev/null || echo dev)
|
||||
LDFLAGS := -s -w -X main.version=$(VERSION)
|
||||
|
||||
# The Go worker agent: a static coordinator client that executes SDK
|
||||
# workloads in a Python subprocess per claimed task.
|
||||
agent:
|
||||
CGO_ENABLED=0 go build -trimpath -ldflags="$(LDFLAGS)" -o bin/worker-agent ./cmd/worker-agent
|
||||
@printf '%s\n' 'Built bin/worker-agent. Configure via environment:' ' COORDINATOR_URL, WORKER_AUTH_TOKEN, WORK_DIR, CPU_COUNT, MEMORY_MB,' ' POLL_INTERVAL, REQUEST_TIMEOUT, HEARTBEAT_INTERVAL, CAPABILITIES,' ' TASK_RUNNER, MAX_TASKS, EXIT_WHEN_IDLE, WORKER_NAME, WORKER_ID'
|
||||
|
||||
# The coordinator server as a static binary, the same way the Docker image
|
||||
# builds it (CGO_ENABLED=0, trimmed). Requires PostgreSQL at runtime.
|
||||
coordinator:
|
||||
CGO_ENABLED=0 go build -trimpath -ldflags="$(LDFLAGS)" -o bin/coordinator ./cmd/coordinator
|
||||
@printf '%s\n' \
|
||||
'Built bin/coordinator. Configure via environment:' \
|
||||
' DATABASE_URL, COORDINATOR_ADDR, COORDINATOR_TOKEN, UI_AUTH_TOKEN,' \
|
||||
' COORDINATOR_STORAGE_DIR, SCIMESH_DOCS_DIR, JWT_SECRET, USERSERVICE_URL' \
|
||||
' (embedded schema migrations run on startup; AUTO_MIGRATE=false disables)'
|
||||
|
||||
# Interactive wizard: checks the database, creates it when missing (via
|
||||
# POSTGRES_ADMIN_URL or --admin-db), applies the embedded schema, generates a
|
||||
# JWT_SECRET, and writes a .env file. Non-interactive: SETUP_ARGS=--yes.
|
||||
setup: coordinator
|
||||
./bin/coordinator setup $(SETUP_ARGS)
|
||||
|
||||
workloads-export:
|
||||
cd .. && .venv/bin/scimesh workload export -o coordinator/$(WORKLOADS_JSON)
|
||||
|
||||
help:
|
||||
@printf '%s\n' \
|
||||
'SciMesh coordinator commands:' \
|
||||
' make up / make down Start or stop the normal coordinator stack.' \
|
||||
' make demo-ui [WORKERS=3] Start isolated UI demo services and local workers.' \
|
||||
' make demo-logs Follow coordinator logs for the UI demo.' \
|
||||
' make demo-down Stop the demo services and workers.' \
|
||||
' make demo-reset Stop the demo and wipe its data volumes.' \
|
||||
' make workloads-export Regenerate the embedded UI workload catalog.' \
|
||||
' make setup Interactive one-shot provisioning wizard.' \
|
||||
' make test / make vet Run Go verification.' \
|
||||
'' \
|
||||
'Demo UI: http://localhost:18080/ui (login page; admin root@scimesh.local / rootpassword).'
|
||||
|
||||
demo-ui:
|
||||
@DEMO_PROJECT="$(DEMO_PROJECT)" \
|
||||
DEMO_POSTGRES_PORT="$(DEMO_POSTGRES_PORT)" \
|
||||
DEMO_COORDINATOR_PORT="$(DEMO_COORDINATOR_PORT)" \
|
||||
DEMO_UI_TOKEN="$(DEMO_UI_TOKEN)" \
|
||||
DEMO_WORKER_TOKEN="$(DEMO_WORKER_TOKEN)" \
|
||||
DEMO_WORKERS="$(WORKERS)" \
|
||||
DEMO_DIR="$(DEMO_DIR)" \
|
||||
./scripts/demo-ui.sh start
|
||||
|
||||
demo-down:
|
||||
@DEMO_PROJECT="$(DEMO_PROJECT)" \
|
||||
DEMO_POSTGRES_PORT="$(DEMO_POSTGRES_PORT)" \
|
||||
DEMO_COORDINATOR_PORT="$(DEMO_COORDINATOR_PORT)" \
|
||||
DEMO_UI_TOKEN="$(DEMO_UI_TOKEN)" \
|
||||
DEMO_WORKER_TOKEN="$(DEMO_WORKER_TOKEN)" \
|
||||
DEMO_DIR="$(DEMO_DIR)" \
|
||||
./scripts/demo-ui.sh stop
|
||||
|
||||
demo-reset:
|
||||
@DEMO_PROJECT="$(DEMO_PROJECT)" \
|
||||
DEMO_POSTGRES_PORT="$(DEMO_POSTGRES_PORT)" \
|
||||
DEMO_COORDINATOR_PORT="$(DEMO_COORDINATOR_PORT)" \
|
||||
DEMO_UI_TOKEN="$(DEMO_UI_TOKEN)" \
|
||||
DEMO_WORKER_TOKEN="$(DEMO_WORKER_TOKEN)" \
|
||||
DEMO_DIR="$(DEMO_DIR)" \
|
||||
./scripts/demo-ui.sh reset
|
||||
|
||||
demo-logs:
|
||||
@DEMO_PROJECT="$(DEMO_PROJECT)" \
|
||||
DEMO_POSTGRES_PORT="$(DEMO_POSTGRES_PORT)" \
|
||||
DEMO_COORDINATOR_PORT="$(DEMO_COORDINATOR_PORT)" \
|
||||
DEMO_UI_TOKEN="$(DEMO_UI_TOKEN)" \
|
||||
DEMO_WORKER_TOKEN="$(DEMO_WORKER_TOKEN)" \
|
||||
DEMO_DIR="$(DEMO_DIR)" \
|
||||
./scripts/demo-ui.sh logs
|
||||
|
||||
# --- build / run ---------------------------------------------------------
|
||||
build:
|
||||
go build ./...
|
||||
@@ -62,10 +163,10 @@ tidy:
|
||||
# DATABASE_URL must be set, e.g.:
|
||||
# export DATABASE_URL='postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable'
|
||||
migrate-up:
|
||||
migrate -path migrations -database "$(DATABASE_URL)" up
|
||||
migrate -path internal/storage/postgres/migrations -database "$(DATABASE_URL)" up
|
||||
|
||||
migrate-down:
|
||||
migrate -path migrations -database "$(DATABASE_URL)" down 1
|
||||
migrate -path internal/storage/postgres/migrations -database "$(DATABASE_URL)" down 1
|
||||
|
||||
# --- docker --------------------------------------------------------------
|
||||
# `up` starts Postgres, applies migrations, then launches the coordinator.
|
||||
|
||||
+39
-3
@@ -76,9 +76,45 @@ UI_AUTH_TOKEN='local-ui-secret' make up
|
||||
```
|
||||
|
||||
The UI is disabled by default and never accepts the worker bearer token.
|
||||
It shows recent jobs, task/worker state, and the per-job partial artifacts.
|
||||
Those files are explicitly diagnostic until the CTX-09 reducer creates a final
|
||||
result; the UI does not present them as final scientific output.
|
||||
The **control room** shows live workers, recent runs, shard state/attempts,
|
||||
safe failures, coordinator artifacts, and the final CSV for completed
|
||||
similarity-search jobs. The job page follows the real stages: TSV accepted →
|
||||
shards execute → workers return CSVs → `reducing` → final deterministic global
|
||||
top-k result. It polls only its own coordinator read-model and never controls
|
||||
or exposes worker processes.
|
||||
|
||||
For a hands-on run, open `/ui`, choose **New similarity search**, select a
|
||||
small ChEMBL-style TSV, then leave one or more `scimesh-worker` processes
|
||||
running in separate terminals. The detail page updates every two seconds and
|
||||
stops polling after a completed, failed, or cancelled job. Use **Preview CSV**
|
||||
to inspect a bounded first page of a partial or completed final result before
|
||||
downloading it. The UI never exposes source datasets or shard inputs; partial
|
||||
CSVs remain available only as diagnostics.
|
||||
|
||||
### One-command manual demo
|
||||
|
||||
From the repository root, create the Python environment once, then start a
|
||||
self-contained UI demo with two local reference workers:
|
||||
|
||||
```sh
|
||||
python3 -m venv .venv
|
||||
.venv/bin/pip install -e '.[dev]'
|
||||
make demo-ui
|
||||
```
|
||||
|
||||
This uses a separate Docker project and ports `18080` (coordinator) and
|
||||
`55432` (PostgreSQL), so it does not conflict with the normal stack. Open
|
||||
`http://localhost:18080/ui`, use username `operator` and password
|
||||
`demo-ui-secret`, upload a small ChEMBL TSV, and observe the workers process
|
||||
it. Change the worker count with `make demo-ui WORKERS=3`; stop all demo
|
||||
services and workers with `make demo-down`.
|
||||
|
||||
The job page shows a live **Processing speed** graph in completed shards per
|
||||
minute. It uses the coordinator snapshots observed by the open browser tab, so
|
||||
it is a transparent local-session measurement rather than a persisted metric.
|
||||
Use **Preview CSV** before downloading a partial diagnostic or completed final
|
||||
result. Run `make help` from either the repository root or this directory for
|
||||
the full list of demo commands.
|
||||
|
||||
`up` starts three services in order: Postgres waits until `pg_isready` passes, a
|
||||
one-shot `migrate` container applies the schema and exits, and only then does the
|
||||
|
||||
@@ -2,6 +2,8 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/signal"
|
||||
@@ -9,13 +11,32 @@ import (
|
||||
"syscall"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/infra"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/metrics"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/storage/blob"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/storage/postgres"
|
||||
httptransport "github.com/emil28092005/SciMesh/coordinator/internal/transport/http"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/workloads"
|
||||
)
|
||||
|
||||
// version is injected at build time (-ldflags "-X main.version=...") and
|
||||
// reported by --version. "dev" marks a local build.
|
||||
var version = "dev"
|
||||
|
||||
func main() {
|
||||
args := os.Args[1:]
|
||||
if len(args) > 0 && args[0] == "setup" {
|
||||
if err := runSetup(args[1:]); err != nil {
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
}
|
||||
showVersion := flag.Bool("version", false, "print the build version and exit")
|
||||
flag.Parse()
|
||||
if *showVersion {
|
||||
fmt.Println("coordinator " + version)
|
||||
return
|
||||
}
|
||||
// All work happens in run() so its defers (pool.Close, log flush, signal
|
||||
// stop) still execute: os.Exit skips deferred calls entirely.
|
||||
if err := run(); err != nil {
|
||||
@@ -52,6 +73,15 @@ func run() error {
|
||||
}
|
||||
defer pool.Close()
|
||||
|
||||
// A downloaded binary provisions its own schema; AUTO_MIGRATE=false keeps
|
||||
// out-of-band migration workflows (the migrate CLI, CI, managed databases).
|
||||
if cfg.AutoMigrate {
|
||||
if err := postgres.Migrate(ctx, cfg.DatabaseURL, log); err != nil {
|
||||
log.Error("apply migrations", "err", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
blobStore, err := blob.NewFSStore(cfg.StorageDir)
|
||||
if err != nil {
|
||||
log.Error("init blob storage", "err", err)
|
||||
@@ -59,35 +89,45 @@ func run() error {
|
||||
}
|
||||
|
||||
var (
|
||||
clk = infra.NewClock()
|
||||
tx = postgres.NewTxManager(pool)
|
||||
taskRepo = postgres.NewTaskRepo(pool)
|
||||
jobRepo = postgres.NewJobRepo(pool)
|
||||
workerRepo = postgres.NewWorkerRepo(pool)
|
||||
artifactRepo = postgres.NewArtifactRepo(pool)
|
||||
uiReadRepo = postgres.NewUIReadRepo(pool)
|
||||
clk = infra.NewClock()
|
||||
tx = postgres.NewTxManager(pool)
|
||||
taskRepo = postgres.NewTaskRepo(pool)
|
||||
jobRepo = postgres.NewJobRepo(pool)
|
||||
workerRepo = postgres.NewWorkerRepo(pool)
|
||||
artifactRepo = postgres.NewArtifactRepo(pool)
|
||||
uiReadRepo = postgres.NewUIReadRepo(pool)
|
||||
taskResultRepo = postgres.NewTaskResultRepo(pool)
|
||||
)
|
||||
|
||||
catalog, err := workloads.Load()
|
||||
if err != nil {
|
||||
log.Error("load workload catalog", "err", err)
|
||||
return err
|
||||
}
|
||||
|
||||
useCases := httptransport.UseCases{
|
||||
RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk),
|
||||
CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk),
|
||||
ClaimTask: usecase.NewClaimTask(taskRepo, clk, cfg.LeaseDuration),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk, cfg.DefaultMaxAttempts, catalog),
|
||||
ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration, catalog),
|
||||
RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
||||
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk),
|
||||
FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk),
|
||||
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, workerRepo, taskResultRepo, tx, clk, cfg.QuorumSize, catalog),
|
||||
ReduceJob: usecase.NewReduceJob(jobRepo, taskRepo, artifactRepo, blobStore, tx, clk, catalog),
|
||||
FailTask: usecase.NewFailTask(taskRepo, jobRepo, workerRepo, tx, clk, catalog),
|
||||
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
|
||||
CancelJob: usecase.NewCancelJob(jobRepo, taskRepo, tx, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(taskRepo, workerRepo, artifactRepo, blobStore, tx, clk),
|
||||
DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore),
|
||||
GetJobResult: usecase.NewGetJobResult(jobRepo, usecase.NewDownloadArtifact(artifactRepo, blobStore)),
|
||||
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
|
||||
Dashboard: usecase.NewDashboard(uiReadRepo),
|
||||
Dashboard: usecase.NewDashboard(uiReadRepo, catalog),
|
||||
PreviewArtifact: usecase.NewPreviewArtifact(uiReadRepo, blobStore),
|
||||
}
|
||||
|
||||
// Background reapers are tracked so shutdown can wait for them. Without this
|
||||
// the process would exit mid-UPDATE, and the deferred pool.Close() would pull
|
||||
// connections out from under them.
|
||||
expireLeases := usecase.NewExpireLeases(taskRepo, clk)
|
||||
expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk, catalog)
|
||||
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
@@ -105,9 +145,18 @@ func run() error {
|
||||
}(r.name, r.fn)
|
||||
}
|
||||
|
||||
// Business metrics: gauges of tasks/jobs/workers by status, sampled from the
|
||||
// database on every Prometheus scrape.
|
||||
statsRepo := postgres.NewStatsRepo(pool)
|
||||
m := metrics.New()
|
||||
m.RegisterBusiness(func(ctx context.Context) (metrics.Stats, error) {
|
||||
tasks, jobs, workers, err := statsRepo.Counts(ctx)
|
||||
return metrics.Stats{Tasks: tasks, Jobs: jobs, Workers: workers}, err
|
||||
})
|
||||
|
||||
// pool.Ping backs /health: readiness means the database answers, not just
|
||||
// that the process is alive.
|
||||
api := httptransport.NewServer(useCases, log, cfg.RequestTimeout, cfg.HeartbeatInterval, cfg.MaxUploadBytes, pool.Ping)
|
||||
api := httptransport.NewServer(useCases, log, cfg.RequestTimeout, cfg.HeartbeatInterval, cfg.MaxUploadBytes, cfg.JWTSecret, cfg.UserserviceURL, m, pool.Ping, cfg.PublicCoordinatorURL, cfg.PublicUserserviceURL, cfg.DocsDir)
|
||||
err = infra.RunServer(ctx, log, cfg.Addr, api.Handler(cfg.Token, cfg.UIToken))
|
||||
|
||||
// Shutdown order matters, and defers alone cannot express it (they run
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/setup"
|
||||
)
|
||||
|
||||
// runSetup implements `coordinator setup` with non-interactive flags and an
|
||||
// interactive fallback for anything still missing.
|
||||
func runSetup(args []string) error {
|
||||
flags := flag.NewFlagSet("setup", flag.ContinueOnError)
|
||||
flags.Usage = func() {
|
||||
_, _ = fmt.Fprintf(flags.Output(), "usage: coordinator setup [options]\n")
|
||||
_, _ = fmt.Fprintf(flags.Output(), "Provisions the coordinator database, schema, and local .env settings.\n\n")
|
||||
flags.PrintDefaults()
|
||||
}
|
||||
var (
|
||||
databaseURL = flags.String("db", "", "coordinator database URL (default: DATABASE_URL)")
|
||||
adminURL = flags.String("admin-db", "", "maintenance URL to create a missing database (default: same host, 'postgres' db)")
|
||||
envFile = flags.String("env-file", "", "settings file to write (default: .env)")
|
||||
force = flags.Bool("force", false, "overwrite an existing settings file")
|
||||
yes = flags.Bool("yes", false, "non-interactive: use defaults, fail on anything missing")
|
||||
)
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() > 0 {
|
||||
return fmt.Errorf("setup takes no positional arguments")
|
||||
}
|
||||
|
||||
databaseURLValue := *databaseURL
|
||||
if databaseURLValue == "" {
|
||||
databaseURLValue = os.Getenv("DATABASE_URL")
|
||||
}
|
||||
envFileValue := *envFile
|
||||
if envFileValue == "" {
|
||||
envFileValue = os.Getenv("ENV_FILE")
|
||||
}
|
||||
adminURLValue := *adminURL
|
||||
if adminURLValue == "" {
|
||||
adminURLValue = os.Getenv("POSTGRES_ADMIN_URL")
|
||||
}
|
||||
|
||||
options := setup.Options{
|
||||
DatabaseURL: databaseURLValue,
|
||||
AdminDatabaseURL: adminURLValue,
|
||||
EnvFile: envFileValue,
|
||||
Force: *force,
|
||||
Yes: *yes,
|
||||
ConnectTimeout: 10 * time.Second,
|
||||
Out: os.Stdout,
|
||||
In: os.Stdin,
|
||||
}
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
log := slog.New(slog.NewTextHandler(os.Stderr, nil))
|
||||
log.Info("setup started", "db", setup.SanitizeDatabaseURL(databaseURLValue), "env_file", envFileValue)
|
||||
|
||||
summary, err := setup.Run(ctx, options)
|
||||
if err != nil {
|
||||
log.Error("setup failed", "err", err)
|
||||
return err
|
||||
}
|
||||
_, _ = fmt.Fprint(os.Stdout, summary)
|
||||
log.Info("setup complete")
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,44 @@
|
||||
// Command worker-agent is the Go worker agent: a coordinator client that
|
||||
// executes SDK workloads in a Python subprocess per claimed task.
|
||||
package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/agent"
|
||||
)
|
||||
|
||||
// version is injected at build time (-ldflags "-X main.version=...") and
|
||||
// reported by --version. "dev" marks a local build.
|
||||
var version = "dev"
|
||||
|
||||
func main() {
|
||||
showVersion := flag.Bool("version", false, "print the build version and exit")
|
||||
flag.Parse()
|
||||
if *showVersion {
|
||||
fmt.Println("worker-agent " + version)
|
||||
return
|
||||
}
|
||||
config, err := agent.LoadConfig()
|
||||
if err != nil {
|
||||
slog.Error("invalid configuration", "error", err)
|
||||
os.Exit(2)
|
||||
}
|
||||
logger := slog.New(slog.NewTextHandler(os.Stderr, nil))
|
||||
tokens := agent.NewTokenProvider(
|
||||
config.WorkerKey,
|
||||
config.UserserviceURL,
|
||||
config.Token,
|
||||
config.RequestTimeout,
|
||||
)
|
||||
client := agent.NewClient(config.CoordinatorURL, tokens, config.RequestTimeout)
|
||||
runner := agent.NewTaskRunner(config.TaskRunner)
|
||||
daemon := agent.NewDaemon(config, client, runner, logger)
|
||||
if err := daemon.RunForever(); err != nil {
|
||||
logger.Error("agent stopped", "error", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
# Demo overlay: Prometheus scrapes the coordinator's /metrics, Grafana shows the
|
||||
# provisioned SciMesh dashboard. Merged by scripts/demo-ui.sh with a third -f.
|
||||
# Both share the coordinator's compose network, so Prometheus reaches it by name.
|
||||
|
||||
services:
|
||||
prometheus:
|
||||
image: prom/prometheus:v2.54.1
|
||||
volumes:
|
||||
- ./monitoring/prometheus.yml:/etc/prometheus/prometheus.yml:ro
|
||||
ports:
|
||||
- "${PROMETHEUS_PORT:-19090}:9090"
|
||||
restart: unless-stopped
|
||||
|
||||
grafana:
|
||||
image: grafana/grafana:11.2.0
|
||||
depends_on:
|
||||
- prometheus
|
||||
environment:
|
||||
GF_SECURITY_ADMIN_USER: admin
|
||||
GF_SECURITY_ADMIN_PASSWORD: ${GRAFANA_PASSWORD:-admin}
|
||||
# Anonymous viewing so the demo dashboard opens without a login.
|
||||
GF_AUTH_ANONYMOUS_ENABLED: "true"
|
||||
GF_AUTH_ANONYMOUS_ORG_ROLE: Viewer
|
||||
GF_USERS_DEFAULT_THEME: dark
|
||||
volumes:
|
||||
- ./monitoring/grafana/provisioning:/etc/grafana/provisioning:ro
|
||||
- ./monitoring/grafana/dashboards:/var/lib/grafana/dashboards:ro
|
||||
ports:
|
||||
- "${GRAFANA_PORT:-13000}:3000"
|
||||
restart: unless-stopped
|
||||
@@ -0,0 +1,71 @@
|
||||
# Demo overlay: adds the userservice (its own Postgres + migrations) alongside
|
||||
# the coordinator and wires the two together with a shared JWT secret, so the
|
||||
# operator UI authenticates through userservice login/registration.
|
||||
#
|
||||
# Used only by scripts/demo-ui.sh, merged onto docker-compose.yml with a second
|
||||
# -f. Not part of the plain `make up` stack.
|
||||
|
||||
services:
|
||||
postgres-users:
|
||||
image: postgres:16-alpine
|
||||
environment:
|
||||
POSTGRES_USER: ${POSTGRES_USER:-scimesh}
|
||||
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:-scimesh}
|
||||
POSTGRES_DB: scimesh_users
|
||||
healthcheck:
|
||||
test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER:-scimesh} -d scimesh_users"]
|
||||
interval: 5s
|
||||
timeout: 3s
|
||||
retries: 10
|
||||
start_period: 5s
|
||||
|
||||
migrate-users:
|
||||
image: migrate/migrate:v4.17.1
|
||||
depends_on:
|
||||
postgres-users:
|
||||
condition: service_healthy
|
||||
volumes:
|
||||
- ../users/migrations:/migrations:ro
|
||||
command:
|
||||
- -path=/migrations
|
||||
- -database=postgres://${POSTGRES_USER:-scimesh}:${POSTGRES_PASSWORD:-scimesh}@postgres-users:5432/scimesh_users?sslmode=disable
|
||||
- up
|
||||
restart: on-failure
|
||||
|
||||
userservice:
|
||||
build:
|
||||
context: ../users
|
||||
depends_on:
|
||||
postgres-users:
|
||||
condition: service_healthy
|
||||
migrate-users:
|
||||
condition: service_completed_successfully
|
||||
environment:
|
||||
USERSERVICE_ADDR: ":8081"
|
||||
DATABASE_URL: postgres://${POSTGRES_USER:-scimesh}:${POSTGRES_PASSWORD:-scimesh}@postgres-users:5432/scimesh_users?sslmode=disable
|
||||
JWT_SECRET: ${JWT_SECRET}
|
||||
# Seeds the first admin the very first time it boots (idempotent after).
|
||||
BOOTSTRAP_ADMIN_EMAIL: ${BOOTSTRAP_ADMIN_EMAIL:-root@scimesh.local}
|
||||
BOOTSTRAP_ADMIN_PASSWORD: ${BOOTSTRAP_ADMIN_PASSWORD}
|
||||
LOG_LEVEL: ${LOG_LEVEL:-info}
|
||||
ports:
|
||||
- "${USERSERVICE_PORT:-18081}:8081"
|
||||
healthcheck:
|
||||
test: ["CMD", "wget", "-qO-", "http://127.0.0.1:8081/health"]
|
||||
interval: 10s
|
||||
timeout: 3s
|
||||
retries: 3
|
||||
start_period: 5s
|
||||
restart: unless-stopped
|
||||
|
||||
# Turn the coordinator UI into session mode: the same shared secret verifies
|
||||
# userservice tokens locally, and USERSERVICE_URL is where login/register proxy.
|
||||
coordinator:
|
||||
environment:
|
||||
JWT_SECRET: ${JWT_SECRET}
|
||||
USERSERVICE_URL: http://userservice:8081
|
||||
# Browser/host-facing URLs for the "add your machine" command. A user's
|
||||
# worker runs on the host, so it reaches the published ports on localhost,
|
||||
# not the in-cluster service names.
|
||||
PUBLIC_COORDINATOR_URL: http://localhost:${COORDINATOR_PORT:-8080}
|
||||
PUBLIC_USERSERVICE_URL: http://localhost:${USERSERVICE_PORT:-8081}
|
||||
@@ -20,20 +20,8 @@ services:
|
||||
retries: 10
|
||||
start_period: 5s
|
||||
|
||||
# One-shot: applies migrations, then exits. Schema changes stay an explicit
|
||||
# deployment step — the coordinator binary never migrates on startup.
|
||||
migrate:
|
||||
image: migrate/migrate:v4.17.1
|
||||
depends_on:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
volumes:
|
||||
- ./migrations:/migrations:ro
|
||||
command:
|
||||
- -path=/migrations
|
||||
- -database=postgres://${POSTGRES_USER:-scimesh}:${POSTGRES_PASSWORD:-scimesh}@postgres:5432/${POSTGRES_DB:-scimesh}?sslmode=disable
|
||||
- up
|
||||
restart: on-failure
|
||||
# The coordinator applies its embedded schema migrations on startup
|
||||
# (AUTO_MIGRATE, on by default), so no separate migration step is needed.
|
||||
|
||||
coordinator:
|
||||
build:
|
||||
@@ -41,9 +29,6 @@ services:
|
||||
depends_on:
|
||||
postgres:
|
||||
condition: service_healthy
|
||||
# Start only once the schema exists, otherwise the first query fails.
|
||||
migrate:
|
||||
condition: service_completed_successfully
|
||||
environment:
|
||||
COORDINATOR_ADDR: ":8080"
|
||||
# Host is the service name: compose resolves it on the project network.
|
||||
@@ -51,6 +36,9 @@ services:
|
||||
WORKER_AUTH_TOKEN: ${WORKER_AUTH_TOKEN:-dev-token}
|
||||
# Empty disables /ui. Set this separately from the worker token.
|
||||
UI_AUTH_TOKEN: ${UI_AUTH_TOKEN:-}
|
||||
# Directory of the built MkDocs site served at /ui/docs/ (empty disables
|
||||
# the docs route; the demo mounts ./site automatically).
|
||||
SCIMESH_DOCS_DIR: ${SCIMESH_DOCS_DIR:-}
|
||||
DB_MAX_CONNS: "10"
|
||||
REQUEST_TIMEOUT: "15s"
|
||||
LEASE_DURATION: "2m"
|
||||
|
||||
+13
-3
@@ -1,23 +1,33 @@
|
||||
module github.com/emil28092005/SciMesh/coordinator
|
||||
|
||||
go 1.22
|
||||
go 1.25.0
|
||||
|
||||
require (
|
||||
github.com/Masterminds/squirrel v1.5.4
|
||||
github.com/cenkalti/backoff/v4 v4.3.0
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1
|
||||
github.com/google/uuid v1.6.0
|
||||
github.com/jackc/pgx/v5 v5.6.0
|
||||
github.com/joho/godotenv v1.5.1
|
||||
github.com/prometheus/client_golang v1.19.1
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/beorn7/perks v1.0.1 // indirect
|
||||
github.com/cespare/xxhash/v2 v2.3.0 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
|
||||
github.com/jackc/puddle/v2 v2.2.1 // indirect
|
||||
github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 // indirect
|
||||
github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 // indirect
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
|
||||
github.com/prometheus/client_model v0.6.2 // indirect
|
||||
github.com/prometheus/common v0.55.0 // indirect
|
||||
github.com/prometheus/procfs v0.21.1 // indirect
|
||||
golang.org/x/crypto v0.17.0 // indirect
|
||||
golang.org/x/sync v0.1.0 // indirect
|
||||
golang.org/x/text v0.14.0 // indirect
|
||||
golang.org/x/sync v0.22.0 // indirect
|
||||
golang.org/x/sys v0.47.0 // indirect
|
||||
golang.org/x/text v0.40.0 // indirect
|
||||
google.golang.org/protobuf v1.36.11 // indirect
|
||||
)
|
||||
|
||||
+28
-6
@@ -1,10 +1,18 @@
|
||||
github.com/Masterminds/squirrel v1.5.4 h1:uUcX/aBc8O7Fg9kaISIUsHXdKuqehiXAMQTYX8afzqM=
|
||||
github.com/Masterminds/squirrel v1.5.4/go.mod h1:NNaOrjSoIDfDA40n7sr2tPNZRfjzjA400rg+riTZj10=
|
||||
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
|
||||
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
|
||||
github.com/cenkalti/backoff/v4 v4.3.0 h1:MyRJ/UdXutAwSAT+s3wNd7MfTIcy71VQueUuFK343L8=
|
||||
github.com/cenkalti/backoff/v4 v4.3.0/go.mod h1:Y3VNntkOUPxTVeUxJ/G5vcM//AlwfmyYozVcomhLiZE=
|
||||
github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UFvs=
|
||||
github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
|
||||
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1 h1:kYf81DTWFe7t+1VvL7eS+jKFVWaUnK9cB1qbwn63YCY=
|
||||
github.com/golang-jwt/jwt/v5 v5.3.1/go.mod h1:fxCRLWMO43lRc8nhHWY6LGqRcf+1gQWArsqaEUEa5bE=
|
||||
github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8=
|
||||
github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU=
|
||||
github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
|
||||
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM=
|
||||
@@ -21,20 +29,34 @@ github.com/lann/builder v0.0.0-20180802200727-47ae307949d0 h1:SOEGU9fKiNWd/HOJuq
|
||||
github.com/lann/builder v0.0.0-20180802200727-47ae307949d0/go.mod h1:dXGbAdH5GtBTC4WfIxhKZfyBF/HBFgRZSWwZ9g/He9o=
|
||||
github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0 h1:P6pPBnrTSX3DEVR4fDembhRWSsG5rVo6hYhAB/ADZrk=
|
||||
github.com/lann/ps v0.0.0-20150810152359-62de8c46ede0/go.mod h1:vmVJ0l/dxyfGW6FmdpVm2joNMFikkuWg0EoCKLGUMNw=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 h1:C3w9PqII01/Oq1c1nUAm88MOHcQC9l5mIlSMApZMrHA=
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8mK8HuQTcFwEeznm/DIxMOiR9yIdICNftLE1DvQ=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
|
||||
github.com/prometheus/client_golang v1.19.1 h1:wZWJDwK+NameRJuPGDhlnFgx8e8HN3XHQeLaYJFJBOE=
|
||||
github.com/prometheus/client_golang v1.19.1/go.mod h1:mP78NwGzrVks5S2H6ab8+ZZGJLZUq1hoULYBAYBw1Ho=
|
||||
github.com/prometheus/client_model v0.6.2 h1:oBsgwpGs7iVziMvrGhE53c/GrLUsZdHnqNwqPLxwZyk=
|
||||
github.com/prometheus/client_model v0.6.2/go.mod h1:y3m2F6Gdpfy6Ut/GBsUqTWZqCUvMVzSfMLjcu6wAwpE=
|
||||
github.com/prometheus/common v0.55.0 h1:KEi6DK7lXW/m7Ig5i47x0vRzuBsHuvJdi5ee6Y3G1dc=
|
||||
github.com/prometheus/common v0.55.0/go.mod h1:2SECS4xJG1kd8XF9IcM1gMX6510RAEL65zxzNImwdc8=
|
||||
github.com/prometheus/procfs v0.21.1 h1:GljZCt+zSTS+NZq88cyQ1LjZ+RCHp3uVuabBWA5+OJI=
|
||||
github.com/prometheus/procfs v0.21.1/go.mod h1:aB55Cww9pdSJVHk0hUf0inxWyyjPogFIjmHKYgMKmtY=
|
||||
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
|
||||
github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs=
|
||||
github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
|
||||
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
|
||||
github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk=
|
||||
github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
|
||||
github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
|
||||
github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
|
||||
golang.org/x/crypto v0.17.0 h1:r8bRNjWL3GshPW3gkd+RpvzWrZAwPS49OmTGZ/uhM4k=
|
||||
golang.org/x/crypto v0.17.0/go.mod h1:gCAAfMLgwOJRpTjQ2zCCt2OcSfYMTeZVSRtQlPC7Nq4=
|
||||
golang.org/x/sync v0.1.0 h1:wsuoTGHzEhffawBOhz5CYhcrV4IdKZbEyZjBMuTp12o=
|
||||
golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/text v0.14.0 h1:ScX5w1eTa3QqT8oi6+ziP7dTV1S2+ALU0bI+0zXKWiQ=
|
||||
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/sync v0.22.0 h1:SZjpbeLmrCk4xhRSZFNZW5gFUeCeFgjekvI/+gfScek=
|
||||
golang.org/x/sync v0.22.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0=
|
||||
golang.org/x/sys v0.47.0 h1:o7XGOvZQCADBQQ4Y7VNq2dRWQR7JmOUW8Kxx4ZsNgWs=
|
||||
golang.org/x/sys v0.47.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw=
|
||||
golang.org/x/text v0.40.0 h1:Ub2Z6/xjgF1WrYQz2nuITOEegKFtiIy+rieRJ5lHZKs=
|
||||
golang.org/x/text v0.40.0/go.mod h1:hpnzDAfGV753zIKo+wk3u1bVKCGPbrnF7+7LBF/UHVY=
|
||||
google.golang.org/protobuf v1.36.11 h1:fV6ZwhNocDyBLK0dj+fg8ektcVegBBuEolpbTQyBNVE=
|
||||
google.golang.org/protobuf v1.36.11/go.mod h1:HTf+CrKn2C3g5S8VImy6tdcUvCska2kB7j23XfzDpco=
|
||||
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1 h1:bBRl1b0OH9s/DuPhuXpNl+VtCaJXFZ5/uEFST95x9zc=
|
||||
gopkg.in/natefinch/lumberjack.v2 v2.2.1/go.mod h1:YD8tP3GAjkrDg1eZH7EGmyESg/lsYskCTPBJVb9jqSc=
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TokenProvider supplies the current bearer token. A static token is served
|
||||
// forever; a worker key is exchanged at the userservice for short-lived JWTs
|
||||
// and refreshed before they expire (mirroring the former Python worker).
|
||||
type TokenProvider interface {
|
||||
Token() (string, error)
|
||||
Refresh() error
|
||||
}
|
||||
|
||||
// StaticToken serves a fixed token forever; empty means no Authorization.
|
||||
type StaticToken struct{ token string }
|
||||
|
||||
func (s *StaticToken) Token() (string, error) { return s.token, nil }
|
||||
func (s *StaticToken) Refresh() error { return nil }
|
||||
|
||||
// WorkerKeyToken exchanges a long-lived worker key for short-lived JWTs.
|
||||
type WorkerKeyToken struct {
|
||||
userserviceURL string
|
||||
workerKey string
|
||||
timeout time.Duration
|
||||
leeway float64
|
||||
mu sync.Mutex
|
||||
token string
|
||||
refreshAt time.Time
|
||||
}
|
||||
|
||||
func NewWorkerKeyToken(userserviceURL, workerKey string, timeout time.Duration) *WorkerKeyToken {
|
||||
return &WorkerKeyToken{
|
||||
userserviceURL: strings.TrimRight(userserviceURL, "/"),
|
||||
workerKey: workerKey,
|
||||
timeout: timeout,
|
||||
leeway: 0.2,
|
||||
}
|
||||
}
|
||||
|
||||
// Token returns the current token, exchanging first when missing or stale.
|
||||
func (p *WorkerKeyToken) Token() (string, error) {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if p.token == "" || time.Now().After(p.refreshAt) {
|
||||
if err := p.exchangeLocked(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return p.token, nil
|
||||
}
|
||||
|
||||
// Refresh forces an immediate exchange.
|
||||
func (p *WorkerKeyToken) Refresh() error {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return p.exchangeLocked()
|
||||
}
|
||||
|
||||
func (p *WorkerKeyToken) exchangeLocked() error {
|
||||
payload, err := json.Marshal(map[string]string{"key": p.workerKey})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodPost, p.userserviceURL+"/worker-tokens/exchange", bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
client := &http.Client{Timeout: p.timeout}
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return fmt.Errorf("worker key exchange request failed")
|
||||
}
|
||||
defer func() { _ = response.Body.Close() }()
|
||||
if response.StatusCode != http.StatusOK {
|
||||
return fmt.Errorf("worker key exchange rejected with status %d", response.StatusCode)
|
||||
}
|
||||
raw, err := io.ReadAll(io.LimitReader(response.Body, 1<<20))
|
||||
if err != nil {
|
||||
return fmt.Errorf("worker key exchange request failed")
|
||||
}
|
||||
var data map[string]any
|
||||
if err := json.Unmarshal(raw, &data); err != nil {
|
||||
return fmt.Errorf("worker key exchange response is invalid")
|
||||
}
|
||||
token, _ := data["token"].(string)
|
||||
if token == "" {
|
||||
return fmt.Errorf("worker key exchange response is missing a token")
|
||||
}
|
||||
var ttl time.Duration
|
||||
switch value := data["expires_in"].(type) {
|
||||
case float64:
|
||||
ttl = time.Duration(value * float64(time.Second))
|
||||
case int:
|
||||
ttl = time.Duration(value) * time.Second
|
||||
}
|
||||
p.token = token
|
||||
p.refreshAt = time.Time{}
|
||||
if ttl > 0 {
|
||||
p.refreshAt = time.Now().Add(time.Duration(float64(ttl) * (1.0 - p.leeway)))
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// NewTokenProvider picks the strategy: a worker key (with userservice) wins
|
||||
// over a static bearer token.
|
||||
func NewTokenProvider(workerKey, userserviceURL, bearerToken string, timeout time.Duration) TokenProvider {
|
||||
if workerKey != "" && userserviceURL != "" {
|
||||
return NewWorkerKeyToken(userserviceURL, workerKey, timeout)
|
||||
}
|
||||
return &StaticToken{token: bearerToken}
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestWorkerKeyTokenExchangesAndCaches(t *testing.T) {
|
||||
var exchanges atomic.Int64
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/worker-tokens/exchange" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
var payload map[string]string
|
||||
if err := json.NewDecoder(r.Body).Decode(&payload); err != nil || payload["key"] != "scimesh_wk_live_x" {
|
||||
http.Error(w, "bad key", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
exchanges.Add(1)
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"token": "jwt-1",
|
||||
"expires_in": 100,
|
||||
})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider := NewWorkerKeyToken(server.URL, "scimesh_wk_live_x", 5*time.Second)
|
||||
token, err := provider.Token()
|
||||
if err != nil || token != "jwt-1" {
|
||||
t.Fatalf("token = %q, err = %v", token, err)
|
||||
}
|
||||
// The second call within the TTL reuses the cache.
|
||||
again, err := provider.Token()
|
||||
if err != nil || again != "jwt-1" {
|
||||
t.Fatalf("cached token = %q, err = %v", again, err)
|
||||
}
|
||||
if exchanges.Load() != 1 {
|
||||
t.Errorf("exchanges = %d, want 1", exchanges.Load())
|
||||
}
|
||||
// An explicit refresh re-exchanges.
|
||||
if err := provider.Refresh(); err != nil {
|
||||
t.Fatalf("Refresh: %v", err)
|
||||
}
|
||||
if exchanges.Load() != 2 {
|
||||
t.Errorf("exchanges after refresh = %d, want 2", exchanges.Load())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkerKeyTokenRejectsBadKey(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider := NewWorkerKeyToken(server.URL, "bad", 5*time.Second)
|
||||
if _, err := provider.Token(); err == nil {
|
||||
t.Error("expected exchange failure for a rejected key")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewTokenProviderSelectsStrategy(t *testing.T) {
|
||||
if _, ok := NewTokenProvider("", "", "static", time.Second).(*StaticToken); !ok {
|
||||
t.Error("expected a static token provider")
|
||||
}
|
||||
if _, ok := NewTokenProvider("key", "http://users", "", time.Second).(*WorkerKeyToken); !ok {
|
||||
t.Error("expected a worker-key provider")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientRefreshesTokenOnceOn401(t *testing.T) {
|
||||
var attempts atomic.Int64
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
attempts.Add(1)
|
||||
if attempts.Load() == 1 {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{})
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
provider := &StaticToken{token: "t"}
|
||||
client := NewClient(server.URL, provider, 5*time.Second)
|
||||
status, _, err := client.requestJSON(http.MethodGet, "/ok", map[string]any{})
|
||||
if err != nil {
|
||||
t.Fatalf("request: %v", err)
|
||||
}
|
||||
if status != http.StatusOK {
|
||||
t.Errorf("status = %d", status)
|
||||
}
|
||||
if attempts.Load() != 2 {
|
||||
t.Errorf("attempts = %d, want 2 (401 then retry)", attempts.Load())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,374 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// CoordinatorError is a non-retriable coordinator response.
|
||||
type CoordinatorError struct{ msg string }
|
||||
|
||||
func (e *CoordinatorError) Error() string { return e.msg }
|
||||
|
||||
// TransientError is a timeout, connection error, or 5xx response.
|
||||
type TransientError struct{ msg string }
|
||||
|
||||
func (e *TransientError) Error() string { return e.msg }
|
||||
|
||||
// ConflictError means the worker no longer owns the task lease.
|
||||
type ConflictError struct{ msg string }
|
||||
|
||||
func (e *ConflictError) Error() string { return e.msg }
|
||||
|
||||
// Client speaks the v1 worker contract over HTTP with a token provider.
|
||||
//
|
||||
// API calls never follow redirects (a redirect is a contract violation); the
|
||||
// artifact download follows redirects but strips the Authorization header on
|
||||
// cross-origin hops, matching the Python worker's SameOriginAuthRedirectHandler.
|
||||
// A 401 response refreshes the token exactly once and retries, so a lapsed JWT
|
||||
// does not fail an in-flight task.
|
||||
type Client struct {
|
||||
baseURL string
|
||||
tokens TokenProvider
|
||||
timeout time.Duration
|
||||
apiClient *http.Client
|
||||
dlClient *http.Client
|
||||
}
|
||||
|
||||
func NewClient(baseURL string, tokens TokenProvider, timeout time.Duration) *Client {
|
||||
return &Client{
|
||||
baseURL: strings.TrimRight(baseURL, "/"),
|
||||
tokens: tokens,
|
||||
timeout: timeout,
|
||||
apiClient: &http.Client{
|
||||
Timeout: timeout,
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse },
|
||||
},
|
||||
dlClient: &http.Client{
|
||||
Timeout: timeout,
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= 10 {
|
||||
return fmt.Errorf("too many redirects")
|
||||
}
|
||||
// Go strips Authorization on cross-host redirects by default;
|
||||
// strip it explicitly on any origin change to be safe.
|
||||
if len(via) > 0 && origin(req.URL) != origin(via[0].URL) {
|
||||
req.Header.Del("Authorization")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func origin(u *url.URL) string {
|
||||
return u.Scheme + "://" + u.Host
|
||||
}
|
||||
|
||||
func (c *Client) authHeaders() (map[string]string, error) {
|
||||
token, err := c.tokens.Token()
|
||||
if err != nil {
|
||||
return nil, &CoordinatorError{msg: "token refresh failed: " + err.Error()}
|
||||
}
|
||||
if token == "" {
|
||||
return map[string]string{}, nil
|
||||
}
|
||||
return map[string]string{"Authorization": "Bearer " + token}, nil
|
||||
}
|
||||
|
||||
func (c *Client) refreshAndRetry() bool {
|
||||
return c.tokens.Refresh() == nil
|
||||
}
|
||||
|
||||
// Register advertises the worker and returns its identity and heartbeat policy.
|
||||
func (c *Client) Register(name string, capabilities []string, cpuCount int, memoryMB int) (*RegisteredWorker, error) {
|
||||
payload := map[string]any{
|
||||
"name": name,
|
||||
"capabilities": capabilities,
|
||||
"cpu_count": cpuCount,
|
||||
}
|
||||
if memoryMB > 0 {
|
||||
payload["memory_mb"] = memoryMB
|
||||
}
|
||||
status, body, err := c.requestJSON("POST", "/workers/register", payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status != http.StatusCreated {
|
||||
return nil, &CoordinatorError{msg: fmt.Sprintf("worker registration rejected with status %d", status)}
|
||||
}
|
||||
return ParseRegistered(body)
|
||||
}
|
||||
|
||||
// Claim leases one compatible task, or returns nil when the queue is empty.
|
||||
func (c *Client) Claim(workerID string, capabilities []string) (*Task, error) {
|
||||
status, body, err := c.requestJSON("POST", "/tasks/claim", map[string]any{
|
||||
"worker_id": workerID,
|
||||
"capabilities": capabilities,
|
||||
"max_concurrency": 1,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if status == http.StatusNoContent {
|
||||
return nil, nil
|
||||
}
|
||||
if status != http.StatusOK {
|
||||
return nil, &CoordinatorError{msg: fmt.Sprintf("unexpected claim status %d", status)}
|
||||
}
|
||||
return ParseTask(body)
|
||||
}
|
||||
|
||||
// Heartbeat renews the lease and returns the new deadline.
|
||||
func (c *Client) Heartbeat(task *Task, workerID string) (time.Time, error) {
|
||||
status, body, err := c.requestJSON("POST", "/tasks/"+task.TaskID+"/heartbeat", map[string]any{
|
||||
"worker_id": workerID,
|
||||
"attempt": task.Attempt,
|
||||
})
|
||||
if err != nil {
|
||||
return time.Time{}, err
|
||||
}
|
||||
if status != http.StatusOK {
|
||||
if status == http.StatusConflict {
|
||||
return time.Time{}, &ConflictError{msg: "heartbeat rejected because the task lease was lost"}
|
||||
}
|
||||
return time.Time{}, &CoordinatorError{msg: fmt.Sprintf("heartbeat rejected with status %d", status)}
|
||||
}
|
||||
raw, ok := body["lease_expires_at"].(string)
|
||||
if !ok {
|
||||
return time.Time{}, &CoordinatorError{msg: "heartbeat response is missing lease_expires_at"}
|
||||
}
|
||||
lease, err := time.Parse(time.RFC3339, raw)
|
||||
if err != nil {
|
||||
return time.Time{}, &CoordinatorError{msg: "heartbeat returned an invalid lease_expires_at"}
|
||||
}
|
||||
task.LeaseExpiresAt = lease
|
||||
task.leaseExpiresRaw = raw
|
||||
return lease, nil
|
||||
}
|
||||
|
||||
// Submit completes a task with the uploaded coordinator-owned artifact.
|
||||
func (c *Client) Submit(task *Task, workerID string, uploaded *Uploaded, metrics map[string]any) error {
|
||||
status, _, err := c.requestJSON("POST", "/tasks/"+task.TaskID+"/result", map[string]any{
|
||||
"worker_id": workerID,
|
||||
"attempt": task.Attempt,
|
||||
"result": map[string]any{"artifact_id": uploaded.ArtifactID},
|
||||
"metrics": metrics,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if status != http.StatusOK && status != http.StatusCreated && status != http.StatusAccepted {
|
||||
if status == http.StatusConflict {
|
||||
return &ConflictError{msg: "result rejected because the task lease was lost"}
|
||||
}
|
||||
return &CoordinatorError{msg: fmt.Sprintf("result rejected with status %d", status)}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Fail reports a sanitized failure.
|
||||
func (c *Client) Fail(task *Task, workerID string, code, message string, retryable bool) error {
|
||||
status, _, err := c.requestJSON("POST", "/tasks/"+task.TaskID+"/failure", map[string]any{
|
||||
"worker_id": workerID,
|
||||
"attempt": task.Attempt,
|
||||
"error_code": code,
|
||||
"error_message": message,
|
||||
"retryable": retryable,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if status != http.StatusOK && status != http.StatusCreated && status != http.StatusAccepted {
|
||||
if status == http.StatusConflict {
|
||||
return &ConflictError{msg: "failure rejected because the task lease was lost"}
|
||||
}
|
||||
return &CoordinatorError{msg: fmt.Sprintf("failure report rejected with status %d", status)}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Download streams the task input to destination and returns its SHA-256.
|
||||
func (c *Client) Download(uri, destination string) (string, error) {
|
||||
resolved, err := url.Parse(uri)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("invalid input URI: %w", err)
|
||||
}
|
||||
if !resolved.IsAbs() {
|
||||
base, parseErr := url.Parse(c.baseURL)
|
||||
if parseErr != nil {
|
||||
return "", fmt.Errorf("invalid coordinator URL")
|
||||
}
|
||||
resolved = base.ResolveReference(resolved)
|
||||
}
|
||||
if resolved.Scheme != "http" && resolved.Scheme != "https" {
|
||||
return "", fmt.Errorf("input URI must be an HTTP(S) URL")
|
||||
}
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodGet, resolved.String(), nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
headers, err := c.authHeaders()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
for name, value := range headers {
|
||||
request.Header.Set(name, value)
|
||||
}
|
||||
response, err := c.dlClient.Do(request)
|
||||
if err != nil {
|
||||
return "", &TransientError{msg: "input download failed"}
|
||||
}
|
||||
defer func() { _ = response.Body.Close() }()
|
||||
if response.StatusCode == http.StatusUnauthorized && c.refreshAndRetry() {
|
||||
return c.Download(uri, destination)
|
||||
}
|
||||
if response.StatusCode != http.StatusOK {
|
||||
return "", &CoordinatorError{msg: fmt.Sprintf("input download rejected with status %d", response.StatusCode)}
|
||||
}
|
||||
if err := os.MkdirAll(filepath.Dir(destination), 0o750); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// #nosec G304 -- destination is the worker's own attempt directory file.
|
||||
target, err := os.OpenFile(destination, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
digest := sha256.New()
|
||||
_, copyErr := io.Copy(io.MultiWriter(target, digest), response.Body)
|
||||
closeErr := target.Close()
|
||||
if copyErr != nil {
|
||||
_ = os.Remove(destination)
|
||||
return "", &TransientError{msg: "input download interrupted"}
|
||||
}
|
||||
if closeErr != nil {
|
||||
return "", closeErr
|
||||
}
|
||||
return hex.EncodeToString(digest.Sum(nil)), nil
|
||||
}
|
||||
|
||||
// Upload streams a partial artifact and verifies the returned metadata.
|
||||
func (c *Client) Upload(task *Task, workerID string, path, contentType string) (*Uploaded, error) {
|
||||
// #nosec G304 -- the upload path is this worker's own artifact file.
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
info, err := file.Stat()
|
||||
if err != nil {
|
||||
_ = file.Close()
|
||||
return nil, err
|
||||
}
|
||||
digest := sha256.New()
|
||||
if _, err := io.Copy(digest, file); err != nil {
|
||||
_ = file.Close()
|
||||
return nil, err
|
||||
}
|
||||
if _, err := file.Seek(0, io.SeekStart); err != nil {
|
||||
_ = file.Close()
|
||||
return nil, err
|
||||
}
|
||||
localSHA := hex.EncodeToString(digest.Sum(nil))
|
||||
uploadURL := c.baseURL + "/tasks/" + url.PathEscape(task.TaskID) + "/artifacts/" + url.PathEscape(filepath.Base(path))
|
||||
request, err := http.NewRequestWithContext(context.Background(), http.MethodPut, uploadURL, file)
|
||||
if err != nil {
|
||||
_ = file.Close()
|
||||
return nil, err
|
||||
}
|
||||
request.ContentLength = info.Size()
|
||||
request.Header.Set("Content-Type", contentType)
|
||||
request.Header.Set("X-Worker-ID", workerID)
|
||||
request.Header.Set("X-Task-Attempt", strconv.Itoa(task.Attempt))
|
||||
headers, err := c.authHeaders()
|
||||
if err != nil {
|
||||
_ = file.Close()
|
||||
return nil, err
|
||||
}
|
||||
for name, value := range headers {
|
||||
request.Header.Set(name, value)
|
||||
}
|
||||
response, err := c.apiClient.Do(request)
|
||||
_ = file.Close()
|
||||
if err != nil {
|
||||
return nil, &TransientError{msg: "artifact upload failed"}
|
||||
}
|
||||
defer func() { _ = response.Body.Close() }()
|
||||
raw, err := io.ReadAll(io.LimitReader(response.Body, 1<<20))
|
||||
if err != nil {
|
||||
return nil, &TransientError{msg: "artifact upload interrupted"}
|
||||
}
|
||||
if response.StatusCode == http.StatusConflict {
|
||||
return nil, &ConflictError{msg: "artifact upload rejected because the task lease was lost"}
|
||||
}
|
||||
if response.StatusCode == http.StatusUnauthorized && c.refreshAndRetry() {
|
||||
return c.Upload(task, workerID, path, contentType)
|
||||
}
|
||||
if response.StatusCode != http.StatusOK {
|
||||
return nil, &CoordinatorError{msg: fmt.Sprintf("artifact upload rejected with status %d", response.StatusCode)}
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal(raw, &payload); err != nil {
|
||||
return nil, &CoordinatorError{msg: "artifact upload returned invalid metadata"}
|
||||
}
|
||||
uploaded, err := ParseUploaded(payload)
|
||||
if err != nil {
|
||||
return nil, &CoordinatorError{msg: "artifact upload returned invalid metadata"}
|
||||
}
|
||||
if uploaded.SHA256 != localSHA || uploaded.SizeBytes != info.Size() {
|
||||
return nil, &CoordinatorError{msg: "artifact upload metadata does not match local artifact"}
|
||||
}
|
||||
return uploaded, nil
|
||||
}
|
||||
|
||||
func (c *Client) requestJSON(method, path string, payload any) (int, map[string]any, error) {
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
request, err := http.NewRequestWithContext(context.Background(), method, c.baseURL+path, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
headers, err := c.authHeaders()
|
||||
if err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
for name, value := range headers {
|
||||
request.Header.Set(name, value)
|
||||
}
|
||||
response, err := c.apiClient.Do(request)
|
||||
if err != nil {
|
||||
return 0, nil, &TransientError{msg: "coordinator request failed"}
|
||||
}
|
||||
defer func() { _ = response.Body.Close() }()
|
||||
raw, err := io.ReadAll(io.LimitReader(response.Body, 1<<20))
|
||||
if err != nil {
|
||||
return 0, nil, &TransientError{msg: "coordinator request interrupted"}
|
||||
}
|
||||
if response.StatusCode == http.StatusUnauthorized && c.refreshAndRetry() {
|
||||
return c.requestJSON(method, path, payload)
|
||||
}
|
||||
if response.StatusCode >= 500 {
|
||||
return response.StatusCode, nil, &TransientError{msg: fmt.Sprintf("coordinator returned %d", response.StatusCode)}
|
||||
}
|
||||
var decoded map[string]any
|
||||
if len(raw) > 0 {
|
||||
if err := json.Unmarshal(raw, &decoded); err != nil {
|
||||
return response.StatusCode, nil, &CoordinatorError{msg: "coordinator returned invalid JSON"}
|
||||
}
|
||||
}
|
||||
return response.StatusCode, decoded, nil
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func newTestClient(t *testing.T, server *httptest.Server) *Client {
|
||||
t.Helper()
|
||||
return NewClient(server.URL, &StaticToken{token: "test-token"}, 5*time.Second)
|
||||
}
|
||||
|
||||
func TestClientRegisterClaimHeartbeat(t *testing.T) {
|
||||
var registered, claimed, heartbeated bool
|
||||
var server *httptest.Server //nolint:staticcheck // the handler closure references server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("Authorization") != "Bearer test-token" {
|
||||
http.Error(w, "missing token", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
switch {
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/workers/register":
|
||||
registered = true
|
||||
writeJSON(w, http.StatusCreated, map[string]any{
|
||||
"worker_id": "22222222-2222-4222-8222-222222222222",
|
||||
"heartbeat_interval_seconds": 15,
|
||||
})
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/tasks/claim":
|
||||
claimed = true
|
||||
writeJSON(w, http.StatusOK, validTaskPayload())
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/tasks/11111111-1111-4111-8111-111111111111/heartbeat":
|
||||
heartbeated = true
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"lease_expires_at": time.Now().Add(2 * time.Minute).UTC().Format(time.RFC3339),
|
||||
})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client := newTestClient(t, server)
|
||||
|
||||
registeredWorker, err := client.Register("test-worker", []string{"similarity-search"}, 2, 1024)
|
||||
if err != nil {
|
||||
t.Fatalf("Register: %v", err)
|
||||
}
|
||||
if registeredWorker.WorkerID != "22222222-2222-4222-8222-222222222222" {
|
||||
t.Errorf("worker id = %q", registeredWorker.WorkerID)
|
||||
}
|
||||
|
||||
task, err := client.Claim("22222222-2222-4222-8222-222222222222", []string{"similarity-search"})
|
||||
if err != nil {
|
||||
t.Fatalf("Claim: %v", err)
|
||||
}
|
||||
if task == nil || task.Workload != "similarity-search" {
|
||||
t.Fatalf("claim = %+v", task)
|
||||
}
|
||||
|
||||
renewed, err := client.Heartbeat(task, "22222222-2222-4222-8222-222222222222")
|
||||
if err != nil {
|
||||
t.Fatalf("Heartbeat: %v", err)
|
||||
}
|
||||
if renewed.Before(time.Now()) {
|
||||
t.Error("renewed lease is in the past")
|
||||
}
|
||||
if !registered || !claimed || !heartbeated {
|
||||
t.Error("some endpoints were not hit")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientClaimEmptyAndConflict(t *testing.T) {
|
||||
var server *httptest.Server //nolint:staticcheck // the handler closure references server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/tasks/claim":
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
case "/tasks/11111111-1111-4111-8111-111111111111/heartbeat":
|
||||
w.WriteHeader(http.StatusConflict)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client := newTestClient(t, server)
|
||||
|
||||
task, err := client.Claim("worker", []string{"similarity-search"})
|
||||
if err != nil {
|
||||
t.Fatalf("Claim: %v", err)
|
||||
}
|
||||
if task != nil {
|
||||
t.Error("expected no task for 204")
|
||||
}
|
||||
|
||||
claimed, err := ParseTask(validTaskPayload())
|
||||
if err != nil {
|
||||
t.Fatalf("ParseTask: %v", err)
|
||||
}
|
||||
if _, err := client.Heartbeat(claimed, "worker"); err == nil {
|
||||
t.Error("expected conflict error")
|
||||
} else if !errors.As(err, &conflictError) {
|
||||
t.Errorf("error type = %T", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientUploadSubmitFail(t *testing.T) {
|
||||
var uploadedPath string
|
||||
var server *httptest.Server //nolint:staticcheck // the handler closure references server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.Method == http.MethodPut && strings.HasPrefix(r.URL.Path, "/tasks/11111111-1111-4111-8111-111111111111/artifacts/"):
|
||||
if r.Header.Get("X-Worker-ID") != "worker" || r.Header.Get("X-Task-Attempt") != "1" {
|
||||
t.Errorf("missing identity headers: %+v", r.Header)
|
||||
}
|
||||
uploadedPath = r.URL.Path
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"artifact_id": "33333333-3333-4333-8333-333333333333",
|
||||
"uri": server.URL + "/artifacts/333/download",
|
||||
"sha256": sha256Of(t, "partial body"),
|
||||
"size_bytes": int64(len("partial body")),
|
||||
})
|
||||
case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/result"):
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{})
|
||||
case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/failure"):
|
||||
writeJSON(w, http.StatusAccepted, map[string]any{})
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
defer server.Close()
|
||||
client := newTestClient(t, server)
|
||||
task, _ := ParseTask(validTaskPayload())
|
||||
|
||||
dir := t.TempDir()
|
||||
partial := filepath.Join(dir, "result.csv")
|
||||
if err := os.WriteFile(partial, []byte("partial body"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
uploaded, err := client.Upload(task, "worker", partial, "text/csv")
|
||||
if err != nil {
|
||||
t.Fatalf("Upload: %v", err)
|
||||
}
|
||||
if uploaded.SizeBytes != int64(len("partial body")) {
|
||||
t.Errorf("size = %d", uploaded.SizeBytes)
|
||||
}
|
||||
if !strings.Contains(uploadedPath, "result.csv") {
|
||||
t.Errorf("upload path = %q", uploadedPath)
|
||||
}
|
||||
|
||||
if err := client.Submit(task, "worker", uploaded, map[string]any{"rows": 1}); err != nil {
|
||||
t.Fatalf("Submit: %v", err)
|
||||
}
|
||||
if err := client.Fail(task, "worker", "ValueError", "bad input", false); err != nil {
|
||||
t.Fatalf("Fail: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientDownloadVerifiesChecksumAndStripsAuthOnRedirect(t *testing.T) {
|
||||
var redirectedAuth string
|
||||
bucket := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
redirectedAuth = r.Header.Get("Authorization")
|
||||
_, _ = w.Write([]byte("input bytes"))
|
||||
}))
|
||||
defer bucket.Close()
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/tasks/11111111-1111-4111-8111-111111111111/input" {
|
||||
http.Redirect(w, r, bucket.URL+"/presigned", http.StatusFound)
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer server.Close()
|
||||
client := newTestClient(t, server)
|
||||
|
||||
destination := filepath.Join(t.TempDir(), "input")
|
||||
digest, err := client.Download(server.URL+"/tasks/11111111-1111-4111-8111-111111111111/input", destination)
|
||||
if err != nil {
|
||||
t.Fatalf("Download: %v", err)
|
||||
}
|
||||
if digest != sha256Of(t, "input bytes") {
|
||||
t.Errorf("digest = %q", digest)
|
||||
}
|
||||
if redirectedAuth != "" {
|
||||
t.Error("Authorization must be stripped on the redirected download")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeErrorMessageRedactsPaths(t *testing.T) {
|
||||
message := SanitizeErrorMessage(
|
||||
"failed at /home/alice/work/attempts/1/input and /private/secret.txt",
|
||||
"/home/alice/work",
|
||||
)
|
||||
for _, forbidden := range []string{"/home/alice", "/private/secret.txt"} {
|
||||
if strings.Contains(message, forbidden) {
|
||||
t.Errorf("message leaks %q: %q", forbidden, message)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(message, "<worker-dir>") {
|
||||
t.Errorf("work dir not redacted: %q", message)
|
||||
}
|
||||
long := SanitizeErrorMessage(strings.Repeat("x", 500), "/tmp")
|
||||
if len(long) != 300 {
|
||||
t.Errorf("truncated length = %d", len(long))
|
||||
}
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, payload any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(payload)
|
||||
}
|
||||
|
||||
func sha256Of(t *testing.T, value string) string {
|
||||
t.Helper()
|
||||
digest := sha256.Sum256([]byte(value))
|
||||
return fmt.Sprintf("%x", digest)
|
||||
}
|
||||
@@ -0,0 +1,177 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/workloads"
|
||||
)
|
||||
|
||||
// Config is read only from the environment, mirroring the former Python
|
||||
// worker's configuration surface.
|
||||
type Config struct {
|
||||
CoordinatorURL string
|
||||
Token string
|
||||
WorkerKey string
|
||||
UserserviceURL string
|
||||
WorkerName string
|
||||
WorkerID string // set after registration; overridable for tests
|
||||
WorkDir string
|
||||
CPUCount int
|
||||
MemoryMB int // 0 = not advertised
|
||||
PollInterval time.Duration
|
||||
RequestTimeout time.Duration
|
||||
Heartbeat time.Duration
|
||||
CleanupAfter time.Duration // 0 = keep attempt directories
|
||||
Capabilities []string
|
||||
TaskRunner []string // command + args; defaults to python -m scimesh.worker.task
|
||||
MaxTasks int // 0 = unlimited
|
||||
ExitWhenIdle bool
|
||||
}
|
||||
|
||||
func envList(name string) ([]string, error) {
|
||||
raw := os.Getenv(name)
|
||||
if raw == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var items []string
|
||||
if err := json.Unmarshal([]byte(raw), &items); err != nil {
|
||||
return nil, fmt.Errorf("%s must be a JSON array", name)
|
||||
}
|
||||
for _, item := range items {
|
||||
if strings.TrimSpace(item) == "" {
|
||||
return nil, fmt.Errorf("%s must not contain empty entries", name)
|
||||
}
|
||||
}
|
||||
return items, nil
|
||||
}
|
||||
|
||||
// LoadConfig validates the environment and fails fast on invalid values.
|
||||
func LoadConfig() (*Config, error) {
|
||||
url := os.Getenv("COORDINATOR_URL")
|
||||
if url == "" {
|
||||
return nil, fmt.Errorf("COORDINATOR_URL is required")
|
||||
}
|
||||
if !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "https://") {
|
||||
return nil, fmt.Errorf("COORDINATOR_URL must be an absolute HTTP(S) URL")
|
||||
}
|
||||
workDir := os.Getenv("WORK_DIR")
|
||||
if workDir == "" {
|
||||
workDir = "./scimesh-agent-data"
|
||||
}
|
||||
cpu := 1
|
||||
if raw := os.Getenv("CPU_COUNT"); raw != "" {
|
||||
parsed, err := strconv.Atoi(raw)
|
||||
if err != nil || parsed < 1 {
|
||||
return nil, fmt.Errorf("CPU_COUNT must be a positive integer")
|
||||
}
|
||||
cpu = parsed
|
||||
}
|
||||
memoryMB := 0
|
||||
if raw := os.Getenv("MEMORY_MB"); raw != "" {
|
||||
parsed, err := strconv.Atoi(raw)
|
||||
if err != nil || parsed < 1 {
|
||||
return nil, fmt.Errorf("MEMORY_MB must be a positive integer")
|
||||
}
|
||||
memoryMB = parsed
|
||||
}
|
||||
poll, err := durationEnv("POLL_INTERVAL", 2*time.Second)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
timeout, err := durationEnv("REQUEST_TIMEOUT", 30*time.Second)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
heartbeat, err := durationEnv("HEARTBEAT_INTERVAL", 15*time.Second)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
cleanup, err := durationEnv("CLEANUP_AFTER_SECONDS", 0)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
capabilities, err := envList("CAPABILITIES")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(capabilities) == 0 {
|
||||
capabilities = defaultCapabilities()
|
||||
}
|
||||
runner, err := envList("TASK_RUNNER")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(runner) == 0 {
|
||||
runner = []string{"python", "-m", "scimesh.worker.task"}
|
||||
}
|
||||
maxTasks := 0
|
||||
if raw := os.Getenv("MAX_TASKS"); raw != "" {
|
||||
parsed, err := strconv.Atoi(raw)
|
||||
if err != nil || parsed < 1 {
|
||||
return nil, fmt.Errorf("MAX_TASKS must be a positive integer")
|
||||
}
|
||||
maxTasks = parsed
|
||||
}
|
||||
name := os.Getenv("WORKER_NAME")
|
||||
if name == "" {
|
||||
host, _ := os.Hostname()
|
||||
name = host
|
||||
}
|
||||
absWorkDir, err := filepath.Abs(workDir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("WORK_DIR must be an absolute path")
|
||||
}
|
||||
return &Config{
|
||||
CoordinatorURL: strings.TrimRight(url, "/"),
|
||||
Token: os.Getenv("WORKER_AUTH_TOKEN"),
|
||||
WorkerKey: os.Getenv("WORKER_KEY"),
|
||||
UserserviceURL: strings.TrimRight(os.Getenv("USERSERVICE_URL"), "/"),
|
||||
WorkerName: name,
|
||||
WorkerID: os.Getenv("WORKER_ID"),
|
||||
WorkDir: absWorkDir,
|
||||
CPUCount: cpu,
|
||||
MemoryMB: memoryMB,
|
||||
PollInterval: poll,
|
||||
RequestTimeout: timeout,
|
||||
Heartbeat: heartbeat,
|
||||
CleanupAfter: cleanup,
|
||||
Capabilities: capabilities,
|
||||
TaskRunner: runner,
|
||||
MaxTasks: maxTasks,
|
||||
ExitWhenIdle: os.Getenv("EXIT_WHEN_IDLE") == "1",
|
||||
}, nil
|
||||
}
|
||||
|
||||
func durationEnv(name string, fallback time.Duration) (time.Duration, error) {
|
||||
raw := os.Getenv(name)
|
||||
if raw == "" {
|
||||
return fallback, nil
|
||||
}
|
||||
parsed, err := time.ParseDuration(raw)
|
||||
if err != nil || parsed < 0 {
|
||||
return 0, fmt.Errorf("%s must be a non-negative duration", name)
|
||||
}
|
||||
return parsed, nil
|
||||
}
|
||||
|
||||
// defaultCapabilities derives the worker's advertised capabilities from the
|
||||
// embedded workload catalog, so an agent is workload-agnostic out of the box:
|
||||
// it claims whatever enabled workloads the coordinator library declares.
|
||||
// Explicit CAPABILITIES still overrides this for operators who want a subset.
|
||||
func defaultCapabilities() []string {
|
||||
catalog, err := workloads.Load()
|
||||
if err != nil {
|
||||
return []string{"similarity-search"}
|
||||
}
|
||||
names := make([]string, 0, len(catalog.Enabled()))
|
||||
for _, workload := range catalog.Enabled() {
|
||||
names = append(names, workload.Name)
|
||||
}
|
||||
return names
|
||||
}
|
||||
@@ -0,0 +1,336 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Outcome reports whether a claim was made and whether it completed.
|
||||
type Outcome struct {
|
||||
Claimed bool
|
||||
Completed bool
|
||||
}
|
||||
|
||||
// Daemon is the agent state machine: register, claim, execute via the Python
|
||||
// task runner, upload, and submit — mirroring the Python worker's lifecycle.
|
||||
// conflictError is the errors.As target for lease conflicts.
|
||||
var conflictError *ConflictError
|
||||
|
||||
type Daemon struct {
|
||||
config *Config
|
||||
client *Client
|
||||
runner *TaskRunner
|
||||
log *slog.Logger
|
||||
workerID string
|
||||
registered bool
|
||||
completed int
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func NewDaemon(config *Config, client *Client, runner *TaskRunner, log *slog.Logger) *Daemon {
|
||||
return &Daemon{config: config, client: client, runner: runner, log: log}
|
||||
}
|
||||
|
||||
// RunForever loops until interrupted, idle-exit, or max tasks.
|
||||
func (d *Daemon) RunForever() error {
|
||||
failures := 0
|
||||
for {
|
||||
if !d.registered {
|
||||
if err := d.register(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
d.cleanupExpiredDirectories()
|
||||
outcome, err := d.runOnce()
|
||||
if err != nil {
|
||||
failures++
|
||||
d.log.Warn("agent cycle failed", "error", err)
|
||||
backoff := d.config.PollInterval
|
||||
for i := 0; i < failures && i < 6; i++ {
|
||||
backoff *= 2
|
||||
}
|
||||
if backoff > 60*time.Second {
|
||||
backoff = 60 * time.Second
|
||||
}
|
||||
time.Sleep(backoff)
|
||||
continue
|
||||
}
|
||||
failures = 0
|
||||
if outcome.Claimed && outcome.Completed {
|
||||
d.completed++
|
||||
if d.config.MaxTasks > 0 && d.completed >= d.config.MaxTasks {
|
||||
d.log.Info("max tasks reached")
|
||||
return nil
|
||||
}
|
||||
}
|
||||
if !outcome.Claimed && d.config.ExitWhenIdle {
|
||||
d.log.Info("queue empty, exiting")
|
||||
return nil
|
||||
}
|
||||
if outcome.Claimed && d.config.ExitWhenIdle {
|
||||
d.log.Info("one claim processed, exiting")
|
||||
return nil
|
||||
}
|
||||
if !outcome.Claimed {
|
||||
time.Sleep(d.config.PollInterval)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Daemon) register() error {
|
||||
registered, err := d.client.Register(
|
||||
d.config.WorkerName,
|
||||
d.config.Capabilities,
|
||||
d.config.CPUCount,
|
||||
d.config.MemoryMB,
|
||||
)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
d.mu.Lock()
|
||||
if d.config.WorkerID != "" {
|
||||
d.workerID = d.config.WorkerID
|
||||
} else {
|
||||
d.workerID = registered.WorkerID
|
||||
}
|
||||
d.registered = true
|
||||
d.mu.Unlock()
|
||||
d.log.Info("registered", "worker_id", d.workerID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *Daemon) workerIDOrEmpty() string {
|
||||
d.mu.Lock()
|
||||
defer d.mu.Unlock()
|
||||
return d.workerID
|
||||
}
|
||||
|
||||
// cleanupExpiredDirectories removes task attempt directories older than the
|
||||
// configured retention, mirroring the former Python worker's cleanup.
|
||||
func (d *Daemon) cleanupExpiredDirectories() {
|
||||
if d.config.CleanupAfter <= 0 {
|
||||
return
|
||||
}
|
||||
cutoff := time.Now().Add(-d.config.CleanupAfter)
|
||||
tasks, err := os.ReadDir(d.config.WorkDir)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
for _, taskEntry := range tasks {
|
||||
if !taskEntry.IsDir() {
|
||||
continue
|
||||
}
|
||||
taskDir := filepath.Join(d.config.WorkDir, taskEntry.Name())
|
||||
attempts, err := os.ReadDir(taskDir)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
for _, attemptEntry := range attempts {
|
||||
if !attemptEntry.IsDir() {
|
||||
continue
|
||||
}
|
||||
info, err := attemptEntry.Info()
|
||||
if err == nil && info.ModTime().Before(cutoff) {
|
||||
_ = os.RemoveAll(filepath.Join(taskDir, attemptEntry.Name()))
|
||||
}
|
||||
}
|
||||
if entries, err := os.ReadDir(taskDir); err == nil && len(entries) == 0 {
|
||||
_ = os.Remove(taskDir)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (d *Daemon) runOnce() (Outcome, error) {
|
||||
workerID := d.workerIDOrEmpty()
|
||||
if workerID == "" {
|
||||
return Outcome{}, fmt.Errorf("agent is not registered")
|
||||
}
|
||||
task, err := d.client.Claim(workerID, d.config.Capabilities)
|
||||
if err != nil {
|
||||
return Outcome{}, err
|
||||
}
|
||||
if task == nil {
|
||||
return Outcome{Claimed: false}, nil
|
||||
}
|
||||
started := time.Now()
|
||||
taskDir := filepath.Join(d.config.WorkDir, task.TaskID, fmt.Sprint(task.Attempt))
|
||||
if err := os.MkdirAll(taskDir, 0o750); err != nil {
|
||||
return Outcome{Claimed: true}, err
|
||||
}
|
||||
heartbeat := newLeaseHeartbeat(task, workerID, d.client, d.config.Heartbeat)
|
||||
completed := false
|
||||
err = heartbeat.Start()
|
||||
if err != nil {
|
||||
if errors.As(err, &conflictError) {
|
||||
d.log.Warn("lease lost", "task_id", task.TaskID)
|
||||
return Outcome{Claimed: true}, nil
|
||||
}
|
||||
return Outcome{Claimed: true}, err
|
||||
}
|
||||
defer heartbeat.Stop()
|
||||
|
||||
// Attempt directory cleanup is deliberately minimal in the prototype:
|
||||
// attempt directories are retained under the work directory.
|
||||
failure := d.executeTask(task, workerID, taskDir, started, heartbeat)
|
||||
if failure != nil {
|
||||
if errors.As(failure, &conflictError) {
|
||||
d.log.Warn("lease lost", "task_id", task.TaskID)
|
||||
return Outcome{Claimed: true}, nil
|
||||
}
|
||||
if err := heartbeat.RaiseIfFailed(); err != nil {
|
||||
return Outcome{Claimed: true}, nil
|
||||
}
|
||||
d.reportFailure(task, workerID, failure)
|
||||
return Outcome{Claimed: true}, nil
|
||||
}
|
||||
if err := heartbeat.RaiseIfFailed(); err != nil {
|
||||
return Outcome{Claimed: true}, nil
|
||||
}
|
||||
completed = true
|
||||
d.log.Info("task completed", "task_id", task.TaskID, "elapsed_seconds", time.Since(started).Seconds())
|
||||
return Outcome{Claimed: true, Completed: completed}, nil
|
||||
}
|
||||
|
||||
// executeTask returns nil on success or a classified failure.
|
||||
func (d *Daemon) executeTask(task *Task, workerID, taskDir string, started time.Time, heartbeat *leaseHeartbeat) error {
|
||||
inputPath := filepath.Join(taskDir, "input")
|
||||
// Downloads use the coordinator-provided URI verbatim; a relative path is
|
||||
// resolved against the coordinator by the client.
|
||||
actualSHA, err := d.client.Download(task.Input.URI, inputPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !strings.EqualFold(actualSHA, task.Input.SHA256) {
|
||||
return &CoordinatorError{msg: "input checksum mismatch"}
|
||||
}
|
||||
if err := heartbeat.RaiseIfFailed(); err != nil {
|
||||
return err
|
||||
}
|
||||
manifestPath := filepath.Join(taskDir, "manifest.json")
|
||||
manifest, err := d.runner.Run(task, taskDir, manifestPath, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := heartbeat.RaiseIfFailed(); err != nil {
|
||||
return err
|
||||
}
|
||||
uploaded, err := d.client.Upload(task, workerID, manifest.ArtifactPath, manifest.ContentType)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := heartbeat.RaiseIfFailed(); err != nil {
|
||||
return err
|
||||
}
|
||||
metrics := map[string]any{"elapsed_seconds": roundSeconds(time.Since(started).Seconds())}
|
||||
for name, value := range manifest.Metrics {
|
||||
metrics[name] = value
|
||||
}
|
||||
return d.client.Submit(task, workerID, uploaded, metrics)
|
||||
}
|
||||
|
||||
func (d *Daemon) reportFailure(task *Task, workerID string, failure error) {
|
||||
var code string
|
||||
var coordinatorErr *CoordinatorError
|
||||
if errors.As(failure, &coordinatorErr) {
|
||||
code = "ValueError"
|
||||
} else {
|
||||
code = "TaskRunnerFailed"
|
||||
}
|
||||
retryable := IsRetryableError(failure)
|
||||
message := SanitizeErrorMessage(failure.Error(), d.config.WorkDir)
|
||||
d.log.Warn("task failed", "task_id", task.TaskID, "error_code", code, "retryable", retryable)
|
||||
if err := d.client.Fail(task, workerID, code, message, retryable); err != nil {
|
||||
if errors.As(err, &conflictError) {
|
||||
d.log.Warn("lease lost while reporting failure", "task_id", task.TaskID)
|
||||
return
|
||||
}
|
||||
d.log.Warn("failure report rejected", "task_id", task.TaskID, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// leaseHeartbeat renews the lease from the returned deadline at less than
|
||||
// half of the remaining TTL, mirroring the Python worker.
|
||||
type leaseHeartbeat struct {
|
||||
task *Task
|
||||
workerID string
|
||||
client *Client
|
||||
interval time.Duration
|
||||
stop chan struct{}
|
||||
once sync.Once
|
||||
mu sync.Mutex
|
||||
lease time.Time
|
||||
failed error
|
||||
}
|
||||
|
||||
func newLeaseHeartbeat(task *Task, workerID string, client *Client, interval time.Duration) *leaseHeartbeat {
|
||||
return &leaseHeartbeat{
|
||||
task: task,
|
||||
workerID: workerID,
|
||||
client: client,
|
||||
interval: interval,
|
||||
stop: make(chan struct{}),
|
||||
lease: task.LeaseExpiresAt,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *leaseHeartbeat) Start() error {
|
||||
if _, err := h.client.Heartbeat(h.task, h.workerID); err != nil {
|
||||
return err
|
||||
}
|
||||
go h.loop()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (h *leaseHeartbeat) Stop() {
|
||||
h.once.Do(func() { close(h.stop) })
|
||||
}
|
||||
|
||||
func (h *leaseHeartbeat) RaiseIfFailed() error {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
return h.failed
|
||||
}
|
||||
|
||||
func (h *leaseHeartbeat) loop() {
|
||||
for {
|
||||
delay := h.nextDelay()
|
||||
select {
|
||||
case <-h.stop:
|
||||
return
|
||||
case <-time.After(delay):
|
||||
}
|
||||
renewed, err := h.client.Heartbeat(h.task, h.workerID)
|
||||
h.mu.Lock()
|
||||
if err != nil {
|
||||
h.failed = err
|
||||
h.mu.Unlock()
|
||||
return
|
||||
}
|
||||
h.lease = renewed
|
||||
h.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func (h *leaseHeartbeat) nextDelay() time.Duration {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
remaining := time.Until(h.lease)
|
||||
if remaining <= 0 {
|
||||
return 0
|
||||
}
|
||||
half := remaining / 2
|
||||
if h.interval < half {
|
||||
return h.interval
|
||||
}
|
||||
return half
|
||||
}
|
||||
|
||||
func roundSeconds(seconds float64) float64 {
|
||||
return float64(int64(seconds*1000)) / 1000
|
||||
}
|
||||
@@ -0,0 +1,296 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func sha256HexOf(value string) string {
|
||||
sum := sha256.Sum256([]byte(value))
|
||||
return fmt.Sprintf("%x", sum)
|
||||
}
|
||||
|
||||
// fakeRunnerScript writes a result manifest for --output and exits with the
|
||||
// given code.
|
||||
func fakeRunnerScript(t *testing.T, dir string, exitCode int) string {
|
||||
t.Helper()
|
||||
script := filepath.Join(dir, "fake-runner.sh")
|
||||
content := `#!/bin/sh
|
||||
out=""
|
||||
task_dir=""
|
||||
while [ "$#" -gt 0 ]; do
|
||||
case "$1" in
|
||||
--output) out="$2"; shift 2;;
|
||||
--task-dir) task_dir="$2"; shift 2;;
|
||||
*) shift;;
|
||||
esac
|
||||
done
|
||||
printf 'id,score\n' > "$task_dir/result.csv"
|
||||
printf '{"artifact_path":"%s/result.csv","content_type":"text/csv","metrics":{"rows":1}}' "$task_dir" > "$out"
|
||||
exit ` + fmt.Sprint(exitCode) + "\n"
|
||||
if err := os.WriteFile(script, []byte(content), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return script
|
||||
}
|
||||
|
||||
// fakeCoordinator implements the v1 contract over HTTP and records calls.
|
||||
type fakeCoordinator struct {
|
||||
server *httptest.Server
|
||||
task map[string]any
|
||||
submits []map[string]any
|
||||
failures []map[string]any
|
||||
heartbeats int
|
||||
uploadSHA string
|
||||
uploadSize int64
|
||||
inputBytes []byte
|
||||
conflict bool // 409 on heartbeat/upload/result
|
||||
}
|
||||
|
||||
func newFakeCoordinator(t *testing.T, task map[string]any) *fakeCoordinator {
|
||||
t.Helper()
|
||||
fake := &fakeCoordinator{task: task, inputBytes: []byte("input fixture")}
|
||||
fake.uploadSHA = sha256HexOf(string(fake.inputBytes))
|
||||
fake.uploadSize = int64(len(fake.inputBytes))
|
||||
var server *httptest.Server
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/workers/register":
|
||||
writeJSON(w, http.StatusCreated, map[string]any{
|
||||
"worker_id": "22222222-2222-4222-8222-222222222222",
|
||||
"heartbeat_interval_seconds": 15.0,
|
||||
})
|
||||
case r.Method == http.MethodPost && r.URL.Path == "/tasks/claim":
|
||||
if fake.task == nil {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, fake.task)
|
||||
case r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/input"):
|
||||
_, _ = w.Write(fake.inputBytes)
|
||||
case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/heartbeat"):
|
||||
fake.heartbeats++
|
||||
if fake.conflict {
|
||||
w.WriteHeader(http.StatusConflict)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"lease_expires_at": time.Now().Add(2 * time.Minute).UTC().Format(time.RFC3339),
|
||||
})
|
||||
case r.Method == http.MethodPut && strings.Contains(r.URL.Path, "/artifacts/"):
|
||||
raw, _ := io.ReadAll(r.Body)
|
||||
fake.uploadSize = int64(len(raw))
|
||||
fake.uploadSHA = fmt.Sprintf("%x", sha256.Sum256(raw))
|
||||
if fake.conflict {
|
||||
w.WriteHeader(http.StatusConflict)
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]any{
|
||||
"artifact_id": "33333333-3333-4333-8333-333333333333",
|
||||
"uri": server.URL + "/artifacts/333/download",
|
||||
"sha256": fake.uploadSHA,
|
||||
"size_bytes": fake.uploadSize,
|
||||
})
|
||||
case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/result"):
|
||||
var payload map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&payload)
|
||||
fake.submits = append(fake.submits, payload)
|
||||
if fake.conflict {
|
||||
w.WriteHeader(http.StatusConflict)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/failure"):
|
||||
var payload map[string]any
|
||||
_ = json.NewDecoder(r.Body).Decode(&payload)
|
||||
fake.failures = append(fake.failures, payload)
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
default:
|
||||
http.NotFound(w, r)
|
||||
}
|
||||
}))
|
||||
fake.server = server
|
||||
return fake
|
||||
}
|
||||
|
||||
func (f *fakeCoordinator) close() { f.server.Close() }
|
||||
|
||||
func testDaemon(t *testing.T, fake *fakeCoordinator, script string) *Daemon {
|
||||
t.Helper()
|
||||
config := &Config{
|
||||
CoordinatorURL: fake.server.URL,
|
||||
WorkerName: "test-worker",
|
||||
WorkerID: "22222222-2222-4222-8222-222222222222",
|
||||
WorkDir: t.TempDir(),
|
||||
CPUCount: 1,
|
||||
PollInterval: time.Millisecond,
|
||||
RequestTimeout: 5 * time.Second,
|
||||
Heartbeat: 15 * time.Second,
|
||||
Capabilities: []string{"similarity-search"},
|
||||
TaskRunner: []string{script},
|
||||
}
|
||||
client := NewClient(fake.server.URL, &StaticToken{token: "test-token"}, 5*time.Second)
|
||||
runner := NewTaskRunner(config.TaskRunner)
|
||||
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
daemon := NewDaemon(config, client, runner, logger)
|
||||
if err := daemon.register(); err != nil {
|
||||
t.Fatalf("register: %v", err)
|
||||
}
|
||||
return daemon
|
||||
}
|
||||
|
||||
func validClaimedTaskPayload() map[string]any {
|
||||
return map[string]any{
|
||||
"task_id": "11111111-1111-4111-8111-111111111111",
|
||||
"attempt": 1.0,
|
||||
"lease_expires_at": time.Now().Add(time.Minute).UTC().Format(time.RFC3339),
|
||||
"workload": "similarity-search",
|
||||
"input": map[string]any{
|
||||
"uri": "/tasks/11111111-1111-4111-8111-111111111111/input",
|
||||
"sha256": sha256HexOf("input fixture"),
|
||||
},
|
||||
"parameters": map[string]any{"query_smiles": "CCO"},
|
||||
}
|
||||
}
|
||||
|
||||
func TestDaemonCompletesAClaimedTask(t *testing.T) {
|
||||
fake := newFakeCoordinator(t, validClaimedTaskPayload())
|
||||
defer fake.close()
|
||||
daemon := testDaemon(t, fake, fakeRunnerScript(t, t.TempDir(), 0))
|
||||
|
||||
outcome, err := daemon.runOnce()
|
||||
if err != nil {
|
||||
t.Fatalf("runOnce: %v", err)
|
||||
}
|
||||
if !outcome.Claimed || !outcome.Completed {
|
||||
t.Fatalf("outcome = %+v", outcome)
|
||||
}
|
||||
if len(fake.submits) != 1 {
|
||||
t.Fatalf("submits = %d", len(fake.submits))
|
||||
}
|
||||
result := fake.submits[0]["result"].(map[string]any)
|
||||
if result["artifact_id"] != "33333333-3333-4333-8333-333333333333" {
|
||||
t.Errorf("result artifact = %v", result)
|
||||
}
|
||||
metrics := fake.submits[0]["metrics"].(map[string]any)
|
||||
if metrics["rows"] != float64(1) {
|
||||
t.Errorf("metrics = %v", metrics)
|
||||
}
|
||||
if _, ok := metrics["elapsed_seconds"].(float64); !ok {
|
||||
t.Errorf("missing elapsed_seconds: %v", metrics)
|
||||
}
|
||||
if fake.heartbeats < 1 {
|
||||
t.Error("expected at least one heartbeat")
|
||||
}
|
||||
if len(fake.failures) != 0 {
|
||||
t.Errorf("unexpected failures: %v", fake.failures)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDaemonReportsChecksumMismatchAsPermanentFailure(t *testing.T) {
|
||||
payload := validClaimedTaskPayload()
|
||||
payload["input"].(map[string]any)["sha256"] = strings.Repeat("b", 64)
|
||||
fake := newFakeCoordinator(t, payload)
|
||||
defer fake.close()
|
||||
daemon := testDaemon(t, fake, fakeRunnerScript(t, t.TempDir(), 0))
|
||||
|
||||
outcome, err := daemon.runOnce()
|
||||
if err != nil {
|
||||
t.Fatalf("runOnce: %v", err)
|
||||
}
|
||||
if outcome.Completed {
|
||||
t.Fatal("task must not complete on checksum mismatch")
|
||||
}
|
||||
if len(fake.failures) != 1 {
|
||||
t.Fatalf("failures = %d", len(fake.failures))
|
||||
}
|
||||
failure := fake.failures[0]
|
||||
if failure["error_code"] != "ValueError" || failure["retryable"] != false {
|
||||
t.Errorf("failure = %v", failure)
|
||||
}
|
||||
if !strings.Contains(failure["error_message"].(string), "checksum") {
|
||||
t.Errorf("message = %v", failure["error_message"])
|
||||
}
|
||||
if len(fake.submits) != 0 {
|
||||
t.Error("no submission expected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDaemonReportsPermanentRunnerFailure(t *testing.T) {
|
||||
fake := newFakeCoordinator(t, validClaimedTaskPayload())
|
||||
defer fake.close()
|
||||
daemon := testDaemon(t, fake, fakeRunnerScript(t, t.TempDir(), ExitPermanent))
|
||||
|
||||
outcome, err := daemon.runOnce()
|
||||
if err != nil {
|
||||
t.Fatalf("runOnce: %v", err)
|
||||
}
|
||||
if outcome.Completed {
|
||||
t.Fatal("task must not complete")
|
||||
}
|
||||
if len(fake.failures) != 1 || fake.failures[0]["retryable"] != false {
|
||||
t.Fatalf("failures = %v", fake.failures)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDaemonReportsRetryableRunnerFailure(t *testing.T) {
|
||||
fake := newFakeCoordinator(t, validClaimedTaskPayload())
|
||||
defer fake.close()
|
||||
daemon := testDaemon(t, fake, fakeRunnerScript(t, t.TempDir(), 1))
|
||||
|
||||
outcome, err := daemon.runOnce()
|
||||
if err != nil {
|
||||
t.Fatalf("runOnce: %v", err)
|
||||
}
|
||||
if outcome.Completed {
|
||||
t.Fatal("task must not complete")
|
||||
}
|
||||
if len(fake.failures) != 1 || fake.failures[0]["retryable"] != true {
|
||||
t.Fatalf("failures = %v", fake.failures)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDaemonLeaseConflictStopsWithoutFailureReport(t *testing.T) {
|
||||
fake := newFakeCoordinator(t, validClaimedTaskPayload())
|
||||
fake.conflict = true
|
||||
defer fake.close()
|
||||
daemon := testDaemon(t, fake, fakeRunnerScript(t, t.TempDir(), 0))
|
||||
|
||||
outcome, err := daemon.runOnce()
|
||||
if err != nil {
|
||||
t.Fatalf("runOnce: %v", err)
|
||||
}
|
||||
if !outcome.Claimed {
|
||||
t.Fatal("task was claimed")
|
||||
}
|
||||
if len(fake.failures) != 0 {
|
||||
t.Errorf("no failure report expected after lease loss: %v", fake.failures)
|
||||
}
|
||||
if len(fake.submits) != 0 {
|
||||
t.Errorf("no submission expected after lease loss: %v", fake.submits)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDaemonIdleClaimIsNotCompleted(t *testing.T) {
|
||||
fake := newFakeCoordinator(t, nil)
|
||||
defer fake.close()
|
||||
daemon := testDaemon(t, fake, fakeRunnerScript(t, t.TempDir(), 0))
|
||||
|
||||
outcome, err := daemon.runOnce()
|
||||
if err != nil {
|
||||
t.Fatalf("runOnce: %v", err)
|
||||
}
|
||||
if outcome.Claimed || outcome.Completed {
|
||||
t.Fatalf("outcome = %+v", outcome)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
// Package agent implements a Go worker agent: a coordinator client and
|
||||
// task-lifecycle supervisor that executes SDK workloads in a Python
|
||||
// subprocess. It mirrors the Python worker's v1 wire contract exactly; the
|
||||
// Python worker remains the reference implementation.
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
uuidPattern = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$`)
|
||||
sha256Pattern = regexp.MustCompile(`^[0-9a-fA-F]{64}$`)
|
||||
workloadPattern = regexp.MustCompile(`^[a-z][a-z0-9]*(?:[_-][a-z0-9]+)*$`)
|
||||
)
|
||||
|
||||
// RegisteredWorker is the coordinator's answer to /workers/register.
|
||||
type RegisteredWorker struct {
|
||||
WorkerID string
|
||||
HeartbeatIntervalSeconds float64
|
||||
}
|
||||
|
||||
// Input is the claimed task's input artifact.
|
||||
type Input struct {
|
||||
URI string
|
||||
SHA256 string
|
||||
}
|
||||
|
||||
// Task is one claimed, leased task.
|
||||
type Task struct {
|
||||
TaskID string
|
||||
Attempt int
|
||||
LeaseExpiresAt time.Time
|
||||
Workload string
|
||||
Input Input
|
||||
Parameters map[string]any
|
||||
leaseExpiresRaw string
|
||||
}
|
||||
|
||||
// Uploaded is the coordinator-owned metadata returned after artifact upload.
|
||||
type Uploaded struct {
|
||||
ArtifactID string
|
||||
URI string
|
||||
SHA256 string
|
||||
SizeBytes int64
|
||||
}
|
||||
|
||||
func requireString(value any, field string) (string, error) {
|
||||
text, ok := value.(string)
|
||||
if !ok || strings.TrimSpace(text) == "" {
|
||||
return "", fmt.Errorf("%s must be a non-empty string", field)
|
||||
}
|
||||
return text, nil
|
||||
}
|
||||
|
||||
func safeCoordinatorURI(value any, field string) (string, error) {
|
||||
uri, err := requireString(value, field)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.HasPrefix(uri, "/") {
|
||||
// A network-path reference (//host/path) or dot segments would
|
||||
// resolve to another origin; reject both.
|
||||
if strings.HasPrefix(uri, "//") {
|
||||
return "", fmt.Errorf("%s must be a safe coordinator path", field)
|
||||
}
|
||||
for _, segment := range strings.Split(uri, "/") {
|
||||
if segment == ".." {
|
||||
return "", fmt.Errorf("%s must be a safe coordinator path", field)
|
||||
}
|
||||
}
|
||||
return uri, nil
|
||||
}
|
||||
parsed, err := url.Parse(uri)
|
||||
if err != nil || (parsed.Scheme != "http" && parsed.Scheme != "https") || parsed.Host == "" {
|
||||
return "", fmt.Errorf("%s must be an absolute HTTP(S) URL or coordinator path", field)
|
||||
}
|
||||
return uri, nil
|
||||
}
|
||||
|
||||
func sha256Hex(value any, field string) (string, error) {
|
||||
digest, err := requireString(value, field)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
digest = strings.ToLower(digest)
|
||||
if !sha256Pattern.MatchString(digest) {
|
||||
return "", fmt.Errorf("%s must be a SHA-256 hex digest", field)
|
||||
}
|
||||
return digest, nil
|
||||
}
|
||||
|
||||
func uuid(value any, field string) (string, error) {
|
||||
text, err := requireString(value, field)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !uuidPattern.MatchString(text) {
|
||||
return "", fmt.Errorf("%s must be a UUID", field)
|
||||
}
|
||||
return strings.ToLower(text), nil
|
||||
}
|
||||
|
||||
// ParseTask validates a claimed-task response with the same strictness as the
|
||||
// Python worker's ClaimedTask.from_json.
|
||||
func ParseTask(payload map[string]any) (*Task, error) {
|
||||
rawInput, ok := payload["input"].(map[string]any)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("input must be an object")
|
||||
}
|
||||
rawAttempt, ok := payload["attempt"].(float64)
|
||||
if !ok || rawAttempt < 1 || rawAttempt != float64(int(rawAttempt)) {
|
||||
return nil, fmt.Errorf("attempt must be a positive integer")
|
||||
}
|
||||
taskID, err := uuid(payload["task_id"], "task_id")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid claimed-task response: %w", err)
|
||||
}
|
||||
rawLease, err := requireString(payload["lease_expires_at"], "lease_expires_at")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid claimed-task response: %w", err)
|
||||
}
|
||||
lease, err := time.Parse(time.RFC3339, rawLease)
|
||||
if err != nil || lease.Location() == nil {
|
||||
return nil, fmt.Errorf("lease_expires_at must include a timezone")
|
||||
}
|
||||
workload, err := requireString(payload["workload"], "workload")
|
||||
if err != nil || !workloadPattern.MatchString(workload) {
|
||||
return nil, fmt.Errorf("workload must be a canonical name")
|
||||
}
|
||||
uri, err := safeCoordinatorURI(rawInput["uri"], "input.uri")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid claimed-task response: %w", err)
|
||||
}
|
||||
digest, err := sha256Hex(rawInput["sha256"], "input.sha256")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid claimed-task response: %w", err)
|
||||
}
|
||||
parameters, ok := payload["parameters"].(map[string]any)
|
||||
if !ok {
|
||||
parameters = map[string]any{}
|
||||
}
|
||||
return &Task{
|
||||
TaskID: taskID,
|
||||
Attempt: int(rawAttempt),
|
||||
LeaseExpiresAt: lease,
|
||||
leaseExpiresRaw: rawLease,
|
||||
Workload: workload,
|
||||
Input: Input{URI: uri, SHA256: digest},
|
||||
Parameters: parameters,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// LeaseExpiresRaw returns the original lease timestamp string for
|
||||
// round-tripping in heartbeat deadlines.
|
||||
func (t *Task) LeaseExpiresRaw() string { return t.leaseExpiresRaw }
|
||||
|
||||
// ParseRegistered validates a registration response.
|
||||
func ParseRegistered(payload map[string]any) (*RegisteredWorker, error) {
|
||||
workerID, err := uuid(payload["worker_id"], "worker_id")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid worker registration response: %w", err)
|
||||
}
|
||||
interval, ok := payload["heartbeat_interval_seconds"].(float64)
|
||||
if !ok || interval <= 0 {
|
||||
return nil, fmt.Errorf("heartbeat_interval_seconds must be positive")
|
||||
}
|
||||
return &RegisteredWorker{WorkerID: workerID, HeartbeatIntervalSeconds: interval}, nil
|
||||
}
|
||||
|
||||
// ParseUploaded validates an artifact upload response.
|
||||
func ParseUploaded(payload map[string]any) (*Uploaded, error) {
|
||||
artifactID, err := uuid(payload["artifact_id"], "artifact_id")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid artifact upload response: %w", err)
|
||||
}
|
||||
uri, err := safeCoordinatorURI(payload["uri"], "uri")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid artifact upload response: %w", err)
|
||||
}
|
||||
digest, err := sha256Hex(payload["sha256"], "sha256")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("invalid artifact upload response: %w", err)
|
||||
}
|
||||
rawSize, ok := payload["size_bytes"].(float64)
|
||||
if !ok || rawSize < 0 || rawSize != float64(int64(rawSize)) {
|
||||
return nil, fmt.Errorf("artifact size_bytes must be a non-negative integer")
|
||||
}
|
||||
return &Uploaded{ArtifactID: artifactID, URI: uri, SHA256: digest, SizeBytes: int64(rawSize)}, nil
|
||||
}
|
||||
|
||||
// TaskRunnerManifest is what the Python task entry writes on success.
|
||||
type TaskRunnerManifest struct {
|
||||
ArtifactPath string `json:"artifact_path"`
|
||||
ContentType string `json:"content_type"`
|
||||
Metrics map[string]any `json:"metrics"`
|
||||
}
|
||||
|
||||
// Encode serializes a claim payload for /tasks/claim.
|
||||
func Encode(v any) ([]byte, error) { return json.Marshal(v) }
|
||||
@@ -0,0 +1,111 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func validTaskPayload() map[string]any {
|
||||
return map[string]any{
|
||||
"task_id": "11111111-1111-4111-8111-111111111111",
|
||||
"attempt": 1.0,
|
||||
"lease_expires_at": "2026-08-02T00:00:00Z",
|
||||
"workload": "similarity-search",
|
||||
"input": map[string]any{
|
||||
"uri": "/tasks/11111111-1111-4111-8111-111111111111/input",
|
||||
"sha256": "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
|
||||
},
|
||||
"parameters": map[string]any{"query_smiles": "CCO"},
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseTaskAcceptsValidPayload(t *testing.T) {
|
||||
task, err := ParseTask(validTaskPayload())
|
||||
if err != nil {
|
||||
t.Fatalf("ParseTask: %v", err)
|
||||
}
|
||||
if task.TaskID != "11111111-1111-4111-8111-111111111111" {
|
||||
t.Errorf("task id = %q", task.TaskID)
|
||||
}
|
||||
if task.Attempt != 1 || task.Workload != "similarity-search" {
|
||||
t.Errorf("attempt/workload = %d/%q", task.Attempt, task.Workload)
|
||||
}
|
||||
if task.Input.SHA256 != "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" {
|
||||
t.Errorf("sha256 = %q", task.Input.SHA256)
|
||||
}
|
||||
if task.LeaseExpiresAt.IsZero() {
|
||||
t.Error("lease must parse")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseTaskRejectsInvalidPayloads(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
mutate func(map[string]any)
|
||||
}{
|
||||
{"non-uuid task id", func(p map[string]any) { p["task_id"] = "../outside" }},
|
||||
{"zero attempt", func(p map[string]any) { p["attempt"] = 0 }},
|
||||
{"naive lease", func(p map[string]any) { p["lease_expires_at"] = "2026-08-02T00:00:00" }},
|
||||
{"network-path uri", func(p map[string]any) {
|
||||
p["input"].(map[string]any)["uri"] = "//outside.example/input"
|
||||
}},
|
||||
{"dot-segment uri", func(p map[string]any) {
|
||||
p["input"].(map[string]any)["uri"] = "/tasks/../outside/input"
|
||||
}},
|
||||
{"short sha256", func(p map[string]any) {
|
||||
p["input"].(map[string]any)["sha256"] = "abc"
|
||||
}},
|
||||
}
|
||||
for _, test := range tests {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
payload := validTaskPayload()
|
||||
test.mutate(payload)
|
||||
if _, err := ParseTask(payload); err == nil {
|
||||
t.Error("expected ParseTask to reject the payload")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseRegisteredAndUploaded(t *testing.T) {
|
||||
registered, err := ParseRegistered(map[string]any{
|
||||
"worker_id": "22222222-2222-4222-8222-222222222222",
|
||||
"heartbeat_interval_seconds": 15.0,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("ParseRegistered: %v", err)
|
||||
}
|
||||
if registered.HeartbeatIntervalSeconds != 15 {
|
||||
t.Errorf("interval = %v", registered.HeartbeatIntervalSeconds)
|
||||
}
|
||||
|
||||
uploaded, err := ParseUploaded(map[string]any{
|
||||
"artifact_id": "33333333-3333-4333-8333-333333333333",
|
||||
"uri": "https://coordinator.example/artifacts/333/download",
|
||||
"sha256": "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb",
|
||||
"size_bytes": 12.0,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("ParseUploaded: %v", err)
|
||||
}
|
||||
if uploaded.SizeBytes != 12 {
|
||||
t.Errorf("size = %d", uploaded.SizeBytes)
|
||||
}
|
||||
|
||||
if _, err := ParseUploaded(map[string]any{"artifact_id": "missing"}); err == nil {
|
||||
t.Error("expected invalid upload metadata to fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLeaseHeartbeatDelayIsBelowHalfTTL(t *testing.T) {
|
||||
task, err := ParseTask(validTaskPayload())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
task.LeaseExpiresAt = time.Now().Add(60 * time.Second)
|
||||
heartbeat := newLeaseHeartbeat(task, "worker", nil, 15*time.Second)
|
||||
delay := heartbeat.nextDelay()
|
||||
if delay > 30*time.Second || delay <= 0 {
|
||||
t.Errorf("delay = %v, want < 30s", delay)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var (
|
||||
// Go's regexp (RE2) has no lookbehind, so these patterns conservatively
|
||||
// anchor on the characters that typically precede a local path:
|
||||
// whitespace, quotes, parens, brackets, or the start of the message.
|
||||
windowsPathPattern = regexp.MustCompile(`[A-Za-z]:\\[^\s'"\],)]+`)
|
||||
posixPathPattern = regexp.MustCompile(`(^|[\s'"(\[=])/(?:[^\s'"\],)]+)`)
|
||||
)
|
||||
|
||||
// SanitizeErrorMessage keeps coordinator-visible failures useful without
|
||||
// exposing local paths. It mirrors the Python worker's sanitizer: local work
|
||||
// directories and absolute paths are redacted, and the message is truncated
|
||||
// to 300 characters.
|
||||
func SanitizeErrorMessage(message string, workDir string) string {
|
||||
message = strings.ReplaceAll(message, workDir, "<worker-dir>")
|
||||
message = windowsPathPattern.ReplaceAllString(message, "<path>")
|
||||
message = posixPathPattern.ReplaceAllString(message, "${1}<path>")
|
||||
if len(message) > 300 {
|
||||
message = message[:300]
|
||||
}
|
||||
return message
|
||||
}
|
||||
|
||||
// IsRetryableError classifies failures for the coordinator. Invalid scientific
|
||||
// input and missing local tools are permanent; everything else (transient
|
||||
// transport errors, unexpected runner failures) may be retried.
|
||||
func IsRetryableError(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var coordinatorErr *CoordinatorError
|
||||
var pathErr *os.PathError
|
||||
if errors.As(err, &coordinatorErr) || errors.As(err, &pathErr) {
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// TaskRunnerExit classifies subprocess exits.
|
||||
const (
|
||||
ExitPermanent = 3 // runner classified the failure as invalid input
|
||||
)
|
||||
|
||||
func runnerExitError(exit int, stderr string) error {
|
||||
if exit == ExitPermanent {
|
||||
return &CoordinatorError{msg: stderr}
|
||||
}
|
||||
return fmt.Errorf("task runner failed with exit code %d: %s", exit, stderr)
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// TaskRunner spawns the Python task entry and returns the sealed partial
|
||||
// artifact manifest it produced.
|
||||
type TaskRunner struct {
|
||||
command []string
|
||||
}
|
||||
|
||||
func NewTaskRunner(command []string) *TaskRunner {
|
||||
return &TaskRunner{command: command}
|
||||
}
|
||||
|
||||
// Run executes one task: the task payload is written as JSON into the attempt
|
||||
// directory, the Python entry computes and seals the partial, and the written
|
||||
// manifest is parsed back. stderr is captured for failure reporting.
|
||||
func (r *TaskRunner) Run(task *Task, taskDir string, manifestPath string, extraEnv []string) (*TaskRunnerManifest, error) {
|
||||
if err := os.MkdirAll(taskDir, 0o750); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
payload := map[string]any{
|
||||
"task_id": task.TaskID,
|
||||
"attempt": task.Attempt,
|
||||
"lease_expires_at": task.leaseExpiresRaw,
|
||||
"workload": task.Workload,
|
||||
"input": map[string]any{
|
||||
"uri": task.Input.URI,
|
||||
"sha256": task.Input.SHA256,
|
||||
},
|
||||
"parameters": task.Parameters,
|
||||
}
|
||||
payloadBytes, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
taskJSONPath := filepath.Join(taskDir, "task.json")
|
||||
if err := os.WriteFile(taskJSONPath, payloadBytes, 0o600); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
args := append([]string{}, r.command[1:]...)
|
||||
args = append(args,
|
||||
"--task-json", taskJSONPath,
|
||||
"--task-dir", taskDir,
|
||||
"--output", manifestPath,
|
||||
)
|
||||
// #nosec G204 -- the command comes from the operator-configured TASK_RUNNER.
|
||||
command := exec.CommandContext(context.Background(), r.command[0], args...)
|
||||
command.Dir = taskDir
|
||||
command.Env = append(os.Environ(), extraEnv...)
|
||||
var stderr bytes.Buffer
|
||||
command.Stderr = &stderr
|
||||
if err := command.Run(); err != nil {
|
||||
var exitErr *exec.ExitError
|
||||
if errors.As(err, &exitErr) {
|
||||
return nil, runnerExitError(exitErr.ExitCode(), stderr.String())
|
||||
}
|
||||
return nil, fmt.Errorf("task runner could not be started: %w", err)
|
||||
}
|
||||
// #nosec G304 -- the manifest path is inside the worker's own task directory.
|
||||
raw, err := os.ReadFile(manifestPath)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("task runner produced no result manifest")
|
||||
}
|
||||
var manifest TaskRunnerManifest
|
||||
if err := json.Unmarshal(raw, &manifest); err != nil {
|
||||
return nil, fmt.Errorf("task runner produced an invalid result manifest")
|
||||
}
|
||||
if manifest.ArtifactPath == "" || manifest.ContentType == "" {
|
||||
return nil, fmt.Errorf("task runner produced an incomplete result manifest")
|
||||
}
|
||||
info, err := os.Stat(manifest.ArtifactPath)
|
||||
if err != nil || !info.Mode().IsRegular() {
|
||||
return nil, fmt.Errorf("task runner produced no artifact file")
|
||||
}
|
||||
return &manifest, nil
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
// Package authctx carries the authenticated requester across the transport and
|
||||
// use-case layers without either one importing the other. The HTTP middleware
|
||||
// stamps a Requester after verifying a user's JWT; the job use cases read it to
|
||||
// record ownership and to enforce that a non-admin only touches their own jobs.
|
||||
package authctx
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Requester is the identity behind a request, derived from a verified JWT.
|
||||
// A request authenticated only by the shared worker/service token carries no
|
||||
// Requester at all (From returns ok=false), which is how worker traffic and
|
||||
// legacy unauthenticated-user traffic stay owner-less.
|
||||
type Requester struct {
|
||||
UserID uuid.UUID
|
||||
Role string
|
||||
Verified bool
|
||||
}
|
||||
|
||||
// IsAdmin reports whether the requester may act on any user's jobs.
|
||||
func (r Requester) IsAdmin() bool { return r.Role == "admin" }
|
||||
|
||||
// IsTrusted reports whether workers this requester registers produce results
|
||||
// the coordinator accepts without quorum. Admins and verified contributors are
|
||||
// trusted; a plain unverified user is not.
|
||||
func (r Requester) IsTrusted() bool { return r.IsAdmin() || r.Verified }
|
||||
|
||||
type ctxKey struct{}
|
||||
|
||||
// With returns a copy of ctx carrying r.
|
||||
func With(ctx context.Context, r Requester) context.Context {
|
||||
return context.WithValue(ctx, ctxKey{}, r)
|
||||
}
|
||||
|
||||
// From returns the requester stamped by the middleware, or ok=false when the
|
||||
// request was not authenticated as a user.
|
||||
func From(ctx context.Context) (Requester, bool) {
|
||||
r, ok := ctx.Value(ctxKey{}).(Requester)
|
||||
return r, ok
|
||||
}
|
||||
@@ -7,12 +7,19 @@ import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ErrNoRows is returned when the input has a header but no data rows: a job with
|
||||
// zero tasks could never complete, so it is rejected at the source.
|
||||
var ErrNoRows = fmt.Errorf("input has no data rows")
|
||||
|
||||
// maxShardBytes bounds the coordinator memory used by one in-progress shard.
|
||||
// The uploaded file may be much larger: it is first stored on disk, then split
|
||||
// in small bounded pieces. Operators can lower rowsPerShard when this limit is
|
||||
// reached rather than exhausting the coordinator process.
|
||||
const maxShardBytes = 64 << 20 // 64 MiB
|
||||
|
||||
// SplitTSV reads a header-plus-rows text stream and cuts it into shards of at
|
||||
// most rowsPerShard data rows. Every shard repeats the header, so a worker can
|
||||
// parse its shard in isolation. emit is called once per shard, in order, with a
|
||||
@@ -26,13 +33,25 @@ var ErrNoRows = fmt.Errorf("input has no data rows")
|
||||
// Only one shard is buffered at a time, so memory is bounded by shard size (a
|
||||
// worker-sized slice of the data), not by the size of the whole dataset.
|
||||
func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error {
|
||||
return SplitTSVLimit(r, rowsPerShard, 0, emit)
|
||||
return splitTSVLimit(r, rowsPerShard, 0, nil, emit)
|
||||
}
|
||||
|
||||
// SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows.
|
||||
// A maxRows value of zero means unlimited. This lets an operator make a small,
|
||||
// representative pipeline check without materialising a second dataset file.
|
||||
func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, maxRows, nil, emit)
|
||||
}
|
||||
|
||||
// SplitChEMBLTSVLimit is the coordinator's scientific-upload splitter. It
|
||||
// validates the two columns every local SciMesh workload requires before any
|
||||
// shard task is persisted, while generic SplitTSV remains reusable for future
|
||||
// non-chemistry workloads.
|
||||
func SplitChEMBLTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
|
||||
return splitTSVLimit(r, rowsPerShard, maxRows, validateChEMBLHeader, emit)
|
||||
}
|
||||
|
||||
func splitTSVLimit(r io.Reader, rowsPerShard, maxRows int, validateHeader func([]byte) error, emit func(index int, shard io.Reader) error) error {
|
||||
if rowsPerShard <= 0 {
|
||||
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
|
||||
}
|
||||
@@ -51,6 +70,11 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
|
||||
return ErrNoRows // completely empty input
|
||||
}
|
||||
header := append([]byte(nil), sc.Bytes()...)
|
||||
if validateHeader != nil {
|
||||
if err := validateHeader(header); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
var (
|
||||
buf bytes.Buffer
|
||||
@@ -71,9 +95,15 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
|
||||
|
||||
for sc.Scan() {
|
||||
if rows == 0 {
|
||||
if len(header)+1 > maxShardBytes {
|
||||
return fmt.Errorf("TSV header exceeds maximum shard size of %d bytes", maxShardBytes)
|
||||
}
|
||||
buf.Write(header)
|
||||
buf.WriteByte('\n')
|
||||
}
|
||||
if buf.Len()+len(sc.Bytes())+1 > maxShardBytes {
|
||||
return fmt.Errorf("shard exceeds maximum size of %d bytes; lower rowsPerShard", maxShardBytes)
|
||||
}
|
||||
buf.Write(sc.Bytes())
|
||||
buf.WriteByte('\n')
|
||||
rows++
|
||||
@@ -103,3 +133,17 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func validateChEMBLHeader(header []byte) error {
|
||||
seen := make(map[string]struct{})
|
||||
for _, field := range strings.Split(strings.TrimPrefix(string(header), "\ufeff"), "\t") {
|
||||
seen[field] = struct{}{}
|
||||
}
|
||||
if _, ok := seen["chembl_id"]; !ok {
|
||||
return fmt.Errorf("TSV is missing required column chembl_id")
|
||||
}
|
||||
if _, ok := seen["canonical_smiles"]; !ok {
|
||||
return fmt.Errorf("TSV is missing required column canonical_smiles")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -119,6 +119,14 @@ func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestChEMBLSplitRejectsMissingRequiredColumns(t *testing.T) {
|
||||
err := SplitChEMBLTSVLimit(strings.NewReader("id\tsmiles\nA\tCC\n"), 1, 0,
|
||||
func(int, io.Reader) error { return nil })
|
||||
if err == nil || !strings.Contains(err.Error(), "chembl_id") {
|
||||
t.Errorf("err = %v, want missing-column error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The scanned bytes are reused by bufio; the shard buffer must copy them, or a
|
||||
// later row would corrupt an earlier one. This guards that copy.
|
||||
func TestSplitDoesNotAliasScannerBuffer(t *testing.T) {
|
||||
|
||||
@@ -11,6 +11,7 @@ type JobStatus string
|
||||
const (
|
||||
JobPending JobStatus = "pending"
|
||||
JobRunning JobStatus = "running"
|
||||
JobReducing JobStatus = "reducing"
|
||||
JobCompleted JobStatus = "completed"
|
||||
JobFailed JobStatus = "failed"
|
||||
JobCancelled JobStatus = "cancelled"
|
||||
@@ -18,14 +19,22 @@ const (
|
||||
|
||||
// Job is one user submission that fans out into one or more tasks.
|
||||
type Job struct {
|
||||
ID uuid.UUID
|
||||
Workload string
|
||||
InputURI string // external input URI; empty for uploaded datasets
|
||||
InputArtifactID *uuid.UUID // uploaded input artifact; nil for URI submissions
|
||||
Parameters map[string]any
|
||||
Status JobStatus
|
||||
CreatedAt time.Time
|
||||
CompletedAt *time.Time
|
||||
ID uuid.UUID
|
||||
// OwnerID is the userservice user who submitted the job (JWT `sub`). nil
|
||||
// when the job was created without user authentication. Not a foreign key:
|
||||
// users live in a separate service/database.
|
||||
OwnerID *uuid.UUID
|
||||
Workload string
|
||||
InputURI string // external input URI; empty for uploaded datasets
|
||||
InputArtifactID *uuid.UUID // uploaded input artifact; nil for URI submissions
|
||||
ResultArtifactID *uuid.UUID
|
||||
Parameters map[string]any
|
||||
Status JobStatus
|
||||
CreatedAt time.Time
|
||||
CompletedAt *time.Time
|
||||
ReducerStartedAt *time.Time
|
||||
ErrorCode *string
|
||||
ErrorMessage *string
|
||||
}
|
||||
|
||||
// NewUploadedJob builds a job whose input was uploaded to the coordinator. The
|
||||
@@ -114,6 +123,12 @@ func (p JobProgress) DeriveStatus() JobStatus {
|
||||
switch {
|
||||
case p.Job.Status == JobCancelled:
|
||||
return JobCancelled
|
||||
case p.Job.Status == JobFailed:
|
||||
// A reducer may fail after every shard has completed. That terminal
|
||||
// failure must not be overwritten by an otherwise-complete task count.
|
||||
return JobFailed
|
||||
case p.Job.Status == JobReducing:
|
||||
return JobReducing
|
||||
case p.Total == 0:
|
||||
return JobPending
|
||||
case p.Done == p.Total:
|
||||
|
||||
@@ -82,6 +82,7 @@ func TestDeriveStatus(t *testing.T) {
|
||||
{"done and failed", JobProgress{Total: 3, Done: 2, Failed: 1}, JobFailed},
|
||||
{"failed but work remains", JobProgress{Total: 3, Pending: 1, Failed: 2}, JobRunning},
|
||||
{"cancelled job wins over task histogram", JobProgress{Job: Job{Status: JobCancelled}, Total: 3, Done: 1, Cancelled: 2}, JobCancelled},
|
||||
{"persisted reducer failure wins over completed tasks", JobProgress{Job: Job{Status: JobFailed}, Total: 3, Done: 3}, JobFailed},
|
||||
}
|
||||
for _, c := range cases {
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
|
||||
@@ -24,6 +24,10 @@ const (
|
||||
// ErrCodeLeaseExpired marks tasks failed by the reaper rather than by a worker.
|
||||
const ErrCodeLeaseExpired = "lease_expired"
|
||||
|
||||
// ErrCodeQuorumFailed marks a task whose untrusted results never reached a
|
||||
// verifying quorum before its attempts ran out.
|
||||
const ErrCodeQuorumFailed = "quorum_failed"
|
||||
|
||||
// Task is one independently executable chunk of a job.
|
||||
//
|
||||
// Nullable columns are pointers so "no lease" stays distinguishable from
|
||||
@@ -216,6 +220,33 @@ func (t *Task) CompleteWith(resultArtifactID uuid.UUID, metrics map[string]any,
|
||||
return nil
|
||||
}
|
||||
|
||||
// ReleaseAfterVote returns an untrusted worker's task to the queue after its
|
||||
// result was recorded as a quorum vote but quorum was not yet reached, so a
|
||||
// different owner can compute it independently. When no attempts remain the task
|
||||
// fails: its untrusted results could not be verified.
|
||||
func (t *Task) ReleaseAfterVote(worker string, attempt int, now time.Time) error {
|
||||
if t.Status == TaskCompleted {
|
||||
return nil // settled by a concurrent quorum
|
||||
}
|
||||
if err := t.verifyLease(worker, attempt, now); err != nil {
|
||||
return err
|
||||
}
|
||||
t.LeaseOwner = nil
|
||||
t.LeaseExpiresAt = nil
|
||||
t.Version++
|
||||
|
||||
if t.CanRetry() {
|
||||
t.Status = TaskPending
|
||||
return nil
|
||||
}
|
||||
code, msg := ErrCodeQuorumFailed, "untrusted results did not reach quorum"
|
||||
t.ErrorCode = &code
|
||||
t.ErrorMessage = &msg
|
||||
t.Status = TaskFailed
|
||||
t.CompletedAt = &now
|
||||
return nil
|
||||
}
|
||||
|
||||
// Fail records a worker-reported failure. A retryable failure with attempts
|
||||
// left returns the task to the queue; otherwise it terminates as failed.
|
||||
func (t *Task) Fail(worker string, attempt int, code, message string, retryable bool, now time.Time) error {
|
||||
|
||||
@@ -14,14 +14,30 @@ const (
|
||||
WorkerOffline WorkerStatus = "offline"
|
||||
)
|
||||
|
||||
// WorkerTrust says whether a worker's results are accepted directly or must
|
||||
// clear quorum cross-checking.
|
||||
type WorkerTrust string
|
||||
|
||||
const (
|
||||
// WorkerTrusted — lab machine (shared token) or a verified/admin contributor.
|
||||
WorkerTrusted WorkerTrust = "trusted"
|
||||
// WorkerUntrusted — a plain enthusiast; results are quarantined until quorum.
|
||||
WorkerUntrusted WorkerTrust = "untrusted"
|
||||
)
|
||||
|
||||
// Worker is a registered process/machine allowed to claim tasks. Its
|
||||
// capabilities are the allowlisted workload names it can run; the coordinator
|
||||
// never hands it a task outside that set.
|
||||
type Worker struct {
|
||||
ID uuid.UUID
|
||||
Name string
|
||||
Capabilities []string
|
||||
Status WorkerStatus
|
||||
ID uuid.UUID
|
||||
Name string
|
||||
Capabilities []string
|
||||
Status WorkerStatus
|
||||
// OwnerID is the userservice user who registered this worker; nil for a
|
||||
// worker registered with the shared service token.
|
||||
OwnerID *uuid.UUID
|
||||
// TrustLevel decides whether this worker's results need quorum.
|
||||
TrustLevel WorkerTrust
|
||||
LastHeartbeatAt time.Time
|
||||
CreatedAt time.Time
|
||||
UpdatedAt time.Time
|
||||
@@ -29,6 +45,9 @@ type Worker struct {
|
||||
|
||||
// NewWorker registers a worker. A worker with no capabilities could never be
|
||||
// handed a task, so an empty set is rejected rather than silently stored.
|
||||
//
|
||||
// Trust defaults to WorkerTrusted (the shared-token lab worker); the caller
|
||||
// overrides it for a volunteer registered through the userservice.
|
||||
func NewWorker(name string, capabilities []string, now time.Time) (*Worker, error) {
|
||||
if len(capabilities) == 0 {
|
||||
return nil, ErrInvalidInput
|
||||
@@ -38,6 +57,7 @@ func NewWorker(name string, capabilities []string, now time.Time) (*Worker, erro
|
||||
Name: name,
|
||||
Capabilities: capabilities,
|
||||
Status: WorkerOnline,
|
||||
TrustLevel: WorkerTrusted,
|
||||
LastHeartbeatAt: now,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
|
||||
@@ -26,6 +26,24 @@ type Config struct {
|
||||
Token string
|
||||
// Local operator UI credential. Empty disables the embedded UI entirely.
|
||||
UIToken string
|
||||
// Shared HS256 secret used to verify userservice-issued JWTs. When set, a
|
||||
// submitter may authenticate with a JWT (in addition to workers using the
|
||||
// shared token) and their jobs are stamped with owner_id. Empty disables
|
||||
// user-JWT auth entirely — the pre-userservice behaviour. Must match the
|
||||
// userservice's JWT_SECRET.
|
||||
JWTSecret string
|
||||
// Base URL of the userservice, e.g. http://userservice:8081. When set
|
||||
// together with JWTSecret, the operator UI authenticates via userservice
|
||||
// login/registration (cookie session) instead of the static UI_AUTH_TOKEN
|
||||
// basic auth. Empty keeps the basic-auth UI.
|
||||
UserserviceURL string
|
||||
// Browser-facing base URLs used to render the "add your machine" command on
|
||||
// the UI. They must be reachable from a user's own machine, which is not
|
||||
// necessarily the in-cluster address the coordinator uses for UserserviceURL.
|
||||
// PublicCoordinatorURL empty lets the page fall back to its own origin;
|
||||
// PublicUserserviceURL empty falls back to UserserviceURL.
|
||||
PublicCoordinatorURL string
|
||||
PublicUserserviceURL string
|
||||
|
||||
// Minimum log level: debug, info, warn, error.
|
||||
LogLevel string
|
||||
@@ -33,6 +51,9 @@ type Config struct {
|
||||
LogFile string
|
||||
// Directory where artifact bytes are stored.
|
||||
StorageDir string
|
||||
// Directory of the built MkDocs site (site/) served at /ui/docs/. Empty
|
||||
// disables the docs route; the UI shows a hint page instead.
|
||||
DocsDir string
|
||||
// Upper bound on an uploaded dataset or artifact body, in bytes.
|
||||
MaxUploadBytes int64
|
||||
|
||||
@@ -50,10 +71,17 @@ type Config struct {
|
||||
LeaseDuration time.Duration
|
||||
// Default attempt ceiling for newly created tasks.
|
||||
DefaultMaxAttempts int
|
||||
// How many distinct owners must agree on an untrusted result before it is
|
||||
// accepted (trusted workers are accepted directly).
|
||||
QuorumSize int
|
||||
// How often the background lease-reaper runs.
|
||||
ReaperInterval time.Duration
|
||||
// A worker silent for longer than this is marked offline by the reaper.
|
||||
WorkerOfflineAfter time.Duration
|
||||
// Whether the binary applies its embedded schema migrations on startup.
|
||||
// On by default so a downloaded binary provisions its own database; set
|
||||
// AUTO_MIGRATE=false when an operator manages migrations out of band.
|
||||
AutoMigrate bool
|
||||
}
|
||||
|
||||
// Load reads the environment and fails fast on anything required-but-missing
|
||||
@@ -78,20 +106,26 @@ func LoadConfig() (Config, error) {
|
||||
DatabaseURL: os.Getenv("DATABASE_URL"),
|
||||
// COORDINATOR_TOKEN is the contract name; WORKER_AUTH_TOKEN is the
|
||||
// former name, still honoured so existing .env files keep working.
|
||||
Token: getEnv("COORDINATOR_TOKEN", os.Getenv("WORKER_AUTH_TOKEN")),
|
||||
UIToken: os.Getenv("UI_AUTH_TOKEN"),
|
||||
LogLevel: getEnv("LOG_LEVEL", "info"),
|
||||
LogFile: os.Getenv("LOG_FILE"),
|
||||
StorageDir: getEnv("COORDINATOR_STORAGE_DIR", "./data"),
|
||||
MaxUploadBytes: 1 << 30, // 1 GiB
|
||||
DBMaxConns: 10,
|
||||
DBConnectTimeout: 30 * time.Second,
|
||||
RequestTimeout: 15 * time.Second,
|
||||
HeartbeatInterval: 15 * time.Second,
|
||||
LeaseDuration: 2 * time.Minute,
|
||||
DefaultMaxAttempts: 3,
|
||||
ReaperInterval: 30 * time.Second,
|
||||
WorkerOfflineAfter: 1 * time.Minute,
|
||||
Token: getEnv("COORDINATOR_TOKEN", os.Getenv("WORKER_AUTH_TOKEN")),
|
||||
UIToken: os.Getenv("UI_AUTH_TOKEN"),
|
||||
JWTSecret: os.Getenv("JWT_SECRET"),
|
||||
UserserviceURL: os.Getenv("USERSERVICE_URL"),
|
||||
PublicCoordinatorURL: os.Getenv("PUBLIC_COORDINATOR_URL"),
|
||||
PublicUserserviceURL: getEnv("PUBLIC_USERSERVICE_URL", os.Getenv("USERSERVICE_URL")),
|
||||
LogLevel: getEnv("LOG_LEVEL", "info"),
|
||||
LogFile: os.Getenv("LOG_FILE"),
|
||||
StorageDir: getEnv("COORDINATOR_STORAGE_DIR", "./data"),
|
||||
DocsDir: os.Getenv("SCIMESH_DOCS_DIR"),
|
||||
MaxUploadBytes: 1 << 30, // 1 GiB
|
||||
DBMaxConns: 10,
|
||||
DBConnectTimeout: 30 * time.Second,
|
||||
RequestTimeout: 15 * time.Second,
|
||||
HeartbeatInterval: 15 * time.Second,
|
||||
LeaseDuration: 2 * time.Minute,
|
||||
DefaultMaxAttempts: 3,
|
||||
QuorumSize: 2,
|
||||
ReaperInterval: 30 * time.Second,
|
||||
WorkerOfflineAfter: 1 * time.Minute,
|
||||
}
|
||||
|
||||
if cfg.DatabaseURL == "" {
|
||||
@@ -100,6 +134,11 @@ func LoadConfig() (Config, error) {
|
||||
if cfg.UIToken != "" && cfg.Token != "" && cfg.UIToken == cfg.Token {
|
||||
return Config{}, fmt.Errorf("UI_AUTH_TOKEN must differ from the worker auth token")
|
||||
}
|
||||
// A short secret makes the HMAC brute-forceable; refuse a weak one rather
|
||||
// than verify tokens against it.
|
||||
if cfg.JWTSecret != "" && len(cfg.JWTSecret) < 32 {
|
||||
return Config{}, fmt.Errorf("JWT_SECRET must be at least 32 bytes")
|
||||
}
|
||||
|
||||
var err error
|
||||
if cfg.DBMaxConns, err = getEnvInt32("DB_MAX_CONNS", cfg.DBMaxConns); err != nil {
|
||||
@@ -129,6 +168,23 @@ func LoadConfig() (Config, error) {
|
||||
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
if cfg.QuorumSize, err = getEnvInt("QUORUM_SIZE", cfg.QuorumSize); err != nil {
|
||||
return Config{}, err
|
||||
}
|
||||
if cfg.QuorumSize < 1 {
|
||||
return Config{}, fmt.Errorf("QUORUM_SIZE must be positive")
|
||||
}
|
||||
if cfg.DefaultMaxAttempts < 1 {
|
||||
return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive")
|
||||
}
|
||||
cfg.AutoMigrate = true
|
||||
if raw := os.Getenv("AUTO_MIGRATE"); raw != "" {
|
||||
parsed, err := strconv.ParseBool(raw)
|
||||
if err != nil {
|
||||
return Config{}, fmt.Errorf("AUTO_MIGRATE must be true or false")
|
||||
}
|
||||
cfg.AutoMigrate = parsed
|
||||
}
|
||||
|
||||
return cfg, nil
|
||||
}
|
||||
|
||||
@@ -32,3 +32,14 @@ func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) {
|
||||
t.Fatalf("unexpected tokens: %+v", cfg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadConfigRejectsNonPositiveDefaultMaxAttempts(t *testing.T) {
|
||||
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
|
||||
t.Setenv("DATABASE_URL", "postgres://test")
|
||||
t.Setenv("DEFAULT_MAX_ATTEMPTS", "0")
|
||||
|
||||
_, err := LoadConfig()
|
||||
if err == nil || !strings.Contains(err.Error(), "DEFAULT_MAX_ATTEMPTS") {
|
||||
t.Fatalf("LoadConfig error = %v, want default-attempt validation", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -160,17 +160,18 @@ func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time
|
||||
return cancelled, nil
|
||||
}
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var n int64
|
||||
affected := make([]uuid.UUID, 0)
|
||||
for _, t := range r.tasks {
|
||||
if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
|
||||
if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) &&
|
||||
t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
|
||||
t.ExpireLease(now)
|
||||
n++
|
||||
affected = append(affected, t.JobID)
|
||||
}
|
||||
}
|
||||
return n, nil
|
||||
return affected, nil
|
||||
}
|
||||
|
||||
// --- JobRepo -------------------------------------------------------------
|
||||
@@ -215,6 +216,51 @@ func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID, status domain.
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *JobRepo) ClaimReduction(_ context.Context, id uuid.UUID, startedAt time.Time) (bool, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
j, ok := r.jobs[id]
|
||||
if !ok {
|
||||
return false, domain.ErrJobNotFound
|
||||
}
|
||||
if j.Status != domain.JobReducing || j.ReducerStartedAt != nil {
|
||||
return false, nil
|
||||
}
|
||||
j.ReducerStartedAt = &startedAt
|
||||
return true, nil
|
||||
}
|
||||
|
||||
func (r *JobRepo) CompleteWithResult(_ context.Context, id, resultArtifactID uuid.UUID, completedAt time.Time) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
j, ok := r.jobs[id]
|
||||
if !ok {
|
||||
return domain.ErrJobNotFound
|
||||
}
|
||||
j.ResultArtifactID = &resultArtifactID
|
||||
j.Status = domain.JobCompleted
|
||||
j.CompletedAt = &completedAt
|
||||
j.ReducerStartedAt = nil
|
||||
j.ErrorCode = nil
|
||||
j.ErrorMessage = nil
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *JobRepo) FailReduction(_ context.Context, id uuid.UUID, code, message string, completedAt time.Time) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
j, ok := r.jobs[id]
|
||||
if !ok {
|
||||
return domain.ErrJobNotFound
|
||||
}
|
||||
j.Status = domain.JobFailed
|
||||
j.CompletedAt = &completedAt
|
||||
j.ReducerStartedAt = nil
|
||||
j.ErrorCode = &code
|
||||
j.ErrorMessage = &message
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- WorkerRepo ----------------------------------------------------------
|
||||
|
||||
type WorkerRepo struct {
|
||||
@@ -298,6 +344,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
|
||||
return &cp, nil
|
||||
}
|
||||
|
||||
func (r *ArtifactRepo) FindPartialResult(_ context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
for _, a := range r.arts {
|
||||
if a.TaskID != nil && *a.TaskID == taskID && a.Kind == domain.ArtifactPartialResult &&
|
||||
a.Attempt != nil && *a.Attempt == attempt {
|
||||
return cloneArtifact(a), nil
|
||||
}
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func cloneArtifact(a *domain.Artifact) *domain.Artifact {
|
||||
cp := *a
|
||||
return &cp
|
||||
}
|
||||
|
||||
// --- BlobStore -----------------------------------------------------------
|
||||
|
||||
type BlobStore struct {
|
||||
@@ -354,3 +417,36 @@ func contains(ss []string, s string) bool {
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// TaskResultRepo is an in-memory usecase.TaskResultRepository: one vote per
|
||||
// (task, owner).
|
||||
type TaskResultRepo struct {
|
||||
mu sync.Mutex
|
||||
votes map[uuid.UUID]map[uuid.UUID]string // taskID -> ownerID -> sha256
|
||||
}
|
||||
|
||||
func NewTaskResultRepo() *TaskResultRepo {
|
||||
return &TaskResultRepo{votes: make(map[uuid.UUID]map[uuid.UUID]string)}
|
||||
}
|
||||
|
||||
func (r *TaskResultRepo) RecordVote(_ context.Context, taskID, ownerID uuid.UUID, sha256 string, _ uuid.UUID) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if r.votes[taskID] == nil {
|
||||
r.votes[taskID] = make(map[uuid.UUID]string)
|
||||
}
|
||||
r.votes[taskID][ownerID] = sha256
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *TaskResultRepo) CountAgreeing(_ context.Context, taskID uuid.UUID, sha256 string) (int, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
n := 0
|
||||
for _, s := range r.votes[taskID] {
|
||||
if s == sha256 {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
@@ -27,7 +27,7 @@ var _ usecase.UIReadRepository = (*UIReadRepo)(nil)
|
||||
func (r *UIReadRepo) GetJob(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
|
||||
return r.jobs.Get(ctx, id)
|
||||
}
|
||||
func (r *UIReadRepo) ListJobs(_ context.Context, limit int) ([]domain.Job, error) {
|
||||
func (r *UIReadRepo) ListJobs(_ context.Context, owner *uuid.UUID, limit int) ([]domain.Job, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
@@ -35,6 +35,9 @@ func (r *UIReadRepo) ListJobs(_ context.Context, limit int) ([]domain.Job, error
|
||||
defer r.jobs.mu.Unlock()
|
||||
out := make([]domain.Job, 0, len(r.jobs.jobs))
|
||||
for _, job := range r.jobs.jobs {
|
||||
if owner != nil && (job.OwnerID == nil || *job.OwnerID != *owner) {
|
||||
continue
|
||||
}
|
||||
out = append(out, *job)
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
@@ -60,6 +63,18 @@ func (r *UIReadRepo) ListTasksByJob(_ context.Context, jobID uuid.UUID) ([]domai
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
|
||||
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
|
||||
for _, id := range jobIDs {
|
||||
tasks, err := r.ListTasksByJob(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[id] = tasks
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
@@ -72,16 +87,44 @@ func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker,
|
||||
copy.Capabilities = append([]string(nil), worker.Capabilities...)
|
||||
out = append(out, copy)
|
||||
}
|
||||
sortWorkers(out)
|
||||
if len(out) > limit {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListWorkersByOwner(_ context.Context, owner uuid.UUID, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
r.workers.mu.Lock()
|
||||
defer r.workers.mu.Unlock()
|
||||
out := []domain.Worker{}
|
||||
for _, worker := range r.workers.workers {
|
||||
if worker.OwnerID == nil || *worker.OwnerID != owner {
|
||||
continue
|
||||
}
|
||||
copy := *worker
|
||||
copy.Capabilities = append([]string(nil), worker.Capabilities...)
|
||||
out = append(out, copy)
|
||||
}
|
||||
sortWorkers(out)
|
||||
if len(out) > limit {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// sortWorkers orders workers most-recently-seen first, breaking ties on id so
|
||||
// the order is deterministic across calls.
|
||||
func sortWorkers(out []domain.Worker) {
|
||||
sort.Slice(out, func(i, j int) bool {
|
||||
if out[i].LastHeartbeatAt.Equal(out[j].LastHeartbeatAt) {
|
||||
return out[i].ID.String() > out[j].ID.String()
|
||||
}
|
||||
return out[i].LastHeartbeatAt.After(out[j].LastHeartbeatAt)
|
||||
})
|
||||
if len(out) > limit {
|
||||
out = out[:limit]
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (r *UIReadRepo) ListArtifactsByJob(_ context.Context, jobID uuid.UUID) ([]domain.Artifact, error) {
|
||||
r.artifacts.mu.Lock()
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
package metrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
)
|
||||
|
||||
// Stats is a point-in-time snapshot of the coordinator's domain state: counts of
|
||||
// tasks, jobs, and workers keyed by their status. Maps are expected to be
|
||||
// zero-filled by the provider so every known status is always present, giving
|
||||
// the dashboard flat zero lines instead of gaps.
|
||||
type Stats struct {
|
||||
Tasks map[string]int
|
||||
Jobs map[string]int
|
||||
Workers map[string]int
|
||||
}
|
||||
|
||||
// StatsFunc returns the current snapshot. It is called on every scrape, so it
|
||||
// must be a cheap aggregate query.
|
||||
type StatsFunc func(context.Context) (Stats, error)
|
||||
|
||||
// RegisterBusiness registers a collector that reports domain-state gauges
|
||||
// (scimesh_tasks/jobs/workers by status) sourced from collect on each scrape.
|
||||
// Deriving the gauges at scrape time keeps them fresh without a background
|
||||
// goroutine, and a failed query simply yields no samples for that scrape.
|
||||
func (m *Metrics) RegisterBusiness(collect StatsFunc) {
|
||||
m.reg.MustRegister(&businessCollector{
|
||||
collect: collect,
|
||||
tasks: prometheus.NewDesc("scimesh_tasks", "Tasks by status.", []string{"status"}, nil),
|
||||
jobs: prometheus.NewDesc("scimesh_jobs", "Jobs by status.", []string{"status"}, nil),
|
||||
workers: prometheus.NewDesc("scimesh_workers", "Workers by status.", []string{"status"}, nil),
|
||||
})
|
||||
}
|
||||
|
||||
type businessCollector struct {
|
||||
collect StatsFunc
|
||||
tasks, jobs, workers *prometheus.Desc
|
||||
}
|
||||
|
||||
func (c *businessCollector) Describe(ch chan<- *prometheus.Desc) {
|
||||
ch <- c.tasks
|
||||
ch <- c.jobs
|
||||
ch <- c.workers
|
||||
}
|
||||
|
||||
func (c *businessCollector) Collect(ch chan<- prometheus.Metric) {
|
||||
// A bounded query so one slow scrape cannot stall Prometheus.
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
defer cancel()
|
||||
|
||||
s, err := c.collect(ctx)
|
||||
if err != nil {
|
||||
return // no samples this scrape; Prometheus keeps the last value
|
||||
}
|
||||
emit(ch, c.tasks, s.Tasks)
|
||||
emit(ch, c.jobs, s.Jobs)
|
||||
emit(ch, c.workers, s.Workers)
|
||||
}
|
||||
|
||||
func emit(ch chan<- prometheus.Metric, desc *prometheus.Desc, counts map[string]int) {
|
||||
for status, n := range counts {
|
||||
ch <- prometheus.MustNewConstMetric(desc, prometheus.GaugeValue, float64(n), status)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,51 @@
|
||||
package metrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func scrape(t *testing.T, m *Metrics) string {
|
||||
t.Helper()
|
||||
rec := httptest.NewRecorder()
|
||||
req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, "/metrics", nil)
|
||||
m.Handler().ServeHTTP(rec, req)
|
||||
return rec.Body.String()
|
||||
}
|
||||
|
||||
func TestBusinessCollectorEmitsGauges(t *testing.T) {
|
||||
m := New()
|
||||
m.RegisterBusiness(func(context.Context) (Stats, error) {
|
||||
return Stats{
|
||||
Tasks: map[string]int{"pending": 3, "running": 1, "completed": 0},
|
||||
Jobs: map[string]int{"running": 2},
|
||||
Workers: map[string]int{"online": 4},
|
||||
}, nil
|
||||
})
|
||||
|
||||
body := scrape(t, m)
|
||||
for _, want := range []string{
|
||||
`scimesh_tasks{status="pending"} 3`,
|
||||
`scimesh_tasks{status="completed"} 0`,
|
||||
`scimesh_jobs{status="running"} 2`,
|
||||
`scimesh_workers{status="online"} 4`,
|
||||
} {
|
||||
if !strings.Contains(body, want) {
|
||||
t.Errorf("metrics missing %q\n%s", want, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBusinessCollectorSkipsOnError(t *testing.T) {
|
||||
m := New()
|
||||
m.RegisterBusiness(func(context.Context) (Stats, error) {
|
||||
return Stats{}, errors.New("db down")
|
||||
})
|
||||
if strings.Contains(scrape(t, m), "scimesh_tasks") {
|
||||
t.Error("a failed snapshot must emit no business samples")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
// Package metrics exposes Prometheus instrumentation for the coordinator: an
|
||||
// HTTP RED middleware (rate, errors, duration) plus the standard Go runtime and
|
||||
// process collectors, all on a private registry so nothing leaks in from global
|
||||
// state.
|
||||
package metrics
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/prometheus/client_golang/prometheus/collectors"
|
||||
"github.com/prometheus/client_golang/prometheus/promhttp"
|
||||
)
|
||||
|
||||
type Metrics struct {
|
||||
reg *prometheus.Registry
|
||||
requests *prometheus.CounterVec
|
||||
duration *prometheus.HistogramVec
|
||||
}
|
||||
|
||||
// New builds the registry and registers the runtime, process, and HTTP metrics.
|
||||
func New() *Metrics {
|
||||
reg := prometheus.NewRegistry()
|
||||
reg.MustRegister(
|
||||
collectors.NewGoCollector(),
|
||||
collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
|
||||
)
|
||||
|
||||
requests := prometheus.NewCounterVec(prometheus.CounterOpts{
|
||||
Namespace: "scimesh",
|
||||
Subsystem: "http",
|
||||
Name: "requests_total",
|
||||
Help: "HTTP requests, labelled by method, normalized route, and status.",
|
||||
}, []string{"method", "route", "status"})
|
||||
|
||||
duration := prometheus.NewHistogramVec(prometheus.HistogramOpts{
|
||||
Namespace: "scimesh",
|
||||
Subsystem: "http",
|
||||
Name: "request_duration_seconds",
|
||||
Help: "HTTP request duration in seconds.",
|
||||
Buckets: prometheus.DefBuckets,
|
||||
}, []string{"method", "route"})
|
||||
|
||||
reg.MustRegister(requests, duration)
|
||||
return &Metrics{reg: reg, requests: requests, duration: duration}
|
||||
}
|
||||
|
||||
// Handler serves the metrics in Prometheus text format.
|
||||
func (m *Metrics) Handler() http.Handler {
|
||||
return promhttp.HandlerFor(m.reg, promhttp.HandlerOpts{})
|
||||
}
|
||||
|
||||
// Registry exposes the registry so callers can register extra collectors.
|
||||
func (m *Metrics) Registry() *prometheus.Registry { return m.reg }
|
||||
|
||||
// Middleware records one request into the RED metrics. It normalizes the path
|
||||
// so per-id routes collapse to a single low-cardinality label.
|
||||
func (m *Metrics) Middleware(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
start := time.Now()
|
||||
rec := &statusRecorder{ResponseWriter: w, status: http.StatusOK}
|
||||
next.ServeHTTP(rec, r)
|
||||
|
||||
route := normalizeRoute(r.URL.Path)
|
||||
m.requests.WithLabelValues(r.Method, route, strconv.Itoa(rec.status)).Inc()
|
||||
m.duration.WithLabelValues(r.Method, route).Observe(time.Since(start).Seconds())
|
||||
})
|
||||
}
|
||||
|
||||
type statusRecorder struct {
|
||||
http.ResponseWriter
|
||||
status int
|
||||
}
|
||||
|
||||
func (s *statusRecorder) WriteHeader(code int) {
|
||||
s.status = code
|
||||
s.ResponseWriter.WriteHeader(code)
|
||||
}
|
||||
|
||||
var uuidRe = regexp.MustCompile(`^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$`)
|
||||
|
||||
// normalizeRoute collapses uuid and numeric path segments to {id}, keeping the
|
||||
// route label cardinality bounded (otherwise every job/task id would be its own
|
||||
// time series).
|
||||
func normalizeRoute(path string) string {
|
||||
if path == "" {
|
||||
return "/"
|
||||
}
|
||||
segs := strings.Split(path, "/")
|
||||
for i, s := range segs {
|
||||
if s == "" {
|
||||
continue
|
||||
}
|
||||
if uuidRe.MatchString(s) || isAllDigits(s) {
|
||||
segs[i] = "{id}"
|
||||
}
|
||||
}
|
||||
return strings.Join(segs, "/")
|
||||
}
|
||||
|
||||
func isAllDigits(s string) bool {
|
||||
for _, r := range s {
|
||||
if r < '0' || r > '9' {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return s != ""
|
||||
}
|
||||
@@ -0,0 +1,47 @@
|
||||
package metrics
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNormalizeRoute(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"/health": "/health",
|
||||
"/jobs/3f2504e0-4f89-41d3-9a0c-0305e82c3301": "/jobs/{id}",
|
||||
"/tasks/3f2504e0-4f89-41d3-9a0c-0305e82c3301/result": "/tasks/{id}/result",
|
||||
"/ui/jobs/12345": "/ui/jobs/{id}",
|
||||
"/": "/",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := normalizeRoute(in); got != want {
|
||||
t.Errorf("normalizeRoute(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMiddlewareAndHandler(t *testing.T) {
|
||||
m := New()
|
||||
h := m.Middleware(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
}))
|
||||
|
||||
req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, "/jobs/3f2504e0-4f89-41d3-9a0c-0305e82c3301", nil)
|
||||
h.ServeHTTP(httptest.NewRecorder(), req)
|
||||
|
||||
// Scrape and confirm the request was recorded under the normalized route.
|
||||
rec := httptest.NewRecorder()
|
||||
greq, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, "/metrics", nil)
|
||||
m.Handler().ServeHTTP(rec, greq)
|
||||
|
||||
body := rec.Body.String()
|
||||
if !strings.Contains(body, `scimesh_http_requests_total{method="POST",route="/jobs/{id}",status="201"}`) {
|
||||
t.Errorf("requests_total not recorded as expected; body:\n%s", body)
|
||||
}
|
||||
if !strings.Contains(body, "go_goroutines") {
|
||||
t.Error("Go runtime collector not registered")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
// Package reducer contains deterministic, coordinator-side result reductions.
|
||||
package reducer
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/csv"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
)
|
||||
|
||||
// ReduceOrderedConcat concatenates worker partial tables in shard order into a
|
||||
// single table with one header. Every partial must carry the same header as the
|
||||
// first partial and rows of the same width; anything else fails the job closed.
|
||||
func ReduceOrderedConcat(partials []io.Reader) ([]byte, error) {
|
||||
var out bytes.Buffer
|
||||
writer := csv.NewWriter(&out)
|
||||
var firstHeader []string
|
||||
for _, partial := range partials {
|
||||
reader := csv.NewReader(partial)
|
||||
header, err := reader.Read()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read partial header: %w", err)
|
||||
}
|
||||
if len(header) == 0 {
|
||||
return nil, fmt.Errorf("partial result has an empty header")
|
||||
}
|
||||
if firstHeader == nil {
|
||||
firstHeader = header
|
||||
if err := writer.Write(header); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
} else if !equalStrings(header, firstHeader) {
|
||||
return nil, fmt.Errorf("partial result has an inconsistent header")
|
||||
}
|
||||
for {
|
||||
row, err := reader.Read()
|
||||
if errors.Is(err, io.EOF) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read partial row: %w", err)
|
||||
}
|
||||
if len(row) != len(header) {
|
||||
return nil, fmt.Errorf("partial result has a row with an inconsistent width")
|
||||
}
|
||||
if err := writer.Write(row); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
writer.Flush()
|
||||
if err := writer.Error(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out.Bytes(), nil
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package reducer
|
||||
|
||||
import (
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestReduceOrderedConcatJoinsPartialsInOrderWithOneHeader(t *testing.T) {
|
||||
first := strings.NewReader("chembl_id,canonical_smiles\nA,CC\nB,CCC\n")
|
||||
second := strings.NewReader("chembl_id,canonical_smiles\nC,CCCC\n")
|
||||
|
||||
output, err := ReduceOrderedConcat([]io.Reader{first, second})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "chembl_id,canonical_smiles\nA,CC\nB,CCC\nC,CCCC\n"
|
||||
if string(output) != want {
|
||||
t.Fatalf("output = %q, want %q", output, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceOrderedConcatIsDeterministicAcrossInputOrder(t *testing.T) {
|
||||
left := strings.NewReader("id,rows\nA,1\nB,2\n")
|
||||
right := strings.NewReader("id,rows\nC,3\n")
|
||||
|
||||
first, err := ReduceOrderedConcat([]io.Reader{left, right})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
left, right = strings.NewReader("id,rows\nA,1\nB,2\n"), strings.NewReader("id,rows\nC,3\n")
|
||||
second, err := ReduceOrderedConcat([]io.Reader{left, right})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(first) != string(second) {
|
||||
t.Fatalf("concat is not deterministic: %q != %q", first, second)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceOrderedConcatRejectsInconsistentHeaders(t *testing.T) {
|
||||
first := strings.NewReader("a,b\n1,2\n")
|
||||
second := strings.NewReader("a,c\n1,2\n")
|
||||
if _, err := ReduceOrderedConcat([]io.Reader{first, second}); err == nil {
|
||||
t.Fatal("inconsistent headers must fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceOrderedConcatRejectsRaggedRows(t *testing.T) {
|
||||
partial := strings.NewReader("a,b\n1,2,3\n")
|
||||
if _, err := ReduceOrderedConcat([]io.Reader{partial}); err == nil {
|
||||
t.Fatal("ragged rows must fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceOrderedConcatEmptyPartials(t *testing.T) {
|
||||
output, err := ReduceOrderedConcat(nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(output) != 0 {
|
||||
t.Fatalf("empty input must produce empty output, got %q", output)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
// Package reducer contains deterministic, coordinator-side result reductions.
|
||||
package reducer
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/csv"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"sort"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
var searchHeader = []string{"rank", "chembl_id", "canonical_smiles", "similarity"}
|
||||
|
||||
type similarityMatch struct {
|
||||
similarity float64
|
||||
id string
|
||||
smiles string
|
||||
}
|
||||
|
||||
// ReduceSimilaritySearch streams worker-local top-k CSVs into the exact global
|
||||
// top-k. Each partial is validated before it can affect the final artifact.
|
||||
func ReduceSimilaritySearch(partials []io.Reader, parameters map[string]any) ([]byte, error) {
|
||||
topK, err := positiveInt(parameters["top_k"], 20)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
direction, err := thresholdDirection(parameters["threshold_direction"])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
h := &matchHeap{direction: direction}
|
||||
for _, partial := range partials {
|
||||
if err := readPartial(partial, direction, func(match similarityMatch) {
|
||||
if len(h.items) < topK {
|
||||
heapPush(h, match)
|
||||
return
|
||||
}
|
||||
if better(match, h.items[0], direction) {
|
||||
h.items[0] = match
|
||||
heapDown(h, 0)
|
||||
}
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
matches := append([]similarityMatch(nil), h.items...)
|
||||
sort.Slice(matches, func(i, j int) bool { return better(matches[i], matches[j], direction) })
|
||||
var out bytes.Buffer
|
||||
writer := csv.NewWriter(&out)
|
||||
if err := writer.Write(searchHeader); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
for index, match := range matches {
|
||||
if err := writer.Write([]string{
|
||||
strconv.Itoa(index + 1), match.id, match.smiles, fmt.Sprintf("%.6f", match.similarity),
|
||||
}); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
writer.Flush()
|
||||
if err := writer.Error(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return out.Bytes(), nil
|
||||
}
|
||||
|
||||
func readPartial(input io.Reader, direction string, consume func(similarityMatch)) error {
|
||||
reader := csv.NewReader(input)
|
||||
header, err := reader.Read()
|
||||
if err != nil {
|
||||
return fmt.Errorf("read partial header: %w", err)
|
||||
}
|
||||
if !equalStrings(header, searchHeader) {
|
||||
return fmt.Errorf("partial result has an invalid CSV header")
|
||||
}
|
||||
var previous *similarityMatch
|
||||
for rank := 1; ; rank++ {
|
||||
row, err := reader.Read()
|
||||
if errors.Is(err, io.EOF) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("read partial row: %w", err)
|
||||
}
|
||||
if len(row) != len(searchHeader) || row[0] != strconv.Itoa(rank) {
|
||||
return fmt.Errorf("partial result has an invalid rank")
|
||||
}
|
||||
score, err := strconv.ParseFloat(row[3], 64)
|
||||
if err != nil || math.IsNaN(score) || math.IsInf(score, 0) || score < 0 || score > 1 {
|
||||
return fmt.Errorf("partial result has an invalid similarity")
|
||||
}
|
||||
match := similarityMatch{similarity: score, id: row[1], smiles: row[2]}
|
||||
if previous != nil && better(match, *previous, direction) {
|
||||
return fmt.Errorf("partial result is not sorted deterministically")
|
||||
}
|
||||
previous = &match
|
||||
consume(match)
|
||||
}
|
||||
}
|
||||
|
||||
func positiveInt(value any, fallback int) (int, error) {
|
||||
if value == nil {
|
||||
return fallback, nil
|
||||
}
|
||||
switch n := value.(type) {
|
||||
case int:
|
||||
if n > 0 {
|
||||
return n, nil
|
||||
}
|
||||
case int64:
|
||||
if n > 0 && n <= math.MaxInt {
|
||||
return int(n), nil
|
||||
}
|
||||
case float64:
|
||||
if n > 0 && n == math.Trunc(n) && n <= math.MaxInt {
|
||||
return int(n), nil
|
||||
}
|
||||
}
|
||||
return 0, fmt.Errorf("top_k must be a positive integer")
|
||||
}
|
||||
|
||||
func thresholdDirection(value any) (string, error) {
|
||||
if value == nil {
|
||||
return "greater", nil
|
||||
}
|
||||
direction, ok := value.(string)
|
||||
if !ok || (direction != "greater" && direction != "less") {
|
||||
return "", fmt.Errorf("threshold_direction must be greater or less")
|
||||
}
|
||||
return direction, nil
|
||||
}
|
||||
|
||||
func better(left, right similarityMatch, direction string) bool {
|
||||
if left.similarity != right.similarity {
|
||||
if direction == "less" {
|
||||
return left.similarity < right.similarity
|
||||
}
|
||||
return left.similarity > right.similarity
|
||||
}
|
||||
if left.id != right.id {
|
||||
return left.id < right.id
|
||||
}
|
||||
return left.smiles < right.smiles
|
||||
}
|
||||
|
||||
func equalStrings(left, right []string) bool {
|
||||
if len(left) != len(right) {
|
||||
return false
|
||||
}
|
||||
for index := range left {
|
||||
if left[index] != right[index] {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// matchHeap keeps the worst retained match at index zero.
|
||||
type matchHeap struct {
|
||||
items []similarityMatch
|
||||
direction string
|
||||
}
|
||||
|
||||
func heapPush(h *matchHeap, value similarityMatch) {
|
||||
h.items = append(h.items, value)
|
||||
for child := len(h.items) - 1; child > 0; {
|
||||
parent := (child - 1) / 2
|
||||
if !better(h.items[parent], h.items[child], h.direction) {
|
||||
break
|
||||
}
|
||||
h.items[parent], h.items[child] = h.items[child], h.items[parent]
|
||||
child = parent
|
||||
}
|
||||
}
|
||||
|
||||
func heapDown(h *matchHeap, parent int) {
|
||||
for {
|
||||
child := parent*2 + 1
|
||||
if child >= len(h.items) {
|
||||
return
|
||||
}
|
||||
if right := child + 1; right < len(h.items) && better(h.items[child], h.items[right], h.direction) {
|
||||
child = right
|
||||
}
|
||||
if !better(h.items[parent], h.items[child], h.direction) {
|
||||
return
|
||||
}
|
||||
h.items[parent], h.items[child] = h.items[child], h.items[parent]
|
||||
parent = child
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
package reducer
|
||||
|
||||
import (
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestReduceSimilaritySearchKeepsExactCrossShardRanking(t *testing.T) {
|
||||
first := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,B,CCC,0.50000048\n2,C,CCCC,0.1\n")
|
||||
second := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.50000049\n")
|
||||
|
||||
output, err := ReduceSimilaritySearch([]io.Reader{first, second}, map[string]any{"top_k": 2})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.500000\n2,B,CCC,0.500000\n"
|
||||
if string(output) != want {
|
||||
t.Fatalf("output = %q, want %q", output, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceSimilaritySearchSupportsLeastSimilarDirection(t *testing.T) {
|
||||
partial := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.1\n2,B,CCC,0.8\n")
|
||||
output, err := ReduceSimilaritySearch([]io.Reader{partial}, map[string]any{
|
||||
"top_k": 1, "threshold_direction": "less",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got, want := string(output), "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.100000\n"; got != want {
|
||||
t.Fatalf("output = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReduceSimilaritySearchRejectsMalformedPartial(t *testing.T) {
|
||||
partial := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n2,A,CC,0.1\n")
|
||||
if _, err := ReduceSimilaritySearch([]io.Reader{partial}, nil); err == nil {
|
||||
t.Fatal("expected malformed rank error")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,239 @@
|
||||
// Package setup implements the `coordinator setup` wizard: database reachability
|
||||
// and creation, embedded schema migration, secret generation, and .env writing.
|
||||
// The wizard never logs or echoes secrets.
|
||||
package setup
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgconn"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/storage/postgres"
|
||||
)
|
||||
|
||||
// Options configures one wizard run.
|
||||
type Options struct {
|
||||
// DatabaseURL is the target coordinator database (pgx/libpq URL).
|
||||
DatabaseURL string
|
||||
// AdminDatabaseURL, when set, is used to create a missing target database.
|
||||
// Defaults to the target URL with the database name replaced by "postgres".
|
||||
AdminDatabaseURL string
|
||||
// EnvFile is where the generated settings are written (default ".env").
|
||||
EnvFile string
|
||||
// Force overwrites an existing EnvFile.
|
||||
Force bool
|
||||
// Yes disables interactive prompts; missing values fail instead.
|
||||
Yes bool
|
||||
// ConnectTimeout bounds the reachability check.
|
||||
ConnectTimeout time.Duration
|
||||
// Out receives progress and summary output; In feeds interactive answers.
|
||||
Out io.Writer
|
||||
In io.Reader
|
||||
}
|
||||
|
||||
// Run executes the wizard and returns a summary of what was done.
|
||||
func Run(ctx context.Context, options Options) (string, error) {
|
||||
if options.DatabaseURL == "" {
|
||||
return "", fmt.Errorf("DATABASE_URL is required (or pass --db)")
|
||||
}
|
||||
if options.EnvFile == "" {
|
||||
options.EnvFile = ".env"
|
||||
}
|
||||
if options.ConnectTimeout <= 0 {
|
||||
options.ConnectTimeout = 5 * time.Second
|
||||
}
|
||||
if options.Out == nil {
|
||||
options.Out = os.Stdout
|
||||
}
|
||||
|
||||
report := func(format string, args ...any) {
|
||||
_, _ = fmt.Fprintf(options.Out, format+"\n", args...)
|
||||
}
|
||||
|
||||
report("SciMesh coordinator setup")
|
||||
report("")
|
||||
|
||||
// 1. Reachability, with optional database creation.
|
||||
target, err := pgx.ParseConfig(options.DatabaseURL)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("DATABASE_URL is not a valid postgres URL: %w", err)
|
||||
}
|
||||
if err := probeDatabase(ctx, target, options.ConnectTimeout); err != nil {
|
||||
var pgErr *pgconn.PgError
|
||||
if !errors.As(err, &pgErr) || pgErr.Code != "3D000" {
|
||||
return "", fmt.Errorf("cannot reach the coordinator database: %w", err)
|
||||
}
|
||||
report("database %q does not exist yet", target.Database)
|
||||
admin, err := resolveAdminConfig(options, target)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := createDatabase(ctx, admin, target.Database, options.ConnectTimeout); err != nil {
|
||||
return "", fmt.Errorf("cannot create database %q: %w", target.Database, err)
|
||||
}
|
||||
report("created database %q", target.Database)
|
||||
}
|
||||
report("database %q is reachable", target.Database)
|
||||
|
||||
// 2. Apply the embedded schema migrations (idempotent).
|
||||
if err := postgres.Migrate(ctx, options.DatabaseURL, nil); err != nil {
|
||||
return "", fmt.Errorf("apply schema migrations: %w", err)
|
||||
}
|
||||
report("schema migrations applied")
|
||||
|
||||
// 3. JWT secret: reuse the environment value when strong, else generate.
|
||||
secret := os.Getenv("JWT_SECRET")
|
||||
if secret != "" && len(secret) < 32 {
|
||||
return "", fmt.Errorf("JWT_SECRET must be at least 32 bytes")
|
||||
}
|
||||
if secret == "" {
|
||||
generated, err := generateSecret()
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("generate JWT_SECRET: %w", err)
|
||||
}
|
||||
secret = generated
|
||||
report("generated a fresh JWT_SECRET")
|
||||
}
|
||||
|
||||
// 4. Write the .env file.
|
||||
storageDir := os.Getenv("COORDINATOR_STORAGE_DIR")
|
||||
if storageDir == "" {
|
||||
storageDir = "./data"
|
||||
}
|
||||
if err := writeEnvFile(options, secret, storageDir); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
// 5. Summary.
|
||||
var summary strings.Builder
|
||||
fmt.Fprintf(&summary, "Setup complete.\n\n")
|
||||
fmt.Fprintf(&summary, "Ready:\n")
|
||||
fmt.Fprintf(&summary, " - database %s is reachable and migrated\n", target.Database)
|
||||
fmt.Fprintf(&summary, " - settings written to %s (chmod 0600)\n", options.EnvFile)
|
||||
fmt.Fprintf(&summary, "\nStart the coordinator:\n")
|
||||
fmt.Fprintf(&summary, " ENV_FILE=%s ./coordinator\n", options.EnvFile)
|
||||
fmt.Fprintf(&summary, "\nOptional — userservice for UI logins (must share JWT_SECRET):\n")
|
||||
fmt.Fprintf(&summary, " cd users && JWT_SECRET=%q docker compose up -d\n", secret)
|
||||
fmt.Fprintf(&summary, " then set USERSERVICE_URL=http://localhost:8081 and BOOTSTRAP_ADMIN_EMAIL/PASSWORD\n")
|
||||
fmt.Fprintf(&summary, "\nThe wizard cannot run PostgreSQL or the userservice for you; the\n")
|
||||
fmt.Fprintf(&summary, "commands above are the supported way to start them.\n")
|
||||
return summary.String(), nil
|
||||
}
|
||||
|
||||
// probeDatabase verifies the target database accepts connections.
|
||||
func probeDatabase(ctx context.Context, config *pgx.ConnConfig, timeout time.Duration) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
conn, err := pgx.ConnectConfig(ctx, config)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return conn.Close(ctx)
|
||||
}
|
||||
|
||||
// resolveAdminConfig picks the maintenance connection used to create
|
||||
// databases. pgx's ConnConfig.ConnString() caches the original URL, so the
|
||||
// config itself (not a re-rendered string) is what the caller connects with.
|
||||
func resolveAdminConfig(options Options, target *pgx.ConnConfig) (*pgx.ConnConfig, error) {
|
||||
if options.AdminDatabaseURL != "" {
|
||||
config, err := pgx.ParseConfig(options.AdminDatabaseURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("--admin-db is not a valid postgres URL: %w", err)
|
||||
}
|
||||
return config, nil
|
||||
}
|
||||
admin := *target
|
||||
admin.Database = "postgres"
|
||||
return &admin, nil
|
||||
}
|
||||
|
||||
// createDatabase creates the named database through the maintenance connection.
|
||||
func createDatabase(ctx context.Context, admin *pgx.ConnConfig, name string, timeout time.Duration) error {
|
||||
ctx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
conn, err := pgx.ConnectConfig(ctx, admin)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = conn.Close(ctx) }()
|
||||
quoted := `"` + strings.ReplaceAll(name, `"`, `""`) + `"`
|
||||
if _, err := conn.Exec(ctx, "CREATE DATABASE "+quoted); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// generateSecret returns 32 random bytes as lowercase hex.
|
||||
func generateSecret() (string, error) {
|
||||
buffer := make([]byte, 32)
|
||||
if _, err := rand.Read(buffer); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(buffer), nil
|
||||
}
|
||||
|
||||
// writeEnvFile writes the settings, refusing to clobber without --force.
|
||||
func writeEnvFile(options Options, secret, storageDir string) error {
|
||||
path := filepath.Clean(options.EnvFile)
|
||||
if _, err := os.Stat(path); err == nil && !options.Force {
|
||||
return fmt.Errorf("%s already exists (use --force to overwrite)", path)
|
||||
}
|
||||
content := strings.Join([]string{
|
||||
"DATABASE_URL=" + options.DatabaseURL,
|
||||
"JWT_SECRET=" + secret,
|
||||
"COORDINATOR_STORAGE_DIR=" + storageDir,
|
||||
"", // trailing newline
|
||||
}, "\n")
|
||||
// #nosec G703 -- the env file path is operator-supplied (--env-file / ENV_FILE).
|
||||
if err := os.WriteFile(path, []byte(content), 0o600); err != nil {
|
||||
return fmt.Errorf("write %s: %w", path, err)
|
||||
}
|
||||
if err := os.Chmod(path, 0o600); err != nil {
|
||||
return fmt.Errorf("chmod %s: %w", path, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SanitizeDatabaseURL hides the password for logging.
|
||||
func SanitizeDatabaseURL(raw string) string {
|
||||
at := strings.LastIndex(raw, "@")
|
||||
if at < 0 {
|
||||
return raw
|
||||
}
|
||||
start := 0
|
||||
if strings.HasPrefix(raw, "postgres://") || strings.HasPrefix(raw, "postgresql://") {
|
||||
start = len("postgres://")
|
||||
}
|
||||
colon := strings.Index(raw[start:at], ":")
|
||||
if colon < 0 {
|
||||
return raw
|
||||
}
|
||||
colon += start
|
||||
return raw[:colon] + ":***@" + raw[at+1:]
|
||||
}
|
||||
|
||||
// prompt asks a question and returns the trimmed answer ("" on EOF).
|
||||
func prompt(options Options, question, fallback string) string {
|
||||
_, _ = fmt.Fprintf(options.Out, "%s [%s]: ", question, fallback)
|
||||
reader := bufio.NewReader(options.In)
|
||||
line, err := reader.ReadString('\n')
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
return fallback
|
||||
}
|
||||
answer := strings.TrimSpace(line)
|
||||
if answer == "" {
|
||||
return fallback
|
||||
}
|
||||
return answer
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
//go:build integration
|
||||
|
||||
package setup
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
func TestRunProvisionsDatabaseSchemaAndEnvFile(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
base := os.Getenv("TEST_DATABASE_URL")
|
||||
if base == "" {
|
||||
t.Skip("TEST_DATABASE_URL is not set")
|
||||
}
|
||||
// The wizard must create a *missing* database through the admin URL.
|
||||
slash := strings.LastIndex(base, "/")
|
||||
target := base[:slash+1] + "scimesh_setup_test"
|
||||
admin := base[:slash+1] + "postgres"
|
||||
|
||||
cleanup := func() {
|
||||
conn, err := pgx.Connect(ctx, admin)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = conn.Close(ctx) }()
|
||||
_, _ = conn.Exec(ctx, `DROP DATABASE IF EXISTS "scimesh_setup_test"`)
|
||||
}
|
||||
cleanup()
|
||||
t.Cleanup(cleanup)
|
||||
|
||||
envPath := filepath.Join(t.TempDir(), ".env")
|
||||
var output strings.Builder
|
||||
options := Options{
|
||||
DatabaseURL: target,
|
||||
AdminDatabaseURL: admin,
|
||||
EnvFile: envPath,
|
||||
Force: true,
|
||||
Yes: true,
|
||||
ConnectTimeout: 10 * time.Second,
|
||||
Out: &output,
|
||||
In: strings.NewReader(""),
|
||||
}
|
||||
summary, err := Run(ctx, options)
|
||||
if err != nil {
|
||||
t.Fatalf("setup run: %v", err)
|
||||
}
|
||||
for _, expected := range []string{"created database", "schema migrations applied"} {
|
||||
if !strings.Contains(output.String(), expected) {
|
||||
t.Errorf("progress output is missing %q:\n%s", expected, output.String())
|
||||
}
|
||||
}
|
||||
for _, expected := range []string{"Setup complete", "Start the coordinator"} {
|
||||
if !strings.Contains(summary, expected) {
|
||||
t.Errorf("summary is missing %q:\n%s", expected, summary)
|
||||
}
|
||||
}
|
||||
|
||||
// The database now exists, is migrated, and the env file is written.
|
||||
conn, err := pgx.Connect(ctx, target)
|
||||
if err != nil {
|
||||
t.Fatalf("connect to provisioned database: %v", err)
|
||||
}
|
||||
defer func() { _ = conn.Close(ctx) }()
|
||||
var watermark int64
|
||||
if err := conn.QueryRow(ctx, "SELECT COALESCE(MAX(version), 0) FROM schema_migrations").Scan(&watermark); err != nil {
|
||||
t.Fatalf("read schema_migrations: %v", err)
|
||||
}
|
||||
if watermark < 13 {
|
||||
t.Errorf("schema watermark = %d, want >= 13", watermark)
|
||||
}
|
||||
envContent, err := os.ReadFile(envPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
env := string(envContent)
|
||||
if !strings.Contains(env, "DATABASE_URL="+target) || !strings.Contains(env, "JWT_SECRET=") {
|
||||
t.Errorf("env file is incomplete:\n%s", env)
|
||||
}
|
||||
|
||||
// A second run is idempotent: no create-database error, same outcome.
|
||||
if _, err := Run(ctx, options); err != nil {
|
||||
t.Fatalf("second setup run: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package setup
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGenerateSecretIsRandomAndStrong(t *testing.T) {
|
||||
first, err := generateSecret()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
second, err := generateSecret()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(first) != 64 || len(second) != 64 {
|
||||
t.Fatalf("secrets must be 32 random bytes as hex, got %d and %d", len(first), len(second))
|
||||
}
|
||||
if first == second {
|
||||
t.Fatal("two generated secrets must differ")
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteEnvFileContentsAndPermissions(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), ".env")
|
||||
if err := writeEnvFile(Options{
|
||||
EnvFile: path,
|
||||
DatabaseURL: "postgres://scimesh@localhost/scimesh",
|
||||
}, "s3cr3t", "./data"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
content, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := "DATABASE_URL=postgres://scimesh@localhost/scimesh\nJWT_SECRET=s3cr3t\nCOORDINATOR_STORAGE_DIR=./data\n"
|
||||
if got := string(content); got != want {
|
||||
t.Fatalf("env file = %q, want %q", got, want)
|
||||
}
|
||||
info, err := os.Stat(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if info.Mode().Perm() != 0o600 {
|
||||
t.Errorf("env file mode = %o, want 0600", info.Mode().Perm())
|
||||
}
|
||||
}
|
||||
|
||||
func TestWriteEnvFileRefusesWithoutForce(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), ".env")
|
||||
if err := writeEnvFile(Options{EnvFile: path, DatabaseURL: "postgres://x@localhost/a"}, "a", "./data"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := writeEnvFile(Options{EnvFile: path, DatabaseURL: "postgres://x@localhost/a"}, "b", "./data"); err == nil {
|
||||
t.Fatal("second write without --force must fail")
|
||||
}
|
||||
if err := writeEnvFile(Options{EnvFile: path, Force: true, DatabaseURL: "postgres://x@localhost/a"}, "b", "./data"); err != nil {
|
||||
t.Fatalf("write with --force: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPromptReadsAnswer(t *testing.T) {
|
||||
var out bytes.Buffer
|
||||
answer := prompt(Options{Out: &out, In: strings.NewReader("postgres://custom\n")}, "Database URL", "default")
|
||||
if answer != "postgres://custom" {
|
||||
t.Fatalf("answer = %q, want the typed value", answer)
|
||||
}
|
||||
if !strings.Contains(out.String(), "Database URL [default]:") {
|
||||
t.Fatalf("prompt output = %q", out.String())
|
||||
}
|
||||
fallback := prompt(Options{Out: &out, In: strings.NewReader("\n")}, "Question", "fb")
|
||||
if fallback != "fb" {
|
||||
t.Fatalf("empty answer must fall back, got %q", fallback)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSanitizeDatabaseURL(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"postgres://scimesh:hunter2@localhost:5432/scimesh?sslmode=disable": "postgres://scimesh:***@localhost:5432/scimesh?sslmode=disable",
|
||||
"postgresql://scimesh@localhost/scimesh": "postgresql://scimesh@localhost/scimesh",
|
||||
"not-a-url": "not-a-url",
|
||||
}
|
||||
for raw, want := range cases {
|
||||
got := SanitizeDatabaseURL(raw)
|
||||
if got != want {
|
||||
t.Errorf("sanitize(%q) = %q, want %q", raw, got, want)
|
||||
}
|
||||
if strings.Contains(got, "hunter2") {
|
||||
t.Errorf("sanitize(%q) leaked the password", raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -70,3 +70,33 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
func (r *ArtifactRepo) FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
|
||||
sql, args, err := psql.Select(artifactColumns...).
|
||||
From("artifacts").
|
||||
Where(sq.Eq{
|
||||
"task_id": taskID,
|
||||
"attempt": attempt,
|
||||
"kind": string(domain.ArtifactPartialResult),
|
||||
}).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var (
|
||||
a domain.Artifact
|
||||
kind string
|
||||
)
|
||||
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
|
||||
&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey,
|
||||
&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("find partial result: %w", err)
|
||||
}
|
||||
a.Kind = domain.ArtifactKind(kind)
|
||||
return &a, nil
|
||||
}
|
||||
|
||||
@@ -24,6 +24,7 @@ import (
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/workloads"
|
||||
)
|
||||
|
||||
func testPool(t *testing.T) *pgxpool.Pool {
|
||||
@@ -90,6 +91,76 @@ func TestCreateJobPersistsEveryTask(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimReductionIsAtomic(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
job, _ := seedJob(t, pool, 1)
|
||||
repo := NewJobRepo(pool)
|
||||
ctx := context.Background()
|
||||
if err := repo.UpdateStatus(ctx, job.ID, domain.JobReducing, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var (
|
||||
wg sync.WaitGroup
|
||||
mu sync.Mutex
|
||||
claimed int
|
||||
)
|
||||
for range 8 {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
ok, err := repo.ClaimReduction(context.Background(), job.ID, time.Now().UTC())
|
||||
if err != nil {
|
||||
t.Errorf("claim reduction: %v", err)
|
||||
return
|
||||
}
|
||||
if ok {
|
||||
mu.Lock()
|
||||
claimed++
|
||||
mu.Unlock()
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
if claimed != 1 {
|
||||
t.Fatalf("reducer claims = %d, want 1", claimed)
|
||||
}
|
||||
stored, err := repo.Get(ctx, job.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.Status != domain.JobReducing || stored.ReducerStartedAt == nil {
|
||||
t.Fatalf("stored reduction state = %+v", stored)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIReadRepoListsReducerFields(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
job, _ := seedJob(t, pool, 1)
|
||||
jobs := NewJobRepo(pool)
|
||||
ctx := context.Background()
|
||||
if err := jobs.UpdateStatus(ctx, job.ID, domain.JobReducing, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if claimed, err := jobs.ClaimReduction(ctx, job.ID, time.Now().UTC()); err != nil || !claimed {
|
||||
t.Fatalf("claim reduction = (%v, %v)", claimed, err)
|
||||
}
|
||||
listed, err := NewUIReadRepo(pool).ListJobs(ctx, nil, 20)
|
||||
if err != nil {
|
||||
t.Fatalf("list UI jobs: %v", err)
|
||||
}
|
||||
for _, item := range listed {
|
||||
if item.ID != job.ID {
|
||||
continue
|
||||
}
|
||||
if item.Status != domain.JobReducing || item.ReducerStartedAt == nil {
|
||||
t.Fatalf("UI reducer projection = %+v", item)
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatalf("seeded job %s is missing from UI list", job.ID)
|
||||
}
|
||||
|
||||
// A job must land whole or not at all: a half-created job leaves chunks no
|
||||
// worker could ever complete.
|
||||
func TestCreateJobRollsBackOnFailure(t *testing.T) {
|
||||
@@ -336,8 +407,9 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) {
|
||||
job, _ := seedJob(t, pool, 1)
|
||||
|
||||
tasks, jobs, artifacts, tx := NewTaskRepo(pool), NewJobRepo(pool), NewArtifactRepo(pool), NewTxManager(pool)
|
||||
workers, results := NewWorkerRepo(pool), NewTaskResultRepo(pool)
|
||||
clk := fixedClock{now: time.Now().UTC()}
|
||||
uc := usecase.NewCompleteTask(tasks, jobs, artifacts, tx, clk)
|
||||
uc := usecase.NewCompleteTask(tasks, jobs, artifacts, workers, results, tx, clk, 2, integrationCatalog())
|
||||
|
||||
claimed, err := tasks.ClaimNext(ctx, usecase.ClaimFilter{
|
||||
Owner: "worker-1", Now: clk.now, LeaseUntil: clk.now.Add(time.Minute),
|
||||
@@ -360,12 +432,23 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) {
|
||||
t.Errorf("replay must be idempotent, got %v", err)
|
||||
}
|
||||
|
||||
// A different result artifact for the same task is a genuine conflict.
|
||||
art2 := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult)
|
||||
other := in
|
||||
other.ResultArtifactID = art2.ID
|
||||
if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) {
|
||||
t.Errorf("err = %v, want ErrResultConflict", err)
|
||||
}
|
||||
|
||||
func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
job, tasks := seedJob(t, pool, 1)
|
||||
taskID := tasks[0].ID
|
||||
first := seedArtifact(t, pool, job.ID, &taskID, domain.ArtifactPartialResult)
|
||||
second, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, "retry.csv", "text/csv", time.Now().UTC())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
attempt := 1
|
||||
second.Attempt = &attempt
|
||||
second.SetContent("other-sha", 5)
|
||||
if err := NewArtifactRepo(pool).Insert(ctx, second); err == nil {
|
||||
t.Fatalf("second partial artifact for %s/%d was accepted after %s", taskID, attempt, first.ID)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -575,3 +658,46 @@ func TestExpireLeasesRequeuesElapsedTasks(t *testing.T) {
|
||||
t.Errorf("pending = %d, want 1 — a dead worker must not strand its task", counts[domain.TaskPending])
|
||||
}
|
||||
}
|
||||
|
||||
func TestMigrateProvisionsAndIsIdempotent(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
url := os.Getenv("TEST_DATABASE_URL")
|
||||
if url == "" {
|
||||
t.Skip("TEST_DATABASE_URL is not set")
|
||||
}
|
||||
if err := Migrate(ctx, url, nil); err != nil {
|
||||
t.Fatalf("first migrate: %v", err)
|
||||
}
|
||||
if err := Migrate(ctx, url, nil); err != nil {
|
||||
t.Fatalf("second migrate (idempotent): %v", err)
|
||||
}
|
||||
pool := testPool(t)
|
||||
migrations, err := listMigrations()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var watermark int64
|
||||
if err := pool.QueryRow(ctx, "SELECT COALESCE(MAX(version), 0) FROM schema_migrations").Scan(&watermark); err != nil {
|
||||
t.Fatalf("read schema_migrations: %v", err)
|
||||
}
|
||||
if watermark != int64(len(migrations)) {
|
||||
t.Errorf("schema watermark = %d, want %d", watermark, len(migrations))
|
||||
}
|
||||
var hasJobs bool
|
||||
if err := pool.QueryRow(ctx,
|
||||
"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'jobs')",
|
||||
).Scan(&hasJobs); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !hasJobs {
|
||||
t.Error("jobs table was not created by the embedded migrations")
|
||||
}
|
||||
}
|
||||
|
||||
func integrationCatalog() *workloads.Catalog {
|
||||
catalog, err := workloads.Load()
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
return catalog
|
||||
}
|
||||
|
||||
@@ -25,14 +25,18 @@ func NewJobRepo(pool *pgxpool.Pool) *JobRepo {
|
||||
|
||||
var _ usecase.JobRepository = (*JobRepo)(nil)
|
||||
|
||||
var jobColumns = []string{"id", "workload", "input_uri", "parameters", "status", "created_at", "completed_at"}
|
||||
var jobColumns = []string{
|
||||
"id", "workload", "input_uri", "parameters", "status", "created_at", "completed_at",
|
||||
"input_artifact_id", "result_artifact_id", "error_code", "error_message", "reducer_started_at",
|
||||
"owner_id",
|
||||
}
|
||||
|
||||
// Insert runs inside the caller's transaction, alongside the job's tasks — that
|
||||
// is what makes "all tasks or none" hold.
|
||||
func (r *JobRepo) Insert(ctx context.Context, j *domain.Job) error {
|
||||
sql, args, err := psql.Insert("jobs").
|
||||
Columns("id", "workload", "input_uri", "parameters", "status", "created_at").
|
||||
Values(j.ID, j.Workload, j.InputURI, jsonbOrEmpty(j.Parameters), string(j.Status), j.CreatedAt).
|
||||
Columns("id", "workload", "input_uri", "parameters", "status", "created_at", "owner_id").
|
||||
Values(j.ID, j.Workload, j.InputURI, jsonbOrEmpty(j.Parameters), string(j.Status), j.CreatedAt, j.OwnerID).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return err
|
||||
@@ -55,7 +59,9 @@ func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
|
||||
status string
|
||||
)
|
||||
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
|
||||
&j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt)
|
||||
&j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt,
|
||||
&j.InputArtifactID, &j.ResultArtifactID, &j.ErrorCode, &j.ErrorMessage, &j.ReducerStartedAt,
|
||||
&j.OwnerID)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, domain.ErrJobNotFound
|
||||
}
|
||||
@@ -66,6 +72,64 @@ func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
|
||||
return &j, nil
|
||||
}
|
||||
|
||||
func (r *JobRepo) ClaimReduction(ctx context.Context, id uuid.UUID, startedAt time.Time) (bool, error) {
|
||||
sql, args, err := psql.Update("jobs").
|
||||
Set("reducer_started_at", startedAt).
|
||||
Where(sq.Eq{"id": id, "status": string(domain.JobReducing), "reducer_started_at": nil}).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return tag.RowsAffected() == 1, nil
|
||||
}
|
||||
|
||||
func (r *JobRepo) CompleteWithResult(ctx context.Context, id, resultArtifactID uuid.UUID, completedAt time.Time) error {
|
||||
sql, args, err := psql.Update("jobs").
|
||||
SetMap(map[string]any{
|
||||
"status": string(domain.JobCompleted),
|
||||
"result_artifact_id": resultArtifactID,
|
||||
"completed_at": completedAt,
|
||||
"reducer_started_at": nil,
|
||||
"error_code": nil,
|
||||
"error_message": nil,
|
||||
}).
|
||||
Where(sq.Eq{"id": id, "status": string(domain.JobReducing)}).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return domain.ErrJobNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *JobRepo) FailReduction(ctx context.Context, id uuid.UUID, code, message string, completedAt time.Time) error {
|
||||
sql, args, err := psql.Update("jobs").
|
||||
SetMap(map[string]any{
|
||||
"status": string(domain.JobFailed),
|
||||
"completed_at": completedAt,
|
||||
"error_code": code,
|
||||
"error_message": message,
|
||||
"reducer_started_at": nil,
|
||||
}).
|
||||
Where(sq.Eq{"id": id, "status": string(domain.JobReducing)}).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = conn(ctx, r.pool).Exec(ctx, sql, args...)
|
||||
return err
|
||||
}
|
||||
|
||||
func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID,
|
||||
status domain.JobStatus, completedAt *time.Time) error {
|
||||
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"embed"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"regexp"
|
||||
"sort"
|
||||
"strconv"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
//go:embed migrations/*.sql
|
||||
var migrationFiles embed.FS
|
||||
|
||||
var migrationNamePattern = regexp.MustCompile(`^([0-9]+)_[a-z0-9_]+\.(up|down)\.sql$`)
|
||||
|
||||
// migration is one parsed embedded migration file.
|
||||
type migration struct {
|
||||
version int
|
||||
name string
|
||||
sql string
|
||||
}
|
||||
|
||||
// listMigrations parses and orders the embedded .up.sql files by version.
|
||||
func listMigrations() ([]migration, error) {
|
||||
entries, err := migrationFiles.ReadDir("migrations")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read embedded migrations: %w", err)
|
||||
}
|
||||
up := map[int]migration{}
|
||||
for _, entry := range entries {
|
||||
match := migrationNamePattern.FindStringSubmatch(entry.Name())
|
||||
if match == nil {
|
||||
continue
|
||||
}
|
||||
if match[2] != "up" {
|
||||
continue
|
||||
}
|
||||
version, err := strconv.Atoi(match[1])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("migration %q has an invalid version: %w", entry.Name(), err)
|
||||
}
|
||||
if _, duplicate := up[version]; duplicate {
|
||||
return nil, fmt.Errorf("migration version %d is duplicated", version)
|
||||
}
|
||||
body, err := migrationFiles.ReadFile("migrations/" + entry.Name())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read migration %q: %w", entry.Name(), err)
|
||||
}
|
||||
up[version] = migration{version: version, name: entry.Name(), sql: string(body)}
|
||||
}
|
||||
if len(up) == 0 {
|
||||
return nil, fmt.Errorf("no .up.sql migrations are embedded")
|
||||
}
|
||||
versions := make([]int, 0, len(up))
|
||||
for version := range up {
|
||||
versions = append(versions, version)
|
||||
}
|
||||
sort.Ints(versions)
|
||||
migrations := make([]migration, 0, len(versions))
|
||||
for _, version := range versions {
|
||||
migrations = append(migrations, up[version])
|
||||
}
|
||||
for index, item := range migrations {
|
||||
if item.version != index+1 {
|
||||
return nil, fmt.Errorf("embedded migrations are not contiguous: version %d at position %d", item.version, index+1)
|
||||
}
|
||||
}
|
||||
return migrations, nil
|
||||
}
|
||||
|
||||
// Migrate applies every embedded migration above the recorded schema version,
|
||||
// so the binary provisions its own schema. It is idempotent and interoperates
|
||||
// with the golang-migrate CLI: both tools use the same schema_migrations
|
||||
// watermark table (single row: version + dirty flag), and a PostgreSQL
|
||||
// advisory lock serializes concurrent migrators. Each migration file runs as
|
||||
// its own transaction (the files carry explicit BEGIN/COMMIT, matching the
|
||||
// golang-migrate format the CLI and CI still use).
|
||||
func Migrate(ctx context.Context, databaseURL string, log *slog.Logger) error {
|
||||
migrations, err := listMigrations()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
connConfig, err := pgx.ParseConfig(databaseURL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("parse database url: %w", err)
|
||||
}
|
||||
// Migration files contain multiple statements (BEGIN...COMMIT), which the
|
||||
// extended query protocol rejects; run them with the simple protocol.
|
||||
connConfig.DefaultQueryExecMode = pgx.QueryExecModeSimpleProtocol
|
||||
conn, err := pgx.ConnectConfig(ctx, connConfig)
|
||||
if err != nil {
|
||||
return fmt.Errorf("connect for migration: %w", err)
|
||||
}
|
||||
defer func() { _ = conn.Close(ctx) }()
|
||||
|
||||
if _, err := conn.Exec(ctx, "SELECT pg_advisory_lock(82473911)"); err != nil {
|
||||
return fmt.Errorf("acquire migration lock: %w", err)
|
||||
}
|
||||
defer func() { _, _ = conn.Exec(ctx, "SELECT pg_advisory_unlock(82473911)") }()
|
||||
|
||||
if _, err := conn.Exec(ctx,
|
||||
"CREATE TABLE IF NOT EXISTS schema_migrations (version bigint PRIMARY KEY, dirty boolean NOT NULL DEFAULT false)",
|
||||
); err != nil {
|
||||
return fmt.Errorf("ensure schema_migrations: %w", err)
|
||||
}
|
||||
|
||||
var applied int64
|
||||
if err := conn.QueryRow(ctx,
|
||||
"SELECT COALESCE(MAX(version), 0) FROM schema_migrations",
|
||||
).Scan(&applied); err != nil {
|
||||
return fmt.Errorf("read applied schema version: %w", err)
|
||||
}
|
||||
|
||||
for _, item := range migrations {
|
||||
if int64(item.version) <= applied {
|
||||
continue
|
||||
}
|
||||
if log != nil {
|
||||
log.Info("applying migration", "version", item.version, "file", item.name)
|
||||
}
|
||||
if _, err := conn.Exec(ctx, item.sql); err != nil {
|
||||
return fmt.Errorf("apply migration %s: %w", item.name, err)
|
||||
}
|
||||
// Advance the watermark to the single-row golang-migrate layout.
|
||||
tag, err := conn.Exec(ctx,
|
||||
"UPDATE schema_migrations SET version = $1, dirty = false", item.version)
|
||||
if err != nil {
|
||||
return fmt.Errorf("record migration %s: %w", item.name, err)
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
if _, err := conn.Exec(ctx,
|
||||
"INSERT INTO schema_migrations (version, dirty) VALUES ($1, false)",
|
||||
item.version,
|
||||
); err != nil {
|
||||
return fmt.Errorf("record migration %s: %w", item.name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestListMigrationsParsesAndOrdersEmbeddedFiles(t *testing.T) {
|
||||
migrations, err := listMigrations()
|
||||
if err != nil {
|
||||
t.Fatalf("list migrations: %v", err)
|
||||
}
|
||||
if len(migrations) == 0 {
|
||||
t.Fatal("no embedded migrations")
|
||||
}
|
||||
for index, item := range migrations {
|
||||
if item.version != index+1 {
|
||||
t.Errorf("migration %d has version %d, want contiguous ordering", index, item.version)
|
||||
}
|
||||
if item.name != expectedMigrationName(item.version) {
|
||||
t.Errorf("migration %d file is %q, want %q", item.version, item.name, expectedMigrationName(item.version))
|
||||
}
|
||||
if strings.TrimSpace(item.sql) == "" {
|
||||
t.Errorf("migration %d is empty", item.version)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func expectedMigrationName(version int) string {
|
||||
switch version {
|
||||
case 1:
|
||||
return "0001_init.up.sql"
|
||||
case 2:
|
||||
return "0002_workers.up.sql"
|
||||
case 3:
|
||||
return "0003_artifacts.up.sql"
|
||||
case 4:
|
||||
return "0004_result_artifact.up.sql"
|
||||
case 5:
|
||||
return "0005_uploaded_input.up.sql"
|
||||
case 6:
|
||||
return "0006_task_running_enum.up.sql"
|
||||
case 7:
|
||||
return "0007_task_running_lease.up.sql"
|
||||
case 8:
|
||||
return "0008_artifact_attempt.up.sql"
|
||||
case 9:
|
||||
return "0009_unique_partial_result_attempt.up.sql"
|
||||
case 10:
|
||||
return "0010_job_reduction.up.sql"
|
||||
case 11:
|
||||
return "0011_job_owner.up.sql"
|
||||
case 12:
|
||||
return "0012_worker_trust.up.sql"
|
||||
case 13:
|
||||
return "0013_task_results.up.sql"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
+5
@@ -0,0 +1,5 @@
|
||||
BEGIN;
|
||||
|
||||
DROP INDEX IF EXISTS uq_partial_result_task_attempt;
|
||||
|
||||
COMMIT;
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
BEGIN;
|
||||
|
||||
-- Old deployments can contain more than one partial result because earlier
|
||||
-- versions accepted repeated PUTs. Preserve the one referenced by a completed
|
||||
-- task and discard stale rows; unfinished tasks must upload again after a
|
||||
-- deploy, just as they do after a lost lease.
|
||||
DELETE FROM artifacts AS a
|
||||
USING tasks AS t
|
||||
WHERE a.task_id = t.id
|
||||
AND a.kind = 'partial_result'::artifact_kind
|
||||
AND t.status <> 'completed'::task_status;
|
||||
|
||||
DELETE FROM artifacts AS a
|
||||
USING tasks AS t
|
||||
WHERE a.task_id = t.id
|
||||
AND a.kind = 'partial_result'::artifact_kind
|
||||
AND t.status = 'completed'::task_status
|
||||
AND a.id <> t.result_artifact_id;
|
||||
|
||||
-- One lease attempt has one durable partial result. This makes an upload retry
|
||||
-- idempotent and prevents repeated uploads from accumulating orphan artifacts.
|
||||
CREATE UNIQUE INDEX uq_partial_result_task_attempt
|
||||
ON artifacts (task_id, attempt)
|
||||
WHERE kind = 'partial_result'::artifact_kind;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,7 @@
|
||||
BEGIN;
|
||||
|
||||
ALTER TABLE jobs DROP COLUMN IF EXISTS error_message;
|
||||
ALTER TABLE jobs DROP COLUMN IF EXISTS error_code;
|
||||
ALTER TABLE jobs DROP COLUMN IF EXISTS reducer_started_at;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,7 @@
|
||||
-- PostgreSQL enum values must be committed before they are used by a later
|
||||
-- transaction, so this migration intentionally has no BEGIN/COMMIT wrapper.
|
||||
ALTER TYPE job_status ADD VALUE IF NOT EXISTS 'reducing';
|
||||
|
||||
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS error_code text;
|
||||
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS error_message text;
|
||||
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS reducer_started_at timestamptz;
|
||||
@@ -0,0 +1,6 @@
|
||||
BEGIN;
|
||||
|
||||
DROP INDEX IF EXISTS ix_jobs_owner;
|
||||
ALTER TABLE jobs DROP COLUMN IF EXISTS owner_id;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,14 @@
|
||||
BEGIN;
|
||||
|
||||
-- Who submitted this job. Equals users.id from the userservice, taken from the
|
||||
-- JWT `sub` claim. NOT a foreign key: users live in a separate service/database,
|
||||
-- so integrity is guaranteed by the signed token, not by the DB.
|
||||
--
|
||||
-- Nullable because rows created before auth existed have no owner; new inserts
|
||||
-- must supply it (enforced in the app, not the schema, during the MVP).
|
||||
ALTER TABLE jobs ADD COLUMN owner_id uuid;
|
||||
|
||||
-- "List my jobs" / "admin filters by owner" scans by owner.
|
||||
CREATE INDEX ix_jobs_owner ON jobs (owner_id);
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,8 @@
|
||||
BEGIN;
|
||||
|
||||
DROP INDEX IF EXISTS ix_workers_owner;
|
||||
ALTER TABLE workers DROP COLUMN IF EXISTS trust_level;
|
||||
ALTER TABLE workers DROP COLUMN IF EXISTS owner_id;
|
||||
DROP TYPE IF EXISTS worker_trust;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,18 @@
|
||||
BEGIN;
|
||||
|
||||
-- Whether a worker's results are accepted directly or must clear quorum.
|
||||
-- 'trusted' — lab machine (shared token) or a verified/admin contributor.
|
||||
-- 'untrusted' — a plain enthusiast; results are quarantined until quorum (C2).
|
||||
CREATE TYPE worker_trust AS ENUM ('trusted', 'untrusted');
|
||||
|
||||
-- Who registered this worker (userservice user id, from the JWT sub). NULL for
|
||||
-- workers registered with the shared service token. Not a foreign key: users
|
||||
-- live in a separate service/database.
|
||||
ALTER TABLE workers ADD COLUMN owner_id uuid;
|
||||
|
||||
-- Existing rows were all shared-token lab workers, hence 'trusted'.
|
||||
ALTER TABLE workers ADD COLUMN trust_level worker_trust NOT NULL DEFAULT 'trusted';
|
||||
|
||||
CREATE INDEX ix_workers_owner ON workers (owner_id);
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,5 @@
|
||||
BEGIN;
|
||||
|
||||
DROP TABLE IF EXISTS task_results;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,23 @@
|
||||
BEGIN;
|
||||
|
||||
-- Quorum votes for a task computed by untrusted (volunteer) workers. A trusted
|
||||
-- worker's result completes the task directly and never lands here; an untrusted
|
||||
-- result is recorded as one vote, and the task is only completed once enough
|
||||
-- distinct owners submit the same result_sha256.
|
||||
--
|
||||
-- One vote per (task, owner): a single volunteer cannot stuff the ballot by
|
||||
-- running many workers under one account. A resubmission updates their vote.
|
||||
CREATE TABLE task_results (
|
||||
task_id uuid NOT NULL REFERENCES tasks(id) ON DELETE CASCADE,
|
||||
owner_id uuid NOT NULL,
|
||||
result_sha256 text NOT NULL,
|
||||
result_artifact_id uuid NOT NULL REFERENCES artifacts(id) ON DELETE CASCADE,
|
||||
created_at timestamptz NOT NULL DEFAULT now(),
|
||||
|
||||
PRIMARY KEY (task_id, owner_id)
|
||||
);
|
||||
|
||||
-- Quorum check groups a task's votes by result_sha256.
|
||||
CREATE INDEX ix_task_results_quorum ON task_results (task_id, result_sha256);
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,65 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
)
|
||||
|
||||
// Known statuses per entity, so counts are zero-filled and every status is
|
||||
// always present in the metrics (a flat 0 line beats a gap on the dashboard).
|
||||
var (
|
||||
taskStatuses = []string{string(domain.TaskPending), string(domain.TaskLeased), string(domain.TaskRunning), string(domain.TaskCompleted), string(domain.TaskFailed), string(domain.TaskCancelled)}
|
||||
jobStatuses = []string{string(domain.JobPending), string(domain.JobRunning), string(domain.JobReducing), string(domain.JobCompleted), string(domain.JobFailed), string(domain.JobCancelled)}
|
||||
workerStatuses = []string{string(domain.WorkerOnline), string(domain.WorkerBusy), string(domain.WorkerOffline)}
|
||||
)
|
||||
|
||||
// StatsRepo answers the aggregate status counts the business metrics report. It
|
||||
// runs one cheap GROUP BY per entity; the collector calls this on every scrape.
|
||||
type StatsRepo struct {
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
func NewStatsRepo(pool *pgxpool.Pool) *StatsRepo {
|
||||
return &StatsRepo{pool: pool}
|
||||
}
|
||||
|
||||
// Counts returns status->count maps for tasks, jobs, and workers, each
|
||||
// zero-filled across its known statuses.
|
||||
func (r *StatsRepo) Counts(ctx context.Context) (tasks, jobs, workers map[string]int, err error) {
|
||||
if tasks, err = r.countByStatus(ctx, "tasks", taskStatuses); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
if jobs, err = r.countByStatus(ctx, "jobs", jobStatuses); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
if workers, err = r.countByStatus(ctx, "workers", workerStatuses); err != nil {
|
||||
return nil, nil, nil, err
|
||||
}
|
||||
return tasks, jobs, workers, nil
|
||||
}
|
||||
|
||||
func (r *StatsRepo) countByStatus(ctx context.Context, table string, known []string) (map[string]int, error) {
|
||||
out := make(map[string]int, len(known))
|
||||
for _, s := range known {
|
||||
out[s] = 0 // zero-fill
|
||||
}
|
||||
// table is a fixed internal constant, never user input — safe to format.
|
||||
rows, err := r.pool.Query(ctx, fmt.Sprintf("SELECT status, count(*) FROM %s GROUP BY status", table))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("count %s by status: %w", table, err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var status string
|
||||
var n int
|
||||
if err := rows.Scan(&status, &n); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[status] = n // an unknown status still shows up, which is a useful signal
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
@@ -84,6 +84,9 @@ WITH candidate AS (
|
||||
WHERE status = 'pending'
|
||||
AND attempt < max_attempts
|
||||
AND (cardinality($1::text[]) = 0 OR workload = ANY($1))
|
||||
AND ($5::uuid IS NULL OR NOT EXISTS (
|
||||
SELECT 1 FROM task_results tr
|
||||
WHERE tr.task_id = tasks.id AND tr.owner_id = $5))
|
||||
ORDER BY created_at, chunk_index
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT 1
|
||||
@@ -108,7 +111,7 @@ func (r *TaskRepo) ClaimNext(ctx context.Context, f usecase.ClaimFilter) (*domai
|
||||
|
||||
var task *domain.Task
|
||||
err := withRetry(ctx, func(ctx context.Context) error {
|
||||
row := conn(ctx, r.pool).QueryRow(ctx, claimNextSQL, workloads, f.Owner, f.LeaseUntil, f.Now)
|
||||
row := conn(ctx, r.pool).QueryRow(ctx, claimNextSQL, workloads, f.Owner, f.LeaseUntil, f.Now, f.VoterOwner)
|
||||
t, err := scanTask(row)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
task = nil
|
||||
@@ -337,17 +340,26 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta
|
||||
ELSE error_message END,
|
||||
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
|
||||
version = version + 1
|
||||
WHERE status IN ('leased','running') AND lease_expires_at < $1`
|
||||
WHERE status IN ('leased','running') AND lease_expires_at < $1
|
||||
RETURNING job_id`
|
||||
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
|
||||
var affected int64
|
||||
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
|
||||
var affected []uuid.UUID
|
||||
err := withRetry(ctx, func(ctx context.Context) error {
|
||||
tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
affected = tag.RowsAffected()
|
||||
return nil
|
||||
defer rows.Close()
|
||||
affected = affected[:0]
|
||||
for rows.Next() {
|
||||
var jobID uuid.UUID
|
||||
if err := rows.Scan(&jobID); err != nil {
|
||||
return err
|
||||
}
|
||||
affected = append(affected, jobID)
|
||||
}
|
||||
return rows.Err()
|
||||
})
|
||||
return affected, err
|
||||
}
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
)
|
||||
|
||||
// TaskResultRepo records and tallies quorum votes for untrusted task results.
|
||||
type TaskResultRepo struct {
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
func NewTaskResultRepo(pool *pgxpool.Pool) *TaskResultRepo {
|
||||
return &TaskResultRepo{pool: pool}
|
||||
}
|
||||
|
||||
// RecordVote stores (or replaces) one owner's vote for a task's result.
|
||||
func (r *TaskResultRepo) RecordVote(ctx context.Context, taskID, ownerID uuid.UUID, sha256 string, artifactID uuid.UUID) error {
|
||||
const sql = `
|
||||
INSERT INTO task_results (task_id, owner_id, result_sha256, result_artifact_id)
|
||||
VALUES ($1, $2, $3, $4)
|
||||
ON CONFLICT (task_id, owner_id) DO UPDATE
|
||||
SET result_sha256 = EXCLUDED.result_sha256,
|
||||
result_artifact_id = EXCLUDED.result_artifact_id,
|
||||
created_at = now()`
|
||||
if _, err := conn(ctx, r.pool).Exec(ctx, sql, taskID, ownerID, sha256, artifactID); err != nil {
|
||||
return fmt.Errorf("record vote: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CountAgreeing returns how many distinct owners have voted for the given result
|
||||
// hash on this task — the size of the agreeing set the quorum is measured
|
||||
// against.
|
||||
func (r *TaskResultRepo) CountAgreeing(ctx context.Context, taskID uuid.UUID, sha256 string) (int, error) {
|
||||
const sql = `SELECT count(DISTINCT owner_id) FROM task_results WHERE task_id = $1 AND result_sha256 = $2`
|
||||
var n int
|
||||
if err := conn(ctx, r.pool).QueryRow(ctx, sql, taskID, sha256).Scan(&n); err != nil {
|
||||
return 0, fmt.Errorf("count agreeing: %w", err)
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
@@ -24,11 +24,15 @@ func (r *UIReadRepo) GetJob(ctx context.Context, id uuid.UUID) (*domain.Job, err
|
||||
return job, err
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListJobs(ctx context.Context, limit int) ([]domain.Job, error) {
|
||||
func (r *UIReadRepo) ListJobs(ctx context.Context, owner *uuid.UUID, limit int) ([]domain.Job, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
sql, args, err := psql.Select(jobColumns...).From("jobs").OrderBy("created_at DESC", "id DESC").Limit(uint64(limit)).ToSql()
|
||||
q := psql.Select(jobColumns...).From("jobs")
|
||||
if owner != nil {
|
||||
q = q.Where(sq.Eq{"owner_id": *owner})
|
||||
}
|
||||
sql, args, err := q.OrderBy("created_at DESC", "id DESC").Limit(uint64(limit)).ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -41,13 +45,13 @@ func (r *UIReadRepo) ListJobs(ctx context.Context, limit int) ([]domain.Job, err
|
||||
for rows.Next() {
|
||||
var j domain.Job
|
||||
var status string
|
||||
var inputURI *string
|
||||
if err := rows.Scan(&j.ID, &j.Workload, &inputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt); err != nil {
|
||||
if err := rows.Scan(
|
||||
&j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt,
|
||||
&j.InputArtifactID, &j.ResultArtifactID, &j.ErrorCode, &j.ErrorMessage, &j.ReducerStartedAt,
|
||||
&j.OwnerID,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if inputURI != nil {
|
||||
j.InputURI = *inputURI
|
||||
}
|
||||
j.Status = domain.JobStatus(status)
|
||||
jobs = append(jobs, j)
|
||||
}
|
||||
@@ -75,6 +79,31 @@ func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]dom
|
||||
return tasks, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
|
||||
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
|
||||
if len(jobIDs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
sql, args, err := psql.Select(taskColumns...).From("tasks").
|
||||
Where(sq.Eq{"job_id": jobIDs}).OrderBy("job_id ASC", "chunk_index ASC").ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list tasks by jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
task, err := scanTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[task.JobID] = append(out[task.JobID], *task)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
@@ -99,6 +128,32 @@ func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worke
|
||||
return workers, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListWorkersByOwner(ctx context.Context, owner uuid.UUID, limit int) ([]domain.Worker, error) {
|
||||
if limit < 1 || limit > 100 {
|
||||
return nil, domain.ErrInvalidInput
|
||||
}
|
||||
sql, args, err := psql.Select(workerColumns...).From("workers").
|
||||
Where(sq.Eq{"owner_id": owner}).
|
||||
OrderBy("last_heartbeat_at DESC", "id DESC").Limit(uint64(limit)).ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list workers by owner: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
workers := make([]domain.Worker, 0)
|
||||
for rows.Next() {
|
||||
worker, err := scanWorker(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
workers = append(workers, *worker)
|
||||
}
|
||||
return workers, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error) {
|
||||
sql, args, err := psql.Select(artifactColumns...).From("artifacts").Where(sq.Eq{"job_id": jobID}).OrderBy("created_at ASC", "id ASC").ToSql()
|
||||
if err != nil {
|
||||
|
||||
@@ -23,13 +23,13 @@ func NewWorkerRepo(pool *pgxpool.Pool) *WorkerRepo {
|
||||
return &WorkerRepo{pool: pool}
|
||||
}
|
||||
|
||||
var workerColumns = []string{"id", "name", "capabilities", "status", "last_heartbeat_at", "created_at", "updated_at"}
|
||||
var workerColumns = []string{"id", "name", "capabilities", "status", "owner_id", "trust_level", "last_heartbeat_at", "created_at", "updated_at"}
|
||||
|
||||
func (r *WorkerRepo) Insert(ctx context.Context, w *domain.Worker) error {
|
||||
sql, args, err := psql.Insert("workers").
|
||||
Columns(workerColumns...).
|
||||
// capabilities is a jsonb column; pgx marshals the []string to a JSON array.
|
||||
Values(w.ID, w.Name, w.Capabilities, string(w.Status),
|
||||
Values(w.ID, w.Name, w.Capabilities, string(w.Status), w.OwnerID, string(w.TrustLevel),
|
||||
w.LastHeartbeatAt, w.CreatedAt, w.UpdatedAt).
|
||||
ToSql()
|
||||
if err != nil {
|
||||
@@ -95,11 +95,13 @@ func scanWorker(row pgx.Row) (*domain.Worker, error) {
|
||||
var (
|
||||
w domain.Worker
|
||||
status string
|
||||
trust string
|
||||
)
|
||||
if err := row.Scan(&w.ID, &w.Name, &w.Capabilities, &status,
|
||||
if err := row.Scan(&w.ID, &w.Name, &w.Capabilities, &status, &w.OwnerID, &trust,
|
||||
&w.LastHeartbeatAt, &w.CreatedAt, &w.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
w.Status = domain.WorkerStatus(status)
|
||||
w.TrustLevel = domain.WorkerTrust(trust)
|
||||
return &w, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,60 @@
|
||||
// Package token verifies the HS256 JWTs minted by the userservice. The
|
||||
// coordinator only ever *verifies* — it never issues — so this is a deliberately
|
||||
// small counterpart to the userservice's issuer. Verification is local: the
|
||||
// shared secret is enough, with no runtime call back to the userservice.
|
||||
package token
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
// Claims is the subset of a userservice token the coordinator cares about.
|
||||
type Claims struct {
|
||||
UserID uuid.UUID
|
||||
Role string
|
||||
Verified bool
|
||||
}
|
||||
|
||||
// Verifier checks tokens against the shared HS256 secret.
|
||||
type Verifier struct {
|
||||
secret []byte
|
||||
}
|
||||
|
||||
// NewVerifier returns a Verifier, or nil when secret is empty — a nil Verifier
|
||||
// means user-JWT auth is disabled and only the shared service token is accepted.
|
||||
func NewVerifier(secret string) *Verifier {
|
||||
if secret == "" {
|
||||
return nil
|
||||
}
|
||||
return &Verifier{secret: []byte(secret)}
|
||||
}
|
||||
|
||||
type claims struct {
|
||||
Role string `json:"role"`
|
||||
Verified bool `json:"verified"`
|
||||
jwt.RegisteredClaims
|
||||
}
|
||||
|
||||
// Verify checks the signature and expiry and returns the identity. It pins the
|
||||
// algorithm to HMAC, rejecting a token that asks for "none" or an RS256 public
|
||||
// key — the classic algorithm-substitution attack.
|
||||
func (v *Verifier) Verify(raw string) (Claims, error) {
|
||||
var c claims
|
||||
_, err := jwt.ParseWithClaims(raw, &c, func(t *jwt.Token) (any, error) {
|
||||
if _, ok := t.Method.(*jwt.SigningMethodHMAC); !ok {
|
||||
return nil, fmt.Errorf("unexpected signing method: %v", t.Header["alg"])
|
||||
}
|
||||
return v.secret, nil
|
||||
})
|
||||
if err != nil {
|
||||
return Claims{}, err
|
||||
}
|
||||
id, err := uuid.Parse(c.Subject)
|
||||
if err != nil {
|
||||
return Claims{}, fmt.Errorf("token subject is not a uuid: %w", err)
|
||||
}
|
||||
return Claims{UserID: id, Role: c.Role, Verified: c.Verified}, nil
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
package token
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang-jwt/jwt/v5"
|
||||
"github.com/google/uuid"
|
||||
)
|
||||
|
||||
const secret = "coordinator-verify-secret-32-bytes!!"
|
||||
|
||||
func sign(t *testing.T, method jwt.SigningMethod, key any, sub, role string, exp time.Time) string {
|
||||
t.Helper()
|
||||
return signVerified(t, method, key, sub, role, false, exp)
|
||||
}
|
||||
|
||||
func signVerified(t *testing.T, method jwt.SigningMethod, key any, sub, role string, verified bool, exp time.Time) string {
|
||||
t.Helper()
|
||||
tok := jwt.NewWithClaims(method, claims{
|
||||
Role: role,
|
||||
Verified: verified,
|
||||
RegisteredClaims: jwt.RegisteredClaims{
|
||||
Subject: sub,
|
||||
ExpiresAt: jwt.NewNumericDate(exp),
|
||||
},
|
||||
})
|
||||
raw, err := tok.SignedString(key)
|
||||
if err != nil {
|
||||
t.Fatalf("sign: %v", err)
|
||||
}
|
||||
return raw
|
||||
}
|
||||
|
||||
func TestVerifyCarriesVerifiedClaim(t *testing.T) {
|
||||
v := NewVerifier(secret)
|
||||
raw := signVerified(t, jwt.SigningMethodHS256, []byte(secret), uuid.New().String(), "user", true, time.Now().Add(time.Hour))
|
||||
|
||||
claims, err := v.Verify(raw)
|
||||
if err != nil {
|
||||
t.Fatalf("verify: %v", err)
|
||||
}
|
||||
if !claims.Verified {
|
||||
t.Error("verified claim not read from token")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewVerifierNilWhenNoSecret(t *testing.T) {
|
||||
if NewVerifier("") != nil {
|
||||
t.Error("empty secret must yield a nil verifier (auth disabled)")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVerifyRoundTrip(t *testing.T) {
|
||||
v := NewVerifier(secret)
|
||||
id := uuid.New()
|
||||
raw := sign(t, jwt.SigningMethodHS256, []byte(secret), id.String(), "admin", time.Now().Add(time.Hour))
|
||||
|
||||
claims, err := v.Verify(raw)
|
||||
if err != nil {
|
||||
t.Fatalf("verify: %v", err)
|
||||
}
|
||||
if claims.UserID != id {
|
||||
t.Errorf("UserID = %v, want %v", claims.UserID, id)
|
||||
}
|
||||
if claims.Role != "admin" {
|
||||
t.Errorf("Role = %q, want admin", claims.Role)
|
||||
}
|
||||
}
|
||||
|
||||
func TestVerifyRejectsExpired(t *testing.T) {
|
||||
v := NewVerifier(secret)
|
||||
raw := sign(t, jwt.SigningMethodHS256, []byte(secret), uuid.New().String(), "user", time.Now().Add(-time.Minute))
|
||||
if _, err := v.Verify(raw); err == nil {
|
||||
t.Error("expired token accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVerifyRejectsWrongSecret(t *testing.T) {
|
||||
raw := sign(t, jwt.SigningMethodHS256, []byte(secret), uuid.New().String(), "user", time.Now().Add(time.Hour))
|
||||
if _, err := NewVerifier("another-secret-also-at-least-32-byte").Verify(raw); err == nil {
|
||||
t.Error("token verified under the wrong secret")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVerifyRejectsNoneAlg(t *testing.T) {
|
||||
raw := sign(t, jwt.SigningMethodNone, jwt.UnsafeAllowNoneSignatureType, uuid.New().String(), "admin", time.Now().Add(time.Hour))
|
||||
if _, err := NewVerifier(secret).Verify(raw); err == nil {
|
||||
t.Error("none-signed token accepted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestVerifyRejectsNonUUIDSubject(t *testing.T) {
|
||||
raw := sign(t, jwt.SigningMethodHS256, []byte(secret), "not-a-uuid", "user", time.Now().Add(time.Hour))
|
||||
if _, err := NewVerifier(secret).Verify(raw); err == nil {
|
||||
t.Error("non-uuid subject accepted")
|
||||
}
|
||||
}
|
||||
@@ -119,6 +119,8 @@ type jobProgressResponse struct {
|
||||
Done int `json:"completed"`
|
||||
Failed int `json:"failed"`
|
||||
Cancelled int `json:"cancelled"`
|
||||
ResultURI string `json:"result_uri,omitempty"`
|
||||
ErrorCode string `json:"error_code,omitempty"`
|
||||
}
|
||||
|
||||
type uploadArtifactResponse struct {
|
||||
@@ -153,7 +155,7 @@ func toClaimedTaskResponse(c domain.ClaimedTask) claimedTaskResponse {
|
||||
}
|
||||
|
||||
func toJobProgressResponse(p domain.JobProgress) jobProgressResponse {
|
||||
return jobProgressResponse{
|
||||
out := jobProgressResponse{
|
||||
ID: p.Job.ID,
|
||||
Status: string(p.DeriveStatus()),
|
||||
Total: p.Total,
|
||||
@@ -163,4 +165,11 @@ func toJobProgressResponse(p domain.JobProgress) jobProgressResponse {
|
||||
Failed: p.Failed,
|
||||
Cancelled: p.Cancelled,
|
||||
}
|
||||
if p.Job.ResultArtifactID != nil && out.Status == string(domain.JobCompleted) {
|
||||
out.ResultURI = "/jobs/" + p.Job.ID.String() + "/result"
|
||||
}
|
||||
if p.Job.ErrorCode != nil {
|
||||
out.ErrorCode = *p.Job.ErrorCode
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
)
|
||||
|
||||
func TestJobProgressResponseExposesFinalResultOnlyWhenCompleted(t *testing.T) {
|
||||
id := uuid.New()
|
||||
result := uuid.New()
|
||||
progress := domain.JobProgress{Job: domain.Job{
|
||||
ID: id, Status: domain.JobCompleted, ResultArtifactID: &result,
|
||||
}, Total: 1, Done: 1}
|
||||
if got, want := toJobProgressResponse(progress).ResultURI, "/jobs/"+id.String()+"/result"; got != want {
|
||||
t.Fatalf("result URI = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
progress.Job.Status = domain.JobReducing
|
||||
if got := toJobProgressResponse(progress).ResultURI; got != "" {
|
||||
t.Fatalf("reducing job exposes result URI %q", got)
|
||||
}
|
||||
}
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/authctx"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
)
|
||||
@@ -56,10 +57,24 @@ func (s *Server) handleRegister(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
worker, err := s.uc.RegisterWorker.Execute(ctx, usecase.RegisterWorkerInput{
|
||||
// Resolve the worker's trust tier from how the caller authenticated:
|
||||
// - shared service token (no requester) -> trusted lab worker
|
||||
// - verified/admin user JWT -> trusted volunteer
|
||||
// - plain user JWT -> untrusted (quarantined)
|
||||
in := usecase.RegisterWorkerInput{
|
||||
Name: req.Name,
|
||||
Capabilities: req.Capabilities,
|
||||
})
|
||||
TrustLevel: domain.WorkerTrusted,
|
||||
}
|
||||
if requester, ok := authctx.From(ctx); ok {
|
||||
id := requester.UserID
|
||||
in.OwnerID = &id
|
||||
if !requester.IsTrusted() {
|
||||
in.TrustLevel = domain.WorkerUntrusted
|
||||
}
|
||||
}
|
||||
|
||||
worker, err := s.uc.RegisterWorker.Execute(ctx, in)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
@@ -79,6 +94,10 @@ func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
if _, err := uuid.Parse(req.WorkerID); err != nil {
|
||||
s.writeError(w, r, domain.ErrInvalidInput)
|
||||
return
|
||||
}
|
||||
|
||||
claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{
|
||||
WorkerID: req.WorkerID,
|
||||
@@ -146,6 +165,12 @@ func (s *Server) handleResult(w http.ResponseWriter, r *http.Request) {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
if s.uc.ReduceJob != nil {
|
||||
if err := s.uc.ReduceJob.Execute(ctx, task.JobID); err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, taskResponse{ID: task.ID, JobID: task.JobID, Status: string(task.Status)})
|
||||
}
|
||||
|
||||
@@ -395,6 +420,25 @@ func (s *Server) handleGetJob(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, toJobProgressResponse(progress))
|
||||
}
|
||||
|
||||
func (s *Server) handleGetJobResult(w http.ResponseWriter, r *http.Request) {
|
||||
ctx, cancel := s.reqCtx(r)
|
||||
defer cancel()
|
||||
jobID, ok := s.pathUUID(w, r, "job_id")
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
art, body, err := s.uc.GetJobResult.Execute(ctx, jobID)
|
||||
if err != nil {
|
||||
s.writeError(w, r, err)
|
||||
return
|
||||
}
|
||||
defer func() { _ = body.Close() }()
|
||||
w.Header().Set("Content-Type", art.ContentType)
|
||||
w.Header().Set("Content-Length", strconv.FormatInt(art.SizeBytes, 10))
|
||||
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", art.Filename))
|
||||
_, _ = io.Copy(w, body)
|
||||
}
|
||||
|
||||
// handleCancelJob stops all non-terminal shards for an operator-requested job.
|
||||
// It is available to both the bearer API and the separately authenticated UI.
|
||||
func (s *Server) handleCancelJob(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
@@ -9,6 +9,9 @@ import (
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/authctx"
|
||||
tokenpkg "github.com/emil28092005/SciMesh/coordinator/internal/token"
|
||||
)
|
||||
|
||||
type ctxKey string
|
||||
@@ -41,25 +44,46 @@ func newRequestID() string {
|
||||
|
||||
// withAuth enforces the shared bearer token every worker presents.
|
||||
// An empty token disables the check (local development only).
|
||||
func withAuth(token string) func(http.Handler) http.Handler {
|
||||
// withAuth authenticates a request one of two ways. Workers (and legacy
|
||||
// submitters) present the shared service token. When user-JWT auth is enabled
|
||||
// (verifier != nil), a submitter may instead present a userservice JWT; on
|
||||
// success the requester is stamped into the context so the job use cases can
|
||||
// record owner_id and enforce ownership. An empty token with no verifier
|
||||
// disables auth entirely (dev only).
|
||||
func withAuth(token string, verifier *tokenpkg.Verifier) func(http.Handler) http.Handler {
|
||||
return func(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if token == "" {
|
||||
if token == "" && verifier == nil {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
presented := strings.TrimPrefix(r.Header.Get("Authorization"), "Bearer ")
|
||||
// Constant-time compare: a byte-by-byte early exit would let an
|
||||
// attacker recover the token by timing responses.
|
||||
if subtle.ConstantTimeCompare([]byte(presented), []byte(token)) != 1 {
|
||||
w.Header().Set("WWW-Authenticate", "Bearer")
|
||||
writeJSON(w, http.StatusUnauthorized, errorResponse{
|
||||
Error: "unauthorized",
|
||||
RequestID: requestIDFrom(r.Context()),
|
||||
})
|
||||
|
||||
// Shared service token: constant-time compare so a byte-by-byte
|
||||
// early exit cannot leak the token through response timing.
|
||||
if token != "" && subtle.ConstantTimeCompare([]byte(presented), []byte(token)) == 1 {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
|
||||
// Otherwise try a user JWT, if that path is configured.
|
||||
if verifier != nil && presented != "" {
|
||||
if claims, err := verifier.Verify(presented); err == nil {
|
||||
ctx := authctx.With(r.Context(), authctx.Requester{
|
||||
UserID: claims.UserID,
|
||||
Role: claims.Role,
|
||||
Verified: claims.Verified,
|
||||
})
|
||||
next.ServeHTTP(w, r.WithContext(ctx))
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
w.Header().Set("WWW-Authenticate", "Bearer")
|
||||
writeJSON(w, http.StatusUnauthorized, errorResponse{
|
||||
Error: "unauthorized",
|
||||
RequestID: requestIDFrom(r.Context()),
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,8 +7,11 @@ import (
|
||||
"context"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/metrics"
|
||||
tokenpkg "github.com/emil28092005/SciMesh/coordinator/internal/token"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
)
|
||||
|
||||
@@ -22,13 +25,16 @@ type UseCases struct {
|
||||
ClaimTask *usecase.ClaimTask
|
||||
RenewLease *usecase.RenewLease
|
||||
CompleteTask *usecase.CompleteTask
|
||||
ReduceJob *usecase.ReduceJob
|
||||
FailTask *usecase.FailTask
|
||||
GetJobStatus *usecase.GetJobStatus
|
||||
GetJobResult *usecase.GetJobResult
|
||||
CancelJob *usecase.CancelJob
|
||||
UploadArtifact *usecase.UploadArtifact
|
||||
DownloadArtifact *usecase.DownloadArtifact
|
||||
GetTaskInput *usecase.GetTaskInput
|
||||
Dashboard *usecase.Dashboard
|
||||
PreviewArtifact *usecase.PreviewArtifact
|
||||
}
|
||||
|
||||
type Server struct {
|
||||
@@ -37,21 +43,69 @@ type Server struct {
|
||||
requestTimeout time.Duration
|
||||
heartbeatInterval time.Duration
|
||||
maxUploadBytes int64
|
||||
// verifier validates userservice JWTs. nil disables user-JWT auth, leaving
|
||||
// only the shared service token — the pre-userservice behaviour.
|
||||
verifier *tokenpkg.Verifier
|
||||
// userserviceURL is the base URL the UI proxies login/registration to. Empty
|
||||
// keeps the static basic-auth UI.
|
||||
userserviceURL string
|
||||
// publicCoordinatorURL / publicUserserviceURL are the browser-facing URLs
|
||||
// rendered into the worker-enrollment command. Either may be empty; the
|
||||
// template falls back (own origin / userserviceURL respectively).
|
||||
publicCoordinatorURL string
|
||||
publicUserserviceURL string
|
||||
// httpClient makes the login/register calls to the userservice.
|
||||
httpClient *http.Client
|
||||
// docsDir serves the built MkDocs site at /ui/docs/. Empty disables it.
|
||||
docsDir string
|
||||
// metrics holds the Prometheus registry and HTTP instrumentation.
|
||||
metrics *metrics.Metrics
|
||||
// ready probes downstream dependencies (the database) for /health. Kept as
|
||||
// a func so the transport layer never imports pgx.
|
||||
ready func(context.Context) error
|
||||
}
|
||||
|
||||
func NewServer(uc UseCases, log *slog.Logger, requestTimeout, heartbeatInterval time.Duration,
|
||||
maxUploadBytes int64, ready func(context.Context) error) *Server {
|
||||
return &Server{
|
||||
uc: uc,
|
||||
log: log,
|
||||
requestTimeout: requestTimeout,
|
||||
heartbeatInterval: heartbeatInterval,
|
||||
maxUploadBytes: maxUploadBytes,
|
||||
ready: ready,
|
||||
maxUploadBytes int64, jwtSecret, userserviceURL string, m *metrics.Metrics, ready func(context.Context) error,
|
||||
publicURLs ...string) *Server {
|
||||
if m == nil {
|
||||
m = metrics.New()
|
||||
}
|
||||
// publicURLs is variadic so existing callers/tests need no change: [0] is the
|
||||
// public coordinator URL, [1] the public userservice URL; both optional.
|
||||
var publicCoordinatorURL, publicUserserviceURL string
|
||||
if len(publicURLs) > 0 {
|
||||
publicCoordinatorURL = strings.TrimRight(publicURLs[0], "/")
|
||||
}
|
||||
if len(publicURLs) > 1 {
|
||||
publicUserserviceURL = strings.TrimRight(publicURLs[1], "/")
|
||||
}
|
||||
docsDir := ""
|
||||
if len(publicURLs) > 2 {
|
||||
docsDir = publicURLs[2]
|
||||
}
|
||||
return &Server{
|
||||
uc: uc,
|
||||
log: log,
|
||||
requestTimeout: requestTimeout,
|
||||
heartbeatInterval: heartbeatInterval,
|
||||
maxUploadBytes: maxUploadBytes,
|
||||
verifier: tokenpkg.NewVerifier(jwtSecret),
|
||||
userserviceURL: strings.TrimRight(userserviceURL, "/"),
|
||||
publicCoordinatorURL: publicCoordinatorURL,
|
||||
publicUserserviceURL: publicUserserviceURL,
|
||||
docsDir: docsDir,
|
||||
httpClient: &http.Client{Timeout: 10 * time.Second},
|
||||
metrics: m,
|
||||
ready: ready,
|
||||
}
|
||||
}
|
||||
|
||||
// uiSessionMode reports whether the operator UI authenticates via userservice
|
||||
// login (cookie session) rather than the static basic-auth token. It needs both
|
||||
// a verifier (to check the JWT locally) and a userservice URL (to issue it).
|
||||
func (s *Server) uiSessionMode() bool {
|
||||
return s.verifier != nil && s.userserviceURL != ""
|
||||
}
|
||||
|
||||
// Handler builds the router. Go 1.22's ServeMux matches on method and path
|
||||
@@ -62,6 +116,7 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler {
|
||||
protected.HandleFunc("POST /jobs", s.handleCreateJob)
|
||||
protected.HandleFunc("POST /jobs/upload", s.handleUploadDataset)
|
||||
protected.HandleFunc("GET /jobs/{job_id}", s.handleGetJob)
|
||||
protected.HandleFunc("GET /jobs/{job_id}/result", s.handleGetJobResult)
|
||||
protected.HandleFunc("POST /jobs/{job_id}/cancel", s.handleCancelJob)
|
||||
protected.HandleFunc("POST /tasks/claim", s.handleClaim)
|
||||
protected.HandleFunc("GET /tasks/{task_id}/input", s.handleGetTaskInput)
|
||||
@@ -73,17 +128,68 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler {
|
||||
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /health", s.handleHealth)
|
||||
if len(uiToken) > 0 && uiToken[0] != "" && s.uc.Dashboard != nil {
|
||||
// Unauthenticated like /health, so a Prometheus scraper needs no credential.
|
||||
mux.Handle("GET /metrics", s.metrics.Handler())
|
||||
|
||||
hasBasicAuth := len(uiToken) > 0 && uiToken[0] != ""
|
||||
if s.uc.Dashboard != nil && (s.uiSessionMode() || hasBasicAuth) {
|
||||
ui := http.NewServeMux()
|
||||
ui.HandleFunc("GET /ui", s.handleUIHome)
|
||||
ui.HandleFunc("GET /ui/jobs/new", s.handleUINewJob)
|
||||
ui.HandleFunc("GET /ui/jobs/{job_id}", s.handleUIJob)
|
||||
ui.HandleFunc("GET /ui/api/jobs/{job_id}", s.handleUIJobJSON)
|
||||
ui.HandleFunc("POST /ui/api/jobs/{job_id}/cancel", s.handleCancelJob)
|
||||
ui.HandleFunc("POST /ui/api/jobs/upload", s.handleUploadDataset)
|
||||
ui.HandleFunc("GET /ui/jobs/{job_id}/artifacts/{artifact_id}", s.handleUIArtifactDownload)
|
||||
mux.Handle("/ui", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
|
||||
mux.Handle("/ui/", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
|
||||
|
||||
// The operator application routes, all requiring an authenticated caller.
|
||||
app := []struct {
|
||||
pattern string
|
||||
handler http.HandlerFunc
|
||||
}{
|
||||
{"GET /ui", s.handleUIHome},
|
||||
{"GET /ui/jobs/new", s.handleUINewJob},
|
||||
{"GET /ui/workloads", s.handleUIWorkloads},
|
||||
{"GET /ui/docs", s.handleUIDocsIndex},
|
||||
{"GET /ui/docs/{path...}", s.handleUIDocs},
|
||||
{"GET /ui/jobs/{job_id}", s.handleUIJob},
|
||||
{"GET /ui/api/overview", s.handleUIOverviewJSON},
|
||||
{"GET /ui/api/jobs/{job_id}", s.handleUIJobJSON},
|
||||
{"POST /ui/api/jobs/{job_id}/cancel", s.handleCancelJob},
|
||||
{"POST /ui/api/jobs/upload", s.handleUploadDataset},
|
||||
{"GET /ui/jobs/{job_id}/artifacts/{artifact_id}", s.handleUIArtifactDownload},
|
||||
{"GET /ui/jobs/{job_id}/artifacts/{artifact_id}/preview", s.handleUIArtifactPreview},
|
||||
}
|
||||
|
||||
if s.uiSessionMode() {
|
||||
// Public auth pages — reachable without a session so a user can log in.
|
||||
ui.HandleFunc("GET /ui/login", s.handleUILoginForm)
|
||||
ui.HandleFunc("POST /ui/login", s.handleUILogin)
|
||||
ui.HandleFunc("GET /ui/register", s.handleUIRegisterForm)
|
||||
ui.HandleFunc("POST /ui/register", s.handleUIRegister)
|
||||
ui.HandleFunc("POST /ui/logout", s.handleUILogout)
|
||||
gate := withUISession(s.verifier)
|
||||
for _, rt := range app {
|
||||
ui.Handle(rt.pattern, gate(rt.handler))
|
||||
}
|
||||
ui.Handle("GET /ui/profile", gate(http.HandlerFunc(s.handleUIProfile)))
|
||||
// Worker enrollment: a user creates/lists/revokes their own worker keys
|
||||
// and copies a ready-to-run command. Session-only — it proxies to the
|
||||
// userservice with the caller's token, so it has no meaning under basic
|
||||
// auth (which has no userservice).
|
||||
ui.Handle("GET /ui/workers/new", gate(http.HandlerFunc(s.handleUIAddWorker)))
|
||||
ui.Handle("GET /ui/api/worker-keys", gate(http.HandlerFunc(s.handleUIWorkerKeysList)))
|
||||
ui.Handle("POST /ui/api/worker-keys", gate(http.HandlerFunc(s.handleUIWorkerKeyCreate)))
|
||||
ui.Handle("POST /ui/api/worker-keys/{id}/revoke", gate(http.HandlerFunc(s.handleUIWorkerKeyRevoke)))
|
||||
// Admin panel: session + admin role.
|
||||
ui.Handle("GET /ui/admin", chain(http.HandlerFunc(s.handleUIAdmin), gate, requireAdmin))
|
||||
ui.Handle("POST /ui/admin/user-action", chain(http.HandlerFunc(s.handleUIAdminUserAction), gate, requireAdmin))
|
||||
} else {
|
||||
for _, rt := range app {
|
||||
ui.HandleFunc(rt.pattern, rt.handler)
|
||||
}
|
||||
}
|
||||
|
||||
common := []func(http.Handler) http.Handler{withRequestID, withAccessLog(s.log)}
|
||||
if !s.uiSessionMode() {
|
||||
common = append(common, withBasicAuth(uiToken[0]))
|
||||
}
|
||||
common = append(common, withSameOrigin)
|
||||
mux.Handle("/ui", chain(ui, common...))
|
||||
mux.Handle("/ui/", chain(ui, common...))
|
||||
} else {
|
||||
// More specific than the protected catch-all: UI absence is not an auth
|
||||
// failure and does not disclose that a UI feature is configured elsewhere.
|
||||
@@ -93,9 +199,10 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler {
|
||||
mux.Handle("/", chain(protected,
|
||||
withRequestID, // outermost: every response gets an ID,
|
||||
withAccessLog(s.log), // including the 401s below
|
||||
withAuth(token),
|
||||
withAuth(token, s.verifier),
|
||||
))
|
||||
return mux
|
||||
// Measure every request once, outermost, with a normalized route label.
|
||||
return s.metrics.Middleware(mux)
|
||||
}
|
||||
|
||||
// handleHealth reports readiness. It probes the database so an orchestrator
|
||||
|
||||
@@ -23,8 +23,9 @@ const token = "secret"
|
||||
const uiToken = "ui-secret"
|
||||
|
||||
type env struct {
|
||||
ts *httptest.Server
|
||||
blobs *memstore.BlobStore
|
||||
ts *httptest.Server
|
||||
blobs *memstore.BlobStore
|
||||
workerID string
|
||||
}
|
||||
|
||||
func newEnv(t *testing.T, ready func(context.Context) error) *env {
|
||||
@@ -41,26 +42,36 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur
|
||||
clk := memstore.NewClock(time.Date(2026, 7, 21, 12, 0, 0, 0, time.UTC))
|
||||
tx := memstore.Tx{}
|
||||
lease := 2 * time.Minute
|
||||
downloadArtifact := usecase.NewDownloadArtifact(arts, blobs)
|
||||
|
||||
uc := coordhttp.UseCases{
|
||||
RegisterWorker: usecase.NewRegisterWorker(work, clk),
|
||||
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk),
|
||||
ClaimTask: usecase.NewClaimTask(tasks, clk, lease),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk, 3, testCatalog()),
|
||||
ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease, testCatalog()),
|
||||
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
|
||||
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
|
||||
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
|
||||
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, work, memstore.NewTaskResultRepo(), tx, clk, 2, testCatalog()),
|
||||
ReduceJob: usecase.NewReduceJob(jobs, tasks, arts, blobs, tx, clk, testCatalog()),
|
||||
FailTask: usecase.NewFailTask(tasks, jobs, work, tx, clk, testCatalog()),
|
||||
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
|
||||
GetJobResult: usecase.NewGetJobResult(jobs, downloadArtifact),
|
||||
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
|
||||
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, clk),
|
||||
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
|
||||
UploadArtifact: usecase.NewUploadArtifact(tasks, work, arts, blobs, tx, clk),
|
||||
DownloadArtifact: downloadArtifact,
|
||||
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
|
||||
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
|
||||
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts), testCatalog()),
|
||||
PreviewArtifact: usecase.NewPreviewArtifact(memstore.NewUIReadRepo(jobs, tasks, work, arts), blobs),
|
||||
}
|
||||
srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, 1<<30, ready)
|
||||
worker, err := uc.RegisterWorker.Execute(context.Background(), usecase.RegisterWorkerInput{
|
||||
Name: "test-worker", Capabilities: []string{"w", "similarity-search"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("register test worker: %v", err)
|
||||
}
|
||||
srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, 1<<30, "", "", nil, ready)
|
||||
ts := httptest.NewServer(srv.Handler(token, configuredUIToken))
|
||||
t.Cleanup(ts.Close)
|
||||
return &env{ts: ts, blobs: blobs}
|
||||
return &env{ts: ts, blobs: blobs, workerID: worker.ID.String()}
|
||||
}
|
||||
|
||||
func healthy(context.Context) error { return nil }
|
||||
@@ -68,6 +79,7 @@ func healthy(context.Context) error { return nil }
|
||||
// do sends an authenticated JSON request and returns status + decoded body.
|
||||
func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) {
|
||||
t.Helper()
|
||||
body = strings.ReplaceAll(body, `"worker_id":"w1"`, `"worker_id":"`+e.workerID+`"`)
|
||||
req, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
if body != "" {
|
||||
@@ -141,11 +153,36 @@ func TestUIRequiresDistinctCredentialAndRendersDashboard(t *testing.T) {
|
||||
t.Fatalf("UI status: %d", resp.StatusCode)
|
||||
}
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if !strings.Contains(string(body), "SciMesh operator dashboard") {
|
||||
if !strings.Contains(string(body), "SciMesh control room") {
|
||||
t.Errorf("dashboard body missing title")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIOverviewReturnsLiveSafeProjection(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("create job: %d", code)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/api/overview", nil)
|
||||
req.SetBasicAuth("operator", uiToken)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
var overview map[string]any
|
||||
if err := json.NewDecoder(resp.Body).Decode(&overview); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK || overview["active_jobs"].(float64) != 1 || overview["online_workers"].(float64) != 1 {
|
||||
t.Fatalf("overview = (%d, %v)", resp.StatusCode, overview)
|
||||
}
|
||||
if _, leaked := overview["worker_auth_token"]; leaked {
|
||||
t.Fatal("overview must not expose authentication configuration")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIDisabledReturnsNotFound(t *testing.T) {
|
||||
e := newEnvWithUIToken(t, healthy, "")
|
||||
resp := e.get(t, "/ui")
|
||||
@@ -330,6 +367,26 @@ func TestRegisterRejectsNoCapabilities(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestClaimRequiresRegisteredWorkerAndUsesStoredCapabilities(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"not-a-uuid"}`); code != http.StatusBadRequest {
|
||||
t.Fatalf("invalid worker id claim = %d, want 400", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"11111111-1111-4111-8111-111111111111"}`); code != http.StatusNotFound {
|
||||
t.Fatalf("unregistered worker claim = %d, want 404", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`); code != http.StatusCreated {
|
||||
t.Fatalf("create job = %d", code)
|
||||
}
|
||||
code, worker := e.do(t, "POST", "/workers/register", `{"name":"search-only","capabilities":["similarity-search"]}`)
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("register = %d", code)
|
||||
}
|
||||
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"`+worker["worker_id"].(string)+`","capabilities":["w"]}`); code != http.StatusNoContent {
|
||||
t.Fatalf("forged capability claim = %d, want 204", code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFullLifecycle(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
|
||||
@@ -372,6 +429,113 @@ func TestFullLifecycle(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestSimilaritySearchLifecyclePublishesFinalResult(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
code, job := e.uploadDataset(t, "similarity-search", 10, "chembl_id\tcanonical_smiles\nA\tCC\n")
|
||||
if code != http.StatusCreated {
|
||||
t.Fatalf("upload job: %d", code)
|
||||
}
|
||||
jobID := job["job_id"].(string)
|
||||
code, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["similarity-search"]}`)
|
||||
if code != http.StatusOK {
|
||||
t.Fatalf("claim: %d", code)
|
||||
}
|
||||
taskID := claim["task_id"].(string)
|
||||
attempt := int(claim["attempt"].(float64))
|
||||
artifactID := e.putArtifact(t, taskID, "w1", attempt, "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.900000\n")
|
||||
if code, _ := e.do(t, "POST", "/tasks/"+taskID+"/result",
|
||||
`{"worker_id":"w1","attempt":`+itoa(attempt)+`,"result":{"artifact_id":"`+artifactID+`"}}`); code != http.StatusOK {
|
||||
t.Fatalf("complete: %d", code)
|
||||
}
|
||||
|
||||
code, progress := e.do(t, "GET", "/jobs/"+jobID, "")
|
||||
if code != http.StatusOK || progress["status"] != "completed" || progress["result_uri"] != "/jobs/"+jobID+"/result" {
|
||||
t.Fatalf("progress = (%d, %v)", code, progress)
|
||||
}
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+progress["result_uri"].(string), nil)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(resp.Body)
|
||||
if resp.StatusCode != http.StatusOK || string(body) != "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.900000\n" {
|
||||
t.Fatalf("final result = (%d, %q)", resp.StatusCode, body)
|
||||
}
|
||||
|
||||
uiRequest, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+jobID, nil)
|
||||
uiRequest.SetBasicAuth("operator", uiToken)
|
||||
uiResponse, err := http.DefaultClient.Do(uiRequest)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer uiResponse.Body.Close()
|
||||
uiBody, _ := io.ReadAll(uiResponse.Body)
|
||||
if uiResponse.StatusCode != http.StatusOK || !strings.Contains(string(uiBody), "Final result ready") || !strings.Contains(string(uiBody), "Preview CSV") || !strings.Contains(string(uiBody), "Processing speed") {
|
||||
t.Fatalf("final UI = (%d, %q)", uiResponse.StatusCode, uiBody)
|
||||
}
|
||||
|
||||
jsonRequest, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/api/jobs/"+jobID, nil)
|
||||
jsonRequest.SetBasicAuth("operator", uiToken)
|
||||
jsonResponse, err := http.DefaultClient.Do(jsonRequest)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer jsonResponse.Body.Close()
|
||||
var detail map[string]any
|
||||
if err := json.NewDecoder(jsonResponse.Body).Decode(&detail); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if jsonResponse.StatusCode != http.StatusOK || detail["final_result_available"] != true {
|
||||
t.Fatalf("final UI JSON = (%d, %v)", jsonResponse.StatusCode, detail)
|
||||
}
|
||||
artifacts := detail["artifacts"].([]any)
|
||||
var finalID string
|
||||
for _, raw := range artifacts {
|
||||
artifact := raw.(map[string]any)
|
||||
if artifact["kind"] == "final_result" && artifact["downloadable"] == true {
|
||||
finalID = artifact["id"].(string)
|
||||
break
|
||||
}
|
||||
}
|
||||
if finalID == "" {
|
||||
t.Fatalf("artifacts = %v, want downloadable final result", artifacts)
|
||||
}
|
||||
var inputID string
|
||||
for _, raw := range artifacts {
|
||||
artifact := raw.(map[string]any)
|
||||
if artifact["kind"] == "input" {
|
||||
inputID = artifact["id"].(string)
|
||||
break
|
||||
}
|
||||
}
|
||||
if inputID == "" {
|
||||
t.Fatalf("artifacts = %v, want input artifact", artifacts)
|
||||
}
|
||||
inputRequest, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+jobID+"/artifacts/"+inputID, nil)
|
||||
inputRequest.SetBasicAuth("operator", uiToken)
|
||||
inputResponse, err := http.DefaultClient.Do(inputRequest)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer inputResponse.Body.Close()
|
||||
if inputResponse.StatusCode != http.StatusNotFound {
|
||||
t.Fatalf("UI input download = %d, want 404", inputResponse.StatusCode)
|
||||
}
|
||||
previewRequest, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+jobID+"/artifacts/"+finalID+"/preview", nil)
|
||||
previewRequest.SetBasicAuth("operator", uiToken)
|
||||
previewResponse, err := http.DefaultClient.Do(previewRequest)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer previewResponse.Body.Close()
|
||||
previewBody, _ := io.ReadAll(previewResponse.Body)
|
||||
if previewResponse.StatusCode != http.StatusOK || !strings.Contains(string(previewBody), "Final result preview") || !strings.Contains(string(previewBody), "0.900000") {
|
||||
t.Fatalf("final preview = (%d, %q)", previewResponse.StatusCode, previewBody)
|
||||
}
|
||||
}
|
||||
|
||||
func TestForeignArtifactResultConflict(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in",
|
||||
@@ -393,9 +557,9 @@ func TestForeignArtifactResultConflict(t *testing.T) {
|
||||
|
||||
func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
|
||||
|
||||
code, body := e.uploadDataset(t, "w", 2, tsv)
|
||||
code, body := e.uploadDataset(t, "similarity-search", 2, tsv)
|
||||
if code != 201 {
|
||||
t.Fatalf("upload: status = %d", code)
|
||||
}
|
||||
@@ -421,7 +585,7 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) {
|
||||
t.Fatalf("get input: status = %d", resp.StatusCode)
|
||||
}
|
||||
shard, _ := io.ReadAll(resp.Body)
|
||||
if !strings.HasPrefix(string(shard), "id\tsmiles\n") {
|
||||
if !strings.HasPrefix(string(shard), "chembl_id\tcanonical_smiles\n") {
|
||||
t.Errorf("shard missing header: %q", shard)
|
||||
}
|
||||
}
|
||||
@@ -430,11 +594,12 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "w")
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "2")
|
||||
_ = mw.WriteField("max_rows", "3")
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
|
||||
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
|
||||
_ = mw.Close()
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
@@ -451,6 +616,29 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadDatasetRejectsMissingChEMBLColumns(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "2")
|
||||
fw, _ := mw.CreateFormFile("file", "not-chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
|
||||
_ = mw.Close()
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", mw.FormDataContentType())
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusBadRequest {
|
||||
t.Errorf("missing ChEMBL columns = %d, want 400", resp.StatusCode)
|
||||
}
|
||||
}
|
||||
|
||||
func TestErrorMappings(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
zero := "00000000-0000-0000-0000-000000000000"
|
||||
@@ -478,10 +666,11 @@ func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) {
|
||||
e := newEnv(t, healthy)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", "w")
|
||||
_ = mw.WriteField("workload", "similarity-search")
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", "not-a-number")
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
|
||||
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"))
|
||||
_ = mw.Close()
|
||||
|
||||
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
|
||||
@@ -505,6 +694,9 @@ func (e *env) putArtifact(t *testing.T, taskID, worker string, attempt int, data
|
||||
e.ts.URL+"/tasks/"+taskID+"/artifacts/r.csv", strings.NewReader(data))
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
req.Header.Set("Content-Type", "text/csv")
|
||||
if worker == "w1" {
|
||||
worker = e.workerID
|
||||
}
|
||||
req.Header.Set("X-Worker-ID", worker)
|
||||
req.Header.Set("X-Task-Attempt", itoa(attempt))
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
@@ -526,6 +718,7 @@ func (e *env) uploadDataset(t *testing.T, workload string, rows int, tsv string)
|
||||
var buf bytes.Buffer
|
||||
mw := multipart.NewWriter(&buf)
|
||||
_ = mw.WriteField("workload", workload)
|
||||
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
|
||||
_ = mw.WriteField("chunk_rows", itoa(rows))
|
||||
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
|
||||
_, _ = io.Copy(fw, strings.NewReader(tsv))
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
{{define "add-worker.html"}}
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>Add your machine · SciMesh</title>
|
||||
<style>
|
||||
:root{color:#e5efff;background:#08111f;font:16px/1.5 Inter,ui-sans-serif,system-ui,sans-serif;color-scheme:dark}*{box-sizing:border-box}body{margin:0;background:radial-gradient(circle at 10% -8%,#183f77 0,transparent 32rem),#08111f}.page{max-width:980px;margin:auto;padding:28px 22px 64px}a{color:#94bdff}.back{text-decoration:none}.eyebrow{margin:28px 0 4px;color:#7baaff;font-size:.78rem;font-weight:800;letter-spacing:.14em;text-transform:uppercase}h1{margin:0;color:#f4f8ff;font-size:clamp(2rem,5vw,3.25rem);letter-spacing:-.055em}.lead{max-width:720px;margin:10px 0 0;color:#aabed9;font-size:1.06rem}.layout{display:grid;grid-template-columns:1.45fr .8fr;gap:15px;margin-top:28px}.card,.aside,.notice{border:1px solid #294662;border-radius:15px;background:#0d1a2cdd;box-shadow:0 20px 45px #00000021}.card{padding:22px}.aside,.notice{padding:18px}.aside h2,.notice h2{margin:0;color:#f1f6ff;font-size:1rem}.aside p,.notice p{color:#9fb3cf}.aside ol{margin:13px 0 0;padding-left:20px;color:#aebfda}.aside li{margin:10px 0}label{display:block;margin:18px 0 5px;color:#eaf2ff;font-weight:750}input{width:100%;border:1px solid #42617f;border-radius:9px;padding:10px 11px;background:#0a1626;color:#e5efff;font:inherit}input:focus{outline:2px solid #5d97f5;outline-offset:1px}.hint{margin:5px 0;color:#92a9c6;font-size:.87rem}.button{display:inline-flex;margin-top:16px;border:0;border-radius:10px;padding:11px 15px;background:#67e3b8;color:#062018;font:inherit;font-weight:850;cursor:pointer}.button.secondary{background:#23344d;color:#dce8ff}.button:disabled{opacity:.6;cursor:wait}.command{margin-top:18px;border:1px solid #2c8c70;border-radius:11px;padding:14px;background:#0c2b2a;color:#a8f1d0}.command strong{color:#e6fff4}.command pre{margin:10px 0 0;padding:12px;overflow-x:auto;border-radius:8px;background:#061a19;color:#c8ffe8;font:.82rem/1.5 ui-monospace,SFMono-Regular,monospace;white-space:pre;word-break:normal}.keys{margin-top:14px;display:grid;gap:9px}.key{display:flex;align-items:center;justify-content:space-between;gap:12px;border:1px solid #294662;border-radius:11px;padding:12px 14px;background:#0a1626}.key .kn{color:#f3f7ff;font-weight:700}.key .kp{margin-top:3px;color:#8196b3;font-family:ui-monospace,SFMono-Regular,monospace;font-size:.76rem}.key .kd{color:#8fa6c3;font-size:.8rem}.revoke{border:1px solid #6a2a3a;border-radius:8px;padding:7px 11px;background:#2a1420;color:#ff9bad;font:inherit;font-weight:700;cursor:pointer}.empty{padding:20px;border:1px dashed #35516f;border-radius:12px;color:#9ab0cb;text-align:center}.error{margin:12px 0 0;color:#ffacba}.warn{color:#ffd08a}.hidden{display:none}.cap{display:inline-block;margin:4px 5px 0 0;border:1px solid #375978;border-radius:5px;padding:2px 6px;color:#b5d3f5;font:.76rem ui-monospace,SFMono-Regular,monospace}@media(max-width:720px){.layout{grid-template-columns:1fr}.page{padding:22px 14px}}
|
||||
</style>
|
||||
</head>
|
||||
<body data-coordinator="{{.CoordinatorURL}}" data-userservice="{{.UserserviceURL}}">
|
||||
<main class="page">
|
||||
<a class="back" href="/ui">← Back to control room</a><p class="eyebrow">Contribute compute</p><h1>Turn this computer into a worker</h1><p class="lead">Create a key, install the worker, and run one command. The worker binds to your account and pulls tasks whenever it is online.</p>
|
||||
<div class="layout">
|
||||
<section class="card">
|
||||
<h2 style="margin:0 0 4px;color:#f1f6ff">Your worker keys</h2>
|
||||
<p class="hint">A key is long-lived and does not expire like a login. The worker trades it for short-lived tokens automatically. Revoke a key to stop its machines.</p>
|
||||
<form id="create" novalidate>
|
||||
<label for="key-name">Name this machine <small>(optional)</small></label>
|
||||
<input id="key-name" name="name" maxlength="100" placeholder="e.g. home-desktop" autocomplete="off">
|
||||
<button class="button" id="create-btn" type="submit">Create key →</button>
|
||||
<p id="error" class="error" role="alert"></p>
|
||||
</form>
|
||||
<div id="command" class="command hidden"></div>
|
||||
<div id="keys" class="keys"></div>
|
||||
</section>
|
||||
<aside class="aside">
|
||||
<h2>Set it up</h2>
|
||||
<ol>
|
||||
<li><strong>Create a key</strong><br>Use the form; copy the command it generates.</li>
|
||||
<li><strong>Paste it in a terminal</strong><br>The command clones the project, sets up a Python environment, installs the worker, and starts it. The machine then appears under <a href="/ui">My machines</a>.</li>
|
||||
</ol>
|
||||
<h2 style="margin-top:24px">Will my results count?</h2>
|
||||
<p>Your worker is <strong>untrusted</strong> by default: its results are cross-checked and accepted once a second independent worker computes the same answer (quorum), or once an admin marks your account <strong>verified</strong> — then your workers are trusted and results count immediately.</p>
|
||||
<p><span class="cap">similarity-search</span> and other SDK workloads from the library run on volunteer workers.</p>
|
||||
</aside>
|
||||
</div>
|
||||
</main>
|
||||
<script>
|
||||
const coord=(document.body.dataset.coordinator||location.origin).replace(/\/+$/,'');
|
||||
const users=(document.body.dataset.userservice||'').replace(/\/+$/,'');
|
||||
const keysBox=document.querySelector('#keys'),cmdBox=document.querySelector('#command'),form=document.querySelector('#create'),nameInput=document.querySelector('#key-name'),createBtn=document.querySelector('#create-btn'),error=document.querySelector('#error');
|
||||
const node=(tag,text,cls)=>{const n=document.createElement(tag);if(text!==undefined)n.textContent=text;if(cls)n.className=cls;return n};
|
||||
const shq=s=>"'"+String(s).replace(/'/g,"'\\''")+"'";
|
||||
const buildCommand=(key,name)=>['git clone https://github.com/emil28092005/SciMesh.git','cd SciMesh','python -m venv .venv','source .venv/bin/activate','pip install -e .','','SCIMESH_COORDINATOR_URL='+coord+' \\','SCIMESH_USERSERVICE_URL='+(users||'<your userservice URL>')+' \\','SCIMESH_WORKER_KEY='+key+' \\','scimesh-worker --worker-name '+shq(name||'my-machine')].join('\n');
|
||||
const showCommand=(key,name)=>{cmdBox.replaceChildren();cmdBox.append(node('strong','Key created — copy it now, it is shown only once.'));const pre=node('pre',buildCommand(key,name));cmdBox.append(pre);const copy=node('button','Copy command','button secondary');copy.type='button';copy.addEventListener('click',()=>{navigator.clipboard&&navigator.clipboard.writeText(buildCommand(key,name)).then(()=>{copy.textContent='Copied ✓'},()=>{copy.textContent='Press Ctrl+C to copy'})});cmdBox.append(document.createElement('br'),copy);if(!users){cmdBox.append(node('p','Set SCIMESH_USERSERVICE_URL to a userservice URL your machine can reach.','warn'))}cmdBox.classList.remove('hidden')};
|
||||
const revoke=async id=>{const r=await fetch('/ui/api/worker-keys/'+encodeURIComponent(id)+'/revoke',{method:'POST'});if(r.status===204||r.ok){loadKeys()}else{error.textContent='Could not revoke the key.'}};
|
||||
const renderKeys=keys=>{keysBox.replaceChildren();if(!keys.length){keysBox.append(node('div','No keys yet. Create one above to connect a machine.','empty'));return}for(const k of keys){const row=node('div',undefined,'key'),left=node('div');left.append(node('div',k.name||'unnamed','kn'),node('div',k.prefix+'…','kp'),node('div','Created '+new Date(k.created_at).toLocaleString()+(k.last_used_at?' · last used '+new Date(k.last_used_at).toLocaleString():' · never used'),'kd'));const btn=node('button','Revoke','revoke');btn.type='button';btn.addEventListener('click',()=>revoke(k.id));row.append(left,btn);keysBox.append(row)}};
|
||||
const loadKeys=async()=>{try{const r=await fetch('/ui/api/worker-keys',{headers:{Accept:'application/json'}});if(!r.ok)throw Error();const data=await r.json();renderKeys(data.worker_keys||[])}catch(_){keysBox.replaceChildren(node('div','Could not load your keys.','empty'))}};
|
||||
form.addEventListener('submit',async event=>{event.preventDefault();error.textContent='';createBtn.disabled=true;const machineName=nameInput.value.trim();try{const r=await fetch('/ui/api/worker-keys',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({name:machineName})}),data=await r.json().catch(()=>({}));if(!r.ok)throw Error(data.error||'Could not create the key.');showCommand(data.key,machineName);nameInput.value='';loadKeys()}catch(err){error.textContent=err.message}finally{createBtn.disabled=false}});
|
||||
loadKeys();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
{{end}}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user