From 983c5843ec8b1158e720bd6cf5eb92c42ae29a41 Mon Sep 17 00:00:00 2001 From: Emil Date: Thu, 23 Jul 2026 21:33:37 +0300 Subject: [PATCH 1/2] Fix coordinator worker integration --- STATUS.md | 40 ++--- coordinator/Makefile | 22 ++- coordinator/README.md | 11 ++ coordinator/internal/domain/task.go | 19 ++- coordinator/internal/domain/task_test.go | 29 +++- .../storage/postgres/integration_test.go | 2 +- coordinator/internal/transport/http/errors.go | 9 +- .../internal/transport/http/handlers.go | 44 +++++- .../internal/transport/http/server_test.go | 31 ++++ coordinator/internal/usecase/artifact.go | 2 +- coordinator/internal/usecase/task.go | 3 +- docs/building-workers.md | 5 +- scimesh/worker/artifacts.py | 73 ++++----- scimesh/worker/cli.py | 28 +++- scimesh/worker/config.py | 95 ++++++++--- scimesh/worker/coordinator.py | 45 +++++- scimesh/worker/daemon.py | 82 +++++++--- scimesh/worker/models.py | 114 +++++++++++++- scimesh/worker/runners.py | 10 +- scimesh/worker/transport.py | 51 ++++++ scimesh/workloads/similarity_graph.py | 14 +- scimesh/workloads/similarity_search.py | 1 + tests/test_similarity_graph.py | 17 ++ tests/test_similarity_search.py | 12 +- tests/test_worker_daemon.py | 147 +++++++++++++++++- 25 files changed, 748 insertions(+), 158 deletions(-) create mode 100644 scimesh/worker/transport.py diff --git a/STATUS.md b/STATUS.md index 7682527..eb33b5b 100644 --- a/STATUS.md +++ b/STATUS.md @@ -1,7 +1,7 @@ # SciMesh Status **Updated:** 2026-07-23 -**Branch baseline:** `planning` at `13f9a0b` +**Branch baseline:** `main` at `b4a89dd` (coordinator merge) ## Current state @@ -15,39 +15,43 @@ the reference behaviour for future distributed execution: - Python Worker skeleton: claim, heartbeat, input checksum validation, artifact upload, completion and failure reporting. -The Go coordinator, PostgreSQL schema, coordinator artifact storage, planner, -reducer, and end-to-end distributed execution are **not implemented yet**. +The Go coordinator and its PostgreSQL-backed task lifecycle are implemented: +registration, atomic claiming, lease renewal, artifact storage, dataset +chunking, result/failure reporting, and job progress. The Python worker now +uses the live coordinator contract; its HTTP path was exercised against a real +Docker PostgreSQL stack on 2026-07-23. ## Milestone tracker | CTX | Status | Notes | | --- | --- | --- | -| CTX-00 API and error contract | Ready to implement | `docs/api-contract.md` created; needs owner review/freeze. | -| CTX-01 Go coordinator bootstrap | Not started | Depends on CTX-00. | -| CTX-02 PostgreSQL migrations | Not started | Depends on CTX-00 and CTX-01. | -| CTX-03 Transactional queue | Not started | Depends on CTX-02. | -| CTX-04 Worker registry and HTTP API | Not started | Depends on CTX-03. | -| CTX-05 Artifact storage | Not started | Depends on CTX-02 and CTX-04. | -| CTX-06 Python Worker live-contract alignment | Partially prepared | Worker skeleton exists; needs real Go contract tests. | +| CTX-00 API and error contract | Implemented | Contract, OpenAPI, and request examples are in `docs/`. | +| CTX-01 Go coordinator bootstrap | Implemented | Go service and Docker runtime in `coordinator/`. | +| CTX-02 PostgreSQL migrations | Implemented | Applied by the Compose migration service. | +| CTX-03 Transactional queue | Implemented | Real-PostgreSQL integration tests cover atomic claims and concurrency. | +| CTX-04 Worker registry and HTTP API | Implemented | Registration, claim, heartbeat, result, failure, and status endpoints. | +| CTX-05 Artifact storage | Implemented | Coordinator-owned inputs/results, checksum verification, and upload flow. | +| CTX-06 Python Worker live-contract alignment | Implemented | Worker completed a real uploaded shard via HTTP on 2026-07-23. | | CTX-07 Distributed workload protocol | Not started | Depends on artifact and Worker contracts. | | CTX-08 Distributed similarity-search | Not started | Local reference exists. | | CTX-09 Reducer and final-result API | Not started | Depends on CTX-07 and CTX-08. | | CTX-10 Distributed similarity-graph | Not started | Local reference exists. | | CTX-11 Dashboard/operator view | Not started | Deferred until API and reducer work. | -| CTX-12 Reliability, security, CI | Not started | Final milestone. | +| CTX-12 Reliability, security, CI | In progress | Unit, race, PostgreSQL integration, and smoke checks exist; CI hardening remains. | ## Next recommended assignment -Assign **CTX-00** to the coordinator role in `.agents/coordinator.md`: review -and freeze `docs/api-contract.md` against `PLAN.md`. Do not begin coordinator -or Worker API implementation until the contract owner accepts it. +Assign **CTX-07** to the workload role: define distributed job planning and +reduction boundaries before implementing distributed search or graph execution. ## Known constraints -- Distributed execution is not available; use the local `scimesh` CLI. -- No Go module, PostgreSQL migrations, runtime configuration, or integration - environment exists yet. -- Local worker unit tests do not prove interoperability with a live coordinator. +- Planner/reducer semantics are not implemented; use the local `scimesh` CLI + for complete workload results. +- The worker/coordinator flow currently accepts both underscore API workload + names and hyphenated CLI names while the contract is consolidated. +- A real-stack worker test uses a small `query_smiles` shard. Resolving a + `query_id` once and sharing it across shards belongs to CTX-07. ## Update rule diff --git a/coordinator/Makefile b/coordinator/Makefile index bdc1313..564b2cf 100644 --- a/coordinator/Makefile +++ b/coordinator/Makefile @@ -1,5 +1,15 @@ .PHONY: build run test test-integration vet lint tidy check migrate-up migrate-down up down down-clean logs ps rebuild psql smoke +# `check` deliberately uses its own Compose project and host ports. This keeps +# it from connecting to or replacing a developer's local PostgreSQL instance. +CHECK_PROJECT ?= scimesh-check +CHECK_POSTGRES_PORT ?= 55432 +CHECK_COORDINATOR_PORT ?= 18080 +CHECK_HOST ?= http://localhost:$(CHECK_COORDINATOR_PORT) +CHECK_TOKEN ?= dev-token +CHECK_DATABASE_URL ?= postgres://scimesh:scimesh@localhost:$(CHECK_POSTGRES_PORT)/scimesh?sslmode=disable +CHECK_COMPOSE = POSTGRES_PORT=$(CHECK_POSTGRES_PORT) COORDINATOR_PORT=$(CHECK_COORDINATOR_PORT) docker compose -p $(CHECK_PROJECT) + # --- build / run --------------------------------------------------------- build: go build ./... @@ -23,12 +33,16 @@ vet: # Needs Docker. Hand this to a reviewer. check: vet lint go test -race ./... - docker compose up -d --build + $(CHECK_COMPOSE) up -d --build @echo "waiting for the coordinator to be ready..." - @sleep 6 - TEST_DATABASE_URL='postgres://scimesh:scimesh@localhost:5432/scimesh?sslmode=disable' \ + @attempt=0; until curl -fsS "$(CHECK_HOST)/health" >/dev/null; do \ + attempt=$$((attempt + 1)); \ + if [ $$attempt -ge 30 ]; then $(CHECK_COMPOSE) logs coordinator; exit 1; fi; \ + sleep 1; \ + done + TEST_DATABASE_URL="$(CHECK_DATABASE_URL)" \ go test -tags=integration ./internal/storage/postgres/ -v - ./scripts/smoke.sh + HOST="$(CHECK_HOST)" TOKEN="$(CHECK_TOKEN)" ./scripts/smoke.sh @echo "\nall checks passed ✓" # Runs golangci-lint without installing it system-wide. Install it for speed: diff --git a/coordinator/README.md b/coordinator/README.md index c190806..b79b671 100644 --- a/coordinator/README.md +++ b/coordinator/README.md @@ -169,3 +169,14 @@ make test-integration TEST_DATABASE_URL='postgres://scimesh:scimesh@localhost:54 CI (`.github/workflows/coordinator.yml`) runs vet, gofmt, race tests, lint, and the integration suite against a Postgres service on every push and PR. + +For the complete local verification, including an isolated Docker PostgreSQL +and the HTTP smoke flow, run: + +```sh +make check +``` + +It uses Compose project `scimesh-check` and ports `55432`/`18080` by default, +so it does not connect to a PostgreSQL already running on `5432`. Override +`CHECK_POSTGRES_PORT`, `CHECK_COORDINATOR_PORT`, or `CHECK_PROJECT` if needed. diff --git a/coordinator/internal/domain/task.go b/coordinator/internal/domain/task.go index 4fadca2..306fed5 100644 --- a/coordinator/internal/domain/task.go +++ b/coordinator/internal/domain/task.go @@ -117,8 +117,10 @@ const DefaultMaxAttempts = 3 func (t *Task) CanRetry() bool { return t.Attempt < t.MaxAttempts } // IsLeaseHeldBy reports whether worker currently holds this task at attempt. -func (t *Task) IsLeaseHeldBy(worker string, attempt int) bool { - return t.LeaseOwner != nil && *t.LeaseOwner == worker && t.Attempt == attempt +func (t *Task) IsLeaseHeldBy(worker string, attempt int, now time.Time) bool { + return t.LeaseOwner != nil && t.LeaseExpiresAt != nil && now.Before(*t.LeaseExpiresAt) && + *t.LeaseOwner == worker && t.Attempt == attempt && + (t.Status == TaskLeased || t.Status == TaskRunning) } // AsClaimed projects the task into the trimmed view handed to a worker: @@ -146,7 +148,7 @@ func (t *Task) AsClaimed() ClaimedTask { // verifyLease is the guard every worker-driven transition shares: the caller // must own the lease and reference the attempt it was granted. -func (t *Task) verifyLease(worker string, attempt int) error { +func (t *Task) verifyLease(worker string, attempt int, now time.Time) error { // A task is worker-owned while leased or running: the first heartbeat moves // it from leased to running, but ownership rules are identical for both. if t.Status != TaskLeased && t.Status != TaskRunning { @@ -158,13 +160,16 @@ func (t *Task) verifyLease(worker string, attempt int) error { if t.Attempt != attempt { return ErrStaleAttempt } + if t.LeaseExpiresAt == nil || !now.Before(*t.LeaseExpiresAt) { + return ErrLeaseConflict + } return nil } // RenewLease extends the lease of the worker that holds it. The first heartbeat // also acknowledges start, moving the task from leased to running. -func (t *Task) RenewLease(worker string, attempt int, until time.Time) error { - if err := t.verifyLease(worker, attempt); err != nil { +func (t *Task) RenewLease(worker string, attempt int, now, until time.Time) error { + if err := t.verifyLease(worker, attempt, now); err != nil { return err } t.LeaseExpiresAt = &until @@ -195,7 +200,7 @@ func (t *Task) CompleteWith(resultArtifactID uuid.UUID, metrics map[string]any, return ErrResultConflict } - if err := t.verifyLease(worker, attempt); err != nil { + if err := t.verifyLease(worker, attempt, now); err != nil { return err } @@ -214,7 +219,7 @@ func (t *Task) CompleteWith(resultArtifactID uuid.UUID, metrics map[string]any, // Fail records a worker-reported failure. A retryable failure with attempts // left returns the task to the queue; otherwise it terminates as failed. func (t *Task) Fail(worker string, attempt int, code, message string, retryable bool, now time.Time) error { - if err := t.verifyLease(worker, attempt); err != nil { + if err := t.verifyLease(worker, attempt, now); err != nil { return err } t.ErrorCode = &code diff --git a/coordinator/internal/domain/task_test.go b/coordinator/internal/domain/task_test.go index a8fa352..9c2ff98 100644 --- a/coordinator/internal/domain/task_test.go +++ b/coordinator/internal/domain/task_test.go @@ -172,14 +172,14 @@ func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) { task := leasedTask(1, 3) until := testLater.Add(time.Hour) - if err := task.RenewLease(testWorker, 1, until); err != nil { + if err := task.RenewLease(testWorker, 1, testNow, until); err != nil { t.Fatal(err) } if task.Status != TaskRunning { t.Errorf("status = %q, want running after first heartbeat", task.Status) } // A second heartbeat keeps it running. - if err := task.RenewLease(testWorker, 1, until); err != nil { + if err := task.RenewLease(testWorker, 1, testNow, until); err != nil { t.Fatal(err) } if task.Status != TaskRunning { @@ -190,15 +190,15 @@ func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) { func TestRunningTaskCanBeCompletedAndExpired(t *testing.T) { // Complete works from running. task := leasedTask(1, 3) - _ = task.RenewLease(testWorker, 1, testLater) // -> running + _ = task.RenewLease(testWorker, 1, testNow, testLater) // -> running if err := task.CompleteWith(testResult, nil, testWorker, 1, testNow); err != nil { t.Errorf("complete from running: %v", err) } // Expire reclaims a running task too. task2 := leasedTask(1, 3) - _ = task2.RenewLease(testWorker, 1, testLater) // -> running - task2.ExpireLease(testNow) + _ = task2.RenewLease(testWorker, 1, testNow, testLater) // -> running + task2.ExpireLease(testLater) if task2.Status != TaskPending { t.Errorf("status = %q, want pending after a running lease expires", task2.Status) } @@ -208,14 +208,29 @@ func TestRenewLeaseExtendsOnlyForHolder(t *testing.T) { task := leasedTask(1, 3) until := testLater.Add(time.Hour) - if err := task.RenewLease(testWorker, 1, until); err != nil { + if err := task.RenewLease(testWorker, 1, testNow, until); err != nil { t.Fatalf("unexpected error: %v", err) } if !task.LeaseExpiresAt.Equal(until) { t.Error("lease must be extended") } - if err := task.RenewLease("worker-2", 1, until); !errors.Is(err, ErrLeaseConflict) { + if err := task.RenewLease("worker-2", 1, testNow, until); !errors.Is(err, ErrLeaseConflict) { t.Errorf("err = %v, want ErrLeaseConflict", err) } } + +func TestExpiredLeaseRejectsRenewalCompletionAndFailure(t *testing.T) { + task := leasedTask(1, 3) + expired := testLater.Add(time.Nanosecond) + + if err := task.RenewLease(testWorker, 1, expired, expired.Add(time.Minute)); !errors.Is(err, ErrLeaseConflict) { + t.Errorf("renew expired lease: err = %v, want ErrLeaseConflict", err) + } + if err := task.CompleteWith(testResult, nil, testWorker, 1, expired); !errors.Is(err, ErrLeaseConflict) { + t.Errorf("complete expired lease: err = %v, want ErrLeaseConflict", err) + } + if err := task.Fail(testWorker, 1, "timeout", "expired", true, expired); !errors.Is(err, ErrLeaseConflict) { + t.Errorf("fail expired lease: err = %v, want ErrLeaseConflict", err) + } +} diff --git a/coordinator/internal/storage/postgres/integration_test.go b/coordinator/internal/storage/postgres/integration_test.go index fc917d6..05a7e90 100644 --- a/coordinator/internal/storage/postgres/integration_test.go +++ b/coordinator/internal/storage/postgres/integration_test.go @@ -223,7 +223,7 @@ func TestUpdateRejectsStaleVersion(t *testing.T) { if err != nil { return err } - if err := fresh.RenewLease("worker-1", fresh.Attempt, now.Add(2*time.Minute)); err != nil { + if err := fresh.RenewLease("worker-1", fresh.Attempt, now, now.Add(2*time.Minute)); err != nil { return err } return repo.Update(ctx, fresh) diff --git a/coordinator/internal/transport/http/errors.go b/coordinator/internal/transport/http/errors.go index f808ad7..358704e 100644 --- a/coordinator/internal/transport/http/errors.go +++ b/coordinator/internal/transport/http/errors.go @@ -3,6 +3,7 @@ package http import ( "encoding/json" "errors" + "io" "net/http" "github.com/emil28092005/SciMesh/coordinator/internal/domain" @@ -24,7 +25,13 @@ func decodeJSON(r *http.Request, dst any) error { // Reject unknown fields: silently ignoring a misspelled "worker_ID" would // surface later as a baffling validation failure. dec.DisallowUnknownFields() - return dec.Decode(dst) + if err := dec.Decode(dst); err != nil { + return err + } + if err := dec.Decode(&struct{}{}); !errors.Is(err, io.EOF) { + return errors.New("request body must contain exactly one JSON value") + } + return nil } // writeError translates domain errors into status codes. This mapping is the diff --git a/coordinator/internal/transport/http/handlers.go b/coordinator/internal/transport/http/handlers.go index cdd9385..df61461 100644 --- a/coordinator/internal/transport/http/handlers.go +++ b/coordinator/internal/transport/http/handlers.go @@ -194,11 +194,14 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) { } var ( - workload string - params map[string]any - rows = defaultChunkRows - result usecase.SubmitDatasetResult - gotDataset bool + workload string + params map[string]any + rows = defaultChunkRows + result usecase.SubmitDatasetResult + gotDataset bool + gotWorkload bool + gotParams bool + gotRows bool ) for { @@ -213,9 +216,18 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) { switch part.FormName() { case "workload": + if gotDataset || gotWorkload { + s.writeError(w, r, domain.ErrInvalidInput) + return + } b, _ := io.ReadAll(io.LimitReader(part, 1<<10)) workload = strings.TrimSpace(string(b)) + gotWorkload = true case "parameters": + if gotDataset || gotParams { + s.writeError(w, r, domain.ErrInvalidInput) + return + } b, _ := io.ReadAll(io.LimitReader(part, 1<<16)) if len(b) > 0 { if err := json.Unmarshal(b, ¶ms); err != nil { @@ -223,12 +235,25 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) { return } } + gotParams = true case "chunk_rows": - b, _ := io.ReadAll(io.LimitReader(part, 32)) - if n, err := strconv.Atoi(strings.TrimSpace(string(b))); err == nil { - rows = n + if gotDataset || gotRows { + s.writeError(w, r, domain.ErrInvalidInput) + return } + b, _ := io.ReadAll(io.LimitReader(part, 32)) + n, err := strconv.Atoi(strings.TrimSpace(string(b))) + if err != nil || n < 1 { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + rows = n + gotRows = true case "file", "dataset": + if gotDataset || workload == "" { + s.writeError(w, r, domain.ErrInvalidInput) + return + } filename := part.FileName() if filename == "" { filename = "dataset" @@ -246,6 +271,9 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) { return } gotDataset = true + default: + s.writeError(w, r, domain.ErrInvalidInput) + return } _ = part.Close() } diff --git a/coordinator/internal/transport/http/server_test.go b/coordinator/internal/transport/http/server_test.go index 0f00024..9a0c0f0 100644 --- a/coordinator/internal/transport/http/server_test.go +++ b/coordinator/internal/transport/http/server_test.go @@ -259,6 +259,37 @@ func TestErrorMappings(t *testing.T) { } } +func TestJSONRejectsTrailingValue(t *testing.T) { + e := newEnv(t, healthy) + if code, _ := e.do(t, "POST", "/workers/register", + `{"name":"lab","capabilities":["w"]} {}`); code != http.StatusBadRequest { + t.Errorf("status = %d, want 400", code) + } +} + +func TestUploadDatasetRejectsAmbiguousMultipartInput(t *testing.T) { + e := newEnv(t, healthy) + var buf bytes.Buffer + mw := multipart.NewWriter(&buf) + _ = mw.WriteField("workload", "w") + _ = mw.WriteField("chunk_rows", "not-a-number") + fw, _ := mw.CreateFormFile("file", "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("status = %d, want 400", resp.StatusCode) + } +} + // --- helpers ------------------------------------------------------------- func (e *env) putArtifact(t *testing.T, taskID, worker string, attempt int, data string) string { diff --git a/coordinator/internal/usecase/artifact.go b/coordinator/internal/usecase/artifact.go index a515d09..dcbd64d 100644 --- a/coordinator/internal/usecase/artifact.go +++ b/coordinator/internal/usecase/artifact.go @@ -29,7 +29,7 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) ( } // Only the worker holding the current lease at this attempt may upload the // task's output — the coordinator never trusts an ownership claim on faith. - if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt) { + if !task.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) { return nil, domain.ErrLeaseConflict } diff --git a/coordinator/internal/usecase/task.go b/coordinator/internal/usecase/task.go index 0c2e4f0..7d01b25 100644 --- a/coordinator/internal/usecase/task.go +++ b/coordinator/internal/usecase/task.go @@ -91,7 +91,8 @@ func (uc *RenewLease) Execute(ctx context.Context, in RenewLeaseInput) (*domain. if err != nil { return err } - if err := task.RenewLease(in.WorkerID, in.Attempt, uc.clock.Now().Add(uc.leaseDuration)); err != nil { + now := uc.clock.Now() + if err := task.RenewLease(in.WorkerID, in.Attempt, now, now.Add(uc.leaseDuration)); err != nil { return err } if err := uc.tasks.Update(ctx, task); err != nil { diff --git a/docs/building-workers.md b/docs/building-workers.md index f55d034..4fa7bbd 100644 --- a/docs/building-workers.md +++ b/docs/building-workers.md @@ -56,6 +56,8 @@ Response: `{ "worker_id": "", "heartbeat_interval_seconds": 15 }`. - **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. ## 2. Claim a task @@ -185,7 +187,8 @@ next claim. Per the worker contract, at minimum: - `SCIMESH_COORDINATOR_URL` (e.g. `http://coordinator:8080`) -- `SCIMESH_WORKER_ID` (or derive from hostname) +- worker name (the coordinator returns its `worker_id` at registration; + `SCIMESH_WORKER_ID` is only a legacy/test override) - the bearer token - poll interval and request timeout - a working directory for downloaded inputs and generated outputs diff --git a/scimesh/worker/artifacts.py b/scimesh/worker/artifacts.py index b075c42..a16ba43 100644 --- a/scimesh/worker/artifacts.py +++ b/scimesh/worker/artifacts.py @@ -7,38 +7,23 @@ import http.client import json from pathlib import Path from typing import Protocol -from urllib.parse import quote, urlsplit -from urllib.request import HTTPRedirectHandler, Request, build_opener +from urllib.parse import quote, urljoin, urlsplit +from urllib.request import Request, build_opener -from .models import ClaimedTask, ProducedArtifact +from .coordinator import CoordinatorConflictError +from .models import ClaimedTask, ProducedArtifact, UploadedArtifact +from .transport import SameOriginAuthRedirectHandler, origin - -def _origin(uri: str) -> tuple[str, str, int | None]: - parsed = urlsplit(uri) - scheme = parsed.scheme.lower() - default_port = {"http": 80, "https": 443}.get(scheme) - return scheme, (parsed.hostname or "").lower(), parsed.port or default_port - - -class _SameOriginAuthRedirectHandler(HTTPRedirectHandler): - """Do not forward the coordinator token when a download changes origin.""" - - def __init__(self, coordinator_origin: tuple[str, str, int | None]) -> None: - super().__init__() - self.coordinator_origin = coordinator_origin - - def redirect_request(self, req: Request, fp: object, code: int, msg: str, headers: object, newurl: str) -> Request | None: - redirected = super().redirect_request(req, fp, code, msg, headers, newurl) - if redirected and _origin(newurl) != self.coordinator_origin: - redirected.remove_header("Authorization") - return redirected +# Compatibility aliases for focused transport tests. +_SameOriginAuthRedirectHandler = SameOriginAuthRedirectHandler +_origin = origin class ArtifactClient(Protocol): def download(self, uri: str, destination: Path) -> None: ... def upload( self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact - ) -> str: ... + ) -> UploadedArtifact: ... class HttpArtifactClient: @@ -48,18 +33,21 @@ class HttpArtifactClient: self.coordinator_url = coordinator_url.rstrip("/") self.timeout = timeout self.bearer_token = bearer_token - self.coordinator_origin = _origin(coordinator_url) - self._opener = build_opener(_SameOriginAuthRedirectHandler(self.coordinator_origin)) + self.coordinator_origin = origin(coordinator_url) + self._opener = build_opener(SameOriginAuthRedirectHandler(self.coordinator_origin)) def download(self, uri: str, destination: Path) -> None: destination.parent.mkdir(parents=True, exist_ok=True) - request = Request(uri, headers=self._auth_headers_for(uri)) + resolved_uri = urljoin(f"{self.coordinator_url}/", uri) + request = Request(resolved_uri, headers=self._auth_headers_for(resolved_uri)) with self._opener.open(request, timeout=self.timeout) as response, destination.open("wb") as target: while chunk := response.read(1024 * 1024): target.write(chunk) - def upload(self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact) -> str: - """Stream one result artifact to the coordinator and return its stable URI.""" + def upload( + self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact + ) -> UploadedArtifact: + """Stream an artifact and require durable coordinator-owned metadata.""" url = ( f"{self.coordinator_url}/tasks/{quote(task.task_id, safe='')}/artifacts/" f"{quote(artifact.path.name, safe='')}" @@ -71,11 +59,13 @@ class HttpArtifactClient: http.client.HTTPSConnection if parsed.scheme == "https" else http.client.HTTPConnection ) connection = connection_class(parsed.hostname, parsed.port, timeout=self.timeout) + local_size = artifact.path.stat().st_size + local_sha256 = sha256_file(artifact.path) try: path = parsed.path + (f"?{parsed.query}" if parsed.query else "") connection.putrequest("PUT", path) connection.putheader("Content-Type", artifact.content_type) - connection.putheader("Content-Length", str(artifact.path.stat().st_size)) + connection.putheader("Content-Length", str(local_size)) connection.putheader("X-Worker-ID", worker_id) connection.putheader("X-Task-Attempt", str(task.attempt)) for name, value in self._auth_headers_for(url).items(): @@ -86,23 +76,24 @@ class HttpArtifactClient: connection.send(chunk) response = connection.getresponse() body = response.read() - if not 200 <= response.status < 300: + if response.status == 409: + raise CoordinatorConflictError("artifact upload rejected because the task lease was lost") + if response.status != 200: raise RuntimeError(f"artifact upload rejected with status {response.status}") - if body: - try: - response_data = json.loads(body) - except json.JSONDecodeError as error: - raise RuntimeError("artifact upload returned invalid JSON") from error - response_uri = response_data.get("uri") if isinstance(response_data, dict) else None - if isinstance(response_uri, str) and response_uri: - return response_uri - return url + try: + response_data = json.loads(body) + uploaded = UploadedArtifact.from_json(response_data) + except (ValueError, json.JSONDecodeError) as error: + raise RuntimeError("artifact upload returned invalid metadata") from error + if uploaded.sha256 != local_sha256 or uploaded.size_bytes != local_size: + raise RuntimeError("artifact upload metadata does not match local artifact") + return uploaded finally: connection.close() def _auth_headers_for(self, uri: str) -> dict[str, str]: """Only coordinator-owned URLs receive the coordinator bearer token.""" - if self.bearer_token and _origin(uri) == self.coordinator_origin: + if self.bearer_token and origin(uri) == self.coordinator_origin: return {"Authorization": f"Bearer {self.bearer_token}"} return {} diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py index 7d3bf1f..13e1fb5 100644 --- a/scimesh/worker/cli.py +++ b/scimesh/worker/cli.py @@ -14,22 +14,42 @@ from .runners import SciMeshRunner def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser(prog="scimesh-worker") + 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." + ), + ) parser.add_argument("--coordinator-url") parser.add_argument("--worker-id") parser.add_argument("--work-dir") + parser.add_argument("--worker-name") + parser.add_argument("--cpu-count", type=int) + parser.add_argument("--memory-mb", type=int) parser.add_argument("--poll-interval", type=float) parser.add_argument("--request-timeout", type=float) parser.add_argument("--heartbeat-interval", type=float) + parser.add_argument("--cleanup-after-seconds", type=float) args = parser.parse_args(argv) - config = WorkerConfig.from_environment() overrides = {key: value for key, value in vars(args).items() if value is not None} if "work_dir" in overrides: overrides["work_dir"] = Path(overrides["work_dir"]) - config = WorkerConfig(**{**config.__dict__, **overrides}) + try: + config = WorkerConfig.from_environment(overrides) + except (TypeError, ValueError) as error: + 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(config, client, HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token), SciMeshRunner()).run_forever() + WorkerDaemon( + config, + client, + HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token), + SciMeshRunner(), + ).run_forever() return 0 diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py index 0b5ca91..3732212 100644 --- a/scimesh/worker/config.py +++ b/scimesh/worker/config.py @@ -3,44 +3,99 @@ from __future__ import annotations from dataclasses import dataclass +from math import isfinite from pathlib import Path import os +import socket +from typing import Mapping +from urllib.parse import urlsplit + + +def _positive_number(value: object, name: str, *, allow_zero: bool = False) -> None: + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not isfinite(value) + or value < 0 + or (not allow_zero and value == 0) + ): + qualifier = "non-negative" if allow_zero else "positive" + raise ValueError(f"{name} must be {qualifier}") @dataclass(frozen=True) class WorkerConfig: coordinator_url: str - worker_id: str + worker_id: str | None work_dir: Path + worker_name: str = "scimesh-worker" + cpu_count: int = 1 + memory_mb: int | None = None poll_interval: float = 2.0 request_timeout: float = 30.0 heartbeat_interval: float = 15.0 bearer_token: str | None = None cleanup_after_seconds: float | None = None - capabilities: tuple[str, ...] = ("similarity-search", "similarity-graph") + # The local CLI uses hyphens; the first coordinator contract used + # underscores. Advertise both stable spellings while jobs are migrated. + capabilities: tuple[str, ...] = ( + "similarity-search", + "similarity-graph", + "similarity_search", + "similarity_graph", + ) def __post_init__(self) -> None: - if self.poll_interval <= 0: - raise ValueError("poll_interval must be positive") - if self.request_timeout <= 0: - raise ValueError("request_timeout must be positive") - if self.heartbeat_interval <= 0: - raise ValueError("heartbeat_interval must be positive") + parsed = urlsplit(self.coordinator_url) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise ValueError("coordinator_url must be an absolute HTTP(S) URL") + if not isinstance(self.worker_name, str) or not self.worker_name.strip(): + raise ValueError("worker_name must be non-empty") + if isinstance(self.cpu_count, bool) or not isinstance(self.cpu_count, int) or self.cpu_count < 1: + raise ValueError("cpu_count must be positive") + if self.worker_id is not None and not isinstance(self.worker_id, str): + raise ValueError("worker_id must be a string when set") + if self.memory_mb is not None and ( + isinstance(self.memory_mb, bool) + or not isinstance(self.memory_mb, int) + or self.memory_mb < 1 + ): + raise ValueError("memory_mb must be positive when set") + _positive_number(self.poll_interval, "poll_interval") + _positive_number(self.request_timeout, "request_timeout") + _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 not self.capabilities: + raise ValueError("capabilities cannot be empty") @classmethod - def from_environment(cls) -> "WorkerConfig": - url = os.getenv("SCIMESH_COORDINATOR_URL") - worker_id = os.getenv("SCIMESH_WORKER_ID") - if not url or not worker_id: - raise ValueError("SCIMESH_COORDINATOR_URL and SCIMESH_WORKER_ID are required") - cleanup = os.getenv("SCIMESH_CLEANUP_AFTER_SECONDS") + def from_environment( + cls, overrides: Mapping[str, object] | None = None + ) -> "WorkerConfig": + """Build config from environment, allowing typed CLI values to override it.""" + values = overrides or {} + + def value(name: str, environment: str, default: object | None = None) -> object | None: + override = values.get(name) + return override if override is not None else os.getenv(environment, default) + + url = value("coordinator_url", "SCIMESH_COORDINATOR_URL") + if not isinstance(url, str) or not url: + raise ValueError("SCIMESH_COORDINATOR_URL or --coordinator-url is required") + 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") return cls( coordinator_url=url.rstrip("/"), - worker_id=worker_id, - work_dir=Path(os.getenv("SCIMESH_WORK_DIR", "./scimesh-worker-data")), - poll_interval=float(os.getenv("SCIMESH_POLL_INTERVAL", "2")), - request_timeout=float(os.getenv("SCIMESH_REQUEST_TIMEOUT", "30")), - heartbeat_interval=float(os.getenv("SCIMESH_HEARTBEAT_INTERVAL", "15")), - bearer_token=os.getenv("SCIMESH_BEARER_TOKEN"), + worker_id=value("worker_id", "SCIMESH_WORKER_ID"), + work_dir=Path(value("work_dir", "SCIMESH_WORK_DIR", "./scimesh-worker-data")), + worker_name=str(value("worker_name", "SCIMESH_WORKER_NAME", socket.gethostname())), + cpu_count=int(cpu_count), + memory_mb=int(memory_mb) if memory_mb is not None else None, + poll_interval=float(value("poll_interval", "SCIMESH_POLL_INTERVAL", "2")), + request_timeout=float(value("request_timeout", "SCIMESH_REQUEST_TIMEOUT", "30")), + 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, ) diff --git a/scimesh/worker/coordinator.py b/scimesh/worker/coordinator.py index 5fccd33..7b575b7 100644 --- a/scimesh/worker/coordinator.py +++ b/scimesh/worker/coordinator.py @@ -5,9 +5,10 @@ from __future__ import annotations import json from typing import Any, Protocol from urllib.error import HTTPError, URLError -from urllib.request import Request, urlopen +from urllib.request import Request, build_opener -from .models import ClaimedTask +from .models import ClaimedTask, RegisteredWorker +from .transport import NoRedirectHandler class CoordinatorError(RuntimeError): @@ -18,7 +19,15 @@ class CoordinatorTransientError(CoordinatorError): """A timeout, connection error, or 5xx coordinator response.""" +class CoordinatorConflictError(CoordinatorError): + """The worker no longer owns the task lease or attempted a conflicting mutation.""" + + class CoordinatorClient(Protocol): + def register( + self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None + ) -> RegisteredWorker: ... + def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: ... def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ... @@ -33,6 +42,25 @@ class HttpCoordinatorClient: self.base_url = base_url.rstrip("/") self.timeout = timeout self.bearer_token = bearer_token + self._opener = build_opener(NoRedirectHandler()) + + def register( + self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None + ) -> RegisteredWorker: + payload: dict[str, Any] = { + "name": name, + "capabilities": list(capabilities), + "cpu_count": cpu_count, + } + if memory_mb is not None: + payload["memory_mb"] = memory_mb + status, body = self._request("POST", "/workers/register", payload) + if status != 201: + raise CoordinatorError(f"worker registration rejected with status {status}") + try: + return RegisteredWorker.from_json(body) + except ValueError as error: + raise CoordinatorError("invalid worker registration response") from error def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: status, body = self._request("POST", "/tasks/claim", { @@ -48,11 +76,15 @@ class HttpCoordinatorClient: status, _ = self._request("POST", f"/tasks/{task.task_id}/result", payload) # 200/201/202 include a successful or idempotent duplicate result response. if status not in (200, 201, 202): + if status == 409: + raise CoordinatorConflictError("result rejected because the task lease was lost") raise CoordinatorError(f"result rejected with status {status}") def fail(self, task: ClaimedTask, payload: dict[str, Any]) -> None: status, _ = self._request("POST", f"/tasks/{task.task_id}/failure", payload) if status not in (200, 201, 202): + if status == 409: + raise CoordinatorConflictError("failure rejected because the task lease was lost") raise CoordinatorError(f"failure report rejected with status {status}") def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: @@ -61,6 +93,8 @@ class HttpCoordinatorClient: {"worker_id": worker_id, "attempt": task.attempt}, ) if status != 200: + if status == 409: + raise CoordinatorConflictError("heartbeat rejected because the task lease was lost") raise CoordinatorError(f"heartbeat rejected with status {status}") lease_expires_at = body.get("lease_expires_at") if not isinstance(lease_expires_at, str): @@ -73,9 +107,12 @@ class HttpCoordinatorClient: headers={"Content-Type": "application/json", **self._auth_header()}, ) try: - with urlopen(request, timeout=self.timeout) as response: + with self._opener.open(request, timeout=self.timeout) as response: raw = response.read() - return response.status, json.loads(raw) if raw else {} + try: + return response.status, json.loads(raw) if raw else {} + except json.JSONDecodeError as error: + raise CoordinatorError("coordinator returned invalid JSON") from error except HTTPError as error: if error.code >= 500: raise CoordinatorTransientError(f"coordinator returned {error.code}") from error diff --git a/scimesh/worker/daemon.py b/scimesh/worker/daemon.py index 3200f6b..dbf035b 100644 --- a/scimesh/worker/daemon.py +++ b/scimesh/worker/daemon.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +from dataclasses import replace from pathlib import Path import random import shutil @@ -12,8 +13,8 @@ from datetime import datetime, timezone from .artifacts import ArtifactClient, sha256_file from .config import WorkerConfig -from .coordinator import CoordinatorClient, CoordinatorTransientError -from .models import ClaimedTask +from .coordinator import CoordinatorClient, CoordinatorConflictError, CoordinatorTransientError +from .models import ClaimedTask, UploadedArtifact from .runners import Runner @@ -32,6 +33,7 @@ class LeaseHeartbeat: self._lease_expires_at = self.coordinator.heartbeat( self.task, self.config.worker_id ) + self._next_delay() self._thread = threading.Thread(target=self._run, name=f"lease-{self.task.task_id}", daemon=True) self._thread.start() @@ -45,18 +47,19 @@ class LeaseHeartbeat: raise self._error def _run(self) -> None: - delay = min(self.config.heartbeat_interval, self._seconds_until_expiry() / 2) + delay = self._next_delay() while not self._stop.wait(max(delay, 0.01)): try: self._lease_expires_at = self.coordinator.heartbeat( self.task, self.config.worker_id ) + delay = self._next_delay() except Exception as error: # Surface the lease loss in the main state machine. self._error = error return - delay = min( - self.config.heartbeat_interval, self._seconds_until_expiry() / 2 - ) + + def _next_delay(self) -> float: + return min(self.config.heartbeat_interval, self._seconds_until_expiry() / 2) def _seconds_until_expiry(self) -> float: try: @@ -72,12 +75,16 @@ class LeaseHeartbeat: 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 + self.worker_id = config.worker_id + self._registered = False self.log = logging.getLogger("scimesh.worker") def run_forever(self) -> None: failures = 0 while True: try: + if not self._registered: + self._register_worker() self._cleanup_expired_directories() claimed = self.run_once() failures = 0 @@ -89,16 +96,17 @@ class WorkerDaemon: self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0)) def run_once(self) -> bool: + worker_id = self._worker_id() self._log("claiming") - task = self.coordinator.claim(self.config.worker_id, self.config.capabilities) + task = self.coordinator.claim(worker_id, self.config.capabilities) if task is None: self._log("idle") return False started = time.monotonic() task_dir = self.config.work_dir / task.task_id / str(task.attempt) - task_dir.mkdir(parents=True, exist_ok=False) heartbeat = LeaseHeartbeat(task, self.coordinator, self.config) try: + task_dir.mkdir(parents=True, exist_ok=False) heartbeat.start() self._log("downloading", task) input_path = task_dir / "input" @@ -108,20 +116,28 @@ class WorkerDaemon: self._log("running", task) result = self.runner.run(task, task_dir) heartbeat.raise_if_failed() - manifests = [ - { - "uri": self.artifacts.upload(task, self.config.worker_id, artifact), - "sha256": sha256_file(artifact.path), - "content_type": artifact.content_type, - } - for artifact in result.artifacts - ] - if not manifests: - raise ValueError("runner produced no artifacts") + if len(result.artifacts) != 1: + raise ValueError("runner must produce exactly one result artifact") + artifact = result.artifacts[0] + uploaded = self.artifacts.upload(task, worker_id, artifact) + manifest = self._result_manifest(uploaded) self._log("submitting", task) heartbeat.raise_if_failed() - self.coordinator.submit(task, {"worker_id": self.config.worker_id, "attempt": task.attempt, "status": "completed", "result": manifests[0], "artifacts": manifests, "metrics": {**result.metrics, "elapsed_seconds": round(time.monotonic() - started, 3)}}) + self.coordinator.submit( + task, + { + "worker_id": worker_id, + "attempt": task.attempt, + "result": manifest, + "metrics": { + **result.metrics, + "elapsed_seconds": round(time.monotonic() - started, 3), + }, + }, + ) self._log("idle", task, elapsed_seconds=round(time.monotonic() - started, 3)) + except CoordinatorConflictError as error: + self._log("lease_lost", task, error_type=type(error).__name__) except Exception as error: self._log("failed", task, error_type=type(error).__name__) self._report_failure(task, error) @@ -132,12 +148,38 @@ class WorkerDaemon: def _report_failure(self, task: ClaimedTask, error: Exception) -> None: message = str(error).replace(str(self.config.work_dir), "")[:300] try: - self.coordinator.fail(task, {"worker_id": self.config.worker_id, "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message}) + self.coordinator.fail(task, {"worker_id": self._worker_id(), "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message}) except CoordinatorTransientError: raise except Exception: self._log("failed", task, error_type="FailureReportError") + def _register_worker(self) -> None: + registered = self.coordinator.register( + self.config.worker_name, + self.config.capabilities, + self.config.cpu_count, + self.config.memory_mb, + ) + self.worker_id = registered.worker_id + self.config = replace( + self.config, + worker_id=registered.worker_id, + heartbeat_interval=registered.heartbeat_interval_seconds, + ) + self._registered = True + self._log("registered") + + def _worker_id(self) -> str: + if not self.worker_id: + raise ValueError("worker is not registered") + return self.worker_id + + @staticmethod + def _result_manifest(uploaded: UploadedArtifact) -> dict[str, object]: + """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: 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) diff --git a/scimesh/worker/models.py b/scimesh/worker/models.py index 97de6a1..9cf4f97 100644 --- a/scimesh/worker/models.py +++ b/scimesh/worker/models.py @@ -3,8 +3,40 @@ from __future__ import annotations from dataclasses import dataclass +from datetime import datetime +from math import isfinite from pathlib import Path from typing import Any +from urllib.parse import urlsplit +from uuid import UUID + + +def _required_string(value: object, field: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{field} must be a non-empty string") + return value + + +def _coordinator_uri(value: object, field: str) -> str: + uri = _required_string(value, field) + parsed = urlsplit(uri) + if uri.startswith("/"): + # ``//host/path`` is a network-path reference: urljoin would resolve + # it to another origin. Dot segments are rejected for the same reason + # we reject unsafe local task identifiers. + if parsed.netloc or any(segment == ".." for segment in parsed.path.split("/")): + raise ValueError(f"{field} must be a safe coordinator path") + return uri + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise ValueError(f"{field} must be an absolute HTTP(S) URL or coordinator path") + return uri + + +def _sha256(value: object, field: str) -> str: + digest = _required_string(value, field).lower() + if len(digest) != 64 or any(character not in "0123456789abcdef" for character in digest): + raise ValueError(f"{field} must be a SHA-256 hex digest") + return digest @dataclass(frozen=True) @@ -26,13 +58,28 @@ class ClaimedTask: def from_json(cls, data: dict[str, Any]) -> "ClaimedTask": try: input_data = data["input"] + if not isinstance(input_data, dict): + raise ValueError("input must be an object") + raw_attempt = data["attempt"] + if isinstance(raw_attempt, bool) or not isinstance(raw_attempt, int) or raw_attempt < 1: + raise ValueError("attempt must be a positive integer") + task_id = str(UUID(_required_string(data["task_id"], "task_id"))) + lease_expires_at = _required_string(data["lease_expires_at"], "lease_expires_at") + if datetime.fromisoformat(lease_expires_at.replace("Z", "+00:00")).tzinfo is None: + raise ValueError("lease_expires_at must include a timezone") + parameters = data.get("parameters", {}) + if not isinstance(parameters, dict): + raise ValueError("parameters must be an object") return cls( - task_id=str(data["task_id"]), - attempt=int(data["attempt"]), - lease_expires_at=str(data["lease_expires_at"]), - workload=str(data["workload"]), - input=InputArtifact(uri=str(input_data["uri"]), sha256=str(input_data["sha256"])), - parameters=dict(data.get("parameters", {})), + task_id=task_id, + attempt=raw_attempt, + lease_expires_at=lease_expires_at, + workload=_required_string(data["workload"], "workload"), + input=InputArtifact( + uri=_coordinator_uri(input_data["uri"], "input.uri"), + sha256=_sha256(input_data["sha256"], "input.sha256"), + ), + parameters=parameters, ) except (KeyError, TypeError, ValueError) as error: raise ValueError("invalid claimed-task response") from error @@ -44,6 +91,61 @@ class ProducedArtifact: content_type: str +@dataclass(frozen=True) +class UploadedArtifact: + """Coordinator-owned artifact metadata returned after a successful upload.""" + + artifact_id: str + uri: str + sha256: str + size_bytes: int + + @classmethod + def from_json(cls, data: object) -> "UploadedArtifact": + if not isinstance(data, dict): + raise ValueError("artifact upload response must be an object") + raw_size = data.get("size_bytes") + if isinstance(raw_size, bool) or not isinstance(raw_size, int) or raw_size < 0: + raise ValueError("artifact size_bytes must be a non-negative integer") + try: + return cls( + artifact_id=str(UUID(_required_string(data.get("artifact_id"), "artifact_id"))), + uri=_coordinator_uri(data.get("uri"), "uri"), + sha256=_sha256(data.get("sha256"), "sha256"), + size_bytes=raw_size, + ) + except ValueError as error: + raise ValueError("invalid artifact upload response") from error + + +@dataclass(frozen=True) +class RegisteredWorker: + """Identity and heartbeat policy returned by worker registration.""" + + worker_id: str + heartbeat_interval_seconds: float + + @classmethod + def from_json(cls, data: object) -> "RegisteredWorker": + if not isinstance(data, dict): + raise ValueError("worker registration response must be an object") + raw_interval = data.get("heartbeat_interval_seconds") + if ( + isinstance(raw_interval, bool) + or not isinstance(raw_interval, (int, float)) + or not isfinite(raw_interval) + or raw_interval <= 0 + ): + raise ValueError("heartbeat_interval_seconds must be positive") + try: + return cls( + worker_id=str(UUID(_required_string(data.get("worker_id"), "worker_id"))), + heartbeat_interval_seconds=float(raw_interval), + ) + except ValueError as error: + raise ValueError("invalid worker registration response") from error + + @dataclass(frozen=True) class RunResult: artifacts: tuple[ProducedArtifact, ...] diff --git a/scimesh/worker/runners.py b/scimesh/worker/runners.py index 46d74af..0ef3e77 100644 --- a/scimesh/worker/runners.py +++ b/scimesh/worker/runners.py @@ -20,9 +20,13 @@ class SciMeshRunner: def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: input_path = task_dir / "input" output_path = task_dir / "result.csv" - command = [sys.executable, "-m", "scimesh.cli", task.workload, str(input_path)] + # The coordinator contract historically used underscores while the + # public SciMesh CLI uses hyphens. Accept both spellings at this narrow + # boundary so an API job cannot turn into an opaque worker failure. + workload = task.workload.replace("_", "-") + command = [sys.executable, "-m", "scimesh.cli", workload, str(input_path)] params = task.parameters - if task.workload == "similarity-search": + if workload == "similarity-search": self._reject_unknown(params, {"query_id", "query_smiles", "top_k", "threshold", "threshold_direction", "max_rows", "progress_every"}) query_id, query_smiles = params.get("query_id"), params.get("query_smiles") if (query_id is None) == (query_smiles is None): @@ -31,7 +35,7 @@ class SciMeshRunner: command += ["--query-id", self._string(params, "query_id")] if query_id is not None else ["--query-smiles", self._string(params, "query_smiles")] command += ["--top-k", str(top_k)] self._append_common_options(command, params) - elif task.workload == "similarity-graph": + elif workload == "similarity-graph": self._reject_unknown(params, {"threshold", "threshold_direction", "block_size", "max_rows", "progress_every"}) threshold = self._number(params, "threshold") command += ["--threshold", str(threshold)] diff --git a/scimesh/worker/transport.py b/scimesh/worker/transport.py new file mode 100644 index 0000000..3f975f8 --- /dev/null +++ b/scimesh/worker/transport.py @@ -0,0 +1,51 @@ +"""Small HTTP transport helpers shared by coordinator and artifact clients.""" + +from __future__ import annotations + +from urllib.request import HTTPRedirectHandler, Request +from urllib.parse import urlsplit + + +def origin(uri: str) -> tuple[str, str, int | None]: + """Return a normalized HTTP origin for authorization decisions.""" + parsed = urlsplit(uri) + scheme = parsed.scheme.lower() + default_port = {"http": 80, "https": 443}.get(scheme) + return scheme, (parsed.hostname or "").lower(), parsed.port or default_port + + +class SameOriginAuthRedirectHandler(HTTPRedirectHandler): + """Strip coordinator authorization when an artifact redirect changes origin.""" + + def __init__(self, coordinator_origin: tuple[str, str, int | None]) -> None: + super().__init__() + self.coordinator_origin = coordinator_origin + + def redirect_request( + self, + req: Request, + fp: object, + code: int, + msg: str, + headers: object, + newurl: str, + ) -> Request | None: + redirected = super().redirect_request(req, fp, code, msg, headers, newurl) + if redirected and origin(newurl) != self.coordinator_origin: + redirected.remove_header("Authorization") + return redirected + + +class NoRedirectHandler(HTTPRedirectHandler): + """Reject redirects for mutating coordinator API calls.""" + + def redirect_request( + self, + req: Request, + fp: object, + code: int, + msg: str, + headers: object, + newurl: str, + ) -> Request | None: + return None diff --git a/scimesh/workloads/similarity_graph.py b/scimesh/workloads/similarity_graph.py index 4c6fa0d..08df9a2 100644 --- a/scimesh/workloads/similarity_graph.py +++ b/scimesh/workloads/similarity_graph.py @@ -47,10 +47,15 @@ def _fingerprinted_molecules( tsv_path: Path, max_rows: int | None ) -> tuple[list[GraphMolecule], DatasetStats]: stats = DatasetStats() - molecules = [ - GraphMolecule(record.molecule_id, fingerprint(record.molecule)) - for record in iter_valid_molecules(tsv_path, stats, max_rows=max_rows) - ] + molecules: list[GraphMolecule] = [] + seen_ids: set[str] = set() + for record in iter_valid_molecules(tsv_path, stats, max_rows=max_rows): + if not record.molecule_id: + raise ValueError("Dataset contains an empty chembl_id") + if record.molecule_id in seen_ids: + raise ValueError(f"Dataset contains a duplicate chembl_id: {record.molecule_id}") + seen_ids.add(record.molecule_id) + molecules.append(GraphMolecule(record.molecule_id, fingerprint(record.molecule))) return molecules, stats @@ -119,6 +124,7 @@ def build_similarity_graph( def write_graph_edges(output_path: Path, edges: list[SimilarityEdge]) -> None: """Write a deterministic sparse edge list CSV.""" + output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("w", encoding="utf-8", newline="") as destination: writer = csv.DictWriter(destination, fieldnames=["source_id", "target_id", "similarity"]) writer.writeheader() diff --git a/scimesh/workloads/similarity_search.py b/scimesh/workloads/similarity_search.py index d3c92de..574749a 100644 --- a/scimesh/workloads/similarity_search.py +++ b/scimesh/workloads/similarity_search.py @@ -135,6 +135,7 @@ def search_similar( def write_search_results(output_path: Path, matches: list[SimilarityMatch]) -> None: """Write ranked matches to a deterministic CSV file.""" + output_path.parent.mkdir(parents=True, exist_ok=True) with output_path.open("w", encoding="utf-8", newline="") as destination: writer = csv.DictWriter( destination, diff --git a/tests/test_similarity_graph.py b/tests/test_similarity_graph.py index 9341e01..2399ac2 100644 --- a/tests/test_similarity_graph.py +++ b/tests/test_similarity_graph.py @@ -2,6 +2,7 @@ from __future__ import annotations from pathlib import Path +import pytest from rdkit import DataStructs from scimesh.chemistry.dataset import DatasetStats, iter_valid_molecules @@ -61,3 +62,19 @@ def test_graph_supports_less_than_threshold_direction(small_dataset: Path) -> No ) assert all(edge.similarity <= 0.15 for edge in result.edges) + + +def test_graph_rejects_duplicate_identifiers(tmp_path: Path) -> None: + dataset = tmp_path / "duplicate_ids.tsv" + dataset.write_text( + "chembl_id\tcanonical_smiles\nDUP\tCCO\nDUP\tCCC\n", encoding="utf-8" + ) + + with pytest.raises(ValueError, match="duplicate chembl_id"): + build_similarity_graph(dataset, threshold=0.1, block_size=1) + + +def test_graph_writer_creates_missing_output_directory(tmp_path: Path) -> None: + output = tmp_path / "nested" / "edges.csv" + write_graph_edges(output, []) + assert output.read_text(encoding="utf-8").startswith("source_id,target_id") diff --git a/tests/test_similarity_search.py b/tests/test_similarity_search.py index 4af1c62..66eb5e2 100644 --- a/tests/test_similarity_search.py +++ b/tests/test_similarity_search.py @@ -6,7 +6,11 @@ from rdkit import Chem, DataStructs from scimesh.chemistry.dataset import DatasetStats, find_molecule_by_id, iter_valid_molecules from scimesh.chemistry.fingerprints import fingerprint -from scimesh.workloads.similarity_search import SimilarityMatch, search_similar +from scimesh.workloads.similarity_search import ( + SimilarityMatch, + search_similar, + write_search_results, +) def test_search_matches_full_sorting_and_skips_query_and_invalid( @@ -56,3 +60,9 @@ def test_search_can_rank_and_filter_least_similar_molecules( assert result.matches == sorted( result.matches, key=lambda match: match.sort_key("less") ) + + +def test_search_writer_creates_missing_output_directory(tmp_path: Path) -> None: + output = tmp_path / "nested" / "results.csv" + write_search_results(output, []) + assert output.read_text(encoding="utf-8").startswith("rank,chembl_id") diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py index 583ee93..6ad73c9 100644 --- a/tests/test_worker_daemon.py +++ b/tests/test_worker_daemon.py @@ -11,9 +11,17 @@ import pytest from scimesh.worker.config import WorkerConfig from scimesh.worker.coordinator import CoordinatorTransientError from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon -from scimesh.worker.models import ClaimedTask, InputArtifact, ProducedArtifact, RunResult +from scimesh.worker.models import ( + ClaimedTask, + InputArtifact, + ProducedArtifact, + RegisteredWorker, + RunResult, + UploadedArtifact, +) from scimesh.worker.artifacts import HttpArtifactClient, _SameOriginAuthRedirectHandler, _origin from scimesh.worker.runners import SciMeshRunner +from scimesh.worker.transport import NoRedirectHandler class FakeCoordinator: @@ -24,6 +32,11 @@ class FakeCoordinator: task, self.task = self.task, None return task + def register( + self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None + ) -> RegisteredWorker: + return RegisteredWorker("11111111-1111-4111-8111-111111111111", 15) + def submit(self, task: ClaimedTask, payload: dict) -> None: self.submissions.append(payload) @@ -42,9 +55,17 @@ class FakeArtifacts: def download(self, uri: str, destination: Path) -> None: destination.write_bytes(self.content) - def upload(self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact) -> str: + def upload( + self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact + ) -> UploadedArtifact: self.uploaded.append((task.task_id, worker_id, artifact.path)) - return f"https://example.test/tasks/{task.task_id}/artifacts/{artifact.path.name}" + content = artifact.path.read_bytes() + return UploadedArtifact( + "22222222-2222-4222-8222-222222222222", + f"https://example.test/tasks/{task.task_id}/artifacts/{artifact.path.name}", + hashlib.sha256(content).hexdigest(), + len(content), + ) class FakeRunner: def __init__(self) -> None: @@ -75,9 +96,10 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None: assert runner.calls == 1 assert len(artifacts.uploaded) == 1 assert coordinator.heartbeats == [("task-1", 1, "worker-1")] - assert coordinator.submissions[0]["status"] == "completed" - assert coordinator.submissions[0]["result"]["content_type"] == "text/csv" - assert coordinator.submissions[0]["result"]["uri"].startswith("https://example.test/tasks/task-1/artifacts/") + assert "status" not in coordinator.submissions[0] + assert coordinator.submissions[0]["result"] == { + "artifact_id": "22222222-2222-4222-8222-222222222222" + } def test_no_task_does_not_create_directory(tmp_path: Path) -> None: @@ -95,6 +117,14 @@ def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None: assert not coordinator.submissions +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 coordinator.failures[0]["error_code"] == "FileExistsError" + + def test_transient_claim_error_is_propagated_for_bounded_backoff(tmp_path: Path) -> None: class UnavailableCoordinator(FakeCoordinator): def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: @@ -123,6 +153,15 @@ def test_input_token_is_sent_only_to_the_coordinator_origin() -> None: assert client._auth_headers_for("https://bucket.example/presigned") == {} +def test_relative_input_uri_is_resolved_against_the_coordinator() -> None: + client = HttpArtifactClient("https://coordinator.example/api", 10, "secret") + assert client._auth_headers_for("https://coordinator.example/tasks/1/input") == { + "Authorization": "Bearer secret" + } + # The coordinator's contract returns root-relative artifact paths. + assert client.coordinator_url == "https://coordinator.example/api" + + def test_redirect_to_external_storage_strips_authorization() -> None: handler = _SameOriginAuthRedirectHandler(_origin("https://coordinator.example")) source = Request( @@ -133,6 +172,12 @@ def test_redirect_to_external_storage_strips_authorization() -> None: assert redirected.get_header("Authorization") is None +def test_api_requests_never_follow_redirects() -> None: + handler = NoRedirectHandler() + request = Request("https://coordinator.example/tasks/claim", headers={"Authorization": "Bearer secret"}) + assert handler.redirect_request(request, None, 302, "Found", {}, "https://other.example") is None + + def test_lease_is_renewed_while_a_runner_is_still_working(tmp_path: Path) -> None: content = b"input fixture" worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content) @@ -184,3 +229,93 @@ def test_runner_maps_graph_and_smiles_search_parameters(tmp_path: Path, monkeypa assert "--block-size" in commands[0] and "42" in commands[0] assert "--max-rows" in commands[0] and "7" in commands[0] assert "--query-smiles" in commands[1] and "CCO" in commands[1] + + +def test_runner_accepts_coordinator_workload_names(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + commands: list[list[str]] = [] + + def fake_run(command: list[str], **_: object) -> None: + commands.append(command) + output = Path(command[command.index("--output") + 1]) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text("a,b\\n", encoding="utf-8") + + monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run) + task = ClaimedTask( + "search", 1, "2026-07-30T00:00:00Z", "similarity_search", + InputArtifact("https://example/input", "a" * 64), {"query_smiles": "CCO"}, + ) + SciMeshRunner().run(task, tmp_path / "search") + assert commands[0][3] == "similarity-search" + + +def test_claimed_task_rejects_path_traversal_and_invalid_metadata() -> None: + payload = { + "task_id": "../outside", + "attempt": 1, + "lease_expires_at": "2026-07-30T00:00:00Z", + "workload": "similarity-search", + "input": {"uri": "https://example.test/input", "sha256": "a" * 64}, + "parameters": {}, + } + with pytest.raises(ValueError, match="invalid claimed-task response"): + ClaimedTask.from_json(payload) + + payload["task_id"] = "11111111-1111-4111-8111-111111111111" + payload["input"] = {"uri": "//outside.example/input", "sha256": "a" * 64} + with pytest.raises(ValueError, match="invalid claimed-task response"): + ClaimedTask.from_json(payload) + + payload["input"] = {"uri": "/tasks/../outside/input", "sha256": "a" * 64} + with pytest.raises(ValueError, match="invalid claimed-task response"): + ClaimedTask.from_json(payload) + + +def test_claimed_task_accepts_a_coordinator_relative_input_path() -> None: + task = ClaimedTask.from_json( + { + "task_id": "11111111-1111-4111-8111-111111111111", + "attempt": 1, + "lease_expires_at": "2026-07-30T00:00:00Z", + "workload": "similarity_search", + "input": {"uri": "/tasks/11111111-1111-4111-8111-111111111111/input", "sha256": "a" * 64}, + "parameters": {}, + } + ) + assert task.input.uri.startswith("/tasks/") + + +def test_uploaded_artifact_requires_complete_durable_metadata() -> None: + artifact = UploadedArtifact.from_json( + { + "artifact_id": "22222222-2222-4222-8222-222222222222", + "uri": "https://coordinator.example/artifacts/222/download", + "sha256": "a" * 64, + "size_bytes": 12, + } + ) + assert artifact.size_bytes == 12 + with pytest.raises(ValueError, match="artifact size_bytes"): + UploadedArtifact.from_json({"artifact_id": "missing"}) + + +def test_environment_overrides_allow_cli_only_configuration(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + monkeypatch.delenv("SCIMESH_COORDINATOR_URL", raising=False) + config = WorkerConfig.from_environment( + { + "coordinator_url": "https://coordinator.example", + "work_dir": tmp_path, + "worker_name": "test-worker", + } + ) + assert config.coordinator_url == "https://coordinator.example" + assert config.worker_id is None + assert "similarity-search" in config.capabilities + assert "similarity_search" in config.capabilities + + +def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None: + worker, _, _, _, _ = daemon(tmp_path, None, b"") + worker._register_worker() + assert worker.worker_id == "11111111-1111-4111-8111-111111111111" + assert worker.config.heartbeat_interval == 15 From 484ecd0dfa9b4896abe8bf8323b0b02ade419041 Mon Sep 17 00:00:00 2001 From: Emil Date: Thu, 23 Jul 2026 21:44:44 +0300 Subject: [PATCH 2/2] Bind result artifacts to lease attempts --- coordinator/internal/domain/artifact.go | 1 + .../storage/postgres/artifact_repo.go | 6 +-- .../storage/postgres/integration_test.go | 33 ++++++++++++ coordinator/internal/usecase/artifact.go | 15 ++++++ coordinator/internal/usecase/task.go | 6 +-- coordinator/internal/usecase/usecase_test.go | 52 +++++++++++++++++++ .../migrations/0008_artifact_attempt.down.sql | 7 +++ .../migrations/0008_artifact_attempt.up.sql | 34 ++++++++++++ 8 files changed, 148 insertions(+), 6 deletions(-) create mode 100644 coordinator/migrations/0008_artifact_attempt.down.sql create mode 100644 coordinator/migrations/0008_artifact_attempt.up.sql diff --git a/coordinator/internal/domain/artifact.go b/coordinator/internal/domain/artifact.go index 43bc219..dbe287f 100644 --- a/coordinator/internal/domain/artifact.go +++ b/coordinator/internal/domain/artifact.go @@ -23,6 +23,7 @@ type Artifact struct { ID uuid.UUID JobID uuid.UUID TaskID *uuid.UUID // nil for a job-level input + Attempt *int // required for a partial result; nil for non-worker artifacts Kind ArtifactKind Filename string StorageKey string diff --git a/coordinator/internal/storage/postgres/artifact_repo.go b/coordinator/internal/storage/postgres/artifact_repo.go index 771b5ee..f42abac 100644 --- a/coordinator/internal/storage/postgres/artifact_repo.go +++ b/coordinator/internal/storage/postgres/artifact_repo.go @@ -26,14 +26,14 @@ func NewArtifactRepo(pool *pgxpool.Pool) *ArtifactRepo { var _ usecase.ArtifactRepository = (*ArtifactRepo)(nil) var artifactColumns = []string{ - "id", "job_id", "task_id", "kind", "filename", "storage_key", + "id", "job_id", "task_id", "attempt", "kind", "filename", "storage_key", "content_type", "size_bytes", "sha256", "created_at", } func (r *ArtifactRepo) Insert(ctx context.Context, a *domain.Artifact) error { sql, args, err := psql.Insert("artifacts"). Columns(artifactColumns...). - Values(a.ID, a.JobID, a.TaskID, string(a.Kind), a.Filename, a.StorageKey, + Values(a.ID, a.JobID, a.TaskID, a.Attempt, string(a.Kind), a.Filename, a.StorageKey, a.ContentType, a.SizeBytes, a.SHA256, a.CreatedAt). ToSql() if err != nil { @@ -59,7 +59,7 @@ func (r *ArtifactRepo) Get(ctx context.Context, id uuid.UUID) (*domain.Artifact, kind string ) err = conn(ctx, r.pool).QueryRow(ctx, sql, args...).Scan( - &a.ID, &a.JobID, &a.TaskID, &kind, &a.Filename, &a.StorageKey, + &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, domain.ErrArtifactNotFound diff --git a/coordinator/internal/storage/postgres/integration_test.go b/coordinator/internal/storage/postgres/integration_test.go index 05a7e90..fd461bd 100644 --- a/coordinator/internal/storage/postgres/integration_test.go +++ b/coordinator/internal/storage/postgres/integration_test.go @@ -259,6 +259,8 @@ func TestListCompletedIsOrderedByChunkIndex(t *testing.T) { return err } art.SetContent(fmt.Sprintf("rsha-%d", task.ChunkIndex), 1) + attempt := 1 + art.Attempt = &attempt if err := artifacts.Insert(ctx, art); err != nil { return err } @@ -269,6 +271,7 @@ func TestListCompletedIsOrderedByChunkIndex(t *testing.T) { } owner := "worker-1" fresh.Status = domain.TaskLeased + fresh.Attempt = attempt fresh.LeaseOwner = &owner expires := now.Add(time.Minute) fresh.LeaseExpiresAt = &expires @@ -350,6 +353,10 @@ func seedArtifact(t *testing.T, pool *pgxpool.Pool, jobID uuid.UUID, taskID *uui t.Fatalf("build artifact: %v", err) } art.SetContent(fmt.Sprintf("sha-%s", art.ID), 3) + if kind == domain.ArtifactPartialResult { + attempt := 1 + art.Attempt = &attempt + } if err := NewArtifactRepo(pool).Insert(context.Background(), art); err != nil { t.Fatalf("insert artifact: %v", err) } @@ -446,6 +453,32 @@ func TestArtifactRepoRoundTrip(t *testing.T) { } } +func TestPartialResultArtifactRoundTripsAttempt(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, tasks := seedJob(t, pool, 1) + taskID := tasks[0].ID + art, err := domain.NewArtifact(job.ID, &taskID, domain.ArtifactPartialResult, + "result.csv", "text/csv", time.Now().UTC()) + if err != nil { + t.Fatal(err) + } + attempt := 2 + art.Attempt = &attempt + art.SetContent("sha", 3) + repo := NewArtifactRepo(pool) + if err := repo.Insert(ctx, art); err != nil { + t.Fatalf("insert: %v", err) + } + got, err := repo.Get(ctx, art.ID) + if err != nil { + t.Fatalf("get: %v", err) + } + if got.Attempt == nil || *got.Attempt != attempt { + t.Fatalf("attempt = %v, want %d", got.Attempt, attempt) + } +} + // A shard task stores its input as an artifact and no URI: this exercises the // nullable input_uri column, the input_artifact_id round-trip, and the // ck_tasks_has_input check that requires one or the other. diff --git a/coordinator/internal/usecase/artifact.go b/coordinator/internal/usecase/artifact.go index dcbd64d..5cb8b8b 100644 --- a/coordinator/internal/usecase/artifact.go +++ b/coordinator/internal/usecase/artifact.go @@ -39,6 +39,8 @@ func (uc *UploadArtifact) Execute(ctx context.Context, in UploadArtifactInput) ( if err != nil { return nil, err } + attempt := in.Attempt + art.Attempt = &attempt // Stream to storage first: size and checksum are measured here, by us, not // taken from the worker. A large shard never sits in memory. @@ -48,6 +50,19 @@ 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) + if err != nil { + _ = uc.blobs.Delete(ctx, art.StorageKey) + return nil, err + } + if !current.IsLeaseHeldBy(in.WorkerID, in.Attempt, uc.clk.Now()) { + _ = 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) diff --git a/coordinator/internal/usecase/task.go b/coordinator/internal/usecase/task.go index 7d01b25..131cffe 100644 --- a/coordinator/internal/usecase/task.go +++ b/coordinator/internal/usecase/task.go @@ -144,7 +144,7 @@ func (uc *CompleteTask) Execute(ctx context.Context, in CompleteTaskInput) (*dom } // Rule 10: never trust a worker-supplied artifact reference. The result // must be an artifact the coordinator itself stored for *this* task. - if err := uc.verifyResultArtifact(ctx, in.TaskID, in.ResultArtifactID); err != nil { + if err := uc.verifyResultArtifact(ctx, in.TaskID, in.Attempt, in.ResultArtifactID); err != nil { return err } now := uc.clock.Now() @@ -176,12 +176,12 @@ func (uc *CompleteTask) Execute(ctx context.Context, in CompleteTaskInput) (*dom // verifyResultArtifact enforces that the referenced artifact was stored by the // coordinator for this exact task. It stops a worker from completing task B with // an artifact it uploaded for task A, and from naming an id that isn't a result. -func (uc *CompleteTask) verifyResultArtifact(ctx context.Context, taskID, artifactID uuid.UUID) error { +func (uc *CompleteTask) verifyResultArtifact(ctx context.Context, taskID uuid.UUID, attempt int, artifactID uuid.UUID) error { art, err := uc.artifacts.Get(ctx, artifactID) if err != nil { return err } - if art.TaskID == nil || *art.TaskID != taskID || art.Kind != domain.ArtifactPartialResult { + if art.TaskID == nil || *art.TaskID != taskID || art.Attempt == nil || *art.Attempt != attempt || art.Kind != domain.ArtifactPartialResult { return domain.ErrResultConflict } return nil diff --git a/coordinator/internal/usecase/usecase_test.go b/coordinator/internal/usecase/usecase_test.go index 254a8ef..4786b53 100644 --- a/coordinator/internal/usecase/usecase_test.go +++ b/coordinator/internal/usecase/usecase_test.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "io" "strings" "testing" "time" @@ -19,6 +20,17 @@ var ctx = context.Background() const lease = 2 * time.Minute +type expiringBlobStore struct { + *memstore.BlobStore + clock *memstore.Clock +} + +func (s expiringBlobStore) Put(ctx context.Context, key string, body io.Reader) (string, int64, error) { + sum, size, err := s.BlobStore.Put(ctx, key, body) + s.clock.Advance(lease + time.Second) + return sum, size, err +} + // harness wires every use case to in-memory stores so orchestration can be // tested without a database. type harness struct { @@ -239,6 +251,46 @@ func TestCompleteRejectsForeignArtifact(t *testing.T) { } } +func TestCompleteRejectsArtifactFromExpiredAttempt(t *testing.T) { + h := newHarness() + h.seedJob(t, "w", 1) + taskID, attemptOne := h.leaseOne(t, "w1", "w") + staleArtifact := h.uploadResult(t, taskID, "w1", attemptOne) + + h.clk.Advance(lease + time.Second) + if _, err := h.expire.Execute(ctx); err != nil { + t.Fatalf("expire lease: %v", err) + } + _, attemptTwo := h.leaseOne(t, "w2", "w") + if attemptTwo != attemptOne+1 { + t.Fatalf("attempt = %d, want %d", attemptTwo, attemptOne+1) + } + + _, err := h.complete.Execute(ctx, usecase.CompleteTaskInput{ + TaskID: taskID, WorkerID: "w2", Attempt: attemptTwo, ResultArtifactID: staleArtifact, + }) + if !errors.Is(err, domain.ErrResultConflict) { + t.Errorf("stale-attempt artifact: err = %v, want ErrResultConflict", err) + } +} + +func TestUploadRejectsLeaseThatExpiresDuringStreaming(t *testing.T) { + h := newHarness() + 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, + ) + + _, err := h.uploadArt.Execute(ctx, usecase.UploadArtifactInput{ + TaskID: taskID, WorkerID: "w1", Attempt: attempt, + Filename: "result.csv", ContentType: "text/csv", Body: strings.NewReader("result"), + }) + if !errors.Is(err, domain.ErrLeaseConflict) { + t.Errorf("upload after lease expiry: err = %v, want ErrLeaseConflict", err) + } +} + func TestCompleteIsIdempotentOnReplay(t *testing.T) { h := newHarness() h.seedJob(t, "w", 1) diff --git a/coordinator/migrations/0008_artifact_attempt.down.sql b/coordinator/migrations/0008_artifact_attempt.down.sql new file mode 100644 index 0000000..f0731bb --- /dev/null +++ b/coordinator/migrations/0008_artifact_attempt.down.sql @@ -0,0 +1,7 @@ +BEGIN; + +ALTER TABLE artifacts DROP CONSTRAINT IF EXISTS ck_artifact_attempt_positive; +ALTER TABLE artifacts DROP CONSTRAINT IF EXISTS ck_partial_result_attempt; +ALTER TABLE artifacts DROP COLUMN IF EXISTS attempt; + +COMMIT; diff --git a/coordinator/migrations/0008_artifact_attempt.up.sql b/coordinator/migrations/0008_artifact_attempt.up.sql new file mode 100644 index 0000000..a0edae7 --- /dev/null +++ b/coordinator/migrations/0008_artifact_attempt.up.sql @@ -0,0 +1,34 @@ +BEGIN; + +-- A partial result belongs to the lease attempt that uploaded it. Without this +-- binding a worker holding a later retry could complete a task with stale bytes +-- uploaded by an expired attempt of that same task. +ALTER TABLE artifacts ADD COLUMN attempt integer; + +-- A completed task never gets a later lease, so its current attempt is also +-- the attempt that produced the stored result. +UPDATE artifacts AS a +SET attempt = t.attempt +FROM tasks AS t +WHERE a.task_id = t.id + AND a.kind = 'partial_result'::artifact_kind + AND t.status = 'completed'::task_status + AND a.attempt IS NULL; + +-- For unfinished tasks the old schema cannot tell which attempt uploaded a +-- partial result. Keeping it would let a later retry claim stale bytes, so the +-- worker must upload again. Blob garbage is harmless and follows the existing +-- coordinator-owned storage cleanup policy. +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.attempt IS NULL; + +ALTER TABLE artifacts ADD CONSTRAINT ck_partial_result_attempt + CHECK (kind <> 'partial_result'::artifact_kind OR attempt IS NOT NULL); +ALTER TABLE artifacts ADD CONSTRAINT ck_artifact_attempt_positive + CHECK (attempt IS NULL OR attempt > 0); + +COMMIT;