diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml new file mode 100644 index 0000000..da35008 --- /dev/null +++ b/.github/workflows/python.yml @@ -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 diff --git a/.gitignore b/.gitignore index c7419b5..93ddbb4 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,7 @@ results/ *_similarities.csv test_results.csv test_structures/ + +# Local coordinator-worker execution state +worker-data*/ +scimesh-worker-data/ diff --git a/README.md b/README.md index be03abe..b51582f 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,12 @@ # SciMesh SciMesh is a scientific-workload framework for molecular datasets. Its public CLI -currently runs exact similarity search and sparse similarity-graph construction -locally in one Python process; it creates no dense similarity matrix. A Python -Worker client and the planned Go/PostgreSQL coordinator contract are tracked in -the repository, but distributed execution is not available yet; see -[`STATUS.md`](STATUS.md). +runs exact similarity search and sparse similarity-graph construction locally in +one Python process; it creates no dense similarity matrix. The Go/PostgreSQL +coordinator and Python worker can run a diagnostic, shard-based +`similarity-search` pipeline locally. Its CSV artifacts are not a global result +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`. diff --git a/STATUS.md b/STATUS.md index f6a7a23..b3b842d 100644 --- a/STATUS.md +++ b/STATUS.md @@ -36,7 +36,7 @@ Docker PostgreSQL stack on 2026-07-23. | CTX-08 Distributed similarity-search | Not started | Local reference exists. | | CTX-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. | | CTX-10 Distributed similarity-graph | Not started | Local reference exists. | -| CTX-11 Dashboard/operator view | In progress | `feat/web-interface` adds a protected local view: job/task/worker status, dataset upload, diagnostic partial-artifact download, and polling. Final-result reduction remains CTX-09. | +| CTX-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. | ## Next recommended assignment @@ -53,6 +53,9 @@ reduction boundaries before implementing distributed search or graph execution. names and hyphenated CLI names while the contract is consolidated. - A real-stack worker test uses a small `query_smiles` shard. Resolving a `query_id` once and sharing it across shards belongs to CTX-07. +- The 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 diff --git a/coordinator/cmd/coordinator/main.go b/coordinator/cmd/coordinator/main.go index 1134257..8447414 100644 --- a/coordinator/cmd/coordinator/main.go +++ b/coordinator/cmd/coordinator/main.go @@ -71,14 +71,14 @@ func run() error { useCases := httptransport.UseCases{ RegisterWorker: usecase.NewRegisterWorker(workerRepo, clk), CreateJob: usecase.NewCreateJob(jobRepo, taskRepo, tx, clk), - SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk), - ClaimTask: usecase.NewClaimTask(taskRepo, clk, cfg.LeaseDuration), + SubmitDataset: usecase.NewSubmitDataset(blobStore, artifactRepo, jobRepo, taskRepo, tx, clk, cfg.DefaultMaxAttempts), + ClaimTask: usecase.NewClaimTask(taskRepo, jobRepo, workerRepo, tx, clk, cfg.LeaseDuration), RenewLease: usecase.NewRenewLease(taskRepo, workerRepo, tx, clk, cfg.LeaseDuration), CompleteTask: usecase.NewCompleteTask(taskRepo, jobRepo, artifactRepo, tx, clk), FailTask: usecase.NewFailTask(taskRepo, jobRepo, tx, clk), GetJobStatus: usecase.NewGetJobStatus(jobRepo, taskRepo), 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), GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore), Dashboard: usecase.NewDashboard(uiReadRepo), @@ -87,7 +87,7 @@ func run() error { // Background reapers are tracked so shutdown can wait for them. Without this // the process would exit mid-UPDATE, and the deferred pool.Close() would pull // connections out from under them. - expireLeases := usecase.NewExpireLeases(taskRepo, clk) + expireLeases := usecase.NewExpireLeases(taskRepo, jobRepo, tx, clk) markOffline := usecase.NewMarkWorkersOffline(workerRepo, clk, cfg.WorkerOfflineAfter) var wg sync.WaitGroup diff --git a/coordinator/internal/chunk/tsv.go b/coordinator/internal/chunk/tsv.go index f020810..b62e8b5 100644 --- a/coordinator/internal/chunk/tsv.go +++ b/coordinator/internal/chunk/tsv.go @@ -7,12 +7,19 @@ import ( "bytes" "fmt" "io" + "strings" ) // ErrNoRows is returned when the input has a header but no data rows: a job with // zero tasks could never complete, so it is rejected at the source. var ErrNoRows = fmt.Errorf("input has no data rows") +// maxShardBytes bounds the coordinator memory used by one in-progress shard. +// The uploaded file may be much larger: it is first stored on disk, then split +// in small bounded pieces. Operators can lower rowsPerShard when this limit is +// reached rather than exhausting the coordinator process. +const maxShardBytes = 64 << 20 // 64 MiB + // SplitTSV reads a header-plus-rows text stream and cuts it into shards of at // most rowsPerShard data rows. Every shard repeats the header, so a worker can // parse its shard in isolation. emit is called once per shard, in order, with a @@ -26,13 +33,25 @@ var ErrNoRows = fmt.Errorf("input has no data rows") // Only one shard is buffered at a time, so memory is bounded by shard size (a // worker-sized slice of the data), not by the size of the whole dataset. func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reader) error) error { - return SplitTSVLimit(r, rowsPerShard, 0, emit) + return splitTSVLimit(r, rowsPerShard, 0, nil, emit) } // SplitTSVLimit behaves like SplitTSV but emits no more than maxRows data rows. // A maxRows value of zero means unlimited. This lets an operator make a small, // representative pipeline check without materialising a second dataset file. func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error { + return splitTSVLimit(r, rowsPerShard, maxRows, nil, emit) +} + +// SplitChEMBLTSVLimit is the coordinator's scientific-upload splitter. It +// validates the two columns every local SciMesh workload requires before any +// shard task is persisted, while generic SplitTSV remains reusable for future +// non-chemistry workloads. +func SplitChEMBLTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, shard io.Reader) error) error { + return splitTSVLimit(r, rowsPerShard, maxRows, validateChEMBLHeader, emit) +} + +func splitTSVLimit(r io.Reader, rowsPerShard, maxRows int, validateHeader func([]byte) error, emit func(index int, shard io.Reader) error) error { if rowsPerShard <= 0 { return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard) } @@ -51,6 +70,11 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, return ErrNoRows // completely empty input } header := append([]byte(nil), sc.Bytes()...) + if validateHeader != nil { + if err := validateHeader(header); err != nil { + return err + } + } var ( buf bytes.Buffer @@ -71,9 +95,15 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, for sc.Scan() { if rows == 0 { + if len(header)+1 > maxShardBytes { + return fmt.Errorf("TSV header exceeds maximum shard size of %d bytes", maxShardBytes) + } buf.Write(header) buf.WriteByte('\n') } + if buf.Len()+len(sc.Bytes())+1 > maxShardBytes { + return fmt.Errorf("shard exceeds maximum size of %d bytes; lower rowsPerShard", maxShardBytes) + } buf.Write(sc.Bytes()) buf.WriteByte('\n') rows++ @@ -103,3 +133,17 @@ func SplitTSVLimit(r io.Reader, rowsPerShard, maxRows int, emit func(index int, } return nil } + +func validateChEMBLHeader(header []byte) error { + seen := make(map[string]struct{}) + for _, field := range strings.Split(strings.TrimPrefix(string(header), "\ufeff"), "\t") { + seen[field] = struct{}{} + } + if _, ok := seen["chembl_id"]; !ok { + return fmt.Errorf("TSV is missing required column chembl_id") + } + if _, ok := seen["canonical_smiles"]; !ok { + return fmt.Errorf("TSV is missing required column canonical_smiles") + } + return nil +} diff --git a/coordinator/internal/chunk/tsv_test.go b/coordinator/internal/chunk/tsv_test.go index 585b35a..ea3dc04 100644 --- a/coordinator/internal/chunk/tsv_test.go +++ b/coordinator/internal/chunk/tsv_test.go @@ -119,6 +119,14 @@ func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) { } } +func TestChEMBLSplitRejectsMissingRequiredColumns(t *testing.T) { + err := SplitChEMBLTSVLimit(strings.NewReader("id\tsmiles\nA\tCC\n"), 1, 0, + func(int, io.Reader) error { return nil }) + if err == nil || !strings.Contains(err.Error(), "chembl_id") { + t.Errorf("err = %v, want missing-column error", err) + } +} + // The scanned bytes are reused by bufio; the shard buffer must copy them, or a // later row would corrupt an earlier one. This guards that copy. func TestSplitDoesNotAliasScannerBuffer(t *testing.T) { diff --git a/coordinator/internal/infra/config.go b/coordinator/internal/infra/config.go index 5414866..7d6db0e 100644 --- a/coordinator/internal/infra/config.go +++ b/coordinator/internal/infra/config.go @@ -129,6 +129,9 @@ func LoadConfig() (Config, error) { if cfg.DefaultMaxAttempts, err = getEnvInt("DEFAULT_MAX_ATTEMPTS", cfg.DefaultMaxAttempts); err != nil { return Config{}, err } + if cfg.DefaultMaxAttempts < 1 { + return Config{}, fmt.Errorf("DEFAULT_MAX_ATTEMPTS must be positive") + } return cfg, nil } diff --git a/coordinator/internal/infra/config_test.go b/coordinator/internal/infra/config_test.go index f9c0945..953887c 100644 --- a/coordinator/internal/infra/config_test.go +++ b/coordinator/internal/infra/config_test.go @@ -32,3 +32,14 @@ func TestLoadConfigAllowsDistinctUIAndWorkerTokens(t *testing.T) { t.Fatalf("unexpected tokens: %+v", cfg) } } + +func TestLoadConfigRejectsNonPositiveDefaultMaxAttempts(t *testing.T) { + t.Setenv("ENV_FILE", filepath.Join(t.TempDir(), "missing.env")) + t.Setenv("DATABASE_URL", "postgres://test") + t.Setenv("DEFAULT_MAX_ATTEMPTS", "0") + + _, err := LoadConfig() + if err == nil || !strings.Contains(err.Error(), "DEFAULT_MAX_ATTEMPTS") { + t.Fatalf("LoadConfig error = %v, want default-attempt validation", err) + } +} diff --git a/coordinator/internal/memstore/memstore.go b/coordinator/internal/memstore/memstore.go index 44db68e..c7f1076 100644 --- a/coordinator/internal/memstore/memstore.go +++ b/coordinator/internal/memstore/memstore.go @@ -160,17 +160,18 @@ func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time return cancelled, nil } -func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) { +func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) { r.mu.Lock() defer r.mu.Unlock() - var n int64 + affected := make([]uuid.UUID, 0) for _, t := range r.tasks { - if t.Status == domain.TaskLeased && t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) { + if (t.Status == domain.TaskLeased || t.Status == domain.TaskRunning) && + t.LeaseExpiresAt != nil && t.LeaseExpiresAt.Before(now) { t.ExpireLease(now) - n++ + affected = append(affected, t.JobID) } } - return n, nil + return affected, nil } // --- JobRepo ------------------------------------------------------------- @@ -298,6 +299,23 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, return &cp, nil } +func (r *ArtifactRepo) FindPartialResult(_ context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) { + r.mu.Lock() + defer r.mu.Unlock() + for _, a := range r.arts { + if a.TaskID != nil && *a.TaskID == taskID && a.Kind == domain.ArtifactPartialResult && + a.Attempt != nil && *a.Attempt == attempt { + return cloneArtifact(a), nil + } + } + return nil, nil +} + +func cloneArtifact(a *domain.Artifact) *domain.Artifact { + cp := *a + return &cp +} + // --- BlobStore ----------------------------------------------------------- type BlobStore struct { diff --git a/coordinator/internal/memstore/ui_read.go b/coordinator/internal/memstore/ui_read.go index ce36e5c..23bff9d 100644 --- a/coordinator/internal/memstore/ui_read.go +++ b/coordinator/internal/memstore/ui_read.go @@ -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 }) return out, nil } + +func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) { + out := make(map[uuid.UUID][]domain.Task, len(jobIDs)) + for _, id := range jobIDs { + tasks, err := r.ListTasksByJob(ctx, id) + if err != nil { + return nil, err + } + out[id] = tasks + } + return out, nil +} func (r *UIReadRepo) ListWorkers(_ context.Context, limit int) ([]domain.Worker, error) { if limit < 1 || limit > 100 { return nil, domain.ErrInvalidInput diff --git a/coordinator/internal/storage/postgres/artifact_repo.go b/coordinator/internal/storage/postgres/artifact_repo.go index f42abac..8f807f8 100644 --- a/coordinator/internal/storage/postgres/artifact_repo.go +++ b/coordinator/internal/storage/postgres/artifact_repo.go @@ -70,3 +70,33 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, a.Kind = domain.ArtifactKind(kind) return &a, nil } + +func (r *ArtifactRepo) FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) { + sql, args, err := psql.Select(artifactColumns...). + From("artifacts"). + Where(sq.Eq{ + "task_id": taskID, + "attempt": attempt, + "kind": string(domain.ArtifactPartialResult), + }). + ToSql() + if err != nil { + return nil, err + } + + var ( + a domain.Artifact + kind string + ) + err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan( + &a.ID, &a.JobID, &a.TaskID, &a.Attempt, &kind, &a.Filename, &a.StorageKey, + &a.ContentType, &a.SizeBytes, &a.SHA256, &a.CreatedAt) + if errors.Is(err, pgx.ErrNoRows) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("find partial result: %w", err) + } + a.Kind = domain.ArtifactKind(kind) + return &a, nil +} diff --git a/coordinator/internal/storage/postgres/integration_test.go b/coordinator/internal/storage/postgres/integration_test.go index 6b3395a..b4b18ac 100644 --- a/coordinator/internal/storage/postgres/integration_test.go +++ b/coordinator/internal/storage/postgres/integration_test.go @@ -360,12 +360,23 @@ func TestCompleteTaskReplayIsIdempotent(t *testing.T) { t.Errorf("replay must be idempotent, got %v", err) } - // A different result artifact for the same task is a genuine conflict. - art2 := seedArtifact(t, pool, job.ID, &claimed.ID, domain.ArtifactPartialResult) - other := in - other.ResultArtifactID = art2.ID - if _, err := uc.Execute(ctx, other); !errors.Is(err, domain.ErrResultConflict) { - t.Errorf("err = %v, want ErrResultConflict", err) +} + +func TestPartialResultIsUniquePerTaskAttempt(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, tasks := seedJob(t, pool, 1) + taskID := tasks[0].ID + first := seedArtifact(t, pool, job.ID, &taskID, domain.ArtifactPartialResult) + second, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, "retry.csv", "text/csv", time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + attempt := 1 + second.Attempt = &attempt + second.SetContent("other-sha", 5) + if err := NewArtifactRepo(pool).Insert(ctx, second); err == nil { + t.Fatalf("second partial artifact for %s/%d was accepted after %s", taskID, attempt, first.ID) } } diff --git a/coordinator/internal/storage/postgres/task_repo.go b/coordinator/internal/storage/postgres/task_repo.go index 18b2646..11b0616 100644 --- a/coordinator/internal/storage/postgres/task_repo.go +++ b/coordinator/internal/storage/postgres/task_repo.go @@ -337,17 +337,26 @@ SET status = CASE WHEN attempt < max_attempts THEN 'pending'::task_sta ELSE error_message END, completed_at = CASE WHEN attempt >= max_attempts THEN $1 ELSE completed_at END, version = version + 1 -WHERE status IN ('leased','running') AND lease_expires_at < $1` +WHERE status IN ('leased','running') AND lease_expires_at < $1 +RETURNING job_id` -func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) { - var affected int64 +func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) { + var affected []uuid.UUID err := withRetry(ctx, func(ctx context.Context) error { - tag, err := conn(ctx, r.pool).Exec(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired) + rows, err := conn(ctx, r.pool).Query(ctx, expireLeasesSQL, now, domain.ErrCodeLeaseExpired) if err != nil { return err } - affected = tag.RowsAffected() - return nil + defer rows.Close() + affected = affected[:0] + for rows.Next() { + var jobID uuid.UUID + if err := rows.Scan(&jobID); err != nil { + return err + } + affected = append(affected, jobID) + } + return rows.Err() }) return affected, err } diff --git a/coordinator/internal/storage/postgres/ui_read_repo.go b/coordinator/internal/storage/postgres/ui_read_repo.go index 4a24ed2..81270b1 100644 --- a/coordinator/internal/storage/postgres/ui_read_repo.go +++ b/coordinator/internal/storage/postgres/ui_read_repo.go @@ -75,6 +75,31 @@ func (r *UIReadRepo) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]dom return tasks, rows.Err() } +func (r *UIReadRepo) ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) { + out := make(map[uuid.UUID][]domain.Task, len(jobIDs)) + if len(jobIDs) == 0 { + return out, nil + } + sql, args, err := psql.Select(taskColumns...).From("tasks"). + Where(sq.Eq{"job_id": jobIDs}).OrderBy("job_id ASC", "chunk_index ASC").ToSql() + if err != nil { + return nil, err + } + rows, err := conn(ctx, r.pool).Query(ctx, sql, args...) + if err != nil { + return nil, fmt.Errorf("list tasks by jobs: %w", err) + } + defer rows.Close() + for rows.Next() { + task, err := scanTask(rows) + if err != nil { + return nil, err + } + out[task.JobID] = append(out[task.JobID], *task) + } + return out, rows.Err() +} + func (r *UIReadRepo) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) { if limit < 1 || limit > 100 { return nil, domain.ErrInvalidInput diff --git a/coordinator/internal/transport/http/handlers.go b/coordinator/internal/transport/http/handlers.go index 4bbf5c8..33ad06b 100644 --- a/coordinator/internal/transport/http/handlers.go +++ b/coordinator/internal/transport/http/handlers.go @@ -79,6 +79,10 @@ func (s *Server) handleClaim(w http.ResponseWriter, r *http.Request) { s.writeError(w, r, domain.ErrInvalidInput) return } + if _, err := uuid.Parse(req.WorkerID); err != nil { + s.writeError(w, r, domain.ErrInvalidInput) + return + } claimed, err := s.uc.ClaimTask.Execute(ctx, usecase.ClaimTaskInput{ WorkerID: req.WorkerID, diff --git a/coordinator/internal/transport/http/server_test.go b/coordinator/internal/transport/http/server_test.go index 261197c..52c453f 100644 --- a/coordinator/internal/transport/http/server_test.go +++ b/coordinator/internal/transport/http/server_test.go @@ -23,8 +23,9 @@ const token = "secret" const uiToken = "ui-secret" type env struct { - ts *httptest.Server - blobs *memstore.BlobStore + ts *httptest.Server + blobs *memstore.BlobStore + workerID string } func newEnv(t *testing.T, ready func(context.Context) error) *env { @@ -45,22 +46,28 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur uc := coordhttp.UseCases{ RegisterWorker: usecase.NewRegisterWorker(work, clk), CreateJob: usecase.NewCreateJob(jobs, tasks, tx, clk), - SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk), - ClaimTask: usecase.NewClaimTask(tasks, clk, lease), + SubmitDataset: usecase.NewSubmitDataset(blobs, arts, jobs, tasks, tx, clk, 3), + ClaimTask: usecase.NewClaimTask(tasks, jobs, work, tx, clk, lease), RenewLease: usecase.NewRenewLease(tasks, work, tx, clk, lease), CompleteTask: usecase.NewCompleteTask(tasks, jobs, arts, tx, clk), FailTask: usecase.NewFailTask(tasks, jobs, tx, clk), GetJobStatus: usecase.NewGetJobStatus(jobs, tasks), 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), GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs), Dashboard: usecase.NewDashboard(memstore.NewUIReadRepo(jobs, tasks, work, arts)), } + worker, err := uc.RegisterWorker.Execute(context.Background(), usecase.RegisterWorkerInput{ + Name: "test-worker", Capabilities: []string{"w", "similarity-search"}, + }) + if err != nil { + t.Fatalf("register test worker: %v", err) + } srv := coordhttp.NewServer(uc, slog.New(slog.NewTextHandler(io.Discard, nil)), 5*time.Second, 15*time.Second, 1<<30, ready) ts := httptest.NewServer(srv.Handler(token, configuredUIToken)) t.Cleanup(ts.Close) - return &env{ts: ts, blobs: blobs} + return &env{ts: ts, blobs: blobs, workerID: worker.ID.String()} } func healthy(context.Context) error { return nil } @@ -68,6 +75,7 @@ func healthy(context.Context) error { return nil } // do sends an authenticated JSON request and returns status + decoded body. func (e *env) do(t *testing.T, method, path, body string) (int, map[string]any) { t.Helper() + body = strings.ReplaceAll(body, `"worker_id":"w1"`, `"worker_id":"`+e.workerID+`"`) req, _ := http.NewRequestWithContext(context.Background(), method, e.ts.URL+path, strings.NewReader(body)) req.Header.Set("Authorization", "Bearer "+token) if body != "" { @@ -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) { e := newEnv(t, healthy) @@ -393,9 +421,9 @@ func TestForeignArtifactResultConflict(t *testing.T) { func TestUploadDatasetChunksAndServesInput(t *testing.T) { e := newEnv(t, healthy) - tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" + tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" - code, body := e.uploadDataset(t, "w", 2, tsv) + code, body := e.uploadDataset(t, "similarity-search", 2, tsv) if code != 201 { t.Fatalf("upload: status = %d", code) } @@ -421,7 +449,7 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) { t.Fatalf("get input: status = %d", resp.StatusCode) } shard, _ := io.ReadAll(resp.Body) - if !strings.HasPrefix(string(shard), "id\tsmiles\n") { + if !strings.HasPrefix(string(shard), "chembl_id\tcanonical_smiles\n") { t.Errorf("shard missing header: %q", shard) } } @@ -430,11 +458,12 @@ func TestUploadDatasetLimitsRows(t *testing.T) { e := newEnv(t, healthy) var buf bytes.Buffer mw := multipart.NewWriter(&buf) - _ = mw.WriteField("workload", "w") + _ = mw.WriteField("workload", "similarity-search") + _ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`) _ = mw.WriteField("chunk_rows", "2") _ = mw.WriteField("max_rows", "3") fw, _ := mw.CreateFormFile("file", "chembl.tsv") - _, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n")) + _, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\n")) _ = mw.Close() req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf) req.Header.Set("Authorization", "Bearer "+token) @@ -451,6 +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) { e := newEnv(t, healthy) zero := "00000000-0000-0000-0000-000000000000" @@ -478,10 +530,11 @@ func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) { e := newEnv(t, healthy) var buf bytes.Buffer mw := multipart.NewWriter(&buf) - _ = mw.WriteField("workload", "w") + _ = mw.WriteField("workload", "similarity-search") + _ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`) _ = mw.WriteField("chunk_rows", "not-a-number") fw, _ := mw.CreateFormFile("file", "chembl.tsv") - _, _ = io.Copy(fw, strings.NewReader("id\tsmiles\nA\tCC\n")) + _, _ = io.Copy(fw, strings.NewReader("chembl_id\tcanonical_smiles\nA\tCC\n")) _ = mw.Close() req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf) @@ -505,6 +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)) req.Header.Set("Authorization", "Bearer "+token) req.Header.Set("Content-Type", "text/csv") + if worker == "w1" { + worker = e.workerID + } req.Header.Set("X-Worker-ID", worker) req.Header.Set("X-Task-Attempt", itoa(attempt)) resp, err := http.DefaultClient.Do(req) @@ -526,6 +582,7 @@ func (e *env) uploadDataset(t *testing.T, workload string, rows int, tsv string) var buf bytes.Buffer mw := multipart.NewWriter(&buf) _ = mw.WriteField("workload", workload) + _ = mw.WriteField("parameters", `{"query_smiles":"CCO"}`) _ = mw.WriteField("chunk_rows", itoa(rows)) fw, _ := mw.CreateFormFile("file", "chembl.tsv") _, _ = io.Copy(fw, strings.NewReader(tsv)) diff --git a/coordinator/internal/transport/http/templates/job.html b/coordinator/internal/transport/http/templates/job.html index f983e2d..b985a2e 100644 --- a/coordinator/internal/transport/http/templates/job.html +++ b/coordinator/internal/transport/http/templates/job.html @@ -22,7 +22,7 @@ diff --git a/coordinator/internal/usecase/artifact.go b/coordinator/internal/usecase/artifact.go index 5cb8b8b..7a0187f 100644 --- a/coordinator/internal/usecase/artifact.go +++ b/coordinator/internal/usecase/artifact.go @@ -14,12 +14,13 @@ type UploadArtifact struct { tasks TaskRepository artifacts ArtifactRepository blobs BlobStore + tx TxManager clk Clock } func NewUploadArtifact(tasks TaskRepository, artifacts ArtifactRepository, - blobs BlobStore, clk Clock) *UploadArtifact { - return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, clk: clk} + blobs BlobStore, tx TxManager, clk Clock) *UploadArtifact { + return &UploadArtifact{tasks: tasks, artifacts: artifacts, blobs: blobs, tx: tx, clk: clk} } func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) (*domain.Artifact, error) { @@ -32,6 +33,15 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) ( if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) { return nil, domain.ErrLeaseConflict } + // A client can retry a PUT after losing the response. Return the one durable + // result for this lease attempt instead of storing duplicate artifacts. + existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt) + if err != nil { + return nil, err + } + if existing != nil { + return existing, nil + } taskID := task.ID art, err := domain.NewArtifact(task.JobID, &taskID, domain.ArtifactPartialResult, @@ -50,25 +60,41 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) ( } art.SetContent(sum, size) - // The stream may take longer than the lease. Re-check after it finishes so - // an expired worker cannot leave a durable result record behind. Completion - // performs the same ownership check under its transaction. - current, err := uc.tasks.Get(ctx, in.TaskID) + // The stream may take longer than the lease. Lock the task while re-checking + // ownership and inserting metadata: completion or another upload cannot race + // this final decision. The database unique index is a second line of defence. + var durable *domain.Artifact + err = uc.tx.WithinTx(ctx, func(ctx context.Context) error { + current, err := uc.tasks.GetForUpdate(ctx, in.TaskID) + if err != nil { + return err + } + if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) { + return domain.ErrLeaseConflict + } + existing, err := uc.artifacts.FindPartialResult(ctx, in.TaskID, in.Attempt) + if err != nil { + return err + } + if existing != nil { + durable = existing + return nil + } + if err := uc.artifacts.Insert(ctx, art); err != nil { + return err + } + durable = art + return nil + }) if err != nil { _ = uc.blobs.Delete(ctx, art.StorageKey) return nil, err } - if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) { + if durable != art { + // Another request won the race while this stream was being written. _ = uc.blobs.Delete(ctx, art.StorageKey) - return nil, domain.ErrLeaseConflict } - - // Persist the record. If that fails the blob would be an orphan, so remove it. - if err := uc.artifacts.Insert(ctx, art); err != nil { - _ = uc.blobs.Delete(ctx, art.StorageKey) - return nil, err - } - return art, nil + return durable, nil } // DownloadArtifact returns an artifact's metadata together with a reader over diff --git a/coordinator/internal/usecase/job.go b/coordinator/internal/usecase/job.go index 748f021..f6e0497 100644 --- a/coordinator/internal/usecase/job.go +++ b/coordinator/internal/usecase/job.go @@ -33,6 +33,17 @@ func NewCreateJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock // The all-or-none guarantee comes from TxManager: a half-created job would // leave chunks no worker could ever complete. func (uc *CreateJob) Execute(ctx context.Context, in CreateJobInput) (*domain.Job, error) { + if in.Workload == "similarity-graph" || in.Workload == "similarity_graph" { + // CTX-10 must plan triangular block pairs; ordinary independent input + // chunks would silently omit every cross-chunk molecular pair. + return nil, domain.ErrInvalidInput + } + if (in.Workload == "similarity-search" || in.Workload == "similarity_search") && + len(in.Chunks) > 1 && in.Parameters["query_id"] != nil { + // Resolving once against the source dataset belongs to CTX-07. Letting + // each shard resolve it would make most tasks fail or use inconsistent data. + return nil, domain.ErrInvalidInput + } chunks := make([]domain.ChunkSpec, 0, len(in.Chunks)) for _, c := range in.Chunks { chunks = append(chunks, domain.ChunkSpec(c)) @@ -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 { 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) if err != nil { return err @@ -212,3 +234,18 @@ func syncJobStatus(ctx context.Context, jobs JobRepository, tasks TaskRepository } 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 +} diff --git a/coordinator/internal/usecase/ports.go b/coordinator/internal/usecase/ports.go index 7b56211..119a0c9 100644 --- a/coordinator/internal/usecase/ports.go +++ b/coordinator/internal/usecase/ports.go @@ -60,9 +60,9 @@ type TaskRepository interface { // lease. It returns how many tasks changed. CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error) - // ExpireLeases applies the lease-expiry rule to every elapsed task and - // reports how many were affected. - ExpireLeases(ctx context.Context, now time.Time) (int64, error) + // ExpireLeases applies the lease-expiry rule to every elapsed task and returns + // the distinct jobs whose aggregate status may have changed. + ExpireLeases(ctx context.Context, now time.Time) ([]uuid.UUID, error) } // JobRepository persists jobs. @@ -89,6 +89,9 @@ type WorkerRepository interface { type ArtifactRepository interface { Insert(ctx context.Context, a *domain.Artifact) error Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, error) + // FindPartialResult returns the durable result already uploaded for one task + // attempt. A nil artifact means the attempt has not uploaded one yet. + FindPartialResult(ctx context.Context, taskID uuid.UUID, attempt int) (*domain.Artifact, error) } // BlobStore holds artifact bytes, addressed by an opaque storage key. It streams diff --git a/coordinator/internal/usecase/task.go b/coordinator/internal/usecase/task.go index 131cffe..1af5bf8 100644 --- a/coordinator/internal/usecase/task.go +++ b/coordinator/internal/usecase/task.go @@ -21,12 +21,15 @@ import ( type ClaimTask struct { tasks TaskRepository + jobs JobRepository + workers WorkerRepository + tx TxManager clock Clock leaseDuration time.Duration } -func NewClaimTask(tasks TaskRepository, clock Clock, leaseDuration time.Duration) *ClaimTask { - return &ClaimTask{tasks: tasks, clock: clock, leaseDuration: leaseDuration} +func NewClaimTask(tasks TaskRepository, jobs JobRepository, workers WorkerRepository, tx TxManager, clock Clock, leaseDuration time.Duration) *ClaimTask { + return &ClaimTask{tasks: tasks, jobs: jobs, workers: workers, tx: tx, clock: clock, leaseDuration: leaseDuration} } // Execute reclaims elapsed leases first, then hands out one task. @@ -42,27 +45,46 @@ func (uc *ClaimTask) Execute(ctx context.Context, in ClaimTaskInput) (*domain.Cl if in.WorkerID == "" { return nil, domain.ErrInvalidInput } - now := uc.clock.Now() - - if _, err := uc.tasks.ExpireLeases(ctx, now); err != nil { - return nil, err + workloads := in.Workloads + if workerID, err := uuid.Parse(in.WorkerID); err == nil { + worker, err := uc.workers.Get(ctx, workerID) + if err != nil { + return nil, err + } + // Never trust caller-supplied capabilities: registration is the durable + // worker identity and its allowlist. + workloads = worker.Capabilities } + var claimed *domain.ClaimedTask + err := uc.tx.WithinTx(ctx, func(ctx context.Context) error { + now := uc.clock.Now() + affectedJobs, err := uc.tasks.ExpireLeases(ctx, now) + if err != nil { + return err + } + if err := syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affectedJobs, now); err != nil { + return err + } - task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{ - Workloads: in.Workloads, - Owner: in.WorkerID, - Now: now, - LeaseUntil: now.Add(uc.leaseDuration), + task, err := uc.tasks.ClaimNext(ctx, ClaimFilter{ + Workloads: workloads, + Owner: in.WorkerID, + Now: now, + LeaseUntil: now.Add(uc.leaseDuration), + }) + if err != nil { + return err + } + if task != nil { + value := task.AsClaimed() + claimed = &value + } + return nil }) if err != nil { return nil, err } - if task == nil { - return nil, nil // empty queue is a normal state, not an error - } - - claimed := task.AsClaimed() - return &claimed, nil + return claimed, nil // nil means an empty queue } // --- RenewLease ---------------------------------------------------------- @@ -230,18 +252,30 @@ func (uc *FailTask) Execute(ctx context.Context, in FailTaskInput) (*domain.Task type ExpireLeases struct { tasks TaskRepository + jobs JobRepository + tx TxManager clock Clock } -func NewExpireLeases(tasks TaskRepository, clock Clock) *ExpireLeases { - return &ExpireLeases{tasks: tasks, clock: clock} +func NewExpireLeases(tasks TaskRepository, jobs JobRepository, tx TxManager, clock Clock) *ExpireLeases { + return &ExpireLeases{tasks: tasks, jobs: jobs, tx: tx, clock: clock} } -// Execute reports how many tasks were reclaimed. +// Execute reclaims elapsed tasks and persists the state of every affected job. // // The sweep is one set-based statement rather than a load-decide-save loop: // several coordinators run it concurrently, and a single atomic UPDATE makes // the duplicate work harmless — the loser simply updates 0 rows. func (uc *ExpireLeases) Execute(ctx context.Context) (int64, error) { - return uc.tasks.ExpireLeases(ctx, uc.clock.Now()) + var affected []uuid.UUID + err := uc.tx.WithinTx(ctx, func(ctx context.Context) error { + now := uc.clock.Now() + var err error + affected, err = uc.tasks.ExpireLeases(ctx, now) + if err != nil { + return err + } + return syncExpiredJobStatuses(ctx, uc.jobs, uc.tasks, affected, now) + }) + return int64(len(affected)), err } diff --git a/coordinator/internal/usecase/ui.go b/coordinator/internal/usecase/ui.go index c7283c3..dac5fe0 100644 --- a/coordinator/internal/usecase/ui.go +++ b/coordinator/internal/usecase/ui.go @@ -15,6 +15,7 @@ type UIReadRepository interface { GetJob(ctx context.Context, jobID uuid.UUID) (*domain.Job, error) ListJobs(ctx context.Context, limit int) ([]domain.Job, error) ListTasksByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Task, error) + ListTasksByJobs(ctx context.Context, jobIDs []uuid.UUID) (map[uuid.UUID][]domain.Task, error) ListWorkers(ctx context.Context, limit int) ([]domain.Worker, error) ListArtifactsByJob(ctx context.Context, jobID uuid.UUID) ([]domain.Artifact, error) } @@ -88,12 +89,16 @@ func (d *Dashboard) Overview(ctx context.Context, limit int) (DashboardView, err return DashboardView{}, err } out := DashboardView{Jobs: make([]JobCard, 0, len(jobs)), Workers: make([]WorkerCard, 0, len(workers))} + jobIDs := make([]uuid.UUID, 0, len(jobs)) for _, job := range jobs { - tasks, err := d.read.ListTasksByJob(ctx, job.ID) - if err != nil { - return DashboardView{}, err - } - out.Jobs = append(out.Jobs, jobCard(job, tasks)) + jobIDs = append(jobIDs, job.ID) + } + tasksByJob, err := d.read.ListTasksByJobs(ctx, jobIDs) + if err != nil { + return DashboardView{}, err + } + for _, job := range jobs { + out.Jobs = append(out.Jobs, jobCard(job, tasksByJob[job.ID])) } 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}) diff --git a/coordinator/internal/usecase/upload.go b/coordinator/internal/usecase/upload.go index 20fb899..ad394db 100644 --- a/coordinator/internal/usecase/upload.go +++ b/coordinator/internal/usecase/upload.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "io" + "math" "github.com/google/uuid" @@ -15,20 +16,27 @@ import ( // creates the job with one task per shard — the coordinator-side counterpart of // a client submitting pre-chunked URIs. type SubmitDataset struct { - blobs BlobStore - artifacts ArtifactRepository - jobs JobRepository - tasks TaskRepository - tx TxManager - clk Clock + blobs BlobStore + artifacts ArtifactRepository + jobs JobRepository + tasks TaskRepository + tx TxManager + clk Clock + maxAttempts int } func NewSubmitDataset(blobs BlobStore, artifacts ArtifactRepository, jobs JobRepository, - tasks TaskRepository, tx TxManager, clk Clock) *SubmitDataset { - return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk} + tasks TaskRepository, tx TxManager, clk Clock, maxAttempts int) *SubmitDataset { + return &SubmitDataset{blobs: blobs, artifacts: artifacts, jobs: jobs, tasks: tasks, tx: tx, clk: clk, maxAttempts: maxAttempts} } func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (SubmitDatasetResult, error) { + if err := validateUploadedWorkload(in.Workload, in.Parameters); err != nil { + return SubmitDatasetResult{}, err + } + if uc.maxAttempts < 1 { + return SubmitDatasetResult{}, domain.ErrInvalidInput + } now := uc.clk.Now() job, err := domain.NewUploadedJob(in.Workload, in.Parameters, now) @@ -64,7 +72,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su cleanup() return SubmitDatasetResult{}, err } - splitErr := chunk.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, fmt.Sprintf("shard-%d.tsv", index), in.ContentType, now) if err != nil { @@ -77,7 +85,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su putKeys = append(putKeys, art.StorageKey) art.SetContent(ssum, ssize) - task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, 0, now) + task, err := domain.NewShardTask(job.ID, index, in.Workload, art.ID, ssum, in.Parameters, uc.maxAttempts, now) if err != nil { return err } @@ -88,7 +96,8 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su _ = rc.Close() if splitErr != nil { cleanup() - return SubmitDatasetResult{}, splitErr + // Dataset shape is caller input, not an internal coordinator failure. + return SubmitDatasetResult{}, domain.ErrInvalidInput } // 3. Persist job + all artifacts + all tasks atomically. @@ -118,6 +127,76 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su }, nil } +// validateUploadedWorkload is deliberately narrow until CTX-07/08/10 adds a +// typed distributed-workload registry. In particular, running similarity-graph +// independently per TSV shard is scientifically wrong: cross-shard pairs would +// be absent from the apparent graph. +func validateUploadedWorkload(workload string, parameters map[string]any) error { + if workload != "similarity-search" { + return domain.ErrInvalidInput + } + allowed := map[string]struct{}{ + "query_smiles": {}, "top_k": {}, "threshold": {}, + "threshold_direction": {}, "progress_every": {}, + } + for key := range parameters { + if _, ok := allowed[key]; !ok { + return domain.ErrInvalidInput + } + } + query, ok := parameters["query_smiles"].(string) + if !ok || query == "" || len(query) > 200 { + return domain.ErrInvalidInput + } + if value, ok := parameters["top_k"]; ok && !isPositiveJSONInteger(value) { + return domain.ErrInvalidInput + } + if value, ok := parameters["progress_every"]; ok && !isNonNegativeJSONInteger(value) { + return domain.ErrInvalidInput + } + if value, ok := parameters["threshold"]; ok && !isUnitIntervalNumber(value) { + return domain.ErrInvalidInput + } + if value, ok := parameters["threshold_direction"]; ok && value != "greater" && value != "less" { + return domain.ErrInvalidInput + } + return nil +} + +func isPositiveJSONInteger(value any) bool { return isJSONInteger(value, false) } +func isNonNegativeJSONInteger(value any) bool { return isJSONInteger(value, true) } + +func isJSONInteger(value any, allowZero bool) bool { + var n int64 + switch v := value.(type) { + case int: + n = int64(v) + case int64: + n = v + case float64: + if math.Trunc(v) != v || v > math.MaxInt64 || v < math.MinInt64 { + return false + } + n = int64(v) + default: + return false + } + return n >= 0 && (allowZero || n > 0) +} + +func isUnitIntervalNumber(value any) bool { + switch v := value.(type) { + case float64: + return !math.IsNaN(v) && !math.IsInf(v, 0) && v >= 0 && v <= 1 + case int: + return v >= 0 && v <= 1 + case int64: + return v >= 0 && v <= 1 + default: + return false + } +} + // GetTaskInput resolves a task's input shard and opens it for streaming. The // caller closes the reader. type GetTaskInput struct { diff --git a/coordinator/internal/usecase/usecase_test.go b/coordinator/internal/usecase/usecase_test.go index 68b94ac..2083c5c 100644 --- a/coordinator/internal/usecase/usecase_test.go +++ b/coordinator/internal/usecase/usecase_test.go @@ -68,18 +68,18 @@ func newHarness() *harness { } tx := memstore.Tx{} h.createJob = usecase.NewCreateJob(h.jobs, h.tasks, tx, h.clk) - h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk) - h.claim = usecase.NewClaimTask(h.tasks, h.clk, lease) + h.submit = usecase.NewSubmitDataset(h.blobs, h.arts, h.jobs, h.tasks, tx, h.clk, 3) + h.claim = usecase.NewClaimTask(h.tasks, h.jobs, h.work, tx, h.clk, lease) h.renew = usecase.NewRenewLease(h.tasks, h.work, tx, h.clk, lease) h.complete = usecase.NewCompleteTask(h.tasks, h.jobs, h.arts, tx, h.clk) h.fail = usecase.NewFailTask(h.tasks, h.jobs, tx, h.clk) h.status = usecase.NewGetJobStatus(h.jobs, h.tasks) h.results = usecase.NewListResults(h.tasks) h.register = usecase.NewRegisterWorker(h.work, h.clk) - h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, h.clk) + h.uploadArt = usecase.NewUploadArtifact(h.tasks, h.arts, h.blobs, tx, h.clk) h.downloadArt = usecase.NewDownloadArtifact(h.arts, h.blobs) h.getInput = usecase.NewGetTaskInput(h.tasks, h.arts, h.blobs) - h.expire = usecase.NewExpireLeases(h.tasks, h.clk) + h.expire = usecase.NewExpireLeases(h.tasks, h.jobs, tx, h.clk) h.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk) 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) { h := newHarness() 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) taskID, attempt := h.leaseOne(t, "w1", "w") h.uploadArt = usecase.NewUploadArtifact( - h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, h.clk, + h.tasks, h.arts, expiringBlobStore{BlobStore: h.blobs, clock: h.clk}, memstore.Tx{}, h.clk, ) _, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{ @@ -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) { h := newHarness() h.seedJob(t, "w", 1) @@ -434,10 +489,10 @@ func TestDownloadArtifactRoundTrips(t *testing.T) { func TestSubmitDatasetChunksAndServesInput(t *testing.T) { h := newHarness() - tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" + tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{ - Workload: "w", RowsPerShard: 2, Filename: "chembl.tsv", + Workload: "similarity-search", Parameters: map[string]any{"query_smiles": "CCO"}, RowsPerShard: 2, Filename: "chembl.tsv", ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv), }) if err != nil { @@ -453,7 +508,7 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) { t.Errorf("job total = %d, want 3", prog.Total) } - c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}}) + c, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"similarity-search"}}) if err != nil || c == nil { t.Fatalf("claim shard: %v", err) } @@ -472,9 +527,9 @@ func TestSubmitDatasetChunksAndServesInput(t *testing.T) { func TestSubmitDatasetLimitsRowsBeforeCreatingShards(t *testing.T) { h := newHarness() - tsv := "id\tsmiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" + tsv := "chembl_id\tcanonical_smiles\nA\tCC\nB\tCCC\nC\tCCCC\nD\tCCCCC\nE\tCCCCCC\n" res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{ - Workload: "w", RowsPerShard: 2, 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), }) 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) { h := newHarness() 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) } } + +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) + } +} diff --git a/coordinator/migrations/0009_unique_partial_result_attempt.down.sql b/coordinator/migrations/0009_unique_partial_result_attempt.down.sql new file mode 100644 index 0000000..a724fb8 --- /dev/null +++ b/coordinator/migrations/0009_unique_partial_result_attempt.down.sql @@ -0,0 +1,5 @@ +BEGIN; + +DROP INDEX IF EXISTS uq_partial_result_task_attempt; + +COMMIT; diff --git a/coordinator/migrations/0009_unique_partial_result_attempt.up.sql b/coordinator/migrations/0009_unique_partial_result_attempt.up.sql new file mode 100644 index 0000000..e6ce5e0 --- /dev/null +++ b/coordinator/migrations/0009_unique_partial_result_attempt.up.sql @@ -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; diff --git a/coordinator/scripts/smoke.sh b/coordinator/scripts/smoke.sh index 2d9a1f2..ac018d6 100755 --- a/coordinator/scripts/smoke.sh +++ b/coordinator/scripts/smoke.sh @@ -51,6 +51,13 @@ check "register worker" 201 -X POST "${HOST}/workers/register" " -d '{"name":"smoke-worker","capabilities":["similarity_search"],"cpu_count":4,"memory_mb":8192}' check "register without capabilities → 400" 400 -X POST "${HOST}/workers/register" "${auth[@]}" \ -d '{"name":"bad"}' +registration=$(curl -sS "${auth[@]}" -X POST "${HOST}/workers/register" \ + -d '{"name":"smoke-worker-active","capabilities":["similarity_search","similarity-search"],"cpu_count":4}') +worker_id=$(printf '%s' "$registration" | python3 -c 'import json,sys;print(json.load(sys.stdin)["worker_id"])' 2>/dev/null) +if [[ -z "${worker_id:-}" ]]; then + echo " ✗ could not register active worker: $registration" + exit 1 +fi echo echo "job lifecycle" @@ -75,7 +82,7 @@ declare -A our_chunks task_id="" attempt="" for _ in $(seq 1 40); do - claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d '{"worker_id":"w1"}') + claim=$(curl -sS "${auth[@]}" -X POST "${HOST}/tasks/claim" -d "{\"worker_id\":\"${worker_id}\"}") [[ -z "$claim" ]] && break # 204: queue drained read -r c_job c_task c_chunk c_attempt < <(printf '%s' "$claim" | @@ -100,7 +107,7 @@ else fi check "heartbeat" 200 -X POST "${HOST}/tasks/${task_id}/heartbeat" "${auth[@]}" \ - -d "{\"worker_id\":\"w1\",\"attempt\":${attempt}}" + -d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt}}" # --- artifacts + result (uploads happen while the task is still leased) --- bearer=(-H "Authorization: Bearer ${TOKEN}") @@ -108,43 +115,47 @@ bearer=(-H "Authorization: Bearer ${TOKEN}") # upload -> prints the artifact_id upload() { curl -sS -X PUT "${HOST}/tasks/${task_id}/artifacts/$1" "${bearer[@]}" \ - -H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \ + -H 'Content-Type: text/csv' -H "X-Worker-ID: ${worker_id}" -H "X-Task-Attempt: ${attempt}" \ --data-binary $'query,match,score\nA,B,0.9\n' | python3 -c 'import json,sys;print(json.load(sys.stdin)["artifact_id"])' 2>/dev/null } check "upload artifact" 200 -X PUT "${HOST}/tasks/${task_id}/artifacts/result.csv" "${bearer[@]}" \ - -H 'Content-Type: text/csv' -H 'X-Worker-ID: w1' -H "X-Task-Attempt: ${attempt}" \ + -H 'Content-Type: text/csv' -H "X-Worker-ID: ${worker_id}" -H "X-Task-Attempt: ${attempt}" \ --data-binary $'query,match,score\nA,B,0.9\n' check "foreign worker upload → 409" 409 -X PUT "${HOST}/tasks/${task_id}/artifacts/x.csv" "${bearer[@]}" \ -H 'Content-Type: text/csv' -H 'X-Worker-ID: impostor' -H "X-Task-Attempt: ${attempt}" \ --data-binary 'x' -# Two result artifacts, uploaded now while the lease is held: one to complete -# with, a second to prove a different manifest is rejected after completion. +# A retry of a PUT returns the same durable artifact for the task attempt. art_id=$(upload primary.csv) art_id2=$(upload secondary.csv) +if [[ "$art_id" == "$art_id2" ]]; then + printf ' \033[32m✓\033[0m %-46s %s\n' "duplicate upload is idempotent" "$art_id" + pass=$((pass + 1)) +else + printf ' \033[31m✗\033[0m %-46s got %s and %s\n' "duplicate upload is idempotent" "$art_id" "$art_id2" + fail=$((fail + 1)) +fi check "download artifact" 200 "${HOST}/artifacts/${art_id}/download" "${bearer[@]}" check "foreign worker submits → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \ -d "{\"worker_id\":\"impostor\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}" check "submit result" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \ - -d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}" + -d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}" check "replay same result → idempotent" 200 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \ - -d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}" -check "different result → 409" 409 -X POST "${HOST}/tasks/${task_id}/result" "${auth[@]}" \ - -d "{\"worker_id\":\"w1\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id2}\"}}" + -d "{\"worker_id\":\"${worker_id}\",\"attempt\":${attempt},\"result\":{\"artifact_id\":\"${art_id}\"}}" check "GET /jobs/{id}" 200 "${HOST}/jobs/${job_id}" "${auth[@]}" echo echo "input validation" check "malformed uuid → 400" 400 -X POST "${HOST}/tasks/not-a-uuid/result" "${auth[@]}" \ - -d '{"worker_id":"w1","attempt":1,"result":{"artifact_id":"00000000-0000-0000-0000-000000000000"}}' + -d "{\"worker_id\":\"${worker_id}\",\"attempt\":1,\"result\":{\"artifact_id\":\"00000000-0000-0000-0000-000000000000\"}}" # Note: Go's encoding/json matches field names case-insensitively, so # "worker_ID" would be accepted as "worker_id". Only a genuinely unknown key # trips DisallowUnknownFields. check "unknown json field → 400" 400 -X POST "${HOST}/tasks/claim" "${auth[@]}" \ - -d '{"worker_id":"w1","totally_unknown":1}' + -d "{\"worker_id\":\"${worker_id}\",\"totally_unknown\":1}" check "unknown job → 404" 404 "${HOST}/jobs/00000000-0000-0000-0000-000000000000" "${auth[@]}" echo @@ -152,11 +163,11 @@ echo "dataset upload → chunking" # Upload a 5-row TSV split at 2 rows/shard → expect 3 shard tasks. The text # fields precede the file part, which the coordinator streams. up=$(curl -sS "${bearer[@]}" -X POST "${HOST}/jobs/upload" \ - -F 'workload=similarity_search' \ - -F 'parameters={"top_k":10}' \ + -F 'workload=similarity-search' \ + -F 'parameters={"query_smiles":"CCO","top_k":10}' \ -F 'chunk_rows=2' \ -F 'file=@-;filename=chembl.tsv;type=text/tab-separated-values' <<'TSV' -id smiles +chembl_id canonical_smiles A CC B CCC C CCCC @@ -180,7 +191,7 @@ fi up_input="" for _ in $(seq 1 30); do c=$(curl -sS "${bearer[@]}" -H 'Content-Type: application/json' -X POST "${HOST}/tasks/claim" \ - -d '{"worker_id":"up-w","capabilities":["similarity_search"]}') + -d "{\"worker_id\":\"${worker_id}\"}") [[ -z "$c" ]] && break cj=$(printf '%s' "$c" | python3 -c 'import json,sys;print(json.load(sys.stdin)["job_id"])' 2>/dev/null) [[ "$cj" != "$up_job" ]] && continue diff --git a/docs/api-contract.md b/docs/api-contract.md index 50abf4d..103b2d1 100644 --- a/docs/api-contract.md +++ b/docs/api-contract.md @@ -107,7 +107,8 @@ Content-Type: application/json ``` `cpu_count`/`memory_mb` are accepted for forward compatibility and not yet -persisted. `capabilities` must be non-empty (an allowlisted workload set). +persisted. `capabilities` must be non-empty. A claim uses the capabilities +stored at registration; the request cannot broaden them. ## Claim task @@ -120,7 +121,8 @@ Content-Type: application/json ``` - `204 No Content`: no compatible task. -- `200 OK`: a task is leased atomically. +- `200 OK`: a task is leased atomically. `worker_id` must be a registered UUID; + its persisted capabilities, rather than this request field, decide eligibility. ```json { diff --git a/docs/building-workers.md b/docs/building-workers.md index 4fa7bbd..08f8d6a 100644 --- a/docs/building-workers.md +++ b/docs/building-workers.md @@ -46,24 +46,25 @@ coordinator was started with. Never log it, never send it in an error body. ```http POST /workers/register -{ "name": "lab-worker-01", "capabilities": ["similarity_search"] } +{ "name": "lab-worker-01", "capabilities": ["similarity-search"] } ``` Response: `{ "worker_id": "", "heartbeat_interval_seconds": 15 }`. -- `capabilities` are the workload names you can run — the coordinator only hands - you matching tasks. +- `capabilities` are fixed at registration — the coordinator only hands you + matching tasks and a later claim cannot broaden that set. - **Keep `worker_id`**. Use it as your identity in every later call. Using the registered UUID is what lets the coordinator track your liveness (it marks workers offline after they go silent). -- Current coordinator jobs use `similarity_search` / `similarity_graph`; the - reference Python worker also accepts the public CLI spellings with hyphens. +- Current diagnostic uploads use `similarity-search` with `query_smiles`. The + reference worker accepts the legacy `similarity_search` spelling too. Do not + advertise `similarity-graph` until CTX-10 implements cross-shard pair planning. ## 2. Claim a task ```http POST /tasks/claim -{ "worker_id": "", "capabilities": ["similarity_search"] } +{ "worker_id": "", "capabilities": ["similarity-search"] } ``` - `200` → a leased task (below). @@ -74,9 +75,9 @@ POST /tasks/claim "task_id": "", "attempt": 1, "lease_expires_at": "2026-07-22T12:05:00Z", - "workload": "similarity_search", + "workload": "similarity-search", "input": { "uri": "/tasks//input", "sha256": "" }, - "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 - a working directory for downloaded inputs and generated outputs +## Run the reference worker locally + +Use one terminal per worker and a distinct work directory for each process: + +```sh +SCIMESH_COORDINATOR_URL=http://localhost:8080 \ +SCIMESH_BEARER_TOKEN=dev-token \ +SCIMESH_WORKER_NAME=worker-1 \ +scimesh-worker --work-dir "$PWD/worker-data-1" +``` + +For a bounded manual check, use one of these lifecycle modes: + +```sh +# Make exactly one claim; exit immediately when no task is available. +scimesh-worker --work-dir "$PWD/worker-data-check" --once + +# Keep polling until two tasks complete successfully, then exit. +scimesh-worker --work-dir "$PWD/worker-data-check" --max-tasks 2 +``` + +`SCIMESH_MAX_TASKS` provides the same limit through the environment. Pressing +`Ctrl+C` stops the reference worker cleanly. If it interrupts an active task, +the worker reports a sanitized retriable failure first, emits no traceback, and +exits with status `130`. + ## Generate a client from the spec Instead of hand-writing request code, generate it: diff --git a/docs/ctx-07-distributed-workload-protocol.md b/docs/ctx-07-distributed-workload-protocol.md new file mode 100644 index 0000000..c8022c2 --- /dev/null +++ b/docs/ctx-07-distributed-workload-protocol.md @@ -0,0 +1,207 @@ +# CTX-07: distributed workload protocol and planner contract + +## Status and scope + +This document is the implementation contract for CTX-07. It does not implement +a 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. diff --git a/docs/openapi.yaml b/docs/openapi.yaml index 7117a78..4e574df 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -96,7 +96,9 @@ paths: description: > multipart/form-data. The text fields (`workload`, `parameters`, `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: required: true content: @@ -388,7 +390,7 @@ components: type: array minItems: 1 items: { type: string } - example: [similarity_search, similarity_graph] + example: [similarity-search] cpu_count: type: integer description: Accepted for forward compatibility; not yet persisted. @@ -419,7 +421,7 @@ components: type: object required: [workload, input_uri, chunks] properties: - workload: { type: string, example: similarity_search } + workload: { type: string, example: similarity-search } input_uri: { type: string } parameters: { type: object, additionalProperties: true } chunks: @@ -437,11 +439,11 @@ components: type: object required: [workload, file] properties: - workload: { type: string, example: similarity_search } + workload: { type: string, enum: [similarity-search], example: similarity-search } parameters: type: string description: JSON object, sent as a string form field. - example: '{"top_k":10}' + example: '{"query_smiles":"CCO","top_k":10}' chunk_rows: type: integer description: Data rows per shard. Default 1000. @@ -486,11 +488,11 @@ components: type: object required: [worker_id] properties: - worker_id: { type: string } + worker_id: { type: string, format: uuid, description: Registered worker identity. } capabilities: type: array items: { type: string } - description: Workloads this worker can run. Empty means "any". + description: Accepted for compatibility only; registration capabilities decide eligibility. max_concurrency: type: integer description: Accepted; the coordinator leases one task per call. diff --git a/docs/worker-daemon-task.md b/docs/worker-daemon-task.md index c5e7445..e4d2f07 100644 --- a/docs/worker-daemon-task.md +++ b/docs/worker-daemon-task.md @@ -59,7 +59,7 @@ Content-Type: application/json { "name": "lab-worker-01", - "capabilities": ["similarity-search", "similarity-graph"], + "capabilities": ["similarity-search"], "cpu_count": 8, "memory_mb": 16384 } @@ -74,8 +74,8 @@ POST /tasks/claim Content-Type: application/json { - "worker_id": "worker-01", - "capabilities": ["similarity-search", "similarity-graph"], + "worker_id": "", + "capabilities": ["similarity-search"], "max_concurrency": 1 } ``` @@ -95,7 +95,7 @@ When a task is available, it returns `200 OK`: "sha256": "..." }, "parameters": { - "query_id": "CHEMBL939", + "query_smiles": "CCO", "top_k": 20 } } diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py index 13e1fb5..422c178 100644 --- a/scimesh/worker/cli.py +++ b/scimesh/worker/cli.py @@ -13,15 +13,17 @@ from .daemon import WorkerDaemon from .runners import SciMeshRunner -def main(argv: list[str] | None = None) -> int: +def build_parser() -> argparse.ArgumentParser: + """Build the worker CLI parser for command-line use and focused tests.""" parser = argparse.ArgumentParser( prog="scimesh-worker", epilog=( "Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, " "SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, " "SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, " - "SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, and " - "SCIMESH_BEARER_TOKEN. SCIMESH_WORKER_ID is a legacy/test override." + "SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, " + "SCIMESH_MAX_TASKS, and SCIMESH_BEARER_TOKEN. " + "SCIMESH_WORKER_ID is a legacy/test override." ), ) parser.add_argument("--coordinator-url") @@ -34,8 +36,31 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--request-timeout", type=float) parser.add_argument("--heartbeat-interval", type=float) parser.add_argument("--cleanup-after-seconds", type=float) + lifecycle = parser.add_mutually_exclusive_group() + lifecycle.add_argument( + "--once", + action="store_true", + help="Claim at most one task, then exit; exit immediately when the queue is empty", + ) + lifecycle.add_argument( + "--max-tasks", + type=int, + help="Process this many claimed tasks, then exit", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() args = parser.parse_args(argv) - overrides = {key: value for key, value in vars(args).items() if value is not None} + overrides = { + key: value + for key, value in vars(args).items() + if value is not None and key != "once" + } + if args.once: + overrides["max_tasks"] = 1 + overrides["exit_when_idle"] = True if "work_dir" in overrides: overrides["work_dir"] = Path(overrides["work_dir"]) try: @@ -44,13 +69,13 @@ def main(argv: list[str] | None = None) -> int: parser.error(str(error)) logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token) - WorkerDaemon( + completed_without_interruption = WorkerDaemon( config, client, HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token), SciMeshRunner(), ).run_forever() - return 0 + return 0 if completed_without_interruption else 130 if __name__ == "__main__": diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py index a0420eb..84a1157 100644 --- a/scimesh/worker/config.py +++ b/scimesh/worker/config.py @@ -36,13 +36,16 @@ class WorkerConfig: heartbeat_interval: float = 15.0 bearer_token: str | None = None cleanup_after_seconds: float | None = None + max_tasks: int | None = None + exit_when_idle: bool = False + # Distributed similarity-graph requires triangular block-pair planning and + # is deliberately not advertised until CTX-10. A normal worker must never + # make a multi-shard graph job appear scientifically complete. # The local CLI uses hyphens; the first coordinator contract used - # underscores. Advertise both stable spellings while jobs are migrated. + # underscores, so retain the search alias during migration. capabilities: tuple[str, ...] = ( "similarity-search", - "similarity-graph", "similarity_search", - "similarity_graph", ) def __post_init__(self) -> None: @@ -66,6 +69,15 @@ class WorkerConfig: _positive_number(self.heartbeat_interval, "heartbeat_interval") if self.cleanup_after_seconds is not None: _positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True) + if self.max_tasks is not None: + if ( + isinstance(self.max_tasks, bool) + or not isinstance(self.max_tasks, int) + or self.max_tasks < 1 + ): + raise ValueError("max_tasks must be positive when set") + if not isinstance(self.exit_when_idle, bool): + raise ValueError("exit_when_idle must be a boolean") if not self.capabilities: raise ValueError("capabilities cannot be empty") # Runner subprocesses use a task directory as their cwd. Keep the @@ -90,6 +102,7 @@ class WorkerConfig: cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS") cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1) memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB") + max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS") return cls( coordinator_url=url.rstrip("/"), worker_id=value("worker_id", "SCIMESH_WORKER_ID"), @@ -102,4 +115,6 @@ class WorkerConfig: heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")), bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"), cleanup_after_seconds=float(cleanup) if cleanup else None, + max_tasks=int(max_tasks) if max_tasks is not None else None, + exit_when_idle=bool(values.get("exit_when_idle", False)), ) diff --git a/scimesh/worker/daemon.py b/scimesh/worker/daemon.py index dbf035b..f8e51a7 100644 --- a/scimesh/worker/daemon.py +++ b/scimesh/worker/daemon.py @@ -4,8 +4,10 @@ from __future__ import annotations import logging from dataclasses import replace +from dataclasses import dataclass from pathlib import Path import random +import re import shutil import threading import time @@ -72,6 +74,14 @@ class LeaseHeartbeat: return seconds +@dataclass(frozen=True) +class RunOnceOutcome: + """Whether a claim was made and whether that claimed task completed.""" + + claimed: bool + completed: bool + + class WorkerDaemon: def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None: self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner @@ -79,32 +89,68 @@ class WorkerDaemon: self._registered = False self.log = logging.getLogger("scimesh.worker") - def run_forever(self) -> None: + def run_forever(self) -> bool: + """Run until stopped; return false only when interrupted by the operator.""" failures = 0 - while True: - try: - if not self._registered: - self._register_worker() - self._cleanup_expired_directories() - claimed = self.run_once() - failures = 0 - if not claimed: - self._sleep(self.config.poll_interval) - except CoordinatorTransientError as error: - failures += 1 - self._log("failed", error_type=type(error).__name__) - self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0)) + completed_tasks = 0 + self._log( + "started", + max_tasks=self.config.max_tasks, + exit_when_idle=self.config.exit_when_idle, + ) + try: + while True: + try: + if not self._registered: + self._register_worker() + self._cleanup_expired_directories() + outcome = self.run_once() + failures = 0 + if outcome.claimed: + if outcome.completed: + completed_tasks += 1 + if self.config.exit_when_idle: + self._log( + "stopped", + reason="one_claim_processed", + completed_tasks=completed_tasks, + ) + return True + if ( + outcome.completed + and self.config.max_tasks is not None + and completed_tasks >= self.config.max_tasks + ): + self._log( + "stopped", + reason="max_tasks_reached", + completed_tasks=completed_tasks, + ) + return True + elif self.config.exit_when_idle: + self._log("stopped", reason="queue_empty", completed_tasks=completed_tasks) + return True + else: + self._sleep(self.config.poll_interval) + except CoordinatorTransientError as error: + failures += 1 + self._log("failed", error_type=type(error).__name__) + self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0)) + except KeyboardInterrupt: + self._log("stopped", reason="interrupted", completed_tasks=completed_tasks) + return False - def run_once(self) -> bool: + def run_once(self) -> RunOnceOutcome: worker_id = self._worker_id() - self._log("claiming") + self._log("claiming", log_level=logging.DEBUG) task = self.coordinator.claim(worker_id, self.config.capabilities) if task is None: - self._log("idle") - return False + self._log("idle", log_level=logging.DEBUG) + return RunOnceOutcome(claimed=False, completed=False) started = time.monotonic() task_dir = self.config.work_dir / task.task_id / str(task.attempt) heartbeat = LeaseHeartbeat(task, self.coordinator, self.config) + completed = False try: task_dir.mkdir(parents=True, exist_ok=False) heartbeat.start() @@ -135,7 +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: self._log("lease_lost", task, error_type=type(error).__name__) except Exception as error: @@ -143,10 +197,10 @@ class WorkerDaemon: self._report_failure(task, error) finally: heartbeat.stop() - return True + return RunOnceOutcome(claimed=True, completed=completed) def _report_failure(self, task: ClaimedTask, error: Exception) -> None: - message = str(error).replace(str(self.config.work_dir), "")[:300] + message = self._sanitize_error_message(error) try: self.coordinator.fail(task, {"worker_id": self._worker_id(), "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message}) except CoordinatorTransientError: @@ -154,6 +208,16 @@ class WorkerDaemon: except Exception: 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), "") + # 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"(?", message) + message = re.sub(r"(?", message) + return message[:300] + def _register_worker(self) -> None: registered = self.coordinator.register( self.config.worker_name, @@ -180,9 +244,16 @@ class WorkerDaemon: """Keep completion payload exact: coordinator owns all artifact metadata.""" return {"artifact_id": uploaded.artifact_id} - def _log(self, state: str, task: ClaimedTask | None = None, **extra: object) -> None: + def _log( + self, + state: str, + task: ClaimedTask | None = None, + *, + log_level: int = logging.INFO, + **extra: object, + ) -> None: fields = {"worker_id": self.config.worker_id, "task_id": task.task_id if task else None, "attempt": task.attempt if task else None, "state": state, **extra} - self.log.info("worker_event %s", fields) + self.log.log(log_level, "worker_event %s", fields) def _cleanup_expired_directories(self) -> None: """Remove only old task attempt directories when retention was configured.""" diff --git a/scripts/two-worker-smoke.sh b/scripts/two-worker-smoke.sh new file mode 100755 index 0000000..c8d96a5 --- /dev/null +++ b/scripts/two-worker-smoke.sh @@ -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" diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py index 36f6cf6..7bf6f65 100644 --- a/tests/test_worker_daemon.py +++ b/tests/test_worker_daemon.py @@ -1,7 +1,9 @@ from __future__ import annotations import hashlib +import logging from pathlib import Path +import subprocess import time from datetime import datetime, timedelta, timezone from urllib.request import Request @@ -9,8 +11,10 @@ from urllib.request import Request import pytest from scimesh.worker.config import WorkerConfig +from scimesh.worker import cli as worker_cli +from scimesh.worker.cli import build_parser from scimesh.worker.coordinator import CoordinatorTransientError -from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon +from scimesh.worker.daemon import LeaseHeartbeat, RunOnceOutcome, WorkerDaemon from scimesh.worker.models import ( ClaimedTask, InputArtifact, @@ -92,7 +96,7 @@ def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes): def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None: content = b"input fixture" worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content) - assert worker.run_once() is True + assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True) assert runner.calls == 1 assert len(artifacts.uploaded) == 1 assert coordinator.heartbeats == [("task-1", 1, "worker-1")] @@ -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: worker, _, _, runner, config = daemon(tmp_path, None, b"") - assert worker.run_once() is False + assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False) assert runner.calls == 0 assert not config.work_dir.exists() +def test_once_worker_exits_after_an_empty_claim(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + caplog.set_level(logging.INFO, logger="scimesh.worker") + worker, _, _, runner, _ = daemon(tmp_path, None, b"") + worker.config = WorkerConfig(**{**worker.config.__dict__, "exit_when_idle": True, "max_tasks": 1}) + assert worker.run_forever() is True + assert runner.calls == 0 + assert "queue_empty" in caplog.text + + +def test_worker_stops_after_the_configured_number_of_claims(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + caplog.set_level(logging.INFO, logger="scimesh.worker") + content = b"input fixture" + worker, _, _, runner, _ = daemon(tmp_path, make_task(content), content) + worker.config = WorkerConfig(**{**worker.config.__dict__, "max_tasks": 1}) + assert worker.run_forever() is True + assert runner.calls == 1 + assert "max_tasks_reached" in caplog.text + + +def test_keyboard_interrupt_stops_worker_without_propagating(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + caplog.set_level(logging.INFO, logger="scimesh.worker") + class InterruptingCoordinator(FakeCoordinator): + def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + raise KeyboardInterrupt + + worker, _, _, _, _ = daemon(tmp_path, None, b"") + worker.coordinator = InterruptingCoordinator(None) + assert worker.run_forever() is False + assert "interrupted" in caplog.text + + +def test_interrupting_an_active_task_reports_a_sanitized_failure(tmp_path: Path) -> None: + content = b"input fixture" + worker, coordinator, _, _, _ = daemon(tmp_path, make_task(content), content) + + class InterruptingRunner(FakeRunner): + def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: + raise KeyboardInterrupt + + worker.runner = InterruptingRunner() + with pytest.raises(KeyboardInterrupt): + worker.run_once() + assert coordinator.failures == [ + { + "worker_id": "worker-1", + "attempt": 1, + "error_code": "InterruptedError", + "error_message": "worker interrupted by operator", + } + ] + + +def test_max_tasks_counts_successes_not_failed_claims(tmp_path: Path) -> None: + successful_content = b"successful input" + + class SequencedCoordinator(FakeCoordinator): + def __init__(self) -> None: + super().__init__(None) + self.tasks = [ + make_task(b"bad input", "wrong-checksum"), + ClaimedTask( + "task-2", + 1, + (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(), + "similarity-search", + InputArtifact( + "https://example.test/input", + hashlib.sha256(successful_content).hexdigest(), + ), + {"query_id": "CHEMBL1"}, + ), + ] + + def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + return self.tasks.pop(0) if self.tasks else None + + coordinator = SequencedCoordinator() + artifacts, runner = FakeArtifacts(successful_content), FakeRunner() + config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work", max_tasks=1) + worker = WorkerDaemon(config, coordinator, artifacts, runner) + assert worker.run_forever() is True + assert len(coordinator.failures) == 1 + assert len(coordinator.submissions) == 1 + assert runner.calls == 1 + + +def test_worker_cli_lifecycle_options_are_explicit_and_exclusive() -> None: + parser = build_parser() + assert parser.parse_args(["--once"]).once is True + assert parser.parse_args(["--max-tasks", "2"]).max_tasks == 2 + with pytest.raises(SystemExit): + parser.parse_args(["--once", "--max-tasks", "2"]) + + +def test_worker_cli_uses_a_nonzero_exit_code_for_interruption( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + class InterruptedDaemon: + def __init__(self, *_: object) -> None: + pass + + def run_forever(self) -> bool: + return False + + monkeypatch.setattr(worker_cli, "WorkerDaemon", InterruptedDaemon) + assert worker_cli.main( + ["--coordinator-url", "https://example.test", "--work-dir", str(tmp_path)] + ) == 130 + + +@pytest.mark.parametrize("value", [0, -1, True]) +def test_max_tasks_must_be_positive(value: object, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="max_tasks"): + WorkerConfig("https://example.test", None, tmp_path, max_tasks=value) # type: ignore[arg-type] + + def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None: worker, coordinator, _, runner, _ = daemon(tmp_path, make_task(b"actual", "not-the-hash"), b"actual") - assert worker.run_once() is True + assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False) assert runner.calls == 0 assert coordinator.failures[0]["error_code"] == "ValueError" assert 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 "" in message + + def test_directory_creation_failure_is_reported(tmp_path: Path) -> None: content = b"input fixture" worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content) (config.work_dir / "task-1" / "1").mkdir(parents=True) - assert worker.run_once() is True + assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False) assert coordinator.failures[0]["error_code"] == "FileExistsError" @@ -312,6 +445,7 @@ def test_environment_overrides_allow_cli_only_configuration(monkeypatch: pytest. assert config.worker_id is None assert "similarity-search" in config.capabilities assert "similarity_search" in config.capabilities + assert "similarity-graph" not in config.capabilities def test_relative_work_dir_is_normalized_for_runner_subprocesses(