Compare commits

...
Author SHA1 Message Date
Emil d0aeb7fc95 Upgrade pipeline observability UI
coordinator / test (push) Waiting to run
2026-07-24 15:13:59 +03:00
Emil a055473706 Add final result reduction
coordinator / test (push) Waiting to run
2026-07-24 14:58:18 +03:00
Emil 6e67daa9eb Merge distributed similarity search 2026-07-24 14:43:57 +03:00
Emil 0f3a2d92d8 Add distributed similarity search
python / test (push) Waiting to run
2026-07-24 14:38:11 +03:00
Emil 0bef7604fd Merge distributed workload protocol 2026-07-24 14:26:58 +03:00
Emil 6ef92908a1 Add distributed workload protocol
python / test (push) Waiting to run
2026-07-24 14:25:38 +03:00
Emil f953112cfd Merge distributed pipeline hardening 2026-07-24 14:16:42 +03:00
Emil 19cbf7f113 Harden distributed pipeline 2026-07-24 14:16:42 +03:00
Emil 9ec8f50313 Fix worker interruption handling 2026-07-24 13:19:40 +03:00
Emil 08f5478a66 Define distributed workload protocol 2026-07-24 13:09:25 +03:00
Emil bde6cdb4ba Improve worker lifecycle controls 2026-07-24 13:08:10 +03:00
Emil 43ceec1f77 Add two-worker smoke test 2026-07-24 13:05:34 +03:00
Emil f5b16b057f Add UI development task briefs 2026-07-24 00:13:55 +03:00
Emil f8de0b2b9d Merge feat/web-interface 2026-07-23 23:43:57 +03:00
Emil 7547a30bde Add job cancellation and dataset row limit
coordinator / test (push) Waiting to run
2026-07-23 23:14:25 +03:00
Emil 6bac7dad3c Clarify worker failures in operator UI 2026-07-23 22:59:51 +03:00
Emil c7956c4683 Fix relative worker work directory 2026-07-23 22:48:27 +03:00
Emil d648beede2 Use English operator UI copy 2026-07-23 22:45:36 +03:00
Emil ac9b921401 Clarify operator UI workflow 2026-07-23 22:41:14 +03:00
Emil 5be87ad762 Fix UI dataset upload field order 2026-07-23 22:20:18 +03:00
Emil e83e0b5e1f Add local operator web interface 2026-07-23 22:08:58 +03:00
74 changed files with 5569 additions and 298 deletions
+28
View File
@@ -0,0 +1,28 @@
name: python
on:
push:
paths:
- "scimesh/**"
- "tests/**"
- "pyproject.toml"
- ".github/workflows/python.yml"
pull_request:
paths:
- "scimesh/**"
- "tests/**"
- "pyproject.toml"
- ".github/workflows/python.yml"
jobs:
test:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: "3.11"
cache: pip
- run: python -m pip install --upgrade pip
- run: python -m pip install -e '.[dev]'
- run: pytest -q
+4
View File
@@ -11,3 +11,7 @@ results/
*_similarities.csv
test_results.csv
test_structures/
# Local coordinator-worker execution state
worker-data*/
scimesh-worker-data/
+5 -4
View File
@@ -1,10 +1,11 @@
# SciMesh
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
currently runs exact similarity search and sparse similarity-graph construction
locally in one Python process; it creates no dense similarity matrix. A Python
Worker client and the planned Go/PostgreSQL coordinator contract are tracked in
the repository, but distributed execution is not available yet; see
runs exact similarity search and sparse similarity-graph construction locally in
one Python process; it creates no dense similarity matrix. The Go/PostgreSQL
coordinator and 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).
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`.
+18 -14
View File
@@ -1,7 +1,7 @@
# SciMesh Status
**Updated:** 2026-07-23
**Branch baseline:** `main` at `b4a89dd` (coordinator merge)
**Updated:** 2026-07-24
**Branch baseline:** `main` at `6e67daa` (distributed similarity-search)
## Current state
@@ -18,8 +18,10 @@ the reference behaviour for future distributed execution:
The Go coordinator and its PostgreSQL-backed task lifecycle are implemented:
registration, atomic claiming, lease renewal, artifact storage, dataset
chunking, result/failure reporting, and job progress. The Python worker now
uses the live coordinator contract; its HTTP path was exercised against a real
Docker PostgreSQL stack on 2026-07-23.
uses the live coordinator contract. Completed similarity-search shard results
are reduced once into a checksum-protected final CSV, which is downloadable
through the coordinator. The full Go checks (including a fresh migration and
real PostgreSQL smoke test) passed on 2026-07-24.
## Milestone tracker
@@ -32,26 +34,28 @@ Docker PostgreSQL stack on 2026-07-23.
| CTX-04 Worker registry and HTTP API | Implemented | Registration, claim, heartbeat, result, failure, and status endpoints. |
| CTX-05 Artifact storage | Implemented | Coordinator-owned inputs/results, checksum verification, and upload flow. |
| CTX-06 Python Worker live-contract alignment | Implemented | Worker completed a real uploaded shard via HTTP on 2026-07-23. |
| CTX-07 Distributed workload protocol | Not started | Depends on artifact and Worker contracts. |
| CTX-08 Distributed similarity-search | Not started | Local reference exists. |
| CTX-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. |
| CTX-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 | Not started | Deferred until API and reducer work. |
| CTX-11 Dashboard/operator view | Implemented | Protected live control room: recent-run/worker overview, real pipeline-stage visualization, shard attempts and safe failures, validated similarity-search upload, coordinator artifacts, final-result download, and bounded polling. |
| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
## Next recommended assignment
Assign **CTX-07** to the workload role: define distributed job planning and
reduction boundaries before implementing distributed search or graph execution.
Assign **CTX-10** to the distributed-science role: implement deterministic
block-pair planning and reduction for `similarity-graph`.
## Known constraints
- Planner/reducer semantics are not implemented; use the local `scimesh` CLI
for complete workload results.
- The worker/coordinator flow currently accepts both underscore API workload
names and hyphenated CLI names while the contract is consolidated.
- A real-stack worker test uses a small `query_smiles` shard. Resolving a
`query_id` once and sharing it across shards belongs to CTX-07.
- 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.
## Update rule
+4
View File
@@ -6,6 +6,10 @@ DATABASE_URL=postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable
# Shared bearer token every worker must present. Leave empty to disable auth (dev only).
WORKER_AUTH_TOKEN=change-me
# Optional local operator UI. Use a separate value; never reuse the worker token.
# When empty, /ui is disabled.
UI_AUTH_TOKEN=
# Logging. LOG_LEVEL: debug|info|warn|error. LOG_FILE empty = stdout only;
# set a path to also write a size-rotated file (kept across restarts).
LOG_LEVEL=info
+22
View File
@@ -68,6 +68,28 @@ make logs # follow the coordinator
make down # stop (add down-clean to drop the DB volume)
```
To enable the local operator UI, set a separate credential before starting:
```sh
UI_AUTH_TOKEN='local-ui-secret' make up
# Open http://localhost:8080/ui and use any username with this value as password.
```
The UI is disabled by default and never accepts the worker bearer token.
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. Download the
`final_result` artifact only after the job reaches **Completed**; shard partial
CSVs remain available as diagnostics.
`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
coordinator start — so it never queries a database that has no tables.
+10 -5
View File
@@ -65,26 +65,31 @@ func run() error {
jobRepo = postgres.NewJobRepo(pool)
workerRepo = postgres.NewWorkerRepo(pool)
artifactRepo = postgres.NewArtifactRepo(pool)
uiReadRepo = postgres.NewUIReadRepo(pool)
)
useCases := httptransport.UseCases{
RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk),
CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk),
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk),
ClaimTask: usecase.NewClaimTask(taskRepo, clk, cfg.LeaseDuration),
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk, cfg.DefaultMaxAttempts),
ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration),
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),
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, clk),
CancelJob: usecase.NewCancelJob(jobRepo, taskRepo, tx, clk),
UploadArtifact: usecase.NewUploadArtifact(taskRepo, 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),
}
// Background reapers are tracked so shutdown can wait for them. Without this
// the process would exit mid-UPDATE, and the deferred pool.Close() would pull
// connections out from under them.
expireLeases := usecase.NewExpireLeases(taskRepo, clk)
expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk)
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
var wg sync.WaitGroup
@@ -105,7 +110,7 @@ func run() error {
// pool.Ping backs /health: readiness means the database answers, not just
// that the process is alive.
api := httptransport.NewServer(useCases, log, cfg.RequestTimeout, cfg.HeartbeatInterval, cfg.MaxUploadBytes, pool.Ping)
err = infra.RunServer(ctx, log, cfg.Addr, api.Handler(cfg.Token))
err = infra.RunServer(ctx, log, cfg.Addr, api.Handler(cfg.Token, cfg.UIToken))
// Shutdown order matters, and defers alone cannot express it (they run
// LIFO, so the deferred stop() would fire *after* the wait below).
+2
View File
@@ -49,6 +49,8 @@ services:
# Host is the service name: compose resolves it on the project network.
DATABASE_URL: postgres://${POSTGRES_USER:-scimesh}:${POSTGRES_PASSWORD:-scimesh}@postgres:5432/${POSTGRES_DB:-scimesh}?sslmode=disable
WORKER_AUTH_TOKEN: ${WORKER_AUTH_TOKEN:-dev-token}
# Empty disables /ui. Set this separately from the worker token.
UI_AUTH_TOKEN: ${UI_AUTH_TOKEN:-}
DB_MAX_CONNS: "10"
REQUEST_TIMEOUT: "15s"
LEASE_DURATION: "2m"
+57
View File
@@ -7,12 +7,19 @@ import (
"bytes"
"fmt"
"io"
"strings"
)
// ErrNoRows is returned when the input has a header but no data rows: a job with
// zero tasks could never complete, so it is rejected at the source.
var ErrNoRows = fmt.Errorf("input has no data rows")
// maxShardBytes bounds the coordinator memory used by one in-progress shard.
// The uploaded file may be much larger: it is first stored on disk, then split
// in small bounded pieces. Operators can lower rowsPerShard when this limit is
// reached rather than exhausting the coordinator process.
const maxShardBytes = 64 << 20 // 64 MiB
// SplitTSV reads a header-plus-rows text stream and cuts it into shards of at
// most rowsPerShard data rows. Every shard repeats the header, so a worker can
// parse its shard in isolation. emit is called once per shard, in order, with a
@@ -26,9 +33,31 @@ var ErrNoRows = fmt.Errorf("input has no data rows")
// Only one shard is buffered at a time, so memory is bounded by shard size (a
// worker-sized slice of the data), not by the size of the whole dataset.
func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error {
return splitTSVLimit(r, rowsPerShard, 0, nil, emit)
}
// SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows.
// A maxRows value of zero means unlimited. This lets an operator make a small,
// representative pipeline check without materialising a second dataset file.
func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
return splitTSVLimit(r, rowsPerShard, maxRows, nil, emit)
}
// SplitChEMBLTSVLimit is the coordinator's scientific-upload splitter. It
// validates the two columns every local SciMesh workload requires before any
// shard task is persisted, while generic SplitTSV remains reusable for future
// non-chemistry workloads.
func SplitChEMBLTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error {
return splitTSVLimit(r, rowsPerShard, maxRows, validateChEMBLHeader, emit)
}
func splitTSVLimit(r io.Reader, rowsPerShard, maxRows int, validateHeader func([]byte) error, emit func(index int, shard io.Reader) error) error {
if rowsPerShard <= 0 {
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
}
if maxRows < 0 {
return fmt.Errorf("maxRows must be non-negative, got %d", maxRows)
}
sc := bufio.NewScanner(r)
// Allow long lines: a SMILES row can be far wider than bufio's 64 KB default.
@@ -41,6 +70,11 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
return ErrNoRows // completely empty input
}
header := append([]byte(nil), sc.Bytes()...)
if validateHeader != nil {
if err := validateHeader(header); err != nil {
return err
}
}
var (
buf bytes.Buffer
@@ -61,9 +95,15 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
for sc.Scan() {
if rows == 0 {
if len(header)+1 > maxShardBytes {
return fmt.Errorf("TSV header exceeds maximum shard size of %d bytes", maxShardBytes)
}
buf.Write(header)
buf.WriteByte('\n')
}
if buf.Len()+len(sc.Bytes())+1 > maxShardBytes {
return fmt.Errorf("shard exceeds maximum size of %d bytes; lower rowsPerShard", maxShardBytes)
}
buf.Write(sc.Bytes())
buf.WriteByte('\n')
rows++
@@ -73,6 +113,9 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
return err
}
}
if maxRows > 0 && index*rowsPerShard+rows == maxRows {
break
}
}
if err := sc.Err(); err != nil {
return fmt.Errorf("read rows: %w", err)
@@ -90,3 +133,17 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade
}
return nil
}
func validateChEMBLHeader(header []byte) error {
seen := make(map[string]struct{})
for _, field := range strings.Split(strings.TrimPrefix(string(header), "\ufeff"), "\t") {
seen[field] = struct{}{}
}
if _, ok := seen["chembl_id"]; !ok {
return fmt.Errorf("TSV is missing required column chembl_id")
}
if _, ok := seen["canonical_smiles"]; !ok {
return fmt.Errorf("TSV is missing required column canonical_smiles")
}
return nil
}
+24
View File
@@ -103,6 +103,30 @@ func TestSplitSingleShardWhenSizeExceedsRows(t *testing.T) {
}
}
func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) {
input := "h\nr1\nr2\nr3\nr4\nr5\n"
var shards []string
err := SplitTSVLimit(strings.NewReader(input), 2, 3, func(_ int, shard io.Reader) error {
b, _ := io.ReadAll(shard)
shards = append(shards, string(b))
return nil
})
if err != nil {
t.Fatal(err)
}
if got, want := strings.Join(shards, ""), "h\nr1\nr2\nh\nr3\n"; got != want {
t.Errorf("limited shards = %q, want %q", got, want)
}
}
func TestChEMBLSplitRejectsMissingRequiredColumns(t *testing.T) {
err := SplitChEMBLTSVLimit(strings.NewReader("id\tsmiles\nA\tCC\n"), 1, 0,
func(int, io.Reader) error { return nil })
if err == nil || !strings.Contains(err.Error(), "chembl_id") {
t.Errorf("err = %v, want missing-column error", err)
}
}
// The scanned bytes are reused by bufio; the shard buffer must copy them, or a
// later row would corrupt an earlier one. This guards that copy.
func TestSplitDoesNotAliasScannerBuffer(t *testing.T) {
+10 -9
View File
@@ -8,13 +8,14 @@ import "errors"
//
// Always compare with errors.Is — outer layers may wrap these with %w.
var (
ErrJobNotFound = errors.New("job not found")
ErrTaskNotFound = errors.New("task not found")
ErrWorkerNotFound = errors.New("worker not found")
ErrArtifactNotFound = errors.New("artifact not found")
ErrLeaseConflict = errors.New("task leased to another worker")
ErrStaleAttempt = errors.New("attempt does not match lease")
ErrResultConflict = errors.New("different result already recorded")
ErrInvalidInput = errors.New("invalid input")
ErrTaskNotLeased = errors.New("task is not currently leased")
ErrJobNotFound = errors.New("job not found")
ErrTaskNotFound = errors.New("task not found")
ErrWorkerNotFound = errors.New("worker not found")
ErrArtifactNotFound = errors.New("artifact not found")
ErrJobNotCancellable = errors.New("job cannot be cancelled")
ErrLeaseConflict = errors.New("task leased to another worker")
ErrStaleAttempt = errors.New("attempt does not match lease")
ErrResultConflict = errors.New("different result already recorded")
ErrInvalidInput = errors.New("invalid input")
ErrTaskNotLeased = errors.New("task is not currently leased")
)
+24 -14
View File
@@ -11,6 +11,7 @@ type JobStatus string
const (
JobPending JobStatus = "pending"
JobRunning JobStatus = "running"
JobReducing JobStatus = "reducing"
JobCompleted JobStatus = "completed"
JobFailed JobStatus = "failed"
JobCancelled JobStatus = "cancelled"
@@ -18,14 +19,18 @@ const (
// Job is one user submission that fans out into one or more tasks.
type Job struct {
ID uuid.UUID
Workload string
InputURI string // external input URI; empty for uploaded datasets
InputArtifactID *uuid.UUID // uploaded input artifact; nil for URI submissions
Parameters map[string]any
Status JobStatus
CreatedAt time.Time
CompletedAt *time.Time
ID uuid.UUID
Workload string
InputURI string // external input URI; empty for uploaded datasets
InputArtifactID *uuid.UUID // uploaded input artifact; nil for URI submissions
ResultArtifactID *uuid.UUID
Parameters map[string]any
Status JobStatus
CreatedAt time.Time
CompletedAt *time.Time
ReducerStartedAt *time.Time
ErrorCode *string
ErrorMessage *string
}
// NewUploadedJob builds a job whose input was uploaded to the coordinator. The
@@ -99,18 +104,23 @@ func NewJobWithTasks(workload, inputURI string, params map[string]any,
// JobProgress is the aggregate view of a job and the state of its tasks.
type JobProgress struct {
Job Job
Total int
Pending int
Leased int
Done int
Failed int
Job Job
Total int
Pending int
Leased int
Done int
Failed int
Cancelled int
}
// DeriveStatus computes what the job's status should be from its task counts,
// so the rule lives here rather than in a SQL trigger or a handler.
func (p JobProgress) DeriveStatus() JobStatus {
switch {
case p.Job.Status == JobCancelled:
return JobCancelled
case p.Job.Status == JobReducing:
return JobReducing
case p.Total == 0:
return JobPending
case p.Done == p.Total:
+1
View File
@@ -81,6 +81,7 @@ func TestDeriveStatus(t *testing.T) {
{"all done", JobProgress{Total: 3, Done: 3}, JobCompleted},
{"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},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
+18
View File
@@ -259,6 +259,24 @@ func (t *Task) ExpireLease(now time.Time) {
t.CompletedAt = &now
}
// Cancel prevents any further worker transition for a task that has not
// reached a terminal result. A cancelled lease deliberately becomes invalid:
// a worker still running locally must not upload or complete after its job was
// stopped by the operator.
func (t *Task) Cancel(now time.Time) bool {
if t.Status == TaskCompleted || t.Status == TaskFailed || t.Status == TaskCancelled {
return false
}
t.Status = TaskCancelled
t.LeaseOwner = nil
t.LeaseExpiresAt = nil
t.ErrorCode = nil
t.ErrorMessage = nil
t.CompletedAt = &now
t.Version++
return true
}
// ClaimedTask is the worker-facing projection of a leased task. Input is either
// an external URI or a coordinator-stored shard (InputArtifactID set); the
// transport turns the latter into a coordinator download URL.
+17
View File
@@ -168,6 +168,23 @@ func TestExpireLeaseIgnoresUnleasedTasks(t *testing.T) {
}
}
func TestCancelInvalidatesLeaseButPreservesTerminalTask(t *testing.T) {
task := leasedTask(1, 3)
if !task.Cancel(testNow) {
t.Fatal("leased task should be cancelled")
}
if task.Status != TaskCancelled || task.LeaseOwner != nil || task.LeaseExpiresAt != nil {
t.Errorf("cancelled task = %+v", task)
}
if task.Cancel(testLater) {
t.Error("cancelled task must not be changed twice")
}
completed := &Task{Status: TaskCompleted}
if completed.Cancel(testNow) {
t.Error("completed task must remain terminal")
}
}
func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) {
task := leasedTask(1, 3)
until := testLater.Add(time.Hour)
+9
View File
@@ -24,6 +24,8 @@ type Config struct {
DatabaseURL string
// Shared bearer token workers must present. Empty disables auth (dev only).
Token string
// Local operator UI credential. Empty disables the embedded UI entirely.
UIToken string
// Minimum log level: debug, info, warn, error.
LogLevel string
@@ -77,6 +79,7 @@ func LoadConfig() (Config, error) {
// 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"),
@@ -94,6 +97,9 @@ func LoadConfig() (Config, error) {
if cfg.DatabaseURL == "" {
return Config{}, fmt.Errorf("DATABASE_URL is required")
}
if cfg.UIToken != "" && cfg.Token != "" && cfg.UIToken == cfg.Token {
return Config{}, fmt.Errorf("UI_AUTH_TOKEN must differ from the worker auth token")
}
var err error
if cfg.DBMaxConns, err = getEnvInt32("DB_MAX_CONNS", cfg.DBMaxConns); err != nil {
@@ -123,6 +129,9 @@ func LoadConfig() (Config, error) {
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
return Config{}, err
}
if cfg.DefaultMaxAttempts < 1 {
return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive")
}
return cfg, nil
}
+45
View File
@@ -0,0 +1,45 @@
package infra
import (
"path/filepath"
"strings"
"testing"
)
func TestLoadConfigRejectsSharedUIAndWorkerToken(t *testing.T) {
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
t.Setenv("DATABASE_URL", "postgres://test")
t.Setenv("COORDINATOR_TOKEN", "shared-secret")
t.Setenv("UI_AUTH_TOKEN", "shared-secret")
_, err := LoadConfig()
if err == nil || !strings.Contains(err.Error(), "must differ") {
t.Fatalf("LoadConfig error = %v, want distinct-token error", err)
}
}
func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) {
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
t.Setenv("DATABASE_URL", "postgres://test")
t.Setenv("COORDINATOR_TOKEN", "worker-secret")
t.Setenv("UI_AUTH_TOKEN", "ui-secret")
cfg, err := LoadConfig()
if err != nil {
t.Fatalf("LoadConfig: %v", err)
}
if cfg.Token != "worker-secret" || cfg.UIToken != "ui-secret" {
t.Fatalf("unexpected tokens: %+v", cfg)
}
}
func TestLoadConfigRejectsNonPositiveDefaultMaxAttempts(t *testing.T) {
t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env"))
t.Setenv("DATABASE_URL", "postgres://test")
t.Setenv("DEFAULT_MAX_ATTEMPTS", "0")
_, err := LoadConfig()
if err == nil || !strings.Contains(err.Error(), "DEFAULT_MAX_ATTEMPTS") {
t.Fatalf("LoadConfig error = %v, want default-attempt validation", err)
}
}
+82 -7
View File
@@ -148,17 +148,30 @@ func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[doma
return counts, nil
}
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time) (int64, error) {
r.mu.Lock()
defer r.mu.Unlock()
var n int64
for _, t := range r.tasks {
if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
t.ExpireLease(now)
n++
var cancelled int64
for _, task := range r.tasks {
if task.JobID == jobID && task.Cancel(now) {
cancelled++
}
}
return n, nil
return cancelled, nil
}
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
r.mu.Lock()
defer r.mu.Unlock()
affected := make([]uuid.UUID, 0)
for _, t := range r.tasks {
if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) &&
t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
t.ExpireLease(now)
affected = append(affected, t.JobID)
}
}
return affected, nil
}
// --- JobRepo -------------------------------------------------------------
@@ -203,6 +216,51 @@ func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID, status domain.
return nil
}
func (r *JobRepo) ClaimReduction(_ context.Context, id uuid.UUID, startedAt time.Time) (bool, error) {
r.mu.Lock()
defer r.mu.Unlock()
j, ok := r.jobs[id]
if !ok {
return false, domain.ErrJobNotFound
}
if j.Status != domain.JobReducing || j.ReducerStartedAt != nil {
return false, nil
}
j.ReducerStartedAt = &startedAt
return true, nil
}
func (r *JobRepo) CompleteWithResult(_ context.Context, id, resultArtifactID uuid.UUID, completedAt time.Time) error {
r.mu.Lock()
defer r.mu.Unlock()
j, ok := r.jobs[id]
if !ok {
return domain.ErrJobNotFound
}
j.ResultArtifactID = &resultArtifactID
j.Status = domain.JobCompleted
j.CompletedAt = &completedAt
j.ReducerStartedAt = nil
j.ErrorCode = nil
j.ErrorMessage = nil
return nil
}
func (r *JobRepo) FailReduction(_ context.Context, id uuid.UUID, code, message string, completedAt time.Time) error {
r.mu.Lock()
defer r.mu.Unlock()
j, ok := r.jobs[id]
if !ok {
return domain.ErrJobNotFound
}
j.Status = domain.JobFailed
j.CompletedAt = &completedAt
j.ReducerStartedAt = nil
j.ErrorCode = &code
j.ErrorMessage = &message
return nil
}
// --- WorkerRepo ----------------------------------------------------------
type WorkerRepo struct {
@@ -286,6 +344,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
return &cp, nil
}
func (r *ArtifactRepo) FindPartialResult(_ context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
r.mu.Lock()
defer r.mu.Unlock()
for _, a := range r.arts {
if a.TaskID != nil && *a.TaskID == taskID && a.Kind == domain.ArtifactPartialResult &&
a.Attempt != nil && *a.Attempt == attempt {
return cloneArtifact(a), nil
}
}
return nil, nil
}
func cloneArtifact(a *domain.Artifact) *domain.Artifact {
cp := *a
return &cp
}
// --- BlobStore -----------------------------------------------------------
type BlobStore struct {
+114
View File
@@ -0,0 +1,114 @@
package memstore
import (
"context"
"sort"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
"github.com/emil28092005/SciMesh/coordinator/internal/usecase"
)
// UIReadRepo is the in-memory read projection used by HTTP/UI tests.
type UIReadRepo struct {
jobs *JobRepo
tasks *TaskRepo
workers *WorkerRepo
artifacts *ArtifactRepo
}
func NewUIReadRepo(j *JobRepo, t *TaskRepo, w *WorkerRepo, a *ArtifactRepo) *UIReadRepo {
return &UIReadRepo{j, t, w, a}
}
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) {
if limit < 1 || limit > 100 {
return nil, domain.ErrInvalidInput
}
r.jobs.mu.Lock()
defer r.jobs.mu.Unlock()
out := make([]domain.Job, 0, len(r.jobs.jobs))
for _, job := range r.jobs.jobs {
out = append(out, *job)
}
sort.Slice(out, func(i, j int) bool {
if out[i].CreatedAt.Equal(out[j].CreatedAt) {
return out[i].ID.String() > out[j].ID.String()
}
return out[i].CreatedAt.After(out[j].CreatedAt)
})
if len(out) > limit {
out = out[:limit]
}
return out, nil
}
func (r *UIReadRepo) ListTasksByJob(_ context.Context, jobID uuid.UUID) ([]domain.Task, error) {
r.tasks.mu.Lock()
defer r.tasks.mu.Unlock()
out := []domain.Task{}
for _, task := range r.tasks.tasks {
if task.JobID == jobID {
out = append(out, *clone(task))
}
}
sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
return out, nil
}
func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) {
out := make(map[uuid.UUID][]domain.Task, len(jobIDs))
for _, id := range jobIDs {
tasks, err := r.ListTasksByJob(ctx, id)
if err != nil {
return nil, err
}
out[id] = tasks
}
return out, nil
}
func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
if limit < 1 || limit > 100 {
return nil, domain.ErrInvalidInput
}
r.workers.mu.Lock()
defer r.workers.mu.Unlock()
out := []domain.Worker{}
for _, worker := range r.workers.workers {
copy := *worker
copy.Capabilities = append([]string(nil), worker.Capabilities...)
out = append(out, copy)
}
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()
defer r.artifacts.mu.Unlock()
out := []domain.Artifact{}
for _, artifact := range r.artifacts.arts {
if artifact.JobID == jobID {
out = append(out, *artifact)
}
}
sort.Slice(out, func(i, j int) bool {
if out[i].CreatedAt.Equal(out[j].CreatedAt) {
return out[i].ID.String() < out[j].ID.String()
}
return out[i].CreatedAt.Before(out[j].CreatedAt)
})
return out, nil
}
@@ -0,0 +1,195 @@
// Package reducer contains deterministic, coordinator-side result reductions.
package reducer
import (
"bytes"
"encoding/csv"
"errors"
"fmt"
"io"
"math"
"sort"
"strconv"
)
var searchHeader = []string{"rank", "chembl_id", "canonical_smiles", "similarity"}
type similarityMatch struct {
similarity float64
id string
smiles string
}
// ReduceSimilaritySearch streams worker-local top-k CSVs into the exact global
// top-k. Each partial is validated before it can affect the final artifact.
func ReduceSimilaritySearch(partials []io.Reader, parameters map[string]any) ([]byte, error) {
topK, err := positiveInt(parameters["top_k"], 20)
if err != nil {
return nil, err
}
direction, err := thresholdDirection(parameters["threshold_direction"])
if err != nil {
return nil, err
}
h := &matchHeap{direction: direction}
for _, partial := range partials {
if err := readPartial(partial, direction, func(match similarityMatch) {
if len(h.items) < topK {
heapPush(h, match)
return
}
if better(match, h.items[0], direction) {
h.items[0] = match
heapDown(h, 0)
}
}); err != nil {
return nil, err
}
}
matches := append([]similarityMatch(nil), h.items...)
sort.Slice(matches, func(i, j int) bool { return better(matches[i], matches[j], direction) })
var out bytes.Buffer
writer := csv.NewWriter(&out)
if err := writer.Write(searchHeader); err != nil {
return nil, err
}
for index, match := range matches {
if err := writer.Write([]string{
strconv.Itoa(index + 1), match.id, match.smiles, fmt.Sprintf("%.6f", match.similarity),
}); err != nil {
return nil, err
}
}
writer.Flush()
if err := writer.Error(); err != nil {
return nil, err
}
return out.Bytes(), nil
}
func readPartial(input io.Reader, direction string, consume func(similarityMatch)) error {
reader := csv.NewReader(input)
header, err := reader.Read()
if err != nil {
return fmt.Errorf("read partial header: %w", err)
}
if !equalStrings(header, searchHeader) {
return fmt.Errorf("partial result has an invalid CSV header")
}
var previous *similarityMatch
for rank := 1; ; rank++ {
row, err := reader.Read()
if errors.Is(err, io.EOF) {
return nil
}
if err != nil {
return fmt.Errorf("read partial row: %w", err)
}
if len(row) != len(searchHeader) || row[0] != strconv.Itoa(rank) {
return fmt.Errorf("partial result has an invalid rank")
}
score, err := strconv.ParseFloat(row[3], 64)
if err != nil || math.IsNaN(score) || math.IsInf(score, 0) || score < 0 || score > 1 {
return fmt.Errorf("partial result has an invalid similarity")
}
match := similarityMatch{similarity: score, id: row[1], smiles: row[2]}
if previous != nil && better(match, *previous, direction) {
return fmt.Errorf("partial result is not sorted deterministically")
}
previous = &match
consume(match)
}
}
func positiveInt(value any, fallback int) (int, error) {
if value == nil {
return fallback, nil
}
switch n := value.(type) {
case int:
if n > 0 {
return n, nil
}
case int64:
if n > 0 && n <= math.MaxInt {
return int(n), nil
}
case float64:
if n > 0 && n == math.Trunc(n) && n <= math.MaxInt {
return int(n), nil
}
}
return 0, fmt.Errorf("top_k must be a positive integer")
}
func thresholdDirection(value any) (string, error) {
if value == nil {
return "greater", nil
}
direction, ok := value.(string)
if !ok || (direction != "greater" && direction != "less") {
return "", fmt.Errorf("threshold_direction must be greater or less")
}
return direction, nil
}
func better(left, right similarityMatch, direction string) bool {
if left.similarity != right.similarity {
if direction == "less" {
return left.similarity < right.similarity
}
return left.similarity > right.similarity
}
if left.id != right.id {
return left.id < right.id
}
return left.smiles < right.smiles
}
func equalStrings(left, right []string) bool {
if len(left) != len(right) {
return false
}
for index := range left {
if left[index] != right[index] {
return false
}
}
return true
}
// matchHeap keeps the worst retained match at index zero.
type matchHeap struct {
items []similarityMatch
direction string
}
func heapPush(h *matchHeap, value similarityMatch) {
h.items = append(h.items, value)
for child := len(h.items) - 1; child > 0; {
parent := (child - 1) / 2
if !better(h.items[parent], h.items[child], h.direction) {
break
}
h.items[parent], h.items[child] = h.items[child], h.items[parent]
child = parent
}
}
func heapDown(h *matchHeap, parent int) {
for {
child := parent*2 + 1
if child >= len(h.items) {
return
}
if right := child + 1; right < len(h.items) && better(h.items[child], h.items[right], h.direction) {
child = right
}
if !better(h.items[parent], h.items[child], h.direction) {
return
}
h.items[parent], h.items[child] = h.items[child], h.items[parent]
parent = child
}
}
@@ -0,0 +1,41 @@
package reducer
import (
"io"
"strings"
"testing"
)
func TestReduceSimilaritySearchKeepsExactCrossShardRanking(t *testing.T) {
first := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,B,CCC,0.50000048\n2,C,CCCC,0.1\n")
second := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.50000049\n")
output, err := ReduceSimilaritySearch([]io.Reader{first, second}, map[string]any{"top_k": 2})
if err != nil {
t.Fatal(err)
}
want := "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.500000\n2,B,CCC,0.500000\n"
if string(output) != want {
t.Fatalf("output = %q, want %q", output, want)
}
}
func TestReduceSimilaritySearchSupportsLeastSimilarDirection(t *testing.T) {
partial := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.1\n2,B,CCC,0.8\n")
output, err := ReduceSimilaritySearch([]io.Reader{partial}, map[string]any{
"top_k": 1, "threshold_direction": "less",
})
if err != nil {
t.Fatal(err)
}
if got, want := string(output), "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.100000\n"; got != want {
t.Fatalf("output = %q, want %q", got, want)
}
}
func TestReduceSimilaritySearchRejectsMalformedPartial(t *testing.T) {
partial := strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n2,A,CC,0.1\n")
if _, err := ReduceSimilaritySearch([]io.Reader{partial}, nil); err == nil {
t.Fatal("expected malformed rank error")
}
}
@@ -70,3 +70,33 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
a.Kind = domain.ArtifactKind(kind)
return &a, nil
}
func (r *ArtifactRepo) FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) {
sql, args, err := psql.Select(artifactColumns...).
From("artifacts").
Where(sq.Eq{
"task_id": taskID,
"attempt": attempt,
"kind": string(domain.ArtifactPartialResult),
}).
ToSql()
if err != nil {
return nil, err
}
var (
a domain.Artifact
kind string
)
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey,
&a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt)
if errors.Is(err, pgx.ErrNoRows) {
return nil, nil
}
if err != nil {
return nil, fmt.Errorf("find partial result: %w", err)
}
a.Kind = domain.ArtifactKind(kind)
return &a, nil
}
@@ -90,6 +90,76 @@ func TestCreateJobPersistsEveryTask(t *testing.T) {
}
}
func TestClaimReductionIsAtomic(t *testing.T) {
pool := testPool(t)
job, _ := seedJob(t, pool, 1)
repo := NewJobRepo(pool)
ctx := context.Background()
if err := repo.UpdateStatus(ctx, job.ID, domain.JobReducing, nil); err != nil {
t.Fatal(err)
}
var (
wg sync.WaitGroup
mu sync.Mutex
claimed int
)
for range 8 {
wg.Add(1)
go func() {
defer wg.Done()
ok, err := repo.ClaimReduction(context.Background(), job.ID, time.Now().UTC())
if err != nil {
t.Errorf("claim reduction: %v", err)
return
}
if ok {
mu.Lock()
claimed++
mu.Unlock()
}
}()
}
wg.Wait()
if claimed != 1 {
t.Fatalf("reducer claims = %d, want 1", claimed)
}
stored, err := repo.Get(ctx, job.ID)
if err != nil {
t.Fatal(err)
}
if stored.Status != domain.JobReducing || stored.ReducerStartedAt == nil {
t.Fatalf("stored reduction state = %+v", stored)
}
}
func TestUIReadRepoListsReducerFields(t *testing.T) {
pool := testPool(t)
job, _ := seedJob(t, pool, 1)
jobs := NewJobRepo(pool)
ctx := context.Background()
if err := jobs.UpdateStatus(ctx, job.ID, domain.JobReducing, nil); err != nil {
t.Fatal(err)
}
if claimed, err := jobs.ClaimReduction(ctx, job.ID, time.Now().UTC()); err != nil || !claimed {
t.Fatalf("claim reduction = (%v, %v)", claimed, err)
}
listed, err := NewUIReadRepo(pool).ListJobs(ctx, 20)
if err != nil {
t.Fatalf("list UI jobs: %v", err)
}
for _, item := range listed {
if item.ID != job.ID {
continue
}
if item.Status != domain.JobReducing || item.ReducerStartedAt == nil {
t.Fatalf("UI reducer projection = %+v", item)
}
return
}
t.Fatalf("seeded job %s is missing from UI list", job.ID)
}
// A job must land whole or not at all: a half-created job leaves chunks no
// worker could ever complete.
func TestCreateJobRollsBackOnFailure(t *testing.T) {
@@ -137,30 +207,37 @@ func TestConcurrentClaimGivesEachTaskToExactlyOneWorker(t *testing.T) {
claimed = make(map[uuid.UUID]string)
wg sync.WaitGroup
)
// More workers than tasks, so the surplus must come back empty rather than
// steal an already-leased row.
// More workers than tasks. With SKIP LOCKED, a concurrent caller can
// transiently see no eligible row while every remaining row is locked by a
// different claim statement. Poll briefly, as a real worker does, before
// treating the queue as empty. This verifies the actual contract: tasks are
// unique and all eventually become claimable without lock contention.
for i := 0; i < tasks*2; i++ {
wg.Add(1)
go func(n int) {
defer wg.Done()
task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{
Owner: fmt.Sprintf("worker-%d", n),
Now: now,
LeaseUntil: now.Add(time.Minute),
})
if err != nil {
t.Errorf("claim: %v", err)
for attempt := 0; attempt < 20; attempt++ {
task, err := repo.ClaimNext(context.Background(), usecase.ClaimFilter{
Owner: fmt.Sprintf("worker-%d", n),
Now: now,
LeaseUntil: now.Add(time.Minute),
})
if err != nil {
t.Errorf("claim: %v", err)
return
}
if task == nil || task.JobID != job.ID {
time.Sleep(time.Millisecond)
continue
}
mu.Lock()
if prev, dup := claimed[task.ID]; dup {
t.Errorf("task %s handed to both %s and worker-%d", task.ID, prev, n)
}
claimed[task.ID] = fmt.Sprintf("worker-%d", n)
mu.Unlock()
return
}
if task == nil || task.JobID != job.ID {
return // empty queue, or a task from another test's job
}
mu.Lock()
defer mu.Unlock()
if prev, dup := claimed[task.ID]; dup {
t.Errorf("task %s handed to both %s and worker-%d", task.ID, prev, n)
}
claimed[task.ID] = fmt.Sprintf("worker-%d", n)
}(i)
}
wg.Wait()
@@ -199,6 +276,27 @@ func TestClaimNextReturnsNilOnEmptyQueue(t *testing.T) {
}
}
func TestCancelJobCancelsEveryUnfinishedTask(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
job, _ := seedJob(t, pool, 3)
clk := fixedClock{now: time.Now().UTC()}
uc := usecase.NewCancelJob(NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool), clk)
cancelled, err := uc.Execute(ctx, job.ID)
if err != nil || cancelled != 3 {
t.Fatalf("cancel = (%d, %v), want (3, nil)", cancelled, err)
}
stored, err := NewJobRepo(pool).Get(ctx, job.ID)
if err != nil || stored.Status != domain.JobCancelled {
t.Fatalf("job after cancel = (%+v, %v)", stored, err)
}
counts, err := NewTaskRepo(pool).CountByStatus(ctx, job.ID)
if err != nil || counts[domain.TaskCancelled] != 3 {
t.Fatalf("cancelled tasks = %d, err = %v", counts[domain.TaskCancelled], err)
}
}
func TestUpdateRejectsStaleVersion(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
@@ -332,12 +430,23 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) {
t.Errorf("replay must be idempotent, got %v", err)
}
// A different result artifact for the same task is a genuine conflict.
art2 := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult)
other := in
other.ResultArtifactID = art2.ID
if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) {
t.Errorf("err = %v, want ErrResultConflict", err)
}
func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
pool := testPool(t)
ctx := context.Background()
job, tasks := seedJob(t, pool, 1)
taskID := tasks[0].ID
first := seedArtifact(t, pool, job.ID, &taskID, domain.ArtifactPartialResult)
second, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, "retry.csv", "text/csv", time.Now().UTC())
if err != nil {
t.Fatal(err)
}
attempt := 1
second.Attempt = &attempt
second.SetContent("other-sha", 5)
if err := NewArtifactRepo(pool).Insert(ctx, second); err == nil {
t.Fatalf("second partial artifact for %s/%d was accepted after %s", taskID, attempt, first.ID)
}
}
@@ -25,7 +25,10 @@ func NewJobRepo(pool *pgxpool.Pool) *JobRepo {
var _ usecase.JobRepository = (*JobRepo)(nil)
var jobColumns = []string{"id", "workload", "input_uri", "parameters", "status", "created_at", "completed_at"}
var jobColumns = []string{
"id", "workload", "input_uri", "parameters", "status", "created_at", "completed_at",
"input_artifact_id", "result_artifact_id", "error_code", "error_message", "reducer_started_at",
}
// Insert runs inside the caller's transaction, alongside the job's tasks — that
// is what makes "all tasks or none" hold.
@@ -55,7 +58,8 @@ func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
status string
)
err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan(
&j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt)
&j.ID, &j.Workload, &j.InputURI, &j.Parameters, &status, &j.CreatedAt, &j.CompletedAt,
&j.InputArtifactID, &j.ResultArtifactID, &j.ErrorCode, &j.ErrorMessage, &j.ReducerStartedAt)
if errors.Is(err, pgx.ErrNoRows) {
return nil, domain.ErrJobNotFound
}
@@ -66,6 +70,64 @@ func (r *JobRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
return &j, nil
}
func (r *JobRepo) ClaimReduction(ctx context.Context, id uuid.UUID, startedAt time.Time) (bool, error) {
sql, args, err := psql.Update("jobs").
Set("reducer_started_at", startedAt).
Where(sq.Eq{"id": id, "status": string(domain.JobReducing), "reducer_started_at": nil}).
ToSql()
if err != nil {
return false, err
}
tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...)
if err != nil {
return false, err
}
return tag.RowsAffected() == 1, nil
}
func (r *JobRepo) CompleteWithResult(ctx context.Context, id, resultArtifactID uuid.UUID, completedAt time.Time) error {
sql, args, err := psql.Update("jobs").
SetMap(map[string]any{
"status": string(domain.JobCompleted),
"result_artifact_id": resultArtifactID,
"completed_at": completedAt,
"reducer_started_at": nil,
"error_code": nil,
"error_message": nil,
}).
Where(sq.Eq{"id": id, "status": string(domain.JobReducing)}).
ToSql()
if err != nil {
return err
}
tag, err := conn(ctx, r.pool).Exec(ctx, sql, args...)
if err != nil {
return err
}
if tag.RowsAffected() == 0 {
return domain.ErrJobNotFound
}
return nil
}
func (r *JobRepo) FailReduction(ctx context.Context, id uuid.UUID, code, message string, completedAt time.Time) error {
sql, args, err := psql.Update("jobs").
SetMap(map[string]any{
"status": string(domain.JobFailed),
"completed_at": completedAt,
"error_code": code,
"error_message": message,
"reducer_started_at": nil,
}).
Where(sq.Eq{"id": id, "status": string(domain.JobReducing)}).
ToSql()
if err != nil {
return err
}
_, err = conn(ctx, r.pool).Exec(ctx, sql, args...)
return err
}
func (r *JobRepo) UpdateStatus(ctx context.Context, id uuid.UUID,
status domain.JobStatus, completedAt *time.Time) error {
@@ -295,6 +295,29 @@ func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[doma
return counts, rows.Err()
}
// cancelByJobSQL mirrors domain.Task.Cancel in one set-based update. It runs in
// the same transaction as the job-status update, so no claimable shard remains
// after an operator receives a successful cancellation response.
const cancelByJobSQL = `
UPDATE tasks
SET status = 'cancelled'::task_status,
lease_owner = NULL,
lease_expires_at = NULL,
error_code = NULL,
error_message = NULL,
completed_at = $2,
version = version + 1
WHERE job_id = $1
AND status IN ('pending','leased','running')`
func (r *TaskRepo) CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error) {
tag, err := conn(ctx, r.pool).Exec(ctx, cancelByJobSQL, jobID, now)
if err != nil {
return 0, err
}
return tag.RowsAffected(), nil
}
// expireLeasesSQL applies the lease-expiry rule set-based, mirroring
// domain.Task.ExpireLease: requeue while attempts remain, otherwise fail.
//
@@ -314,17 +337,26 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta
ELSE error_message END,
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
version = version + 1
WHERE status IN ('leased','running') AND lease_expires_at < $1`
WHERE status IN ('leased','running') AND lease_expires_at < $1
RETURNING job_id`
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) {
var affected int64
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
var affected []uuid.UUID
err := withRetry(ctx, func(ctx context.Context) error {
tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
rows, err := conn(ctx, r.pool).Query(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired)
if err != nil {
return err
}
affected = tag.RowsAffected()
return nil
defer rows.Close()
affected = affected[:0]
for rows.Next() {
var jobID uuid.UUID
if err := rows.Scan(&jobID); err != nil {
return err
}
affected = append(affected, jobID)
}
return rows.Err()
})
return affected, err
}
@@ -0,0 +1,147 @@
package postgres
import (
"context"
"fmt"
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"
)
// UIReadRepo contains bounded, deterministic read queries for the operator UI.
type UIReadRepo struct{ pool *pgxpool.Pool }
func NewUIReadRepo(pool *pgxpool.Pool) *UIReadRepo { return &UIReadRepo{pool: pool} }
var _ usecase.UIReadRepository = (*UIReadRepo)(nil)
func (r *UIReadRepo) GetJob(ctx context.Context, id uuid.UUID) (*domain.Job, error) {
job, err := NewJobRepo(r.pool).Get(ctx, id)
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
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,
); err != nil {
return nil, err
}
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 {
return nil, err
}
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
if err != nil {
return nil, fmt.Errorf("list tasks: %w", err)
}
defer rows.Close()
tasks := make([]domain.Task, 0)
for rows.Next() {
task, err := scanTask(rows)
if err != nil {
return nil, err
}
tasks = append(tasks, *task)
}
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()
if err != nil {
return nil, err
}
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
if err != nil {
return nil, fmt.Errorf("list workers: %w", err)
}
defer rows.Close()
workers := make([]domain.Worker, 0)
for rows.Next() {
worker, err := scanWorker(rows)
if err != nil {
return nil, err
}
workers = append(workers, *worker)
}
return workers, rows.Err()
}
func (r *UIReadRepo) ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error) {
sql, args, err := psql.Select(artifactColumns...).From("artifacts").Where(sq.Eq{"job_id": jobID}).OrderBy("created_at ASC", "id ASC").ToSql()
if err != nil {
return nil, err
}
rows, err := conn(ctx, r.pool).Query(ctx, sql, args...)
if err != nil {
return nil, fmt.Errorf("list artifacts: %w", err)
}
defer rows.Close()
artifacts := make([]domain.Artifact, 0)
for rows.Next() {
var a domain.Artifact
var kind string
if err := rows.Scan(&a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey, &a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt); err != nil {
return nil, err
}
a.Kind = domain.ArtifactKind(kind)
artifacts = append(artifacts, a)
}
return artifacts, rows.Err()
}
+26 -15
View File
@@ -111,13 +111,16 @@ type uploadJobResponse struct {
}
type jobProgressResponse struct {
ID uuid.UUID `json:"id"`
Status string `json:"status"`
Total int `json:"total"`
Pending int `json:"pending"`
Leased int `json:"leased"`
Done int `json:"completed"`
Failed int `json:"failed"`
ID uuid.UUID `json:"id"`
Status string `json:"status"`
Total int `json:"total"`
Pending int `json:"pending"`
Leased int `json:"leased"`
Done int `json:"completed"`
Failed int `json:"failed"`
Cancelled int `json:"cancelled"`
ResultURI string `json:"result_uri,omitempty"`
ErrorCode string `json:"error_code,omitempty"`
}
type uploadArtifactResponse struct {
@@ -152,13 +155,21 @@ func toClaimedTaskResponse(c domain.ClaimedTask) claimedTaskResponse {
}
func toJobProgressResponse(p domain.JobProgress) jobProgressResponse {
return jobProgressResponse{
ID: p.Job.ID,
Status: string(p.DeriveStatus()),
Total: p.Total,
Pending: p.Pending,
Leased: p.Leased,
Done: p.Done,
Failed: p.Failed,
out := jobProgressResponse{
ID: p.Job.ID,
Status: string(p.DeriveStatus()),
Total: p.Total,
Pending: p.Pending,
Leased: p.Leased,
Done: p.Done,
Failed: p.Failed,
Cancelled: p.Cancelled,
}
if p.Job.ResultArtifactID != nil && out.Status == string(domain.JobCompleted) {
out.ResultURI = "/jobs/" + p.Job.ID.String() + "/result"
}
if p.Job.ErrorCode != nil {
out.ErrorCode = *p.Job.ErrorCode
}
return out
}
@@ -0,0 +1,25 @@
package http
import (
"testing"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
)
func TestJobProgressResponseExposesFinalResultOnlyWhenCompleted(t *testing.T) {
id := uuid.New()
result := uuid.New()
progress := domain.JobProgress{Job: domain.Job{
ID: id, Status: domain.JobCompleted, ResultArtifactID: &result,
}, Total: 1, Done: 1}
if got, want := toJobProgressResponse(progress).ResultURI, "/jobs/"+id.String()+"/result"; got != want {
t.Fatalf("result URI = %q, want %q", got, want)
}
progress.Job.Status = domain.JobReducing
if got := toJobProgressResponse(progress).ResultURI; got != "" {
t.Fatalf("reducing job exposes result URI %q", got)
}
}
@@ -50,7 +50,8 @@ func (s *Server) writeError(w http.ResponseWriter, r *http.Request, err error) {
case errors.Is(err, domain.ErrLeaseConflict),
errors.Is(err, domain.ErrStaleAttempt),
errors.Is(err, domain.ErrResultConflict),
errors.Is(err, domain.ErrTaskNotLeased):
errors.Is(err, domain.ErrTaskNotLeased),
errors.Is(err, domain.ErrJobNotCancellable):
status = http.StatusConflict
case errors.Is(err, usecase.ErrNotImplemented):
status = http.StatusNotImplemented
@@ -79,6 +79,10 @@ func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
if _, err := uuid.Parse(req.WorkerID); err != nil {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{
WorkerID: req.WorkerID,
@@ -146,6 +150,12 @@ func (s *Server) handleResult(w http.ResponseWriter, r *http.Request) {
s.writeError(w, r, err)
return
}
if s.uc.ReduceJob != nil {
if err := s.uc.ReduceJob.Execute(ctx, task.JobID); err != nil {
s.writeError(w, r, err)
return
}
}
writeJSON(w, http.StatusOK, taskResponse{ID: task.ID, JobID: task.JobID, Status: string(task.Status)})
}
@@ -182,7 +192,7 @@ func (s *Server) handleFailure(w http.ResponseWriter, r *http.Request) {
const defaultChunkRows = 1000
// handleUploadDataset accepts a multipart submission — the dataset file plus the
// workload/parameters/chunk_rows fields — and hands the file, streamed, to the
// workload/parameters/chunk_rows/max_rows fields — and hands the file, streamed, to the
// chunker. The text fields MUST precede the file part: the file is streamed, not
// buffered, so by the time it arrives the other fields are already parsed.
func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
@@ -197,11 +207,13 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
workload string
params map[string]any
rows = defaultChunkRows
maxRows int
result usecase.SubmitDatasetResult
gotDataset bool
gotWorkload bool
gotParams bool
gotRows bool
gotMaxRows bool
)
for {
@@ -249,6 +261,19 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
}
rows = n
gotRows = true
case "max_rows":
if gotDataset || gotMaxRows {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
b, _ := io.ReadAll(io.LimitReader(part, 32))
n, err := strconv.Atoi(strings.TrimSpace(string(b)))
if err != nil || n < 1 {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
maxRows = n
gotMaxRows = true
case "file", "dataset":
if gotDataset || workload == "" {
s.writeError(w, r, domain.ErrInvalidInput)
@@ -262,6 +287,7 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) {
Workload: workload,
Parameters: params,
RowsPerShard: rows,
MaxRows: maxRows,
Filename: filename,
ContentType: part.Header.Get("Content-Type"),
Body: part,
@@ -379,6 +405,46 @@ func (s *Server) handleGetJob(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, toJobProgressResponse(progress))
}
func (s *Server) handleGetJobResult(w http.ResponseWriter, r *http.Request) {
ctx, cancel := s.reqCtx(r)
defer cancel()
jobID, ok := s.pathUUID(w, r, "job_id")
if !ok {
return
}
art, body, err := s.uc.GetJobResult.Execute(ctx, jobID)
if err != nil {
s.writeError(w, r, err)
return
}
defer func() { _ = body.Close() }()
w.Header().Set("Content-Type", art.ContentType)
w.Header().Set("Content-Length", strconv.FormatInt(art.SizeBytes, 10))
w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%q", art.Filename))
_, _ = io.Copy(w, body)
}
// handleCancelJob stops all non-terminal shards for an operator-requested job.
// It is available to both the bearer API and the separately authenticated UI.
func (s *Server) handleCancelJob(w http.ResponseWriter, r *http.Request) {
ctx, cancel := s.reqCtx(r)
defer cancel()
jobID, ok := s.pathUUID(w, r, "job_id")
if !ok {
return
}
cancelled, err := s.uc.CancelJob.Execute(ctx, jobID)
if err != nil {
s.writeError(w, r, err)
return
}
writeJSON(w, http.StatusOK, map[string]any{
"job_id": jobID,
"status": domain.JobCancelled,
"cancelled_tasks": cancelled,
})
}
// --- helpers ---
func (s *Server) reqCtx(r *http.Request) (context.Context, context.CancelFunc) {
@@ -64,6 +64,48 @@ func withAuth(token string) func(http.Handler) http.Handler {
}
}
// withBasicAuth protects the local operator UI with a credential distinct from
// the worker bearer token. The username is intentionally ignored; the password
// is the configured UI token. Basic Auth is suitable only for localhost or a
// TLS-terminating trusted reverse proxy.
func withBasicAuth(token string) func(http.Handler) http.Handler {
return func(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, password, ok := r.BasicAuth()
if !ok || subtle.ConstantTimeCompare([]byte(password), []byte(token)) != 1 {
w.Header().Set("WWW-Authenticate", `Basic realm="SciMesh UI", charset="UTF-8"`)
writeJSON(w, http.StatusUnauthorized, errorResponse{Error: "unauthorized", RequestID: requestIDFrom(r.Context())})
return
}
next.ServeHTTP(w, r)
})
}
}
// withSameOrigin rejects browser form/fetch writes initiated by another origin.
// A missing Origin is allowed for direct local tools; authenticated UI pages use
// the browser-supplied Origin header on state-changing requests.
func withSameOrigin(next http.Handler) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodGet || r.Method == http.MethodHead || r.Method == http.MethodOptions {
next.ServeHTTP(w, r)
return
}
origin := r.Header.Get("Origin")
if origin != "" {
scheme := "http"
if r.TLS != nil {
scheme = "https"
}
if origin != scheme+"://"+r.Host {
writeJSON(w, http.StatusForbidden, errorResponse{Error: "cross-origin request rejected", RequestID: requestIDFrom(r.Context())})
return
}
}
next.ServeHTTP(w, r)
})
}
// statusRecorder captures the status code for the access log.
type statusRecorder struct {
http.ResponseWriter
+25 -1
View File
@@ -22,11 +22,15 @@ type UseCases struct {
ClaimTask *usecase.ClaimTask
RenewLease *usecase.RenewLease
CompleteTask *usecase.CompleteTask
ReduceJob *usecase.ReduceJob
FailTask *usecase.FailTask
GetJobStatus *usecase.GetJobStatus
GetJobResult *usecase.GetJobResult
CancelJob *usecase.CancelJob
UploadArtifact *usecase.UploadArtifact
DownloadArtifact *usecase.DownloadArtifact
GetTaskInput *usecase.GetTaskInput
Dashboard *usecase.Dashboard
}
type Server struct {
@@ -54,12 +58,14 @@ func NewServer(uc UseCases, log *slog.Logger, requestTimeout, heartbeatInterval
// Handler builds the router. Go 1.22's ServeMux matches on method and path
// wildcards, so no third-party router is needed.
func (s *Server) Handler(token string) http.Handler {
func (s *Server) Handler(token string, uiToken ...string) http.Handler {
protected := http.NewServeMux()
protected.HandleFunc("POST /workers/register", s.handleRegister)
protected.HandleFunc("POST /jobs", s.handleCreateJob)
protected.HandleFunc("POST /jobs/upload", s.handleUploadDataset)
protected.HandleFunc("GET /jobs/{job_id}", s.handleGetJob)
protected.HandleFunc("GET /jobs/{job_id}/result", s.handleGetJobResult)
protected.HandleFunc("POST /jobs/{job_id}/cancel", s.handleCancelJob)
protected.HandleFunc("POST /tasks/claim", s.handleClaim)
protected.HandleFunc("GET /tasks/{task_id}/input", s.handleGetTaskInput)
protected.HandleFunc("POST /tasks/{task_id}/heartbeat", s.handleHeartbeat)
@@ -70,6 +76,24 @@ func (s *Server) Handler(token string) http.Handler {
mux := http.NewServeMux()
mux.HandleFunc("GET /health", s.handleHealth)
if len(uiToken) > 0 && uiToken[0] != "" && s.uc.Dashboard != nil {
ui := http.NewServeMux()
ui.HandleFunc("GET /ui", s.handleUIHome)
ui.HandleFunc("GET /ui/jobs/new", s.handleUINewJob)
ui.HandleFunc("GET /ui/jobs/{job_id}", s.handleUIJob)
ui.HandleFunc("GET /ui/api/overview", s.handleUIOverviewJSON)
ui.HandleFunc("GET /ui/api/jobs/{job_id}", s.handleUIJobJSON)
ui.HandleFunc("POST /ui/api/jobs/{job_id}/cancel", s.handleCancelJob)
ui.HandleFunc("POST /ui/api/jobs/upload", s.handleUploadDataset)
ui.HandleFunc("GET /ui/jobs/{job_id}/artifacts/{artifact_id}", s.handleUIArtifactDownload)
mux.Handle("/ui", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
mux.Handle("/ui/", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin))
} else {
// More specific than the protected catch-all: UI absence is not an auth
// failure and does not disclose that a UI feature is configured elsewhere.
mux.HandleFunc("/ui", http.NotFound)
mux.HandleFunc("/ui/", http.NotFound)
}
mux.Handle("/", chain(protected,
withRequestID, // outermost: every response gets an ID,
withAccessLog(s.log), // including the 401s below
@@ -20,13 +20,19 @@ import (
)
const token = "secret"
const uiToken = "ui-secret"
type env struct {
ts *httptest.Server
blobs *memstore.BlobStore
ts *httptest.Server
blobs *memstore.BlobStore
workerID string
}
func newEnv(t *testing.T, ready func(context.Context) error) *env {
return newEnvWithUIToken(t, ready, uiToken)
}
func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configuredUIToken string) *env {
t.Helper()
tasks := memstore.NewTaskRepo()
jobs := memstore.NewJobRepo()
@@ -36,24 +42,35 @@ func newEnv(t *testing.T, ready func(context.Context) error) *env {
clk := memstore.NewClock(time.Date(2026, 7, 21, 12, 0, 0, 0, time.UTC))
tx := memstore.Tx{}
lease := 2 * time.Minute
downloadArtifact := usecase.NewDownloadArtifact(arts, blobs)
uc := coordhttp.UseCases{
RegisterWorker: usecase.NewRegisterWorker(work, clk),
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk),
ClaimTask: usecase.NewClaimTask(tasks, clk, lease),
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk, 3),
ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease),
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
ReduceJob: usecase.NewReduceJob(jobs, tasks, arts, blobs, tx, clk),
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, clk),
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
GetJobResult: usecase.NewGetJobResult(jobs, downloadArtifact),
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, tx, clk),
DownloadArtifact: downloadArtifact,
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)),
}
worker, err := uc.RegisterWorker.Execute(context.Background(), usecase.RegisterWorkerInput{
Name: "test-worker", Capabilities: []string{"w", "similarity-search"},
})
if err != nil {
t.Fatalf("register test worker: %v", err)
}
srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, 1<<30, ready)
ts := httptest.NewServer(srv.Handler(token))
ts := httptest.NewServer(srv.Handler(token, configuredUIToken))
t.Cleanup(ts.Close)
return &env{ts: ts, blobs: blobs}
return &env{ts: ts, blobs: blobs, workerID: worker.ID.String()}
}
func healthy(context.Context) error { return nil }
@@ -61,6 +78,7 @@ func healthy(context.Context) error { return nil }
// do sends an authenticated JSON request and returns status + decoded body.
func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) {
t.Helper()
body = strings.ReplaceAll(body, `"worker_id":"w1"`, `"worker_id":"`+e.workerID+`"`)
req, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token)
if body != "" {
@@ -97,6 +115,206 @@ func TestHealthOK(t *testing.T) {
}
}
func TestUIRequiresDistinctCredentialAndRendersDashboard(t *testing.T) {
e := newEnv(t, healthy)
request := func() *http.Request {
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui", nil)
return req
}
resp, err := http.DefaultClient.Do(request())
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusUnauthorized {
t.Fatalf("no UI auth: %d", resp.StatusCode)
}
req := request()
req.Header.Set("Authorization", "Bearer "+token)
resp, err = http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusUnauthorized {
t.Fatalf("worker token authorized UI: %d", resp.StatusCode)
}
req = request()
req.SetBasicAuth("operator", uiToken)
resp, err = http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("UI status: %d", resp.StatusCode)
}
body, _ := io.ReadAll(resp.Body)
if !strings.Contains(string(body), "SciMesh control room") {
t.Errorf("dashboard body missing title")
}
}
func TestUIOverviewReturnsLiveSafeProjection(t *testing.T) {
e := newEnv(t, healthy)
code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("create job: %d", code)
}
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/api/overview", nil)
req.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
var overview map[string]any
if err := json.NewDecoder(resp.Body).Decode(&overview); err != nil {
t.Fatal(err)
}
if resp.StatusCode != http.StatusOK || overview["active_jobs"].(float64) != 1 || overview["online_workers"].(float64) != 1 {
t.Fatalf("overview = (%d, %v)", resp.StatusCode, overview)
}
if _, leaked := overview["worker_auth_token"]; leaked {
t.Fatal("overview must not expose authentication configuration")
}
}
func TestUIDisabledReturnsNotFound(t *testing.T) {
e := newEnvWithUIToken(t, healthy, "")
resp := e.get(t, "/ui")
defer resp.Body.Close()
if resp.StatusCode != http.StatusNotFound {
t.Fatalf("disabled UI = %d, want 404", resp.StatusCode)
}
}
func TestUIRejectsCrossOriginUpload(t *testing.T) {
e := newEnv(t, healthy)
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/upload", strings.NewReader("dataset=x"))
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
req.Header.Set("Origin", "https://attacker.example")
req.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusForbidden {
t.Fatalf("cross-origin upload = %d, want 403", resp.StatusCode)
}
}
func TestUIUploadDatasetCreatesJob(t *testing.T) {
e := newEnv(t, healthy)
var body bytes.Buffer
mw := multipart.NewWriter(&body)
_ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO","top_k":20,"progress_every":0}`)
_ = mw.WriteField("chunk_rows", "1000")
file, err := mw.CreateFormFile("file", "chembl.tsv")
if err != nil {
t.Fatal(err)
}
_, _ = io.WriteString(file, "chembl_id\tcanonical_smiles\nCHEMBL1\tCCO\n")
if err := mw.Close(); err != nil {
t.Fatal(err)
}
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/upload", &body)
req.Header.Set("Content-Type", mw.FormDataContentType())
req.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusCreated {
result, _ := io.ReadAll(resp.Body)
t.Fatalf("UI upload = %d: %s", resp.StatusCode, result)
}
}
func TestCancelJobStopsUnfinishedTasks(t *testing.T) {
e := newEnv(t, healthy)
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"sha"},{"chunk_index":1,"input_uri":"s3://c1","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("create: %d", code)
}
jobID := job["id"].(string)
if code, body := e.do(t, "POST", "/jobs/"+jobID+"/cancel", ""); code != http.StatusOK || body["cancelled_tasks"].(float64) != 2 {
t.Fatalf("cancel = (%d, %v)", code, body)
}
if code, progress := e.do(t, "GET", "/jobs/"+jobID, ""); code != http.StatusOK || progress["status"] != "cancelled" || progress["cancelled"].(float64) != 2 {
t.Fatalf("cancelled job progress = (%d, %v)", code, progress)
}
}
func TestUICancelJobUsesOperatorCredential(t *testing.T) {
e := newEnv(t, healthy)
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("create: %d", code)
}
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/"+job["id"].(string)+"/cancel", nil)
req.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("UI cancel = %d", resp.StatusCode)
}
}
func TestUIJobAndArtifactAreScopedToTheirJob(t *testing.T) {
e := newEnv(t, healthy)
code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("create: %d", code)
}
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+job["id"].(string), nil)
req.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
t.Fatalf("detail: %d", resp.StatusCode)
}
if got := resp.Header.Get("Content-Security-Policy"); got == "" {
t.Error("missing UI CSP")
}
}
func TestUIArtifactDownloadRejectsAnotherJobsArtifact(t *testing.T) {
e := newEnv(t, healthy)
code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("first job: %d", code)
}
_, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["w"]}`)
artifactID := e.putArtifact(t, claim["task_id"].(string), "w1", int(claim["attempt"].(float64)), "result")
code, second := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`)
if code != http.StatusCreated {
t.Fatalf("second job: %d", code)
}
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+second["id"].(string)+"/artifacts/"+artifactID, nil)
req.SetBasicAuth("operator", uiToken)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusNotFound {
t.Errorf("cross-job artifact = %d, want 404", resp.StatusCode)
}
}
func TestHealthUnavailableWhenDBDown(t *testing.T) {
e := newEnv(t, func(context.Context) error { return context.DeadlineExceeded })
resp := e.get(t, "/health")
@@ -148,6 +366,26 @@ func TestRegisterRejectsNoCapabilities(t *testing.T) {
}
}
func TestClaimRequiresRegisteredWorkerAndUsesStoredCapabilities(t *testing.T) {
e := newEnv(t, healthy)
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"not-a-uuid"}`); code != http.StatusBadRequest {
t.Fatalf("invalid worker id claim = %d, want 400", code)
}
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"11111111-1111-4111-8111-111111111111"}`); code != http.StatusNotFound {
t.Fatalf("unregistered worker claim = %d, want 404", code)
}
if code, _ := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`); code != http.StatusCreated {
t.Fatalf("create job = %d", code)
}
code, worker := e.do(t, "POST", "/workers/register", `{"name":"search-only","capabilities":["similarity-search"]}`)
if code != http.StatusCreated {
t.Fatalf("register = %d", code)
}
if code, _ := e.do(t, "POST", "/tasks/claim", `{"worker_id":"`+worker["worker_id"].(string)+`","capabilities":["w"]}`); code != http.StatusNoContent {
t.Fatalf("forged capability claim = %d, want 204", code)
}
}
func TestFullLifecycle(t *testing.T) {
e := newEnv(t, healthy)
@@ -190,6 +428,69 @@ func TestFullLifecycle(t *testing.T) {
}
}
func TestSimilaritySearchLifecyclePublishesFinalResult(t *testing.T) {
e := newEnv(t, healthy)
code, job := e.uploadDataset(t, "similarity-search", 10, "chembl_id\tcanonical_smiles\nA\tCC\n")
if code != http.StatusCreated {
t.Fatalf("upload job: %d", code)
}
jobID := job["job_id"].(string)
code, claim := e.do(t, "POST", "/tasks/claim", `{"worker_id":"w1","capabilities":["similarity-search"]}`)
if code != http.StatusOK {
t.Fatalf("claim: %d", code)
}
taskID := claim["task_id"].(string)
attempt := int(claim["attempt"].(float64))
artifactID := e.putArtifact(t, taskID, "w1", attempt, "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.900000\n")
if code, _ := e.do(t, "POST", "/tasks/"+taskID+"/result",
`{"worker_id":"w1","attempt":`+itoa(attempt)+`,"result":{"artifact_id":"`+artifactID+`"}}`); code != http.StatusOK {
t.Fatalf("complete: %d", code)
}
code, progress := e.do(t, "GET", "/jobs/"+jobID, "")
if code != http.StatusOK || progress["status"] != "completed" || progress["result_uri"] != "/jobs/"+jobID+"/result" {
t.Fatalf("progress = (%d, %v)", code, progress)
}
req, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+progress["result_uri"].(string), nil)
req.Header.Set("Authorization", "Bearer "+token)
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
body, _ := io.ReadAll(resp.Body)
if resp.StatusCode != http.StatusOK || string(body) != "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.900000\n" {
t.Fatalf("final result = (%d, %q)", resp.StatusCode, body)
}
uiRequest, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/jobs/"+jobID, nil)
uiRequest.SetBasicAuth("operator", uiToken)
uiResponse, err := http.DefaultClient.Do(uiRequest)
if err != nil {
t.Fatal(err)
}
defer uiResponse.Body.Close()
uiBody, _ := io.ReadAll(uiResponse.Body)
if uiResponse.StatusCode != http.StatusOK || !strings.Contains(string(uiBody), "Final result ready") {
t.Fatalf("final UI = (%d, %q)", uiResponse.StatusCode, uiBody)
}
jsonRequest, _ := http.NewRequestWithContext(context.Background(), "GET", e.ts.URL+"/ui/api/jobs/"+jobID, nil)
jsonRequest.SetBasicAuth("operator", uiToken)
jsonResponse, err := http.DefaultClient.Do(jsonRequest)
if err != nil {
t.Fatal(err)
}
defer jsonResponse.Body.Close()
var detail map[string]any
if err := json.NewDecoder(jsonResponse.Body).Decode(&detail); err != nil {
t.Fatal(err)
}
if jsonResponse.StatusCode != http.StatusOK || detail["final_result_available"] != true {
t.Fatalf("final UI JSON = (%d, %v)", jsonResponse.StatusCode, detail)
}
}
func TestForeignArtifactResultConflict(t *testing.T) {
e := newEnv(t, healthy)
e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in",
@@ -211,9 +512,9 @@ func TestForeignArtifactResultConflict(t *testing.T) {
func TestUploadDatasetChunksAndServesInput(t *testing.T) {
e := newEnv(t, healthy)
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
code, body := e.uploadDataset(t, "w", 2, tsv)
code, body := e.uploadDataset(t, "similarity-search", 2, tsv)
if code != 201 {
t.Fatalf("upload: status = %d", code)
}
@@ -239,11 +540,60 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) {
t.Fatalf("get input: status = %d", resp.StatusCode)
}
shard, _ := io.ReadAll(resp.Body)
if !strings.HasPrefix(string(shard), "id\tsmiles\n") {
if !strings.HasPrefix(string(shard), "chembl_id\tcanonical_smiles\n") {
t.Errorf("shard missing header: %q", shard)
}
}
func TestUploadDatasetLimitsRows(t *testing.T) {
e := newEnv(t, healthy)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", "2")
_ = mw.WriteField("max_rows", "3")
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
_ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", mw.FormDataContentType())
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
var result map[string]any
_ = json.NewDecoder(resp.Body).Decode(&result)
if resp.StatusCode != http.StatusCreated || result["task_count"].(float64) != 2 {
t.Fatalf("limited upload = (%d, %v)", resp.StatusCode, result)
}
}
func TestUploadDatasetRejectsMissingChEMBLColumns(t *testing.T) {
e := newEnv(t, healthy)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", "2")
fw, _ := mw.CreateFormFile("file", "not-chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
_ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", mw.FormDataContentType())
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusBadRequest {
t.Errorf("missing ChEMBL columns = %d, want 400", resp.StatusCode)
}
}
func TestErrorMappings(t *testing.T) {
e := newEnv(t, healthy)
zero := "00000000-0000-0000-0000-000000000000"
@@ -271,10 +621,11 @@ func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) {
e := newEnv(t, healthy)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", "w")
_ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", "not-a-number")
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
_, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"))
_ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
@@ -298,6 +649,9 @@ func (e *env) putArtifact(t *testing.T, taskID, worker string, attempt int, data
e.ts.URL+"/tasks/"+taskID+"/artifacts/r.csv", strings.NewReader(data))
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", "text/csv")
if worker == "w1" {
worker = e.workerID
}
req.Header.Set("X-Worker-ID", worker)
req.Header.Set("X-Task-Attempt", itoa(attempt))
resp, err := http.DefaultClient.Do(req)
@@ -319,6 +673,7 @@ func (e *env) uploadDataset(t *testing.T, workload string, rows int, tsv string)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", workload)
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", itoa(rows))
fw, _ := mw.CreateFormFile("file", "chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader(tsv))
@@ -0,0 +1,39 @@
{{define "dashboard.html"}}
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>SciMesh control room</title>
<style>
:root{color:#dce8ff;background:#08111f;font:16px/1.5 Inter,ui-sans-serif,system-ui,sans-serif;color-scheme:dark}*{box-sizing:border-box}body{margin:0;background:radial-gradient(circle at 15% -10%,#163d77 0,transparent 32rem),radial-gradient(circle at 95% 5%,#123e39 0,transparent 29rem),#08111f}.page{max-width:1240px;margin:auto;padding:28px 22px 64px}.top{display:flex;align-items:flex-start;justify-content:space-between;gap:22px}.eyebrow{margin:0;color:#78a9ff;font-size:.77rem;font-weight:750;letter-spacing:.14em;text-transform:uppercase}.title{max-width:700px;margin:8px 0;font-size:clamp(2rem,5vw,3.6rem);line-height:1.04;letter-spacing:-.055em}.lead{max-width:690px;margin:0;color:#aabbd5;font-size:1.05rem}.button{display:inline-flex;align-items:center;gap:8px;border:0;border-radius:10px;padding:12px 16px;background:#4f8cff;color:#071224;font:inherit;font-weight:800;text-decoration:none;box-shadow:0 12px 30px #163d7766}.live{display:inline-flex;align-items:center;gap:7px;margin-top:18px;color:#8ba2c2;font-size:.87rem}.pulse{width:8px;height:8px;border-radius:50%;background:#5ee6a6;box-shadow:0 0 0 5px #5ee6a622}.summary{display:grid;grid-template-columns:1.4fr repeat(3,1fr);gap:13px;margin:32px 0}.panel,.metric,.flow-step,.job,.worker{border:1px solid #26415f;border-radius:15px;background:#0d1a2cdd;box-shadow:0 20px 45px #0000001f}.panel{padding:20px}.pipeline{display:grid;grid-template-columns:repeat(4,1fr);gap:10px;margin-top:16px}.flow-step{position:relative;padding:14px;min-height:116px}.flow-step:not(:last-child):after{content:"";position:absolute;top:35px;right:-10px;width:10px;height:2px;background:#3c5d83}.flow-step b{display:block;color:#f2f7ff}.flow-step small{display:block;margin-top:6px;color:#91a8c6}.flow-step .dot{display:inline-block;width:9px;height:9px;margin-right:7px;border-radius:50%;background:#5ee6a6}.metric{padding:16px}.metric b{display:block;margin-top:7px;color:#f5f8ff;font-size:2rem;line-height:1}.metric span{color:#9bb0cc;font-size:.84rem}.section-head{display:flex;align-items:baseline;justify-content:space-between;gap:12px;margin:36px 0 12px}.section-head h2{margin:0;color:#f3f7ff;font-size:1.18rem}.section-head p{margin:0;color:#8fa6c3;font-size:.9rem}.jobs{display:grid;gap:10px}.job{display:grid;grid-template-columns:minmax(210px,1.35fr) minmax(150px,.9fr) minmax(180px,1fr) auto;gap:18px;align-items:center;padding:17px 18px;text-decoration:none;color:inherit;transition:border-color .15s,transform .15s}.job:hover{border-color:#5d96ee;transform:translateY(-1px)}.job-name{color:#f3f7ff;font-weight:750}.job-id{margin-top:3px;color:#8196b3;font-family:ui-monospace,SFMono-Regular,monospace;font-size:.76rem}.badge{display:inline-flex;align-items:center;border-radius:999px;padding:4px 9px;font-size:.78rem;font-weight:800}.badge-waiting{background:#23344d;color:#b9cce9}.badge-active{background:#173d77;color:#9fc7ff}.badge-success{background:#123f34;color:#76efb5}.badge-danger{background:#552334;color:#ff9bad}.bar{height:7px;margin-top:8px;overflow:hidden;border-radius:999px;background:#20344e}.bar span{display:block;height:100%;border-radius:inherit;background:linear-gradient(90deg,#5a92ff,#61e1bd)}.numbers{color:#afc0d9;font-size:.88rem}.arrow{color:#79aaff;font-size:1.35rem}.workers{display:grid;grid-template-columns:repeat(3,1fr);gap:11px}.worker{padding:15px}.worker-head{display:flex;justify-content:space-between;gap:8px}.worker strong{color:#f3f7ff}.worker small,.worker p{color:#95a9c4}.worker p{margin:12px 0 0}.cap{display:inline-block;margin:4px 5px 0 0;border:1px solid #365576;border-radius:5px;padding:2px 6px;color:#a9c9f4;font:.75rem ui-monospace,SFMono-Regular,monospace}.empty{padding:30px;border:1px dashed #35516f;border-radius:14px;color:#9ab0cb;text-align:center}.offline{color:#faafbd}.sr{position:absolute;width:1px;height:1px;clip:rect(0,0,0,0);overflow:hidden;white-space:nowrap}@media(max-width:820px){.top,.section-head{display:block}.button{margin-top:18px}.summary{grid-template-columns:1fr 1fr}.panel{grid-column:span 2}.pipeline{grid-template-columns:1fr 1fr}.flow-step:not(:last-child):after{display:none}.job{grid-template-columns:1fr 1fr}.arrow{display:none}.workers{grid-template-columns:1fr 1fr}}@media(max-width:540px){.page{padding:22px 14px}.summary,.workers{grid-template-columns:1fr}.panel{grid-column:auto}.pipeline,.job{grid-template-columns:1fr}.title{font-size:2.35rem}}
</style>
</head>
<body>
<main class="page">
<header class="top">
<div><p class="eyebrow">Local scientific compute</p><h1 class="title">SciMesh control room</h1><p class="lead">Follow the real path from a molecular TSV to a globally reduced similarity result—without reading coordinator logs.</p><div class="live"><i class="pulse"></i><span id="refresh-state">Live overview · refreshes every 2 seconds</span></div></div>
<a class="button" href="/ui/jobs/new"> New similarity search</a>
</header>
<section class="summary" aria-label="Pipeline summary">
<div class="panel"><strong>How a search becomes a result</strong><div class="pipeline"><div class="flow-step"><span><i class="dot"></i>01</span><b>Upload TSV</b><small>The coordinator validates and slices the dataset.</small></div><div class="flow-step"><span><i class="dot"></i>02</span><b>Run shards</b><small>Workers fingerprint molecules and return shard top-k CSVs.</small></div><div class="flow-step"><span><i class="dot"></i>03</span><b>Merge exactly</b><small>The coordinator ranks retained candidates deterministically.</small></div><div class="flow-step"><span><i class="dot"></i>04</span><b>Download CSV</b><small>A checksum-protected global result is ready.</small></div></div></div>
<div class="metric"><span>Active runs</span><b id="active-jobs">{{.ActiveJobs}}</b><small>waiting, running, or merging</small></div>
<div class="metric"><span>Available workers</span><b id="online-workers">{{.OnlineWorkers}}</b><small>recently registered</small></div>
<div class="metric"><span>Finished runs</span><b id="finished-jobs">{{.FinishedJobs}}</b><small>in the latest 20</small></div>
</section>
<section><div class="section-head"><h2>Recent computations</h2><p id="job-count">{{len .Jobs}} shown · newest first</p></div><div id="jobs" class="jobs">{{range .Jobs}}<a class="job" href="/ui/jobs/{{.ID}}"><div><div class="job-name">{{workloadLabel .Workload}}</div><div class="job-id">{{.ID}}</div></div><div><span class="badge badge-{{statusClass .Status}}">{{statusLabel .Status}}</span><div class="job-id">{{statusHint .Status}}</div></div><div><div class="numbers"><b>{{.Completed}}</b> / {{.Total}} shards complete{{if gt .Failed 0}} · <span class="offline">{{.Failed}} failed</span>{{end}}</div><div class="bar"><span style="width:{{progressPercent .Completed .Failed .Cancelled .Total}}%"></span></div></div><span class="arrow" aria-hidden="true"></span></a>{{else}}<div class="empty"><strong>No computations yet.</strong><br>Start a small similarity search, then keep one or more workers running to watch this dashboard come alive.</div>{{end}}</div></section>
<section><div class="section-head"><h2>Worker fleet</h2><p>Workers register themselves; this page never controls their processes.</p></div><div id="workers" class="workers">{{range .Workers}}<article class="worker"><div class="worker-head"><strong>{{.Name}}</strong><span class="badge badge-{{workerStatusClass .Status}}">{{workerStatusLabel .Status}}</span></div><small>{{.ID}}</small><p>{{range .Capabilities}}<span class="cap">{{.}}</span>{{end}}</p><p>Last signal · {{time .LastHeartbeatAt}}</p></article>{{else}}<div class="empty"><strong>No worker is registered.</strong><br>Start <code>scimesh-worker</code> in another terminal, then return here.</div>{{end}}</div></section>
</main>
<script>
const statusInfo={pending:['Waiting for a worker','waiting'],leased:['Assigned to a worker','active'],running:['Running','active'],reducing:['Merging results','active'],completed:['Completed','success'],failed:['Needs attention','danger'],cancelled:['Stopped','waiting']};
const pct=j=>j.total?Math.min(100,Math.floor((j.completed+j.failed+j.cancelled)*100/j.total)):0;
const node=(tag,text,cls)=>{const n=document.createElement(tag);if(text!==undefined)n.textContent=text;if(cls)n.className=cls;return n};
const renderJobs=jobs=>{const box=document.querySelector('#jobs');box.replaceChildren();if(!jobs.length){const empty=node('div',undefined,'empty');empty.append(node('strong','No computations yet.'),document.createElement('br'),document.createTextNode('Start a small similarity search, then keep one or more workers running to watch this dashboard come alive.'));box.append(empty);return}for(const job of jobs){const info=statusInfo[job.status]||[job.status,'waiting'],link=node('a',undefined,'job');link.href='/ui/jobs/'+encodeURIComponent(job.id);const intro=node('div');intro.append(node('div',job.workload==='similarity-search'?'Molecule similarity search':job.workload,'job-name'),node('div',job.id,'job-id'));const state=node('div');state.append(node('span',info[0],'badge badge-'+info[1]),node('div',job.status==='reducing'?'Every shard is complete; coordinator is ranking the global top-k.':'Live coordinator state','job-id'));const progress=node('div'),numbers=node('div',undefined,'numbers');numbers.append(node('b',String(job.completed)),document.createTextNode(' / '+job.total+' shards complete'));if(job.failed){numbers.append(document.createTextNode(' · '),node('span',job.failed+' failed','offline'))}const bar=node('div',undefined,'bar'),fill=node('span');fill.style.width=pct(job)+'%';bar.append(fill);progress.append(numbers,bar);link.append(intro,state,progress,node('span','→','arrow'));box.append(link)}};
const renderWorkers=workers=>{const box=document.querySelector('#workers');box.replaceChildren();if(!workers.length){const empty=node('div',undefined,'empty');empty.append(node('strong','No worker is registered.'),document.createElement('br'),document.createTextNode('Start scimesh-worker in another terminal, then return here.'));box.append(empty);return}for(const worker of workers){const card=node('article',undefined,'worker'),head=node('div',undefined,'worker-head'),left=node('div'),workerInfo=worker.status==='online'?['Available','success']:worker.status==='busy'?['Busy','active']:['Offline','waiting'];left.append(node('strong',worker.name),node('small',worker.id));head.append(left,node('span',workerInfo[0],'badge badge-'+workerInfo[1]));const caps=node('p');for(const capability of worker.capabilities||[])caps.append(node('span',capability,'cap'));card.append(head,caps,node('p','Last signal · '+new Date(worker.last_heartbeat_at).toLocaleString()));box.append(card)}};
let timer;const refresh=async()=>{try{const response=await fetch('/ui/api/overview',{headers:{Accept:'application/json'}});if(!response.ok)throw Error();const view=await response.json();document.querySelector('#active-jobs').textContent=view.active_jobs;document.querySelector('#online-workers').textContent=view.online_workers;document.querySelector('#finished-jobs').textContent=view.finished_jobs;document.querySelector('#job-count').textContent=view.jobs.length+' shown · newest first';renderJobs(view.jobs);renderWorkers(view.workers);document.querySelector('#refresh-state').textContent='Live overview · updated just now'}catch(_){document.querySelector('#refresh-state').textContent='Connection interrupted · trying again automatically'}};
const start=()=>{if(!timer&&!document.hidden)timer=setInterval(refresh,2000)};document.addEventListener('visibilitychange',()=>{if(document.hidden&&timer){clearInterval(timer);timer=undefined}else start()});start();
</script>
</body>
</html>
{{end}}
File diff suppressed because one or more lines are too long
@@ -0,0 +1,24 @@
{{define "new-job.html"}}
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>New similarity search · SciMesh</title>
<style>
:root{color:#e5efff;background:#08111f;font:16px/1.5 Inter,ui-sans-serif,system-ui,sans-serif;color-scheme:dark}*{box-sizing:border-box}body{margin:0;background:radial-gradient(circle at 10% -8%,#183f77 0,transparent 32rem),#08111f}.page{max-width:980px;margin:auto;padding:28px 22px 64px}a{color:#94bdff}.back{text-decoration:none}.eyebrow{margin:28px 0 4px;color:#7baaff;font-size:.78rem;font-weight:800;letter-spacing:.14em;text-transform:uppercase}h1{margin:0;color:#f4f8ff;font-size:clamp(2rem,5vw,3.25rem);letter-spacing:-.055em}.lead{max-width:720px;margin:10px 0 0;color:#aabed9;font-size:1.06rem}.layout{display:grid;grid-template-columns:1.45fr .8fr;gap:15px;margin-top:28px}.card,.aside,.notice{border:1px solid #294662;border-radius:15px;background:#0d1a2cdd;box-shadow:0 20px 45px #00000021}.card{padding:22px}.aside,.notice{padding:18px}.aside h2,.notice h2{margin:0;color:#f1f6ff;font-size:1rem}.aside p,.notice p{color:#9fb3cf}.aside ol{margin:13px 0 0;padding-left:20px;color:#aebfda}.aside li{margin:10px 0}label{display:block;margin:18px 0 5px;color:#eaf2ff;font-weight:750}input,select{width:100%;border:1px solid #42617f;border-radius:9px;padding:10px 11px;background:#0a1626;color:#e5efff;font:inherit}input:focus,select:focus{outline:2px solid #5d97f5;outline-offset:1px}input[type=file]{padding:8px}.hint{margin:5px 0;color:#92a9c6;font-size:.87rem}.split{display:grid;grid-template-columns:1fr 1fr;gap:12px}.run-preview{margin-top:18px;border:1px solid #2c8c70;border-radius:11px;padding:13px;background:#0c2b2a;color:#a8f1d0}.run-preview strong{color:#e6fff4}.button{display:inline-flex;margin-top:22px;border:0;border-radius:10px;padding:12px 16px;background:#67e3b8;color:#062018;font:inherit;font-weight:850;cursor:pointer}.button:disabled{opacity:.6;cursor:wait}.working{margin:14px 0 0;color:#9fc5ff}.error{margin:12px 0 0;color:#ffacba}.hidden{display:none}.cap{display:inline-block;margin:4px 5px 0 0;border:1px solid #375978;border-radius:5px;padding:2px 6px;color:#b5d3f5;font:.76rem ui-monospace,SFMono-Regular,monospace}@media(max-width:720px){.layout,.split{grid-template-columns:1fr}.page{padding:22px 14px}}
</style>
</head>
<body>
<main class="page">
<a class="back" href="/ui">← Back to control room</a><p class="eyebrow">New computation</p><h1>Similarity search, end to end</h1><p class="lead">Upload a ChEMBL-style TSV. Workers calculate shard candidates; after every shard succeeds, SciMesh merges the exact global top-k into one final CSV.</p>
<div class="layout"><form id="run" class="card" novalidate><label for="file">ChEMBL TSV</label><input id="file" type="file" name="file" required accept=".tsv,.txt,text/tab-separated-values"><p class="hint">Required columns: <code>chembl_id</code> and <code>canonical_smiles</code>.</p><label for="query-smiles">Target molecule (SMILES)</label><input id="query-smiles" name="query_smiles" required maxlength="200" value="CCO" autocomplete="off"><p class="hint">Use a valid SMILES. The coordinator shares this exact query with every shard.</p><div class="split"><div><label for="top-k">Global top-k</label><input id="top-k" name="top_k" type="number" min="1" max="100000" value="20" required><p class="hint">How many final molecules to retain.</p></div><div><label for="chunk-rows">Rows per shard</label><input id="chunk-rows" name="chunk_rows" type="number" min="1" max="100000" value="1000" required><p class="hint">Smaller shards make more visible tasks.</p></div></div><div class="split"><div><label for="threshold">Similarity threshold <small>(optional)</small></label><input id="threshold" name="threshold" type="number" min="0" max="1" step="0.01" placeholder="For example: 0.70"><p class="hint">Leave blank to rank every valid candidate.</p></div><div><label for="direction">Keep molecules</label><select id="direction" name="threshold_direction"><option value="greater">more similar (≥ threshold)</option><option value="less">less similar (≤ threshold)</option></select><p class="hint">“Less” helps explore dissimilar molecules.</p></div></div><label for="max-rows">Maximum dataset rows <small>(optional quick run)</small></label><input id="max-rows" name="max_rows" type="number" min="1" max="10000000" placeholder="For example: 500"><p class="hint">Only the first N data rows become shards; the original upload remains stored by the coordinator.</p><div id="preview" class="run-preview"><strong>Ready to plan a run.</strong><br>Select a TSV to see the file that will be sent to the coordinator.</div><button class="button" id="submit" type="submit">Create pipeline run →</button><p id="working" class="working hidden" aria-live="polite">Uploading TSV and creating coordinator-owned shards…</p><p id="error" class="error" role="alert"></p></form><aside class="aside"><h2>What you will observe</h2><ol><li><strong>Input accepted</strong><br>TSV is validated and split into durable shard artifacts.</li><li><strong>Workers claim tasks</strong><br>Each worker downloads one shard, fingerprints it, and uploads a partial CSV.</li><li><strong>Global reduction</strong><br>The coordinator compares exact scores from all partial results.</li><li><strong>Final download</strong><br>The result page exposes a checksum-protected global CSV.</li></ol><h2 style="margin-top:26px">Before you submit</h2><p>Keep at least one worker running in another terminal. The browser cannot start or control worker processes.</p><p><span class="cap">similarity-search</span> is currently the only distributed workload available here.</p></aside></div>
</main>
<script>
const form=document.querySelector('#run'),button=document.querySelector('#submit'),working=document.querySelector('#working'),error=document.querySelector('#error'),preview=document.querySelector('#preview'),fileInput=document.querySelector('#file');
const bytes=size=>size<1024?size+' B':size<1024*1024?(size/1024).toFixed(1)+' KiB':(size/(1024*1024)).toFixed(1)+' MiB';fileInput.addEventListener('change',()=>{const file=fileInput.files&&fileInput.files[0];preview.replaceChildren();if(!file){preview.append(document.createTextNode('Select a TSV to see the file that will be sent to the coordinator.'));return}const strong=document.createElement('strong');strong.textContent='Source ready: '+file.name;preview.append(strong,document.createElement('br'),document.createTextNode(bytes(file.size)+' · the coordinator will validate its header before creating tasks.'))});
form.addEventListener('submit',async event=>{event.preventDefault();error.textContent='';const fields=new FormData(form),file=fields.get('file'),query=String(fields.get('query_smiles')||'').trim(),topK=Number(fields.get('top_k')),chunkRows=Number(fields.get('chunk_rows')),threshold=String(fields.get('threshold')||'').trim(),maxRows=String(fields.get('max_rows')||'').trim();if(!(file instanceof File)||file.size===0){error.textContent='Choose a non-empty TSV file.';return}if(!query||query.length>200||!Number.isInteger(topK)||topK<1||!Number.isInteger(chunkRows)||chunkRows<1){error.textContent='Enter a target SMILES, a positive global top-k, and a positive rows-per-shard value.';return}if(threshold&&(Number.isNaN(Number(threshold))||Number(threshold)<0||Number(threshold)>1)){error.textContent='Similarity threshold must be between 0 and 1.';return}const parameters={query_smiles:query,top_k:topK,threshold_direction:fields.get('threshold_direction'),progress_every:0};if(threshold)parameters.threshold=Number(threshold);const upload=new FormData();upload.append('workload','similarity-search');upload.append('parameters',JSON.stringify(parameters));upload.append('chunk_rows',String(chunkRows));if(maxRows)upload.append('max_rows',maxRows);upload.append('file',file,file.name);button.disabled=true;working.classList.remove('hidden');try{const response=await fetch('/ui/api/jobs/upload',{method:'POST',body:upload}),data=await response.json();if(!response.ok)throw Error(data.error||'Unable to create the job.');location.href='/ui/jobs/'+encodeURIComponent(data.job_id)}catch(err){error.textContent=err.message==='invalid input'?'The coordinator could not accept this run. Check the TSV columns and form values.':err.message;button.disabled=false;working.classList.add('hidden')}});
</script>
</body>
</html>
{{end}}
+313
View File
@@ -0,0 +1,313 @@
package http
import (
"embed"
"fmt"
"html/template"
"io"
"mime"
"net/http"
"strconv"
"time"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
)
//go:embed templates/*.html
var uiFiles embed.FS
var uiTemplates = template.Must(template.New("ui").Funcs(template.FuncMap{
"time": formatUITime,
"statusLabel": uiStatusLabel,
"statusHint": uiStatusHint,
"statusClass": uiStatusClass,
"taskErrorLabel": uiTaskErrorLabel,
"taskErrorHint": uiTaskErrorHint,
"workerStatusLabel": uiWorkerStatusLabel,
"workerStatusClass": uiWorkerStatusClass,
"workloadLabel": uiWorkloadLabel,
"progressPercent": uiProgressPercent,
"cancellable": uiCancellable,
"bytes": uiBytes,
"add": func(a, b int) int { return a + b },
}).ParseFS(uiFiles, "templates/*.html"))
func formatUITime(t time.Time) string {
if t.IsZero() {
return "—"
}
return t.UTC().Format("02.01.2006 15:04 UTC")
}
func uiStatusLabel(status string) string {
switch status {
case "pending":
return "Waiting for a worker"
case "leased":
return "Assigned to a worker"
case "running":
return "Running"
case "reducing":
return "Merging results"
case "completed":
return "Completed"
case "failed":
return "Needs attention"
case "cancelled":
return "Stopped"
default:
return status
}
}
func uiStatusHint(status string) string {
switch status {
case "pending":
return "Waiting for an available worker with the required capability."
case "leased":
return "A worker has claimed the task and should begin processing shortly."
case "running":
return "A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator."
case "reducing":
return "All shards are complete. The coordinator is merging their candidates into one final CSV."
case "completed":
return "The final result is ready to download."
case "failed":
return "One or more shard tasks failed. Open the task list below for details."
case "cancelled":
return "The operator stopped this job. No new shards can be claimed."
default:
return "Status reported by the coordinator."
}
}
func uiStatusClass(status string) string {
switch status {
case "completed":
return "success"
case "failed":
return "danger"
case "cancelled":
return "waiting"
case "running", "leased", "reducing":
return "active"
default:
return "waiting"
}
}
func uiWorkerStatusLabel(status string) string {
switch status {
case "online":
return "Available"
case "busy":
return "Busy"
case "offline":
return "Offline"
default:
return status
}
}
func uiWorkerStatusClass(status string) string {
switch status {
case "online":
return "success"
case "busy":
return "active"
default:
return "waiting"
}
}
// uiTaskErrorLabel deliberately maps worker implementation errors to an
// operator-facing diagnosis. Raw subprocess commands and local paths belong in
// the worker terminal, not in the web UI.
func uiTaskErrorLabel(errorCode string) string {
switch errorCode {
case "CalledProcessError":
return "Local calculation failed"
case "ValueError":
return "Task input could not be processed"
case "CoordinatorTransientError":
return "Coordinator connection was interrupted"
case "CoordinatorConflictError":
return "Worker lease was no longer valid"
case "FileNotFoundError":
return "Local task file is missing"
default:
return errorCode
}
}
func uiTaskErrorHint(errorCode string) string {
switch errorCode {
case "CalledProcessError":
return "The local SciMesh command stopped before it could upload a result. Check the worker terminal for the original error."
case "ValueError":
return "The coordinator task or its downloaded input did not meet the worker validation rules."
case "CoordinatorTransientError":
return "The worker will retry after the coordinator connection is available again."
case "CoordinatorConflictError":
return "Another worker or a lease timeout changed this task before completion."
case "FileNotFoundError":
return "The worker could not find one of its local task files. Restart it with an absolute --work-dir."
default:
return "Check the worker terminal for the original error details."
}
}
func uiWorkloadLabel(workload string) string {
switch workload {
case "similarity-search", "similarity_search":
return "Molecule similarity search"
case "similarity-graph", "similarity_graph":
return "Molecular similarity graph"
default:
return workload
}
}
func uiCancellable(status string) bool {
return status == "pending" || status == "running"
}
func uiProgressPercent(completed, failed, cancelled, total int) int {
if total <= 0 {
return 0
}
percent := (completed + failed + cancelled) * 100 / total
if percent > 100 {
return 100
}
return percent
}
func uiBytes(n int64) string {
const kib = 1024
if n < kib {
return fmt.Sprintf("%d B", n)
}
if n < kib*kib {
return fmt.Sprintf("%.1f KiB", float64(n)/kib)
}
if n < kib*kib*kib {
return fmt.Sprintf("%.1f MiB", float64(n)/(kib*kib))
}
return fmt.Sprintf("%.1f GiB", float64(n)/(kib*kib*kib))
}
func (s *Server) renderUI(w http.ResponseWriter, name string, data any) {
w.Header().Set("Content-Type", "text/html; charset=utf-8")
w.Header().Set("Cache-Control", "no-store")
w.Header().Set("X-Content-Type-Options", "nosniff")
w.Header().Set("Content-Security-Policy", "default-src 'self'; style-src 'self' 'unsafe-inline'; script-src 'self' 'unsafe-inline'; base-uri 'none'; frame-ancestors 'none'")
if err := uiTemplates.ExecuteTemplate(w, name, data); err != nil {
s.log.Error("render UI", "err", err)
http.Error(w, "internal error", http.StatusInternalServerError)
}
}
func (s *Server) handleUIHome(w http.ResponseWriter, r *http.Request) {
ctx, cancel := s.reqCtx(r)
defer cancel()
view, err := s.uc.Dashboard.Overview(ctx, 20)
if err != nil {
s.writeError(w, r, err)
return
}
s.renderUI(w, "dashboard.html", view)
}
// handleUIOverviewJSON is the bounded polling projection used by the operator
// dashboard. It intentionally returns only the safe UI read model, never
// worker tokens, storage keys, or database entities.
func (s *Server) handleUIOverviewJSON(w http.ResponseWriter, r *http.Request) {
ctx, cancel := s.reqCtx(r)
defer cancel()
view, err := s.uc.Dashboard.Overview(ctx, 20)
if err != nil {
s.writeError(w, r, err)
return
}
writeJSON(w, http.StatusOK, view)
}
func (s *Server) handleUINewJob(w http.ResponseWriter, r *http.Request) {
s.renderUI(w, "new-job.html", nil)
}
func (s *Server) uiJobID(w http.ResponseWriter, r *http.Request) (uuid.UUID, bool) {
return s.pathUUID(w, r, "job_id")
}
func (s *Server) handleUIJob(w http.ResponseWriter, r *http.Request) {
jobID, ok := s.uiJobID(w, r)
if !ok {
return
}
ctx, cancel := s.reqCtx(r)
defer cancel()
view, err := s.uc.Dashboard.JobDetail(ctx, jobID)
if err != nil {
s.writeError(w, r, err)
return
}
s.renderUI(w, "job.html", view)
}
func (s *Server) handleUIJobJSON(w http.ResponseWriter, r *http.Request) {
jobID, ok := s.uiJobID(w, r)
if !ok {
return
}
ctx, cancel := s.reqCtx(r)
defer cancel()
view, err := s.uc.Dashboard.JobDetail(ctx, jobID)
if err != nil {
s.writeError(w, r, err)
return
}
writeJSON(w, http.StatusOK, view)
}
func (s *Server) handleUIArtifactDownload(w http.ResponseWriter, r *http.Request) {
jobID, ok := s.uiJobID(w, r)
if !ok {
return
}
artifactID, err := uuid.Parse(r.PathValue("artifact_id"))
if err != nil {
s.writeError(w, r, domain.ErrInvalidInput)
return
}
ctx, cancel := s.reqCtx(r)
defer cancel()
belongs, err := s.uc.Dashboard.ArtifactBelongsToJob(ctx, jobID, artifactID)
if err != nil {
s.writeError(w, r, err)
return
}
if !belongs {
s.writeError(w, r, domain.ErrArtifactNotFound)
return
}
// Reuse the coordinator-owned blob stream after the job-scoped check above.
art, body, err := s.uc.DownloadArtifact.Execute(ctx, artifactID)
if err != nil {
s.writeError(w, r, err)
return
}
defer func() {
if err := body.Close(); err != nil {
s.log.Warn("close downloaded UI artifact", "artifact_id", artifactID, "err", err)
}
}()
w.Header().Set("Content-Type", art.ContentType)
w.Header().Set("Content-Disposition", mime.FormatMediaType("attachment", map[string]string{"filename": art.Filename}))
w.Header().Set("Content-Length", strconv.FormatInt(art.SizeBytes, 10))
w.Header().Set("X-Checksum-SHA256", art.SHA256)
_, _ = io.Copy(w, body)
}
@@ -0,0 +1,54 @@
package http
import "testing"
func TestUIStatusPresentation(t *testing.T) {
tests := []struct {
status string
label string
class string
}{
{"pending", "Waiting for a worker", "waiting"},
{"running", "Running", "active"},
{"reducing", "Merging results", "active"},
{"completed", "Completed", "success"},
{"failed", "Needs attention", "danger"},
}
for _, test := range tests {
t.Run(test.status, func(t *testing.T) {
if got := uiStatusLabel(test.status); got != test.label {
t.Errorf("label = %q, want %q", got, test.label)
}
if got := uiStatusClass(test.status); got != test.class {
t.Errorf("class = %q, want %q", got, test.class)
}
})
}
}
func TestUIProgressPercent(t *testing.T) {
if got := uiProgressPercent(3, 1, 0, 8); got != 50 {
t.Errorf("progress = %d, want 50", got)
}
if got := uiProgressPercent(1, 1, 0, 0); got != 0 {
t.Errorf("empty progress = %d, want 0", got)
}
}
func TestUITaskErrorPresentationDoesNotExposeCommand(t *testing.T) {
if got := uiTaskErrorLabel("CalledProcessError"); got != "Local calculation failed" {
t.Errorf("error label = %q", got)
}
if got := uiTaskErrorHint("CalledProcessError"); got == "" {
t.Error("error hint must explain the failure")
}
}
func TestUIWorkerStatusPresentation(t *testing.T) {
if got := uiWorkerStatusLabel("busy"); got != "Busy" {
t.Errorf("busy worker label = %q", got)
}
if got := uiWorkerStatusClass("busy"); got != "active" {
t.Errorf("busy worker class = %q", got)
}
}
+41 -15
View File
@@ -14,12 +14,13 @@ type UploadArtifact struct {
tasks TaskRepository
artifacts ArtifactRepository
blobs BlobStore
tx TxManager
clk Clock
}
func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository,
blobs BlobStore, clk Clock) *UploadArtifact {
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, clk: clk}
blobs BlobStore, tx TxManager, clk Clock) *UploadArtifact {
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, clk: clk}
}
func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (*domain.Artifact, error) {
@@ -32,6 +33,15 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (
if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
return nil, domain.ErrLeaseConflict
}
// A client can retry a PUT after losing the response. Return the one durable
// result for this lease attempt instead of storing duplicate artifacts.
existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt)
if err != nil {
return nil, err
}
if existing != nil {
return existing, nil
}
taskID := task.ID
art, err := domain.NewArtifact(task.JobID, &taskID, domain.ArtifactPartialResult,
@@ -50,25 +60,41 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (
}
art.SetContent(sum, size)
// The stream may take longer than the lease. Re-check after it finishes so
// an expired worker cannot leave a durable result record behind. Completion
// performs the same ownership check under its transaction.
current, err := uc.tasks.Get(ctx, in.TaskID)
// The stream may take longer than the lease. Lock the task while re-checking
// ownership and inserting metadata: completion or another upload cannot race
// this final decision. The database unique index is a second line of defence.
var durable *domain.Artifact
err = uc.tx.WithinTx(ctx, func(ctx context.Context) error {
current, err := uc.tasks.GetForUpdate(ctx, in.TaskID)
if err != nil {
return err
}
if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
return domain.ErrLeaseConflict
}
existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt)
if err != nil {
return err
}
if existing != nil {
durable = existing
return nil
}
if err := uc.artifacts.Insert(ctx, art); err != nil {
return err
}
durable = art
return nil
})
if err != nil {
_ = uc.blobs.Delete(ctx, art.StorageKey)
return nil, err
}
if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
if durable != art {
// Another request won the race while this stream was being written.
_ = uc.blobs.Delete(ctx, art.StorageKey)
return nil, domain.ErrLeaseConflict
}
// Persist the record. If that fails the blob would be an orphan, so remove it.
if err := uc.artifacts.Insert(ctx, art); err != nil {
_ = uc.blobs.Delete(ctx, art.StorageKey)
return nil, err
}
return art, nil
return durable, nil
}
// DownloadArtifact returns an artifact's metadata together with a reader over
+6 -3
View File
@@ -53,9 +53,12 @@ type SubmitDatasetInput struct {
Workload string
Parameters map[string]any
RowsPerShard int
Filename string
ContentType string
Body io.Reader
// MaxRows limits how many data rows are turned into shards. Zero means the
// whole uploaded dataset; the input artifact itself remains stored intact.
MaxRows int
Filename string
ContentType string
Body io.Reader
}
type SubmitDatasetResult struct {
+92 -5
View File
@@ -33,6 +33,17 @@ func NewCreateJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock
// The all-or-none guarantee comes from TxManager: a half-created job would
// leave chunks no worker could ever complete.
func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) {
if in.Workload == "similarity-graph" || in.Workload == "similarity_graph" {
// CTX-10 must plan triangular block pairs; ordinary independent input
// chunks would silently omit every cross-chunk molecular pair.
return nil, domain.ErrInvalidInput
}
if (in.Workload == "similarity-search" || in.Workload == "similarity_search") &&
len(in.Chunks) > 1 && in.Parameters["query_id"] != nil {
// Resolving once against the source dataset belongs to CTX-07. Letting
// each shard resolve it would make most tasks fail or use inconsistent data.
return nil, domain.ErrInvalidInput
}
chunks := make([]domain.ChunkSpec, 0, len(in.Chunks))
for _, c := range in.Chunks {
chunks = append(chunks, domain.ChunkSpec(c))
@@ -62,6 +73,56 @@ type GetJobStatus struct {
tasks TaskRepository
}
// --- CancelJob -----------------------------------------------------------
type CancelJob struct {
jobs JobRepository
tasks TaskRepository
tx TxManager
clock Clock
}
func NewCancelJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock Clock) *CancelJob {
return &CancelJob{jobs: jobs, tasks: tasks, tx: tx, clock: clock}
}
// Execute stops a job atomically. Completed and finally failed tasks are kept
// as historical evidence; all other tasks are cancelled, including leased and
// running ones. A repeated cancel of an already cancelled job is idempotent.
func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error) {
now := uc.clock.Now()
var cancelled int64
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
job, err := uc.jobs.Get(ctx, jobID)
if err != nil {
return err
}
if job.Status == domain.JobCancelled {
return nil
}
if job.Status == domain.JobReducing || job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
return domain.ErrJobNotCancellable
}
// The lease reaper can be the transition that exhausted the final task.
// Check the authoritative task histogram as well as the cached job status,
// so a stale status can never turn a failed/completed job into cancelled.
counts, err := uc.tasks.CountByStatus(ctx, jobID)
if err != nil {
return err
}
derived := progressFrom(*job, counts).DeriveStatus()
if derived == domain.JobReducing || derived == domain.JobCompleted || derived == domain.JobFailed {
return domain.ErrJobNotCancellable
}
cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
if err != nil {
return err
}
return uc.jobs.UpdateStatus(ctx, jobID, domain.JobCancelled, &now)
})
return cancelled, err
}
func NewGetJobStatus(jobs JobRepository, tasks TaskRepository) *GetJobStatus {
return &GetJobStatus{jobs: jobs, tasks: tasks}
}
@@ -144,9 +205,10 @@ func progressFrom(job domain.Job, counts map[domain.TaskStatus]int) domain.JobPr
Job: job,
Pending: counts[domain.TaskPending],
// Leased and running are both "in flight" for progress purposes.
Leased: counts[domain.TaskLeased] + counts[domain.TaskRunning],
Done: counts[domain.TaskCompleted],
Failed: counts[domain.TaskFailed],
Leased: counts[domain.TaskLeased] + counts[domain.TaskRunning],
Done: counts[domain.TaskCompleted],
Failed: counts[domain.TaskFailed],
Cancelled: counts[domain.TaskCancelled],
}
for _, n := range counts {
p.Total += n
@@ -164,11 +226,36 @@ func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository
if err != nil {
return err
}
status := progressFrom(domain.Job{}, counts).DeriveStatus()
job, err := jobs.Get(ctx, jobID)
if err != nil {
return err
}
status := progressFrom(*job, counts).DeriveStatus()
// All worker shards being complete means scientific reduction is ready, not
// that the job's final artifact already exists. CTX-09 owns the transition
// from reducing to completed after it persists that artifact.
if status == domain.JobCompleted && job.Workload == "similarity-search" {
status = domain.JobReducing
}
var completedAt *time.Time
if status == domain.JobCompleted || status == domain.JobFailed {
if status == domain.JobFailed {
completedAt = &now
}
return jobs.UpdateStatus(ctx, jobID, status, completedAt)
}
func syncExpiredJobStatuses(ctx context.Context, jobs JobRepository, tasks TaskRepository,
jobIDs []uuid.UUID, now time.Time) error {
seen := make(map[uuid.UUID]struct{}, len(jobIDs))
for _, jobID := range jobIDs {
if _, duplicate := seen[jobID]; duplicate {
continue
}
seen[jobID] = struct{}{}
if err := syncJobStatus(ctx, jobs, tasks, jobID, now); err != nil {
return err
}
}
return nil
}
+13 -3
View File
@@ -56,9 +56,13 @@ type TaskRepository interface {
// CountByStatus aggregates a job's tasks for progress reporting.
CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error)
// ExpireLeases applies the lease-expiry rule to every elapsed task and
// reports how many were affected.
ExpireLeases(ctx context.Context, now time.Time) (int64, error)
// CancelByJob marks every non-terminal task as cancelled and invalidates its
// lease. It returns how many tasks changed.
CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error)
// ExpireLeases applies the lease-expiry rule to every elapsed task and returns
// the distinct jobs whose aggregate status may have changed.
ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error)
}
// JobRepository persists jobs.
@@ -66,6 +70,9 @@ type JobRepository interface {
Insert(ctx context.Context, j *domain.Job) error
Get(ctx context.Context, id uuid.UUID) (*domain.Job, error)
UpdateStatus(ctx context.Context, id uuid.UUID, status domain.JobStatus, completedAt *time.Time) error
ClaimReduction(ctx context.Context, id uuid.UUID, startedAt time.Time) (bool, error)
CompleteWithResult(ctx context.Context, id, resultArtifactID uuid.UUID, completedAt time.Time) error
FailReduction(ctx context.Context, id uuid.UUID, code, message string, completedAt time.Time) error
}
// WorkerRepository persists the worker registry.
@@ -85,6 +92,9 @@ type WorkerRepository interface {
type ArtifactRepository interface {
Insert(ctx context.Context, a *domain.Artifact) error
Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error)
// FindPartialResult returns the durable result already uploaded for one task
// attempt. A nil artifact means the attempt has not uploaded one yet.
FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error)
}
// BlobStore holds artifact bytes, addressed by an opaque storage key. It streams
+133
View File
@@ -0,0 +1,133 @@
package usecase
import (
"bytes"
"context"
"io"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
"github.com/emil28092005/SciMesh/coordinator/internal/reducer"
)
// ReduceJob turns completed coordinator-owned partial artifacts into one final
// artifact. It performs no worker I/O and never trusts a worker URI or path.
type ReduceJob struct {
jobs JobRepository
tasks TaskRepository
artifacts ArtifactRepository
blobs BlobStore
tx TxManager
clock Clock
}
func NewReduceJob(jobs JobRepository, tasks TaskRepository, artifacts ArtifactRepository,
blobs BlobStore, tx TxManager, clock Clock) *ReduceJob {
return &ReduceJob{jobs: jobs, tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, clock: clock}
}
// Execute is idempotent for jobs that are not currently reducing. The worker
// completion path may call it after every result; only the last task changes a
// similarity-search job into reducing state.
func (uc *ReduceJob) Execute(ctx context.Context, jobID uuid.UUID) error {
claimed, err := uc.jobs.ClaimReduction(ctx, jobID, uc.clock.Now())
if err != nil || !claimed {
return err
}
job, err := uc.jobs.Get(ctx, jobID)
if err != nil {
return err
}
if job.Status != domain.JobReducing {
return nil
}
if job.Workload != "similarity-search" {
return uc.fail(ctx, jobID)
}
completed, err := uc.tasks.ListCompleted(ctx, jobID)
if err != nil {
return uc.fail(ctx, jobID)
}
if len(completed) == 0 {
return uc.fail(ctx, jobID)
}
readers := make([]io.Reader, 0, len(completed))
closers := make([]io.Closer, 0, len(completed))
for _, task := range completed {
if task.ResultArtifactID == nil {
closeAll(closers)
return uc.fail(ctx, jobID)
}
artifact, err := uc.artifacts.Get(ctx, *task.ResultArtifactID)
if err != nil || artifact.JobID != jobID || artifact.TaskID == nil || *artifact.TaskID != task.ID || artifact.Kind != domain.ArtifactPartialResult {
closeAll(closers)
return uc.fail(ctx, jobID)
}
body, err := uc.blobs.Open(ctx, artifact.StorageKey)
if err != nil {
closeAll(closers)
return uc.fail(ctx, jobID)
}
readers = append(readers, body)
closers = append(closers, body)
}
output, reduceErr := reducer.ReduceSimilaritySearch(readers, job.Parameters)
closeAll(closers)
if reduceErr != nil {
return uc.fail(ctx, jobID)
}
final, err := domain.NewArtifact(jobID, nil, domain.ArtifactFinalResult, "similarity-search.csv", "text/csv", uc.clock.Now())
if err != nil {
return uc.fail(ctx, jobID)
}
sum, size, err := uc.blobs.Put(ctx, final.StorageKey, bytes.NewReader(output))
if err != nil {
return uc.fail(ctx, jobID)
}
final.SetContent(sum, size)
if err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
if err := uc.artifacts.Insert(ctx, final); err != nil {
return err
}
return uc.jobs.CompleteWithResult(ctx, jobID, final.ID, uc.clock.Now())
}); err != nil {
_ = uc.blobs.Delete(ctx, final.StorageKey)
return err
}
return nil
}
func (uc *ReduceJob) fail(ctx context.Context, jobID uuid.UUID) error {
// The public state carries a stable sanitized failure, never parser/storage
// internals that may include local paths or implementation details.
return uc.jobs.FailReduction(ctx, jobID, "reducer_failed", "final result reduction failed", uc.clock.Now())
}
func closeAll(closers []io.Closer) {
for _, closer := range closers {
_ = closer.Close()
}
}
type GetJobResult struct {
jobs JobRepository
download *DownloadArtifact
}
func NewGetJobResult(jobs JobRepository, download *DownloadArtifact) *GetJobResult {
return &GetJobResult{jobs: jobs, download: download}
}
func (uc *GetJobResult) Execute(ctx context.Context, jobID uuid.UUID) (*domain.Artifact, io.ReadCloser, error) {
job, err := uc.jobs.Get(ctx, jobID)
if err != nil {
return nil, nil, err
}
if job.Status != domain.JobCompleted || job.ResultArtifactID == nil {
return nil, nil, domain.ErrArtifactNotFound
}
return uc.download.Execute(ctx, *job.ResultArtifactID)
}
+55 -21
View File
@@ -21,12 +21,15 @@ import (
type ClaimTask struct {
tasks TaskRepository
jobs JobRepository
workers WorkerRepository
tx TxManager
clock Clock
leaseDuration time.Duration
}
func NewClaimTask(tasks TaskRepository, clock Clock, leaseDuration time.Duration) *ClaimTask {
return &ClaimTask{tasks: tasks, clock: clock, leaseDuration: leaseDuration}
func NewClaimTask(tasks TaskRepository, jobs JobRepository, workers WorkerRepository, tx TxManager, clock Clock, leaseDuration time.Duration) *ClaimTask {
return &ClaimTask{tasks: tasks, jobs: jobs, workers: workers, tx: tx, clock: clock, leaseDuration: leaseDuration}
}
// Execute reclaims elapsed leases first, then hands out one task.
@@ -42,27 +45,46 @@ func (uc *ClaimTask) Execute(ctx context.Context, in ClaimTaskInput) (*domain.Cl
if in.WorkerID == "" {
return nil, domain.ErrInvalidInput
}
now := uc.clock.Now()
if _, err := uc.tasks.ExpireLeases(ctx, now); err != nil {
return nil, err
workloads := in.Workloads
if workerID, err := uuid.Parse(in.WorkerID); err == nil {
worker, err := uc.workers.Get(ctx, workerID)
if err != nil {
return nil, err
}
// Never trust caller-supplied capabilities: registration is the durable
// worker identity and its allowlist.
workloads = worker.Capabilities
}
var claimed *domain.ClaimedTask
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
now := uc.clock.Now()
affectedJobs, err := uc.tasks.ExpireLeases(ctx, now)
if err != nil {
return err
}
if err := syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affectedJobs, now); err != nil {
return err
}
task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
Workloads: in.Workloads,
Owner: in.WorkerID,
Now: now,
LeaseUntil: now.Add(uc.leaseDuration),
task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
Workloads: workloads,
Owner: in.WorkerID,
Now: now,
LeaseUntil: now.Add(uc.leaseDuration),
})
if err != nil {
return err
}
if task != nil {
value := task.AsClaimed()
claimed = &value
}
return nil
})
if err != nil {
return nil, err
}
if task == nil {
return nil, nil // empty queue is a normal state, not an error
}
claimed := task.AsClaimed()
return &claimed, nil
return claimed, nil // nil means an empty queue
}
// --- RenewLease ----------------------------------------------------------
@@ -230,18 +252,30 @@ func (uc *FailTask) Execute(ctx context.Context, in FailTaskInput) (*domain.Task
type ExpireLeases struct {
tasks TaskRepository
jobs JobRepository
tx TxManager
clock Clock
}
func NewExpireLeases(tasks TaskRepository, clock Clock) *ExpireLeases {
return &ExpireLeases{tasks: tasks, clock: clock}
func NewExpireLeases(tasks TaskRepository, jobs JobRepository, tx TxManager, clock Clock) *ExpireLeases {
return &ExpireLeases{tasks: tasks, jobs: jobs, tx: tx, clock: clock}
}
// Execute reports how many tasks were reclaimed.
// Execute reclaims elapsed tasks and persists the state of every affected job.
//
// The sweep is one set-based statement rather than a load-decide-save loop:
// several coordinators run it concurrently, and a single atomic UPDATE makes
// the duplicate work harmless — the loser simply updates 0 rows.
func (uc *ExpireLeases) Execute(ctx context.Context) (int64, error) {
return uc.tasks.ExpireLeases(ctx, uc.clock.Now())
var affected []uuid.UUID
err := uc.tx.WithinTx(ctx, func(ctx context.Context) error {
now := uc.clock.Now()
var err error
affected, err = uc.tasks.ExpireLeases(ctx, now)
if err != nil {
return err
}
return syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affected, now)
})
return int64(len(affected)), err
}
+276
View File
@@ -0,0 +1,276 @@
package usecase
import (
"context"
"fmt"
"time"
"github.com/google/uuid"
"github.com/emil28092005/SciMesh/coordinator/internal/domain"
)
// UIReadRepository is a read-only projection source for the local operator UI.
// It intentionally exposes no storage paths or credentials.
type UIReadRepository interface {
GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error)
ListJobs(ctx context.Context, limit int) ([]domain.Job, error)
ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error)
ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error)
ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error)
ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error)
}
type JobCard struct {
ID string `json:"id"`
Workload string `json:"workload"`
Status string `json:"status"`
CreatedAt time.Time `json:"created_at"`
CompletedAt *time.Time `json:"completed_at,omitempty"`
ReducerStartedAt *time.Time `json:"reducer_started_at,omitempty"`
ErrorCode string `json:"error_code,omitempty"`
ErrorMessage string `json:"error_message,omitempty"`
Total int `json:"total"`
Pending int `json:"pending"`
Leased int `json:"leased"`
Running int `json:"running"`
Completed int `json:"completed"`
Failed int `json:"failed"`
Cancelled int `json:"cancelled"`
}
type TaskCard struct {
ID string `json:"id"`
ChunkIndex int `json:"chunk_index"`
Status string `json:"status"`
Attempt int `json:"attempt"`
MaxAttempts int `json:"max_attempts"`
LeaseOwner string `json:"lease_owner,omitempty"`
LeaseExpiresAt *time.Time `json:"lease_expires_at,omitempty"`
StartedAt *time.Time `json:"started_at,omitempty"`
CompletedAt *time.Time `json:"completed_at,omitempty"`
ErrorCode string `json:"error_code,omitempty"`
ErrorMessage string `json:"error_message,omitempty"`
}
// ParameterCard is an intentionally small allowlist of run configuration that
// helps an operator verify what is being computed without exposing arbitrary
// job payloads to the browser.
type ParameterCard struct {
Label string `json:"label"`
Value string `json:"value"`
}
type ArtifactCard struct {
ID string `json:"id"`
Kind string `json:"kind"`
Filename string `json:"filename"`
SizeBytes int64 `json:"size_bytes"`
SHA256 string `json:"sha256"`
Downloadable bool `json:"downloadable"`
Diagnostic bool `json:"diagnostic"`
}
type WorkerCard struct {
ID string `json:"id"`
Name string `json:"name"`
Status string `json:"status"`
Capabilities []string `json:"capabilities"`
LastHeartbeatAt time.Time `json:"last_heartbeat_at"`
}
type DashboardView struct {
Jobs []JobCard `json:"jobs"`
Workers []WorkerCard `json:"workers"`
ActiveJobs int `json:"active_jobs"`
FinishedJobs int `json:"finished_jobs"`
OnlineWorkers int `json:"online_workers"`
}
type JobDetailView struct {
JobCard
Tasks []TaskCard `json:"tasks"`
Artifacts []ArtifactCard `json:"artifacts"`
Parameters []ParameterCard `json:"parameters"`
FinalResultAvailable bool `json:"final_result_available"`
}
type Dashboard struct{ read UIReadRepository }
func NewDashboard(read UIReadRepository) *Dashboard { return &Dashboard{read: read} }
func (d *Dashboard) Overview(ctx context.Context, limit int) (DashboardView, error) {
jobs, err := d.read.ListJobs(ctx, limit)
if err != nil {
return DashboardView{}, err
}
workers, err := d.read.ListWorkers(ctx, limit)
if err != nil {
return DashboardView{}, err
}
out := DashboardView{Jobs: make([]JobCard, 0, len(jobs)), Workers: make([]WorkerCard, 0, len(workers))}
jobIDs := make([]uuid.UUID, 0, len(jobs))
for _, job := range jobs {
jobIDs = append(jobIDs, job.ID)
}
tasksByJob, err := d.read.ListTasksByJobs(ctx, jobIDs)
if err != nil {
return DashboardView{}, err
}
for _, job := range jobs {
card := jobCard(job, tasksByJob[job.ID])
out.Jobs = append(out.Jobs, card)
switch card.Status {
case string(domain.JobCompleted), string(domain.JobFailed), string(domain.JobCancelled):
out.FinishedJobs++
default:
out.ActiveJobs++
}
}
for _, worker := range workers {
out.Workers = append(out.Workers, WorkerCard{ID: worker.ID.String(), Name: worker.Name, Status: string(worker.Status), Capabilities: worker.Capabilities, LastHeartbeatAt: worker.LastHeartbeatAt})
if worker.Status == domain.WorkerOnline || worker.Status == domain.WorkerBusy {
out.OnlineWorkers++
}
}
return out, nil
}
func (d *Dashboard) JobDetail(ctx context.Context, jobID uuid.UUID) (JobDetailView, error) {
job, err := d.read.GetJob(ctx, jobID)
if err != nil {
return JobDetailView{}, err
}
tasks, err := d.read.ListTasksByJob(ctx, jobID)
if err != nil {
return JobDetailView{}, err
}
artifacts, err := d.read.ListArtifactsByJob(ctx, jobID)
if err != nil {
return JobDetailView{}, err
}
workers, err := d.read.ListWorkers(ctx, 100)
if err != nil {
return JobDetailView{}, err
}
workerNames := make(map[string]string, len(workers))
for _, worker := range workers {
workerNames[worker.ID.String()] = worker.Name
}
out := JobDetailView{
JobCard: jobCard(*job, tasks),
Tasks: make([]TaskCard, 0, len(tasks)),
Artifacts: make([]ArtifactCard, 0, len(artifacts)),
Parameters: uiParameters(job.Parameters),
}
for _, task := range tasks {
card := TaskCard{ID: task.ID.String(), ChunkIndex: task.ChunkIndex, Status: string(task.Status), Attempt: task.Attempt, MaxAttempts: task.MaxAttempts, LeaseExpiresAt: task.LeaseExpiresAt, StartedAt: task.StartedAt, CompletedAt: task.CompletedAt}
if task.LeaseOwner != nil {
card.LeaseOwner = workerNames[*task.LeaseOwner]
if card.LeaseOwner == "" {
card.LeaseOwner = "Worker " + shortID(*task.LeaseOwner)
}
}
if task.ErrorCode != nil {
card.ErrorCode = *task.ErrorCode
}
if task.ErrorMessage != nil {
card.ErrorMessage = *task.ErrorMessage
}
out.Tasks = append(out.Tasks, card)
}
for _, artifact := range artifacts {
diagnostic := artifact.Kind == domain.ArtifactPartialResult
downloadable := diagnostic || (artifact.Kind == domain.ArtifactFinalResult && out.Status == string(domain.JobCompleted))
out.Artifacts = append(out.Artifacts, ArtifactCard{ID: artifact.ID.String(), Kind: string(artifact.Kind), Filename: artifact.Filename, SizeBytes: artifact.SizeBytes, SHA256: artifact.SHA256, Downloadable: downloadable, Diagnostic: diagnostic})
if artifact.Kind == domain.ArtifactFinalResult && downloadable {
out.FinalResultAvailable = true
}
}
return out, nil
}
func (d *Dashboard) ArtifactBelongsToJob(ctx context.Context, jobID, artifactID uuid.UUID) (bool, error) {
artifacts, err := d.read.ListArtifactsByJob(ctx, jobID)
if err != nil {
return false, err
}
for _, a := range artifacts {
if a.ID == artifactID {
return true, nil
}
}
return false, nil
}
func jobCard(job domain.Job, tasks []domain.Task) JobCard {
c := JobCard{ID: job.ID.String(), Workload: job.Workload, CreatedAt: job.CreatedAt, CompletedAt: job.CompletedAt, ReducerStartedAt: job.ReducerStartedAt}
if job.ErrorCode != nil {
c.ErrorCode = *job.ErrorCode
}
if job.ErrorMessage != nil {
c.ErrorMessage = *job.ErrorMessage
}
for _, task := range tasks {
c.Total++
switch task.Status {
case domain.TaskPending:
c.Pending++
case domain.TaskLeased:
c.Leased++
case domain.TaskRunning:
c.Running++
case domain.TaskCompleted:
c.Completed++
case domain.TaskFailed:
c.Failed++
case domain.TaskCancelled:
c.Cancelled++
}
}
p := domain.JobProgress{Job: job, Total: c.Total, Pending: c.Pending, Leased: c.Leased + c.Running, Done: c.Completed, Failed: c.Failed, Cancelled: c.Cancelled}
c.Status = string(p.DeriveStatus())
return c
}
func uiParameters(parameters map[string]any) []ParameterCard {
keys := []struct {
key string
label string
}{
{"query_smiles", "Target SMILES"},
{"query_id", "Target ChEMBL ID"},
{"top_k", "Global top-k"},
{"threshold", "Similarity threshold"},
{"threshold_direction", "Threshold direction"},
}
out := make([]ParameterCard, 0, len(keys))
for _, entry := range keys {
value, ok := parameters[entry.key]
if !ok {
continue
}
formatted, ok := formatUIParameter(value)
if ok {
out = append(out, ParameterCard{Label: entry.label, Value: formatted})
}
}
return out
}
func formatUIParameter(value any) (string, bool) {
switch typed := value.(type) {
case string:
return typed, true
case int, int64, float64, bool:
return fmt.Sprint(typed), true
default:
return "", false
}
}
func shortID(value string) string {
if len(value) <= 8 {
return value
}
return value[:8]
}
@@ -0,0 +1,19 @@
package usecase
import "testing"
func TestUIParametersAreAllowlisted(t *testing.T) {
parameters := uiParameters(map[string]any{
"query_smiles": "CCO",
"top_k": float64(20),
"internal_storage_key": "must-not-reach-browser",
"nested": map[string]any{"secret": "no"},
})
if len(parameters) != 2 {
t.Fatalf("parameters = %#v, want only two allowlisted values", parameters)
}
if parameters[0] != (ParameterCard{Label: "Target SMILES", Value: "CCO"}) ||
parameters[1] != (ParameterCard{Label: "Global top-k", Value: "20"}) {
t.Fatalf("parameters = %#v", parameters)
}
}
+90 -11
View File
@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"io"
"math"
"github.com/google/uuid"
@@ -15,20 +16,27 @@ import (
// creates the job with one task per shard — the coordinator-side counterpart of
// a client submitting pre-chunked URIs.
type SubmitDataset struct {
blobs BlobStore
artifacts ArtifactRepository
jobs JobRepository
tasks TaskRepository
tx TxManager
clk Clock
blobs BlobStore
artifacts ArtifactRepository
jobs JobRepository
tasks TaskRepository
tx TxManager
clk Clock
maxAttempts int
}
func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset {
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk}
tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int) *SubmitDataset {
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk, maxAttempts: maxAttempts}
}
func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) {
if err := validateUploadedWorkload(in.Workload, in.Parameters); err != nil {
return SubmitDatasetResult{}, err
}
if uc.maxAttempts < 1 {
return SubmitDatasetResult{}, domain.ErrInvalidInput
}
now := uc.clk.Now()
job, err := domain.NewUploadedJob(in.Workload, in.Parameters, now)
@@ -64,7 +72,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
cleanup()
return SubmitDatasetResult{}, err
}
splitErr := chunk.SplitTSV(rc, in.RowsPerShard, func(index int, shard io.Reader) error {
splitErr := chunk.SplitChEMBLTSVLimit(rc, in.RowsPerShard, in.MaxRows, func(index int, shard io.Reader) error {
art, err := domain.NewArtifact(job.ID, nil, domain.ArtifactShard,
fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now)
if err != nil {
@@ -77,7 +85,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
putKeys = append(putKeys, art.StorageKey)
art.SetContent(ssum, ssize)
task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, 0, now)
task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, uc.maxAttempts, now)
if err != nil {
return err
}
@@ -88,7 +96,8 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
_ = rc.Close()
if splitErr != nil {
cleanup()
return SubmitDatasetResult{}, splitErr
// Dataset shape is caller input, not an internal coordinator failure.
return SubmitDatasetResult{}, domain.ErrInvalidInput
}
// 3. Persist job + all artifacts + all tasks atomically.
@@ -118,6 +127,76 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
}, nil
}
// validateUploadedWorkload is deliberately narrow until CTX-07/08/10 adds a
// typed distributed-workload registry. In particular, running similarity-graph
// independently per TSV shard is scientifically wrong: cross-shard pairs would
// be absent from the apparent graph.
func validateUploadedWorkload(workload string, parameters map[string]any) error {
if workload != "similarity-search" {
return domain.ErrInvalidInput
}
allowed := map[string]struct{}{
"query_smiles": {}, "top_k": {}, "threshold": {},
"threshold_direction": {}, "progress_every": {},
}
for key := range parameters {
if _, ok := allowed[key]; !ok {
return domain.ErrInvalidInput
}
}
query, ok := parameters["query_smiles"].(string)
if !ok || query == "" || len(query) > 200 {
return domain.ErrInvalidInput
}
if value, ok := parameters["top_k"]; ok && !isPositiveJSONInteger(value) {
return domain.ErrInvalidInput
}
if value, ok := parameters["progress_every"]; ok && !isNonNegativeJSONInteger(value) {
return domain.ErrInvalidInput
}
if value, ok := parameters["threshold"]; ok && !isUnitIntervalNumber(value) {
return domain.ErrInvalidInput
}
if value, ok := parameters["threshold_direction"]; ok && value != "greater" && value != "less" {
return domain.ErrInvalidInput
}
return nil
}
func isPositiveJSONInteger(value any) bool { return isJSONInteger(value, false) }
func isNonNegativeJSONInteger(value any) bool { return isJSONInteger(value, true) }
func isJSONInteger(value any, allowZero bool) bool {
var n int64
switch v := value.(type) {
case int:
n = int64(v)
case int64:
n = v
case float64:
if math.Trunc(v) != v || v > math.MaxInt64 || v < math.MinInt64 {
return false
}
n = int64(v)
default:
return false
}
return n >= 0 && (allowZero || n > 0)
}
func isUnitIntervalNumber(value any) bool {
switch v := value.(type) {
case float64:
return !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 && v <= 1
case int:
return v >= 0 && v <= 1
case int64:
return v >= 0 && v <= 1
default:
return false
}
}
// GetTaskInput resolves a task's input shard and opens it for streaming. The
// caller closes the reader.
type GetTaskInput struct {
+223 -8
View File
@@ -54,6 +54,9 @@ type harness struct {
downloadArt *usecase.DownloadArtifact
getInput *usecase.GetTaskInput
expire *usecase.ExpireLeases
cancel *usecase.CancelJob
reduce *usecase.ReduceJob
jobResult *usecase.GetJobResult
}
func newHarness() *harness {
@@ -67,21 +70,101 @@ func newHarness() *harness {
}
tx := memstore.Tx{}
h.createJob = usecase.NewCreateJob(h.jobs, h.tasks, tx, h.clk)
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk)
h.claim = usecase.NewClaimTask(h.tasks, h.clk, lease)
h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk, 3)
h.claim = usecase.NewClaimTask(h.tasks, h.jobs, h.work, tx, h.clk, lease)
h.renew = usecase.NewRenewLease(h.tasks, h.work, tx, h.clk, lease)
h.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk)
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
h.results = usecase.NewListResults(h.tasks)
h.register = usecase.NewRegisterWorker(h.work, h.clk)
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, h.clk)
h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, tx, h.clk)
h.downloadArt = usecase.NewDownloadArtifact(h.arts, h.blobs)
h.getInput = usecase.NewGetTaskInput(h.tasks, h.arts, h.blobs)
h.expire = usecase.NewExpireLeases(h.tasks, h.clk)
h.expire = usecase.NewExpireLeases(h.tasks, h.jobs, tx, h.clk)
h.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk)
h.reduce = usecase.NewReduceJob(h.jobs, h.tasks, h.arts, h.blobs, tx, h.clk)
h.jobResult = usecase.NewGetJobResult(h.jobs, h.downloadArt)
return h
}
func TestSimilaritySearchReductionCreatesFinalArtifact(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "similarity-search", 2)
_, err := h.jobs.Get(ctx, jobID)
if err != nil {
t.Fatal(err)
}
if err := h.jobs.UpdateStatus(ctx, jobID, domain.JobRunning, nil); err != nil {
t.Fatal(err)
}
partials := []string{
"rank,chembl_id,canonical_smiles,similarity\n1,B,CCC,0.50000048\n",
"rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.50000049\n",
}
for _, partial := range partials {
taskID, attempt := h.leaseOne(t, "w1", "similarity-search")
art, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{TaskID: taskID, WorkerID: "w1", Attempt: attempt, Filename: "partial.csv", ContentType: "text/csv", Body: strings.NewReader(partial)})
if err != nil {
t.Fatal(err)
}
if _, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{TaskID: taskID, WorkerID: "w1", Attempt: attempt, ResultArtifactID: art.ID}); err != nil {
t.Fatal(err)
}
}
if err := h.reduce.Execute(ctx, jobID); err != nil {
t.Fatal(err)
}
progress, err := h.status.Execute(ctx, jobID)
if err != nil || progress.Job.Status != domain.JobCompleted {
t.Fatalf("status=%s err=%v", progress.Job.Status, err)
}
art, body, err := h.jobResult.Execute(ctx, jobID)
if err != nil {
t.Fatal(err)
}
defer body.Close()
bytes, _ := io.ReadAll(body)
if art.Kind != domain.ArtifactFinalResult || string(bytes) != "rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.500000\n2,B,CCC,0.500000\n" {
t.Fatalf("unexpected final %q", bytes)
}
}
func TestSimilaritySearchReductionFailureIsSanitized(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "similarity-search", 1)
if err := h.jobs.UpdateStatus(ctx, jobID, domain.JobRunning, nil); err != nil {
t.Fatal(err)
}
taskID, attempt := h.leaseOne(t, "w1", "similarity-search")
art, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
TaskID: taskID, WorkerID: "w1", Attempt: attempt, Filename: "partial.csv",
ContentType: "text/csv", Body: strings.NewReader("rank,chembl_id,canonical_smiles,similarity\n2,A,CC,0.9\n"),
})
if err != nil {
t.Fatal(err)
}
if _, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{
TaskID: taskID, WorkerID: "w1", Attempt: attempt, ResultArtifactID: art.ID,
}); err != nil {
t.Fatal(err)
}
if err := h.reduce.Execute(ctx, jobID); err != nil {
t.Fatal(err)
}
job, err := h.jobs.Get(ctx, jobID)
if err != nil {
t.Fatal(err)
}
if job.Status != domain.JobFailed || job.ErrorCode == nil || *job.ErrorCode != "reducer_failed" ||
job.ErrorMessage == nil || *job.ErrorMessage != "final result reduction failed" {
t.Fatalf("unexpected failed job: %+v", job)
}
if job.ResultArtifactID != nil {
t.Fatal("failed reduction must not expose a final result")
}
}
// seedJob creates a URI-chunked job with n chunks and returns its id.
func (h *harness) seedJob(t *testing.T, workload string, n int) uuid.UUID {
t.Helper()
@@ -136,6 +219,44 @@ func TestClaimLeasesAndAdvancesAttempt(t *testing.T) {
}
}
func TestRegisteredWorkerCannotBroadenItsCapabilitiesAtClaim(t *testing.T) {
h := newHarness()
h.seedJob(t, "restricted", 1)
worker, err := h.register.Execute(ctx, usecase.RegisterWorkerInput{
Name: "search-only", Capabilities: []string{"similarity-search"},
})
if err != nil {
t.Fatal(err)
}
claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{
WorkerID: worker.ID.String(), Workloads: []string{"restricted"},
})
if err != nil || claimed != nil {
t.Fatalf("claim = (%v, %v), want no compatible task", claimed, err)
}
}
func TestCreateJobRejectsUnsafeDistributedScientificPlans(t *testing.T) {
h := newHarness()
_, err := h.createJob.Execute(ctx, usecase.CreateJobInput{
Workload: "similarity-graph", InputURI: "s3://input",
Chunks: []usecase.ChunkInput{{ChunkIndex: 0, InputURI: "s3://chunk", InputSHA256: "sha"}},
})
if !errors.Is(err, domain.ErrInvalidInput) {
t.Errorf("graph job err = %v, want ErrInvalidInput", err)
}
_, err = h.createJob.Execute(ctx, usecase.CreateJobInput{
Workload: "similarity-search", InputURI: "s3://input", Parameters: map[string]any{"query_id": "CHEMBL1"},
Chunks: []usecase.ChunkInput{
{ChunkIndex: 0, InputURI: "s3://chunk0", InputSHA256: "sha"},
{ChunkIndex: 1, InputURI: "s3://chunk1", InputSHA256: "sha"},
},
})
if !errors.Is(err, domain.ErrInvalidInput) {
t.Errorf("sharded query_id job err = %v, want ErrInvalidInput", err)
}
}
func TestClaimEmptyQueueReturnsNil(t *testing.T) {
h := newHarness()
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
@@ -279,7 +400,7 @@ func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) {
h.seedJob(t, "w", 1)
taskID, attempt := h.leaseOne(t, "w1", "w")
h.uploadArt = usecase.NewUploadArtifact(
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, h.clk,
h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, memstore.Tx{}, h.clk,
)
_, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
@@ -412,6 +533,23 @@ func TestUploadArtifactRejectsForeignWorker(t *testing.T) {
}
}
func TestUploadArtifactIsIdempotentPerTaskAttempt(t *testing.T) {
h := newHarness()
h.seedJob(t, "w", 1)
taskID, attempt := h.leaseOne(t, "w1", "w")
first := h.uploadResult(t, taskID, "w1", attempt)
second, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
TaskID: taskID, WorkerID: "w1", Attempt: attempt,
Filename: "retry.csv", ContentType: "text/csv", Body: strings.NewReader("different bytes"),
})
if err != nil {
t.Fatal(err)
}
if second.ID != first {
t.Errorf("retry artifact = %s, want existing %s", second.ID, first)
}
}
func TestDownloadArtifactRoundTrips(t *testing.T) {
h := newHarness()
h.seedJob(t, "w", 1)
@@ -432,10 +570,10 @@ func TestDownloadArtifactRoundTrips(t *testing.T) {
func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
h := newHarness()
tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "w", RowsPerShard: 2, Filename: "chembl.tsv",
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, Filename: "chembl.tsv",
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
})
if err != nil {
@@ -451,7 +589,7 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
t.Errorf("job total = %d, want 3", prog.Total)
}
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"similarity-search"}})
if err != nil || c == nil {
t.Fatalf("claim shard: %v", err)
}
@@ -468,6 +606,64 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
}
}
func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
h := newHarness()
tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n"
res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv",
ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
})
if err != nil {
t.Fatal(err)
}
if res.TaskCount != 2 {
t.Fatalf("task_count = %d, want 2", res.TaskCount)
}
}
func TestSubmitDatasetRejectsUnsupportedDistributedWorkloads(t *testing.T) {
h := newHarness()
_, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "similarity-graph", Parameters: map[string]any{"threshold": 0.7}, RowsPerShard: 2,
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
})
if !errors.Is(err, domain.ErrInvalidInput) {
t.Errorf("graph submission err = %v, want ErrInvalidInput", err)
}
_, err = h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "similarity-search", Parameters: map[string]any{"query_id": "CHEMBL1"}, RowsPerShard: 2,
Filename: "chembl.tsv", ContentType: "text/tab-separated-values",
Body: strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n"),
})
if !errors.Is(err, domain.ErrInvalidInput) {
t.Errorf("query_id submission err = %v, want ErrInvalidInput", err)
}
}
func TestCancelJobInvalidatesClaimedAndPendingTasks(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "w", 3)
claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}})
if err != nil || claimed == nil {
t.Fatalf("claim: %v", err)
}
cancelled, err := h.cancel.Execute(ctx, jobID)
if err != nil || cancelled != 3 {
t.Fatalf("cancel = (%d, %v), want (3, nil)", cancelled, err)
}
if _, err := h.renew.Execute(ctx, usecase.RenewLeaseInput{TaskID: claimed.TaskID, WorkerID: "w1", Attempt: claimed.Attempt}); !errors.Is(err, domain.ErrTaskNotLeased) {
t.Errorf("cancelled lease heartbeat = %v, want ErrTaskNotLeased", err)
}
progress, err := h.status.Execute(ctx, jobID)
if err != nil || progress.DeriveStatus() != domain.JobCancelled || progress.Cancelled != 3 {
t.Errorf("cancelled progress = %+v, err = %v", progress, err)
}
if cancelled, err := h.cancel.Execute(ctx, jobID); err != nil || cancelled != 0 {
t.Errorf("second cancel = (%d, %v), want (0, nil)", cancelled, err)
}
}
func TestGetTaskInputMissingForURITask(t *testing.T) {
h := newHarness()
h.seedJob(t, "w", 1) // URI-based task, no coordinator-stored input
@@ -491,3 +687,22 @@ func TestExpireLeasesReclaims(t *testing.T) {
t.Errorf("expire = (%d, %v), want (1, nil)", n, err)
}
}
func TestFinalLeaseExpiryPersistsFailedJobAndCannotBeCancelled(t *testing.T) {
h := newHarness()
jobID := h.seedJob(t, "w", 1)
for attempt := 1; attempt <= domain.DefaultMaxAttempts; attempt++ {
h.leaseOne(t, "w1", "w")
h.clk.Advance(lease + time.Second)
if _, err := h.expire.Execute(ctx); err != nil {
t.Fatalf("expire attempt %d: %v", attempt, err)
}
}
progress, err := h.status.Execute(ctx, jobID)
if err != nil || progress.Job.Status != domain.JobFailed || progress.DeriveStatus() != domain.JobFailed {
t.Fatalf("progress = %+v, err = %v; want persisted failed job", progress, err)
}
if _, err := h.cancel.Execute(ctx, jobID); !errors.Is(err, domain.ErrJobNotCancellable) {
t.Errorf("cancel terminal lease failure = %v, want ErrJobNotCancellable", err)
}
}
@@ -0,0 +1,5 @@
BEGIN;
DROP INDEX IF EXISTS uq_partial_result_task_attempt;
COMMIT;
@@ -0,0 +1,26 @@
BEGIN;
-- Old deployments can contain more than one partial result because earlier
-- versions accepted repeated PUTs. Preserve the one referenced by a completed
-- task and discard stale rows; unfinished tasks must upload again after a
-- deploy, just as they do after a lost lease.
DELETE FROM artifacts AS a
USING tasks AS t
WHERE a.task_id = t.id
AND a.kind = 'partial_result'::artifact_kind
AND t.status <> 'completed'::task_status;
DELETE FROM artifacts AS a
USING tasks AS t
WHERE a.task_id = t.id
AND a.kind = 'partial_result'::artifact_kind
AND t.status = 'completed'::task_status
AND a.id <> t.result_artifact_id;
-- One lease attempt has one durable partial result. This makes an upload retry
-- idempotent and prevents repeated uploads from accumulating orphan artifacts.
CREATE UNIQUE INDEX uq_partial_result_task_attempt
ON artifacts (task_id, attempt)
WHERE kind = 'partial_result'::artifact_kind;
COMMIT;
@@ -0,0 +1,7 @@
BEGIN;
ALTER TABLE jobs DROP COLUMN IF EXISTS error_message;
ALTER TABLE jobs DROP COLUMN IF EXISTS error_code;
ALTER TABLE jobs DROP COLUMN IF EXISTS reducer_started_at;
COMMIT;
@@ -0,0 +1,7 @@
-- PostgreSQL enum values must be committed before they are used by a later
-- transaction, so this migration intentionally has no BEGIN/COMMIT wrapper.
ALTER TYPE job_status ADD VALUE IF NOT EXISTS 'reducing';
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS error_code text;
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS error_message text;
ALTER TABLE jobs ADD COLUMN IF NOT EXISTS reducer_started_at timestamptz;
+27 -16
View File
@@ -51,6 +51,13 @@ check "register worker" 201 -X POST "${HOST}/workers/register" "
-d '{"name":"smoke-worker","capabilities":["similarity_search"],"cpu_count":4,"memory_mb":8192}'
check "register without capabilities → 400" 400 -X POST "${HOST}/workers/register" "${auth[@]}" \
-d '{"name":"bad"}'
registration=$(curl -sS "${auth[@]}" -X POST "${HOST}/workers/register" \
-d '{"name":"smoke-worker-active","capabilities":["similarity_search","similarity-search"],"cpu_count":4}')
worker_id=$(printf '%s' "$registration" | python3 -c 'import json,sys;print(json.load(sys.stdin)["worker_id"])' 2>/dev/null)
if [[ -z "${worker_id:-}" ]]; then
echo " ✗ could not register active worker: $registration"
exit 1
fi
echo
echo "job lifecycle"
@@ -75,7 +82,7 @@ declare -A our_chunks
task_id=""
attempt=""
for _ in $(seq 1 40); do
claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d '{"worker_id":"w1"}')
claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d "{\"worker_id\":\"${worker_id}\"}")
[[ -z "$claim" ]] && break # 204: queue drained
read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" |
@@ -100,7 +107,7 @@ else
fi
check "heartbeat" 200 -X POST "${HOST}/tasks/${task_id}/heartbeat" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt}}"
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt}}"
# --- artifacts + result (uploads happen while the task is still leased) ---
bearer=(-H "Authorization: Bearer ${TOKEN}")
@@ -108,43 +115,47 @@ bearer=(-H "Authorization: Bearer ${TOKEN}")
# upload <filename> -> prints the artifact_id
upload() {
curl -sS -X PUT "${HOST}/tasks/${task_id}/artifacts/$1" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \
-H 'Content-Type: text/csv' -H "X-Worker-ID: ${worker_id}" -H "X-Task-Attempt: ${attempt}" \
--data-binary $'query,match,score\nA,B,0.9\n' |
python3 -c 'import json,sys;print(json.load(sys.stdin)["artifact_id"])' 2>/dev/null
}
check "upload artifact" 200 -X PUT "${HOST}/tasks/${task_id}/artifacts/result.csv" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \
-H 'Content-Type: text/csv' -H "X-Worker-ID: ${worker_id}" -H "X-Task-Attempt: ${attempt}" \
--data-binary $'query,match,score\nA,B,0.9\n'
check "foreign worker upload → 409" 409 -X PUT "${HOST}/tasks/${task_id}/artifacts/x.csv" "${bearer[@]}" \
-H 'Content-Type: text/csv' -H 'X-Worker-ID: impostor' -H "X-Task-Attempt: ${attempt}" \
--data-binary 'x'
# Two result artifacts, uploaded now while the lease is held: one to complete
# with, a second to prove a different manifest is rejected after completion.
# A retry of a PUT returns the same durable artifact for the task attempt.
art_id=$(upload primary.csv)
art_id2=$(upload secondary.csv)
if [[ "$art_id" == "$art_id2" ]]; then
printf ' \033[32m✓\033[0m %-46s %s\n' "duplicate upload is idempotent" "$art_id"
pass=$((pass + 1))
else
printf ' \033[31m✗\033[0m %-46s got %s and %s\n' "duplicate upload is idempotent" "$art_id" "$art_id2"
fail=$((fail + 1))
fi
check "download artifact" 200 "${HOST}/artifacts/${art_id}/download" "${bearer[@]}"
check "foreign worker submits → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"impostor\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "submit result" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "replay same result → idempotent" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "different result → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \
-d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id2}\"}}"
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}"
echo
echo "input validation"
check "malformed uuid → 400" 400 -X POST "${HOST}/tasks/not-a-uuid/result" "${auth[@]}" \
-d '{"worker_id":"w1","attempt":1,"result":{"artifact_id":"00000000-0000-0000-0000-000000000000"}}'
-d "{\"worker_id\":\"${worker_id}\",\"attempt\":1,\"result\":{\"artifact_id\":\"00000000-0000-0000-0000-000000000000\"}}"
# Note: Go's encoding/json matches field names case-insensitively, so
# "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key
# trips DisallowUnknownFields.
check "unknown json field → 400" 400 -X POST "${HOST}/tasks/claim" "${auth[@]}" \
-d '{"worker_id":"w1","totally_unknown":1}'
-d "{\"worker_id\":\"${worker_id}\",\"totally_unknown\":1}"
check "unknown job → 404" 404 "${HOST}/jobs/00000000-0000-0000-0000-000000000000" "${auth[@]}"
echo
@@ -152,11 +163,11 @@ echo "dataset upload → chunking"
# Upload a 5-row TSV split at 2 rows/shard → expect 3 shard tasks. The text
# fields precede the file part, which the coordinator streams.
up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \
-F 'workload=similarity_search' \
-F 'parameters={"top_k":10}' \
-F 'workload=similarity-search' \
-F 'parameters={"query_smiles":"CCO","top_k":10}' \
-F 'chunk_rows=2' \
-F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV'
id smiles
chembl_id canonical_smiles
A CC
B CCC
C CCCC
@@ -180,7 +191,7 @@ fi
up_input=""
for _ in $(seq 1 30); do
c=$(curl -sS "${bearer[@]}" -H 'Content-Type: application/json' -X POST "${HOST}/tasks/claim" \
-d '{"worker_id":"up-w","capabilities":["similarity_search"]}')
-d "{\"worker_id\":\"${worker_id}\"}")
[[ -z "$c" ]] && break
cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null)
[[ "$cj" != "$up_job" ]] && continue
+58 -6
View File
@@ -26,7 +26,8 @@ must be updated in the same change as any behaviour it describes.
| `POST /tasks/{id}/heartbeat` | renew lease | ✅ done |
| `POST /tasks/{id}/result` | complete | ✅ done, references `artifact_id` |
| `POST /tasks/{id}/failure` | fail | ✅ done |
| `GET /jobs/{id}` | progress | ✅ done |
| `GET /jobs/{id}` | progress and final-result URI | ✅ done |
| `GET /jobs/{id}/result` | download final CSV | ✅ done |
| `PUT /tasks/{id}/artifacts/{name}` | upload partial | ✅ done |
| `GET /artifacts/{id}/download` | download by id | ✅ done |
| `POST /jobs/upload` | upload dataset, coordinator chunks it | ✅ done |
@@ -52,9 +53,11 @@ Content-Type: multipart/form-data
```
Fields, in order (text fields first, file last — the file is streamed):
`workload`, `parameters` (JSON), `chunk_rows` (int, default 1000), and the file
part `file`. The coordinator stores the input, splits the TSV into shard
artifacts (header repeated per shard), and creates one task per shard.
`workload`, `parameters` (JSON), `chunk_rows` (int, default 1000), optional
`max_rows` (positive int), and the file part `file`. `max_rows` limits the
leading data rows that become shards; it does not change the stored source
artifact. The coordinator splits the selected TSV rows into shard artifacts
(header repeated per shard) and creates one task per shard.
`201`:
@@ -65,6 +68,53 @@ artifacts (header repeated per shard), and creates one task per shard.
Each resulting task's claim response carries `input.uri = /tasks/{id}/input`,
served by §5.4.
## Job progress and final result
```http
GET /jobs/{job_id}
Authorization: Bearer <token>
```
The response contains task counters and a derived status. A successful
`similarity-search` enters `reducing` after the last shard completes, then
becomes `completed` only after the coordinator stores its deterministic final
CSV. At that point `result_uri` is present:
```json
{
"id": "uuid",
"status": "completed",
"total": 3,
"pending": 0,
"leased": 0,
"completed": 3,
"failed": 0,
"cancelled": 0,
"result_uri": "/jobs/uuid/result"
}
```
`GET /jobs/{job_id}/result` downloads that final coordinator-owned CSV. Before
the job is completed (or for a failed/cancelled job), it returns `404`.
## Stop a job
```http
POST /jobs/{job_id}/cancel
Authorization: Bearer <token>
```
The coordinator transactionally marks every pending, leased, or running shard
as `cancelled`, invalidates its lease, and marks the job `cancelled`. Completed
and terminally failed shards remain as history. Repeating a cancellation of an
already cancelled job is safe.
`200`:
```json
{ "job_id": "uuid", "status": "cancelled", "cancelled_tasks": 12 }
```
## Register worker
```http
@@ -87,7 +137,8 @@ Content-Type: application/json
```
`cpu_count`/`memory_mb` are accepted for forward compatibility and not yet
persisted. `capabilities` must be non-empty (an allowlisted workload set).
persisted. `capabilities` must be non-empty. A claim uses the capabilities
stored at registration; the request cannot broaden them.
## Claim task
@@ -100,7 +151,8 @@ Content-Type: application/json
```
- `204 No Content`: no compatible task.
- `200 OK`: a task is leased atomically.
- `200 OK`: a task is leased atomically. `worker_id` must be a registered UUID;
its persisted capabilities, rather than this request field, decide eligibility.
```json
{
+81
View File
@@ -0,0 +1,81 @@
# Task: safely preview partial CSV artifacts in the UI
## Assignment
You are the junior developer implementing one contained UI feature: an
authenticated operator can preview a small portion of a CSV artifact belonging
to the job they are viewing. This is a diagnostic aid, not a final-results
page.
## Read first
1. `AGENTS.md`
2. `.agents/coordinator.md`
3. `docs/web-interface-plan.md`
4. `docs/api-contract.md`
5. `coordinator/internal/transport/http/ui.go` and its tests
## Current baseline
The coordinator serves an authenticated local web UI. Job detail pages already
list partial result artifacts and provide job-scoped downloads. The browser has
no worker token and must not learn storage paths. A partial shard CSV is never
a global or final molecular-search result.
## Scope
Add a **Preview** action next to eligible CSV artifacts on a job detail page.
- Preview only artifacts that belong to the requested job.
- Show at most the first **30 rows** and read at most **64 KiB** from storage.
- State clearly when content was truncated.
- Preserve the existing download action.
- For a non-CSV artifact, return a friendly, sanitized explanation rather than
attempting to render bytes as text.
- Use an existing UI route pattern or add a small UI-authenticated endpoint;
keep it separate from worker API routes.
## Security rules
- Require UI Basic Auth for every preview request.
- Verify job ownership in the coordinator before opening the artifact; an
artifact ID from another job must not be previewable.
- Never expose `storage_key`, filesystem paths, database errors, bearer tokens,
or worker-local information.
- Do not use `innerHTML` for CSV fields. Use `html/template` escaping or
`textContent` so strings such as `<script>alert(1)</script>` are displayed as
data, not executed.
- Do not load the complete artifact into memory.
## Out of scope
- Charts, molecule imagery, RDKit rendering, client-side CSV libraries, React,
and a new frontend service.
- Changing job/task state, retrying tasks, or implementing reducer output.
- Redesigning the broader dashboard or job-creation workflow; that belongs to
`docs/user-space-task.md`.
## Suggested implementation shape
Keep UI transport, use case, and storage responsibilities separate. Return a
small view model containing artifact metadata, column headers, rows, and a
`truncated` flag. Reuse existing coordinator-owned artifact access rather than
reading a path supplied by the browser. Keep the handler streaming/limited.
## Acceptance criteria
- A valid partial CSV can be previewed from its own job detail page.
- The result shows no more than 30 data rows and marks 64 KiB/row truncation.
- Empty and malformed CSV content fail safely with a clear message.
- A non-CSV artifact is rejected safely.
- Unauthenticated access is rejected; a cross-job artifact request is not
disclosed or served.
- HTML-like values are escaped in the rendered preview.
- Existing artifact downloads still work.
- Add Go tests for all cases above and run `go test ./...` and `go vet ./...`.
## Handoff
Work in one focused branch and one PR. Report files changed, any API impact,
test commands/results, and known limitations. Do not commit datasets, generated
CSV files, Docker volumes, `.venv`, or `worker-data/`.
+36 -8
View File
@@ -46,24 +46,26 @@ coordinator was started with. Never log it, never send it in an error body.
```http
POST /workers/register
{ "name": "lab-worker-01", "capabilities": ["similarity_search"] }
{ "name": "lab-worker-01", "capabilities": ["similarity-search"] }
```
Response: `{ "worker_id": "<uuid>", "heartbeat_interval_seconds": 15 }`.
- `capabilities` are the workload names you can run — the coordinator only hands
you matching tasks.
- `capabilities` are fixed at registration — the coordinator only hands you
matching tasks and a later claim cannot broaden that set.
- **Keep `worker_id`**. Use it as your identity in every later call. Using the
registered UUID is what lets the coordinator track your liveness (it marks
workers offline after they go silent).
- Current coordinator jobs use `similarity_search` / `similarity_graph`; the
reference Python worker also accepts the public CLI spellings with hyphens.
- Current distributed uploads use `similarity-search` with `query_smiles`; the
coordinator merges completed shard candidates into a final CSV. The reference
worker accepts the legacy `similarity_search` spelling too. Do not advertise
`similarity-graph` until CTX-10 implements cross-shard pair planning.
## 2. Claim a task
```http
POST /tasks/claim
{ "worker_id": "<uuid>", "capabilities": ["similarity_search"] }
{ "worker_id": "<uuid>", "capabilities": ["similarity-search"] }
```
- `200` → a leased task (below).
@@ -74,9 +76,9 @@ POST /tasks/claim
"task_id": "<uuid>",
"attempt": 1,
"lease_expires_at": "2026-07-22T12:05:00Z",
"workload": "similarity_search",
"workload": "similarity-search",
"input": { "uri": "/tasks/<uuid>/input", "sha256": "<hex>" },
"parameters": { "query_id": "CHEMBL939", "top_k": 20 }
"parameters": { "query_smiles": "CCO", "top_k": 20 }
}
```
@@ -193,6 +195,32 @@ Per the worker contract, at minimum:
- poll interval and request timeout
- a working directory for downloaded inputs and generated outputs
## Run the reference worker locally
Use one terminal per worker and a distinct work directory for each process:
```sh
SCIMESH_COORDINATOR_URL=http://localhost:8080 \
SCIMESH_BEARER_TOKEN=dev-token \
SCIMESH_WORKER_NAME=worker-1 \
scimesh-worker --work-dir "$PWD/worker-data-1"
```
For a bounded manual check, use one of these lifecycle modes:
```sh
# Make exactly one claim; exit immediately when no task is available.
scimesh-worker --work-dir "$PWD/worker-data-check" --once
# Keep polling until two tasks complete successfully, then exit.
scimesh-worker --work-dir "$PWD/worker-data-check" --max-tasks 2
```
`SCIMESH_MAX_TASKS` provides the same limit through the environment. Pressing
`Ctrl+C` stops the reference worker cleanly. If it interrupts an active task,
the worker reports a sanitized retriable failure first, emits no traceback, and
exits with status `130`.
## Generate a client from the spec
Instead of hand-writing request code, generate it:
@@ -0,0 +1,214 @@
# CTX-07: distributed workload protocol and planner contract
## Status and scope
This document is the implementation contract for CTX-07. Its generic protocol,
registry, strict JSON models, and deterministic reduction ordering are
implemented in `scimesh/distributed/`. CTX-08 implements the molecular
similarity-search planner, worker adapter, and pure reducer on top of it. This
document does not implement a coordinator API endpoint, database migration, or
durable final artifact. Until CTX-09 is complete, shard CSVs remain diagnostic
partial results.
The protocol gives local scientific workloads a coordinator-independent way to
validate a job, plan artifact-backed tasks, and later reduce completed outputs.
The Go coordinator owns durable artifacts, transactions, task rows, leases, and
HTTP. A Python workload must never access PostgreSQL or call the coordinator.
Read `PLAN.md`, `.agents/workloads.md`, and `docs/api-contract.md` before
implementing this CTX.
## Canonical vocabulary
- External workload names are lowercase hyphenated names: `similarity-search`
and, later, `similarity-graph`.
- The existing underscore spellings are a temporary compatibility alias at the
Python worker boundary only. Planners, persisted job/task payloads, and new
API examples use the canonical hyphenated spelling.
- A **plan** contains only JSON-compatible values and coordinator artifact
references. It contains no local filesystem path, worker URI, presigned URL,
database connection, or callable.
- `chunk_index` is a non-negative integer, unique within a plan, and sorted
ascending whenever results are enumerated.
## Python boundary
CTX-07 adds a small `DistributedWorkload` protocol under `scimesh/distributed/`
and a registry separate from the local CLI registry. Names below are proposed
public types; keep concrete implementation details minimal.
```python
class DistributedWorkload(Protocol):
name: str
def validate_job(self, parameters: Mapping[str, object]) -> None: ...
def plan(
self,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan: ...
def reduce(
self,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult: ...
```
`input_path` and `workspace` are temporary files supplied by the coordinator
bridge. They are never serialized. `plan()` returns only a `DistributedPlan`;
the bridge validates it, persists artifact/task rows in one coordinator
transaction, and removes its temporary workspace. If validation or planning
fails, no job or task may be written.
## JSON models
All objects below are schema version `1`. Future incompatible changes require a
new version; never infer a schema from missing fields.
### Artifact reference
```json
{
"artifact_id": "c4273293-f8b4-4ecb-99df-3b9f5a32b6a6",
"sha256": "3b2d...64-lowercase-hex-characters",
"content_type": "text/tab-separated-values"
}
```
The artifact ID is coordinator-owned. The checksum is included so planning and
tests can assert exactly which immutable input was used. A worker receives the
coordinator-generated download URI only through `POST /tasks/claim`.
### Distributed plan
```json
{
"schema_version": 1,
"workload": "similarity-search",
"resolved_parameters": {
"query_smiles": "COc1ccc(Nc2ncnc3cc(OCCCN4CCOCC4)c(OC)c23)cc1",
"query_source": {"kind": "chembl_id", "value": "CHEMBL939"},
"top_k": 20,
"threshold": 0.7,
"threshold_direction": "greater",
"fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}
},
"tasks": [
{
"chunk_index": 0,
"input_artifact": {
"artifact_id": "69e41105-d9fb-4c7f-a2db-7dd9e3ba2c76",
"sha256": "4c92...64-lowercase-hex-characters",
"content_type": "text/tab-separated-values"
},
"parameters": {
"query_smiles": "COc1ccc(Nc2ncnc3cc(OCCCN4CCOCC4)c(OC)c23)cc1",
"top_k": 20,
"threshold": 0.7,
"threshold_direction": "greater",
"fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}
}
}
]
}
```
`resolved_parameters` are immutable job metadata. A task copies only the
values required by its worker runner. The coordinator may add its own durable
task ID and generated input URI; it must not alter scientific parameters.
## Similarity-search planning rules
1. Accept exactly one of `query_id` and `query_smiles` at the public boundary.
2. Validate a supplied SMILES once. For `query_id`, find and validate that
molecule once against the original uploaded TSV **before** creating shards.
3. Persist the resolved canonical query SMILES and the original query source in
`resolved_parameters`. Workers receive `query_smiles`, never `query_id`.
4. Fingerprint settings are fixed to Morgan radius `2` and `fp_size` `2048`.
Reject a request that tries to override them rather than silently changing
scientific semantics.
5. Split source rows in input order. Every shard includes the original TSV
header and has a contiguous, zero-based `chunk_index`.
6. Each shard uses the global `top_k`, not a smaller local limit. A global
reducer cannot recover a candidate discarded by every shard.
7. Preserve `threshold`, `threshold_direction`, and valid `max_rows` semantics
in the resolved plan. A job-level row limit is applied before sharding, not
independently by every worker.
Invalid row SMILES are not planner failures. They remain shard data and are
counted by the worker exactly as the local workload does. An invalid query is a
planning failure.
## Partial-result contract
A completed similarity-search task owns exactly one coordinator-uploaded CSV
artifact with content type `text/csv` and these columns, in this order:
```csv
rank,chembl_id,canonical_smiles,similarity
1,CHEMBL123,CCO,0.875000
```
- `rank` is one-based local rank.
- `similarity` uses a round-trip decimal representation of the computed float
(for Python, `repr(similarity)`). This preserves exact cross-shard ranking;
the reducer writes the user-facing final CSV with the local CLI's six-decimal
display formatting.
- Rows are sorted by `(-similarity, chembl_id, canonical_smiles)` for
`threshold_direction=greater`, or `(similarity, chembl_id,
canonical_smiles)` for `less`.
- The query molecule and every candidate with the same canonical query SMILES
are excluded using the existing local-workload definition.
- Empty valid result files still include the header.
The worker completion metrics must include JSON numbers for `scanned_rows`,
`valid_molecules`, `invalid_smiles`, `matches_emitted`, and
`elapsed_seconds`. Metrics are observability data; the reducer derives final
scientific output exclusively from coordinator-owned partial artifacts.
## Reduction boundary
CTX-09 invokes the registered reducer only after every task is completed. It
passes `CompletedPartial` values ordered by `chunk_index`, each containing its
coordinator artifact reference and validated metrics.
For similarity-search the reducer:
1. reads partial CSVs in `chunk_index` order;
2. validates header, row shape, rank, finite similarity in `[0, 1]`, and sort
order;
3. retains a bounded heap of at most the global `top_k` candidates using the
exact local ranking key;
4. writes the same header and deterministic rank numbering as the local CLI.
It must not deduplicate ordinary records: the local reference keeps input-row
multiplicity. Reduction is independent of worker completion order and uses
`O(top_k + shard_rows)` memory apart from CSV streaming buffers.
## Required tests for the CTX-07 implementation
- unknown workload is rejected before any coordinator job/task write;
- invalid public parameters and invalid `query_id` produce no partial plan;
- `query_id` resolution occurs once, before shard construction;
- the same input, parameters, and shard size generate byte-equivalent
JSON plans and identical shard order;
- every task payload is JSON-serializable and contains only artifact references
and validated scalar/object values;
- a two-shard dummy workload proves coordinator transaction rollback on planner
validation failure;
- completed partial artifacts reach the reducer ordered by `chunk_index`, even
when workers finish in a different order;
- the protocol registry never imports the Go coordinator or database code.
## Deferred work
CTX-09 materializes the planned shard files as coordinator artifacts, invokes
the registered reducer once, and persists its final artifact/job state. CTX-10
defines graph-specific triangular block plans; it must not reuse the search
shard scheme without its pair-coverage invariants.
+68 -11
View File
@@ -95,8 +95,11 @@ paths:
summary: Upload a dataset; the coordinator chunks it into shard tasks
description: >
multipart/form-data. The text fields (`workload`, `parameters`,
`chunk_rows`) MUST precede the `file` part: the file is streamed, not
buffered, so the fields have to be parsed before it arrives.
`chunk_rows`, `max_rows`) MUST precede the `file` part: the file is streamed, not
buffered, so the fields have to be parsed before it arrives. Currently
only `similarity-search` with `parameters.query_smiles` is accepted.
When every shard succeeds, the coordinator merges their candidates into
one final CSV; distributed graph planning is not implemented.
requestBody:
required: true
content:
@@ -122,7 +125,9 @@ paths:
- $ref: "#/components/parameters/JobID"
responses:
"200":
description: Progress counts and derived status.
description: >
Progress counts and derived status. A completed similarity-search
response includes `result_uri` for its final CSV.
content:
application/json:
schema: { $ref: "#/components/schemas/JobProgress" }
@@ -130,6 +135,38 @@ paths:
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
/jobs/{job_id}/result:
get:
tags: [jobs]
summary: Download a completed job's final result
parameters:
- $ref: "#/components/parameters/JobID"
responses:
"200":
description: Final coordinator-owned CSV.
content:
text/csv:
schema: { type: string, format: binary }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
/jobs/{job_id}/cancel:
post:
tags: [jobs]
summary: Cancel a job and invalidate all unfinished task leases
parameters:
- $ref: "#/components/parameters/JobID"
responses:
"200":
description: The job is cancelled. Completed and terminally failed tasks remain unchanged.
content:
application/json:
schema: { $ref: "#/components/schemas/CancelJobResponse" }
"400": { $ref: "#/components/responses/BadRequest" }
"401": { $ref: "#/components/responses/Unauthorized" }
"404": { $ref: "#/components/responses/NotFound" }
"409": { $ref: "#/components/responses/Conflict" }
/tasks/claim:
post:
tags: [tasks]
@@ -371,7 +408,7 @@ components:
type: array
minItems: 1
items: { type: string }
example: [similarity_search, similarity_graph]
example: [similarity-search]
cpu_count:
type: integer
description: Accepted for forward compatibility; not yet persisted.
@@ -402,7 +439,7 @@ components:
type: object
required: [workload, input_uri, chunks]
properties:
workload: { type: string, example: similarity_search }
workload: { type: string, example: similarity-search }
input_uri: { type: string }
parameters: { type: object, additionalProperties: true }
chunks:
@@ -420,15 +457,20 @@ components:
type: object
required: [workload, file]
properties:
workload: { type: string, example: similarity_search }
workload: { type: string, enum: [similarity-search], example: similarity-search }
parameters:
type: string
description: JSON object, sent as a string form field.
example: '{"top_k":10}'
example: '{"query_smiles":"CCO","top_k":10}'
chunk_rows:
type: integer
description: Data rows per shard. Default 1000.
example: 1000
max_rows:
type: integer
minimum: 1
description: Optional leading data-row limit for a small pipeline check.
example: 500
file:
type: string
format: binary
@@ -441,6 +483,13 @@ components:
task_count: { type: integer, example: 3 }
input_artifact_id: { type: string, format: uuid }
CancelJobResponse:
type: object
properties:
job_id: { type: string, format: uuid }
status: { type: string, enum: [cancelled] }
cancelled_tasks: { type: integer }
JobProgress:
type: object
properties:
@@ -451,16 +500,24 @@ components:
leased: { type: integer }
completed: { type: integer }
failed: { type: integer }
cancelled: { type: integer }
result_uri:
type: string
description: Present only when the final result is available.
example: /jobs/5a4c3a7f-ccfc-47d6-b78d-2d1fa565bafd/result
error_code:
type: string
description: Sanitized terminal reducer failure code, when applicable.
ClaimRequest:
type: object
required: [worker_id]
properties:
worker_id: { type: string }
worker_id: { type: string, format: uuid, description: Registered worker identity. }
capabilities:
type: array
items: { type: string }
description: Workloads this worker can run. Empty means "any".
description: Accepted for compatibility only; registration capabilities decide eligibility.
max_concurrency:
type: integer
description: Accepted; the coordinator leases one task per call.
@@ -545,8 +602,8 @@ components:
JobStatus:
type: string
enum: [pending, running, completed, failed, cancelled]
enum: [pending, running, reducing, completed, failed, cancelled]
TaskStatus:
type: string
enum: [pending, leased, completed, failed, cancelled]
enum: [pending, leased, running, completed, failed, cancelled]
+83
View File
@@ -0,0 +1,83 @@
# Task: improve the operator user space
## Assignment
You are the senior developer responsible for the **user space**: the
human-facing local operator interface served by the Go coordinator. In this
task, “user space” means a clear UI and workflow for a trusted local operator;
it does **not** mean public accounts, registration, roles, multi-tenancy, or
remote deployment.
Create a small, coherent improvement to the existing UI so a person can
understand and operate a SciMesh pipeline without reading API payloads or
coordinator logs. Keep all interface copy in English.
## Read first
1. `AGENTS.md`
2. `.agents/coordinator.md` and `.agents/integration.md`
3. `docs/web-interface-plan.md`
4. `docs/api-contract.md`
5. `STATUS.md` and the current `coordinator/internal/transport/http/ui.go`
## Current baseline
`main` already provides local Basic Auth (`UI_AUTH_TOKEN`), a dashboard, job
submission for diagnostic similarity-search runs, task progress, partial CSV
downloads, a stop-job action, and an optional dataset row limit. A job is a
**pipeline check** until CTX-09 adds a reducer; individual shard CSVs are not a
final scientific result.
## Scope
Improve the end-to-end operator journey:
- make the dashboard explain service readiness, workers, jobs, and the next
safe action in plain English;
- make job creation validation and success/failure feedback understandable;
- make job detail clearly distinguish queued, running, failed, stopped, and
completed pipeline checks;
- keep polling and all user-visible states reliable after a page refresh;
- expose actionable, sanitized failure guidance without leaking paths, tokens,
SQL errors, or tracebacks;
- document the workflow in `coordinator/README.md` or `README.md`.
Use server-rendered Go templates, embedded assets, and small vanilla
JavaScript only. Preserve the existing worker API and Basic Auth boundary.
## Explicitly out of scope
- user accounts, sign-up, roles, sessions, OAuth, or multi-tenancy;
- executing or controlling workers from the browser;
- direct browser access to PostgreSQL or worker endpoints;
- final-result reduction, distributed workload planning, or graph execution;
- artifact CSV preview/visualisation. That is assigned independently in
`docs/artifact-preview-task.md`.
## Security and protocol rules
- `UI_AUTH_TOKEN` is never sent to HTML, JavaScript, URLs, logs, or storage.
- Use `html/template`; JavaScript must use `textContent`, never `innerHTML` for
received data.
- UI artifact operations must be job-scoped and coordinator-owned.
- Do not disclose raw worker commands, local paths, bearer tokens, or database
errors.
- Do not change worker/coordinator API contracts silently. Document any
intentional API change in `docs/api-contract.md` and `docs/openapi.yaml`.
## Acceptance criteria
- A new operator can start the stack, authenticate, submit a small run, start
workers, understand live progress, and safely stop a job from the UI.
- The UI never calls or displays worker bearer-token endpoints.
- All partial output is visibly labelled as diagnostic until a reducer exists.
- Disabled UI remains `404`; unauthenticated UI requests remain rejected.
- Go tests cover changed routes and states, including auth and a sanitized
error case.
- `go test ./...`, `go vet ./...`, and the relevant real-PostgreSQL tests pass.
## Handoff
Use one branch and one PR. In the PR description state the user journey that
changed, screenshots if visual layout changed, API impact (`none` if none),
and exact test commands/results. Do not stage local datasets or `worker-data/`.
+32 -17
View File
@@ -12,16 +12,28 @@ for a trusted local team. The coordinator remains the only process with direct
database and artifact-storage access; the browser never calls PostgreSQL and
never receives a worker bearer token.
The first release must be useful before CTX-07--CTX-10 are complete. Therefore
it has two visibly different modes:
## Current delivered scope
The initial operator UI and CTX-09 final reduction are now implemented. The
control room polls a bounded, coordinator-owned read model every two seconds
while a tab is visible. It shows the worker fleet, recent jobs, safe shard
diagnostics, the actual `reducing` phase, and final-result availability. A job
detail page renders the concrete pipeline stages—input accepted, shards,
worker CSVs, reduction, final CSV—from coordinator state and replaces task and
artifact views as work changes. All browser mutations remain limited to
validated dataset upload and operator cancellation.
The interface must distinguish an in-progress distributed search from a run
whose reducer has produced a durable final result:
| Mode | What it proves | What it must not claim |
| --- | --- | --- |
| **Pipeline check** | Upload, task creation, claim, heartbeat, artifact upload, task completion, retries, and downloads work end-to-end. | That multiple shard results have been scientifically reduced into one answer. |
| **Final run** | A reducer has produced a durable final CSV for the full job. | Available only after CTX-09, and for graph only after CTX-10. |
| **In-progress run** | Upload, task creation, claim, heartbeat, artifact upload, task completion, retries, and shard diagnostics work end-to-end. | That the partial CSVs are a global scientific answer. |
| **Final run** | A reducer has produced a durable final CSV for the full job. | Available for `similarity-search` after CTX-09; graph remains unavailable until CTX-10. |
Never label a partial artifact as a final molecular result. The UI must show a
clear `Pipeline check — partial results` badge while a reducer is unavailable.
clear waiting or `reducing` stage until a final artifact exists and the job is
`completed`.
## 2. Constraints and decisions
@@ -74,7 +86,7 @@ clear `Pipeline check — partial results` badge while a reducer is unavailable.
| Worker registration/lease flow | Implemented | Add a read-only worker list; no browser worker controls. |
| Task diagnostics | No public list/detail response | Add sanitized job task list with attempt, status, lease owner, expiry and error. |
| Artifact download | Worker endpoint exists | Add UI-authorized, job-scoped download proxy. |
| Final result | Reducer is not implemented | Gate behind CTX-09; show partial diagnostic artifacts meanwhile. |
| Final result | CTX-09 final artifact and download route exist | Show the `reducing` stage, then make the final CSV prominent only for `completed`. |
| Distributed graph correctness | Planner/reducer unavailable | Do not advertise a multi-shard graph as final until CTX-10. |
## 5. Proposed structure
@@ -173,6 +185,9 @@ Rules:
- TSV file, required, streamed; show expected columns
`chembl_id` and `canonical_smiles`.
- `chunk_rows`: integer 1--100000, default 1000.
- `max_rows`: optional positive integer. The coordinator creates shards only
from the first N data rows, so a user can test a large upload without
creating thousands of tasks. It does not truncate the stored source blob.
- optional human-readable run name is a later schema/API addition; v1 does not
silently store it.
- display file name and client-side size only as convenience; server limits and
@@ -183,11 +198,10 @@ Rules:
Inputs: exactly one `query_smiles` or `query_id`, `top_k`, optional threshold,
threshold direction, `max_rows`, and `progress_every`.
For a runnable manual pipeline check before CTX-08, offer `query_smiles` and
default `chunk_rows` large enough to create one shard. A `query_id` across
multiple shards is disabled with an explanation until CTX-07 resolves it once
before fan-out. The detail page calls an artifact a **partial top-k CSV**, not
a global top-k, until CTX-09 reduction exists.
The current upload form accepts `query_smiles`, because resolving a
cross-shard `query_id` has not yet been connected to coordinator uploads. The
detail page calls an artifact a **partial top-k CSV** until all shards are
complete and CTX-09 reduction stores the final global result.
### 8.3 Similarity graph
@@ -280,7 +294,7 @@ checksum/size metadata display, and prominent partial/final labels.
file; `Content-Disposition` is safe; preview never loads an unbounded CSV; no
final-result button exists before CTX-09.
### WUI-06 — Final-result UX after CTX-09
### WUI-06 — Final-result UX after CTX-09 — implemented
**Depends on:** CTX-09 and WUI-05.
@@ -362,8 +376,9 @@ that the interface exists today.
## 13. Definition of done for the first hand-testable release
WUI-00 through WUI-05 are complete when a clean local checkout can run a
trusted, authenticated local UI; display coordinator readiness, workers, jobs,
tasks and safe errors; submit a valid small search pipeline check; poll it to a
terminal task state; and download/preview the coordinator-owned partial CSV.
The page must make the absence of final reduction impossible to miss.
The hand-testable release is complete when a clean local checkout can run a
trusted, authenticated local UI; display workers, jobs, pipeline stages, tasks
and safe errors; submit a valid small search; poll it through `reducing`; and
download the coordinator-owned final CSV only after completion. The page must
make the distinction between partial diagnostics and the final result
impossible to miss.
+13 -7
View File
@@ -59,7 +59,7 @@ Content-Type: application/json
{
"name": "lab-worker-01",
"capabilities": ["similarity-search", "similarity-graph"],
"capabilities": ["similarity-search"],
"cpu_count": 8,
"memory_mb": 16384
}
@@ -74,8 +74,8 @@ POST /tasks/claim
Content-Type: application/json
{
"worker_id": "worker-01",
"capabilities": ["similarity-search", "similarity-graph"],
"worker_id": "<registered-uuid>",
"capabilities": ["similarity-search"],
"max_concurrency": 1
}
```
@@ -95,7 +95,7 @@ When a task is available, it returns `200 OK`:
"sha256": "..."
},
"parameters": {
"query_id": "CHEMBL939",
"query_smiles": "CCO",
"top_k": 20
}
}
@@ -122,7 +122,10 @@ Content-Type: application/json
},
"metrics": {
"elapsed_seconds": 12.4,
"processed_rows": 10000
"scanned_rows": 10000,
"valid_molecules": 9876,
"invalid_smiles": 124,
"matches_emitted": 20
}
}
```
@@ -189,8 +192,11 @@ class Runner(Protocol):
"""Run one task and return output artifacts plus safe metrics."""
```
`SciMeshRunner` should map `workload` and validated parameters to the existing
SciMesh CLI. For example, a `similarity-search` task invokes:
`SciMeshRunner` maps an allowlisted workload and validated parameters to the
local SciMesh reference functions. A planned `similarity-search` task contains
a resolved `query_smiles` (never `query_id`) and writes one exact local top-k
partial CSV plus the metrics above. Legacy single-shard tasks may still use the
CLI compatibility path:
```text
scimesh similarity-search <local-input> --query-id ... --output <task-dir>/result.csv
+36
View File
@@ -0,0 +1,36 @@
"""Coordinator-independent contracts for distributed SciMesh workloads.
This package defines the typed plan and reduction boundary shared by future
planners, worker adapters, and coordinator bridges. It intentionally has no
network, database, or coordinator imports.
"""
from .models import (
ArtifactReference,
CompletedPartial,
DistributedPlan,
FinalResult,
PlannedTask,
)
from .registry import (
DistributedWorkloadRegistry,
PlanningService,
WorkloadDescription,
default_distributed_registry,
)
from .similarity_search import SimilaritySearchDistributedWorkload
from .workload import DistributedWorkload
__all__ = [
"ArtifactReference",
"CompletedPartial",
"default_distributed_registry",
"DistributedPlan",
"DistributedWorkload",
"DistributedWorkloadRegistry",
"FinalResult",
"PlannedTask",
"PlanningService",
"SimilaritySearchDistributedWorkload",
"WorkloadDescription",
]
+262
View File
@@ -0,0 +1,262 @@
"""Versioned, JSON-safe value objects for distributed workload contracts."""
from __future__ import annotations
import json
import math
import re
from dataclasses import dataclass
from typing import Any, Mapping, Sequence
from uuid import UUID
SCHEMA_VERSION = 1
_WORKLOAD_NAME = re.compile(r"^[a-z][a-z0-9]*(?:-[a-z0-9]+)*$")
def _canonical_uuid(value: object, field: str) -> str:
if not isinstance(value, str):
raise ValueError(f"{field} must be a UUID string")
try:
return str(UUID(value))
except ValueError as error:
raise ValueError(f"{field} must be a UUID string") from error
def _sha256(value: object, field: str) -> str:
if not isinstance(value, str) or not re.fullmatch(r"[0-9a-f]{64}", value):
raise ValueError(f"{field} must be a lowercase SHA-256 hex digest")
return value
def _content_type(value: object, field: str) -> str:
if not isinstance(value, str) or not value or len(value) > 128:
raise ValueError(f"{field} must be a non-empty content type")
if any(character.isspace() or ord(character) < 32 for character in value):
raise ValueError(f"{field} must be a non-empty content type")
return value
def _workload_name(value: object, field: str = "workload") -> str:
if not isinstance(value, str) or not _WORKLOAD_NAME.fullmatch(value):
raise ValueError(f"{field} must be a canonical hyphenated workload name")
return value
def _json_value(value: object, field: str) -> Any:
"""Deep-copy a JSON value and reject non-finite or non-string-key data."""
if value is None or isinstance(value, (bool, int)):
return value
if isinstance(value, str):
# Coordinator-owned artifacts are represented exclusively by
# ArtifactReference. A URI or a local path in a generic JSON payload
# would let a planner accidentally leak a bridge/worker implementation
# detail into durable task metadata.
forbidden_prefixes = ("file://", "worker://", "http://", "https://", "s3://", "/")
is_windows_path = len(value) >= 3 and value[0].isalpha() and value[1:3] in (":/", ":\\")
if value.startswith(forbidden_prefixes) or is_windows_path:
raise ValueError(f"{field} must not contain a URI or local path")
return value
if isinstance(value, float):
if not math.isfinite(value):
raise ValueError(f"{field} must not contain NaN or infinity")
return value
if isinstance(value, Mapping):
copied: dict[str, Any] = {}
for key, child in value.items():
if not isinstance(key, str):
raise ValueError(f"{field} must use string object keys")
copied[key] = _json_value(child, f"{field}.{key}")
return copied
if isinstance(value, (list, tuple)):
return [_json_value(child, f"{field}[]") for child in value]
raise ValueError(f"{field} must contain only JSON-compatible values")
def _json_mapping(value: object, field: str) -> dict[str, Any]:
if not isinstance(value, Mapping):
raise ValueError(f"{field} must be an object")
return _json_value(value, field)
@dataclass(frozen=True)
class ArtifactReference:
"""Immutable coordinator-owned artifact identity used in a plan."""
artifact_id: str
sha256: str
content_type: str
def __post_init__(self) -> None:
object.__setattr__(self, "artifact_id", _canonical_uuid(self.artifact_id, "artifact_id"))
object.__setattr__(self, "sha256", _sha256(self.sha256, "sha256"))
object.__setattr__(self, "content_type", _content_type(self.content_type, "content_type"))
def to_dict(self) -> dict[str, str]:
return {
"artifact_id": self.artifact_id,
"sha256": self.sha256,
"content_type": self.content_type,
}
@classmethod
def from_dict(cls, value: object) -> "ArtifactReference":
if not isinstance(value, Mapping):
raise ValueError("artifact reference must be an object")
_require_exact_keys(value, {"artifact_id", "sha256", "content_type"}, "artifact reference")
return cls(
artifact_id=value["artifact_id"],
sha256=value["sha256"],
content_type=value["content_type"],
)
@dataclass(frozen=True)
class PlannedTask:
"""One deterministically indexed, artifact-backed worker task."""
chunk_index: int
input_artifact: ArtifactReference
parameters: Mapping[str, object]
def __post_init__(self) -> None:
if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0:
raise ValueError("chunk_index must be a non-negative integer")
if not isinstance(self.input_artifact, ArtifactReference):
raise ValueError("input_artifact must be an ArtifactReference")
object.__setattr__(self, "parameters", _json_mapping(self.parameters, "task parameters"))
def to_dict(self) -> dict[str, Any]:
return {
"chunk_index": self.chunk_index,
"input_artifact": self.input_artifact.to_dict(),
"parameters": _json_value(self.parameters, "task parameters"),
}
@classmethod
def from_dict(cls, value: object) -> "PlannedTask":
if not isinstance(value, Mapping):
raise ValueError("planned task must be an object")
_require_exact_keys(value, {"chunk_index", "input_artifact", "parameters"}, "planned task")
return cls(
chunk_index=value["chunk_index"],
input_artifact=ArtifactReference.from_dict(value["input_artifact"]),
parameters=value["parameters"],
)
@dataclass(frozen=True)
class DistributedPlan:
"""The complete schema-versioned output of a distributed planner."""
workload: str
resolved_parameters: Mapping[str, object]
tasks: Sequence[PlannedTask]
schema_version: int = SCHEMA_VERSION
def __post_init__(self) -> None:
if self.schema_version != SCHEMA_VERSION:
raise ValueError(f"schema_version must be {SCHEMA_VERSION}")
object.__setattr__(self, "workload", _workload_name(self.workload))
object.__setattr__(self, "resolved_parameters", _json_mapping(self.resolved_parameters, "resolved_parameters"))
task_list = tuple(self.tasks)
if not task_list:
raise ValueError("plan must contain at least one task")
if any(not isinstance(task, PlannedTask) for task in task_list):
raise ValueError("tasks must contain PlannedTask values")
indexes = [task.chunk_index for task in task_list]
if indexes != sorted(indexes) or len(set(indexes)) != len(indexes):
raise ValueError("tasks must have unique, ascending chunk_index values")
object.__setattr__(self, "tasks", task_list)
def to_dict(self) -> dict[str, Any]:
return {
"schema_version": self.schema_version,
"workload": self.workload,
"resolved_parameters": _json_value(self.resolved_parameters, "resolved_parameters"),
"tasks": [task.to_dict() for task in self.tasks],
}
def to_json(self) -> str:
"""Return stable JSON suitable for hashing, tests, and durable payloads."""
return json.dumps(self.to_dict(), sort_keys=True, separators=(",", ":"), allow_nan=False)
@classmethod
def from_dict(cls, value: object) -> "DistributedPlan":
if not isinstance(value, Mapping):
raise ValueError("distributed plan must be an object")
_require_exact_keys(
value,
{"schema_version", "workload", "resolved_parameters", "tasks"},
"distributed plan",
)
raw_tasks = value["tasks"]
if not isinstance(raw_tasks, list):
raise ValueError("tasks must be an array")
return cls(
schema_version=value["schema_version"],
workload=value["workload"],
resolved_parameters=value["resolved_parameters"],
tasks=tuple(PlannedTask.from_dict(task) for task in raw_tasks),
)
@classmethod
def from_json(cls, value: str) -> "DistributedPlan":
try:
decoded = json.loads(value)
except (TypeError, json.JSONDecodeError) as error:
raise ValueError("distributed plan must be valid JSON") from error
return cls.from_dict(decoded)
@dataclass(frozen=True)
class CompletedPartial:
"""Coordinator-owned partial output supplied to a reducer."""
chunk_index: int
artifact: ArtifactReference
metrics: Mapping[str, int | float]
def __post_init__(self) -> None:
if isinstance(self.chunk_index, bool) or not isinstance(self.chunk_index, int) or self.chunk_index < 0:
raise ValueError("chunk_index must be a non-negative integer")
if not isinstance(self.artifact, ArtifactReference):
raise ValueError("artifact must be an ArtifactReference")
if not isinstance(self.metrics, Mapping):
raise ValueError("metrics must be an object")
metrics: dict[str, int | float] = {}
for name, value in self.metrics.items():
if not isinstance(name, str) or not name:
raise ValueError("metric names must be non-empty strings")
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
raise ValueError("metric values must be finite JSON numbers")
metrics[name] = value
object.__setattr__(self, "metrics", metrics)
@dataclass(frozen=True)
class FinalResult:
"""A reducer's durable output, ready for coordinator persistence."""
artifact: ArtifactReference
metrics: Mapping[str, int | float]
def __post_init__(self) -> None:
if not isinstance(self.artifact, ArtifactReference):
raise ValueError("artifact must be an ArtifactReference")
# Reuse the CompletedPartial metric validation without inventing a fake
# artifact lifecycle or widening the result contract.
object.__setattr__(self, "metrics", CompletedPartial(0, self.artifact, self.metrics).metrics)
def _require_exact_keys(value: Mapping[str, object], expected: set[str], label: str) -> None:
actual = set(value)
if actual != expected:
missing = sorted(expected - actual)
unknown = sorted(actual - expected)
details: list[str] = []
if missing:
details.append(f"missing {', '.join(missing)}")
if unknown:
details.append(f"unknown {', '.join(unknown)}")
raise ValueError(f"{label} has {'; '.join(details)} fields")
+108
View File
@@ -0,0 +1,108 @@
"""Registry and orchestration helpers for distributed workload contracts."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Mapping, Sequence
from .models import CompletedPartial, DistributedPlan, FinalResult, _workload_name
from .workload import DistributedWorkload
@dataclass(frozen=True)
class WorkloadDescription:
"""Safe metadata that a future coordinator or UI may display."""
name: str
description: str
class DistributedWorkloadRegistry:
"""Collect distributed workloads without coupling them to the CLI registry."""
def __init__(self) -> None:
self._workloads: dict[str, DistributedWorkload] = {}
def register(self, workload: DistributedWorkload) -> None:
name = _workload_name(workload.name)
if name in self._workloads:
raise ValueError(f"distributed workload already registered: {name}")
if not isinstance(workload.description, str) or not workload.description.strip():
raise ValueError("distributed workload description must be non-empty")
self._workloads[name] = workload
def require(self, name: str) -> DistributedWorkload:
try:
return self._workloads[_workload_name(name)]
except KeyError as error:
raise ValueError(f"unknown distributed workload: {name}") from error
def descriptions(self) -> tuple[WorkloadDescription, ...]:
return tuple(
WorkloadDescription(name, workload.description)
for name, workload in sorted(self._workloads.items())
)
class PlanningService:
"""Small bridge-safe orchestration around a distributed workload registry.
It writes neither jobs nor artifacts. A Go coordinator bridge can therefore
validate and produce a plan before opening its own all-or-nothing persistence
transaction; CTX-09 will implement that concrete bridge and durable result
orchestration.
"""
def __init__(self, registry: DistributedWorkloadRegistry) -> None:
self._registry = registry
def plan(
self,
workload_name: str,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan:
if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
raise ValueError("shard_rows must be a positive integer")
workload = self._registry.require(workload_name)
workload.validate_job(parameters)
plan = workload.plan(input_path, input_artifact_id, parameters, shard_rows, workspace)
if not isinstance(plan, DistributedPlan):
raise ValueError("distributed planner must return a DistributedPlan")
if plan.workload != workload.name:
raise ValueError("distributed planner returned a plan for another workload")
# Round-trip through the strict wire schema now, before a future bridge
# persists anything. This catches non-JSON values and undeclared fields.
return DistributedPlan.from_json(plan.to_json())
def reduce(
self,
workload_name: str,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult:
workload = self._registry.require(workload_name)
indexes = [partial.chunk_index for partial in partial_results]
if len(indexes) != len(set(indexes)):
raise ValueError("partial results must have unique chunk_index values")
ordered = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index))
result = workload.reduce(ordered, parameters, workspace)
if not isinstance(result, FinalResult):
raise ValueError("distributed reducer must return a FinalResult")
return result
def default_distributed_registry() -> DistributedWorkloadRegistry:
"""Return the currently supported distributed scientific workloads."""
# Delayed import keeps the generic registry independent of concrete RDKit
# workloads and avoids making the contract layer import application setup.
from .similarity_search import SimilaritySearchDistributedWorkload
registry = DistributedWorkloadRegistry()
registry.register(SimilaritySearchDistributedWorkload())
return registry
+380
View File
@@ -0,0 +1,380 @@
"""Distributed planning and reduction for exact molecular similarity search."""
from __future__ import annotations
import csv
import hashlib
import heapq
import math
from pathlib import Path
from typing import Any, Iterator, Mapping, Sequence
from uuid import UUID, uuid5
from rdkit import Chem
from scimesh.chemistry.dataset import MoleculeRecord, find_molecule_by_id, parse_smiles
from scimesh.chemistry.fingerprints import FP_RADIUS, FP_SIZE
from scimesh.workloads.similarity_search import (
SimilarityMatch,
_HeapEntry,
search_similar,
write_search_results,
)
from .models import ArtifactReference, CompletedPartial, DistributedPlan, FinalResult, PlannedTask
_TSV_CONTENT_TYPE = "text/tab-separated-values"
_CSV_CONTENT_TYPE = "text/csv"
_SEARCH_COLUMNS = ("rank", "chembl_id", "canonical_smiles", "similarity")
_REQUIRED_COLUMNS = {"chembl_id", "canonical_smiles"}
def write_similarity_search_partial(output_path: Path, matches: Sequence[SimilarityMatch]) -> None:
"""Write a worker partial with a round-trip score, not display rounding.
The public final CSV continues to use the local CLI's six-decimal display.
A reducer needs the full binary float representation to rank candidates
from separate shards exactly as the single-process reference does.
"""
output_path.parent.mkdir(parents=True, exist_ok=True)
with output_path.open("w", encoding="utf-8", newline="") as destination:
writer = csv.DictWriter(destination, fieldnames=_SEARCH_COLUMNS)
writer.writeheader()
for rank, match in enumerate(matches, start=1):
writer.writerow({
"rank": rank,
"chembl_id": match.molecule_id,
"canonical_smiles": match.smiles,
"similarity": repr(match.similarity),
})
class SimilaritySearchDistributedWorkload:
"""Planner/reducer for exact global top-k Tanimoto similarity search."""
name = "similarity-search"
description = "Exact top-k molecular similarity search over deterministic TSV shards."
def validate_job(self, parameters: Mapping[str, object]) -> None:
allowed = {
"query_id", "query_smiles", "top_k", "threshold",
"threshold_direction", "max_rows", "progress_every",
}
unknown = set(parameters) - allowed
if unknown:
raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}")
query_id = parameters.get("query_id")
query_smiles = parameters.get("query_smiles")
if (query_id is None) == (query_smiles is None):
raise ValueError("exactly one of query_id or query_smiles is required")
if query_id is not None:
self._string(query_id, "query_id")
if query_smiles is not None:
self._string(query_smiles, "query_smiles")
self._positive_int(parameters.get("top_k", 20), "top_k")
if "max_rows" in parameters:
self._positive_int(parameters["max_rows"], "max_rows")
if "progress_every" in parameters:
self._nonnegative_int(parameters["progress_every"], "progress_every")
if "threshold" in parameters:
self._unit_interval(parameters["threshold"], "threshold")
if "threshold_direction" in parameters and parameters["threshold_direction"] not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
def plan(
self,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan:
self.validate_job(parameters)
if not input_path.is_file():
raise ValueError("input_path must be a readable dataset file")
if isinstance(shard_rows, bool) or not isinstance(shard_rows, int) or shard_rows < 1:
raise ValueError("shard_rows must be a positive integer")
try:
input_id = UUID(input_artifact_id)
except ValueError as error:
raise ValueError("input_artifact_id must be a UUID") from error
query_smiles, query_source = self._resolve_query(input_path, parameters)
resolved = self._resolved_parameters(parameters, query_smiles, query_source)
workspace.mkdir(parents=True, exist_ok=True)
shard_paths: list[Path] = []
try:
shard_paths = self._write_shards(input_path, workspace, shard_rows, resolved.get("max_rows"))
tasks = tuple(
PlannedTask(
chunk_index=index,
input_artifact=ArtifactReference(
artifact_id=str(uuid5(input_id, f"scimesh:similarity-search:shard:{index}")),
sha256=_sha256_file(path),
content_type=_TSV_CONTENT_TYPE,
),
parameters=self._task_parameters(resolved),
)
for index, path in enumerate(shard_paths)
)
except Exception:
for path in shard_paths:
path.unlink(missing_ok=True)
raise
return DistributedPlan(self.name, resolved, tasks)
def reduce(
self,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult:
"""Merge materialized partial CSVs into one deterministic final CSV.
The coordinator bridge materializes each downloaded artifact at
``workspace / artifact_id`` before it calls this method. Those local
paths are an ephemeral bridge detail, never present in the plan or task
payload. CTX-09 owns the durable final-artifact upload and job state.
"""
if not partial_results:
raise ValueError("at least one partial result is required")
resolved = self._validate_resolved_parameters(parameters)
top_k = resolved["top_k"]
direction = resolved["threshold_direction"]
heap: list[_HeapEntry] = []
ordered_partials = tuple(sorted(partial_results, key=lambda partial: partial.chunk_index))
indexes = [partial.chunk_index for partial in ordered_partials]
if len(indexes) != len(set(indexes)):
raise ValueError("partial results must have unique chunk_index values")
for partial in ordered_partials:
path = workspace / partial.artifact.artifact_id
if not path.is_file():
raise ValueError("materialized partial result is missing")
if _sha256_file(path) != partial.artifact.sha256:
raise ValueError("materialized partial result checksum does not match its artifact reference")
for match in self._read_partial(path, direction):
rank_key = match.sort_key(direction)
entry = _HeapEntry(match, rank_key)
if len(heap) < top_k:
heapq.heappush(heap, entry)
elif rank_key < heap[0].rank_key:
heapq.heapreplace(heap, entry)
matches = sorted((entry.match for entry in heap), key=lambda match: match.sort_key(direction))
output = workspace / "result.csv"
write_search_results(output, matches)
final_id = uuid5(
UUID(ordered_partials[0].artifact.artifact_id),
"scimesh:similarity-search:final:" + ",".join(
partial.artifact.artifact_id for partial in ordered_partials
),
)
return FinalResult(
ArtifactReference(str(final_id), _sha256_file(output), _CSV_CONTENT_TYPE),
{"matches_emitted": len(matches), "partial_count": len(ordered_partials)},
)
def _resolve_query(
self, input_path: Path, parameters: Mapping[str, object]
) -> tuple[str, dict[str, str]]:
query_id = parameters.get("query_id")
if isinstance(query_id, str):
record = find_molecule_by_id(input_path, query_id)
return Chem.MolToSmiles(record.molecule, canonical=True), {"kind": "chembl_id", "value": query_id}
supplied = parameters["query_smiles"]
assert isinstance(supplied, str) # checked by validate_job
molecule = parse_smiles(supplied)
if molecule is None:
raise ValueError("query_smiles is invalid")
return Chem.MolToSmiles(molecule, canonical=True), {"kind": "smiles", "value": supplied}
def _resolved_parameters(
self, parameters: Mapping[str, object], query_smiles: str, query_source: Mapping[str, str]
) -> dict[str, object]:
resolved: dict[str, object] = {
"query_smiles": query_smiles,
"query_source": dict(query_source),
"top_k": self._positive_int(parameters.get("top_k", 20), "top_k"),
"threshold_direction": parameters.get("threshold_direction", "greater"),
"fingerprint": {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE},
}
if "threshold" in parameters:
resolved["threshold"] = self._unit_interval(parameters["threshold"], "threshold")
if "max_rows" in parameters:
resolved["max_rows"] = self._positive_int(parameters["max_rows"], "max_rows")
if "progress_every" in parameters:
resolved["progress_every"] = self._nonnegative_int(parameters["progress_every"], "progress_every")
return resolved
def _validate_resolved_parameters(self, parameters: Mapping[str, object]) -> dict[str, object]:
query_smiles = self._string(parameters.get("query_smiles"), "query_smiles")
if parse_smiles(query_smiles) is None:
raise ValueError("query_smiles is invalid")
resolved = self._resolved_parameters(
parameters,
Chem.MolToSmiles(parse_smiles(query_smiles), canonical=True),
{"kind": "resolved", "value": query_smiles},
)
# A reducer receives immutable plan metadata, whose query source and
# fixed fingerprint are observational context rather than worker input.
if "fingerprint" in parameters:
fingerprint = parameters["fingerprint"]
if fingerprint != {"algorithm": "morgan", "radius": FP_RADIUS, "fp_size": FP_SIZE}:
raise ValueError("resolved fingerprint does not match SciMesh defaults")
return resolved
@staticmethod
def _task_parameters(resolved: Mapping[str, object]) -> dict[str, object]:
# max_rows is applied before sharding. Passing it to each task would
# silently scan N rows per shard instead of the requested global prefix.
return {
key: value for key, value in resolved.items()
if key in {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
}
def _write_shards(
self, input_path: Path, workspace: Path, shard_rows: int, max_rows: object
) -> list[Path]:
limit = int(max_rows) if isinstance(max_rows, int) else None
paths: list[Path] = []
current: Path | None = None
destination = None
rows_in_shard = 0
seen_rows = 0
try:
with input_path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source, delimiter="\t")
fieldnames = reader.fieldnames or []
if not _REQUIRED_COLUMNS.issubset(set(fieldnames)):
missing = sorted(_REQUIRED_COLUMNS - set(fieldnames))
raise ValueError(f"dataset is missing required columns: {', '.join(missing)}")
for row in reader:
if limit is not None and seen_rows >= limit:
break
if destination is None or rows_in_shard == shard_rows:
if destination is not None:
destination.close()
current = workspace / f"shard-{len(paths)}.tsv"
destination = current.open("w", encoding="utf-8", newline="")
writer = csv.DictWriter(destination, fieldnames=fieldnames, delimiter="\t", lineterminator="\n")
writer.writeheader()
paths.append(current)
rows_in_shard = 0
writer.writerow(row)
rows_in_shard += 1
seen_rows += 1
finally:
if destination is not None:
destination.close()
if not paths:
raise ValueError("dataset has no data rows")
return paths
@staticmethod
def _read_partial(path: Path, direction: object) -> Iterator[SimilarityMatch]:
if not path.is_file():
raise ValueError("materialized partial result is missing")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
previous_key: tuple[float, str, str] | None = None
with path.open("r", encoding="utf-8", newline="") as source:
reader = csv.DictReader(source)
if tuple(reader.fieldnames or ()) != _SEARCH_COLUMNS:
raise ValueError("partial result has an invalid CSV header")
for expected_rank, row in enumerate(reader, start=1):
if set(row) != set(_SEARCH_COLUMNS) or row["rank"] != str(expected_rank):
raise ValueError("partial result has an invalid rank")
try:
similarity = float(row["similarity"])
except (TypeError, ValueError) as error:
raise ValueError("partial result has an invalid similarity") from error
if not math.isfinite(similarity) or not 0 <= similarity <= 1:
raise ValueError("partial result has an invalid similarity")
match = SimilarityMatch(similarity, row["chembl_id"], row["canonical_smiles"])
key = match.sort_key(direction)
if previous_key is not None and key < previous_key:
raise ValueError("partial result is not sorted deterministically")
previous_key = key
yield match
@staticmethod
def _string(value: object, name: str) -> str:
if not isinstance(value, str) or not value.strip() or len(value) > 200:
raise ValueError(f"{name} must be a non-empty string")
return value
@staticmethod
def _positive_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 1:
raise ValueError(f"{name} must be a positive integer")
return value
@staticmethod
def _nonnegative_int(value: object, name: str) -> int:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
raise ValueError(f"{name} must be a non-negative integer")
return value
@staticmethod
def _unit_interval(value: object, name: str) -> float:
if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or not 0 <= value <= 1:
raise ValueError(f"{name} must be a number between 0 and 1")
return float(value)
def run_similarity_search_shard(
input_path: Path, parameters: Mapping[str, object], output_path: Path
) -> dict[str, int]:
"""Run one planned shard using the local reference implementation.
This is the worker adapter used by CTX-08. It deliberately accepts only
resolved ``query_smiles``: resolving an identifier independently in each
shard would make the distributed search scientifically invalid.
"""
allowed = {"query_smiles", "top_k", "threshold", "threshold_direction", "progress_every"}
unknown = set(parameters) - allowed
if unknown:
raise ValueError(f"unsupported similarity-search parameters: {', '.join(sorted(unknown))}")
query_smiles = parameters.get("query_smiles")
if not isinstance(query_smiles, str) or not query_smiles.strip():
raise ValueError("query_smiles is required for a distributed shard")
molecule = parse_smiles(query_smiles)
if molecule is None:
raise ValueError("query_smiles is invalid")
top_k = SimilaritySearchDistributedWorkload._positive_int(parameters.get("top_k", 20), "top_k")
threshold = None
if "threshold" in parameters:
threshold = SimilaritySearchDistributedWorkload._unit_interval(parameters["threshold"], "threshold")
direction = parameters.get("threshold_direction", "greater")
if direction not in {"greater", "less"}:
raise ValueError("threshold_direction must be 'greater' or 'less'")
progress_every = 0
if "progress_every" in parameters:
progress_every = SimilaritySearchDistributedWorkload._nonnegative_int(
parameters["progress_every"], "progress_every"
)
result = search_similar(
input_path,
MoleculeRecord("query", query_smiles, molecule),
top_k=top_k,
progress_every=progress_every,
threshold=threshold,
threshold_direction=direction,
)
write_similarity_search_partial(output_path, result.matches)
return {
"scanned_rows": result.stats.scanned,
"valid_molecules": result.stats.valid,
"invalid_smiles": result.stats.invalid,
"matches_emitted": len(result.matches),
}
def _sha256_file(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as source:
for block in iter(lambda: source.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
+40
View File
@@ -0,0 +1,40 @@
"""Protocol implemented by coordinator-independent distributed workloads."""
from __future__ import annotations
from pathlib import Path
from typing import Mapping, Protocol, Sequence
from .models import CompletedPartial, DistributedPlan, FinalResult
class DistributedWorkload(Protocol):
"""Validate, plan, and reduce one explicit scientific workload.
``input_path`` and ``workspace`` are bridge-provided temporary local paths.
They must never be included in returned plans or persisted task payloads.
"""
name: str
description: str
def validate_job(self, parameters: Mapping[str, object]) -> None:
"""Reject invalid public parameters before the bridge writes metadata."""
def plan(
self,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan:
"""Build a JSON-safe plan containing only coordinator artifact references."""
def reduce(
self,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult:
"""Reduce coordinator-owned partial artifacts in ascending chunk order."""
+31 -6
View File
@@ -13,15 +13,17 @@ from .daemon import WorkerDaemon
from .runners import SciMeshRunner
def main(argv: list[str] | None = None) -> int:
def build_parser() -> argparse.ArgumentParser:
"""Build the worker CLI parser for command-line use and focused tests."""
parser = argparse.ArgumentParser(
prog="scimesh-worker",
epilog=(
"Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, "
"SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, "
"SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, "
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, and "
"SCIMESH_BEARER_TOKEN. SCIMESH_WORKER_ID is a legacy/test override."
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, "
"SCIMESH_MAX_TASKS, and SCIMESH_BEARER_TOKEN. "
"SCIMESH_WORKER_ID is a legacy/test override."
),
)
parser.add_argument("--coordinator-url")
@@ -34,8 +36,31 @@ def main(argv: list[str] | None = None) -> int:
parser.add_argument("--request-timeout", type=float)
parser.add_argument("--heartbeat-interval", type=float)
parser.add_argument("--cleanup-after-seconds", type=float)
lifecycle = parser.add_mutually_exclusive_group()
lifecycle.add_argument(
"--once",
action="store_true",
help="Claim at most one task, then exit; exit immediately when the queue is empty",
)
lifecycle.add_argument(
"--max-tasks",
type=int,
help="Process this many claimed tasks, then exit",
)
return parser
def main(argv: list[str] | None = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
overrides = {key: value for key, value in vars(args).items() if value is not None}
overrides = {
key: value
for key, value in vars(args).items()
if value is not None and key != "once"
}
if args.once:
overrides["max_tasks"] = 1
overrides["exit_when_idle"] = True
if "work_dir" in overrides:
overrides["work_dir"] = Path(overrides["work_dir"])
try:
@@ -44,13 +69,13 @@ def main(argv: list[str] | None = None) -> int:
parser.error(str(error))
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
WorkerDaemon(
completed_without_interruption = WorkerDaemon(
config,
client,
HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token),
SciMeshRunner(),
).run_forever()
return 0
return 0 if completed_without_interruption else 130
if __name__ == "__main__":
+22 -3
View File
@@ -36,13 +36,16 @@ class WorkerConfig:
heartbeat_interval: float = 15.0
bearer_token: str | None = None
cleanup_after_seconds: float | None = None
max_tasks: int | None = None
exit_when_idle: bool = False
# Distributed similarity-graph requires triangular block-pair planning and
# is deliberately not advertised until CTX-10. A normal worker must never
# make a multi-shard graph job appear scientifically complete.
# The local CLI uses hyphens; the first coordinator contract used
# underscores. Advertise both stable spellings while jobs are migrated.
# underscores, so retain the search alias during migration.
capabilities: tuple[str, ...] = (
"similarity-search",
"similarity-graph",
"similarity_search",
"similarity_graph",
)
def __post_init__(self) -> None:
@@ -66,8 +69,21 @@ class WorkerConfig:
_positive_number(self.heartbeat_interval, "heartbeat_interval")
if self.cleanup_after_seconds is not None:
_positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True)
if self.max_tasks is not None:
if (
isinstance(self.max_tasks, bool)
or not isinstance(self.max_tasks, int)
or self.max_tasks < 1
):
raise ValueError("max_tasks must be positive when set")
if not isinstance(self.exit_when_idle, bool):
raise ValueError("exit_when_idle must be a boolean")
if not self.capabilities:
raise ValueError("capabilities cannot be empty")
# Runner subprocesses use a task directory as their cwd. Keep the
# configured root absolute so input/output paths remain valid there
# even when the CLI received a convenient relative --work-dir value.
object.__setattr__(self, "work_dir", self.work_dir.expanduser().resolve())
@classmethod
def from_environment(
@@ -86,6 +102,7 @@ class WorkerConfig:
cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS")
cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS")
return cls(
coordinator_url=url.rstrip("/"),
worker_id=value("worker_id", "SCIMESH_WORKER_ID"),
@@ -98,4 +115,6 @@ class WorkerConfig:
heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")),
bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"),
cleanup_after_seconds=float(cleanup) if cleanup else None,
max_tasks=int(max_tasks) if max_tasks is not None else None,
exit_when_idle=bool(values.get("exit_when_idle", False)),
)
+107 -24
View File
@@ -4,9 +4,12 @@ from __future__ import annotations
import logging
from dataclasses import replace
from dataclasses import dataclass
from pathlib import Path
import random
import re
import shutil
import subprocess
import threading
import time
from datetime import datetime, timezone
@@ -72,6 +75,14 @@ class LeaseHeartbeat:
return seconds
@dataclass(frozen=True)
class RunOnceOutcome:
"""Whether a claim was made and whether that claimed task completed."""
claimed: bool
completed: bool
class WorkerDaemon:
def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None:
self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner
@@ -79,32 +90,68 @@ class WorkerDaemon:
self._registered = False
self.log = logging.getLogger("scimesh.worker")
def run_forever(self) -> None:
def run_forever(self) -> bool:
"""Run until stopped; return false only when interrupted by the operator."""
failures = 0
while True:
try:
if not self._registered:
self._register_worker()
self._cleanup_expired_directories()
claimed = self.run_once()
failures = 0
if not claimed:
self._sleep(self.config.poll_interval)
except CoordinatorTransientError as error:
failures += 1
self._log("failed", error_type=type(error).__name__)
self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
completed_tasks = 0
self._log(
"started",
max_tasks=self.config.max_tasks,
exit_when_idle=self.config.exit_when_idle,
)
try:
while True:
try:
if not self._registered:
self._register_worker()
self._cleanup_expired_directories()
outcome = self.run_once()
failures = 0
if outcome.claimed:
if outcome.completed:
completed_tasks += 1
if self.config.exit_when_idle:
self._log(
"stopped",
reason="one_claim_processed",
completed_tasks=completed_tasks,
)
return True
if (
outcome.completed
and self.config.max_tasks is not None
and completed_tasks >= self.config.max_tasks
):
self._log(
"stopped",
reason="max_tasks_reached",
completed_tasks=completed_tasks,
)
return True
elif self.config.exit_when_idle:
self._log("stopped", reason="queue_empty", completed_tasks=completed_tasks)
return True
else:
self._sleep(self.config.poll_interval)
except CoordinatorTransientError as error:
failures += 1
self._log("failed", error_type=type(error).__name__)
self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
except KeyboardInterrupt:
self._log("stopped", reason="interrupted", completed_tasks=completed_tasks)
return False
def run_once(self) -> bool:
def run_once(self) -> RunOnceOutcome:
worker_id = self._worker_id()
self._log("claiming")
self._log("claiming", log_level=logging.DEBUG)
task = self.coordinator.claim(worker_id, self.config.capabilities)
if task is None:
self._log("idle")
return False
self._log("idle", log_level=logging.DEBUG)
return RunOnceOutcome(claimed=False, completed=False)
started = time.monotonic()
task_dir = self.config.work_dir / task.task_id / str(task.attempt)
heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
completed = False
try:
task_dir.mkdir(parents=True, exist_ok=False)
heartbeat.start()
@@ -135,7 +182,15 @@ class WorkerDaemon:
},
},
)
self._log("idle", task, elapsed_seconds=round(time.monotonic() - started, 3))
completed = True
self._log("completed", task, elapsed_seconds=round(time.monotonic() - started, 3))
except KeyboardInterrupt:
self._log("interrupted", task)
try:
self._report_failure(task, InterruptedError("worker interrupted by operator"))
except CoordinatorTransientError:
self._log("failed", task, error_type="FailureReportError")
raise
except CoordinatorConflictError as error:
self._log("lease_lost", task, error_type=type(error).__name__)
except Exception as error:
@@ -143,17 +198,38 @@ class WorkerDaemon:
self._report_failure(task, error)
finally:
heartbeat.stop()
return True
return RunOnceOutcome(claimed=True, completed=completed)
def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")[:300]
message = self._sanitize_error_message(error)
try:
self.coordinator.fail(task, {"worker_id": self._worker_id(), "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message})
self.coordinator.fail(task, {
"worker_id": self._worker_id(),
"attempt": task.attempt,
"error_code": type(error).__name__,
"error_message": message,
"retryable": self._is_retryable(error),
})
except CoordinatorTransientError:
raise
except Exception:
self._log("failed", task, error_type="FailureReportError")
@staticmethod
def _is_retryable(error: Exception) -> bool:
"""Retry transient worker/transport failures, never invalid scientific input."""
return not isinstance(error, (ValueError, FileNotFoundError, subprocess.CalledProcessError))
def _sanitize_error_message(self, error: Exception) -> str:
"""Keep coordinator-visible failures useful without exposing local paths."""
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")
# CalledProcessError includes the complete argv, including sys.executable
# outside work_dir. Replace POSIX and Windows absolute paths before the
# message reaches the coordinator database or operator UI.
message = re.sub(r"(?<![\w:])[A-Za-z]:\\[^\s'\"\],)]+", "<path>", message)
message = re.sub(r"(?<![\w:])/(?:[^\s'\"\],)]+)", "<path>", message)
return message[:300]
def _register_worker(self) -> None:
registered = self.coordinator.register(
self.config.worker_name,
@@ -180,9 +256,16 @@ class WorkerDaemon:
"""Keep completion payload exact: coordinator owns all artifact metadata."""
return {"artifact_id": uploaded.artifact_id}
def _log(self, state: str, task: ClaimedTask | None = None, **extra: object) -> None:
def _log(
self,
state: str,
task: ClaimedTask | None = None,
*,
log_level: int = logging.INFO,
**extra: object,
) -> None:
fields = {"worker_id": self.config.worker_id, "task_id": task.task_id if task else None, "attempt": task.attempt if task else None, "state": state, **extra}
self.log.info("worker_event %s", fields)
self.log.log(log_level, "worker_event %s", fields)
def _cleanup_expired_directories(self) -> None:
"""Remove only old task attempt directories when retention was configured."""
+11
View File
@@ -7,6 +7,8 @@ import subprocess
import sys
from typing import Protocol
from scimesh.distributed.similarity_search import run_similarity_search_shard
from .models import ClaimedTask, ProducedArtifact, RunResult
@@ -18,6 +20,9 @@ class SciMeshRunner:
"""Allowlisted adapter from coordinator workloads to the local SciMesh CLI."""
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
# The subprocess changes cwd to task_dir. Absolute paths keep a caller
# supplied relative work directory from being resolved twice.
task_dir = task_dir.resolve()
input_path = task_dir / "input"
output_path = task_dir / "result.csv"
# The coordinator contract historically used underscores while the
@@ -31,6 +36,12 @@ class SciMeshRunner:
query_id, query_smiles = params.get("query_id"), params.get("query_smiles")
if (query_id is None) == (query_smiles is None):
raise ValueError("exactly one of query_id or query_smiles is required")
if query_smiles is not None and "max_rows" not in params:
metrics = run_similarity_search_shard(input_path, params, output_path)
return RunResult((ProducedArtifact(output_path, "text/csv"),), metrics)
# Legacy URI jobs may still use query_id or an explicitly task-local
# max_rows value. CTX-08 plans never create those payloads; retain
# CLI execution only for backwards compatibility at this boundary.
top_k = self._positive_int(params, "top_k", default=20)
command += ["--query-id", self._string(params, "query_id")] if query_id is not None else ["--query-smiles", self._string(params, "query_smiles")]
command += ["--top-k", str(top_k)]
+165
View File
@@ -0,0 +1,165 @@
#!/usr/bin/env bash
#
# End-to-end check: two Python workers process separate coordinator shards.
#
# Requires Docker, curl, python3, and an installed scimesh-worker (normally
# from this repository's .venv). The test uses its own Compose project, ports,
# volumes, and temporary worker directories, leaving a developer stack alone.
set -euo pipefail
ROOT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
COORDINATOR_DIR="$ROOT_DIR/coordinator"
COMPOSE_PROJECT=${COMPOSE_PROJECT:-scimesh-two-worker-smoke}
COORDINATOR_PORT=${COORDINATOR_PORT:-18081}
POSTGRES_PORT=${POSTGRES_PORT:-55434}
HOST="http://127.0.0.1:${COORDINATOR_PORT}"
TOKEN=${SCIMESH_SMOKE_TOKEN:-two-worker-smoke-token}
WORK_DIR=$(mktemp -d "${TMPDIR:-/tmp}/scimesh-two-worker-smoke.XXXXXX")
WORKER_ONE_PID=""
WORKER_TWO_PID=""
WORKER_PYTHON=${SCIMESH_WORKER_PYTHON:-"$ROOT_DIR/.venv/bin/python"}
cleanup() {
local exit_code=$?
if [[ "$exit_code" -ne 0 ]]; then
printf '\nTwo-worker smoke failed; worker logs follow.\n' >&2
sed -n '1,200p' "$WORK_DIR/worker-a.log" >&2 || true
sed -n '1,200p' "$WORK_DIR/worker-b.log" >&2 || true
(
cd "$COORDINATOR_DIR"
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
docker compose -p "$COMPOSE_PROJECT" logs coordinator >&2 || true
)
fi
if [[ -n "$WORKER_ONE_PID" ]]; then kill "$WORKER_ONE_PID" 2>/dev/null || true; fi
if [[ -n "$WORKER_TWO_PID" ]]; then kill "$WORKER_TWO_PID" 2>/dev/null || true; fi
if [[ -n "$WORKER_ONE_PID" ]]; then wait "$WORKER_ONE_PID" 2>/dev/null || true; fi
if [[ -n "$WORKER_TWO_PID" ]]; then wait "$WORKER_TWO_PID" 2>/dev/null || true; fi
(
cd "$COORDINATOR_DIR"
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
docker compose -p "$COMPOSE_PROJECT" down -v --remove-orphans >/dev/null 2>&1 || true
)
rm -rf "$WORK_DIR"
exit "$exit_code"
}
trap cleanup EXIT INT TERM
require() {
command -v "$1" >/dev/null || {
printf 'missing required command: %s\n' "$1" >&2
exit 2
}
}
for command in docker curl python3; do require "$command"; done
[[ -x "$WORKER_PYTHON" ]] || {
printf 'worker Python is not executable: %s\n' "$WORKER_PYTHON" >&2
printf 'Set SCIMESH_WORKER_PYTHON to a Python environment with SciMesh and RDKit.\n' >&2
exit 2
}
printf 'Starting isolated coordinator on %s (project %s)\n' "$HOST" "$COMPOSE_PROJECT"
(
cd "$COORDINATOR_DIR"
POSTGRES_PORT="$POSTGRES_PORT" COORDINATOR_PORT="$COORDINATOR_PORT" \
WORKER_AUTH_TOKEN="$TOKEN" UI_AUTH_TOKEN= \
docker compose -p "$COMPOSE_PROJECT" up -d --build
)
for _ in $(seq 1 45); do
if curl -fsS "$HOST/health" >/dev/null; then break; fi
sleep 1
done
curl -fsS "$HOST/health" >/dev/null || {
printf 'coordinator did not become healthy\n' >&2
exit 1
}
start_worker() {
local worker_name=$1
local worker_dir=$2
SCIMESH_COORDINATOR_URL="$HOST" \
SCIMESH_BEARER_TOKEN="$TOKEN" \
SCIMESH_WORKER_NAME="$worker_name" \
SCIMESH_POLL_INTERVAL=0.2 \
"$WORKER_PYTHON" -m scimesh.worker.cli --work-dir "$worker_dir" --max-tasks 2 >"$worker_dir.log" 2>&1 &
STARTED_WORKER_PID=$!
}
start_worker two-worker-smoke-a "$WORK_DIR/worker-a"
WORKER_ONE_PID=$STARTED_WORKER_PID
start_worker two-worker-smoke-b "$WORK_DIR/worker-b"
WORKER_TWO_PID=$STARTED_WORKER_PID
for _ in $(seq 1 30); do
registered=$(docker compose -p "$COMPOSE_PROJECT" -f "$COORDINATOR_DIR/docker-compose.yml" \
exec -T postgres psql -U scimesh -d scimesh -Atc "SELECT count(*) FROM workers" 2>/dev/null || printf '0')
[[ "$registered" == "2" ]] && break
sleep 1
done
[[ "${registered:-0}" == "2" ]] || {
printf 'workers did not register; logs follow\n' >&2
sed -n '1,160p' "$WORK_DIR/worker-a.log" >&2 || true
sed -n '1,160p' "$WORK_DIR/worker-b.log" >&2 || true
exit 1
}
DATASET="$WORK_DIR/fixture.tsv"
printf '%s\n' \
$'chembl_id\tcanonical_smiles' \
$'TEST001\tCC' $'TEST002\tCCC' $'TEST003\tCCCC' $'TEST004\tCCCO' $'TEST005\tCCN' \
$'TEST006\tCCCl' $'TEST007\tCCBr' $'TEST008\tCCF' $'TEST009\tCC=O' $'TEST010\tCC#N' \
$'TEST011\tCO' $'TEST012\tCOC' $'TEST013\tCOCC' $'TEST014\tCN' $'TEST015\tCNC' \
$'TEST016\tO=C=O' $'TEST017\tC1CC1' $'TEST018\tc1ccccc1' $'TEST019\tCC(C)O' $'TEST020\tCC(C)N' \
>"$DATASET"
response=$(curl -fsS -H "Authorization: Bearer $TOKEN" -X POST "$HOST/jobs/upload" \
-F 'workload=similarity-search' \
-F 'parameters={"query_smiles":"CCO","top_k":5,"progress_every":0}' \
-F 'chunk_rows=5' \
-F 'max_rows=20' \
-F "file=@${DATASET};type=text/tab-separated-values")
job_id=$(printf '%s' "$response" | python3 -c 'import json,sys; print(json.load(sys.stdin)["job_id"])')
task_count=$(printf '%s' "$response" | python3 -c 'import json,sys; print(json.load(sys.stdin)["task_count"])')
[[ "$task_count" == "4" ]] || {
printf 'expected four shards, got response: %s\n' "$response" >&2
exit 1
}
printf 'Submitted job %s with four shards\n' "$job_id"
for _ in $(seq 1 90); do
job=$(curl -fsS -H "Authorization: Bearer $TOKEN" "$HOST/jobs/$job_id")
status=$(printf '%s' "$job" | python3 -c 'import json,sys; print(json.load(sys.stdin)["status"])')
[[ "$status" == "completed" || "$status" == "failed" || "$status" == "cancelled" ]] && break
sleep 1
done
printf '%s' "$job" | python3 -c '
import json, sys
job = json.load(sys.stdin)
assert job["status"] == "completed", job
assert job["total"] == 4, job
assert job["completed"] == 4, job
assert job["failed"] == 0, job
'
task_check=$(docker compose -p "$COMPOSE_PROJECT" -f "$COORDINATOR_DIR/docker-compose.yml" \
exec -T postgres psql -U scimesh -d scimesh -Atc \
"SELECT count(*) FROM tasks WHERE job_id = '$job_id'::uuid AND status = 'completed' AND attempt = 1 AND result_artifact_id IS NOT NULL")
[[ "$task_check" == "4" ]] || {
printf 'expected four first-attempt tasks with coordinator artifacts, got %s\n' "$task_check" >&2
exit 1
}
worker_one_results=$(find "$WORK_DIR/worker-a" -name result.csv -type f | wc -l | tr -d ' ')
worker_two_results=$(find "$WORK_DIR/worker-b" -name result.csv -type f | wc -l | tr -d ' ')
[[ "$worker_one_results" -ge 1 && "$worker_two_results" -ge 1 ]] || {
printf 'both workers must process at least one shard (a=%s, b=%s)\n' \
"$worker_one_results" "$worker_two_results" >&2
exit 1
}
printf 'PASS: 4/4 shards completed; worker-a=%s, worker-b=%s\n' \
"$worker_one_results" "$worker_two_results"
+186
View File
@@ -0,0 +1,186 @@
"""Contract tests for the coordinator-independent distributed workload boundary."""
from __future__ import annotations
import json
from pathlib import Path
from typing import Mapping, Sequence
from uuid import NAMESPACE_URL, uuid5
import pytest
from scimesh.distributed import (
ArtifactReference,
CompletedPartial,
DistributedPlan,
DistributedWorkloadRegistry,
FinalResult,
PlannedTask,
PlanningService,
)
def artifact(seed: str, content_type: str = "text/tab-separated-values") -> ArtifactReference:
return ArtifactReference(
artifact_id=str(uuid5(NAMESPACE_URL, seed)),
sha256=(seed.encode("utf-8").hex() * 64)[:64],
content_type=content_type,
)
class DummyWorkload:
"""A deterministic fake workload used to test the generic CTX-07 bridge."""
name = "dummy-workload"
description = "A deterministic test workload."
def __init__(self) -> None:
self.plan_calls = 0
self.received_partials: tuple[CompletedPartial, ...] = ()
def validate_job(self, parameters: Mapping[str, object]) -> None:
if parameters != {"mode": "valid"}:
raise ValueError("mode must be valid")
def plan(
self,
input_path: Path,
input_artifact_id: str,
parameters: Mapping[str, object],
shard_rows: int,
workspace: Path,
) -> DistributedPlan:
self.plan_calls += 1
assert input_path.name == "input.tsv"
assert workspace.name == "workspace"
return DistributedPlan(
workload=self.name,
resolved_parameters={"mode": parameters["mode"], "source": input_artifact_id},
tasks=(
PlannedTask(0, artifact(f"{input_artifact_id}:0"), {"mode": "valid"}),
PlannedTask(1, artifact(f"{input_artifact_id}:1"), {"mode": "valid"}),
),
)
def reduce(
self,
partial_results: Sequence[CompletedPartial],
parameters: Mapping[str, object],
workspace: Path,
) -> FinalResult:
self.received_partials = tuple(partial_results)
return FinalResult(artifact("final", "text/csv"), {"partial_count": len(partial_results)})
def service() -> tuple[PlanningService, DummyWorkload]:
workload = DummyWorkload()
registry = DistributedWorkloadRegistry()
registry.register(workload)
return PlanningService(registry), workload
def test_unknown_workload_is_rejected_before_a_plan_is_written(tmp_path: Path) -> None:
planner, workload = service()
with pytest.raises(ValueError, match="unknown distributed workload"):
planner.plan(
"unknown-workload", tmp_path / "input.tsv", artifact("input").artifact_id,
{"mode": "valid"}, 10, tmp_path / "workspace",
)
assert workload.plan_calls == 0
def test_invalid_job_is_rejected_before_the_planner_runs(tmp_path: Path) -> None:
planner, workload = service()
with pytest.raises(ValueError, match="mode must be valid"):
planner.plan(
"dummy-workload", tmp_path / "input.tsv", artifact("input").artifact_id,
{"mode": "invalid"}, 10, tmp_path / "workspace",
)
assert workload.plan_calls == 0
def test_two_shard_plan_is_deterministic_and_json_serializable(tmp_path: Path) -> None:
planner, _ = service()
input_artifact_id = artifact("input").artifact_id
first = planner.plan(
"dummy-workload", tmp_path / "input.tsv", input_artifact_id,
{"mode": "valid"}, 10, tmp_path / "workspace",
)
second = planner.plan(
"dummy-workload", tmp_path / "input.tsv", input_artifact_id,
{"mode": "valid"}, 10, tmp_path / "workspace",
)
assert first.to_json() == second.to_json()
payload = json.loads(first.to_json())
assert [task["chunk_index"] for task in payload["tasks"]] == [0, 1]
assert all(set(task) == {"chunk_index", "input_artifact", "parameters"} for task in payload["tasks"])
assert DistributedPlan.from_json(first.to_json()) == first
def test_plan_rejects_unsafe_or_non_deterministic_task_payloads() -> None:
with pytest.raises(ValueError, match="unique, ascending"):
DistributedPlan(
workload="dummy-workload",
resolved_parameters={},
tasks=(
PlannedTask(1, artifact("one"), {}),
PlannedTask(0, artifact("zero"), {}),
),
)
with pytest.raises(ValueError, match="JSON-compatible"):
PlannedTask(0, artifact("bad"), {"path": Path("not-serializable")})
with pytest.raises(ValueError, match="URI or local path"):
PlannedTask(0, artifact("uri"), {"input": "file:///tmp/input.tsv"})
with pytest.raises(ValueError, match="canonical hyphenated"):
DistributedPlan("dummy_workload", {}, (PlannedTask(0, artifact("one"), {}),))
def test_reducer_receives_completed_partials_in_chunk_order(tmp_path: Path) -> None:
planner, workload = service()
result = planner.reduce(
"dummy-workload",
(
CompletedPartial(3, artifact("three", "text/csv"), {"scanned_rows": 10}),
CompletedPartial(1, artifact("one", "text/csv"), {"scanned_rows": 10}),
),
{"mode": "valid"},
tmp_path / "workspace",
)
assert [partial.chunk_index for partial in workload.received_partials] == [1, 3]
assert result.metrics == {"partial_count": 2}
def test_reducer_rejects_duplicate_chunk_indexes_before_invocation(tmp_path: Path) -> None:
planner, workload = service()
duplicate = CompletedPartial(0, artifact("partial", "text/csv"), {"scanned_rows": 1})
with pytest.raises(ValueError, match="unique chunk_index"):
planner.reduce("dummy-workload", (duplicate, duplicate), {"mode": "valid"}, tmp_path)
assert workload.received_partials == ()
def test_artifact_references_never_accept_paths_or_uris() -> None:
with pytest.raises(ValueError, match="UUID"):
ArtifactReference("file:///tmp/input.tsv", "a" * 64, "text/csv")
with pytest.raises(ValueError, match="lowercase SHA-256"):
ArtifactReference(str(uuid5(NAMESPACE_URL, "input")), "A" * 64, "text/csv")
def test_registry_descriptions_are_stable_and_duplicate_names_are_rejected() -> None:
registry = DistributedWorkloadRegistry()
first, second = DummyWorkload(), DummyWorkload()
registry.register(first)
assert registry.descriptions()[0].name == "dummy-workload"
with pytest.raises(ValueError, match="already registered"):
registry.register(second)
+184
View File
@@ -0,0 +1,184 @@
"""Scientific reference tests for the CTX-08 distributed search workload."""
from __future__ import annotations
import csv
import hashlib
from pathlib import Path
from uuid import NAMESPACE_URL, uuid5
import pytest
from scimesh.chemistry.dataset import find_molecule_by_id
from scimesh.distributed import (
ArtifactReference,
CompletedPartial,
PlanningService,
default_distributed_registry,
)
from scimesh.distributed.registry import DistributedWorkloadRegistry
from scimesh.distributed.similarity_search import (
SimilaritySearchDistributedWorkload,
run_similarity_search_shard,
write_similarity_search_partial,
)
from scimesh.workloads.similarity_search import search_similar, write_search_results
def make_dataset(path: Path) -> None:
path.write_text(
"chembl_id\tcanonical_smiles\textra\n"
"CHEMBL_QUERY\tCCO\tquery\n"
"CHEMBL_A\tCCCO\ta\n"
"CHEMBL_B\tCCCC\tb\n"
"CHEMBL_INVALID\tnot-a-smiles\tbad\n"
"CHEMBL_DUPLICATE\tCCO\tduplicate\n"
"CHEMBL_C\tCCN\tc\n",
encoding="utf-8",
)
def planner() -> tuple[PlanningService, SimilaritySearchDistributedWorkload]:
workload = SimilaritySearchDistributedWorkload()
registry = DistributedWorkloadRegistry()
registry.register(workload)
return PlanningService(registry), workload
def test_default_registry_exposes_only_supported_distributed_search() -> None:
assert [item.name for item in default_distributed_registry().descriptions()] == ["similarity-search"]
def checksum(path: Path) -> str:
return hashlib.sha256(path.read_bytes()).hexdigest()
def test_query_id_is_resolved_once_before_deterministic_shards(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
dataset = tmp_path / "chembl.tsv"
workspace = tmp_path / "workspace"
make_dataset(dataset)
service, _ = planner()
calls = 0
real_find = find_molecule_by_id
def count_find(path: Path, query_id: str) -> MoleculeRecord:
nonlocal calls
calls += 1
return real_find(path, query_id)
monkeypatch.setattr("scimesh.distributed.similarity_search.find_molecule_by_id", count_find)
input_id = str(uuid5(NAMESPACE_URL, "dataset"))
plan = service.plan(
"similarity-search", dataset, input_id,
{"query_id": "CHEMBL_QUERY", "top_k": 3, "max_rows": 5, "progress_every": 0},
2, workspace,
)
assert calls == 1
assert plan.resolved_parameters["query_smiles"] == "CCO"
assert plan.resolved_parameters["query_source"] == {"kind": "chembl_id", "value": "CHEMBL_QUERY"}
assert [task.chunk_index for task in plan.tasks] == [0, 1, 2]
assert all("query_id" not in task.parameters for task in plan.tasks)
assert all("max_rows" not in task.parameters for task in plan.tasks)
assert all(task.parameters["query_smiles"] == "CCO" for task in plan.tasks)
assert [
sum(1 for _ in path.open(encoding="utf-8")) - 1
for path in sorted(workspace.glob("shard-*.tsv"))
] == [2, 2, 1]
def test_distributed_reduction_matches_single_process_reference(tmp_path: Path) -> None:
dataset = tmp_path / "chembl.tsv"
workspace = tmp_path / "workspace"
make_dataset(dataset)
service, workload = planner()
plan = service.plan(
"similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
{"query_smiles": "CCO", "top_k": 3, "threshold": 0.0}, 2, workspace,
)
partials: list[CompletedPartial] = []
# Worker two finishes the latter shards first. Worker one loses its first
# attempt for shard zero, then retries it last. The reducer must remain
# independent of both completion and retry order.
for task in reversed(plan.tasks):
shard = workspace / f"shard-{task.chunk_index}.tsv"
temporary_partial = workspace / f"worker-output-{task.chunk_index}.csv"
metrics = run_similarity_search_shard(shard, task.parameters, temporary_partial)
partial_id = str(uuid5(NAMESPACE_URL, f"partial:{task.chunk_index}"))
partials.append(
CompletedPartial(
task.chunk_index,
ArtifactReference(
partial_id, checksum(temporary_partial), "text/csv",
),
metrics,
)
)
# The reducer materializes result files under their own coordinator IDs,
# not shard input IDs. Keep this fixture faithful to that boundary.
temporary_partial.rename(workspace / partial_id)
final = workload.reduce(tuple(partials), plan.resolved_parameters, workspace)
reference = tmp_path / "reference.csv"
query_record = find_molecule_by_id(dataset, "CHEMBL_QUERY")
write_search_results(reference, search_similar(dataset, query_record, top_k=3, threshold=0.0).matches)
assert (workspace / "result.csv").read_bytes() == reference.read_bytes()
assert final.metrics == {"matches_emitted": 3, "partial_count": 3}
rows = list(csv.DictReader((workspace / "result.csv").open(encoding="utf-8")))
assert {row["chembl_id"] for row in rows}.isdisjoint({"CHEMBL_QUERY", "CHEMBL_DUPLICATE"})
def test_reducer_rejects_unsorted_or_invalid_partial_csv(tmp_path: Path) -> None:
workspace = tmp_path / "workspace"
workspace.mkdir()
workload = SimilaritySearchDistributedWorkload()
artifact_id = str(uuid5(NAMESPACE_URL, "bad"))
partial_path = workspace / artifact_id
partial_path.write_text(
"rank,chembl_id,canonical_smiles,similarity\n1,A,CC,0.1\n2,B,CCC,0.9\n",
encoding="utf-8",
)
artifact = ArtifactReference(artifact_id, checksum(partial_path), "text/csv")
with pytest.raises(ValueError, match="not sorted"):
workload.reduce(
(CompletedPartial(0, artifact, {"scanned_rows": 2}),),
{"query_smiles": "CCO", "top_k": 2, "threshold_direction": "greater", "fingerprint": {"algorithm": "morgan", "radius": 2, "fp_size": 2048}},
workspace,
)
def test_partial_csv_preserves_exact_scores_for_global_ranking(tmp_path: Path) -> None:
partial = tmp_path / "partial.csv"
# Both values look identical in a six-decimal final CSV. The exact value
# must survive shard transport so the global reducer can still rank them.
from scimesh.workloads.similarity_search import SimilarityMatch
write_similarity_search_partial(
partial,
[SimilarityMatch(0.50000049, "A", "CC"), SimilarityMatch(0.50000048, "B", "CCC")],
)
values = list(csv.DictReader(partial.open(encoding="utf-8")))
assert values[0]["similarity"] == repr(0.50000049)
assert values[1]["similarity"] == repr(0.50000048)
def test_planner_rejects_fingerprint_override_and_invalid_query(tmp_path: Path) -> None:
dataset = tmp_path / "chembl.tsv"
make_dataset(dataset)
service, _ = planner()
with pytest.raises(ValueError, match="unsupported similarity-search parameters"):
service.plan(
"similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
{"query_smiles": "CCO", "fingerprint": {"radius": 1}}, 2, tmp_path / "workspace",
)
with pytest.raises(ValueError, match="query_smiles is invalid"):
service.plan(
"similarity-search", dataset, str(uuid5(NAMESPACE_URL, "dataset")),
{"query_smiles": "invalid"}, 2, tmp_path / "workspace",
)
+263 -18
View File
@@ -1,7 +1,9 @@
from __future__ import annotations
import hashlib
import logging
from pathlib import Path
import subprocess
import time
from datetime import datetime, timedelta, timezone
from urllib.request import Request
@@ -9,8 +11,10 @@ from urllib.request import Request
import pytest
from scimesh.worker.config import WorkerConfig
from scimesh.worker import cli as worker_cli
from scimesh.worker.cli import build_parser
from scimesh.worker.coordinator import CoordinatorTransientError
from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon
from scimesh.worker.daemon import LeaseHeartbeat, RunOnceOutcome, WorkerDaemon
from scimesh.worker.models import (
ClaimedTask,
InputArtifact,
@@ -92,7 +96,7 @@ def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes):
def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
content = b"input fixture"
worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content)
assert worker.run_once() is True
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True)
assert runner.calls == 1
assert len(artifacts.uploaded) == 1
assert coordinator.heartbeats == [("task-1", 1, "worker-1")]
@@ -102,26 +106,241 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
}
def test_worker_executes_a_resolved_similarity_search_shard(tmp_path: Path) -> None:
content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\nINVALID\tnot-a-smiles\n"
task = make_task(content)
task = ClaimedTask(
task.task_id, task.attempt, task.lease_expires_at, task.workload, task.input,
{"query_smiles": "CCO", "top_k": 5, "progress_every": 0},
)
worker, coordinator, artifacts, _, _ = daemon(tmp_path, task, content)
worker.runner = SciMeshRunner()
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True)
output = artifacts.uploaded[0][2].read_text(encoding="utf-8")
assert output.startswith("rank,chembl_id,canonical_smiles,similarity\n")
metrics = coordinator.submissions[0]["metrics"]
assert metrics["scanned_rows"] == 3
assert metrics["valid_molecules"] == 2
assert metrics["invalid_smiles"] == 1
assert metrics["matches_emitted"] == 1
assert isinstance(metrics["elapsed_seconds"], float)
def test_two_workers_complete_resolved_shards_after_one_retry(tmp_path: Path) -> None:
content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n"
first = make_task(content)
first = ClaimedTask(
"retry-task", 1, first.lease_expires_at, "similarity-search", first.input,
{"query_smiles": "CCO", "top_k": 5},
)
second = ClaimedTask(
"other-task", 1, first.lease_expires_at, "similarity-search", first.input,
{"query_smiles": "CCO", "top_k": 5},
)
class RetryCoordinator(FakeCoordinator):
def __init__(self) -> None:
super().__init__(None)
self.queue = [first, second]
self.claimants: list[str] = []
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
self.claimants.append(worker_id)
return self.queue.pop(0) if self.queue else None
def fail(self, task: ClaimedTask, payload: dict) -> None:
self.failures.append(payload)
if task.task_id == "retry-task" and task.attempt == 1 and payload["retryable"]:
self.queue.append(
ClaimedTask(
task.task_id, 2, task.lease_expires_at, task.workload, task.input,
task.parameters,
)
)
class FailFirstAttempt:
def __init__(self) -> None:
self.calls = 0
self.delegate = SciMeshRunner()
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
self.calls += 1
if self.calls == 1:
raise RuntimeError("simulated retryable shard failure")
return self.delegate.run(task, task_dir)
coordinator = RetryCoordinator()
artifacts = FakeArtifacts(content)
worker_a = WorkerDaemon(
WorkerConfig("https://example.test", "worker-a", tmp_path / "worker-a"),
coordinator, artifacts, FailFirstAttempt(),
)
worker_b = WorkerDaemon(
WorkerConfig("https://example.test", "worker-b", tmp_path / "worker-b"),
coordinator, artifacts, SciMeshRunner(),
)
assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=False)
assert worker_b.run_once() == RunOnceOutcome(claimed=True, completed=True)
assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=True)
assert coordinator.claimants == ["worker-a", "worker-b", "worker-a"]
assert len(coordinator.failures) == 1
assert coordinator.failures[0]["retryable"] is True
assert len(coordinator.submissions) == 2
def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
worker, _, _, runner, config = daemon(tmp_path, None, b"")
assert worker.run_once() is False
assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False)
assert runner.calls == 0
assert not config.work_dir.exists()
def test_once_worker_exits_after_an_empty_claim(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
caplog.set_level(logging.INFO, logger="scimesh.worker")
worker, _, _, runner, _ = daemon(tmp_path, None, b"")
worker.config = WorkerConfig(**{**worker.config.__dict__, "exit_when_idle": True, "max_tasks": 1})
assert worker.run_forever() is True
assert runner.calls == 0
assert "queue_empty" in caplog.text
def test_worker_stops_after_the_configured_number_of_claims(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
caplog.set_level(logging.INFO, logger="scimesh.worker")
content = b"input fixture"
worker, _, _, runner, _ = daemon(tmp_path, make_task(content), content)
worker.config = WorkerConfig(**{**worker.config.__dict__, "max_tasks": 1})
assert worker.run_forever() is True
assert runner.calls == 1
assert "max_tasks_reached" in caplog.text
def test_keyboard_interrupt_stops_worker_without_propagating(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None:
caplog.set_level(logging.INFO, logger="scimesh.worker")
class InterruptingCoordinator(FakeCoordinator):
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
raise KeyboardInterrupt
worker, _, _, _, _ = daemon(tmp_path, None, b"")
worker.coordinator = InterruptingCoordinator(None)
assert worker.run_forever() is False
assert "interrupted" in caplog.text
def test_interrupting_an_active_task_reports_a_sanitized_failure(tmp_path: Path) -> None:
content = b"input fixture"
worker, coordinator, _, _, _ = daemon(tmp_path, make_task(content), content)
class InterruptingRunner(FakeRunner):
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
raise KeyboardInterrupt
worker.runner = InterruptingRunner()
with pytest.raises(KeyboardInterrupt):
worker.run_once()
assert coordinator.failures == [
{
"worker_id": "worker-1",
"attempt": 1,
"error_code": "InterruptedError",
"error_message": "worker interrupted by operator",
"retryable": True,
}
]
def test_max_tasks_counts_successes_not_failed_claims(tmp_path: Path) -> None:
successful_content = b"successful input"
class SequencedCoordinator(FakeCoordinator):
def __init__(self) -> None:
super().__init__(None)
self.tasks = [
make_task(b"bad input", "wrong-checksum"),
ClaimedTask(
"task-2",
1,
(datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
"similarity-search",
InputArtifact(
"https://example.test/input",
hashlib.sha256(successful_content).hexdigest(),
),
{"query_id": "CHEMBL1"},
),
]
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
return self.tasks.pop(0) if self.tasks else None
coordinator = SequencedCoordinator()
artifacts, runner = FakeArtifacts(successful_content), FakeRunner()
config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work", max_tasks=1)
worker = WorkerDaemon(config, coordinator, artifacts, runner)
assert worker.run_forever() is True
assert len(coordinator.failures) == 1
assert len(coordinator.submissions) == 1
assert runner.calls == 1
def test_worker_cli_lifecycle_options_are_explicit_and_exclusive() -> None:
parser = build_parser()
assert parser.parse_args(["--once"]).once is True
assert parser.parse_args(["--max-tasks", "2"]).max_tasks == 2
with pytest.raises(SystemExit):
parser.parse_args(["--once", "--max-tasks", "2"])
def test_worker_cli_uses_a_nonzero_exit_code_for_interruption(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
class InterruptedDaemon:
def __init__(self, *_: object) -> None:
pass
def run_forever(self) -> bool:
return False
monkeypatch.setattr(worker_cli, "WorkerDaemon", InterruptedDaemon)
assert worker_cli.main(
["--coordinator-url", "https://example.test", "--work-dir", str(tmp_path)]
) == 130
@pytest.mark.parametrize("value", [0, -1, True])
def test_max_tasks_must_be_positive(value: object, tmp_path: Path) -> None:
with pytest.raises(ValueError, match="max_tasks"):
WorkerConfig("https://example.test", None, tmp_path, max_tasks=value) # type: ignore[arg-type]
def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None:
worker, coordinator, _, runner, _ = daemon(tmp_path, make_task(b"actual", "not-the-hash"), b"actual")
assert worker.run_once() is True
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
assert runner.calls == 0
assert coordinator.failures[0]["error_code"] == "ValueError"
assert coordinator.failures[0]["retryable"] is False
assert not coordinator.submissions
def test_failure_reporting_removes_paths_outside_the_worker_directory(tmp_path: Path) -> None:
worker, coordinator, _, _, _ = daemon(tmp_path, make_task(b"input"), b"input")
error = subprocess.CalledProcessError(
1,
["/home/alice/.venv/bin/python", "-m", "scimesh.cli", "/private/input.tsv"],
)
worker._report_failure(make_task(b"input"), error)
message = coordinator.failures[0]["error_message"]
assert "/home/alice" not in message
assert "/private/input.tsv" not in message
assert "<path>" in message
def test_directory_creation_failure_is_reported(tmp_path: Path) -> None:
content = b"input fixture"
worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
(config.work_dir / "task-1" / "1").mkdir(parents=True)
assert worker.run_once() is True
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
assert coordinator.failures[0]["error_code"] == "FileExistsError"
@@ -223,30 +442,35 @@ def test_runner_maps_graph_and_smiles_search_parameters(tmp_path: Path, monkeypa
runner = SciMeshRunner()
graph = ClaimedTask("graph", 1, "2026-07-30T00:00:00Z", "similarity-graph", InputArtifact("https://example/input", "x"), {"threshold": 0.2, "threshold_direction": "less", "block_size": 42, "max_rows": 7, "progress_every": 0})
search = ClaimedTask("search", 1, "2026-07-30T00:00:00Z", "similarity-search", InputArtifact("https://example/input", "x"), {"query_smiles": "CCO", "top_k": 3})
search_dir = tmp_path / "search"
search_dir.mkdir()
(search_dir / "input").write_text(
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
)
runner.run(graph, tmp_path / "graph")
runner.run(search, tmp_path / "search")
result = runner.run(search, search_dir)
assert "--threshold-direction" in commands[0] and "less" in commands[0]
assert "--block-size" in commands[0] and "42" in commands[0]
assert "--max-rows" in commands[0] and "7" in commands[0]
assert "--query-smiles" in commands[1] and "CCO" in commands[1]
assert len(commands) == 1
assert result.metrics == {
"scanned_rows": 2, "valid_molecules": 2, "invalid_smiles": 0, "matches_emitted": 1,
}
def test_runner_accepts_coordinator_workload_names(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
commands: list[list[str]] = []
def fake_run(command: list[str], **_: object) -> None:
commands.append(command)
output = Path(command[command.index("--output") + 1])
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text("a,b\\n", encoding="utf-8")
monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run)
task_dir = tmp_path / "search"
task_dir.mkdir()
(task_dir / "input").write_text(
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
)
task = ClaimedTask(
"search", 1, "2026-07-30T00:00:00Z", "similarity_search",
InputArtifact("https://example/input", "a" * 64), {"query_smiles": "CCO"},
)
SciMeshRunner().run(task, tmp_path / "search")
assert commands[0][3] == "similarity-search"
result = SciMeshRunner().run(task, task_dir)
assert result.metrics["matches_emitted"] == 1
assert (task_dir / "result.csv").is_file()
def test_claimed_task_rejects_path_traversal_and_invalid_metadata() -> None:
@@ -312,6 +536,27 @@ def test_environment_overrides_allow_cli_only_configuration(monkeypatch: pytest.
assert config.worker_id is None
assert "similarity-search" in config.capabilities
assert "similarity_search" in config.capabilities
assert "similarity-graph" not in config.capabilities
def test_relative_work_dir_is_normalized_for_runner_subprocesses(
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
) -> None:
monkeypatch.chdir(tmp_path)
config = WorkerConfig("https://coordinator.example", None, Path("./worker-data"))
assert config.work_dir == tmp_path / "worker-data"
task_dir = config.work_dir / "task" / "1"
task_dir.mkdir(parents=True)
(task_dir / "input").write_text(
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
)
task = ClaimedTask(
"task", 1, "2026-07-30T00:00:00Z", "similarity-search",
InputArtifact("https://example.test/input", "a" * 64), {"query_smiles": "CCO"},
)
SciMeshRunner().run(task, task_dir)
assert (task_dir / "result.csv").is_file()
def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None: