Compare commits

...
Author SHA1 Message Date
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
45 changed files with 2204 additions and 188 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 *_similarities.csv
test_results.csv test_results.csv
test_structures/ test_structures/
# Local coordinator-worker execution state
worker-data*/
scimesh-worker-data/
+6 -5
View File
@@ -1,11 +1,12 @@
# SciMesh # SciMesh
SciMesh is a scientific-workload framework for molecular datasets. Its public CLI SciMesh is a scientific-workload framework for molecular datasets. Its public CLI
currently runs exact similarity search and sparse similarity-graph construction runs exact similarity search and sparse similarity-graph construction locally in
locally in one Python process; it creates no dense similarity matrix. A Python one Python process; it creates no dense similarity matrix. The Go/PostgreSQL
Worker client and the planned Go/PostgreSQL coordinator contract are tracked in coordinator and Python worker can run a diagnostic, shard-based
the repository, but distributed execution is not available yet; see `similarity-search` pipeline locally. Its CSV artifacts are not a global result
[`STATUS.md`](STATUS.md). until CTX-07--09 add planning and reduction; use the local CLI for scientific
results today. 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`. 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`.
+13 -8
View File
@@ -1,7 +1,7 @@
# SciMesh Status # SciMesh Status
**Updated:** 2026-07-23 **Updated:** 2026-07-24
**Branch baseline:** `main` at `b4a89dd` (coordinator merge) **Branch baseline:** `main` at `f953112` (distributed pipeline hardening)
## Current state ## Current state
@@ -32,27 +32,32 @@ Docker PostgreSQL stack on 2026-07-23.
| CTX-04 Worker registry and HTTP API | Implemented | Registration, claim, heartbeat, result, failure, and status endpoints. | | 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-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-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-07 Distributed workload protocol | Implemented | Versioned Python contract models, registry, strict plan validation, and deterministic reduction ordering are in `scimesh/distributed/`. The concrete molecular planner/reducer remains CTX-08/09. |
| CTX-08 Distributed similarity-search | Not started | Local reference exists. | | 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-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. |
| CTX-10 Distributed similarity-graph | Not started | Local reference exists. | | CTX-10 Distributed similarity-graph | Not started | Local reference exists. |
| CTX-11 Dashboard/operator view | In progress | `feat/web-interface` adds a protected local view: job/task/worker status, dataset upload, diagnostic partial-artifact download, and polling. Final-result reduction remains CTX-09. | | CTX-11 Dashboard/operator view | Implemented (diagnostic scope) | Protected local view: job/task/worker status, validated similarity-search upload, diagnostic partial-artifact download, and bounded polling. Final-result reduction remains CTX-09. |
| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. | | CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. |
## Next recommended assignment ## Next recommended assignment
Assign **CTX-07** to the workload role: define distributed job planning and Assign **CTX-08** to the workload role: implement the molecular
reduction boundaries before implementing distributed search or graph execution. `similarity-search` planner and worker adapter on top of the accepted CTX-07
contract.
## Known constraints ## Known constraints
- Planner/reducer semantics are not implemented; the operator UI labels - The CTX-07 protocol is implemented, but no concrete molecular planner or
`partial_result` files as diagnostic and cannot present them as final output. reducer is registered yet; the operator UI labels `partial_result` files as
diagnostic and cannot present them as final output.
Use the local `scimesh` CLI for complete workload results. Use the local `scimesh` CLI for complete workload results.
- The worker/coordinator flow currently accepts both underscore API workload - The worker/coordinator flow currently accepts both underscore API workload
names and hyphenated CLI names while the contract is consolidated. names and hyphenated CLI names while the contract is consolidated.
- A real-stack worker test uses a small `query_smiles` shard. Resolving a - 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. `query_id` once and sharing it across shards belongs to CTX-07.
- 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 ## Update rule
+4 -4
View File
@@ -71,14 +71,14 @@ func run() error {
useCases := httptransport.UseCases{ useCases := httptransport.UseCases{
RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk), RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk),
CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk), CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk),
SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk), SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk, cfg.DefaultMaxAttempts),
ClaimTask: usecase.NewClaimTask(taskRepo, clk, cfg.LeaseDuration), ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration),
RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration), RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration),
CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk), CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk),
FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk), FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk),
GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo), GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo),
CancelJob: usecase.NewCancelJob(jobRepo, taskRepo, tx, clk), CancelJob: usecase.NewCancelJob(jobRepo, taskRepo, tx, clk),
UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, clk), UploadArtifact: usecase.NewUploadArtifact(taskRepo, artifactRepo, blobStore, tx, clk),
DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore), DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore),
GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore), GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore),
Dashboard: usecase.NewDashboard(uiReadRepo), Dashboard: usecase.NewDashboard(uiReadRepo),
@@ -87,7 +87,7 @@ func run() error {
// Background reapers are tracked so shutdown can wait for them. Without this // Background reapers are tracked so shutdown can wait for them. Without this
// the process would exit mid-UPDATE, and the deferred pool.Close() would pull // the process would exit mid-UPDATE, and the deferred pool.Close() would pull
// connections out from under them. // connections out from under them.
expireLeases := usecase.NewExpireLeases(taskRepo, clk) expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk)
markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter) markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter)
var wg sync.WaitGroup var wg sync.WaitGroup
+45 -1
View File
@@ -7,12 +7,19 @@ import (
"bytes" "bytes"
"fmt" "fmt"
"io" "io"
"strings"
) )
// ErrNoRows is returned when the input has a header but no data rows: a job with // 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. // zero tasks could never complete, so it is rejected at the source.
var ErrNoRows = fmt.Errorf("input has no data rows") 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 // 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 // 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 // parse its shard in isolation. emit is called once per shard, in order, with a
@@ -26,13 +33,25 @@ var ErrNoRows = fmt.Errorf("input has no data rows")
// Only one shard is buffered at a time, so memory is bounded by shard size (a // 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. // 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 { func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error {
return SplitTSVLimit(r, rowsPerShard, 0, emit) return splitTSVLimit(r, rowsPerShard, 0, nil, emit)
} }
// SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows. // 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, // A maxRows value of zero means unlimited. This lets an operator make a small,
// representative pipeline check without materialising a second dataset file. // 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 { 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 { if rowsPerShard <= 0 {
return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard) return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard)
} }
@@ -51,6 +70,11 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
return ErrNoRows // completely empty input return ErrNoRows // completely empty input
} }
header := append([]byte(nil), sc.Bytes()...) header := append([]byte(nil), sc.Bytes()...)
if validateHeader != nil {
if err := validateHeader(header); err != nil {
return err
}
}
var ( var (
buf bytes.Buffer buf bytes.Buffer
@@ -71,9 +95,15 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
for sc.Scan() { for sc.Scan() {
if rows == 0 { 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.Write(header)
buf.WriteByte('\n') 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.Write(sc.Bytes())
buf.WriteByte('\n') buf.WriteByte('\n')
rows++ rows++
@@ -103,3 +133,17 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int,
} }
return nil 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
}
+8
View File
@@ -119,6 +119,14 @@ func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) {
} }
} }
func TestChEMBLSplitRejectsMissingRequiredColumns(t *testing.T) {
err := SplitChEMBLTSVLimit(strings.NewReader("id\tsmiles\nA\tCC\n"), 1, 0,
func(int, io.Reader) error { return nil })
if err == nil || !strings.Contains(err.Error(), "chembl_id") {
t.Errorf("err = %v, want missing-column error", err)
}
}
// The scanned bytes are reused by bufio; the shard buffer must copy them, or a // 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. // later row would corrupt an earlier one. This guards that copy.
func TestSplitDoesNotAliasScannerBuffer(t *testing.T) { func TestSplitDoesNotAliasScannerBuffer(t *testing.T) {
+3
View File
@@ -129,6 +129,9 @@ func LoadConfig() (Config, error) {
if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil { if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil {
return Config{}, err return Config{}, err
} }
if cfg.DefaultMaxAttempts < 1 {
return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive")
}
return cfg, nil return cfg, nil
} }
+11
View File
@@ -32,3 +32,14 @@ func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) {
t.Fatalf("unexpected tokens: %+v", cfg) 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)
}
}
+23 -5
View File
@@ -160,17 +160,18 @@ func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time
return cancelled, nil return cancelled, nil
} }
func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) { func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
r.mu.Lock() r.mu.Lock()
defer r.mu.Unlock() defer r.mu.Unlock()
var n int64 affected := make([]uuid.UUID, 0)
for _, t := range r.tasks { for _, t := range r.tasks {
if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) { if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) &&
t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) {
t.ExpireLease(now) t.ExpireLease(now)
n++ affected = append(affected, t.JobID)
} }
} }
return n, nil return affected, nil
} }
// --- JobRepo ------------------------------------------------------------- // --- JobRepo -------------------------------------------------------------
@@ -298,6 +299,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
return &cp, nil 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 ----------------------------------------------------------- // --- BlobStore -----------------------------------------------------------
type BlobStore struct { type BlobStore struct {
+12
View File
@@ -60,6 +60,18 @@ func (r *UIReadRepo) ListTasksByJob(_ context.Context, jobID uuid.UUID) ([]domai
sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex }) sort.Slice(out, func(i, j int) bool { return out[i].ChunkIndex < out[j].ChunkIndex })
return out, nil 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) { func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) {
if limit < 1 || limit > 100 { if limit < 1 || limit > 100 {
return nil, domain.ErrInvalidInput return nil, domain.ErrInvalidInput
@@ -70,3 +70,33 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact,
a.Kind = domain.ArtifactKind(kind) a.Kind = domain.ArtifactKind(kind)
return &a, nil 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
}
@@ -360,12 +360,23 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) {
t.Errorf("replay must be idempotent, got %v", err) 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 func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) {
other.ResultArtifactID = art2.ID pool := testPool(t)
if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) { ctx := context.Background()
t.Errorf("err = %v, want ErrResultConflict", err) 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)
} }
} }
@@ -337,17 +337,26 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta
ELSE error_message END, ELSE error_message END,
completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END, completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END,
version = version + 1 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) { func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) {
var affected int64 var affected []uuid.UUID
err := withRetry(ctx, func(ctx context.Context) error { 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 { if err != nil {
return err return err
} }
affected = tag.RowsAffected() defer rows.Close()
return nil 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 return affected, err
} }
@@ -75,6 +75,31 @@ func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]dom
return tasks, rows.Err() 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) { func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) {
if limit < 1 || limit > 100 { if limit < 1 || limit > 100 {
return nil, domain.ErrInvalidInput return nil, domain.ErrInvalidInput
@@ -79,6 +79,10 @@ func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) {
s.writeError(w, r, domain.ErrInvalidInput) s.writeError(w, r, domain.ErrInvalidInput)
return 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{ claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{
WorkerID: req.WorkerID, WorkerID: req.WorkerID,
@@ -25,6 +25,7 @@ const uiToken = "ui-secret"
type env struct { type env struct {
ts *httptest.Server ts *httptest.Server
blobs *memstore.BlobStore blobs *memstore.BlobStore
workerID string
} }
func newEnv(t *testing.T, ready func(context.Context) error) *env { func newEnv(t *testing.T, ready func(context.Context) error) *env {
@@ -45,22 +46,28 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur
uc := coordhttp.UseCases{ uc := coordhttp.UseCases{
RegisterWorker: usecase.NewRegisterWorker(work, clk), RegisterWorker: usecase.NewRegisterWorker(work, clk),
CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk), CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk),
SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk), SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk, 3),
ClaimTask: usecase.NewClaimTask(tasks, clk, lease), ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease),
RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease), RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease),
CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk), CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk),
FailTask: usecase.NewFailTask(tasks, jobs, tx, clk), FailTask: usecase.NewFailTask(tasks, jobs, tx, clk),
GetJobStatus: usecase.NewGetJobStatus(jobs, tasks), GetJobStatus: usecase.NewGetJobStatus(jobs, tasks),
CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk), CancelJob: usecase.NewCancelJob(jobs, tasks, tx, clk),
UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, clk), UploadArtifact: usecase.NewUploadArtifact(tasks, arts, blobs, tx, clk),
DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs), DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs),
GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs), GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs),
Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)), 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) 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, configuredUIToken)) ts := httptest.NewServer(srv.Handler(token, configuredUIToken))
t.Cleanup(ts.Close) 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 } func healthy(context.Context) error { return nil }
@@ -68,6 +75,7 @@ func healthy(context.Context) error { return nil }
// do sends an authenticated JSON request and returns status + decoded body. // 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) { func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) {
t.Helper() 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, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body))
req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Authorization", "Bearer "+token)
if body != "" { if body != "" {
@@ -330,6 +338,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) { func TestFullLifecycle(t *testing.T) {
e := newEnv(t, healthy) e := newEnv(t, healthy)
@@ -393,9 +421,9 @@ func TestForeignArtifactResultConflict(t *testing.T) {
func TestUploadDatasetChunksAndServesInput(t *testing.T) { func TestUploadDatasetChunksAndServesInput(t *testing.T) {
e := newEnv(t, healthy) 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 { if code != 201 {
t.Fatalf("upload: status = %d", code) t.Fatalf("upload: status = %d", code)
} }
@@ -421,7 +449,7 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) {
t.Fatalf("get input: status = %d", resp.StatusCode) t.Fatalf("get input: status = %d", resp.StatusCode)
} }
shard, _ := io.ReadAll(resp.Body) 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) t.Errorf("shard missing header: %q", shard)
} }
} }
@@ -430,11 +458,12 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
e := newEnv(t, healthy) e := newEnv(t, healthy)
var buf bytes.Buffer var buf bytes.Buffer
mw := multipart.NewWriter(&buf) mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", "w") _ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", "2") _ = mw.WriteField("chunk_rows", "2")
_ = mw.WriteField("max_rows", "3") _ = mw.WriteField("max_rows", "3")
fw, _ := mw.CreateFormFile("file", "chembl.tsv") fw, _ := mw.CreateFormFile("file", "chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n")) _, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n"))
_ = mw.Close() _ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf) req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Authorization", "Bearer "+token)
@@ -451,6 +480,29 @@ func TestUploadDatasetLimitsRows(t *testing.T) {
} }
} }
func TestUploadDatasetRejectsMissingChEMBLColumns(t *testing.T) {
e := newEnv(t, healthy)
var buf bytes.Buffer
mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", "similarity-search")
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", "2")
fw, _ := mw.CreateFormFile("file", "not-chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n"))
_ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", mw.FormDataContentType())
resp, err := http.DefaultClient.Do(req)
if err != nil {
t.Fatal(err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusBadRequest {
t.Errorf("missing ChEMBL columns = %d, want 400", resp.StatusCode)
}
}
func TestErrorMappings(t *testing.T) { func TestErrorMappings(t *testing.T) {
e := newEnv(t, healthy) e := newEnv(t, healthy)
zero := "00000000-0000-0000-0000-000000000000" zero := "00000000-0000-0000-0000-000000000000"
@@ -478,10 +530,11 @@ func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) {
e := newEnv(t, healthy) e := newEnv(t, healthy)
var buf bytes.Buffer var buf bytes.Buffer
mw := multipart.NewWriter(&buf) 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") _ = mw.WriteField("chunk_rows", "not-a-number")
fw, _ := mw.CreateFormFile("file", "chembl.tsv") 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() _ = mw.Close()
req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf) req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf)
@@ -505,6 +558,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)) e.ts.URL+"/tasks/"+taskID+"/artifacts/r.csv", strings.NewReader(data))
req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Authorization", "Bearer "+token)
req.Header.Set("Content-Type", "text/csv") req.Header.Set("Content-Type", "text/csv")
if worker == "w1" {
worker = e.workerID
}
req.Header.Set("X-Worker-ID", worker) req.Header.Set("X-Worker-ID", worker)
req.Header.Set("X-Task-Attempt", itoa(attempt)) req.Header.Set("X-Task-Attempt", itoa(attempt))
resp, err := http.DefaultClient.Do(req) resp, err := http.DefaultClient.Do(req)
@@ -526,6 +582,7 @@ func (e *env) uploadDataset(t *testing.T, workload string, rows int, tsv string)
var buf bytes.Buffer var buf bytes.Buffer
mw := multipart.NewWriter(&buf) mw := multipart.NewWriter(&buf)
_ = mw.WriteField("workload", workload) _ = mw.WriteField("workload", workload)
_ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`)
_ = mw.WriteField("chunk_rows", itoa(rows)) _ = mw.WriteField("chunk_rows", itoa(rows))
fw, _ := mw.CreateFormFile("file", "chembl.tsv") fw, _ := mw.CreateFormFile("file", "chembl.tsv")
_, _ = io.Copy(fw, strings.NewReader(tsv)) _, _ = io.Copy(fw, strings.NewReader(tsv))
@@ -22,7 +22,7 @@
<script> <script>
const id={{printf "%q" .ID}},state={pending:['Waiting for a worker','waiting','Waiting for an available worker with the required capability.'],leased:['Assigned to a worker','active','A worker has claimed the task and should begin processing shortly.'],running:['Running','active','A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator.'],completed:['Tasks complete','success','Every shard task is complete. Files below are still partial results.'],failed:['Needs attention','danger','One or more shard tasks failed. Open the task list below for details.'],cancelled:['Stopped','waiting','The operator stopped this job. No new shards can be claimed.']}; const id={{printf "%q" .ID}},state={pending:['Waiting for a worker','waiting','Waiting for an available worker with the required capability.'],leased:['Assigned to a worker','active','A worker has claimed the task and should begin processing shortly.'],running:['Running','active','A worker is reading a shard, calculating fingerprints, and uploading its result through the coordinator.'],completed:['Tasks complete','success','Every shard task is complete. Files below are still partial results.'],failed:['Needs attention','danger','One or more shard tasks failed. Open the task list below for details.'],cancelled:['Stopped','waiting','The operator stopped this job. No new shards can be claimed.']};
const stop=document.querySelector('#stop-job');if(stop)stop.addEventListener('click',async()=>{if(!confirm('Stop this job? Unfinished shards will be cancelled.'))return;stop.disabled=true;const response=await fetch('/ui/api/jobs/'+id+'/cancel',{method:'POST'});if(!response.ok){stop.disabled=false;alert('Unable to stop this job.');return}location.reload()}); const stop=document.querySelector('#stop-job');if(stop)stop.addEventListener('click',async()=>{if(!confirm('Stop this job? Unfinished shards will be cancelled.'))return;stop.disabled=true;const response=await fetch('/ui/api/jobs/'+id+'/cancel',{method:'POST'});if(!response.ok){stop.disabled=false;alert('Unable to stop this job.');return}location.reload()});
setInterval(async()=>{try{const response=await fetch('/ui/api/jobs/'+id);if(!response.ok)return;const job=await response.json(),info=state[job.status]||[job.status,'waiting','Status reported by the coordinator.'],done=job.completed+job.failed+job.cancelled,percent=job.total?Math.min(100,Math.floor(done*100/job.total)):0,badge=document.querySelector('#status');badge.textContent=info[0];badge.className='status status-'+info[1];document.querySelector('#hint').textContent=info[2];document.querySelector('#progress').textContent=job.completed+' of '+job.total+' tasks complete'+(job.failed?' · failed: '+job.failed:'')+(job.cancelled?' · stopped: '+job.cancelled:'');document.querySelector('#progress-bar').style.width=percent+'%';for(const key of ['total','completed','pending','failed','cancelled'])document.querySelector('#'+key).textContent=job[key];document.querySelector('#active').textContent=job.leased+job.running}catch(_){}} ,2000); const terminal=new Set(['completed','failed','cancelled']);let timer;const poll=async()=>{try{const response=await fetch('/ui/api/jobs/'+id);if(!response.ok)return;const job=await response.json(),info=state[job.status]||[job.status,'waiting','Status reported by the coordinator.'],done=job.completed+job.failed+job.cancelled,percent=job.total?Math.min(100,Math.floor(done*100/job.total)):0,badge=document.querySelector('#status');badge.textContent=info[0];badge.className='status status-'+info[1];document.querySelector('#hint').textContent=info[2];document.querySelector('#progress').textContent=job.completed+' of '+job.total+' tasks complete'+(job.failed?' · failed: '+job.failed:'')+(job.cancelled?' · stopped: '+job.cancelled:'');document.querySelector('#progress-bar').style.width=percent+'%';for(const key of ['total','completed','pending','failed','cancelled'])document.querySelector('#'+key).textContent=job[key];document.querySelector('#active').textContent=job.leased+job.running;if(terminal.has(job.status)&&timer){clearInterval(timer);timer=undefined}}catch(_){}};const start=()=>{if(!timer&&!document.hidden&&!terminal.has(document.querySelector('#status').textContent.toLowerCase()))timer=setInterval(poll,2000)};document.addEventListener('visibilitychange',()=>{if(document.hidden&&timer){clearInterval(timer);timer=undefined}else start()});start();
</script> </script>
</body> </body>
</html> </html>
+41 -15
View File
@@ -14,12 +14,13 @@ type UploadArtifact struct {
tasks TaskRepository tasks TaskRepository
artifacts ArtifactRepository artifacts ArtifactRepository
blobs BlobStore blobs BlobStore
tx TxManager
clk Clock clk Clock
} }
func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository, func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository,
blobs BlobStore, clk Clock) *UploadArtifact { blobs BlobStore, tx TxManager, clk Clock) *UploadArtifact {
return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, clk: clk} return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, clk: clk}
} }
func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (*domain.Artifact, error) { 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()) { if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) {
return nil, domain.ErrLeaseConflict 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 taskID := task.ID
art, err := domain.NewArtifact(task.JobID, &taskID, domain.ArtifactPartialResult, 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) art.SetContent(sum, size)
// The stream may take longer than the lease. Re-check after it finishes so // The stream may take longer than the lease. Lock the task while re-checking
// an expired worker cannot leave a durable result record behind. Completion // ownership and inserting metadata: completion or another upload cannot race
// performs the same ownership check under its transaction. // this final decision. The database unique index is a second line of defence.
current, err := uc.tasks.Get(ctx, in.TaskID) 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 { if err != nil {
_ = uc.blobs.Delete(ctx, art.StorageKey) _ = uc.blobs.Delete(ctx, art.StorageKey)
return nil, err 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) _ = uc.blobs.Delete(ctx, art.StorageKey)
return nil, domain.ErrLeaseConflict
} }
return durable, nil
// 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
} }
// DownloadArtifact returns an artifact's metadata together with a reader over // DownloadArtifact returns an artifact's metadata together with a reader over
+37
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 // The all-or-none guarantee comes from TxManager: a half-created job would
// leave chunks no worker could ever complete. // leave chunks no worker could ever complete.
func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) { 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)) chunks := make([]domain.ChunkSpec, 0, len(in.Chunks))
for _, c := range in.Chunks { for _, c := range in.Chunks {
chunks = append(chunks, domain.ChunkSpec(c)) chunks = append(chunks, domain.ChunkSpec(c))
@@ -92,6 +103,17 @@ func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error
if job.Status == domain.JobCompleted || job.Status == domain.JobFailed { if job.Status == domain.JobCompleted || job.Status == domain.JobFailed {
return domain.ErrJobNotCancellable 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.JobCompleted || derived == domain.JobFailed {
return domain.ErrJobNotCancellable
}
cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now) cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now)
if err != nil { if err != nil {
return err return err
@@ -212,3 +234,18 @@ func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository
} }
return jobs.UpdateStatus(ctx, jobID, status, completedAt) 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
}
+6 -3
View File
@@ -60,9 +60,9 @@ type TaskRepository interface {
// lease. It returns how many tasks changed. // lease. It returns how many tasks changed.
CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error) CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error)
// ExpireLeases applies the lease-expiry rule to every elapsed task and // ExpireLeases applies the lease-expiry rule to every elapsed task and returns
// reports how many were affected. // the distinct jobs whose aggregate status may have changed.
ExpireLeases(ctx context.Context, now time.Time) (int64, error) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error)
} }
// JobRepository persists jobs. // JobRepository persists jobs.
@@ -89,6 +89,9 @@ type WorkerRepository interface {
type ArtifactRepository interface { type ArtifactRepository interface {
Insert(ctx context.Context, a *domain.Artifact) error Insert(ctx context.Context, a *domain.Artifact) error
Get(ctx context.Context, id uuid.UUID) (*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 // BlobStore holds artifact bytes, addressed by an opaque storage key. It streams
+50 -16
View File
@@ -21,12 +21,15 @@ import (
type ClaimTask struct { type ClaimTask struct {
tasks TaskRepository tasks TaskRepository
jobs JobRepository
workers WorkerRepository
tx TxManager
clock Clock clock Clock
leaseDuration time.Duration leaseDuration time.Duration
} }
func NewClaimTask(tasks TaskRepository, clock Clock, leaseDuration time.Duration) *ClaimTask { func NewClaimTask(tasks TaskRepository, jobs JobRepository, workers WorkerRepository, tx TxManager, clock Clock, leaseDuration time.Duration) *ClaimTask {
return &ClaimTask{tasks: tasks, clock: clock, leaseDuration: leaseDuration} return &ClaimTask{tasks: tasks, jobs: jobs, workers: workers, tx: tx, clock: clock, leaseDuration: leaseDuration}
} }
// Execute reclaims elapsed leases first, then hands out one task. // 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 == "" { if in.WorkerID == "" {
return nil, domain.ErrInvalidInput return nil, domain.ErrInvalidInput
} }
now := uc.clock.Now() workloads := in.Workloads
if workerID, err := uuid.Parse(in.WorkerID); err == nil {
if _, err := uc.tasks.ExpireLeases(ctx, now); err != nil { worker, err := uc.workers.Get(ctx, workerID)
if err != nil {
return nil, err 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{ task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{
Workloads: in.Workloads, Workloads: workloads,
Owner: in.WorkerID, Owner: in.WorkerID,
Now: now, Now: now,
LeaseUntil: now.Add(uc.leaseDuration), LeaseUntil: now.Add(uc.leaseDuration),
}) })
if err != nil {
return err
}
if task != nil {
value := task.AsClaimed()
claimed = &value
}
return nil
})
if err != nil { if err != nil {
return nil, err return nil, err
} }
if task == nil { return claimed, nil // nil means an empty queue
return nil, nil // empty queue is a normal state, not an error
}
claimed := task.AsClaimed()
return &claimed, nil
} }
// --- RenewLease ---------------------------------------------------------- // --- RenewLease ----------------------------------------------------------
@@ -230,18 +252,30 @@ func (uc *FailTask) Execute(ctx context.Context, in FailTaskInput) (*domain.Task
type ExpireLeases struct { type ExpireLeases struct {
tasks TaskRepository tasks TaskRepository
jobs JobRepository
tx TxManager
clock Clock clock Clock
} }
func NewExpireLeases(tasks TaskRepository, clock Clock) *ExpireLeases { func NewExpireLeases(tasks TaskRepository, jobs JobRepository, tx TxManager, clock Clock) *ExpireLeases {
return &ExpireLeases{tasks: tasks, clock: clock} 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: // 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 // several coordinators run it concurrently, and a single atomic UPDATE makes
// the duplicate work harmless — the loser simply updates 0 rows. // the duplicate work harmless — the loser simply updates 0 rows.
func (uc *ExpireLeases) Execute(ctx context.Context) (int64, error) { 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
} }
+7 -2
View File
@@ -15,6 +15,7 @@ type UIReadRepository interface {
GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error) GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error)
ListJobs(ctx context.Context, limit int) ([]domain.Job, error) ListJobs(ctx context.Context, limit int) ([]domain.Job, error)
ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, 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) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error)
ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error) ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error)
} }
@@ -88,12 +89,16 @@ func (d *Dashboard) Overview(ctx context.Context, limit int) (DashboardView, err
return DashboardView{}, err return DashboardView{}, err
} }
out := DashboardView{Jobs: make([]JobCard, 0, len(jobs)), Workers: make([]WorkerCard, 0, len(workers))} 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 { for _, job := range jobs {
tasks, err := d.read.ListTasksByJob(ctx, job.ID) jobIDs = append(jobIDs, job.ID)
}
tasksByJob, err := d.read.ListTasksByJobs(ctx, jobIDs)
if err != nil { if err != nil {
return DashboardView{}, err return DashboardView{}, err
} }
out.Jobs = append(out.Jobs, jobCard(job, tasks)) for _, job := range jobs {
out.Jobs = append(out.Jobs, jobCard(job, tasksByJob[job.ID]))
} }
for _, worker := range workers { 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}) out.Workers = append(out.Workers, WorkerCard{ID: worker.ID.String(), Name: worker.Name, Status: string(worker.Status), Capabilities: worker.Capabilities, LastHeartbeatAt: worker.LastHeartbeatAt})
+84 -5
View File
@@ -4,6 +4,7 @@ import (
"context" "context"
"fmt" "fmt"
"io" "io"
"math"
"github.com/google/uuid" "github.com/google/uuid"
@@ -21,14 +22,21 @@ type SubmitDataset struct {
tasks TaskRepository tasks TaskRepository
tx TxManager tx TxManager
clk Clock clk Clock
maxAttempts int
} }
func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository, func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository,
tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset { tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int) *SubmitDataset {
return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk} 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) { 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() now := uc.clk.Now()
job, err := domain.NewUploadedJob(in.Workload, in.Parameters, 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() cleanup()
return SubmitDatasetResult{}, err return SubmitDatasetResult{}, err
} }
splitErr := chunk.SplitTSVLimit(rc, in.RowsPerShard, in.MaxRows, 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, art, err := domain.NewArtifact(job.ID, nil, domain.ArtifactShard,
fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now) fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now)
if err != nil { if err != nil {
@@ -77,7 +85,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
putKeys = append(putKeys, art.StorageKey) putKeys = append(putKeys, art.StorageKey)
art.SetContent(ssum, ssize) 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 { if err != nil {
return err return err
} }
@@ -88,7 +96,8 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
_ = rc.Close() _ = rc.Close()
if splitErr != nil { if splitErr != nil {
cleanup() 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. // 3. Persist job + all artifacts + all tasks atomically.
@@ -118,6 +127,76 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su
}, nil }, 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 // GetTaskInput resolves a task's input shard and opens it for streaming. The
// caller closes the reader. // caller closes the reader.
type GetTaskInput struct { type GetTaskInput struct {
+104 -10
View File
@@ -68,18 +68,18 @@ func newHarness() *harness {
} }
tx := memstore.Tx{} tx := memstore.Tx{}
h.createJob = usecase.NewCreateJob(h.jobs, h.tasks, tx, h.clk) 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.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk, 3)
h.claim = usecase.NewClaimTask(h.tasks, h.clk, lease) 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.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.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk)
h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk) h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk)
h.status = usecase.NewGetJobStatus(h.jobs, h.tasks) h.status = usecase.NewGetJobStatus(h.jobs, h.tasks)
h.results = usecase.NewListResults(h.tasks) h.results = usecase.NewListResults(h.tasks)
h.register = usecase.NewRegisterWorker(h.work, h.clk) 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.downloadArt = usecase.NewDownloadArtifact(h.arts, h.blobs)
h.getInput = usecase.NewGetTaskInput(h.tasks, 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.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk)
return h return h
} }
@@ -138,6 +138,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) { func TestClaimEmptyQueueReturnsNil(t *testing.T) {
h := newHarness() h := newHarness()
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{"w"}})
@@ -281,7 +319,7 @@ func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) {
h.seedJob(t, "w", 1) h.seedJob(t, "w", 1)
taskID, attempt := h.leaseOne(t, "w1", "w") taskID, attempt := h.leaseOne(t, "w1", "w")
h.uploadArt = usecase.NewUploadArtifact( 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{ _, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{
@@ -414,6 +452,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) { func TestDownloadArtifactRoundTrips(t *testing.T) {
h := newHarness() h := newHarness()
h.seedJob(t, "w", 1) h.seedJob(t, "w", 1)
@@ -434,10 +489,10 @@ func TestDownloadArtifactRoundTrips(t *testing.T) {
func TestSubmitDatasetChunksAndServesInput(t *testing.T) { func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
h := newHarness() 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{ 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), ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
}) })
if err != nil { if err != nil {
@@ -453,7 +508,7 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
t.Errorf("job total = %d, want 3", prog.Total) 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 { if err != nil || c == nil {
t.Fatalf("claim shard: %v", err) t.Fatalf("claim shard: %v", err)
} }
@@ -472,9 +527,9 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) {
func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) { func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
h := newHarness() 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{ res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{
Workload: "w", RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv", 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), ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv),
}) })
if err != nil { if err != nil {
@@ -485,6 +540,26 @@ func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) {
} }
} }
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) { func TestCancelJobInvalidatesClaimedAndPendingTasks(t *testing.T) {
h := newHarness() h := newHarness()
jobID := h.seedJob(t, "w", 3) jobID := h.seedJob(t, "w", 3)
@@ -531,3 +606,22 @@ func TestExpireLeasesReclaims(t *testing.T) {
t.Errorf("expire = (%d, %v), want (1, nil)", n, err) 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;
+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}' -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[@]}" \ check "register without capabilities → 400" 400 -X POST "${HOST}/workers/register" "${auth[@]}" \
-d '{"name":"bad"}' -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
echo "job lifecycle" echo "job lifecycle"
@@ -75,7 +82,7 @@ declare -A our_chunks
task_id="" task_id=""
attempt="" attempt=""
for _ in $(seq 1 40); do 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 [[ -z "$claim" ]] && break # 204: queue drained
read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" | read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" |
@@ -100,7 +107,7 @@ else
fi fi
check "heartbeat" 200 -X POST "${HOST}/tasks/${task_id}/heartbeat" "${auth[@]}" \ 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) --- # --- artifacts + result (uploads happen while the task is still leased) ---
bearer=(-H "Authorization: Bearer ${TOKEN}") bearer=(-H "Authorization: Bearer ${TOKEN}")
@@ -108,43 +115,47 @@ bearer=(-H "Authorization: Bearer ${TOKEN}")
# upload <filename> -> prints the artifact_id # upload <filename> -> prints the artifact_id
upload() { upload() {
curl -sS -X PUT "${HOST}/tasks/${task_id}/artifacts/$1" "${bearer[@]}" \ 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' | --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 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[@]}" \ 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' --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[@]}" \ 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}" \ -H 'Content-Type: text/csv' -H 'X-Worker-ID: impostor' -H "X-Task-Attempt: ${attempt}" \
--data-binary 'x' --data-binary 'x'
# Two result artifacts, uploaded now while the lease is held: one to complete # A retry of a PUT returns the same durable artifact for the task attempt.
# with, a second to prove a different manifest is rejected after completion.
art_id=$(upload primary.csv) art_id=$(upload primary.csv)
art_id2=$(upload secondary.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 "download artifact" 200 "${HOST}/artifacts/${art_id}/download" "${bearer[@]}"
check "foreign worker submits → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \ 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}\"}}" -d "{\"worker_id\":\"impostor\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}"
check "submit result" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \ 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[@]}" \ 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}\"}}" -d "{\"worker_id\":\"${worker_id}\",\"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}\"}}"
check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}" check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}"
echo echo
echo "input validation" echo "input validation"
check "malformed uuid → 400" 400 -X POST "${HOST}/tasks/not-a-uuid/result" "${auth[@]}" \ 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 # Note: Go's encoding/json matches field names case-insensitively, so
# "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key # "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key
# trips DisallowUnknownFields. # trips DisallowUnknownFields.
check "unknown json field → 400" 400 -X POST "${HOST}/tasks/claim" "${auth[@]}" \ 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[@]}" check "unknown job → 404" 404 "${HOST}/jobs/00000000-0000-0000-0000-000000000000" "${auth[@]}"
echo 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 # 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. # fields precede the file part, which the coordinator streams.
up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \ up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \
-F 'workload=similarity_search' \ -F 'workload=similarity-search' \
-F 'parameters={"top_k":10}' \ -F 'parameters={"query_smiles":"CCO","top_k":10}' \
-F 'chunk_rows=2' \ -F 'chunk_rows=2' \
-F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV' -F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV'
id smiles chembl_id canonical_smiles
A CC A CC
B CCC B CCC
C CCCC C CCCC
@@ -180,7 +191,7 @@ fi
up_input="" up_input=""
for _ in $(seq 1 30); do for _ in $(seq 1 30); do
c=$(curl -sS "${bearer[@]}" -H 'Content-Type: application/json' -X POST "${HOST}/tasks/claim" \ 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 [[ -z "$c" ]] && break
cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null) cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null)
[[ "$cj" != "$up_job" ]] && continue [[ "$cj" != "$up_job" ]] && continue
+4 -2
View File
@@ -107,7 +107,8 @@ Content-Type: application/json
``` ```
`cpu_count`/`memory_mb` are accepted for forward compatibility and not yet `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 ## Claim task
@@ -120,7 +121,8 @@ Content-Type: application/json
``` ```
- `204 No Content`: no compatible task. - `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 ```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/`.
+35 -8
View File
@@ -46,24 +46,25 @@ coordinator was started with. Never log it, never send it in an error body.
```http ```http
POST /workers/register 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 }`. Response: `{ "worker_id": "<uuid>", "heartbeat_interval_seconds": 15 }`.
- `capabilities` are the workload names you can run — the coordinator only hands - `capabilities` are fixed at registration — the coordinator only hands you
you matching tasks. matching tasks and a later claim cannot broaden that set.
- **Keep `worker_id`**. Use it as your identity in every later call. Using the - **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 registered UUID is what lets the coordinator track your liveness (it marks
workers offline after they go silent). workers offline after they go silent).
- Current coordinator jobs use `similarity_search` / `similarity_graph`; the - Current diagnostic uploads use `similarity-search` with `query_smiles`. The
reference Python worker also accepts the public CLI spellings with hyphens. 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 ## 2. Claim a task
```http ```http
POST /tasks/claim POST /tasks/claim
{ "worker_id": "<uuid>", "capabilities": ["similarity_search"] } { "worker_id": "<uuid>", "capabilities": ["similarity-search"] }
``` ```
- `200` → a leased task (below). - `200` → a leased task (below).
@@ -74,9 +75,9 @@ POST /tasks/claim
"task_id": "<uuid>", "task_id": "<uuid>",
"attempt": 1, "attempt": 1,
"lease_expires_at": "2026-07-22T12:05:00Z", "lease_expires_at": "2026-07-22T12:05:00Z",
"workload": "similarity_search", "workload": "similarity-search",
"input": { "uri": "/tasks/<uuid>/input", "sha256": "<hex>" }, "input": { "uri": "/tasks/<uuid>/input", "sha256": "<hex>" },
"parameters": { "query_id": "CHEMBL939", "top_k": 20 } "parameters": { "query_smiles": "CCO", "top_k": 20 }
} }
``` ```
@@ -193,6 +194,32 @@ Per the worker contract, at minimum:
- poll interval and request timeout - poll interval and request timeout
- a working directory for downloaded inputs and generated outputs - 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 ## Generate a client from the spec
Instead of hand-writing request code, generate it: Instead of hand-writing request code, generate it:
@@ -0,0 +1,209 @@
# 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/`. It does not implement a molecular
planner, reducer, API endpoint, database migration, or final artifact. Until
CTX-08 and CTX-09 are 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 the local CLI's six-decimal 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-08 implements the similarity-search planner, runner adapter, reducer, and
comparison against the local CLI. CTX-09 persists the final artifact and job
state. CTX-10 defines graph-specific triangular block plans; it must not reuse
the search shard scheme without its pair-coverage invariants.
+9 -7
View File
@@ -96,7 +96,9 @@ paths:
description: > description: >
multipart/form-data. The text fields (`workload`, `parameters`, multipart/form-data. The text fields (`workload`, `parameters`,
`chunk_rows`, `max_rows`) MUST precede the `file` part: the file is streamed, not `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. buffered, so the fields have to be parsed before it arrives. Currently
only diagnostic `similarity-search` with `parameters.query_smiles` is
accepted; distributed graph planning is not implemented.
requestBody: requestBody:
required: true required: true
content: content:
@@ -388,7 +390,7 @@ components:
type: array type: array
minItems: 1 minItems: 1
items: { type: string } items: { type: string }
example: [similarity_search, similarity_graph] example: [similarity-search]
cpu_count: cpu_count:
type: integer type: integer
description: Accepted for forward compatibility; not yet persisted. description: Accepted for forward compatibility; not yet persisted.
@@ -419,7 +421,7 @@ components:
type: object type: object
required: [workload, input_uri, chunks] required: [workload, input_uri, chunks]
properties: properties:
workload: { type: string, example: similarity_search } workload: { type: string, example: similarity-search }
input_uri: { type: string } input_uri: { type: string }
parameters: { type: object, additionalProperties: true } parameters: { type: object, additionalProperties: true }
chunks: chunks:
@@ -437,11 +439,11 @@ components:
type: object type: object
required: [workload, file] required: [workload, file]
properties: properties:
workload: { type: string, example: similarity_search } workload: { type: string, enum: [similarity-search], example: similarity-search }
parameters: parameters:
type: string type: string
description: JSON object, sent as a string form field. description: JSON object, sent as a string form field.
example: '{"top_k":10}' example: '{"query_smiles":"CCO","top_k":10}'
chunk_rows: chunk_rows:
type: integer type: integer
description: Data rows per shard. Default 1000. description: Data rows per shard. Default 1000.
@@ -486,11 +488,11 @@ components:
type: object type: object
required: [worker_id] required: [worker_id]
properties: properties:
worker_id: { type: string } worker_id: { type: string, format: uuid, description: Registered worker identity. }
capabilities: capabilities:
type: array type: array
items: { type: string } items: { type: string }
description: Workloads this worker can run. Empty means "any". description: Accepted for compatibility only; registration capabilities decide eligibility.
max_concurrency: max_concurrency:
type: integer type: integer
description: Accepted; the coordinator leases one task per call. description: Accepted; the coordinator leases one task per call.
+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/`.
+4 -4
View File
@@ -59,7 +59,7 @@ Content-Type: application/json
{ {
"name": "lab-worker-01", "name": "lab-worker-01",
"capabilities": ["similarity-search", "similarity-graph"], "capabilities": ["similarity-search"],
"cpu_count": 8, "cpu_count": 8,
"memory_mb": 16384 "memory_mb": 16384
} }
@@ -74,8 +74,8 @@ POST /tasks/claim
Content-Type: application/json Content-Type: application/json
{ {
"worker_id": "worker-01", "worker_id": "<registered-uuid>",
"capabilities": ["similarity-search", "similarity-graph"], "capabilities": ["similarity-search"],
"max_concurrency": 1 "max_concurrency": 1
} }
``` ```
@@ -95,7 +95,7 @@ When a task is available, it returns `200 OK`:
"sha256": "..." "sha256": "..."
}, },
"parameters": { "parameters": {
"query_id": "CHEMBL939", "query_smiles": "CCO",
"top_k": 20 "top_k": 20
} }
} }
+28
View File
@@ -0,0 +1,28 @@
"""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
from .workload import DistributedWorkload
__all__ = [
"ArtifactReference",
"CompletedPartial",
"DistributedPlan",
"DistributedWorkload",
"DistributedWorkloadRegistry",
"FinalResult",
"PlannedTask",
"PlanningService",
"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")
+96
View File
@@ -0,0 +1,96 @@
"""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-08/09 will implement that concrete bridge and reducers.
"""
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
+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 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( parser = argparse.ArgumentParser(
prog="scimesh-worker", prog="scimesh-worker",
epilog=( epilog=(
"Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, " "Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, "
"SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, " "SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, "
"SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, " "SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, "
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, and " "SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, "
"SCIMESH_BEARER_TOKEN. SCIMESH_WORKER_ID is a legacy/test override." "SCIMESH_MAX_TASKS, and SCIMESH_BEARER_TOKEN. "
"SCIMESH_WORKER_ID is a legacy/test override."
), ),
) )
parser.add_argument("--coordinator-url") 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("--request-timeout", type=float)
parser.add_argument("--heartbeat-interval", type=float) parser.add_argument("--heartbeat-interval", type=float)
parser.add_argument("--cleanup-after-seconds", 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) 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: if "work_dir" in overrides:
overrides["work_dir"] = Path(overrides["work_dir"]) overrides["work_dir"] = Path(overrides["work_dir"])
try: try:
@@ -44,13 +69,13 @@ def main(argv: list[str] | None = None) -> int:
parser.error(str(error)) parser.error(str(error))
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token) client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
WorkerDaemon( completed_without_interruption = WorkerDaemon(
config, config,
client, client,
HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token), HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token),
SciMeshRunner(), SciMeshRunner(),
).run_forever() ).run_forever()
return 0 return 0 if completed_without_interruption else 130
if __name__ == "__main__": if __name__ == "__main__":
+18 -3
View File
@@ -36,13 +36,16 @@ class WorkerConfig:
heartbeat_interval: float = 15.0 heartbeat_interval: float = 15.0
bearer_token: str | None = None bearer_token: str | None = None
cleanup_after_seconds: float | 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 # 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, ...] = ( capabilities: tuple[str, ...] = (
"similarity-search", "similarity-search",
"similarity-graph",
"similarity_search", "similarity_search",
"similarity_graph",
) )
def __post_init__(self) -> None: def __post_init__(self) -> None:
@@ -66,6 +69,15 @@ class WorkerConfig:
_positive_number(self.heartbeat_interval, "heartbeat_interval") _positive_number(self.heartbeat_interval, "heartbeat_interval")
if self.cleanup_after_seconds is not None: if self.cleanup_after_seconds is not None:
_positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True) _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: if not self.capabilities:
raise ValueError("capabilities cannot be empty") raise ValueError("capabilities cannot be empty")
# Runner subprocesses use a task directory as their cwd. Keep the # Runner subprocesses use a task directory as their cwd. Keep the
@@ -90,6 +102,7 @@ class WorkerConfig:
cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS") cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS")
cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1) cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB") memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS")
return cls( return cls(
coordinator_url=url.rstrip("/"), coordinator_url=url.rstrip("/"),
worker_id=value("worker_id", "SCIMESH_WORKER_ID"), worker_id=value("worker_id", "SCIMESH_WORKER_ID"),
@@ -102,4 +115,6 @@ class WorkerConfig:
heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")), heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")),
bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"), bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"),
cleanup_after_seconds=float(cleanup) if cleanup else None, 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)),
) )
+83 -12
View File
@@ -4,8 +4,10 @@ from __future__ import annotations
import logging import logging
from dataclasses import replace from dataclasses import replace
from dataclasses import dataclass
from pathlib import Path from pathlib import Path
import random import random
import re
import shutil import shutil
import threading import threading
import time import time
@@ -72,6 +74,14 @@ class LeaseHeartbeat:
return seconds return seconds
@dataclass(frozen=True)
class RunOnceOutcome:
"""Whether a claim was made and whether that claimed task completed."""
claimed: bool
completed: bool
class WorkerDaemon: class WorkerDaemon:
def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None: def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None:
self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner
@@ -79,32 +89,68 @@ class WorkerDaemon:
self._registered = False self._registered = False
self.log = logging.getLogger("scimesh.worker") 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 failures = 0
completed_tasks = 0
self._log(
"started",
max_tasks=self.config.max_tasks,
exit_when_idle=self.config.exit_when_idle,
)
try:
while True: while True:
try: try:
if not self._registered: if not self._registered:
self._register_worker() self._register_worker()
self._cleanup_expired_directories() self._cleanup_expired_directories()
claimed = self.run_once() outcome = self.run_once()
failures = 0 failures = 0
if not claimed: 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) self._sleep(self.config.poll_interval)
except CoordinatorTransientError as error: except CoordinatorTransientError as error:
failures += 1 failures += 1
self._log("failed", error_type=type(error).__name__) self._log("failed", error_type=type(error).__name__)
self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0)) 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() worker_id = self._worker_id()
self._log("claiming") self._log("claiming", log_level=logging.DEBUG)
task = self.coordinator.claim(worker_id, self.config.capabilities) task = self.coordinator.claim(worker_id, self.config.capabilities)
if task is None: if task is None:
self._log("idle") self._log("idle", log_level=logging.DEBUG)
return False return RunOnceOutcome(claimed=False, completed=False)
started = time.monotonic() started = time.monotonic()
task_dir = self.config.work_dir / task.task_id / str(task.attempt) task_dir = self.config.work_dir / task.task_id / str(task.attempt)
heartbeat = LeaseHeartbeat(task, self.coordinator, self.config) heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
completed = False
try: try:
task_dir.mkdir(parents=True, exist_ok=False) task_dir.mkdir(parents=True, exist_ok=False)
heartbeat.start() heartbeat.start()
@@ -135,7 +181,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: except CoordinatorConflictError as error:
self._log("lease_lost", task, error_type=type(error).__name__) self._log("lease_lost", task, error_type=type(error).__name__)
except Exception as error: except Exception as error:
@@ -143,10 +197,10 @@ class WorkerDaemon:
self._report_failure(task, error) self._report_failure(task, error)
finally: finally:
heartbeat.stop() heartbeat.stop()
return True return RunOnceOutcome(claimed=True, completed=completed)
def _report_failure(self, task: ClaimedTask, error: Exception) -> None: 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: 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})
except CoordinatorTransientError: except CoordinatorTransientError:
@@ -154,6 +208,16 @@ class WorkerDaemon:
except Exception: except Exception:
self._log("failed", task, error_type="FailureReportError") self._log("failed", task, error_type="FailureReportError")
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: def _register_worker(self) -> None:
registered = self.coordinator.register( registered = self.coordinator.register(
self.config.worker_name, self.config.worker_name,
@@ -180,9 +244,16 @@ class WorkerDaemon:
"""Keep completion payload exact: coordinator owns all artifact metadata.""" """Keep completion payload exact: coordinator owns all artifact metadata."""
return {"artifact_id": uploaded.artifact_id} 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} 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: def _cleanup_expired_directories(self) -> None:
"""Remove only old task attempt directories when retention was configured.""" """Remove only old task attempt directories when retention was configured."""
+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)
+139 -5
View File
@@ -1,7 +1,9 @@
from __future__ import annotations from __future__ import annotations
import hashlib import hashlib
import logging
from pathlib import Path from pathlib import Path
import subprocess
import time import time
from datetime import datetime, timedelta, timezone from datetime import datetime, timedelta, timezone
from urllib.request import Request from urllib.request import Request
@@ -9,8 +11,10 @@ from urllib.request import Request
import pytest import pytest
from scimesh.worker.config import WorkerConfig 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.coordinator import CoordinatorTransientError
from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon from scimesh.worker.daemon import LeaseHeartbeat, RunOnceOutcome, WorkerDaemon
from scimesh.worker.models import ( from scimesh.worker.models import (
ClaimedTask, ClaimedTask,
InputArtifact, 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: def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
content = b"input fixture" content = b"input fixture"
worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content) 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 runner.calls == 1
assert len(artifacts.uploaded) == 1 assert len(artifacts.uploaded) == 1
assert coordinator.heartbeats == [("task-1", 1, "worker-1")] assert coordinator.heartbeats == [("task-1", 1, "worker-1")]
@@ -104,24 +108,153 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
def test_no_task_does_not_create_directory(tmp_path: Path) -> None: def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
worker, _, _, runner, config = daemon(tmp_path, None, b"") 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 runner.calls == 0
assert not config.work_dir.exists() 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",
}
]
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: 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") 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 runner.calls == 0
assert coordinator.failures[0]["error_code"] == "ValueError" assert coordinator.failures[0]["error_code"] == "ValueError"
assert not coordinator.submissions 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: def test_directory_creation_failure_is_reported(tmp_path: Path) -> None:
content = b"input fixture" content = b"input fixture"
worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content) worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
(config.work_dir / "task-1" / "1").mkdir(parents=True) (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" assert coordinator.failures[0]["error_code"] == "FileExistsError"
@@ -312,6 +445,7 @@ def test_environment_overrides_allow_cli_only_configuration(monkeypatch: pytest.
assert config.worker_id is None assert config.worker_id is None
assert "similarity-search" in config.capabilities assert "similarity-search" in config.capabilities
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( def test_relative_work_dir_is_normalized_for_runner_subprocesses(