Compare commits
130
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6b67326c3b | ||
|
|
d57f8778ac | ||
|
|
5b2d5b5f6e | ||
|
|
b63441bb8c | ||
|
|
91d049b4df | ||
|
|
70c9e567ca | ||
|
|
81e854eb1d | ||
|
|
9a90d59ecf | ||
|
|
2b6531fa7e | ||
|
|
ab922356e9 | ||
|
|
fdea3a7152 | ||
|
|
c90c6b2ef6 | ||
|
|
ea0fb8c595 | ||
|
|
357ed34714 | ||
|
|
7ea7c52325 | ||
|
|
8db1f9846c | ||
|
|
7d1aacb19f | ||
|
|
0be114c8da | ||
|
|
6956ea6b1f | ||
|
|
5790ae78a7 | ||
|
|
40ff688416 | ||
|
|
1ec4b2e60b | ||
|
|
d309145290 | ||
|
|
15dd9651a8 | ||
|
|
30c441a7e9 | ||
|
|
cb51172885 | ||
|
|
029b26e6ae | ||
|
|
12499d2d7b | ||
|
|
c4d88c7ffc | ||
|
|
44247bd94e | ||
|
|
7c1d0dc568 | ||
|
|
82eaec4c55 | ||
|
|
c301499bef | ||
|
|
8b7094b0cb | ||
|
|
bdbe0e3e39 | ||
|
|
c11756c3d8 | ||
|
|
46645b8730 | ||
|
|
e923627ce0 | ||
|
|
7fb1059401 | ||
|
|
41cae546ff | ||
|
|
f8ff0956b9 | ||
|
|
7da9bf8c3a | ||
|
|
69a2d59b23 | ||
|
|
45d5c2eaad | ||
|
|
10a7e8df1d | ||
|
|
76746187eb | ||
|
|
9b0b9b208a | ||
|
|
a339ac853b | ||
|
|
0494b8bc2a | ||
|
|
0f9cbcec48 | ||
|
|
a4abdb970a | ||
|
|
b4fe7aaf72 | ||
|
|
e308085fe2 | ||
|
|
988101afab | ||
|
|
215325a9e0 | ||
|
|
c06d8673ce | ||
|
|
1473bbe2a8 | ||
|
|
9883def0c2 | ||
|
|
079eca071e | ||
|
|
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 |
@@ -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,146 @@
|
||||
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
|
||||
|
||||
wheel:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
# Full history: setuptools_scm derives the package version from the
|
||||
# release tag (v1.1.0-alpha.16 -> 1.1.0a16), so the wheel is
|
||||
# version-locked to the binaries of the same release.
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
fetch-depth: 0
|
||||
fetch-tags: true
|
||||
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: build the scimesh wheel
|
||||
run: |
|
||||
python -m pip install --quiet build
|
||||
python -m build --wheel --outdir dist
|
||||
ls -la dist/
|
||||
|
||||
- uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: wheel
|
||||
path: dist/*.whl
|
||||
if-no-files-found: error
|
||||
|
||||
release:
|
||||
needs: [binaries, wheel]
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- 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
|
||||
|
||||
- uses: actions/download-artifact@v4
|
||||
with:
|
||||
name: wheel
|
||||
path: artifacts
|
||||
|
||||
- name: checksums
|
||||
working-directory: artifacts
|
||||
run: sha256sum * > SHA256SUMS.txt
|
||||
|
||||
- uses: softprops/action-gh-release@v2
|
||||
with:
|
||||
files: |
|
||||
artifacts/*
|
||||
install.sh
|
||||
install.ps1
|
||||
# Pre-release tags (e.g. v1.1.0-alpha.1) publish as pre-releases.
|
||||
prerelease: ${{ contains(github.ref_name, '-alpha') }}
|
||||
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
|
||||
@@ -15,3 +15,6 @@ test_structures/
|
||||
# Local coordinator-worker execution state
|
||||
worker-data*/
|
||||
scimesh-worker-data/
|
||||
coordinator/.demo/
|
||||
site/
|
||||
coordinator/bin/
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
# Session Goal
|
||||
|
||||
выполни план на ночь (полный план ниже — автономное исполнение, пользователь спит)
|
||||
|
||||
## Ночная сессия — полный план (пользователь спит, 2026-08-03)
|
||||
|
||||
### 1. Баги (реальные, найденные в проде)
|
||||
- [ ] **Пустое имя воркера**: регистрация принимает `name=""` (в БД пользователя было 15 таких). Фикс: валидация в `domain.NewWorker` (TrimSpace != "") + регрессионный тест. Частично начат — в `worker_test.go` сломан тест (`fixedTime` vs `testNow`), доделать.
|
||||
- [ ] **`worker-agent --check` не видит venv**: проверяет только системный python3; после Install воркер работает через `~/.scimesh-worker/venv`. Выровнять с визардом (пробовать venv, если он есть).
|
||||
|
||||
### 2. Визуальный долг — рестайлинг старых страниц в дизайн-систему админки (#0b0e13, карточки, pill-статусы, кнопки)
|
||||
- [ ] `new-job.html` (форма запуска вычислений, UIElement-поля сохранить)
|
||||
- [ ] `job.html` (детали джоба: шарды, артефакты, прогресс, JS-логику сохранить)
|
||||
- [ ] `workloads.html`
|
||||
- [ ] `add-worker.html`
|
||||
- [ ] `profile.html`
|
||||
- [ ] Браузерная проверка каждой страницы (playwright).
|
||||
|
||||
### 3. Технический долг (паритет движков, обещанный планом)
|
||||
- [ ] **Postgres integration-тесты для admin-методов**: `SetTrust`, `WorkloadSettings` (Get/List/Set), `ListJobsPaginated`, метрики (`JobCountsByDay`, `TaskStats`, `ArtifactSizeByKind`, `DatabaseSizeBytes`) — сейчас покрыт только sqlite. (build tag integration, docker postgres как в CI.)
|
||||
|
||||
### 4. Фичи из плана (v2-задел)
|
||||
- [ ] **Статистика воркера в визарде**: статусная страница — claimed/completed/failed/heartbeat, парсинг из `worker.log` (обещано в docs/ui-admin-worker-plan.md §4).
|
||||
- [ ] **Admin: prune артефактов** (danger zone): удаление артефактов completed-джобов старше N дней (строки + blob-файлы), освобождает место; кнопка в Settings с подтверждением.
|
||||
- [ ] **Admin: удаление offline-воркеров** (устаревшие, как 15 мусорных) — кнопкой в Workers вместо SQL.
|
||||
|
||||
### 5. Гигиена
|
||||
- [ ] **`pyproject.toml` версия 0.1.0 → setuptools_scm** (версия из git-тегов), убрать sed-костыль из release.yml; локальный `pip install -e .` покажет правильную версию.
|
||||
- [ ] Синхронизация STATUS.md / README с фактическим состоянием (админка, визард, wheel, авто-открытие, restyle).
|
||||
|
||||
### 6. Верификация и релиз
|
||||
- [ ] Полный гейт: `go test -race ./...`, golangci-lint `--build-tags=integration`, pytest.
|
||||
- [ ] Браузерные проверки: админка (все секции), визард (install+start), рестайленные страницы.
|
||||
- [ ] E2E quorum-флоу (untrusted воркер через worker key → джоб требует quorum) в Docker, если останется время.
|
||||
- [ ] Релиз `v1.1.0-alpha.16` (бинарники + wheel), проверка ассетов.
|
||||
- [ ] Обновить этот файл: прогресс по пунктам, COMPLETED в конце.
|
||||
|
||||
## Plan (предыдущая задача — выполнена)
|
||||
1–7. Docker E2E пайплайна «install как человек → serve → визард → воркер → джоб» — выполнено, см. Progress ниже.
|
||||
|
||||
## Progress (утро: similarity-search-parallel)
|
||||
- ✅ Новый workload `similarity-search-parallel@1.0.0` (отдельная версия, как просил пользователь): подкласс `SimilaritySearchSDKWorkload` + параллельное ядро `search_parallel/core.py` — fingerprinting+скоринг шарда через `ThreadPoolExecutor` (параметр `threads`, default CPU count); `pool.map` сохраняет порядок строк, поэтому merge идентичен последовательному (`_HeapEntry`) и результат **байт-в-байт** равен `similarity-search` при любом числе потоков.
|
||||
- ✅ Тесты: байт-в-байт vs эталон для threads 1/2/4 с намеренными связями (изомеры, дубликаты), executor-прогон, валидация параметров, регистрация манифеста; 213 pytest зелёные.
|
||||
- ✅ Экспортирован в каталог координатора (5 ворклоадов), Go-тесты зелёные; release v1.1.0-alpha.19 (бинарники + wheel `scimesh-1.1.0a19`).
|
||||
- ✅ Распределённый E2E в Docker: воркер с wheel a19, джоб similarity-search-parallel (threads=4) → completed → результат байт-в-байт = локальному эталону.
|
||||
- ✅ Бинарники на машине пользователя обновлены до alpha.19 (для появления ворклоада в UI нужен рестарт serve + переустановка рантайма воркера).
|
||||
- Примечание: в CPython потоки не ускоряют чистый RDKit-путь (GIL), но структура готова к ядрам, отпускающим GIL (numpy и т.п.); при желании можно добавить процесс-пул как отдельный workload в будущей версии протокола.
|
||||
|
||||
## Progress (ночная сессия)
|
||||
- ✅ **П.1 Пустое имя воркера**: `domain.NewWorker` нормализует/отклоняет пустое имя + `TestNewWorkerRejectsBlankName` (починен `fixedTime`→`testNow`).
|
||||
- ✅ **П.2 `--check`**: пробует managed venv (если установлен) + реальный пробинг учётки (exchange ключа / claim-пробa) — `CheckAuth` + тесты; на машине пользователя: `✓ auth: credential accepted`, venv python, scimesh installed.
|
||||
- ✅ **П.3 Рестайлинг**: единый CSS-partial `ui-base.html` (дизайн-система админки), все 5 страниц (new-job, job, workloads, add-worker, profile) переведены, проверены в браузере без console-ошибок.
|
||||
- ✅ **П.4 Postgres integration**: admin-методы (SetTrust, ListJobsPaginated, TaskCounts, byDay/byWorkload, TaskStats, ArtifactSize, DB size, WorkloadSettings) + `ensureMigrated` для порядка запуска; весь suite зелёный.
|
||||
- ✅ **П.5 Статистика воркера в визарде**: лог `task claimed` в агента + парсинг registered/claimed/completed/failed → `/api/status.stats` + карточки в статусной странице + тест.
|
||||
- ✅ **П.6 Prune артефактов**: `JobRepository.ListCompletedBefore/Delete` (sqlite+postgres+memstore), usecase `PruneArtifacts` (каскад + blob-файлы), `POST /ui/admin/api/prune`, кнопка в Settings, тесты (sqlite+usecase); E2E: 200, freed bytes.
|
||||
- ✅ **П.7 Удаление offline-воркеров**: `WorkerRepository.Delete` + `Admin.RemoveWorker` (только offline) + `POST /ui/admin/api/workers/{id}/remove` + кнопка в Workers + тест; E2E: 204, строка удалена.
|
||||
- ✅ **П.8 setuptools_scm**: `dynamic = ["version"]`, CI-джоба wheel без sed (fetch-depth 0); локальная проверка: wheel на теге = `scimesh-1.1.0a16-py3-none-any.whl` (совпадает с Go-нормализацией); релизный ассет подтверждён.
|
||||
- ✅ **П.9 Docs**: STATUS.md синхронизирован (админка, визард, wheel, CTX-19/20 implemented).
|
||||
- ✅ **П.10 (доп.) Баг в serve-режиме**: worker-key exchange был недоступен снаружи (userservice на loopback) — добавлен прокси `POST /worker-tokens/exchange` на координаторе, `PublicUserserviceURL=""` + fallback на origin в add-worker. Проверено E2E.
|
||||
- ✅ **П.10 Quorum E2E (Docker)**: 2 untrusted-воркера с разными ключами (alice/bob) → джоб completed 3/3, в task_results по 2 голоса от разных владельцев с одинаковым sha256 → результат байт-в-байт = локальному эталону.
|
||||
- ✅ **Финальный гейт**: `go test -race ./...` ✅, golangci-lint 0 issues ✅, pytest 208 ✅, postgres integration ✅, Windows кросс-сборка ✅.
|
||||
- ✅ **Релиз v1.1.0-alpha.16** (бинарники + wheel `scimesh-1.1.0a16`), все воркфлоу success; бинарники на машине пользователя обновлены до alpha.16.
|
||||
|
||||
## Progress (прошлая работа — выполнено)
|
||||
- ✅ Релизы alpha.12–15: фикс версии визарда, venv task_runner, preflight через venv, MkdirAll при скачивании wheel, кнопка Install в шаблоне.
|
||||
- ✅ Docker E2E: координатор+воркер контейнеры, установка install.sh, визард (config→install→start), воркер online, джоб completed, результат байт-в-байт = локальному эталону.
|
||||
- ✅ На машине пользователя: визард alpha.15, правильный токен, venv из wheel, воркер emil-pc online, 15 пустых воркеров вычищены из БД.
|
||||
- ✅ Гейт: race + lint + pytest 208.
|
||||
|
||||
## Completion
|
||||
COMPLETED — ночной план выполнен полностью (10 пунктов + 2 найденных бага, включая E2E quorum на релизном коде). Все гейты зелёные, релиз v1.1.0-alpha.16 опубликован.
|
||||
|
||||
## Completion (предыдущая задача)
|
||||
COMPLETED — пайплайн доведён до рабочего состояния и проверен на релизных артефактах v1.1.0-alpha.14.
|
||||
@@ -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,212 @@ 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. All steps
|
||||
are implemented: embedded migrations with `AUTO_MIGRATE`, the `setup` wizard,
|
||||
the embedded SQLite storage backend (`SCIMESH_DB=sqlite`), the embedded
|
||||
userservice, and the `coordinator serve` single-binary mode with local worker
|
||||
agents. The standalone `users/` service and the PostgreSQL engine remain for
|
||||
cluster deployments.
|
||||
|
||||
**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.
|
||||
|
||||
---
|
||||
|
||||
### CTX-18 — Single-binary platform (`coordinator serve`)
|
||||
|
||||
**Goal:** A scientist installs one file, runs one command, and gets the whole
|
||||
platform: coordinator, both databases, the userservice, and local workers —
|
||||
no PostgreSQL, no Docker, no Python setup.
|
||||
|
||||
**Depends on:** CTX-17 (embedded migrations, SQLite, embedded userservice).
|
||||
|
||||
**Acceptance criteria:**
|
||||
|
||||
- `coordinator serve` provisions `~/.scimesh` (databases, secrets chmod 0600,
|
||||
generated admin password printed once, artifacts dir) and serves the UI on
|
||||
127.0.0.1:8080 by default; `--open` opens the browser;
|
||||
- `--workers N` spawns N `coordinator agent` subprocesses that claim and
|
||||
execute tasks locally; agents are stopped on shutdown;
|
||||
- the embedded userservice listens on the loopback interface and shares the
|
||||
JWT secret with the coordinator, so UI login/registration work unchanged;
|
||||
- a managed scientific runtime venv (`~/.scimesh/venv`) is bootstrapped on
|
||||
first start; `SCIMESH_PIP_PACKAGE` controls what gets installed (the PyPI
|
||||
name is not ours), and `TASK_RUNNER` points at the venv python;
|
||||
- `install.sh` / `install.ps1` detect the platform, download the release
|
||||
binary, and print the start command; both are release assets;
|
||||
- the PostgreSQL engine and the standalone userservice stay fully supported;
|
||||
- `coordinator serve` passes the full local E2E without any external service:
|
||||
health, login, upload, claim, compute, reduction, byte-exact result.
|
||||
|
||||
---
|
||||
|
||||
### CTX-19 — Coordinator Admin UI
|
||||
|
||||
**Goal:** Replace the demo-level operator pages with a real admin console
|
||||
(`/ui/admin`): system overview, job management with filters and pagination,
|
||||
worker management (trust, capabilities), users/roles and worker keys,
|
||||
workload enable/disable, settings (read-only), and storage metrics.
|
||||
|
||||
**Depends on:** CTX-11 (UI patterns, session/roles), CTX-15 (userservice
|
||||
proxy), the sqlite/postgres storage pair.
|
||||
|
||||
**Status: implemented.** Admin console at `/ui/admin`: system/storage/health,
|
||||
jobs (filter + pagination + owner resolution), workers with trust controls,
|
||||
users & worker keys (userservice admin endpoints), workload enable/disable
|
||||
(`workload_settings` migration in both engines, enforced at submit time and
|
||||
hidden from the job form), metrics (7-day buckets, failure rate), and
|
||||
settings with an audited worker-token reveal. All `/ui/admin/*` routes are
|
||||
admin-only and backed by bounded read models; sqlite + postgres parity;
|
||||
unit/permission/integration tests green.
|
||||
|
||||
### CTX-20 — Worker Setup UI
|
||||
|
||||
**Goal:** A local setup wizard embedded in the **worker-agent** binary
|
||||
(`worker-agent setup`, browser on 127.0.0.1) that turns any machine into a
|
||||
worker: coordinator URL + auth (serve token or worker key), work dir,
|
||||
connection check, start/stop, and a live status page. Runs on machines that
|
||||
have only the worker installed — no coordinator needed.
|
||||
|
||||
**Depends on:** the agent daemon; the worker-key exchange (CTX-15).
|
||||
|
||||
**Status: implemented.** `worker-agent setup` serves the local wizard on
|
||||
127.0.0.1 (default port 12700): coordinator URL + token or worker key,
|
||||
work dir and name, preflight check (coordinator/health, python3, scimesh),
|
||||
config saved to `~/.scimesh-worker/config.json` (0600), start/stop of the
|
||||
worker as a background process, and a live status page with the agent log.
|
||||
The daemon also gained `--config <path>` (environment still wins) and
|
||||
`--check`. End-to-end verified: the wizard started a real worker that
|
||||
registered with a coordinator and appeared in the admin console.
|
||||
|
||||
**Acceptance criteria:** full plan in
|
||||
[`docs/ui-admin-worker-plan.md`](docs/ui-admin-worker-plan.md) (section 4);
|
||||
the agent gains `setup`, `--config <path>`, and `--check`; end-to-end: the
|
||||
wizard starts a real worker that registers and completes a job.
|
||||
|
||||
---
|
||||
|
||||
## 10. Suggested assignment bundles
|
||||
@@ -835,6 +1114,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 +1228,23 @@ 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.
|
||||
- Publish the scimesh Python package to PyPI so the managed venv bootstrap
|
||||
(`SCIMESH_PIP_PACKAGE`) works out of the box on a scientist's machine.
|
||||
- Bundle a Python runtime (python-build-standalone) into the release so local
|
||||
workers need no system Python at all.
|
||||
- Native installers (.msi/.dmg/.deb) built by the release workflow.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -1,14 +1,26 @@
|
||||
# SciMesh
|
||||
|
||||
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
|
||||
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 Python worker 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).
|
||||
SciMesh is a local-first platform for scientific computation on molecular
|
||||
datasets. It turns a scientific run into independent tasks, dispatches them
|
||||
to worker agents, and deterministically combines the partial results into a
|
||||
checksum-protected final artifact.
|
||||
|
||||
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`.
|
||||
- **The Workload SDK (`scimesh.sdk`)** — a strict Python framework for
|
||||
authoring scientific workloads: `similarity-search` (exact top-k Tanimoto),
|
||||
`similarity-graph` (exact sparse graph), `descriptor-batch`, and
|
||||
`molwt-filter`. Workloads are ordinary user scripts built on the SDK; they
|
||||
run locally, in the conformance harness, and on claimed coordinator tasks
|
||||
without touching any other part of the program.
|
||||
- **The coordinator and worker agents** — a Go/PostgreSQL coordinator with an
|
||||
operator UI and Go worker agents that execute SDK workloads in a Python
|
||||
subprocess. The UI is workload-agnostic: the "New computation" form offers
|
||||
every workload from the embedded SDK library, and each workload declares its
|
||||
own form controls (`UIElement`) through the SDK.
|
||||
|
||||
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`. See
|
||||
[`STATUS.md`](STATUS.md) and [`PLAN.md`](PLAN.md).
|
||||
|
||||
## Installation
|
||||
|
||||
@@ -27,6 +39,66 @@ conda install -c conda-forge rdkit
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
## Releases
|
||||
|
||||
Every `v*` tag pushes a GitHub Release with static binaries for `coordinator`
|
||||
and `worker-agent` on linux/darwin/windows × amd64/arm64 (plus SHA-256
|
||||
checksums), the installer scripts, and the `coordinator` image on GHCR:
|
||||
|
||||
```bash
|
||||
docker pull ghcr.io/emil28092005/SciMesh/coordinator:latest
|
||||
```
|
||||
|
||||
For scientists: one command downloads the right binary, starts it, and opens
|
||||
the UI in the browser:
|
||||
|
||||
```bash
|
||||
# Linux / macOS — installs, starts and opens the admin console automatically
|
||||
curl -fsSL https://raw.githubusercontent.com/emil28092005/SciMesh/main/install.sh | bash
|
||||
|
||||
# Windows (PowerShell)
|
||||
powershell -ExecutionPolicy Bypass -c "irm https://raw.githubusercontent.com/emil28092005/SciMesh/main/install.ps1 | iex"
|
||||
```
|
||||
|
||||
Set `SCIMESH_AUTO_START=0` to install without starting anything. The old demo
|
||||
control room was removed: `/ui` is the admin console. A standalone
|
||||
worker is installed the same way (`bash -s worker`, or
|
||||
`SCIMESH_COMPONENT=worker` on Windows); its installer opens the local setup
|
||||
wizard (`worker-agent setup`) in the browser automatically.
|
||||
|
||||
`coordinator serve` is the single-binary mode: it embeds SQLite (coordinator +
|
||||
userservice databases), the userservice itself, and local worker agents
|
||||
(`--workers N`, default 1). On first start it generates secrets and the admin
|
||||
password under `~/.scimesh`, prints the login, and opens the UI. No
|
||||
PostgreSQL, no Docker, no environment variables. The scientific runtime is a
|
||||
managed venv (`~/.scimesh/venv`); point `SCIMESH_PIP_PACKAGE` at your scimesh
|
||||
wheel to install it automatically.
|
||||
|
||||
The coordinator's UI is the **admin console** at `/ui/admin` — cluster health
|
||||
and storage, paginated job table, worker fleet with trust controls, users and
|
||||
worker keys, workload enable/disable, metrics and the worker token. The job
|
||||
form (`/ui/jobs/new`), job detail pages and the workload library complete the
|
||||
operator surface; `/ui` redirects to the console, `serve --open` lands on it,
|
||||
and login returns you to the page you asked for. The **worker** binary (`worker-agent`) carries its own local setup
|
||||
wizard for machines that run only a worker: `worker-agent setup` opens a
|
||||
browser wizard at `127.0.0.1` that collects the coordinator URL and
|
||||
credential, runs a preflight check, saves `~/.scimesh-worker/config.json` and
|
||||
starts/stops the worker with a live log (see the
|
||||
[standalone docs](mkdocs/standalone.md)).
|
||||
|
||||
Manual download and run of a release binary:
|
||||
|
||||
```bash
|
||||
curl -L -o coordinator https://github.com/emil28092005/SciMesh/releases/latest/download/coordinator-linux-amd64
|
||||
chmod +x coordinator
|
||||
./coordinator --version
|
||||
```
|
||||
|
||||
Cluster deployments keep the PostgreSQL engine (`SCIMESH_DB=postgres` with
|
||||
`DATABASE_URL`, or `coordinator setup` to provision it) and the standalone
|
||||
userservice (`users/`). `coordinator agent` runs a worker agent from the same
|
||||
binary.
|
||||
|
||||
## Quick start
|
||||
|
||||
Run the built-in help command for copy-paste examples of both workloads:
|
||||
@@ -44,6 +116,41 @@ scimesh similarity-search --help
|
||||
scimesh similarity-graph --help
|
||||
```
|
||||
|
||||
## Manual pipeline demo
|
||||
|
||||
To inspect the coordinator, Web UI, and distributed 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
|
||||
`root@scimesh.local` and password `rootpassword`. The command starts
|
||||
PostgreSQL, the coordinator, and two Go worker agents (built by `make agent`;
|
||||
each executes the SDK workload in a Python subprocess). The **New computation**
|
||||
form offers every upload-ready workload from the installed library — the
|
||||
controls come from each workload's own SDK declarations. 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.
|
||||
@@ -111,3 +218,64 @@ pytest
|
||||
```
|
||||
|
||||
The package separates common dataset parsing and fingerprints from independent workloads. Add future workloads through the workload registry without changing the main CLI.
|
||||
|
||||
The coordinator and worker agent are Go modules under `coordinator/` and `users/`:
|
||||
|
||||
```bash
|
||||
cd coordinator && make coordinator agent && go test ./...
|
||||
```
|
||||
|
||||
On headless servers (no desktop environment), RDKit needs a few X11
|
||||
libraries that desktops already ship: `sudo apt-get install -y libxrender1
|
||||
libxext6 libxcursor1 libxfixes3 libxi6 libxrandr2`.
|
||||
|
||||
`make check` runs the full gate: vet, lint, race tests, the PostgreSQL
|
||||
integration suite, and the two-worker end-to-end smoke script.
|
||||
|
||||
## 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`, `descriptor-batch`, and
|
||||
`molwt-filter` 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; `scimesh workload export` writes the
|
||||
coordinator's embedded workload catalog, and `scimesh workload allowlist`
|
||||
prints the JSON for `SCIMESH_WORKLOAD_ALLOWLIST`.
|
||||
|
||||
Workloads can also declare how they should appear in the coordinator UI:
|
||||
a tuple of `UIElement`s (`scimesh.sdk.UIElement`) shapes the "New computation"
|
||||
form — widget, label, help, defaults, and ordering — plus the coordinator-side
|
||||
reduction mode (`reduction`: `top-k` or `ordered-concat`) and whether a single
|
||||
uploaded dataset can drive the workload (`upload_ready`). The strict parameter
|
||||
schema stays the authoritative validation contract.
|
||||
|
||||
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,8 @@
|
||||
# SciMesh Status
|
||||
|
||||
**Updated:** 2026-07-24
|
||||
**Branch baseline:** `main` at `6e67daa` (distributed similarity-search)
|
||||
**Updated:** 2026-08-03
|
||||
**Branch baseline:** `main`; this revision adds the admin console, the worker
|
||||
setup wizard, and release-shipped Python wheels.
|
||||
|
||||
## Current state
|
||||
|
||||
@@ -17,11 +18,51 @@ 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. 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.
|
||||
chunking, result/failure reporting, and job progress. The Go worker agent now
|
||||
uses the live coordinator contract. Completed shard results are reduced once
|
||||
into a checksum-protected final CSV, downloadable through the coordinator.
|
||||
|
||||
**Single-binary platform (`coordinator serve`)**: the coordinator now ships
|
||||
an embedded SQLite storage backend (`SCIMESH_DB=sqlite`), an embedded
|
||||
userservice, and `serve`/`agent` subcommands, so one downloaded binary runs
|
||||
the whole platform — coordinator, both databases, UI logins, and local
|
||||
workers — with no PostgreSQL, no Docker, and no environment variables. The
|
||||
first start provisions `~/.scimesh` (secrets, admin password printed once,
|
||||
managed scientific-runtime venv) and opens the UI. `install.sh` / `install.ps1`
|
||||
download the right release binary in one command and are release assets. The
|
||||
PostgreSQL engine, the `setup` wizard, and the standalone `users/` service
|
||||
remain fully supported for cluster deployments. The full E2E passes with zero
|
||||
external services: health, UI login, job upload, local agent compute,
|
||||
reduction, and a byte-exact final CSV.
|
||||
|
||||
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.
|
||||
|
||||
**Admin console** (`/ui/admin`, admin role only): system/storage/health,
|
||||
paginated jobs with owner resolution, worker fleet with trust controls and
|
||||
offline-worker removal, users and worker keys, workload enable/disable
|
||||
(persisted in both engines, enforced at submit time), 7-day metrics, and
|
||||
settings with an audited worker-token reveal and artifact pruning for old
|
||||
finished jobs.
|
||||
|
||||
**Worker setup wizard** (`worker-agent setup`, loopback `127.0.0.1`): browser
|
||||
wizard that configures the coordinator URL and credential, runs a preflight
|
||||
(coordinator reachability, real credential probe, python, scimesh), installs
|
||||
the version-locked `scimesh` wheel from the release into a managed venv
|
||||
(`~/.scimesh-worker/venv`), pins the venv task runner, and starts/stops the
|
||||
worker with live log and parsed claim/completed/failed counters. The release
|
||||
pipeline now also ships the Python wheel (`scimesh-<ver>-py3-none-any.whl`,
|
||||
version derived from the git tag via setuptools_scm), and both `coordinator
|
||||
serve` and the wizard install it automatically.
|
||||
|
||||
The full pipeline was verified end-to-end in containers on release artifacts:
|
||||
installer → serve → wizard → registration → real similarity-search job →
|
||||
byte-exact result vs the local reference.
|
||||
|
||||
## Milestone tracker
|
||||
|
||||
@@ -29,17 +70,27 @@ real PostgreSQL smoke test) passed on 2026-07-24.
|
||||
| --- | --- | --- |
|
||||
| CTX-00 API and error contract | Implemented | Contract, OpenAPI, and request examples are in `docs/`. |
|
||||
| CTX-01 Go coordinator bootstrap | Implemented | Go service and Docker runtime in `coordinator/`. |
|
||||
| CTX-02 PostgreSQL migrations | Implemented | Applied by the Compose migration service. |
|
||||
| CTX-03 Transactional queue | Implemented | Real-PostgreSQL integration tests cover atomic claims and concurrency. |
|
||||
| CTX-02 PostgreSQL migrations | Implemented | Embedded into the binary (`AUTO_MIGRATE`); the CLI path is still available for managed databases. |
|
||||
| CTX-03 Transactional queue | Implemented | Real-PostgreSQL integration tests cover atomic claims and concurrency; the SQLite backend mirrors the semantics. |
|
||||
| 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 | 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. |
|
||||
| CTX-11 Dashboard/operator view | Implemented | Protected local view: job/task/worker status, validated similarity-search upload, partial-artifact diagnostics, final-result download, and bounded polling. |
|
||||
| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
|
||||
| CTX-06 Python Worker live-contract alignment | Superseded | The Python worker daemon was removed; the Go worker agent (`coordinator/internal/agent/` + `cmd/worker-agent`, or `coordinator agent`) implements the lifecycle and executes SDK workloads via `scimesh/worker/task.py`. |
|
||||
| CTX-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering. |
|
||||
| CTX-08 Distributed similarity-search | Implemented | Planner/worker/reducer match the local reference byte-exactly. |
|
||||
| CTX-09 Reducer and final-result API | Implemented | Atomic `reducing` claim, deterministic coordinator-side reducers (`top-k` and `ordered-concat`), sanitized failure, final artifact, `result_uri`. |
|
||||
| CTX-10 Distributed similarity-graph | Not started | Local reference exists; the SDK-built local graph workload enforces the pair-coverage invariant. |
|
||||
| CTX-11 Dashboard/operator view | Implemented | The demo control room was removed; `/ui` lands on the admin console. Job form/detail, workload library, add-machine and profile pages share the admin design system; MkDocs at `/ui/docs/`. |
|
||||
| CTX-19 Coordinator Admin UI | Implemented | `/ui/admin` console (see above), bounded read models, admin-only routes, sqlite+postgres parity with integration tests. |
|
||||
| CTX-20 Worker Setup UI | Implemented | `worker-agent setup` wizard with runtime installer, `--config`, `--check` (credential + venv probing), live status. |
|
||||
| CTX-12 Reliability, security, CI | In progress | vet, gofmt, race tests, golangci-lint (0 issues), PostgreSQL integration, and smoke checks exist. |
|
||||
| CTX-15 User Service and access control | Implemented | User/owner scoping, verified contributors, worker keys, quorum; also embeddable (`coordinator serve`). |
|
||||
| CTX-16 Workload SDK foundation | Implemented | Strict immutable manifests/plans/artifacts, digest/trust-pinned tasks, negotiation, verifier primitives, conformance harness. |
|
||||
| CTX-17 Self-provisioning + setup wizard | Implemented | Embedded migrations, `coordinator setup`, SQLite backend, embedded userservice, `serve` mode. |
|
||||
| CTX-18 Single-binary platform | Implemented | `coordinator serve` (data dir, secrets, admin bootstrap, local agents, managed venv) + `install.sh`/`install.ps1`; full no-external-service E2E green. |
|
||||
| SDK roadmap step 3: `descriptor-batch` | Implemented | Byte-identical local/distributed output, quorum verifier test. |
|
||||
| SDK-built `similarity-search` and `similarity-graph` | Implemented | SDK-built packages, byte-identical to single-process references. |
|
||||
| SDK-built `molwt-filter` | Implemented | Minimal authoring example; also the single-binary E2E workload. |
|
||||
| SDK authoring scaffold | Implemented | `MapReduceWorkload` with `UIElement` declarations, `reduction`, `upload_ready`; generic `scimesh workload list|run|export|allowlist` CLI. |
|
||||
|
||||
## Next recommended assignment
|
||||
|
||||
@@ -48,14 +99,27 @@ block-pair planning and reduction for `similarity-graph`.
|
||||
|
||||
## Known constraints
|
||||
|
||||
- The worker/coordinator flow currently accepts both underscore API workload
|
||||
names and hyphenated CLI names while the contract is consolidated.
|
||||
- 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 --------------------------------------------------------
|
||||
|
||||
+109
-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 serve 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,110 @@ 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)
|
||||
|
||||
# The single-binary mode: everything embedded (sqlite + userservice + local
|
||||
# workers), no PostgreSQL or Docker. SETUP_ARGS=--workers 2 --open.
|
||||
serve: coordinator
|
||||
./bin/coordinator serve $(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 +168,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
-4
@@ -76,10 +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, diagnostic shard artifacts, and the
|
||||
final CSV for completed similarity-search jobs. The coordinator enters
|
||||
`reducing` after the last shard completes, then exposes the final deterministic
|
||||
global top-k result when merging succeeds.
|
||||
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
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/agent"
|
||||
)
|
||||
|
||||
// runAgent implements `coordinator agent`: the worker agent as a subcommand of
|
||||
// the same binary, so one file can serve the whole platform. `serve` spawns
|
||||
// these for its local workers.
|
||||
func runAgent(args []string) error {
|
||||
flags := flag.NewFlagSet("agent", flag.ContinueOnError)
|
||||
flags.Usage = func() {
|
||||
_, _ = fmt.Fprintf(flags.Output(), "usage: coordinator agent [options]\n")
|
||||
_, _ = fmt.Fprintf(flags.Output(), "Runs as a worker agent: claims tasks, executes SDK workloads in a\n")
|
||||
_, _ = fmt.Fprintf(flags.Output(), "Python subprocess, uploads results.\n\n")
|
||||
flags.PrintDefaults()
|
||||
}
|
||||
var (
|
||||
coordinatorURL = flags.String("coordinator-url", os.Getenv("COORDINATOR_URL"), "coordinator base URL")
|
||||
token = flags.String("token", os.Getenv("WORKER_AUTH_TOKEN"), "worker bearer token")
|
||||
workDir = flags.String("work-dir", os.Getenv("WORK_DIR"), "worker work directory")
|
||||
name = flags.String("name", os.Getenv("WORKER_NAME"), "worker name (default: hostname)")
|
||||
workerID = flags.String("worker-id", os.Getenv("WORKER_ID"), "persistent worker id (optional)")
|
||||
cpuCount = flags.Int("cpu", envInt("CPU_COUNT", 1), "advertised CPU cores")
|
||||
memoryMB = flags.Int("memory-mb", envInt("MEMORY_MB", 1024), "advertised memory in MiB")
|
||||
poll = flags.Duration("poll-interval", 2*time.Second, "claim poll interval")
|
||||
taskRunner = flags.String("task-runner", os.Getenv("TASK_RUNNER"), "python command + args that run scimesh.worker.task")
|
||||
maxTasks = flags.Int("max-tasks", envInt("MAX_TASKS", 0), "stop after N completed tasks (0 = unlimited)")
|
||||
exitWhenIdle = flags.Bool("exit-when-idle", false, "exit when the queue is empty")
|
||||
)
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() > 0 {
|
||||
return fmt.Errorf("agent takes no positional arguments")
|
||||
}
|
||||
if *coordinatorURL == "" || *token == "" || *workDir == "" {
|
||||
return fmt.Errorf("--coordinator-url, --token, and --work-dir are required")
|
||||
}
|
||||
if *taskRunner == "" {
|
||||
*taskRunner = "python -m scimesh.worker.task"
|
||||
}
|
||||
|
||||
logger := slog.New(slog.NewTextHandler(os.Stderr, nil))
|
||||
config := agent.Config{
|
||||
CoordinatorURL: strings.TrimRight(*coordinatorURL, "/"),
|
||||
Capabilities: agent.DefaultCapabilities(),
|
||||
Token: *token,
|
||||
WorkerName: *name,
|
||||
WorkerID: *workerID,
|
||||
WorkDir: *workDir,
|
||||
CPUCount: *cpuCount,
|
||||
MemoryMB: *memoryMB,
|
||||
PollInterval: *poll,
|
||||
RequestTimeout: 30 * time.Second,
|
||||
Heartbeat: 15 * time.Second,
|
||||
TaskRunner: strings.Fields(*taskRunner),
|
||||
MaxTasks: *maxTasks,
|
||||
ExitWhenIdle: *exitWhenIdle,
|
||||
}
|
||||
tokens := agent.NewTokenProvider("", "", config.Token, config.RequestTimeout)
|
||||
client := agent.NewClient(config.CoordinatorURL, tokens, config.RequestTimeout)
|
||||
runner := agent.NewTaskRunner(config.TaskRunner)
|
||||
daemon := agent.NewDaemon(&config, client, runner, logger)
|
||||
return daemon.RunForever()
|
||||
}
|
||||
|
||||
func envInt(name string, fallback int) int {
|
||||
raw := os.Getenv(name)
|
||||
if raw == "" {
|
||||
return fallback
|
||||
}
|
||||
var n int
|
||||
if _, err := fmt.Sscanf(raw, "%d", &n); err != nil {
|
||||
return fallback
|
||||
}
|
||||
return n
|
||||
}
|
||||
@@ -2,6 +2,8 @@ package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/signal"
|
||||
@@ -9,13 +11,55 @@ 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"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/storage/sqlite"
|
||||
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 {
|
||||
switch args[0] {
|
||||
case "setup":
|
||||
if err := runSetup(args[1:]); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "setup:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
case "serve":
|
||||
if err := runServe(args[1:]); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "serve:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
case "agent":
|
||||
if err := runAgent(args[1:]); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "agent:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
case "token":
|
||||
if err := runToken(args[1:]); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "token:", err)
|
||||
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 {
|
||||
@@ -23,9 +67,27 @@ func main() {
|
||||
}
|
||||
}
|
||||
|
||||
// storageDeps carries the engine-specific database handles and the repository
|
||||
// implementations. The usecases below only ever see the ports.
|
||||
type storageDeps struct {
|
||||
tx usecase.TxManager
|
||||
taskRepo usecase.TaskRepository
|
||||
jobRepo usecase.JobRepository
|
||||
workerRepo usecase.WorkerRepository
|
||||
artifactRepo usecase.ArtifactRepository
|
||||
uiReadRepo usecase.UIReadRepository
|
||||
adminReadRepo usecase.AdminReadRepository
|
||||
settingsRepo usecase.WorkloadSettingsRepository
|
||||
taskResultRepo usecase.TaskResultRepository
|
||||
statsRepo interface {
|
||||
Counts(ctx context.Context) (tasks, jobs, workers map[string]int, err error)
|
||||
}
|
||||
ready func(ctx context.Context) error
|
||||
migrate func(ctx context.Context, log *slog.Logger) error
|
||||
close func()
|
||||
}
|
||||
|
||||
func run() error {
|
||||
// Bootstrap logger, used only until config says where logs should go. It
|
||||
// writes to stderr so it never contaminates the configured stdout stream.
|
||||
boot := slog.New(slog.NewJSONHandler(os.Stderr, nil))
|
||||
|
||||
cfg, err := infra.LoadConfig()
|
||||
@@ -33,6 +95,16 @@ func run() error {
|
||||
boot.Error("load config", "err", err)
|
||||
return err
|
||||
}
|
||||
return runWithConfig(cfg)
|
||||
}
|
||||
|
||||
// runWithConfig boots the coordinator server with an explicit config. The
|
||||
// `serve` subcommand builds such a config for the single-binary mode; the
|
||||
// plain `coordinator` binary loads it from the environment.
|
||||
func runWithConfig(cfg infra.Config) error {
|
||||
// Bootstrap logger, used only until config says where logs should go. It
|
||||
// writes to stderr so it never contaminates the configured stdout stream.
|
||||
boot := slog.New(slog.NewJSONHandler(os.Stderr, nil))
|
||||
|
||||
// The real logger: stdout plus an optional rotated file (LOG_FILE).
|
||||
log, logCloser, err := infra.NewLogger(cfg)
|
||||
@@ -45,12 +117,29 @@ func run() error {
|
||||
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
pool, err := infra.NewPool(ctx, cfg, log)
|
||||
var deps *storageDeps
|
||||
switch cfg.DatabaseEngine {
|
||||
case "sqlite":
|
||||
deps, err = openSQLite(ctx, cfg, log)
|
||||
case "postgres":
|
||||
deps, err = openPostgres(ctx, cfg, log)
|
||||
default:
|
||||
err = fmt.Errorf("SCIMESH_DB must be sqlite or postgres")
|
||||
}
|
||||
if err != nil {
|
||||
log.Error("connect database", "err", err)
|
||||
log.Error("init storage", "err", err)
|
||||
return err
|
||||
}
|
||||
defer pool.Close()
|
||||
defer deps.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 := deps.migrate(ctx, log); err != nil {
|
||||
log.Error("apply migrations", "err", err)
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
blobStore, err := blob.NewFSStore(cfg.StorageDir)
|
||||
if err != nil {
|
||||
@@ -58,38 +147,55 @@ func run() error {
|
||||
return err
|
||||
}
|
||||
|
||||
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, taskRepo, jobRepo, workerRepo, artifactRepo, uiReadRepo, taskResultRepo :=
|
||||
deps.tx, deps.taskRepo, deps.jobRepo, deps.workerRepo, deps.artifactRepo, deps.uiReadRepo, deps.taskResultRepo
|
||||
|
||||
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, cfg.DefaultMaxAttempts),
|
||||
ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration),
|
||||
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk, cfg.DefaultMaxAttempts, catalog, deps.settingsRepo),
|
||||
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),
|
||||
ReduceJob: usecase.NewReduceJob(jobRepo, taskRepo, artifactRepo, blobStore, 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, tx, 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),
|
||||
PruneArtifacts: usecase.NewPruneArtifacts(jobRepo, uiReadRepo, blobStore, clk),
|
||||
PreviewArtifact: usecase.NewPreviewArtifact(uiReadRepo, blobStore),
|
||||
Admin: usecase.NewAdmin(deps.adminReadRepo, uiReadRepo, workerRepo, deps.settingsRepo, catalog,
|
||||
usecase.AdminNodeInfo{
|
||||
Version: version,
|
||||
StartedAt: clk.Now(),
|
||||
Binary: executablePath(),
|
||||
Addr: cfg.Addr,
|
||||
DataDir: cfg.StorageDir,
|
||||
DBEngine: cfg.DatabaseEngine,
|
||||
PublicURL: cfg.PublicCoordinatorURL,
|
||||
Userservice: cfg.UserserviceURL,
|
||||
WorkerToken: func() string { return cfg.Token },
|
||||
}, deps.ready, clk.Now).
|
||||
WithAuditLog(log, func(ctx context.Context, action, detail string) {
|
||||
log.Info("admin audit", "action", action, "detail", detail)
|
||||
}),
|
||||
}
|
||||
|
||||
// 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
|
||||
// the process would exit mid-UPDATE, and the deferred close() would pull
|
||||
// connections out from under them.
|
||||
expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk)
|
||||
expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk, catalog)
|
||||
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
@@ -107,9 +213,17 @@ func run() error {
|
||||
}(r.name, r.fn)
|
||||
}
|
||||
|
||||
// pool.Ping backs /health: readiness means the database answers, not just
|
||||
// Business metrics: gauges of tasks/jobs/workers by status, sampled from the
|
||||
// database on every Prometheus scrape.
|
||||
m := metrics.New()
|
||||
m.RegisterBusiness(func(ctx context.Context) (metrics.Stats, error) {
|
||||
tasks, jobs, workers, err := deps.statsRepo.Counts(ctx)
|
||||
return metrics.Stats{Tasks: tasks, Jobs: jobs, Workers: workers}, err
|
||||
})
|
||||
|
||||
// deps.ready 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, deps.ready, 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
|
||||
@@ -117,7 +231,7 @@ func run() error {
|
||||
//
|
||||
// 1. stop() cancel the context, telling the reaper to finish
|
||||
// 2. wg.Wait() let it return from its current tick
|
||||
// 3. deferred pool.Close() closes an idle pool, not a busy one
|
||||
// 3. deferred close() closes an idle pool, not a busy one
|
||||
//
|
||||
// Calling stop() here also covers the path where RunServer failed on its
|
||||
// own: the context would never be cancelled otherwise and wg.Wait()
|
||||
@@ -128,3 +242,64 @@ func run() error {
|
||||
|
||||
return err
|
||||
}
|
||||
|
||||
// executablePath resolves the running binary for the admin console's node
|
||||
// information, falling back to the invocation name.
|
||||
func executablePath() string {
|
||||
path, err := os.Executable()
|
||||
if err != nil || path == "" {
|
||||
return os.Args[0]
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// openSQLite opens the embedded database and builds the sqlite repositories.
|
||||
func openSQLite(ctx context.Context, cfg infra.Config, log *slog.Logger) (*storageDeps, error) {
|
||||
if err := os.MkdirAll(cfg.StorageDir, 0o750); err != nil {
|
||||
return nil, fmt.Errorf("create storage dir: %w", err)
|
||||
}
|
||||
db, err := sqlite.Open(cfg.DBPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
closeOnce := &sync.Once{}
|
||||
return &storageDeps{
|
||||
tx: sqlite.NewTxManager(db),
|
||||
taskRepo: sqlite.NewTaskRepo(db),
|
||||
jobRepo: sqlite.NewJobRepo(db),
|
||||
workerRepo: sqlite.NewWorkerRepo(db),
|
||||
artifactRepo: sqlite.NewArtifactRepo(db),
|
||||
uiReadRepo: sqlite.NewUIReadRepo(db),
|
||||
adminReadRepo: sqlite.NewAdminReadRepo(db),
|
||||
settingsRepo: sqlite.NewWorkloadSettingsRepo(db),
|
||||
taskResultRepo: sqlite.NewTaskResultRepo(db),
|
||||
statsRepo: sqlite.NewStatsRepo(db),
|
||||
ready: func(ctx context.Context) error { return db.PingContext(ctx) },
|
||||
migrate: func(ctx context.Context, log *slog.Logger) error { return sqlite.Migrate(ctx, db, log) },
|
||||
close: func() { closeOnce.Do(func() { _ = db.Close() }) },
|
||||
}, nil
|
||||
}
|
||||
|
||||
// openPostgres connects to PostgreSQL and builds the postgres repositories.
|
||||
func openPostgres(ctx context.Context, cfg infra.Config, log *slog.Logger) (*storageDeps, error) {
|
||||
pool, err := infra.NewPool(ctx, cfg, log)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
closeOnce := &sync.Once{}
|
||||
return &storageDeps{
|
||||
tx: postgres.NewTxManager(pool),
|
||||
taskRepo: postgres.NewTaskRepo(pool),
|
||||
jobRepo: postgres.NewJobRepo(pool),
|
||||
workerRepo: postgres.NewWorkerRepo(pool),
|
||||
artifactRepo: postgres.NewArtifactRepo(pool),
|
||||
uiReadRepo: postgres.NewUIReadRepo(pool),
|
||||
adminReadRepo: postgres.NewAdminReadRepo(pool),
|
||||
settingsRepo: postgres.NewWorkloadSettingsRepo(pool),
|
||||
taskResultRepo: postgres.NewTaskResultRepo(pool),
|
||||
statsRepo: postgres.NewStatsRepo(pool),
|
||||
ready: func(ctx context.Context) error { return pool.Ping(ctx) },
|
||||
migrate: func(ctx context.Context, log *slog.Logger) error { return postgres.Migrate(ctx, cfg.DatabaseURL, log) },
|
||||
close: func() { closeOnce.Do(pool.Close) },
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,339 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/agent"
|
||||
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"flag"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/infra"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/userservice"
|
||||
)
|
||||
|
||||
// runServe implements `coordinator serve`: the single-binary mode for a
|
||||
// scientist. It provisions a data directory (default ~/.scimesh) with the
|
||||
// coordinator and userservice sqlite databases, secrets, the admin account,
|
||||
// and optionally local worker agents — then runs the same server run() does.
|
||||
func runServe(args []string) error {
|
||||
flags := flag.NewFlagSet("serve", flag.ContinueOnError)
|
||||
flags.Usage = func() {
|
||||
_, _ = fmt.Fprintf(flags.Output(), "usage: coordinator serve [options]\n")
|
||||
_, _ = fmt.Fprintf(flags.Output(), "Runs the whole SciMesh platform from one binary: embedded databases, the\n")
|
||||
_, _ = fmt.Fprintf(flags.Output(), "userservice, and optional local workers. No PostgreSQL or Docker needed.\n\n")
|
||||
flags.PrintDefaults()
|
||||
}
|
||||
var (
|
||||
dataDir = flags.String("data-dir", defaultDataDir(), "data directory (default: ~/.scimesh)")
|
||||
addr = flags.String("addr", "127.0.0.1:8080", "listen address")
|
||||
workers = flags.Int("workers", 1, "number of local worker agents to spawn")
|
||||
open = flags.Bool("open", false, "open the UI in the browser")
|
||||
docsDir = flags.String("docs-dir", "", "built MkDocs site directory to serve at /ui/docs/")
|
||||
email = flags.String("admin-email", "admin@scimesh.local", "admin account email")
|
||||
password = flags.String("admin-password", "", "admin password (generated on first run when empty)")
|
||||
publicURL = flags.String("public-url", "", "browser/worker-facing coordinator URL (default: http://<addr>)")
|
||||
)
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() > 0 {
|
||||
return fmt.Errorf("serve takes no positional arguments")
|
||||
}
|
||||
if *workers < 0 {
|
||||
return fmt.Errorf("--workers must be >= 0")
|
||||
}
|
||||
|
||||
log := slog.New(slog.NewTextHandler(os.Stderr, nil))
|
||||
if err := os.MkdirAll(*dataDir, 0o750); err != nil {
|
||||
return fmt.Errorf("create data dir: %w", err)
|
||||
}
|
||||
|
||||
// 1. Secrets, persisted in the data dir so restarts keep working.
|
||||
jwtSecret, err := loadOrGenerate(filepath.Join(*dataDir, "jwt.secret"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
workerToken, err := loadOrGenerate(filepath.Join(*dataDir, "worker.token"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 2. Admin account: generated once and printed, remembered for later boots.
|
||||
if *password == "" {
|
||||
*password, err = loadOrGenerate(filepath.Join(*dataDir, "admin.password"))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// 3. Scientific runtime: ensure the managed venv (best effort).
|
||||
venvPython := filepath.Join(*dataDir, "venv", binName("bin/python"))
|
||||
ensureRuntime(log, *dataDir, venvPython)
|
||||
|
||||
// 4. Embedded userservice on the loopback interface.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
usersAddr, closeUsers, err := userservice.Serve(ctx, userservice.Config{
|
||||
DBPath: filepath.Join(*dataDir, "users.db"),
|
||||
JWTSecret: jwtSecret,
|
||||
AdminEmail: *email,
|
||||
AdminPassword: *password,
|
||||
Log: log,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("embedded userservice: %w", err)
|
||||
}
|
||||
defer func() { _ = closeUsers() }()
|
||||
|
||||
// 5. Local worker agents before the server, so they can claim immediately.
|
||||
coordinatorURL := "http://" + *addr
|
||||
agents, err := spawnAgents(ctx, log, *dataDir, *workers, coordinatorURL, workerToken, venvPython)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer stopAgents(agents)
|
||||
|
||||
// 6. The coordinator server itself.
|
||||
coordinatorPublicURL := *publicURL
|
||||
if coordinatorPublicURL == "" {
|
||||
coordinatorPublicURL = "http://" + *addr
|
||||
}
|
||||
cfg := infra.Config{
|
||||
Addr: *addr,
|
||||
DatabaseEngine: "sqlite",
|
||||
DBPath: filepath.Join(*dataDir, "scimesh.db"),
|
||||
Token: workerToken,
|
||||
JWTSecret: jwtSecret,
|
||||
UserserviceURL: "http://" + usersAddr,
|
||||
PublicCoordinatorURL: coordinatorPublicURL,
|
||||
PublicUserserviceURL: "http://" + usersAddr,
|
||||
LogLevel: "info",
|
||||
StorageDir: filepath.Join(*dataDir, "artifacts"),
|
||||
DocsDir: *docsDir,
|
||||
MaxUploadBytes: 1 << 30,
|
||||
DBMaxConns: 4,
|
||||
DBConnectTimeout: 10 * 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,
|
||||
AutoMigrate: true,
|
||||
}
|
||||
if *open {
|
||||
openBrowser("http://" + *addr + "/ui/admin")
|
||||
}
|
||||
|
||||
// Print the login once the server is about to start.
|
||||
fmt.Printf("\nSciMesh is starting at http://%s/ui\n", *addr)
|
||||
fmt.Printf(" admin login: %s / %s\n", *email, *password)
|
||||
if runtimeStatus(venvPython) {
|
||||
fmt.Printf(" scientific runtime: ready (%s)\n", venvPython)
|
||||
} else {
|
||||
fmt.Printf(" scientific runtime: NOT ready — install Python 3, then restart serve\n")
|
||||
}
|
||||
fmt.Printf(" data directory: %s\n\n", *dataDir)
|
||||
|
||||
err = runWithConfig(cfg)
|
||||
cancel()
|
||||
stopAgents(agents)
|
||||
return err
|
||||
}
|
||||
|
||||
// defaultDataDir returns the platform-appropriate data directory.
|
||||
func defaultDataDir() string {
|
||||
if dir := os.Getenv("SCIMESH_DATA_DIR"); dir != "" {
|
||||
return dir
|
||||
}
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil || home == "" {
|
||||
return ".scimesh"
|
||||
}
|
||||
return filepath.Join(home, ".scimesh")
|
||||
}
|
||||
|
||||
// loadOrGenerate reads a secret file, creating it with fresh random content
|
||||
// (chmod 0600) when missing.
|
||||
// #nosec G304 -- the path is an operator-supplied secret file inside the data dir.
|
||||
func loadOrGenerate(path string) (string, error) {
|
||||
if raw, err := os.ReadFile(path); err == nil {
|
||||
return strings.TrimSpace(string(raw)), nil
|
||||
}
|
||||
buffer := make([]byte, 32)
|
||||
if _, err := rand.Read(buffer); err != nil {
|
||||
return "", err
|
||||
}
|
||||
secret := hex.EncodeToString(buffer)
|
||||
if err := os.WriteFile(path, []byte(secret+"\n"), 0o600); err != nil {
|
||||
return "", fmt.Errorf("write %s: %w", path, err)
|
||||
}
|
||||
return secret, nil
|
||||
}
|
||||
|
||||
// spawnAgents starts `coordinator agent` subprocesses that claim tasks from
|
||||
// the coordinator. Each gets its own work directory under the data dir.
|
||||
func spawnAgents(ctx context.Context, log *slog.Logger, dataDir string, count int,
|
||||
coordinatorURL, token, venvPython string) ([]*exec.Cmd, error) {
|
||||
|
||||
var agents []*exec.Cmd
|
||||
for i := 0; i < count; i++ {
|
||||
workDir := filepath.Join(dataDir, "workers", fmt.Sprintf("%d", i))
|
||||
if err := os.MkdirAll(workDir, 0o750); err != nil {
|
||||
return agents, err
|
||||
}
|
||||
taskRunner := defaultTaskRunner(venvPython)
|
||||
// #nosec G204,G702 -- the command is this binary itself with operator flags.
|
||||
cmd := exec.CommandContext(ctx, os.Args[0], "agent",
|
||||
"--coordinator-url", coordinatorURL,
|
||||
"--token", token,
|
||||
"--work-dir", workDir,
|
||||
"--task-runner", taskRunner,
|
||||
)
|
||||
cmd.Stdout = os.Stdout
|
||||
cmd.Stderr = os.Stderr
|
||||
if err := cmd.Start(); err != nil {
|
||||
return agents, fmt.Errorf("start local agent %d: %w", i, err)
|
||||
}
|
||||
agents = append(agents, cmd)
|
||||
log.Info("local worker agent started", "index", i)
|
||||
}
|
||||
return agents, nil
|
||||
}
|
||||
|
||||
// stopAgents terminates the spawned agents and waits briefly for them.
|
||||
func stopAgents(agents []*exec.Cmd) {
|
||||
for _, agent := range agents {
|
||||
if agent.Process != nil {
|
||||
_ = agent.Process.Kill()
|
||||
}
|
||||
}
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
for _, agent := range agents {
|
||||
_, _ = agent.Process.Wait()
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
}
|
||||
}
|
||||
|
||||
// defaultTaskRunner picks the managed venv python when present, else the
|
||||
// system `python`.
|
||||
func defaultTaskRunner(venvPython string) string {
|
||||
if runtimeStatus(venvPython) {
|
||||
return venvPython + " -m scimesh.worker.task"
|
||||
}
|
||||
return "python -m scimesh.worker.task"
|
||||
}
|
||||
|
||||
// ensureRuntime creates the managed venv and installs scimesh into it, unless
|
||||
// it already exists. Best effort: a missing Python only logs a hint.
|
||||
func ensureRuntime(log *slog.Logger, dataDir, venvPython string) {
|
||||
if runtimeStatus(venvPython) {
|
||||
return
|
||||
}
|
||||
python := findPython()
|
||||
if python == "" {
|
||||
log.Warn("python3 not found; local workers need it to run scientific workloads")
|
||||
return
|
||||
}
|
||||
log.Info("creating the scientific runtime venv", "python", python)
|
||||
venvDir := filepath.Dir(filepath.Dir(venvPython))
|
||||
// #nosec G204 -- python comes from PATH and venvDir from the data dir.
|
||||
create := exec.CommandContext(context.Background(), python, "-m", "venv", venvDir)
|
||||
if out, err := create.CombinedOutput(); err != nil {
|
||||
log.Warn("venv creation failed; local workers need a manual Python install", "err", err, "output", string(out))
|
||||
return
|
||||
}
|
||||
pip := filepath.Join(venvDir, binName("bin/pip"))
|
||||
// The scimesh package is installed from an explicit source only: the PyPI
|
||||
// name belongs to an unrelated project, so `pip install scimesh` would
|
||||
// fetch a stranger's package. Default: download the wheel attached to our
|
||||
// own GitHub release for this binary version; SCIMESH_PIP_PACKAGE
|
||||
// overrides with a custom wheel, checkout or index.
|
||||
source := os.Getenv("SCIMESH_PIP_PACKAGE")
|
||||
if source == "" {
|
||||
url, _, err := agent.ReleaseWheelURL(version)
|
||||
if err != nil {
|
||||
log.Warn("scientific runtime venv created, but scimesh is not installed",
|
||||
"hint", "set SCIMESH_PIP_PACKAGE to your wheel or index")
|
||||
return
|
||||
}
|
||||
downloaded, err := agent.DownloadWheel(context.Background(), url, venvDir)
|
||||
if err != nil {
|
||||
log.Warn("could not download the scimesh wheel for this release",
|
||||
"err", err, "hint", "set SCIMESH_PIP_PACKAGE to your wheel or index")
|
||||
return
|
||||
}
|
||||
source = downloaded
|
||||
}
|
||||
// #nosec G204,G702 -- pip and source are operator-configured paths.
|
||||
install := exec.CommandContext(context.Background(), pip, "install", source)
|
||||
if out, err := install.CombinedOutput(); err != nil {
|
||||
log.Warn("pip install failed", "err", err, "output", string(out))
|
||||
return
|
||||
}
|
||||
log.Info("scientific runtime installed", "venv", venvDir)
|
||||
}
|
||||
|
||||
// findPython locates a usable python3.
|
||||
func findPython() string {
|
||||
for _, candidate := range []string{"python3", "python"} {
|
||||
path, err := exec.LookPath(candidate)
|
||||
if err == nil {
|
||||
return path
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// runtimeStatus reports whether the managed venv python exists.
|
||||
func runtimeStatus(venvPython string) bool {
|
||||
info, err := os.Stat(venvPython)
|
||||
return err == nil && !info.IsDir()
|
||||
}
|
||||
|
||||
// binName adapts a relative path to the platform layout.
|
||||
func binName(relative string) string {
|
||||
if runtime.GOOS == "windows" {
|
||||
parts := strings.Split(relative, "/")
|
||||
parts[len(parts)-1] += ".exe"
|
||||
return strings.Join(parts, string(filepath.Separator))
|
||||
}
|
||||
return relative
|
||||
}
|
||||
|
||||
// openBrowser opens the UI in the platform's default browser.
|
||||
func openBrowser(target string) {
|
||||
command := ""
|
||||
switch runtime.GOOS {
|
||||
case "darwin":
|
||||
command = "open"
|
||||
case "windows":
|
||||
command = "rundll32"
|
||||
default:
|
||||
command = "xdg-open"
|
||||
}
|
||||
if command == "rundll32" {
|
||||
// #nosec G204 -- target is the local UI URL the operator asked to open.
|
||||
_ = exec.CommandContext(context.Background(), "rundll32", "url.dll,FileProtocolHandler", target).Start()
|
||||
return
|
||||
}
|
||||
// #nosec G204 -- target is the local UI URL the operator asked to open.
|
||||
_ = exec.CommandContext(context.Background(), command, target).Start()
|
||||
}
|
||||
@@ -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,33 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
)
|
||||
|
||||
// runToken implements `coordinator token`: prints the worker auth token of a
|
||||
// `coordinator serve` instance, so a scientist can join a worker without
|
||||
// hunting through the data directory. The file itself is what serve created.
|
||||
func runToken(args []string) error {
|
||||
flags := flag.NewFlagSet("token", flag.ContinueOnError)
|
||||
flags.Usage = func() {
|
||||
_, _ = fmt.Fprintf(flags.Output(), "usage: coordinator token [options]\n")
|
||||
_, _ = fmt.Fprintf(flags.Output(), "Prints the WORKER_AUTH_TOKEN of this coordinator's serve instance.\n\n")
|
||||
flags.PrintDefaults()
|
||||
}
|
||||
var dataDir = flags.String("data-dir", defaultDataDir(), "data directory (default: ~/.scimesh)")
|
||||
if err := flags.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
if flags.NArg() > 0 {
|
||||
return fmt.Errorf("token takes no positional arguments")
|
||||
}
|
||||
token, err := os.ReadFile(filepath.Join(*dataDir, "worker.token"))
|
||||
if err != nil {
|
||||
return fmt.Errorf("no worker token found in %s — start the coordinator with `coordinator serve` first", *dataDir)
|
||||
}
|
||||
fmt.Print(string(token))
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,238 @@
|
||||
// Command worker-agent is the Go worker agent: a coordinator client that
|
||||
// executes SDK workloads in a Python subprocess per claimed task. It also
|
||||
// carries the local setup wizard (`worker-agent setup`) so a machine that
|
||||
// installs only the worker can configure and start itself without a
|
||||
// coordinator on site.
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"os/signal"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/agent"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/agent/setupui"
|
||||
)
|
||||
|
||||
// version is injected at build time (-ldflags "-X main.version=...") and
|
||||
// reported by --version. "dev" marks a local build.
|
||||
var version = "dev"
|
||||
|
||||
func main() {
|
||||
// The wizard and --check need the injected build version too (they resolve
|
||||
// the release wheel matching this binary), so it is set before dispatch.
|
||||
agent.Version = version
|
||||
|
||||
if len(os.Args) > 1 && os.Args[1] == "setup" {
|
||||
os.Exit(runSetup(os.Args[2:]))
|
||||
}
|
||||
|
||||
fs := flag.NewFlagSet("worker-agent", flag.ExitOnError)
|
||||
showVersion := fs.Bool("version", false, "print the build version and exit")
|
||||
configPath := fs.String("config", "", "path to a JSON config file (SCIMESH_WORKER_CONFIG overrides the default)")
|
||||
checkMode := fs.Bool("check", false, "run the preflight check (coordinator + local runtime) and exit 0/1")
|
||||
checkURL := fs.String("coordinator-url", "", "coordinator URL to probe in --check mode")
|
||||
_ = fs.Parse(os.Args[1:])
|
||||
|
||||
if *showVersion {
|
||||
fmt.Println("worker-agent " + version)
|
||||
return
|
||||
}
|
||||
|
||||
if *checkMode {
|
||||
url := *checkURL
|
||||
if url == "" {
|
||||
url = os.Getenv("COORDINATOR_URL")
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
if url == "" {
|
||||
fmt.Println("check: no coordinator URL (pass --coordinator-url or set COORDINATOR_URL)")
|
||||
os.Exit(1)
|
||||
}
|
||||
// Probe the managed venv when the wizard has installed it: workloads
|
||||
// run with that interpreter, so checking the bare system python3
|
||||
// would report a false negative.
|
||||
checkPython, checkToken, checkKey, checkUsers := "", "", "", ""
|
||||
if configPath := checkConfigPath(); configPath != "" {
|
||||
checkPython = agent.VenvPython(configPath)
|
||||
if config, err := agent.LoadConfigFile(configPath); err == nil {
|
||||
checkToken, checkKey, checkUsers = config.Token, config.WorkerKey, config.UserserviceURL
|
||||
}
|
||||
}
|
||||
report := agent.RunCheck(ctx, url, checkPython, checkToken, checkKey, checkUsers)
|
||||
printCheck(report)
|
||||
if !report.Coordinator.OK || !report.Python.OK || !report.Scimesh.OK {
|
||||
os.Exit(1)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
config, err := loadConfig(*configPath)
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
// checkConfigPath resolves where the wizard's config would be, honouring
|
||||
// SCIMESH_WORKER_CONFIG like the rest of the agent.
|
||||
func checkConfigPath() string {
|
||||
if env := os.Getenv("SCIMESH_WORKER_CONFIG"); env != "" {
|
||||
return env
|
||||
}
|
||||
return agent.DefaultConfigPath()
|
||||
}
|
||||
|
||||
// loadConfig prefers a --config file; environment variables override the file
|
||||
// (see agent.ConfigFile.Config). Without a file, the plain environment path is
|
||||
// used exactly as before.
|
||||
func loadConfig(configPath string) (*agent.Config, error) {
|
||||
if configPath != "" {
|
||||
return agent.LoadConfigFile(configPath)
|
||||
}
|
||||
envPath := os.Getenv("SCIMESH_WORKER_CONFIG")
|
||||
if envPath != "" {
|
||||
if _, err := os.Stat(envPath); err == nil { //nolint:gosec // G703: path is the operator's own env var
|
||||
return agent.LoadConfigFile(envPath)
|
||||
}
|
||||
}
|
||||
return agent.LoadConfig()
|
||||
}
|
||||
|
||||
func printCheck(report agent.CheckReport) {
|
||||
fmt.Printf("worker-agent %s\n", report.Agent)
|
||||
line := func(item agent.CheckItem) string {
|
||||
mark := "✗"
|
||||
if item.OK {
|
||||
mark = "✓"
|
||||
}
|
||||
detail := item.Detail
|
||||
if item.Latency > 0 {
|
||||
detail = fmt.Sprintf("%s (%d ms)", detail, item.Latency)
|
||||
}
|
||||
return fmt.Sprintf(" %s %s: %s", mark, item.Name, detail)
|
||||
}
|
||||
fmt.Println(line(report.Coordinator))
|
||||
fmt.Println(line(report.Auth))
|
||||
fmt.Println(line(report.Python))
|
||||
fmt.Println(line(report.Scimesh))
|
||||
}
|
||||
|
||||
// runSetup serves the local setup wizard until interrupted. It binds the
|
||||
// loopback interface only.
|
||||
func runSetup(args []string) int {
|
||||
fs := flag.NewFlagSet("worker-agent setup", flag.ContinueOnError)
|
||||
port := fs.Int("port", 0, "listen port (default 12700)")
|
||||
noOpen := fs.Bool("no-open", false, "do not open the browser automatically")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return 2
|
||||
}
|
||||
logger := slog.New(slog.NewTextHandler(os.Stderr, nil))
|
||||
|
||||
ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
|
||||
defer stop()
|
||||
|
||||
server := setupui.New(logger, setupui.Options{
|
||||
Port: *port,
|
||||
OpenBrowser: func(url string) {
|
||||
if *noOpen {
|
||||
return
|
||||
}
|
||||
openBrowser(url)
|
||||
},
|
||||
})
|
||||
listener, err := server.Listen()
|
||||
if err != nil {
|
||||
// A wizard may already be running on this port (left open, or a
|
||||
// second terminal). If it is ours, opening the browser is the
|
||||
// friendlier outcome than failing the command.
|
||||
if existing := wizardAlreadyRunning(*port); existing != "" {
|
||||
logger.Info("the setup wizard is already running", "url", existing)
|
||||
if !*noOpen {
|
||||
openBrowser(existing)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
logger.Error("setup wizard could not bind the loopback port", "err", err)
|
||||
return 1
|
||||
}
|
||||
url := "http://" + listener.Addr().String()
|
||||
logger.Info("SciMesh worker setup wizard", "url", url, "press-ctrl-c-to-stop", true)
|
||||
server.OpenBrowser(url)
|
||||
// Block until the signal arrives (never returns an error that matters: a
|
||||
// cancelled context is the normal exit path).
|
||||
err = server.Serve(ctx, listener)
|
||||
if err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
logger.Error("setup wizard stopped", "err", err)
|
||||
return 1
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// openBrowser points the user's default browser at the wizard. Best-effort:
|
||||
// a missing browser must never fail the setup flow.
|
||||
func openBrowser(url string) {
|
||||
for _, candidate := range [][]string{
|
||||
{"xdg-open", url},
|
||||
{"open", url},
|
||||
{"cmd", "/c", "start", url},
|
||||
} {
|
||||
binary, err := exec.LookPath(candidate[0])
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
//nolint:gosec // G204: candidates are our own fixed list; the url is a loopback literal
|
||||
_ = exec.CommandContext(context.Background(), binary, candidate[1:]...).Start()
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// wizardAlreadyRunning probes the requested loopback port and returns its URL
|
||||
// when it serves the setup wizard page, or "" when it does not (another
|
||||
// process, or nothing at all).
|
||||
func wizardAlreadyRunning(port int) string {
|
||||
url := fmt.Sprintf("http://127.0.0.1:%d/", port)
|
||||
client := http.Client{Timeout: 2 * time.Second}
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return ""
|
||||
}
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<16))
|
||||
if err != nil || !strings.Contains(string(body), "SciMesh Worker · Setup") {
|
||||
return ""
|
||||
}
|
||||
return url
|
||||
}
|
||||
@@ -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"
|
||||
|
||||
+21
-3
@@ -1,23 +1,41 @@
|
||||
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/dustin/go-humanize v1.0.1 // 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/mattn/go-isatty v0.0.20 // indirect
|
||||
github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
|
||||
github.com/ncruces/go-strftime v1.0.0 // 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
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // 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
|
||||
modernc.org/libc v1.74.1 // indirect
|
||||
modernc.org/mathutil v1.7.1 // indirect
|
||||
modernc.org/memory v1.11.0 // indirect
|
||||
modernc.org/sqlite v1.55.0 // indirect
|
||||
)
|
||||
|
||||
+45
-6
@@ -1,10 +1,20 @@
|
||||
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/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
|
||||
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
|
||||
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,23 +31,52 @@ 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/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
|
||||
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
|
||||
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/ncruces/go-strftime v1.0.0 h1:HMFp8mLCTPp341M/ZnA4qaf7ZlsbTc+miZjCLOFAw7w=
|
||||
github.com/ncruces/go-strftime v1.0.0/go.mod h1:Fwc5htZGVVkseilnfgOVb9mKy6w1naJmn9CehxcKcls=
|
||||
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/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE=
|
||||
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
|
||||
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.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
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=
|
||||
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
|
||||
modernc.org/libc v1.74.1 h1:bdR4VTKFMC4966QSNZ05XLGI/VwzVa2kTUX51Dm0riQ=
|
||||
modernc.org/libc v1.74.1/go.mod h1:uH4t5bOx3G3g9Xcmj10YKlTcVISlRDwv8VoQJG9n8Os=
|
||||
modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
|
||||
modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
|
||||
modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=
|
||||
modernc.org/memory v1.11.0/go.mod h1:/JP4VbVC+K5sU2wZi9bHoq2MAkCnrt2r98UGeSK7Mjw=
|
||||
modernc.org/sqlite v1.55.0 h1:hIFh0MCH0rGinQ/4KYb5/UbCkRkb+UP+OkLCVWa5MTM=
|
||||
modernc.org/sqlite v1.55.0/go.mod h1:4ntCLuNmnH8+GNqjka1wNg7KJd5/Hi5FYp8K+XQ7GZw=
|
||||
|
||||
@@ -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,213 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"os/exec"
|
||||
"runtime"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
guuid "github.com/google/uuid"
|
||||
)
|
||||
|
||||
// CheckItem is one line of the preflight report the setup wizard shows.
|
||||
type CheckItem struct {
|
||||
Name string `json:"name"`
|
||||
OK bool `json:"ok"`
|
||||
Detail string `json:"detail,omitempty"`
|
||||
Latency int64 `json:"latency_ms,omitempty"`
|
||||
}
|
||||
|
||||
// CheckReport is the full preflight result of `worker-agent --check` and of
|
||||
// the wizard's test step.
|
||||
type CheckReport struct {
|
||||
Coordinator CheckItem `json:"coordinator"`
|
||||
Auth CheckItem `json:"auth"`
|
||||
Python CheckItem `json:"python"`
|
||||
Scimesh CheckItem `json:"scimesh"`
|
||||
Agent string `json:"agent_version"`
|
||||
CoordinatorVersion string `json:"coordinator_version,omitempty"`
|
||||
}
|
||||
|
||||
// checkHTTP runs one GET and reports reachability + latency, with a fallback
|
||||
// detail message when the server answers without JSON.
|
||||
func checkHTTP(ctx context.Context, url string, timeout time.Duration) (CheckItem, string) {
|
||||
started := time.Now()
|
||||
reqCtx, cancel := context.WithTimeout(ctx, timeout)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(reqCtx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return CheckItem{Name: "coordinator", OK: false, Detail: "invalid URL"}, ""
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
detail := err.Error()
|
||||
if strings.Contains(detail, "connection refused") {
|
||||
detail = "no coordinator answering at this address"
|
||||
}
|
||||
return CheckItem{Name: "coordinator", OK: false, Detail: detail}, ""
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
version := ""
|
||||
if resp.StatusCode == http.StatusOK {
|
||||
var body struct {
|
||||
Status string `json:"status"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&body); err == nil && body.Status == "ok" {
|
||||
return CheckItem{Name: "coordinator", OK: true, Latency: time.Since(started).Milliseconds()}, version
|
||||
}
|
||||
}
|
||||
return CheckItem{Name: "coordinator", OK: false, Detail: fmt.Sprintf("HTTP %d", resp.StatusCode)}, version
|
||||
}
|
||||
|
||||
// CheckCoordinator probes the coordinator's /health endpoint.
|
||||
func CheckCoordinator(ctx context.Context, url string, timeout time.Duration) CheckReport {
|
||||
report := CheckReport{Agent: Version}
|
||||
item, _ := checkHTTP(ctx, strings.TrimRight(url, "/")+"/health", timeout)
|
||||
report.Coordinator = item
|
||||
report.Auth = CheckItem{Name: "auth", OK: true, Detail: "no token configured — will be checked at registration"}
|
||||
return report
|
||||
}
|
||||
|
||||
// CheckEnvironment verifies the local runtime against the python3 found on
|
||||
// PATH.
|
||||
func CheckEnvironment(ctx context.Context) CheckReport {
|
||||
python, err := exec.LookPath("python3")
|
||||
if err != nil {
|
||||
return CheckReport{Agent: Version, Python: CheckItem{Name: "python", OK: false, Detail: "python3 not found on PATH"}}
|
||||
}
|
||||
return CheckEnvironmentWithPython(ctx, python)
|
||||
}
|
||||
|
||||
// CheckEnvironmentWithPython verifies the local runtime against a specific
|
||||
// interpreter — the wizard's managed venv python when the runtime installer
|
||||
// has created one, so the preflight reflects what the worker will actually
|
||||
// execute with.
|
||||
func CheckEnvironmentWithPython(ctx context.Context, python string) CheckReport {
|
||||
report := CheckReport{Agent: Version, Python: CheckItem{Name: "python", OK: true, Detail: python}}
|
||||
//nolint:gosec // G204: python is a resolved interpreter path, the argument list is constant
|
||||
cmd := exec.CommandContext(ctx, python, "-c", "import scimesh; print(scimesh.__version__ if hasattr(scimesh, '__version__') else 'installed')")
|
||||
out, err := cmd.Output()
|
||||
if err != nil {
|
||||
// The worker executes workloads by spawning scimesh's task runner, so
|
||||
// the package is a hard requirement, not an optimisation. The PyPI
|
||||
// name belongs to a different project, so the wizard installs from
|
||||
// SCIMESH_PIP_PACKAGE instead of suggesting a bare pip install.
|
||||
report.Scimesh = CheckItem{Name: "scimesh", OK: false, Detail: "the worker runs workloads through scimesh — install it from your wheel or index (SCIMESH_PIP_PACKAGE)"}
|
||||
return report
|
||||
}
|
||||
report.Scimesh = CheckItem{Name: "scimesh", OK: true, Detail: strings.TrimSpace(string(out))}
|
||||
return report
|
||||
}
|
||||
|
||||
// RunCheck combines the coordinator probe and the local environment probe; it
|
||||
// is the body behind `worker-agent --check` and the wizard's test step. A
|
||||
// non-empty python overrides the interpreter probed for the scimesh package
|
||||
// (the managed venv after a runtime install).
|
||||
// RunCheck combines the coordinator probe, a credential probe and the local
|
||||
// environment probe; it is the body behind `worker-agent --check` and the
|
||||
// wizard's test step. A non-empty python overrides the interpreter probed for
|
||||
// the scimesh package (the managed venv after a runtime install).
|
||||
func RunCheck(ctx context.Context, coordinatorURL, python, token, workerKey, userserviceURL string) CheckReport {
|
||||
report := CheckCoordinator(ctx, coordinatorURL, 15*time.Second)
|
||||
report.Auth = CheckAuth(ctx, coordinatorURL, token, workerKey, userserviceURL)
|
||||
var env CheckReport
|
||||
if python != "" {
|
||||
env = CheckEnvironmentWithPython(ctx, python)
|
||||
} else {
|
||||
env = CheckEnvironment(ctx)
|
||||
}
|
||||
report.Python = env.Python
|
||||
report.Scimesh = env.Scimesh
|
||||
return report
|
||||
}
|
||||
|
||||
// Version is the agent build version; main injects it via -ldflags and the
|
||||
// setup wizard mirrors it into the report. "dev" marks a local build.
|
||||
var Version = "dev"
|
||||
|
||||
// Platform is the host platform string shown on the wizard.
|
||||
func Platform() string { return runtime.GOOS + "/" + runtime.GOARCH }
|
||||
|
||||
// CheckAuth verifies the configured credential against the coordinator
|
||||
// without mutating anything: with a worker key it first exchanges it at the
|
||||
// userservice for a short-lived JWT, then it probes /tasks/claim with a
|
||||
// throwaway worker id and no capabilities. A 401 anywhere means the
|
||||
// credential was rejected; any other status proves it was accepted.
|
||||
func CheckAuth(ctx context.Context, url, token, workerKey, userserviceURL string) CheckItem {
|
||||
item := CheckItem{Name: "auth"}
|
||||
if token == "" && workerKey == "" {
|
||||
item.OK = true
|
||||
item.Detail = "no credential configured — will be checked at registration"
|
||||
return item
|
||||
}
|
||||
client := &http.Client{Timeout: 30 * time.Second}
|
||||
if workerKey != "" && userserviceURL != "" {
|
||||
payload, _ := json.Marshal(map[string]string{"key": workerKey})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(userserviceURL, "/")+"/worker-tokens/exchange", strings.NewReader(string(payload)))
|
||||
if err != nil {
|
||||
item.OK = false
|
||||
item.Detail = "invalid userservice URL"
|
||||
return item
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
item.OK = false
|
||||
item.Detail = "userservice unreachable: " + err.Error()
|
||||
return item
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
item.OK = false
|
||||
item.Detail = "worker key rejected by the userservice"
|
||||
return item
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
item.OK = false
|
||||
item.Detail = fmt.Sprintf("userservice exchange: HTTP %d", resp.StatusCode)
|
||||
return item
|
||||
}
|
||||
var exchanged struct {
|
||||
Token string `json:"token"`
|
||||
}
|
||||
if err := json.NewDecoder(resp.Body).Decode(&exchanged); err != nil || exchanged.Token == "" {
|
||||
item.OK = false
|
||||
item.Detail = "userservice exchange returned no token"
|
||||
return item
|
||||
}
|
||||
token = exchanged.Token
|
||||
}
|
||||
if token == "" {
|
||||
item.OK = false
|
||||
item.Detail = "no usable credential after the key exchange"
|
||||
return item
|
||||
}
|
||||
payload, _ := json.Marshal(map[string]any{"worker_id": guuid.NewString(), "capabilities": []string{}})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(url, "/")+"/tasks/claim", strings.NewReader(string(payload)))
|
||||
if err != nil {
|
||||
item.OK = false
|
||||
item.Detail = "invalid coordinator URL"
|
||||
return item
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
item.OK = false
|
||||
item.Detail = "coordinator unreachable: " + err.Error()
|
||||
return item
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode == http.StatusUnauthorized {
|
||||
item.OK = false
|
||||
item.Detail = "token rejected by the coordinator"
|
||||
return item
|
||||
}
|
||||
item.OK = true
|
||||
item.Detail = "credential accepted"
|
||||
return item
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCheckAuthAcceptsToken(t *testing.T) {
|
||||
stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/tasks/claim" && r.Header.Get("Authorization") == "Bearer good-token" {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer stub.Close()
|
||||
|
||||
item := CheckAuth(context.Background(), stub.URL, "good-token", "", "")
|
||||
if !item.OK {
|
||||
t.Errorf("good token: %+v", item)
|
||||
}
|
||||
item = CheckAuth(context.Background(), stub.URL, "bad-token", "", "")
|
||||
if item.OK {
|
||||
t.Error("bad token must fail")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckAuthNoCredentialIsNotAFailure(t *testing.T) {
|
||||
item := CheckAuth(context.Background(), "http://coord:8080", "", "", "")
|
||||
if !item.OK {
|
||||
t.Errorf("no credential must not fail the preflight: %+v", item)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckAuthWorkerKeyExchange(t *testing.T) {
|
||||
var exchanged bool
|
||||
users := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/worker-tokens/exchange" {
|
||||
t.Errorf("unexpected userservice path %q", r.URL.Path)
|
||||
}
|
||||
exchanged = true
|
||||
_, _ = w.Write([]byte(`{"token":"jwt-after-exchange"}`))
|
||||
}))
|
||||
defer users.Close()
|
||||
coord := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Header.Get("Authorization") == "Bearer jwt-after-exchange" {
|
||||
w.WriteHeader(http.StatusNoContent)
|
||||
return
|
||||
}
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer coord.Close()
|
||||
|
||||
item := CheckAuth(context.Background(), coord.URL, "", "smk_key", users.URL)
|
||||
if !item.OK || !exchanged {
|
||||
t.Errorf("key exchange flow: %+v exchanged=%v", item, exchanged)
|
||||
}
|
||||
|
||||
rejected := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
}))
|
||||
defer rejected.Close()
|
||||
if item := CheckAuth(context.Background(), coord.URL, "", "smk_bad", rejected.URL); item.OK {
|
||||
t.Error("rejected key must fail the preflight")
|
||||
}
|
||||
}
|
||||
@@ -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,194 @@
|
||||
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
|
||||
// Concurrency is how many claim→execute→upload loops run in parallel
|
||||
// under one worker id: N shards processed concurrently on one machine,
|
||||
// using the coordinator's own task pipeline as the parallel unit.
|
||||
Concurrency int
|
||||
}
|
||||
|
||||
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",
|
||||
Concurrency: envInt("WORKER_CONCURRENCY", 1),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func envInt(name string, fallback int) int {
|
||||
raw := os.Getenv(name)
|
||||
if raw == "" {
|
||||
return fallback
|
||||
}
|
||||
parsed, err := strconv.Atoi(raw)
|
||||
if err != nil || parsed < 1 {
|
||||
return fallback
|
||||
}
|
||||
return parsed
|
||||
}
|
||||
|
||||
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,156 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ConfigFile is the persisted worker configuration written by the setup
|
||||
// wizard and read back by `worker-agent --config`. Environment variables
|
||||
// still win: the file fills in what the environment left unset.
|
||||
type ConfigFile struct {
|
||||
CoordinatorURL string `json:"coordinator_url"`
|
||||
Token string `json:"token,omitempty"`
|
||||
WorkerKey string `json:"worker_key,omitempty"`
|
||||
UserserviceURL string `json:"userservice_url,omitempty"`
|
||||
WorkDir string `json:"work_dir"`
|
||||
WorkerName string `json:"worker_name,omitempty"`
|
||||
CPUCount int `json:"cpu_count"`
|
||||
MemoryMB int `json:"memory_mb"`
|
||||
Concurrency int `json:"concurrency,omitempty"`
|
||||
TaskRunner []string `json:"task_runner,omitempty"`
|
||||
}
|
||||
|
||||
// DefaultConfigPath is where the setup wizard stores the worker's
|
||||
// configuration. SCIMESH_WORKER_CONFIG overrides it.
|
||||
func DefaultConfigPath() string {
|
||||
if raw := os.Getenv("SCIMESH_WORKER_CONFIG"); raw != "" {
|
||||
return raw
|
||||
}
|
||||
home, err := os.UserHomeDir()
|
||||
if err != nil || home == "" {
|
||||
return filepath.Join(".", ".scimesh-worker", "config.json")
|
||||
}
|
||||
return filepath.Join(home, ".scimesh-worker", "config.json")
|
||||
}
|
||||
|
||||
// LoadConfigFile reads and validates a persisted configuration. The file is
|
||||
// created by the wizard with 0600 permissions, so no credential is exposed to
|
||||
// other local users.
|
||||
func LoadConfigFile(path string) (*Config, error) {
|
||||
//nolint:gosec // G304: path is --config or SCIMESH_WORKER_CONFIG, operator-supplied
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read config file: %w", err)
|
||||
}
|
||||
var file ConfigFile
|
||||
if err := json.Unmarshal(raw, &file); err != nil {
|
||||
return nil, fmt.Errorf("parse config file: %w", err)
|
||||
}
|
||||
return file.Config()
|
||||
}
|
||||
|
||||
// Config turns the file into the daemon configuration. Environment variables
|
||||
// take precedence so operators can still override any value per-process.
|
||||
func (f *ConfigFile) Config() (*Config, error) {
|
||||
config := &Config{}
|
||||
if env := os.Getenv("COORDINATOR_URL"); env != "" {
|
||||
config.CoordinatorURL = env
|
||||
} else {
|
||||
config.CoordinatorURL = strings.TrimSpace(f.CoordinatorURL)
|
||||
}
|
||||
if config.CoordinatorURL == "" {
|
||||
return nil, fmt.Errorf("coordinator_url is required")
|
||||
}
|
||||
if !strings.HasPrefix(config.CoordinatorURL, "http://") && !strings.HasPrefix(config.CoordinatorURL, "https://") {
|
||||
return nil, fmt.Errorf("coordinator_url must be an absolute HTTP(S) URL")
|
||||
}
|
||||
if env := os.Getenv("WORKER_AUTH_TOKEN"); env != "" {
|
||||
config.Token = env
|
||||
} else {
|
||||
config.Token = f.Token
|
||||
}
|
||||
if env := os.Getenv("WORKER_KEY"); env != "" {
|
||||
config.WorkerKey = env
|
||||
} else {
|
||||
config.WorkerKey = f.WorkerKey
|
||||
}
|
||||
if env := os.Getenv("USERSERVICE_URL"); env != "" {
|
||||
config.UserserviceURL = env
|
||||
} else {
|
||||
config.UserserviceURL = f.UserserviceURL
|
||||
}
|
||||
if env := os.Getenv("WORK_DIR"); env != "" {
|
||||
config.WorkDir = env
|
||||
} else if f.WorkDir != "" {
|
||||
config.WorkDir = f.WorkDir
|
||||
} else {
|
||||
config.WorkDir = "./scimesh-agent-data"
|
||||
}
|
||||
if env := os.Getenv("WORKER_NAME"); env != "" {
|
||||
config.WorkerName = env
|
||||
} else {
|
||||
config.WorkerName = f.WorkerName
|
||||
}
|
||||
config.CPUCount = f.CPUCount
|
||||
if config.CPUCount < 1 {
|
||||
config.CPUCount = 1
|
||||
}
|
||||
config.MemoryMB = f.MemoryMB
|
||||
if config.MemoryMB < 0 {
|
||||
config.MemoryMB = 0
|
||||
}
|
||||
config.Concurrency = f.Concurrency
|
||||
if config.Concurrency < 1 {
|
||||
config.Concurrency = 1
|
||||
}
|
||||
if len(f.TaskRunner) > 0 {
|
||||
config.TaskRunner = f.TaskRunner
|
||||
}
|
||||
if len(config.TaskRunner) == 0 {
|
||||
config.TaskRunner = []string{"python", "-m", "scimesh.worker.task"}
|
||||
}
|
||||
config.PollInterval = 2 * time.Second
|
||||
config.RequestTimeout = 30 * time.Second
|
||||
config.Heartbeat = 15 * time.Second
|
||||
config.Capabilities = DefaultCapabilities()
|
||||
return config, nil
|
||||
}
|
||||
|
||||
// Save writes the configuration file, creating the parent directory and
|
||||
// restricting permissions to the owner.
|
||||
func SaveConfigFile(path string, file ConfigFile) error {
|
||||
payload, err := json.MarshalIndent(file, "", " ")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload = append(payload, '\n')
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
|
||||
return fmt.Errorf("create config directory: %w", err)
|
||||
}
|
||||
if err := os.WriteFile(path, payload, 0o600); err != nil {
|
||||
return fmt.Errorf("write config file: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// VenvPython returns the managed venv python next to the given config file,
|
||||
// or "" when the runtime installer has not created one yet. Mirrors the
|
||||
// wizard's probe so `worker-agent --check` and the preflight agree on what
|
||||
// interpreter will actually execute workloads.
|
||||
func VenvPython(configPath string) string {
|
||||
dir := filepath.Dir(configPath)
|
||||
for _, candidate := range []string{
|
||||
filepath.Join(dir, "venv", "bin", "python"),
|
||||
filepath.Join(dir, "venv", "Scripts", "python.exe"),
|
||||
} {
|
||||
if info, err := os.Stat(candidate); err == nil && !info.IsDir() {
|
||||
return candidate
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
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 registers once, then runs the claim→execute→upload loop
|
||||
// concurrently under one worker id. With Concurrency > 1, several shards are
|
||||
// processed in parallel on this machine, using the coordinator's own task
|
||||
// pipeline as the parallel unit.
|
||||
func (d *Daemon) RunForever() error {
|
||||
if !d.registered {
|
||||
if err := d.register(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
workers := d.config.Concurrency
|
||||
if workers < 1 {
|
||||
workers = 1
|
||||
}
|
||||
if workers == 1 {
|
||||
return d.loop()
|
||||
}
|
||||
d.log.Info("agent running concurrently", "loops", workers)
|
||||
var wg sync.WaitGroup
|
||||
errs := make(chan error, workers)
|
||||
for i := 0; i < workers; i++ {
|
||||
wg.Add(1)
|
||||
go func(loop int) {
|
||||
defer wg.Done()
|
||||
if err := d.loop(); err != nil {
|
||||
errs <- err
|
||||
return
|
||||
}
|
||||
errs <- nil
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
close(errs)
|
||||
for err := range errs {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// loop is one claim→execute→upload cycle until interrupted, idle-exit, or
|
||||
// the shared max-tasks budget is consumed.
|
||||
func (d *Daemon) loop() error {
|
||||
failures := 0
|
||||
for {
|
||||
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.mu.Lock()
|
||||
d.completed++
|
||||
done := d.config.MaxTasks > 0 && d.completed >= d.config.MaxTasks
|
||||
d.mu.Unlock()
|
||||
if done {
|
||||
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
|
||||
}
|
||||
d.log.Info("task claimed", "task_id", task.TaskID, "attempt", task.Attempt)
|
||||
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,375 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"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
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
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) {
|
||||
fake.mu.Lock()
|
||||
defer fake.mu.Unlock()
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDaemonConcurrencyProcessesTasksInParallel(t *testing.T) {
|
||||
t.Parallel()
|
||||
marker := filepath.Join(t.TempDir(), "marker")
|
||||
script := filepath.Join(t.TempDir(), "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
|
||||
echo start >> ` + marker + `
|
||||
sleep 1
|
||||
echo end >> ` + marker + `
|
||||
printf 'id,score\n1,1\n' > "$task_dir/result.csv"
|
||||
printf '{"artifact_path":"%s/result.csv","content_type":"text/csv","metrics":{"rows":1}}' "$task_dir" > "$out"
|
||||
exit 0
|
||||
`
|
||||
if err := os.WriteFile(script, []byte(content), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
fake := newFakeCoordinator(t, validClaimedTaskPayload())
|
||||
defer fake.close()
|
||||
config := &Config{
|
||||
CoordinatorURL: fake.server.URL,
|
||||
WorkerName: "concurrent-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},
|
||||
MaxTasks: 3,
|
||||
Concurrency: 3,
|
||||
}
|
||||
client := NewClient(fake.server.URL, &StaticToken{token: "test-token"}, 5*time.Second)
|
||||
logger := slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
daemon := NewDaemon(config, client, NewTaskRunner(config.TaskRunner), logger)
|
||||
if err := daemon.RunForever(); err != nil {
|
||||
t.Fatalf("run: %v", err)
|
||||
}
|
||||
raw, err := os.ReadFile(marker)
|
||||
if err != nil {
|
||||
t.Fatalf("marker: %v", err)
|
||||
}
|
||||
starts := strings.Count(string(raw), "start\n")
|
||||
ends := strings.Count(string(raw), "end\n")
|
||||
if starts < 3 || ends < 3 {
|
||||
t.Fatalf("marker: %d starts / %d ends, want at least 3/3", starts, ends)
|
||||
}
|
||||
// With three loops sleeping 1s each, the marker proves all three ran
|
||||
// concurrently (three starts before the first end completes a 1s sleep).
|
||||
lines := strings.Split(strings.TrimSpace(string(raw)), "\n")
|
||||
concurrent := 0
|
||||
running := 0
|
||||
for _, line := range lines {
|
||||
if line == "start" {
|
||||
running++
|
||||
if running > concurrent {
|
||||
concurrent = running
|
||||
}
|
||||
} else {
|
||||
running--
|
||||
}
|
||||
}
|
||||
if concurrent < 3 {
|
||||
t.Errorf("max concurrent executions = %d, want 3", concurrent)
|
||||
}
|
||||
}
|
||||
@@ -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,73 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ReleaseWheelURL returns the download URL of the scimesh wheel attached to
|
||||
// the GitHub release that matches the given binary version (for example
|
||||
// "1.1.0-alpha.10"), plus the wheel file name. The wheel is version-locked to
|
||||
// the binary so a worker's catalog always matches its task runner.
|
||||
func ReleaseWheelURL(version string) (string, string, error) {
|
||||
if version == "" || version == "dev" {
|
||||
return "", "", fmt.Errorf("no release wheel for build %q", version)
|
||||
}
|
||||
filename := fmt.Sprintf("scimesh-%s-py3-none-any.whl", NormalizePEP440(version))
|
||||
return fmt.Sprintf("https://github.com/emil28092005/SciMesh/releases/download/v%s/%s", version, filename), filename, nil
|
||||
}
|
||||
|
||||
// NormalizePEP440 turns our release tag suffixes into the PEP 440 form
|
||||
// setuptools uses for wheel names: 1.1.0-alpha.10 -> 1.1.0a10,
|
||||
// 1.1.0-beta.2 -> 1.1.0b2, 1.1.0-rc.1 -> 1.1.0rc1. Stable tags pass through.
|
||||
func NormalizePEP440(version string) string {
|
||||
for from, to := range map[string]string{"-alpha.": "a", "-beta.": "b", "-rc.": "rc"} {
|
||||
version = strings.ReplaceAll(version, from, to)
|
||||
}
|
||||
return version
|
||||
}
|
||||
|
||||
// DownloadWheel fetches the release wheel into dir (config directory of the
|
||||
// wizard / serve data dir) and returns the local path. Best-effort download
|
||||
// with a generous timeout: wheels can be several MB.
|
||||
func DownloadWheel(ctx context.Context, url, dir string) (string, error) {
|
||||
target := filepath.Join(dir, wheelNameFromURL(url))
|
||||
if err := os.MkdirAll(dir, 0o700); err != nil {
|
||||
return "", fmt.Errorf("download wheel: %w", err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, 10*time.Minute)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("download wheel: %w", err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("download wheel: HTTP %d", resp.StatusCode)
|
||||
}
|
||||
//nolint:gosec // G304: target is our own config dir + a fixed wheel name
|
||||
out, err := os.Create(target)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("download wheel: %w", err)
|
||||
}
|
||||
defer func() { _ = out.Close() }()
|
||||
if _, err := io.Copy(out, resp.Body); err != nil {
|
||||
return "", fmt.Errorf("download wheel: %w", err)
|
||||
}
|
||||
return target, nil
|
||||
}
|
||||
|
||||
// wheelNameFromURL extracts the trailing file name of a wheel URL.
|
||||
func wheelNameFromURL(url string) string {
|
||||
return url[strings.LastIndex(url, "/")+1:]
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNormalizePEP440(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"1.1.0": "1.1.0",
|
||||
"1.1.0-alpha.10": "1.1.0a10",
|
||||
"1.1.0-beta.2": "1.1.0b2",
|
||||
"1.1.0-rc.1": "1.1.0rc1",
|
||||
"1.0.0": "1.0.0",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := NormalizePEP440(in); got != want {
|
||||
t.Errorf("NormalizePEP440(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestReleaseWheelURL(t *testing.T) {
|
||||
url, name, err := ReleaseWheelURL("1.1.0-alpha.10")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
wantURL := "https://github.com/emil28092005/SciMesh/releases/download/v1.1.0-alpha.10/scimesh-1.1.0a10-py3-none-any.whl"
|
||||
if url != wantURL {
|
||||
t.Errorf("url = %q, want %q", url, wantURL)
|
||||
}
|
||||
if name != "scimesh-1.1.0a10-py3-none-any.whl" {
|
||||
t.Errorf("name = %q", name)
|
||||
}
|
||||
|
||||
// A dev build has no release wheel.
|
||||
if _, _, err := ReleaseWheelURL("dev"); err == nil {
|
||||
t.Error("dev build must not resolve a release wheel")
|
||||
}
|
||||
if _, _, err := ReleaseWheelURL(""); err == nil {
|
||||
t.Error("empty version must not resolve a release wheel")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadWheel(t *testing.T) {
|
||||
payload := []byte("fake wheel bytes")
|
||||
stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write(payload)
|
||||
}))
|
||||
defer stub.Close()
|
||||
|
||||
dir := t.TempDir()
|
||||
path, err := DownloadWheel(context.Background(), stub.URL+"/scimesh-1.1.0a10-py3-none-any.whl", dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.HasSuffix(path, "scimesh-1.1.0a10-py3-none-any.whl") {
|
||||
t.Errorf("path = %q", path)
|
||||
}
|
||||
got, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(got) != string(payload) {
|
||||
t.Error("wheel bytes mismatch")
|
||||
}
|
||||
}
|
||||
|
||||
func TestDownloadWheelReportsHTTPErrors(t *testing.T) {
|
||||
stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer stub.Close()
|
||||
if _, err := DownloadWheel(context.Background(), stub.URL+"/missing.whl", t.TempDir()); err == nil {
|
||||
t.Error("404 must fail the download")
|
||||
}
|
||||
}
|
||||
@@ -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,630 @@
|
||||
// Package setupui serves the local worker setup wizard: a small HTTP server
|
||||
// bound to 127.0.0.1 that writes the worker's config file, runs preflight
|
||||
// checks, and starts/stops the worker as a background process. It is part of
|
||||
// the worker-agent binary so a machine that installs only a worker never needs
|
||||
// a coordinator.
|
||||
package setupui
|
||||
|
||||
import (
|
||||
"context"
|
||||
"embed"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/agent"
|
||||
)
|
||||
|
||||
//go:embed template.html
|
||||
var templateFS embed.FS
|
||||
|
||||
const (
|
||||
defaultPort = 12700
|
||||
pidFileName = "worker.pid"
|
||||
logFileName = "worker.log"
|
||||
)
|
||||
|
||||
// Supervisor starts and stops the worker process and tracks its pid. It is an
|
||||
// interface so tests can substitute a fake.
|
||||
type Supervisor interface {
|
||||
// Start launches `worker-agent --config <path>` detached, writing output
|
||||
// into the log file. Returns the child pid.
|
||||
Start(configPath, logPath string) (int, error)
|
||||
// Stop terminates the process recorded in the pid file.
|
||||
Stop() error
|
||||
// Pid returns the recorded child pid, or 0 when none is recorded.
|
||||
Pid() int
|
||||
// Alive reports whether the recorded child is still running.
|
||||
Alive() bool
|
||||
}
|
||||
|
||||
// PIDSupervisor is the real Supervisor: it spawns the running binary with
|
||||
// --config and manages its pid file. Liveness comes from a Wait goroutine, so
|
||||
// it works on every platform (no signal probing, which Windows lacks).
|
||||
type PIDSupervisor struct {
|
||||
mu sync.Mutex
|
||||
pidPath string
|
||||
proc *os.Process
|
||||
done chan struct{} // closed when the spawned process exits; nil when not started
|
||||
}
|
||||
|
||||
func NewPIDSupervisor(pidPath string) *PIDSupervisor { return &PIDSupervisor{pidPath: pidPath} }
|
||||
|
||||
func (s *PIDSupervisor) Pid() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.readPid()
|
||||
}
|
||||
|
||||
func (s *PIDSupervisor) readPid() int {
|
||||
raw, err := os.ReadFile(s.pidPath)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
pid, err := strconv.Atoi(strings.TrimSpace(string(raw)))
|
||||
if err != nil || pid < 1 {
|
||||
return 0
|
||||
}
|
||||
return pid
|
||||
}
|
||||
|
||||
func (s *PIDSupervisor) Alive() bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.done == nil {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case <-s.done:
|
||||
return false
|
||||
default:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
func (s *PIDSupervisor) Start(configPath, logPath string) (int, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.done != nil {
|
||||
select {
|
||||
case <-s.done:
|
||||
default:
|
||||
return s.readPid(), fmt.Errorf("worker is already running (pid %d)", s.readPid())
|
||||
}
|
||||
}
|
||||
exe, err := os.Executable()
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("resolve worker binary: %w", err)
|
||||
}
|
||||
//nolint:gosec // G304: logPath lives in the wizard's own config directory
|
||||
logFile, err := os.OpenFile(logPath, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0o600)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("open worker log: %w", err)
|
||||
}
|
||||
defer func() { _ = logFile.Close() }()
|
||||
//nolint:gosec // G204: exe is os.Executable, configPath is the wizard's own file;
|
||||
// Background ctx: the child's lifecycle is managed by the supervisor, not the context
|
||||
cmd := exec.CommandContext(context.Background(), exe, "--config", configPath)
|
||||
cmd.Stdout = logFile
|
||||
cmd.Stderr = logFile
|
||||
cmd.Stdin = nil
|
||||
if err := cmd.Start(); err != nil {
|
||||
return 0, fmt.Errorf("start worker: %w", err)
|
||||
}
|
||||
// The child inherits our stdout/stderr descriptors pointing at the log
|
||||
// file, so we can close our copy; the child keeps it open.
|
||||
_ = logFile.Close()
|
||||
s.proc = cmd.Process
|
||||
s.done = make(chan struct{})
|
||||
go func() { _ = cmd.Wait(); close(s.done) }()
|
||||
if err := os.WriteFile(s.pidPath, []byte(strconv.Itoa(cmd.Process.Pid)+"\n"), 0o600); err != nil {
|
||||
_ = cmd.Process.Kill()
|
||||
return 0, fmt.Errorf("write pid file: %w", err)
|
||||
}
|
||||
return cmd.Process.Pid, nil
|
||||
}
|
||||
|
||||
func (s *PIDSupervisor) Stop() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.done == nil {
|
||||
return nil
|
||||
}
|
||||
select {
|
||||
case <-s.done:
|
||||
s.done = nil
|
||||
s.proc = nil
|
||||
_ = os.Remove(s.pidPath)
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
// Ask politely, then force. os.Interrupt terminates on Windows too.
|
||||
_ = s.proc.Signal(os.Interrupt)
|
||||
deadline := time.Now().Add(3 * time.Second)
|
||||
for time.Now().Before(deadline) {
|
||||
select {
|
||||
case <-s.done:
|
||||
s.done = nil
|
||||
s.proc = nil
|
||||
_ = os.Remove(s.pidPath)
|
||||
return nil
|
||||
default:
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
_ = s.proc.Kill()
|
||||
select {
|
||||
case <-s.done:
|
||||
case <-time.After(2 * time.Second):
|
||||
}
|
||||
s.done = nil
|
||||
s.proc = nil
|
||||
_ = os.Remove(s.pidPath)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Server is the wizard HTTP server, bound to the loopback interface only.
|
||||
type Server struct {
|
||||
log *slog.Logger
|
||||
cfgPath string
|
||||
logPath string
|
||||
dir string
|
||||
sup Supervisor
|
||||
openBrowser func(string)
|
||||
port int
|
||||
install func(ctx context.Context, venvPython, pkg string) error
|
||||
downloadWheel func(ctx context.Context, url, dir string) (string, error)
|
||||
}
|
||||
|
||||
// Options customises the wizard for tests and embedding.
|
||||
type Options struct {
|
||||
Port int
|
||||
ConfigPath string
|
||||
OpenBrowser func(url string)
|
||||
Supervisor Supervisor
|
||||
Dir string // directory for pid/log files; defaults to the config dir
|
||||
// InstallScimesh overrides the pip step of the runtime installer (tests
|
||||
// substitute a fake); nil uses the real pip inside the managed venv.
|
||||
InstallScimesh func(ctx context.Context, venvPython, pkg string) error
|
||||
// DownloadWheel overrides the release-wheel download (tests substitute a
|
||||
// fake); nil downloads from the GitHub release matching the agent version.
|
||||
DownloadWheel func(ctx context.Context, url, dir string) (string, error)
|
||||
}
|
||||
|
||||
func New(log *slog.Logger, opts Options) *Server {
|
||||
cfgPath := opts.ConfigPath
|
||||
if cfgPath == "" {
|
||||
cfgPath = agent.DefaultConfigPath()
|
||||
}
|
||||
dir := opts.Dir
|
||||
if dir == "" {
|
||||
dir = filepath.Dir(cfgPath)
|
||||
}
|
||||
sup := opts.Supervisor
|
||||
if sup == nil {
|
||||
sup = NewPIDSupervisor(filepath.Join(dir, pidFileName))
|
||||
}
|
||||
open := opts.OpenBrowser
|
||||
if open == nil {
|
||||
open = func(string) {}
|
||||
}
|
||||
port := opts.Port
|
||||
if port == 0 {
|
||||
port = defaultPort
|
||||
}
|
||||
install := opts.InstallScimesh
|
||||
if install == nil {
|
||||
install = installScimeshWithPip
|
||||
}
|
||||
downloadWheel := opts.DownloadWheel
|
||||
if downloadWheel == nil {
|
||||
downloadWheel = agent.DownloadWheel
|
||||
}
|
||||
return &Server{log: log, cfgPath: cfgPath, logPath: filepath.Join(dir, logFileName), dir: dir, sup: sup, openBrowser: open, port: port, install: install, downloadWheel: downloadWheel}
|
||||
}
|
||||
|
||||
// Listen binds the loopback listener and returns it; Serve runs the server on
|
||||
// it. Split so tests can inspect the actual ephemeral port.
|
||||
func (s *Server) Listen() (net.Listener, error) {
|
||||
return (&net.ListenConfig{}).Listen(context.Background(), "tcp", fmt.Sprintf("127.0.0.1:%d", s.port))
|
||||
}
|
||||
|
||||
// OpenBrowser hands the wizard URL to the configured opener (default: no-op).
|
||||
func (s *Server) OpenBrowser(url string) { s.openBrowser(url) }
|
||||
|
||||
// Serve runs the wizard until ctx is cancelled.
|
||||
func (s *Server) Serve(ctx context.Context, listener net.Listener) error {
|
||||
mux := http.NewServeMux()
|
||||
mux.HandleFunc("GET /", s.handleIndex)
|
||||
mux.HandleFunc("GET /api/status", s.handleStatus)
|
||||
mux.HandleFunc("POST /api/config", s.handleSaveConfig)
|
||||
mux.HandleFunc("POST /api/test", s.handleTest)
|
||||
mux.HandleFunc("POST /api/runtime/install", s.handleInstallRuntime)
|
||||
mux.HandleFunc("POST /api/start", s.handleStart)
|
||||
mux.HandleFunc("POST /api/stop", s.handleStop)
|
||||
mux.HandleFunc("GET /api/logs", s.handleLogs)
|
||||
server := &http.Server{
|
||||
Handler: mux,
|
||||
ReadHeaderTimeout: 10 * time.Second,
|
||||
}
|
||||
go func() {
|
||||
<-ctx.Done()
|
||||
shutdownCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
|
||||
defer cancel()
|
||||
_ = server.Shutdown(shutdownCtx)
|
||||
}()
|
||||
return server.Serve(listener)
|
||||
}
|
||||
|
||||
func (s *Server) handleIndex(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/" {
|
||||
http.NotFound(w, r)
|
||||
return
|
||||
}
|
||||
html, err := templateFS.ReadFile("template.html")
|
||||
if err != nil {
|
||||
http.Error(w, "template unavailable", http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/html; charset=utf-8")
|
||||
_, _ = w.Write(html)
|
||||
}
|
||||
|
||||
func writeJSON(w http.ResponseWriter, status int, v any) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(status)
|
||||
_ = json.NewEncoder(w).Encode(v)
|
||||
}
|
||||
|
||||
// statusView is what the wizard needs to paint the running/stopped state.
|
||||
type statusView struct {
|
||||
ConfigPresent bool `json:"config_present"`
|
||||
ConfigPath string `json:"config_path"`
|
||||
LogPath string `json:"log_path"`
|
||||
Running bool `json:"running"`
|
||||
Pid int `json:"pid"`
|
||||
WorkerName string `json:"worker_name,omitempty"`
|
||||
Coordinator string `json:"coordinator,omitempty"`
|
||||
WorkDir string `json:"work_dir,omitempty"`
|
||||
TokenSet bool `json:"token_set"`
|
||||
Stats WorkerStats `json:"stats"`
|
||||
}
|
||||
|
||||
// WorkerStats is parsed from the worker log: the agent reports each claim,
|
||||
// completion and failure as a structured line, so the wizard can show live
|
||||
// counters without any coordinator access.
|
||||
type WorkerStats struct {
|
||||
Registered bool `json:"registered"`
|
||||
Claimed int `json:"claimed"`
|
||||
Completed int `json:"completed"`
|
||||
Failed int `json:"failed"`
|
||||
}
|
||||
|
||||
// ensureVenvTaskRunner rewrites the saved config so its task runner uses the
|
||||
// managed venv python when one exists and the config does not already pin one.
|
||||
func (s *Server) ensureVenvTaskRunner() {
|
||||
raw, err := os.ReadFile(s.cfgPath)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
var file agent.ConfigFile
|
||||
if json.Unmarshal(raw, &file) != nil || len(file.TaskRunner) > 0 {
|
||||
return
|
||||
}
|
||||
if venv := s.venvPython(); venv != "" {
|
||||
file.TaskRunner = []string{venv, "-m", "scimesh.worker.task"}
|
||||
if payload, err := json.MarshalIndent(file, "", " "); err == nil {
|
||||
_ = os.WriteFile(s.cfgPath, append(payload, '\n'), 0o600)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// parseWorkerStats counts the structured agent events in the worker log.
|
||||
func parseWorkerStats(logPath string) WorkerStats {
|
||||
var stats WorkerStats
|
||||
//nolint:gosec // G304: logPath is the wizard's own log file next to the config
|
||||
raw, err := os.ReadFile(logPath)
|
||||
if err != nil {
|
||||
return stats
|
||||
}
|
||||
for _, line := range strings.Split(string(raw), "\n") {
|
||||
switch {
|
||||
case strings.Contains(line, "msg=registered"):
|
||||
stats.Registered = true
|
||||
case strings.Contains(line, `msg="task claimed"`):
|
||||
stats.Claimed++
|
||||
case strings.Contains(line, `msg="task completed"`):
|
||||
stats.Completed++
|
||||
case strings.Contains(line, `msg="task failed"`):
|
||||
stats.Failed++
|
||||
}
|
||||
}
|
||||
return stats
|
||||
}
|
||||
|
||||
func (s *Server) handleStatus(w http.ResponseWriter, r *http.Request) {
|
||||
view := statusView{ConfigPath: s.cfgPath, LogPath: s.logPath, Running: s.sup.Alive(), Pid: s.sup.Pid(), Stats: parseWorkerStats(s.logPath)}
|
||||
if raw, err := os.ReadFile(s.cfgPath); err == nil {
|
||||
var file agent.ConfigFile
|
||||
if json.Unmarshal(raw, &file) == nil {
|
||||
view.ConfigPresent = true
|
||||
view.WorkerName = file.WorkerName
|
||||
view.Coordinator = file.CoordinatorURL
|
||||
view.WorkDir = file.WorkDir
|
||||
view.TokenSet = file.Token != "" || file.WorkerKey != ""
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, view)
|
||||
}
|
||||
|
||||
type saveConfigRequest struct {
|
||||
CoordinatorURL string `json:"coordinator_url"`
|
||||
Token string `json:"token"`
|
||||
WorkerKey string `json:"worker_key"`
|
||||
UserserviceURL string `json:"userservice_url"`
|
||||
WorkDir string `json:"work_dir"`
|
||||
WorkerName string `json:"worker_name"`
|
||||
CPUCount int `json:"cpu_count"`
|
||||
MemoryMB int `json:"memory_mb"`
|
||||
Concurrency int `json:"concurrency"`
|
||||
TaskRunner []string `json:"task_runner"`
|
||||
}
|
||||
|
||||
func (s *Server) handleSaveConfig(w http.ResponseWriter, r *http.Request) {
|
||||
var req saveConfigRequest
|
||||
if err := json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&req); err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON body"})
|
||||
return
|
||||
}
|
||||
file := agent.ConfigFile{
|
||||
CoordinatorURL: strings.TrimSpace(req.CoordinatorURL),
|
||||
Token: req.Token,
|
||||
WorkerKey: req.WorkerKey,
|
||||
UserserviceURL: strings.TrimSpace(req.UserserviceURL),
|
||||
WorkDir: strings.TrimSpace(req.WorkDir),
|
||||
WorkerName: strings.TrimSpace(req.WorkerName),
|
||||
CPUCount: req.CPUCount,
|
||||
MemoryMB: req.MemoryMB,
|
||||
Concurrency: req.Concurrency,
|
||||
TaskRunner: req.TaskRunner,
|
||||
}
|
||||
if file.CoordinatorURL == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "coordinator_url is required"})
|
||||
return
|
||||
}
|
||||
if !strings.HasPrefix(file.CoordinatorURL, "http://") && !strings.HasPrefix(file.CoordinatorURL, "https://") {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "coordinator_url must be an absolute HTTP(S) URL"})
|
||||
return
|
||||
}
|
||||
if file.WorkDir == "" {
|
||||
file.WorkDir = "./scimesh-agent-data"
|
||||
}
|
||||
if file.WorkerName == "" {
|
||||
if host, err := os.Hostname(); err == nil {
|
||||
file.WorkerName = host
|
||||
} else {
|
||||
file.WorkerName = "worker"
|
||||
}
|
||||
}
|
||||
if file.CPUCount < 1 {
|
||||
file.CPUCount = 1
|
||||
}
|
||||
if file.Concurrency < 1 {
|
||||
file.Concurrency = 1
|
||||
}
|
||||
// The wizard UI bakes the venv python into the runner after an install;
|
||||
// an API-driven or scripted flow may not, so the server guarantees it:
|
||||
// workloads execute through scimesh's task runner, which lives in the venv.
|
||||
if len(file.TaskRunner) == 0 {
|
||||
if venv := s.venvPython(); venv != "" {
|
||||
file.TaskRunner = []string{venv, "-m", "scimesh.worker.task"}
|
||||
}
|
||||
}
|
||||
if err := agent.SaveConfigFile(s.cfgPath, file); err != nil {
|
||||
s.log.Error("save wizard config", "err", err)
|
||||
writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "could not write the config file"})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]bool{"saved": true})
|
||||
}
|
||||
|
||||
func (s *Server) handleTest(w http.ResponseWriter, r *http.Request) {
|
||||
var req saveConfigRequest
|
||||
if err := json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&req); err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON body"})
|
||||
return
|
||||
}
|
||||
url := strings.TrimSpace(req.CoordinatorURL)
|
||||
if url == "" {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "coordinator_url is required"})
|
||||
return
|
||||
}
|
||||
// After the runtime installer created the venv, probe that interpreter:
|
||||
// checking the bare system python3 would keep reporting scimesh as
|
||||
// missing even though the worker would run with the venv.
|
||||
report := agent.RunCheck(r.Context(), url, s.venvPython(), req.Token, req.WorkerKey, req.UserserviceURL)
|
||||
writeJSON(w, http.StatusOK, report)
|
||||
}
|
||||
|
||||
func (s *Server) handleStart(w http.ResponseWriter, r *http.Request) {
|
||||
if _, err := os.Stat(s.cfgPath); err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "no configuration saved yet"})
|
||||
return
|
||||
}
|
||||
s.ensureVenvTaskRunner()
|
||||
pid, err := s.sup.Start(s.cfgPath, s.logPath)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusConflict, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]int{"pid": pid})
|
||||
}
|
||||
|
||||
func (s *Server) handleStop(w http.ResponseWriter, r *http.Request) {
|
||||
if err := s.sup.Stop(); err != nil {
|
||||
writeJSON(w, http.StatusConflict, map[string]string{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]bool{"stopped": true})
|
||||
}
|
||||
|
||||
func (s *Server) handleLogs(w http.ResponseWriter, r *http.Request) {
|
||||
raw, err := os.ReadFile(s.logPath)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusOK, map[string]string{"log": ""})
|
||||
return
|
||||
}
|
||||
lines := strings.Split(strings.TrimRight(string(raw), "\n"), "\n")
|
||||
tail := 200
|
||||
if n, err := strconv.Atoi(r.URL.Query().Get("tail")); err == nil && n > 0 && n < 5000 {
|
||||
tail = n
|
||||
}
|
||||
if len(lines) > tail {
|
||||
lines = lines[len(lines)-tail:]
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]string{"log": strings.Join(lines, "\n")})
|
||||
}
|
||||
|
||||
// ErrCanceled mirrors context.Canceled for callers that treat a cancelled
|
||||
// wizard as a clean exit.
|
||||
var ErrCanceled = errors.New("setup wizard cancelled")
|
||||
|
||||
type installRuntimeRequest struct {
|
||||
// ScimeshPackage overrides where the scimesh wheel comes from: a local
|
||||
// wheel/index path or the PyPI name. Defaults to SCIMESH_PIP_PACKAGE, then
|
||||
// to the PyPI name.
|
||||
ScimeshPackage string `json:"scimesh_package"`
|
||||
}
|
||||
|
||||
type installRuntimeResponse struct {
|
||||
OK bool `json:"ok"`
|
||||
Python string `json:"python,omitempty"` // venv python to use as TASK_RUNNER[0]
|
||||
Installed bool `json:"installed"`
|
||||
}
|
||||
|
||||
// handleInstallRuntime creates a managed venv next to the worker config and
|
||||
// installs the scimesh package into it, so the machine needs no manual pip
|
||||
// step. The venv python path is returned for the wizard to bake into the
|
||||
// task runner.
|
||||
func (s *Server) handleInstallRuntime(w http.ResponseWriter, r *http.Request) {
|
||||
// The wizard may run the install before any config was saved, so the
|
||||
// worker directory (venv, wheel) may not exist yet.
|
||||
if err := os.MkdirAll(s.dir, 0o700); err != nil {
|
||||
writeJSON(w, http.StatusConflict, map[string]string{"error": "could not create the worker directory"})
|
||||
return
|
||||
}
|
||||
var req installRuntimeRequest
|
||||
if err := json.NewDecoder(io.LimitReader(r.Body, 1<<20)).Decode(&req); err != nil {
|
||||
writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON body"})
|
||||
return
|
||||
}
|
||||
pkg := strings.TrimSpace(req.ScimeshPackage)
|
||||
if pkg == "" {
|
||||
pkg = os.Getenv("SCIMESH_PIP_PACKAGE")
|
||||
}
|
||||
if pkg == "" {
|
||||
// No PyPI default on purpose: the PyPI name "scimesh" belongs to an
|
||||
// unrelated project. Instead we ship the wheel in our own GitHub
|
||||
// release, version-locked to this binary, and download it from there.
|
||||
url, _, err := agent.ReleaseWheelURL(agent.Version)
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusConflict, map[string]string{
|
||||
"error": "no scimesh source configured: set SCIMESH_PIP_PACKAGE to your wheel, checkout or index, then retry",
|
||||
})
|
||||
return
|
||||
}
|
||||
local, err := s.downloadWheel(r.Context(), url, s.dir)
|
||||
if err != nil {
|
||||
s.log.Error("download release wheel", "err", err, "url", url)
|
||||
writeJSON(w, http.StatusConflict, map[string]string{
|
||||
"error": "could not download the scimesh wheel for this release: " + err.Error() + ". Set SCIMESH_PIP_PACKAGE to your wheel, checkout or index and retry.",
|
||||
})
|
||||
return
|
||||
}
|
||||
pkg = local
|
||||
}
|
||||
|
||||
python3, err := exec.LookPath("python3")
|
||||
if err != nil {
|
||||
writeJSON(w, http.StatusConflict, map[string]string{"error": "python3 is not installed on this machine"})
|
||||
return
|
||||
}
|
||||
venvDir := filepath.Join(s.dir, "venv")
|
||||
venvPython := filepath.Join(venvDir, "bin", "python")
|
||||
if _, err := os.Stat(venvPython); err != nil {
|
||||
// Windows layout: Scripts/python.exe.
|
||||
if win := filepath.Join(venvDir, "Scripts", "python.exe"); stat(win) {
|
||||
venvPython = win
|
||||
}
|
||||
}
|
||||
if _, err := os.Stat(venvPython); err != nil {
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 3*time.Minute)
|
||||
defer cancel()
|
||||
cmd := exec.CommandContext(ctx, python3, "-m", "venv", venvDir) //nolint:gosec // G204: python3 from LookPath, venvDir is our own dir
|
||||
if out, err := cmd.CombinedOutput(); err != nil {
|
||||
s.log.Error("create runtime venv", "err", err, "out", truncate(string(out), 500))
|
||||
writeJSON(w, http.StatusConflict, map[string]string{"error": "could not create the python venv"})
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if err := s.install(r.Context(), venvPython, pkg); err != nil {
|
||||
s.log.Error("install scimesh runtime", "err", err)
|
||||
writeJSON(w, http.StatusConflict, map[string]string{
|
||||
"error": "pip install " + pkg + " failed: " + err.Error() +
|
||||
". Set SCIMESH_PIP_PACKAGE to your scimesh wheel or index and retry.",
|
||||
})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, installRuntimeResponse{OK: true, Python: venvPython, Installed: true})
|
||||
}
|
||||
|
||||
// venvPython returns the managed venv python when the runtime installer has
|
||||
// created one, so the task runner can be pointed at it automatically.
|
||||
func (s *Server) venvPython() string {
|
||||
for _, candidate := range []string{
|
||||
filepath.Join(s.dir, "venv", "bin", "python"),
|
||||
filepath.Join(s.dir, "venv", "Scripts", "python.exe"),
|
||||
} {
|
||||
if _, err := os.Stat(candidate); err == nil {
|
||||
return candidate
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// installScimeshWithPip installs the package with the venv's own pip,
|
||||
// streaming into the agent log so a long build is not silent.
|
||||
func installScimeshWithPip(ctx context.Context, venvPython, pkg string) error {
|
||||
pip := filepath.Join(filepath.Dir(venvPython), "pip")
|
||||
if _, err := os.Stat(pip); err != nil {
|
||||
pip += ".exe"
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(ctx, 10*time.Minute)
|
||||
defer cancel()
|
||||
cmd := exec.CommandContext(ctx, pip, "install", pkg) //nolint:gosec // G204: pip from our venv, pkg is operator-set or a fixed default
|
||||
return cmd.Run()
|
||||
}
|
||||
|
||||
func stat(path string) bool {
|
||||
_, err := os.Stat(path)
|
||||
return err == nil
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "…"
|
||||
}
|
||||
@@ -0,0 +1,499 @@
|
||||
package setupui
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/agent"
|
||||
)
|
||||
|
||||
func testLogger() *slog.Logger {
|
||||
return slog.New(slog.NewTextHandler(io.Discard, nil))
|
||||
}
|
||||
|
||||
func newTestServer(t *testing.T, sup Supervisor) (*Server, string) {
|
||||
t.Helper()
|
||||
return newTestServerWithInstall(t, sup, nil)
|
||||
}
|
||||
|
||||
func newTestServerWithInstall(t *testing.T, sup Supervisor, install func(ctx context.Context, venvPython, pkg string) error) (*Server, string) {
|
||||
t.Helper()
|
||||
return newTestServerWithInstallAndWheel(t, sup, install, nil)
|
||||
}
|
||||
|
||||
func newTestServerWithInstallAndWheel(t *testing.T, sup Supervisor, install func(ctx context.Context, venvPython, pkg string) error, wheel func(ctx context.Context, url, dir string) (string, error)) (*Server, string) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
server := New(testLogger(), Options{
|
||||
// A distinct random port per test: Port 0 means "the default 12700" in
|
||||
// the server, which would let the shared http.Client pool reuse a stale
|
||||
// keep-alive connection across tests (EOF after a Shutdown).
|
||||
Port: freePort(t),
|
||||
ConfigPath: filepath.Join(dir, "config.json"),
|
||||
Dir: dir,
|
||||
Supervisor: sup,
|
||||
OpenBrowser: func(string) {},
|
||||
InstallScimesh: install,
|
||||
DownloadWheel: wheel,
|
||||
})
|
||||
listener, err := server.Listen()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _ = listener.Close() })
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
t.Cleanup(cancel)
|
||||
go func() { _ = server.Serve(ctx, listener) }()
|
||||
return server, "http://" + listener.Addr().String()
|
||||
}
|
||||
|
||||
type fakeSup struct {
|
||||
mu sync.Mutex
|
||||
started bool
|
||||
stopped bool
|
||||
pid int
|
||||
alive bool
|
||||
}
|
||||
|
||||
func (f *fakeSup) Start(configPath, logPath string) (int, error) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.started = true
|
||||
f.alive = true
|
||||
f.pid = 4242
|
||||
return f.pid, nil
|
||||
}
|
||||
|
||||
func (f *fakeSup) Stop() error {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.stopped = true
|
||||
f.alive = false
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeSup) Pid() int { f.mu.Lock(); defer f.mu.Unlock(); return f.pid }
|
||||
func (f *fakeSup) Alive() bool {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return f.alive
|
||||
}
|
||||
|
||||
func postJSON(t *testing.T, base, path string, body any) (*httptest.ResponseRecorder, map[string]any) {
|
||||
t.Helper()
|
||||
req, err := http.NewRequestWithContext(context.Background(), http.MethodPost, base+path, strings.NewReader(mustJSON(t, body)))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
// No keep-alive pooling: a pooled connection to a shut-down test server
|
||||
// would surface as an EOF instead of a fresh dial.
|
||||
client := http.Client{Transport: &http.Transport{DisableKeepAlives: true}}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
rec := httptest.NewRecorder()
|
||||
rec.Code = resp.StatusCode
|
||||
data := map[string]any{}
|
||||
_ = json.NewDecoder(resp.Body).Decode(&data)
|
||||
return rec, data
|
||||
}
|
||||
|
||||
// freePort reserves an ephemeral port and returns it. The listener is closed
|
||||
// immediately; the tiny reuse window is acceptable for tests and each test
|
||||
// gets a different port, so nothing can collide or share pooled connections.
|
||||
func freePort(t *testing.T) int {
|
||||
t.Helper()
|
||||
listener, err := (&net.ListenConfig{}).Listen(context.Background(), "tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
port := listener.Addr().(*net.TCPAddr).Port
|
||||
_ = listener.Close()
|
||||
return port
|
||||
}
|
||||
|
||||
func mustJSON(t *testing.T, v any) string {
|
||||
t.Helper()
|
||||
raw, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
func TestWizardSavesConfigWithStrictPermissions(t *testing.T) {
|
||||
sup := &fakeSup{}
|
||||
server, base := newTestServer(t, sup)
|
||||
|
||||
rec, _ := postJSON(t, base, "/api/config", map[string]any{
|
||||
"coordinator_url": "http://192.168.1.10:8080",
|
||||
"token": "sm_live_secret",
|
||||
"work_dir": "/home/emil/scimesh-worker",
|
||||
"worker_name": "emil-laptop",
|
||||
"cpu_count": 8,
|
||||
"memory_mb": 16384,
|
||||
})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("save config: got %d, want 200", rec.Code)
|
||||
}
|
||||
info, err := os.Stat(server.cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if perm := info.Mode().Perm(); perm != 0o600 {
|
||||
t.Errorf("config perms = %o, want 600", perm)
|
||||
}
|
||||
config, err := agent.LoadConfigFile(server.cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if config.CoordinatorURL != "http://192.168.1.10:8080" || config.Token != "sm_live_secret" || config.WorkDir != "/home/emil/scimesh-worker" || config.WorkerName != "emil-laptop" {
|
||||
t.Errorf("config = %+v", config)
|
||||
}
|
||||
if config.CPUCount != 8 || config.MemoryMB != 16384 {
|
||||
t.Errorf("resources: cpu=%d mem=%d", config.CPUCount, config.MemoryMB)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWizardRejectsInvalidConfig(t *testing.T) {
|
||||
sup := &fakeSup{}
|
||||
_, base := newTestServer(t, sup)
|
||||
|
||||
for _, body := range []map[string]any{
|
||||
{"coordinator_url": "", "token": "x"},
|
||||
{"coordinator_url": "not-a-url", "token": "x"},
|
||||
} {
|
||||
rec, _ := postJSON(t, base, "/api/config", body)
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Errorf("body %v: got %d, want 400", body, rec.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestWizardStartStopLifecycle(t *testing.T) {
|
||||
sup := &fakeSup{}
|
||||
_, base := newTestServer(t, sup)
|
||||
|
||||
// Starting without a saved config is rejected.
|
||||
rec, _ := postJSON(t, base, "/api/start", map[string]any{})
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Errorf("start without config: got %d, want 400", rec.Code)
|
||||
}
|
||||
|
||||
postJSON(t, base, "/api/config", map[string]any{
|
||||
"coordinator_url": "http://127.0.0.1:8080", "token": "t", "work_dir": ".",
|
||||
})
|
||||
rec, data := postJSON(t, base, "/api/start", map[string]any{})
|
||||
if rec.Code != http.StatusOK || int(data["pid"].(float64)) != 4242 {
|
||||
t.Errorf("start: got %d %v, want 200 pid 4242", rec.Code, data)
|
||||
}
|
||||
if !sup.started {
|
||||
t.Error("supervisor never started the worker")
|
||||
}
|
||||
|
||||
// Status reflects the running state.
|
||||
req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, base+"/api/status", nil)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
var status map[string]any
|
||||
_ = json.NewDecoder(resp.Body).Decode(&status)
|
||||
if status["running"] != true || status["pid"] != float64(4242) {
|
||||
t.Errorf("status = %v, want running pid 4242", status)
|
||||
}
|
||||
|
||||
rec, _ = postJSON(t, base, "/api/stop", map[string]any{})
|
||||
if rec.Code != http.StatusOK || !sup.stopped {
|
||||
t.Errorf("stop: got %d stopped=%v, want 200/true", rec.Code, sup.stopped)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWizardStatusPrefillsSavedConfig(t *testing.T) {
|
||||
sup := &fakeSup{}
|
||||
_, base := newTestServer(t, sup)
|
||||
postJSON(t, base, "/api/config", map[string]any{
|
||||
"coordinator_url": "http://10.0.0.5:8080", "worker_key": "smk_abc", "work_dir": "/w", "worker_name": "n1",
|
||||
})
|
||||
req, _ := http.NewRequestWithContext(context.Background(), http.MethodGet, base+"/api/status", nil)
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
var status struct {
|
||||
ConfigPresent bool `json:"config_present"`
|
||||
WorkerName string `json:"worker_name"`
|
||||
Coordinator string `json:"coordinator"`
|
||||
TokenSet bool `json:"token_set"`
|
||||
}
|
||||
_ = json.NewDecoder(resp.Body).Decode(&status)
|
||||
if !status.ConfigPresent || status.WorkerName != "n1" || status.Coordinator != "http://10.0.0.5:8080" || !status.TokenSet {
|
||||
t.Errorf("status = %+v", status)
|
||||
}
|
||||
// The secret must never appear in the status projection.
|
||||
if strings.Contains(strings.ToLower(mustJSON(t, status)), "smk_abc") {
|
||||
t.Error("status leaks the worker key")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckCoordinatorReachable(t *testing.T) {
|
||||
stub := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/health" {
|
||||
_, _ = w.Write([]byte(`{"status":"ok"}`))
|
||||
return
|
||||
}
|
||||
http.NotFound(w, r)
|
||||
}))
|
||||
defer stub.Close()
|
||||
|
||||
report := agent.CheckCoordinator(context.Background(), stub.URL, 5*time.Second)
|
||||
if !report.Coordinator.OK {
|
||||
t.Errorf("coordinator check = %+v, want ok", report.Coordinator)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCheckCoordinatorUnreachable(t *testing.T) {
|
||||
report := agent.CheckCoordinator(context.Background(), "http://127.0.0.1:1", 2*time.Second)
|
||||
if report.Coordinator.OK {
|
||||
t.Error("unreachable coordinator reported ok")
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfigFileDefaultsAndEnvOverride(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.json")
|
||||
file := agent.ConfigFile{
|
||||
CoordinatorURL: "http://coord:8080",
|
||||
Token: "file-token",
|
||||
WorkDir: "/w",
|
||||
CPUCount: 4,
|
||||
}
|
||||
if err := agent.SaveConfigFile(path, file); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Setenv("COORDINATOR_URL", "http://env:9090")
|
||||
t.Setenv("WORKER_AUTH_TOKEN", "")
|
||||
config, err := agent.LoadConfigFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if config.CoordinatorURL != "http://env:9090" {
|
||||
t.Errorf("env must win: %s", config.CoordinatorURL)
|
||||
}
|
||||
if config.Token != "file-token" {
|
||||
t.Errorf("token = %q, want the file value", config.Token)
|
||||
}
|
||||
if config.CPUCount != 4 {
|
||||
t.Errorf("cpu = %d", config.CPUCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstallRuntimeCreatesVenvAndReportsPython(t *testing.T) {
|
||||
var installedPkg string
|
||||
sup := &fakeSup{}
|
||||
_, base := newTestServerWithInstall(t, sup, func(ctx context.Context, venvPython, pkg string) error {
|
||||
installedPkg = pkg
|
||||
// Prove the venv python path really exists by creating a marker file
|
||||
// where the real venv python would be.
|
||||
_ = os.MkdirAll(filepath.Dir(venvPython), 0o755)
|
||||
_ = os.WriteFile(venvPython, []byte("#!/bin/sh\nexit 0\n"), 0o755)
|
||||
return nil
|
||||
})
|
||||
|
||||
rec, data := postJSON(t, base, "/api/runtime/install", map[string]any{"scimesh_package": "/wheels/scimesh.whl"})
|
||||
if rec.Code != http.StatusOK || data["ok"] != true {
|
||||
t.Fatalf("install: got %d %v, want 200 ok", rec.Code, data)
|
||||
}
|
||||
if installedPkg != "/wheels/scimesh.whl" {
|
||||
t.Errorf("package = %q, want the requested wheel", installedPkg)
|
||||
}
|
||||
if !strings.HasSuffix(data["python"].(string), "venv/bin/python") {
|
||||
t.Errorf("python = %v, want the venv python", data["python"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstallRuntimeRequiresASource(t *testing.T) {
|
||||
sup := &fakeSup{}
|
||||
_, base := newTestServerWithInstall(t, sup, func(ctx context.Context, venvPython, pkg string) error {
|
||||
t.Fatal("install must not run without a package source")
|
||||
return nil
|
||||
})
|
||||
// No source anywhere (SCIMESH_PIP_PACKAGE unset, request empty): 409 with
|
||||
// guidance. The PyPI name is another project, so no silent fallback.
|
||||
rec, data := postJSON(t, base, "/api/runtime/install", map[string]any{})
|
||||
if rec.Code != http.StatusConflict {
|
||||
t.Fatalf("install without source: got %d, want 409", rec.Code)
|
||||
}
|
||||
if !strings.Contains(data["error"].(string), "SCIMESH_PIP_PACKAGE") {
|
||||
t.Errorf("error = %v, want a hint about SCIMESH_PIP_PACKAGE", data["error"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstallRuntimeFailureIsExplained(t *testing.T) {
|
||||
sup := &fakeSup{}
|
||||
_, base := newTestServerWithInstall(t, sup, func(ctx context.Context, venvPython, pkg string) error {
|
||||
return errors.New("no matching distribution found")
|
||||
})
|
||||
rec, data := postJSON(t, base, "/api/runtime/install", map[string]any{"scimesh_package": "/wheels/scimesh.whl"})
|
||||
if rec.Code != http.StatusConflict {
|
||||
t.Fatalf("install failure: got %d, want 409", rec.Code)
|
||||
}
|
||||
if !strings.Contains(data["error"].(string), "SCIMESH_PIP_PACKAGE") {
|
||||
t.Errorf("error = %v, want a hint about SCIMESH_PIP_PACKAGE", data["error"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstallRuntimeDownloadsReleaseWheelWhenNoSource(t *testing.T) {
|
||||
oldVersion := agent.Version
|
||||
agent.Version = "1.1.0-alpha.10"
|
||||
t.Cleanup(func() { agent.Version = oldVersion })
|
||||
|
||||
var downloadedURL, installedPkg string
|
||||
sup := &fakeSup{}
|
||||
_, base := newTestServerWithInstallAndWheel(t, sup,
|
||||
func(ctx context.Context, venvPython, pkg string) error {
|
||||
installedPkg = pkg
|
||||
return nil
|
||||
},
|
||||
func(ctx context.Context, url, dir string) (string, error) {
|
||||
downloadedURL = url
|
||||
return filepath.Join(dir, "scimesh-1.1.0a10-py3-none-any.whl"), nil
|
||||
})
|
||||
|
||||
rec, data := postJSON(t, base, "/api/runtime/install", map[string]any{})
|
||||
if rec.Code != http.StatusOK || data["ok"] != true {
|
||||
t.Fatalf("install: got %d %v, want 200 ok", rec.Code, data)
|
||||
}
|
||||
if !strings.Contains(downloadedURL, "releases/download/v1.1.0-alpha.10/scimesh-1.1.0a10-py3-none-any.whl") {
|
||||
t.Errorf("download url = %q, want the release wheel of this version", downloadedURL)
|
||||
}
|
||||
if !strings.HasSuffix(installedPkg, "scimesh-1.1.0a10-py3-none-any.whl") {
|
||||
t.Errorf("pip received %q, want the downloaded wheel", installedPkg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInstallRuntimeWheelDownloadFailureIsExplained(t *testing.T) {
|
||||
oldVersion := agent.Version
|
||||
agent.Version = "1.1.0-alpha.10"
|
||||
t.Cleanup(func() { agent.Version = oldVersion })
|
||||
|
||||
sup := &fakeSup{}
|
||||
_, base := newTestServerWithInstallAndWheel(t, sup,
|
||||
func(ctx context.Context, venvPython, pkg string) error { t.Fatal("pip must not run"); return nil },
|
||||
func(ctx context.Context, url, dir string) (string, error) {
|
||||
return "", errors.New("HTTP 404")
|
||||
})
|
||||
|
||||
rec, data := postJSON(t, base, "/api/runtime/install", map[string]any{})
|
||||
if rec.Code != http.StatusConflict {
|
||||
t.Fatalf("got %d, want 409", rec.Code)
|
||||
}
|
||||
if !strings.Contains(data["error"].(string), "SCIMESH_PIP_PACKAGE") {
|
||||
t.Errorf("error = %v, want a SCIMESH_PIP_PACKAGE hint", data["error"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestStartPinsTheVenvTaskRunner(t *testing.T) {
|
||||
sup := &fakeSup{}
|
||||
server, base := newTestServer(t, sup)
|
||||
postJSON(t, base, "/api/config", map[string]any{
|
||||
"coordinator_url": "http://coord:8080", "token": "t", "work_dir": ".",
|
||||
})
|
||||
// Simulate the runtime installer: create the venv python marker.
|
||||
venvPython := filepath.Join(server.dir, "venv", "bin", "python")
|
||||
_ = os.MkdirAll(filepath.Dir(venvPython), 0o755)
|
||||
_ = os.WriteFile(venvPython, []byte("#!/bin/sh\nexit 0\n"), 0o755)
|
||||
|
||||
rec, _ := postJSON(t, base, "/api/start", map[string]any{})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("start: got %d, want 200", rec.Code)
|
||||
}
|
||||
config, err := agent.LoadConfigFile(server.cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(config.TaskRunner) != 3 || config.TaskRunner[0] != venvPython || config.TaskRunner[1] != "-m" || config.TaskRunner[2] != "scimesh.worker.task" {
|
||||
t.Errorf("task runner = %v, want the venv python runner", config.TaskRunner)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSaveConfigPinsVenvRunnerWhenPresent(t *testing.T) {
|
||||
server, base := newTestServer(t, &fakeSup{})
|
||||
venvPython := filepath.Join(server.dir, "venv", "bin", "python")
|
||||
_ = os.MkdirAll(filepath.Dir(venvPython), 0o755)
|
||||
_ = os.WriteFile(venvPython, []byte("#!/bin/sh\nexit 0\n"), 0o755)
|
||||
|
||||
rec, _ := postJSON(t, base, "/api/config", map[string]any{
|
||||
"coordinator_url": "http://coord:8080", "token": "t", "work_dir": ".",
|
||||
})
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("config: got %d", rec.Code)
|
||||
}
|
||||
config, err := agent.LoadConfigFile(server.cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(config.TaskRunner) != 3 || config.TaskRunner[0] != venvPython {
|
||||
t.Errorf("task runner = %v, want the venv python", config.TaskRunner)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTestProbesTheVenvPythonAfterInstall(t *testing.T) {
|
||||
sup := &fakeSup{}
|
||||
server, base := newTestServer(t, sup)
|
||||
// The runtime installer leaves a venv python; make it a stub that reports
|
||||
// a fake scimesh version so the preflight goes green through the venv.
|
||||
venvPython := filepath.Join(server.dir, "venv", "bin", "python")
|
||||
_ = os.MkdirAll(filepath.Dir(venvPython), 0o755)
|
||||
_ = os.WriteFile(venvPython, []byte("#!/bin/sh\nif [ \"$1\" = \"-c\" ]; then echo 9.9.9-test; exit 0; fi\nexit 0\n"), 0o755)
|
||||
|
||||
req, _ := http.NewRequestWithContext(context.Background(), http.MethodPost, base+"/api/test", strings.NewReader(`{"coordinator_url":"http://127.0.0.1:1"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
var report agent.CheckReport
|
||||
if err := json.NewDecoder(resp.Body).Decode(&report); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !report.Scimesh.OK || report.Scimesh.Detail != "9.9.9-test" {
|
||||
t.Errorf("scimesh check = %+v, want the venv interpreter reporting 9.9.9-test", report.Scimesh)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseWorkerStats(t *testing.T) {
|
||||
logPath := filepath.Join(t.TempDir(), "worker.log")
|
||||
content := `time=1 level=INFO msg=registered worker_id=w1
|
||||
time=2 level=INFO msg="task claimed" task_id=t1 attempt=0
|
||||
time=3 level=INFO msg="task completed" task_id=t1 elapsed_seconds=2
|
||||
time=4 level=INFO msg="task claimed" task_id=t2 attempt=0
|
||||
time=5 level=WARN msg="task failed" task_id=t2 error_code=X retryable=true
|
||||
time=6 level=WARN msg="agent cycle failed" error="boom"
|
||||
`
|
||||
if err := os.WriteFile(logPath, []byte(content), 0o600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
stats := parseWorkerStats(logPath)
|
||||
if !stats.Registered || stats.Claimed != 2 || stats.Completed != 1 || stats.Failed != 1 {
|
||||
t.Errorf("stats = %+v, want registered claimed=2 completed=1 failed=1", stats)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,338 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>SciMesh Worker · Setup</title>
|
||||
<style>
|
||||
:root{--bg:#0b0e13;--panel:#11151d;--panel-2:#161c27;--border:#1f2634;--border-soft:#181f2b;--text:#e9ecf3;--text-2:#98a2b5;--text-3:#5d6879;--accent:#5b8cff;--accent-soft:#5b8cff1f;--accent-strong:#3f6fe0;--green:#3fce8a;--green-soft:#3fce8a1a;--amber:#e5b64f;--amber-soft:#e5b64f1a;--red:#f2647c;--red-soft:#f2647c1a;--mono:ui-monospace,"SF Mono",SFMono-Regular,Menlo,monospace;color-scheme:dark}
|
||||
*{box-sizing:border-box;margin:0;padding:0}
|
||||
body{background:radial-gradient(900px 500px at 50% -180px,#16233d66,transparent),var(--bg);color:var(--text);font:14px/1.55 Inter,-apple-system,"Segoe UI",Roboto,sans-serif;-webkit-font-smoothing:antialiased;min-height:100vh}
|
||||
button{font:inherit;color:inherit;background:none;border:0;cursor:pointer}
|
||||
input{font:inherit;color:var(--text);background:var(--panel-2);border:1px solid var(--border);border-radius:9px;padding:10px 13px;width:100%;outline:none;transition:border-color .12s,box-shadow .12s}
|
||||
input:focus{border-color:var(--accent);box-shadow:0 0 0 3px var(--accent-soft)}
|
||||
input::placeholder{color:var(--text-3)}
|
||||
code{font-family:var(--mono);font-size:.86em}
|
||||
.shell{max-width:660px;margin:0 auto;padding:44px 22px 70px}
|
||||
.brand{display:flex;align-items:center;justify-content:center;gap:11px;margin-bottom:8px}
|
||||
.brand-mark{display:grid;place-items:center;width:34px;height:34px;border-radius:10px;background:linear-gradient(135deg,#5b8cff,#7c5cff);box-shadow:0 4px 16px #5b8cff40}
|
||||
.brand-mark svg{width:18px;height:18px;stroke:#fff}
|
||||
.brand-name{font-weight:700;font-size:16px;letter-spacing:-.01em}
|
||||
.brand-name span{color:var(--text-3);font-weight:500}
|
||||
.tagline{text-align:center;color:var(--text-3);font-size:12.5px;margin-bottom:34px}
|
||||
.tagline code{color:var(--text-2)}
|
||||
.steps{display:flex;align-items:center;justify-content:center;gap:0;margin-bottom:30px}
|
||||
.step{display:flex;flex-direction:column;align-items:center;gap:7px;width:96px}
|
||||
.step-dot{display:grid;place-items:center;width:30px;height:30px;border-radius:50%;border:1.5px solid var(--border);background:var(--panel);color:var(--text-3);font-size:12.5px;font-weight:700;transition:all .2s}
|
||||
.step-label{font-size:11px;font-weight:600;color:var(--text-3);letter-spacing:.02em}
|
||||
.step.active .step-dot{border-color:var(--accent);background:var(--accent-soft);color:var(--accent);box-shadow:0 0 0 4px #5b8cff14}
|
||||
.step.active .step-label{color:var(--text)}
|
||||
.step.done .step-dot{border-color:var(--green);background:var(--green-soft);color:var(--green)}
|
||||
.step.done .step-label{color:var(--text-2)}
|
||||
.step-line{flex:1;max-width:44px;height:1.5px;background:var(--border);margin:0 6px 22px;position:relative;overflow:hidden}
|
||||
.step-line.done:after{content:"";position:absolute;inset:0;background:var(--green)}
|
||||
.card{background:var(--panel);border:1px solid var(--border-soft);border-radius:15px;padding:26px 28px;box-shadow:0 24px 60px #0000004d}
|
||||
.card h1{font-size:18px;font-weight:700;letter-spacing:-.02em;margin-bottom:4px}
|
||||
.card .sub{color:var(--text-2);font-size:13px;margin-bottom:22px}
|
||||
.field{margin-bottom:16px}
|
||||
.field label{display:block;font-size:12px;font-weight:650;letter-spacing:.04em;text-transform:uppercase;color:var(--text-3);margin-bottom:7px}
|
||||
.field .hint{margin-top:6px;font-size:12px;color:var(--text-3)}
|
||||
.field .hint code{color:var(--text-2)}
|
||||
.radio-grid{display:grid;grid-template-columns:1fr 1fr;gap:10px}
|
||||
.radio-card{border:1px solid var(--border);border-radius:11px;padding:13px 14px;cursor:pointer;transition:all .13s;background:var(--panel-2)}
|
||||
.radio-card:hover{border-color:#2a3446}
|
||||
.radio-card.sel{border-color:var(--accent);background:var(--accent-soft);box-shadow:0 0 0 3px #5b8cff14}
|
||||
.radio-card b{display:flex;align-items:center;gap:8px;font-size:13.5px}
|
||||
.radio-card b svg{width:15px;height:15px;stroke:var(--accent)}
|
||||
.radio-card p{margin-top:4px;font-size:12px;color:var(--text-2)}
|
||||
.check-row{display:flex;align-items:center;gap:12px;padding:11px 14px;border:1px solid var(--border-soft);border-radius:10px;margin-bottom:9px;background:var(--panel-2)}
|
||||
.check-ic{display:grid;place-items:center;width:24px;height:24px;border-radius:50%;flex:none}
|
||||
.check-ic svg{width:13px;height:13px;stroke-width:2.6}
|
||||
.check-ok{background:var(--green-soft)}.check-ok svg{stroke:var(--green)}
|
||||
.check-bad{background:var(--red-soft)}.check-bad svg{stroke:var(--red)}
|
||||
.check-wait{background:#ffffff10}.check-wait svg{stroke:var(--text-3)}
|
||||
.check-row b{font-size:13.5px;font-weight:600}
|
||||
.check-row span{display:block;font-size:12px;color:var(--text-3)}
|
||||
.check-row .ms{margin-left:auto;font:11.5px var(--mono);color:var(--text-3)}
|
||||
.actions{display:flex;align-items:center;justify-content:space-between;gap:12px;margin-top:24px}
|
||||
.btn{display:inline-flex;align-items:center;gap:8px;border-radius:9px;padding:10px 18px;font-weight:650;font-size:13.5px;border:1px solid transparent;transition:all .13s}
|
||||
.btn svg{width:15px;height:15px;stroke:currentColor}
|
||||
.btn-primary{background:var(--accent);color:#0a1222}
|
||||
.btn-primary:hover{background:var(--accent-strong);color:#fff}
|
||||
.btn-ghost{border-color:var(--border);color:var(--text-2);background:var(--panel-2)}
|
||||
.btn-ghost:hover{color:var(--text);border-color:#2a3446}
|
||||
.btn-danger{background:var(--red-soft);color:var(--red)}
|
||||
.btn-lg{padding:12px 24px;font-size:14.5px;border-radius:10px}
|
||||
.link{color:var(--text-3);font-size:13px}
|
||||
.link:hover{color:var(--text)}
|
||||
.status-head{display:flex;align-items:center;gap:14px;margin-bottom:22px}
|
||||
.pulse{position:relative;width:12px;height:12px;border-radius:50%;background:var(--green);flex:none}
|
||||
.pulse:after{content:"";position:absolute;inset:-5px;border-radius:50%;border:2px solid var(--green);opacity:.5;animation:ping 1.6s ease-out infinite}
|
||||
@keyframes ping{from{transform:scale(.6);opacity:.7}to{transform:scale(1.4);opacity:0}}
|
||||
.status-head h1{font-size:19px}
|
||||
.status-head .sub{font-size:12.5px;color:var(--text-3)}
|
||||
.stat-row{display:grid;grid-template-columns:repeat(4,1fr);gap:10px;margin-bottom:18px}
|
||||
.stat{background:var(--panel-2);border:1px solid var(--border-soft);border-radius:11px;padding:12px 14px}
|
||||
.stat b{display:block;font-size:20px;font-weight:700;letter-spacing:-.02em}
|
||||
.stat span{font-size:11px;font-weight:600;letter-spacing:.05em;text-transform:uppercase;color:var(--text-3)}
|
||||
.stat.bad b{color:var(--red)}
|
||||
.logbox{background:#0a0d12;border:1px solid var(--border-soft);border-radius:11px;padding:14px 16px;font:12px/1.7 var(--mono);color:#8fa3bf;max-height:210px;overflow-y:auto;white-space:pre-wrap;word-break:break-word}
|
||||
.meta-line{display:flex;gap:8px;flex-wrap:wrap;margin-bottom:16px}
|
||||
.chip{display:inline-flex;align-items:center;gap:6px;border:1px solid var(--border);border-radius:99px;padding:4px 11px;font-size:12px;color:var(--text-2);background:var(--panel-2)}
|
||||
.chip svg{width:12px;height:12px;stroke:var(--text-3)}
|
||||
.wizard-page{display:none}.wizard-page.active{display:block;animation:fade .18s ease}
|
||||
@keyframes fade{from{opacity:0;transform:translateY(5px)}to{opacity:1}}
|
||||
.error-strip{background:var(--red-soft);border:1px solid #f2647c33;border-radius:9px;padding:9px 12px;font-size:12.5px;color:#ffb3c0;margin-bottom:14px}
|
||||
.spinner{display:inline-block;width:13px;height:13px;border:2px solid var(--text-3);border-top-color:transparent;border-radius:50%;animation:spin .7s linear infinite;vertical-align:-2px;margin-right:7px}
|
||||
@keyframes spin{to{transform:rotate(360deg)}}
|
||||
.checks{padding-bottom:6px}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="shell">
|
||||
<div class="brand">
|
||||
<div class="brand-mark"><svg viewBox="0 0 24 24" fill="none" stroke-width="2" stroke-linecap="round"><circle cx="6" cy="6" r="2.4"/><circle cx="18" cy="7" r="2.4"/><circle cx="12" cy="17" r="2.4"/><path d="M8 7.5l7.6-1M7 8.3l3.7 6.5M16.6 9.2l-3.2 5.6"/></svg></div>
|
||||
<div class="brand-name">SciMesh <span>· Worker setup</span></div>
|
||||
</div>
|
||||
<p class="tagline">Local wizard served by <code>worker-agent setup</code> · <code>127.0.0.1</code></p>
|
||||
|
||||
<!-- ═══ WIZARD VIEW ═══ -->
|
||||
<div id="view-wizard">
|
||||
<div class="steps" id="steps">
|
||||
<div class="step active" id="st1"><div class="step-dot">1</div><div class="step-label">Connect</div></div>
|
||||
<div class="step-line" id="sl1"></div>
|
||||
<div class="step" id="st2"><div class="step-dot">2</div><div class="step-label">Machine</div></div>
|
||||
<div class="step-line" id="sl2"></div>
|
||||
<div class="step" id="st3"><div class="step-dot">3</div><div class="step-label">Check</div></div>
|
||||
<div class="step-line" id="sl3"></div>
|
||||
<div class="step" id="st4"><div class="step-dot">4</div><div class="step-label">Run</div></div>
|
||||
</div>
|
||||
<div id="error-box"></div>
|
||||
|
||||
<!-- step 1: connect -->
|
||||
<div class="wizard-page active" id="wp1">
|
||||
<div class="card">
|
||||
<h1>Connect to a coordinator</h1>
|
||||
<p class="sub">The coordinator hands out work and collects results. Ask your cluster admin for its address.</p>
|
||||
<div class="field"><label>Coordinator URL</label><input id="in-url" placeholder="http://192.168.1.10:8080"><p class="hint">For a served instance this is the address printed by <code>coordinator serve</code>.</p></div>
|
||||
<div class="field"><label>Authentication</label>
|
||||
<div class="radio-grid" id="auth-grid">
|
||||
<div class="radio-card sel" data-mode="token"><b><svg viewBox="0 0 24 24" fill="none" stroke-width="2"><rect x="4" y="10" width="16" height="10" rx="2"/><path d="M8 10V7a4 4 0 0 1 8 0v3"/></svg>Cluster token</b><p>Serve instances: one token for every worker.</p></div>
|
||||
<div class="radio-card" data-mode="key"><b><svg viewBox="0 0 24 24" fill="none" stroke-width="2"><circle cx="8" cy="14" r="4"/><path d="M10.8 11.2L20 2M15 4l3 3"/></svg>Worker key</b><p>Shared clusters: a key tied to your account.</p></div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="field" id="token-field"><label>Token</label><input id="in-token" type="password" placeholder="paste the token from coordinator token"><p class="hint">The admin can copy it from <code>coordinator token</code> on the server.</p></div>
|
||||
<div class="field" id="key-field" style="display:none"><label>Worker key + userservice URL</label><input id="in-key" type="password" placeholder="smk_…"><input id="in-users" placeholder="http://userservice-host:8081" style="margin-top:8px"><p class="hint">Create a key in the coordinator UI: Users → worker keys.</p></div>
|
||||
<div class="actions"><span class="link" id="l1">Step 1 of 4</span><button class="btn btn-primary" id="b1">Continue →</button></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- step 2: machine -->
|
||||
<div class="wizard-page" id="wp2">
|
||||
<div class="card">
|
||||
<h1>This machine</h1>
|
||||
<p class="sub">Where tasks run and how the machine appears in the cluster.</p>
|
||||
<div class="field"><label>Worker name</label><input id="in-name" placeholder="auto-detected"><p class="hint">Shown in the coordinator’s worker list.</p></div>
|
||||
<div class="field"><label>Work directory</label><input id="in-dir" placeholder="./scimesh-agent-data"><p class="hint">Datasets and shard results live here. ~1 GB free space recommended.</p></div>
|
||||
<div class="field"><label>Compute resources advertised</label>
|
||||
<div class="radio-grid">
|
||||
<div class="radio-card sel" data-cpu="auto"><b>Auto</b><p>Detect the machine’s CPU count.</p></div>
|
||||
<div class="radio-card" data-cpu="custom"><b>Custom…</b><p>Limit what this machine advertises.</p></div>
|
||||
</div>
|
||||
</div>
|
||||
<div class="field" id="cpu-field" style="display:none"><label>CPU count</label><input id="in-cpu" type="number" min="1" value="1"></div>
|
||||
<div class="field"><label>Concurrent task loops</label><input id="in-conc" type="number" min="1" max="64" value="1"><p class="hint">Process this many shards in parallel on this machine. Each loop runs its own task runner subprocess.</p></div>
|
||||
<div class="actions"><button class="btn btn-ghost" id="b2b">← Back</button><button class="btn btn-primary" id="b2">Continue →</button></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- step 3: preflight -->
|
||||
<div class="wizard-page" id="wp3">
|
||||
<div class="card">
|
||||
<h1>Preflight check</h1>
|
||||
<p class="sub">Making sure this machine can reach the coordinator and run SciMesh workloads.</p>
|
||||
<div class="checks" id="checks"></div>
|
||||
<div class="actions"><button class="btn btn-ghost" id="b3b">← Back</button><button class="btn btn-primary" id="b3" disabled>Continue anyway →</button></div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- step 4: run -->
|
||||
<div class="wizard-page" id="wp4">
|
||||
<div class="card" style="text-align:center;padding:40px 28px">
|
||||
<div class="brand-mark" style="margin:0 auto 18px;width:46px;height:46px;border-radius:13px"><svg viewBox="0 0 24 24" fill="none" stroke-width="2" stroke-linecap="round" style="width:22px;height:22px"><path d="M6 4l14 8-14 8V4z"/></svg></div>
|
||||
<h1 style="font-size:20px">Ready to join the cluster</h1>
|
||||
<p class="sub" style="max-width:380px;margin:8px auto 26px">Configuration will be saved and the worker started as a background process.</p>
|
||||
<div class="actions" style="justify-content:space-between;margin-top:30px"><button class="btn btn-ghost" id="b4b">← Back</button><button class="btn btn-primary btn-lg" id="b4">Start worker</button></div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- ═══ STATUS VIEW ═══ -->
|
||||
<div id="view-status" style="display:none">
|
||||
<div class="card">
|
||||
<div class="status-head">
|
||||
<div class="pulse" id="st-pulse"></div>
|
||||
<div><h1 id="st-title">Worker is working</h1><div class="sub" id="st-sub">—</div></div>
|
||||
<button class="btn btn-danger" id="st-stop" style="margin-left:auto"><svg viewBox="0 0 24 24" fill="none" stroke-width="2"><rect x="6" y="6" width="12" height="12" rx="2"/></svg>Stop</button>
|
||||
</div>
|
||||
<div class="stat-row">
|
||||
<div class="stat"><b id="st-claimed">—</b><span>claimed</span></div>
|
||||
<div class="stat"><b id="st-completed">—</b><span>completed</span></div>
|
||||
<div class="stat bad"><b id="st-failed">—</b><span>failed</span></div>
|
||||
<div class="stat"><b id="st-registered">—</b><span>registered</span></div>
|
||||
</div>
|
||||
<div class="meta-line" id="st-meta"></div>
|
||||
<div class="logbox" id="st-log"></div>
|
||||
<div class="actions"><span class="link" id="st-cfg">—</span><button class="btn btn-ghost" id="st-reconfig">Reconfigure…</button></div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
const $=id=>document.getElementById(id);
|
||||
let state={mode:'token',cpu:'auto',venvPython:null};
|
||||
let checksOk=false;
|
||||
|
||||
function err(msg){$('error-box').innerHTML=msg?'<div class="error-strip">'+msg+'</div>':''}
|
||||
function goto(n){
|
||||
['wp1','wp2','wp3','wp4'].forEach((id,i)=>$(id).classList.toggle('active',i===n-1));
|
||||
for(let i=1;i<=4;i++){
|
||||
const st=$('st'+i);
|
||||
st.classList.toggle('done',i<n);st.classList.toggle('active',i===n);
|
||||
st.querySelector('.step-dot').textContent=i<n?'✓':i;
|
||||
if(i<4)$('sl'+i).classList.toggle('done',i<n);
|
||||
}
|
||||
err('');
|
||||
}
|
||||
document.querySelectorAll('#auth-grid .radio-card').forEach(c=>c.addEventListener('click',()=>{
|
||||
document.querySelectorAll('#auth-grid .radio-card').forEach(x=>x.classList.remove('sel'));
|
||||
c.classList.add('sel');state.mode=c.dataset.mode;
|
||||
$('token-field').style.display=state.mode==='token'?'':'none';
|
||||
$('key-field').style.display=state.mode==='key'?'':'none';
|
||||
}));
|
||||
document.querySelectorAll('#wp2 .radio-card').forEach(c=>c.addEventListener('click',()=>{
|
||||
c.parentElement.querySelectorAll('.radio-card').forEach(x=>x.classList.remove('sel'));
|
||||
c.classList.add('sel');state.cpu=c.dataset.cpu;
|
||||
$('cpu-field').style.display=state.cpu==='custom'?'':'none';
|
||||
}));
|
||||
|
||||
async function postJSON(path,body){
|
||||
const r=await fetch(path,{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify(body)});
|
||||
const data=await r.json().catch(()=>({}));
|
||||
return {status:r.status,data};
|
||||
}
|
||||
function draftConfig(){
|
||||
const cfg={
|
||||
coordinator_url:$('in-url').value.trim(),
|
||||
token:state.mode==='token'?$('in-token').value.trim():'',
|
||||
worker_key:state.mode==='key'?$('in-key').value.trim():'',
|
||||
userservice_url:state.mode==='key'?$('in-users').value.trim():'',
|
||||
work_dir:$('in-dir').value.trim(),
|
||||
worker_name:$('in-name').value.trim(),
|
||||
cpu_count:state.cpu==='custom'?parseInt($('in-cpu').value||'1',10):0,
|
||||
concurrency:parseInt($('in-conc').value||'1',10)
|
||||
};
|
||||
if(state.venvPython)cfg.task_runner=[state.venvPython,'-m','scimesh.worker.task'];
|
||||
return cfg;
|
||||
}
|
||||
|
||||
$('b1').onclick=()=>{
|
||||
const c=draftConfig();
|
||||
if(!c.coordinator_url){err('Enter the coordinator URL.');return}
|
||||
if(state.mode==='token'&&!c.token){err('Enter the cluster token.');return}
|
||||
if(state.mode==='key'&&!c.worker_key){err('Enter the worker key.');return}
|
||||
goto(2);
|
||||
};
|
||||
$('b2b').onclick=()=>goto(1);
|
||||
$('b2').onclick=()=>{goto(3);runChecks()};
|
||||
$('b3b').onclick=()=>goto(2);
|
||||
$('b3').onclick=()=>goto(4);
|
||||
$('b4b').onclick=()=>goto(3);
|
||||
$('b4').onclick=async()=>{
|
||||
const c=draftConfig();
|
||||
const r=await postJSON('/api/config',c);
|
||||
if(r.status!==200){err('Could not save the configuration: '+(r.data.error||'unknown error'));return}
|
||||
const s=await postJSON('/api/start',{});
|
||||
if(s.status!==200){err('Could not start the worker: '+(s.data.error||'unknown error'));return}
|
||||
showStatus();
|
||||
};
|
||||
|
||||
const checkRow=(name,ok,detail,ms,action)=>
|
||||
'<div class="check-row"><div class="check-ic '+(ok===null?'check-wait':ok?'check-ok':'check-bad')+'"><svg viewBox="0 0 24 24" fill="none" stroke-linecap="round">'+(ok===null?'<path d="M12 7v5l3 3"/>':ok?'<path d="M4 12l5 5L20 6"/>':'<path d="M6 6l12 12M18 6L6 18"/>')+'</svg></div><div><b>'+name+'</b><span>'+(detail||'')+'</span></div>'+(ms?'<span class="ms">'+ms+' ms</span>':'')+(action?action:'')+'</div>';
|
||||
async function runChecks(){
|
||||
const box=$('checks');
|
||||
box.innerHTML=checkRow('Coordinator reachable','',null,null)+checkRow('Python 3','',null,null)+checkRow('scimesh package','',null,null);
|
||||
const r=await postJSON('/api/test',draftConfig());
|
||||
box.innerHTML='';
|
||||
checksOk=true;
|
||||
const items=[r.data.coordinator,r.data.python,r.data.scimesh];
|
||||
for(const item of items){
|
||||
if(item&&!item.ok)checksOk=false;
|
||||
let action='';
|
||||
if(item&&item.name==='scimesh'&&!item.ok){
|
||||
action='<div style="margin-left:auto;display:flex;gap:8px;align-items:center"><input id="in-pkg" placeholder="wheel path / checkout / index URL (optional)" style="width:230px;padding:7px 10px;font-size:12px"><button class="btn btn-primary" style="padding:6px 12px;font-size:12px" id="install-scimesh">Install</button></div>';
|
||||
}
|
||||
box.insertAdjacentHTML('beforeend',checkRow(item?item.name:'?',item?item.ok:null,item?item.detail:'',item?item.latency_ms:null,action));
|
||||
}
|
||||
const btn=$('install-scimesh');
|
||||
if(btn)btn.onclick=async()=>{
|
||||
const src=$('in-pkg')?$('in-pkg').value.trim():'';
|
||||
btn.disabled=true;btn.textContent='Installing…';
|
||||
const ir=await postJSON('/api/runtime/install',{scimesh_package:src});
|
||||
if(ir.status!==200){err('Could not install scimesh: '+(ir.data.error||'unknown error'));btn.disabled=false;btn.textContent='Install';return}
|
||||
state.venvPython=ir.data.python;
|
||||
runChecks();
|
||||
};
|
||||
$('b3').disabled=!checksOk;
|
||||
}
|
||||
|
||||
async function showStatus(){
|
||||
$('view-wizard').style.display='none';
|
||||
$('view-status').style.display='block';
|
||||
await refreshStatus();
|
||||
setInterval(refreshStatus,2000);
|
||||
}
|
||||
async function refreshStatus(){
|
||||
const r=await fetch('/api/status');
|
||||
if(!r.ok)return;
|
||||
const v=await r.json();
|
||||
$('st-pulse').style.background=v.running?'var(--green)':'var(--text-3)';
|
||||
$('st-pulse').style.animation=v.running?'':'none';
|
||||
$('st-title').textContent=v.running?(v.worker_name||'Worker')+' is working':(v.worker_name||'Worker')+' is stopped';
|
||||
$('st-sub').textContent='pid '+(v.pid||'—')+' · config '+(v.config_present?v.config_path:'not saved yet');
|
||||
$('st-cfg').textContent='Configuration: '+(v.config_present?v.config_path:'—');
|
||||
const stats=v.stats||{};
|
||||
$('st-claimed').textContent=stats.claimed!=null?stats.claimed:'—';
|
||||
$('st-completed').textContent=stats.completed!=null?stats.completed:'—';
|
||||
$('st-failed').textContent=stats.failed!=null?stats.failed:'—';
|
||||
$('st-registered').textContent=stats.registered?'yes':'no';
|
||||
const meta=$('st-meta');
|
||||
meta.innerHTML='';
|
||||
if(v.coordinator){const c=document.createElement('span');c.className='chip';c.innerHTML='<svg viewBox="0 0 24 24" fill="none" stroke-width="2"><circle cx="12" cy="12" r="9"/></svg>'+v.coordinator;meta.append(c)}
|
||||
if(v.work_dir){const d=document.createElement('span');d.className='chip';d.innerHTML='<svg viewBox="0 0 24 24" fill="none" stroke-width="2"><path d="M4 6h16M4 12h16M4 18h16"/></svg>'+v.work_dir;meta.append(d)}
|
||||
if(!v.token_set){const t=document.createElement('span');t.className='chip';t.style.color='var(--amber)';t.textContent='no credential set';meta.append(t)}
|
||||
const logs=await fetch('/api/logs?tail=200');
|
||||
const lv=await logs.json();
|
||||
$('st-log').textContent=lv.log||'(no log yet — the worker writes here once started)';
|
||||
}
|
||||
$('st-stop').onclick=async()=>{await postJSON('/api/stop',{});refreshStatus()};
|
||||
$('st-reconfig').onclick=()=>{
|
||||
$('view-status').style.display='none';
|
||||
$('view-wizard').style.display='block';
|
||||
goto(1);
|
||||
};
|
||||
|
||||
// Prefill from a saved configuration, then decide which view to show.
|
||||
(async()=>{
|
||||
const r=await fetch('/api/status');
|
||||
const v=await r.json();
|
||||
if(v.config_present){
|
||||
$('in-url').value=v.coordinator||'';
|
||||
$('in-dir').value=v.work_dir||'';
|
||||
$('in-name').value=v.worker_name||'';
|
||||
}
|
||||
if(v.running)showStatus();
|
||||
})();
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
@@ -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
|
||||
}
|
||||
@@ -18,4 +18,5 @@ var (
|
||||
ErrResultConflict = errors.New("different result already recorded")
|
||||
ErrInvalidInput = errors.New("invalid input")
|
||||
ErrTaskNotLeased = errors.New("task is not currently leased")
|
||||
ErrWorkloadDisabled = errors.New("workload is disabled")
|
||||
)
|
||||
|
||||
@@ -19,7 +19,11 @@ const (
|
||||
|
||||
// Job is one user submission that fans out into one or more tasks.
|
||||
type Job struct {
|
||||
ID uuid.UUID
|
||||
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
|
||||
@@ -119,6 +123,10 @@ 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:
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
package domain
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
@@ -14,14 +15,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,15 +46,23 @@ 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
|
||||
}
|
||||
name = strings.TrimSpace(name)
|
||||
if name == "" {
|
||||
return nil, ErrInvalidInput
|
||||
}
|
||||
return &Worker{
|
||||
ID: uuid.New(),
|
||||
Name: name,
|
||||
Capabilities: capabilities,
|
||||
Status: WorkerOnline,
|
||||
TrustLevel: WorkerTrusted,
|
||||
LastHeartbeatAt: now,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
|
||||
@@ -29,3 +29,11 @@ func TestNewWorkerRejectsNoCapabilities(t *testing.T) {
|
||||
t.Errorf("empty slice: err = %v, want ErrInvalidInput", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewWorkerRejectsBlankName(t *testing.T) {
|
||||
for _, name := range []string{"", " ", "\t\n"} {
|
||||
if _, err := NewWorker(name, []string{"similarity-search"}, testNow); !errors.Is(err, ErrInvalidInput) {
|
||||
t.Errorf("name %q: got %v, want ErrInvalidInput", name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"io/fs"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
@@ -26,6 +27,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 +52,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 +72,23 @@ 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
|
||||
// DatabaseEngine selects the storage backend: "sqlite" (embedded, the
|
||||
// single-binary default) or "postgres" (cluster deployments). The
|
||||
// postgres engine requires DATABASE_URL.
|
||||
DatabaseEngine string
|
||||
// DBPath is the sqlite database file (engine=sqlite only).
|
||||
DBPath string
|
||||
}
|
||||
|
||||
// Load reads the environment and fails fast on anything required-but-missing
|
||||
@@ -78,28 +113,47 @@ 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 == "" {
|
||||
return Config{}, fmt.Errorf("DATABASE_URL is required")
|
||||
cfg.DatabaseEngine = getEnv("SCIMESH_DB", "sqlite")
|
||||
switch cfg.DatabaseEngine {
|
||||
case "sqlite", "postgres":
|
||||
default:
|
||||
return Config{}, fmt.Errorf("SCIMESH_DB must be sqlite or postgres")
|
||||
}
|
||||
cfg.DBPath = getEnv("SCIMESH_DB_PATH", filepath.Join(cfg.StorageDir, "scimesh.db"))
|
||||
|
||||
if cfg.DatabaseEngine == "postgres" && cfg.DatabaseURL == "" {
|
||||
return Config{}, fmt.Errorf("DATABASE_URL is required for the postgres engine")
|
||||
}
|
||||
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,10 +183,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
|
||||
}
|
||||
|
||||
|
||||
@@ -314,6 +314,17 @@ func (r *WorkerRepo) MarkStaleOffline(ctx context.Context, cutoff time.Time) (in
|
||||
return n, nil
|
||||
}
|
||||
|
||||
func (r *WorkerRepo) SetTrust(ctx context.Context, id uuid.UUID, trust domain.WorkerTrust) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
w, ok := r.workers[id]
|
||||
if !ok {
|
||||
return domain.ErrWorkerNotFound
|
||||
}
|
||||
w.TrustLevel = trust
|
||||
return nil
|
||||
}
|
||||
|
||||
// --- ArtifactRepo --------------------------------------------------------
|
||||
|
||||
type ArtifactRepo struct {
|
||||
@@ -417,3 +428,65 @@ 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
|
||||
}
|
||||
|
||||
func (r *JobRepo) ListCompletedBefore(ctx context.Context, cutoff time.Time) ([]domain.Job, error) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
var out []domain.Job
|
||||
for _, j := range r.jobs {
|
||||
if j.CompletedAt != nil && j.CompletedAt.Before(cutoff) {
|
||||
out = append(out, *j)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *JobRepo) Delete(ctx context.Context, id uuid.UUID) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
delete(r.jobs, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *WorkerRepo) Delete(ctx context.Context, id uuid.UUID) error {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if _, ok := r.workers[id]; !ok {
|
||||
return domain.ErrWorkerNotFound
|
||||
}
|
||||
delete(r.workers, id)
|
||||
return 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 {
|
||||
@@ -84,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,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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,278 @@
|
||||
//go:build integration
|
||||
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"os"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
)
|
||||
|
||||
// ensureMigrated applies the embedded schema first: the admin tests run
|
||||
// before the dedicated migration test (file order) and need real tables.
|
||||
func ensureMigrated(t *testing.T) {
|
||||
t.Helper()
|
||||
url := os.Getenv("TEST_DATABASE_URL")
|
||||
if url == "" {
|
||||
t.Skip("TEST_DATABASE_URL is not set")
|
||||
}
|
||||
if err := Migrate(context.Background(), url, nil); err != nil {
|
||||
t.Fatalf("migrate: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminWorkerSetTrust(t *testing.T) {
|
||||
ensureMigrated(t)
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
repo := NewWorkerRepo(pool)
|
||||
|
||||
w, err := domain.NewWorker("trust-lab", []string{"similarity-search"}, time.Now().UTC())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := repo.Insert(ctx, w); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { _, _ = pool.Exec(context.Background(), `DELETE FROM workers WHERE id = $1`, w.ID) })
|
||||
|
||||
if err := repo.SetTrust(ctx, w.ID, domain.WorkerUntrusted); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := repo.Get(ctx, w.ID)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.TrustLevel != domain.WorkerUntrusted {
|
||||
t.Errorf("trust = %q, want untrusted", got.TrustLevel)
|
||||
}
|
||||
if err := repo.SetTrust(ctx, uuid.New(), domain.WorkerTrusted); !errors.Is(err, domain.ErrWorkerNotFound) {
|
||||
t.Errorf("unknown worker: got %v, want ErrWorkerNotFound", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminListJobsPaginatedAndCounts(t *testing.T) {
|
||||
ensureMigrated(t)
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
jobRepo := NewJobRepo(pool)
|
||||
adminRepo := NewAdminReadRepo(pool)
|
||||
|
||||
jobs := make([]*domain.Job, 4)
|
||||
for i := range jobs {
|
||||
job, _ := seedJob(t, pool, 1)
|
||||
jobs[i] = job
|
||||
}
|
||||
for _, j := range jobs {
|
||||
t.Cleanup(func() { _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, j.ID) })
|
||||
}
|
||||
if err := jobRepo.UpdateStatus(ctx, jobs[0].ID, domain.JobCompleted, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := jobRepo.UpdateStatus(ctx, jobs[1].ID, domain.JobCompleted, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := jobRepo.UpdateStatus(ctx, jobs[2].ID, domain.JobRunning, nil); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
all, total, err := adminRepo.ListJobsPaginated(ctx, "", 100, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if total != 4 || len(all) != 4 {
|
||||
t.Errorf("all: total=%d len=%d, want 4/4", total, len(all))
|
||||
}
|
||||
completed, total, err := adminRepo.ListJobsPaginated(ctx, "completed", 100, 0)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if total != 2 || len(completed) != 2 {
|
||||
t.Errorf("completed: total=%d len=%d, want 2/2", total, len(completed))
|
||||
}
|
||||
page, total, err := adminRepo.ListJobsPaginated(ctx, "", 2, 2)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if total != 4 || len(page) != 2 {
|
||||
t.Errorf("page: total=%d len=%d, want 4/2", total, len(page))
|
||||
}
|
||||
|
||||
counts, err := adminRepo.CountJobsByStatus(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if counts["completed"] != 2 || counts["running"] != 1 || counts["pending"] != 1 {
|
||||
t.Errorf("counts = %v", counts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminTaskCountsByJobs(t *testing.T) {
|
||||
ensureMigrated(t)
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
job, tasks := seedJob(t, pool, 3)
|
||||
t.Cleanup(func() { _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, job.ID) })
|
||||
|
||||
resultArtifact := seedArtifact(t, pool, job.ID, &tasks[0].ID, domain.ArtifactPartialResult)
|
||||
if _, err := pool.Exec(ctx, `UPDATE tasks SET status = 'completed', result_artifact_id = $1 WHERE id = $2`,
|
||||
resultArtifact.ID, tasks[0].ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `UPDATE tasks SET status = 'failed' WHERE id = $1`, tasks[1].ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
counts, err := NewAdminReadRepo(pool).TaskCountsByJobs(ctx, []uuid.UUID{job.ID})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := counts[job.ID]
|
||||
if got["completed"] != 1 || got["failed"] != 1 || got["pending"] != 1 {
|
||||
t.Errorf("task counts = %v, want completed=1 failed=1 pending=1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminJobCountsByDayAndWorkload(t *testing.T) {
|
||||
ensureMigrated(t)
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
repo := NewAdminReadRepo(pool)
|
||||
|
||||
job, _ := seedJob(t, pool, 1)
|
||||
t.Cleanup(func() { _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, job.ID) })
|
||||
if _, err := pool.Exec(ctx, `UPDATE jobs SET created_at = now() - interval '2 days' WHERE id = $1`, job.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
other, _ := seedJob(t, pool, 1)
|
||||
t.Cleanup(func() { _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, other.ID) })
|
||||
|
||||
byDay, err := repo.JobCountsByDay(ctx, time.Now().UTC().Add(-6*24*time.Hour))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
today := time.Now().UTC().Format("2006-01-02")
|
||||
twoDays := time.Now().UTC().Add(-48 * time.Hour).Format("2006-01-02")
|
||||
if byDay[today] != 1 || byDay[twoDays] != 1 {
|
||||
t.Errorf("by day = %v, want today=1 twoDays=1", byDay)
|
||||
}
|
||||
|
||||
byWorkload, err := repo.JobCountsByWorkload(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if byWorkload[job.Workload] < 2 {
|
||||
t.Errorf("by workload = %v, want at least 2 for %s", byWorkload, job.Workload)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminTaskStatsAndStorage(t *testing.T) {
|
||||
ensureMigrated(t)
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
repo := NewAdminReadRepo(pool)
|
||||
|
||||
job, tasks := seedJob(t, pool, 2)
|
||||
t.Cleanup(func() { _, _ = pool.Exec(context.Background(), `DELETE FROM jobs WHERE id = $1`, job.ID) })
|
||||
|
||||
start := time.Now().UTC().Add(-2 * time.Minute)
|
||||
done := time.Now().UTC().Add(-90 * time.Second)
|
||||
resultArtifact := seedArtifact(t, pool, job.ID, &tasks[0].ID, domain.ArtifactPartialResult)
|
||||
if _, err := pool.Exec(ctx, `UPDATE tasks SET status = 'completed', result_artifact_id = $1, started_at = $2, completed_at = $3 WHERE id = $4`,
|
||||
resultArtifact.ID, start, done, tasks[0].ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := pool.Exec(ctx, `UPDATE tasks SET status = 'failed' WHERE id = $1`, tasks[1].ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
completed, failed, avg, err := repo.TaskStats(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if completed < 1 || failed < 1 {
|
||||
t.Errorf("stats = completed %d failed %d, want >= 1/1", completed, failed)
|
||||
}
|
||||
if avg < 29 || avg > 31 {
|
||||
t.Errorf("avg duration = %.1fs, want ~30s", avg)
|
||||
}
|
||||
|
||||
for _, kind := range []domain.ArtifactKind{domain.ArtifactInput, domain.ArtifactShard} {
|
||||
seedArtifact(t, pool, job.ID, nil, kind)
|
||||
}
|
||||
sizes, err := repo.ArtifactSizeByKind(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if sizes["input"] < 1 || sizes["shard"] < 1 {
|
||||
t.Errorf("sizes = %v, want input/shard > 0", sizes)
|
||||
}
|
||||
dbBytes, err := repo.DatabaseSizeBytes(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if dbBytes <= 0 {
|
||||
t.Errorf("database size = %d, want > 0", dbBytes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWorkloadSettingsRepoRoundTrip(t *testing.T) {
|
||||
ensureMigrated(t)
|
||||
pool := testPool(t)
|
||||
ctx := context.Background()
|
||||
repo := NewWorkloadSettingsRepo(pool)
|
||||
|
||||
enabled, err := repo.GetEnabled(ctx, "similarity-search")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !enabled {
|
||||
t.Error("workload without an override must be enabled")
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
if err := repo.SetEnabled(ctx, "similarity-search", false, now); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = pool.Exec(context.Background(), `DELETE FROM workload_settings WHERE workload = 'similarity-search'`)
|
||||
})
|
||||
enabled, err = repo.GetEnabled(ctx, "similarity-search")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if enabled {
|
||||
t.Error("workload must be disabled after the override")
|
||||
}
|
||||
if err := repo.SetEnabled(ctx, "similarity-search", true, now.Add(time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
enabled, err = repo.GetEnabled(ctx, "similarity-search")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !enabled {
|
||||
t.Error("workload must be re-enabled after the upsert")
|
||||
}
|
||||
list, err := repo.List(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
found := false
|
||||
for _, s := range list {
|
||||
if s.Workload == "similarity-search" {
|
||||
found = true
|
||||
if !s.Enabled {
|
||||
t.Error("list must reflect the re-enabled state")
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("override missing from the settings list")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
sq "github.com/Masterminds/squirrel"
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
|
||||
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
|
||||
)
|
||||
|
||||
// AdminReadRepo backs the coordinator admin console: paginated jobs, status
|
||||
// counters, metrics buckets and storage figures. Read-only.
|
||||
type AdminReadRepo struct{ pool *pgxpool.Pool }
|
||||
|
||||
func NewAdminReadRepo(pool *pgxpool.Pool) *AdminReadRepo { return &AdminReadRepo{pool: pool} }
|
||||
|
||||
var _ usecase.AdminReadRepository = (*AdminReadRepo)(nil)
|
||||
|
||||
func (r *AdminReadRepo) ListJobsPaginated(ctx context.Context, status string, limit, offset int) ([]domain.Job, int, error) {
|
||||
if limit < 1 || limit > 100 || offset < 0 {
|
||||
return nil, 0, domain.ErrInvalidInput
|
||||
}
|
||||
countQ := psql.Select("COUNT(*)").From("jobs")
|
||||
listQ := psql.Select(jobColumns...).From("jobs")
|
||||
if status != "" {
|
||||
countQ = countQ.Where(sq.Eq{"status": status})
|
||||
listQ = listQ.Where(sq.Eq{"status": status})
|
||||
}
|
||||
countSQL, args, err := countQ.ToSql()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
var total int
|
||||
if err := conn(ctx, r.pool).QueryRow(ctx, countSQL, args...).Scan(&total); err != nil {
|
||||
return nil, 0, fmt.Errorf("count jobs: %w", err)
|
||||
}
|
||||
listSQL, args, err := listQ.OrderBy("created_at DESC", "id DESC").Limit(uint64(limit)).Offset(uint64(offset)).ToSql()
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, listSQL, args...)
|
||||
if err != nil {
|
||||
return nil, 0, fmt.Errorf("list jobs paginated: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
jobs := make([]domain.Job, 0)
|
||||
for rows.Next() {
|
||||
var j domain.Job
|
||||
var statusRaw string
|
||||
if err := rows.Scan(
|
||||
&j.ID, &j.Workload, &j.InputURI, &j.Parameters, &statusRaw, &j.CreatedAt, &j.CompletedAt,
|
||||
&j.InputArtifactID, &j.ResultArtifactID, &j.ErrorCode, &j.ErrorMessage, &j.ReducerStartedAt,
|
||||
&j.OwnerID,
|
||||
); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
j.Status = domain.JobStatus(statusRaw)
|
||||
jobs = append(jobs, j)
|
||||
}
|
||||
return jobs, total, rows.Err()
|
||||
}
|
||||
|
||||
func (r *AdminReadRepo) CountJobsByStatus(ctx context.Context) (map[string]int, error) {
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, "SELECT status, COUNT(*) FROM jobs GROUP BY status")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("count jobs by status: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
out := make(map[string]int)
|
||||
for rows.Next() {
|
||||
var status string
|
||||
var count int
|
||||
if err := rows.Scan(&status, &count); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[status] = count
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *AdminReadRepo) TaskCountsByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID]map[string]int, error) {
|
||||
out := make(map[uuid.UUID]map[string]int, len(jobIDs))
|
||||
if len(jobIDs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
sql, args, err := psql.Select("job_id", "status", "COUNT(*)").From("tasks").
|
||||
Where(sq.Eq{"job_id": jobIDs}).GroupBy("job_id", "status").ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("task counts by jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var jobID uuid.UUID
|
||||
var status string
|
||||
var count int
|
||||
if err := rows.Scan(&jobID, &status, &count); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if out[jobID] == nil {
|
||||
out[jobID] = make(map[string]int)
|
||||
}
|
||||
out[jobID][status] = count
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *AdminReadRepo) JobCountsByDay(ctx context.Context, since time.Time) (map[string]int, error) {
|
||||
rows, err := conn(ctx, r.pool).Query(ctx,
|
||||
"SELECT to_char(date_trunc('day', created_at AT TIME ZONE 'UTC'), 'YYYY-MM-DD') AS day, COUNT(*) FROM jobs WHERE created_at >= $1 GROUP BY 1",
|
||||
since.UTC())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("job counts by day: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
out := make(map[string]int)
|
||||
for rows.Next() {
|
||||
var day string
|
||||
var count int
|
||||
if err := rows.Scan(&day, &count); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[day] = count
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *AdminReadRepo) JobCountsByWorkload(ctx context.Context) (map[string]int, error) {
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, "SELECT workload, COUNT(*) FROM jobs GROUP BY workload")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("job counts by workload: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
out := make(map[string]int)
|
||||
for rows.Next() {
|
||||
var workload string
|
||||
var count int
|
||||
if err := rows.Scan(&workload, &count); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[workload] = count
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *AdminReadRepo) TaskStats(ctx context.Context) (int64, int64, float64, error) {
|
||||
var completed, failed int64
|
||||
var avgSeconds *float64
|
||||
err := conn(ctx, r.pool).QueryRow(ctx, `
|
||||
SELECT
|
||||
COALESCE(SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END), 0),
|
||||
COALESCE(SUM(CASE WHEN status = 'failed' THEN 1 ELSE 0 END), 0),
|
||||
AVG(CASE WHEN status = 'completed' AND started_at IS NOT NULL
|
||||
THEN EXTRACT(EPOCH FROM completed_at - started_at) END)::float8
|
||||
FROM tasks`).Scan(&completed, &failed, &avgSeconds)
|
||||
if err != nil {
|
||||
return 0, 0, 0, fmt.Errorf("task stats: %w", err)
|
||||
}
|
||||
var avg float64
|
||||
if avgSeconds != nil {
|
||||
avg = *avgSeconds
|
||||
}
|
||||
return completed, failed, avg, nil
|
||||
}
|
||||
|
||||
func (r *AdminReadRepo) ArtifactSizeByKind(ctx context.Context) (map[string]int64, error) {
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, "SELECT kind, COALESCE(SUM(size_bytes), 0) FROM artifacts GROUP BY kind")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("artifact sizes: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
out := make(map[string]int64)
|
||||
for rows.Next() {
|
||||
var kind string
|
||||
var size int64
|
||||
if err := rows.Scan(&kind, &size); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[kind] = size
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
|
||||
func (r *AdminReadRepo) DatabaseSizeBytes(ctx context.Context) (int64, error) {
|
||||
var size int64
|
||||
if err := conn(ctx, r.pool).QueryRow(ctx, "SELECT pg_database_size(current_database())").Scan(&size); err != nil {
|
||||
return 0, fmt.Errorf("database size: %w", err)
|
||||
}
|
||||
return size, 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 {
|
||||
@@ -133,6 +134,33 @@ func TestClaimReductionIsAtomic(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
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 := NewAdminReadRepo(pool).ListJobsPaginated(ctx, "", 20, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("list admin jobs: %v", err)
|
||||
}
|
||||
for _, item := range listed {
|
||||
if item.ID != job.ID {
|
||||
continue
|
||||
}
|
||||
if item.Status != domain.JobReducing || item.ReducerStartedAt == nil {
|
||||
t.Fatalf("admin reducer projection = %+v", item)
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatalf("seeded job %s is missing from the admin 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) {
|
||||
@@ -379,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),
|
||||
@@ -629,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
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ package postgres
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
sq "github.com/Masterminds/squirrel"
|
||||
@@ -28,14 +29,15 @@ var _ usecase.JobRepository = (*JobRepo)(nil)
|
||||
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
|
||||
@@ -59,7 +61,8 @@ func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
|
||||
)
|
||||
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).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.InputArtifactID, &j.ResultArtifactID, &j.ErrorCode, &j.ErrorMessage, &j.ReducerStartedAt,
|
||||
&j.OwnerID)
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, domain.ErrJobNotFound
|
||||
}
|
||||
@@ -151,3 +154,46 @@ func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID,
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ListCompletedBefore returns jobs whose completion timestamp is older than
|
||||
// the cutoff (completed and failed both count as finished).
|
||||
func (r *JobRepo) ListCompletedBefore(ctx context.Context, cutoff time.Time) ([]domain.Job, error) {
|
||||
sql, args, err := psql.Select(jobColumns...).From("jobs").
|
||||
Where(sq.NotEq{"completed_at": nil}).
|
||||
Where(sq.Lt{"completed_at": cutoff}).
|
||||
OrderBy("completed_at 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 completed jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
var jobs []domain.Job
|
||||
for rows.Next() {
|
||||
var j domain.Job
|
||||
var status string
|
||||
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
|
||||
}
|
||||
j.Status = domain.JobStatus(status)
|
||||
jobs = append(jobs, j)
|
||||
}
|
||||
return jobs, rows.Err()
|
||||
}
|
||||
|
||||
// Delete removes the job row; tasks, artifacts and task_results cascade.
|
||||
func (r *JobRepo) Delete(ctx context.Context, id uuid.UUID) error {
|
||||
sql, args, err := psql.Delete("jobs").Where(sq.Eq{"id": id}).ToSql()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = conn(ctx, r.pool).Exec(ctx, sql, args...)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -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,62 @@
|
||||
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"
|
||||
case 14:
|
||||
return "0014_workload_settings.up.sql"
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
@@ -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,5 @@
|
||||
BEGIN;
|
||||
|
||||
DROP TABLE IF EXISTS workload_settings;
|
||||
|
||||
COMMIT;
|
||||
@@ -0,0 +1,11 @@
|
||||
BEGIN;
|
||||
|
||||
-- Per-workload enable/disable. Absence of a row means "enabled" (the catalog
|
||||
-- default); a row only exists once an admin flipped a workload off or back on.
|
||||
CREATE TABLE workload_settings (
|
||||
workload text NOT NULL PRIMARY KEY,
|
||||
enabled boolean NOT NULL,
|
||||
updated_at timestamptz NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
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
|
||||
|
||||
@@ -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,36 +24,6 @@ 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) {
|
||||
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()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("list jobs: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
jobs := make([]domain.Job, 0)
|
||||
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 {
|
||||
return nil, err
|
||||
}
|
||||
if inputURI != nil {
|
||||
j.InputURI = *inputURI
|
||||
}
|
||||
j.Status = domain.JobStatus(status)
|
||||
jobs = append(jobs, j)
|
||||
}
|
||||
return jobs, rows.Err()
|
||||
}
|
||||
|
||||
func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error) {
|
||||
sql, args, err := psql.Select(taskColumns...).From("tasks").Where(sq.Eq{"job_id": jobID}).OrderBy("chunk_index ASC").ToSql()
|
||||
if err != nil {
|
||||
@@ -75,36 +45,12 @@ 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
|
||||
}
|
||||
sql, args, err := psql.Select(workerColumns...).From("workers").OrderBy("last_heartbeat_at DESC", "id DESC").Limit(uint64(limit)).ToSql()
|
||||
sql, args, err := psql.Select(workerColumns...).From("workers").
|
||||
OrderBy("last_heartbeat_at DESC", "id DESC").Limit(uint64(limit)).ToSql()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user