From 7547a30bde07d99fe16168074e4175b22c39b742 Mon Sep 17 00:00:00 2001 From: Emil Date: Thu, 23 Jul 2026 23:14:25 +0300 Subject: [PATCH] Add job cancellation and dataset row limit --- coordinator/cmd/coordinator/main.go | 1 + coordinator/internal/chunk/tsv.go | 13 ++++ coordinator/internal/chunk/tsv_test.go | 16 +++++ coordinator/internal/domain/errors.go | 19 +++--- coordinator/internal/domain/job.go | 15 +++-- coordinator/internal/domain/job_test.go | 1 + coordinator/internal/domain/task.go | 18 ++++++ coordinator/internal/domain/task_test.go | 17 ++++++ coordinator/internal/memstore/memstore.go | 12 ++++ .../storage/postgres/integration_test.go | 21 +++++++ .../internal/storage/postgres/task_repo.go | 23 ++++++++ coordinator/internal/transport/http/dto.go | 30 +++++----- coordinator/internal/transport/http/errors.go | 3 +- .../internal/transport/http/handlers.go | 39 +++++++++++- coordinator/internal/transport/http/server.go | 3 + .../internal/transport/http/server_test.go | 59 +++++++++++++++++++ .../transport/http/templates/dashboard.html | 2 +- .../transport/http/templates/job.html | 9 +-- .../transport/http/templates/new-job.html | 3 +- coordinator/internal/transport/http/ui.go | 15 ++++- .../internal/transport/http/ui_test.go | 4 +- coordinator/internal/usecase/dto.go | 9 ++- coordinator/internal/usecase/job.go | 46 ++++++++++++++- coordinator/internal/usecase/ports.go | 4 ++ coordinator/internal/usecase/ui.go | 5 +- coordinator/internal/usecase/upload.go | 2 +- coordinator/internal/usecase/usecase_test.go | 40 +++++++++++++ docs/api-contract.md | 26 +++++++- docs/openapi.yaml | 34 ++++++++++- docs/web-interface-plan.md | 3 + 30 files changed, 438 insertions(+), 54 deletions(-) diff --git a/coordinator/cmd/coordinator/main.go b/coordinator/cmd/coordinator/main.go index be84e11..1134257 100644 --- a/coordinator/cmd/coordinator/main.go +++ b/coordinator/cmd/coordinator/main.go @@ -77,6 +77,7 @@ func run() error { 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), DownloadArtifact: usecase.NewDownloadArtifact(artifactRepo, blobStore), GetTaskInput: usecase.NewGetTaskInput(taskRepo, artifactRepo, blobStore), diff --git a/coordinator/internal/chunk/tsv.go b/coordinator/internal/chunk/tsv.go index 6b968eb..f020810 100644 --- a/coordinator/internal/chunk/tsv.go +++ b/coordinator/internal/chunk/tsv.go @@ -26,9 +26,19 @@ 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) +} + +// 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 { if rowsPerShard <= 0 { return fmt.Errorf("rowsPerShard must be positive, got %d", rowsPerShard) } + if maxRows < 0 { + return fmt.Errorf("maxRows must be non-negative, got %d", maxRows) + } sc := bufio.NewScanner(r) // Allow long lines: a SMILES row can be far wider than bufio's 64 KB default. @@ -73,6 +83,9 @@ func SplitTSV(r io.Reader, rowsPerShard int, emit func(index int, shard io.Reade return err } } + if maxRows > 0 && index*rowsPerShard+rows == maxRows { + break + } } if err := sc.Err(); err != nil { return fmt.Errorf("read rows: %w", err) diff --git a/coordinator/internal/chunk/tsv_test.go b/coordinator/internal/chunk/tsv_test.go index ba12917..585b35a 100644 --- a/coordinator/internal/chunk/tsv_test.go +++ b/coordinator/internal/chunk/tsv_test.go @@ -103,6 +103,22 @@ func TestSplitSingleShardWhenSizeExceedsRows(t *testing.T) { } } +func TestSplitLimitUsesOnlyLeadingDataRows(t *testing.T) { + input := "h\nr1\nr2\nr3\nr4\nr5\n" + var shards []string + err := SplitTSVLimit(strings.NewReader(input), 2, 3, func(_ int, shard io.Reader) error { + b, _ := io.ReadAll(shard) + shards = append(shards, string(b)) + return nil + }) + if err != nil { + t.Fatal(err) + } + if got, want := strings.Join(shards, ""), "h\nr1\nr2\nh\nr3\n"; got != want { + t.Errorf("limited shards = %q, want %q", got, want) + } +} + // 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/domain/errors.go b/coordinator/internal/domain/errors.go index d64b1f6..fca2979 100644 --- a/coordinator/internal/domain/errors.go +++ b/coordinator/internal/domain/errors.go @@ -8,13 +8,14 @@ import "errors" // // Always compare with errors.Is — outer layers may wrap these with %w. var ( - ErrJobNotFound = errors.New("job not found") - ErrTaskNotFound = errors.New("task not found") - ErrWorkerNotFound = errors.New("worker not found") - ErrArtifactNotFound = errors.New("artifact not found") - ErrLeaseConflict = errors.New("task leased to another worker") - ErrStaleAttempt = errors.New("attempt does not match lease") - ErrResultConflict = errors.New("different result already recorded") - ErrInvalidInput = errors.New("invalid input") - ErrTaskNotLeased = errors.New("task is not currently leased") + ErrJobNotFound = errors.New("job not found") + ErrTaskNotFound = errors.New("task not found") + ErrWorkerNotFound = errors.New("worker not found") + ErrArtifactNotFound = errors.New("artifact not found") + ErrJobNotCancellable = errors.New("job cannot be cancelled") + ErrLeaseConflict = errors.New("task leased to another worker") + ErrStaleAttempt = errors.New("attempt does not match lease") + ErrResultConflict = errors.New("different result already recorded") + ErrInvalidInput = errors.New("invalid input") + ErrTaskNotLeased = errors.New("task is not currently leased") ) diff --git a/coordinator/internal/domain/job.go b/coordinator/internal/domain/job.go index 0aa4827..7d6be3d 100644 --- a/coordinator/internal/domain/job.go +++ b/coordinator/internal/domain/job.go @@ -99,18 +99,21 @@ func NewJobWithTasks(workload, inputURI string, params map[string]any, // JobProgress is the aggregate view of a job and the state of its tasks. type JobProgress struct { - Job Job - Total int - Pending int - Leased int - Done int - Failed int + Job Job + Total int + Pending int + Leased int + Done int + Failed int + Cancelled int } // DeriveStatus computes what the job's status should be from its task counts, // so the rule lives here rather than in a SQL trigger or a handler. func (p JobProgress) DeriveStatus() JobStatus { switch { + case p.Job.Status == JobCancelled: + return JobCancelled case p.Total == 0: return JobPending case p.Done == p.Total: diff --git a/coordinator/internal/domain/job_test.go b/coordinator/internal/domain/job_test.go index 0b58902..54a221d 100644 --- a/coordinator/internal/domain/job_test.go +++ b/coordinator/internal/domain/job_test.go @@ -81,6 +81,7 @@ func TestDeriveStatus(t *testing.T) { {"all done", JobProgress{Total: 3, Done: 3}, JobCompleted}, {"done and failed", JobProgress{Total: 3, Done: 2, Failed: 1}, JobFailed}, {"failed but work remains", JobProgress{Total: 3, Pending: 1, Failed: 2}, JobRunning}, + {"cancelled job wins over task histogram", JobProgress{Job: Job{Status: JobCancelled}, Total: 3, Done: 1, Cancelled: 2}, JobCancelled}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { diff --git a/coordinator/internal/domain/task.go b/coordinator/internal/domain/task.go index 306fed5..7864660 100644 --- a/coordinator/internal/domain/task.go +++ b/coordinator/internal/domain/task.go @@ -259,6 +259,24 @@ func (t *Task) ExpireLease(now time.Time) { t.CompletedAt = &now } +// Cancel prevents any further worker transition for a task that has not +// reached a terminal result. A cancelled lease deliberately becomes invalid: +// a worker still running locally must not upload or complete after its job was +// stopped by the operator. +func (t *Task) Cancel(now time.Time) bool { + if t.Status == TaskCompleted || t.Status == TaskFailed || t.Status == TaskCancelled { + return false + } + t.Status = TaskCancelled + t.LeaseOwner = nil + t.LeaseExpiresAt = nil + t.ErrorCode = nil + t.ErrorMessage = nil + t.CompletedAt = &now + t.Version++ + return true +} + // ClaimedTask is the worker-facing projection of a leased task. Input is either // an external URI or a coordinator-stored shard (InputArtifactID set); the // transport turns the latter into a coordinator download URL. diff --git a/coordinator/internal/domain/task_test.go b/coordinator/internal/domain/task_test.go index 9c2ff98..a4c7862 100644 --- a/coordinator/internal/domain/task_test.go +++ b/coordinator/internal/domain/task_test.go @@ -168,6 +168,23 @@ func TestExpireLeaseIgnoresUnleasedTasks(t *testing.T) { } } +func TestCancelInvalidatesLeaseButPreservesTerminalTask(t *testing.T) { + task := leasedTask(1, 3) + if !task.Cancel(testNow) { + t.Fatal("leased task should be cancelled") + } + if task.Status != TaskCancelled || task.LeaseOwner != nil || task.LeaseExpiresAt != nil { + t.Errorf("cancelled task = %+v", task) + } + if task.Cancel(testLater) { + t.Error("cancelled task must not be changed twice") + } + completed := &Task{Status: TaskCompleted} + if completed.Cancel(testNow) { + t.Error("completed task must remain terminal") + } +} + func TestFirstHeartbeatMovesLeasedToRunning(t *testing.T) { task := leasedTask(1, 3) until := testLater.Add(time.Hour) diff --git a/coordinator/internal/memstore/memstore.go b/coordinator/internal/memstore/memstore.go index 7916951..44db68e 100644 --- a/coordinator/internal/memstore/memstore.go +++ b/coordinator/internal/memstore/memstore.go @@ -148,6 +148,18 @@ func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[doma return counts, nil } +func (r *TaskRepo) CancelByJob(_ context.Context, jobID uuid.UUID, now time.Time) (int64, error) { + r.mu.Lock() + defer r.mu.Unlock() + var cancelled int64 + for _, task := range r.tasks { + if task.JobID == jobID && task.Cancel(now) { + cancelled++ + } + } + return cancelled, nil +} + func (r *TaskRepo) ExpireLeases(ctx context.Context, now time.Time) (int64, error) { r.mu.Lock() defer r.mu.Unlock() diff --git a/coordinator/internal/storage/postgres/integration_test.go b/coordinator/internal/storage/postgres/integration_test.go index cb0c610..6b3395a 100644 --- a/coordinator/internal/storage/postgres/integration_test.go +++ b/coordinator/internal/storage/postgres/integration_test.go @@ -206,6 +206,27 @@ func TestClaimNextReturnsNilOnEmptyQueue(t *testing.T) { } } +func TestCancelJobCancelsEveryUnfinishedTask(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + job, _ := seedJob(t, pool, 3) + clk := fixedClock{now: time.Now().UTC()} + uc := usecase.NewCancelJob(NewJobRepo(pool), NewTaskRepo(pool), NewTxManager(pool), clk) + + cancelled, err := uc.Execute(ctx, job.ID) + if err != nil || cancelled != 3 { + t.Fatalf("cancel = (%d, %v), want (3, nil)", cancelled, err) + } + stored, err := NewJobRepo(pool).Get(ctx, job.ID) + if err != nil || stored.Status != domain.JobCancelled { + t.Fatalf("job after cancel = (%+v, %v)", stored, err) + } + counts, err := NewTaskRepo(pool).CountByStatus(ctx, job.ID) + if err != nil || counts[domain.TaskCancelled] != 3 { + t.Fatalf("cancelled tasks = %d, err = %v", counts[domain.TaskCancelled], err) + } +} + func TestUpdateRejectsStaleVersion(t *testing.T) { pool := testPool(t) ctx := context.Background() diff --git a/coordinator/internal/storage/postgres/task_repo.go b/coordinator/internal/storage/postgres/task_repo.go index 2a4480f..18b2646 100644 --- a/coordinator/internal/storage/postgres/task_repo.go +++ b/coordinator/internal/storage/postgres/task_repo.go @@ -295,6 +295,29 @@ func (r *TaskRepo) CountByStatus(ctx context.Context, jobID uuid.UUID) (map[doma return counts, rows.Err() } +// cancelByJobSQL mirrors domain.Task.Cancel in one set-based update. It runs in +// the same transaction as the job-status update, so no claimable shard remains +// after an operator receives a successful cancellation response. +const cancelByJobSQL = ` +UPDATE tasks +SET status = 'cancelled'::task_status, + lease_owner = NULL, + lease_expires_at = NULL, + error_code = NULL, + error_message = NULL, + completed_at = $2, + version = version + 1 +WHERE job_id = $1 + AND status IN ('pending','leased','running')` + +func (r *TaskRepo) CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error) { + tag, err := conn(ctx, r.pool).Exec(ctx, cancelByJobSQL, jobID, now) + if err != nil { + return 0, err + } + return tag.RowsAffected(), nil +} + // expireLeasesSQL applies the lease-expiry rule set-based, mirroring // domain.Task.ExpireLease: requeue while attempts remain, otherwise fail. // diff --git a/coordinator/internal/transport/http/dto.go b/coordinator/internal/transport/http/dto.go index c652b25..f944482 100644 --- a/coordinator/internal/transport/http/dto.go +++ b/coordinator/internal/transport/http/dto.go @@ -111,13 +111,14 @@ type uploadJobResponse struct { } type jobProgressResponse struct { - ID uuid.UUID `json:"id"` - Status string `json:"status"` - Total int `json:"total"` - Pending int `json:"pending"` - Leased int `json:"leased"` - Done int `json:"completed"` - Failed int `json:"failed"` + ID uuid.UUID `json:"id"` + Status string `json:"status"` + Total int `json:"total"` + Pending int `json:"pending"` + Leased int `json:"leased"` + Done int `json:"completed"` + Failed int `json:"failed"` + Cancelled int `json:"cancelled"` } type uploadArtifactResponse struct { @@ -153,12 +154,13 @@ func toClaimedTaskResponse(c domain.ClaimedTask) claimedTaskResponse { func toJobProgressResponse(p domain.JobProgress) jobProgressResponse { return jobProgressResponse{ - ID: p.Job.ID, - Status: string(p.DeriveStatus()), - Total: p.Total, - Pending: p.Pending, - Leased: p.Leased, - Done: p.Done, - Failed: p.Failed, + ID: p.Job.ID, + Status: string(p.DeriveStatus()), + Total: p.Total, + Pending: p.Pending, + Leased: p.Leased, + Done: p.Done, + Failed: p.Failed, + Cancelled: p.Cancelled, } } diff --git a/coordinator/internal/transport/http/errors.go b/coordinator/internal/transport/http/errors.go index 358704e..ae14782 100644 --- a/coordinator/internal/transport/http/errors.go +++ b/coordinator/internal/transport/http/errors.go @@ -50,7 +50,8 @@ func (s *Server) writeError(w http.ResponseWriter, r *http.Request, err error) { case errors.Is(err, domain.ErrLeaseConflict), errors.Is(err, domain.ErrStaleAttempt), errors.Is(err, domain.ErrResultConflict), - errors.Is(err, domain.ErrTaskNotLeased): + errors.Is(err, domain.ErrTaskNotLeased), + errors.Is(err, domain.ErrJobNotCancellable): status = http.StatusConflict case errors.Is(err, usecase.ErrNotImplemented): status = http.StatusNotImplemented diff --git a/coordinator/internal/transport/http/handlers.go b/coordinator/internal/transport/http/handlers.go index df61461..4bbf5c8 100644 --- a/coordinator/internal/transport/http/handlers.go +++ b/coordinator/internal/transport/http/handlers.go @@ -182,7 +182,7 @@ func (s *Server) handleFailure(w http.ResponseWriter, r *http.Request) { const defaultChunkRows = 1000 // handleUploadDataset accepts a multipart submission — the dataset file plus the -// workload/parameters/chunk_rows fields — and hands the file, streamed, to the +// workload/parameters/chunk_rows/max_rows fields — and hands the file, streamed, to the // chunker. The text fields MUST precede the file part: the file is streamed, not // buffered, so by the time it arrives the other fields are already parsed. func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) { @@ -197,11 +197,13 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) { workload string params map[string]any rows = defaultChunkRows + maxRows int result usecase.SubmitDatasetResult gotDataset bool gotWorkload bool gotParams bool gotRows bool + gotMaxRows bool ) for { @@ -249,6 +251,19 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) { } rows = n gotRows = true + case "max_rows": + if gotDataset || gotMaxRows { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + b, _ := io.ReadAll(io.LimitReader(part, 32)) + n, err := strconv.Atoi(strings.TrimSpace(string(b))) + if err != nil || n < 1 { + s.writeError(w, r, domain.ErrInvalidInput) + return + } + maxRows = n + gotMaxRows = true case "file", "dataset": if gotDataset || workload == "" { s.writeError(w, r, domain.ErrInvalidInput) @@ -262,6 +277,7 @@ func (s *Server) handleUploadDataset(w http.ResponseWriter, r *http.Request) { Workload: workload, Parameters: params, RowsPerShard: rows, + MaxRows: maxRows, Filename: filename, ContentType: part.Header.Get("Content-Type"), Body: part, @@ -379,6 +395,27 @@ func (s *Server) handleGetJob(w http.ResponseWriter, r *http.Request) { writeJSON(w, http.StatusOK, toJobProgressResponse(progress)) } +// handleCancelJob stops all non-terminal shards for an operator-requested job. +// It is available to both the bearer API and the separately authenticated UI. +func (s *Server) handleCancelJob(w http.ResponseWriter, r *http.Request) { + ctx, cancel := s.reqCtx(r) + defer cancel() + jobID, ok := s.pathUUID(w, r, "job_id") + if !ok { + return + } + cancelled, err := s.uc.CancelJob.Execute(ctx, jobID) + if err != nil { + s.writeError(w, r, err) + return + } + writeJSON(w, http.StatusOK, map[string]any{ + "job_id": jobID, + "status": domain.JobCancelled, + "cancelled_tasks": cancelled, + }) +} + // --- helpers --- func (s *Server) reqCtx(r *http.Request) (context.Context, context.CancelFunc) { diff --git a/coordinator/internal/transport/http/server.go b/coordinator/internal/transport/http/server.go index 1142836..00e80d8 100644 --- a/coordinator/internal/transport/http/server.go +++ b/coordinator/internal/transport/http/server.go @@ -24,6 +24,7 @@ type UseCases struct { CompleteTask *usecase.CompleteTask FailTask *usecase.FailTask GetJobStatus *usecase.GetJobStatus + CancelJob *usecase.CancelJob UploadArtifact *usecase.UploadArtifact DownloadArtifact *usecase.DownloadArtifact GetTaskInput *usecase.GetTaskInput @@ -61,6 +62,7 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler { protected.HandleFunc("POST /jobs", s.handleCreateJob) protected.HandleFunc("POST /jobs/upload", s.handleUploadDataset) protected.HandleFunc("GET /jobs/{job_id}", s.handleGetJob) + protected.HandleFunc("POST /jobs/{job_id}/cancel", s.handleCancelJob) protected.HandleFunc("POST /tasks/claim", s.handleClaim) protected.HandleFunc("GET /tasks/{task_id}/input", s.handleGetTaskInput) protected.HandleFunc("POST /tasks/{task_id}/heartbeat", s.handleHeartbeat) @@ -77,6 +79,7 @@ func (s *Server) Handler(token string, uiToken ...string) http.Handler { ui.HandleFunc("GET /ui/jobs/new", s.handleUINewJob) ui.HandleFunc("GET /ui/jobs/{job_id}", s.handleUIJob) ui.HandleFunc("GET /ui/api/jobs/{job_id}", s.handleUIJobJSON) + ui.HandleFunc("POST /ui/api/jobs/{job_id}/cancel", s.handleCancelJob) ui.HandleFunc("POST /ui/api/jobs/upload", s.handleUploadDataset) ui.HandleFunc("GET /ui/jobs/{job_id}/artifacts/{artifact_id}", s.handleUIArtifactDownload) mux.Handle("/ui", chain(ui, withRequestID, withAccessLog(s.log), withBasicAuth(uiToken[0]), withSameOrigin)) diff --git a/coordinator/internal/transport/http/server_test.go b/coordinator/internal/transport/http/server_test.go index 01c299f..261197c 100644 --- a/coordinator/internal/transport/http/server_test.go +++ b/coordinator/internal/transport/http/server_test.go @@ -51,6 +51,7 @@ func newEnvWithUIToken(t *testing.T, ready func(context.Context) error, configur 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), DownloadArtifact: usecase.NewDownloadArtifact(arts, blobs), GetTaskInput: usecase.NewGetTaskInput(tasks, arts, blobs), @@ -200,6 +201,39 @@ func TestUIUploadDatasetCreatesJob(t *testing.T) { } } +func TestCancelJobStopsUnfinishedTasks(t *testing.T) { + e := newEnv(t, healthy) + code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"sha"},{"chunk_index":1,"input_uri":"s3://c1","input_sha256":"sha"}]}`) + if code != http.StatusCreated { + t.Fatalf("create: %d", code) + } + jobID := job["id"].(string) + if code, body := e.do(t, "POST", "/jobs/"+jobID+"/cancel", ""); code != http.StatusOK || body["cancelled_tasks"].(float64) != 2 { + t.Fatalf("cancel = (%d, %v)", code, body) + } + if code, progress := e.do(t, "GET", "/jobs/"+jobID, ""); code != http.StatusOK || progress["status"] != "cancelled" || progress["cancelled"].(float64) != 2 { + t.Fatalf("cancelled job progress = (%d, %v)", code, progress) + } +} + +func TestUICancelJobUsesOperatorCredential(t *testing.T) { + e := newEnv(t, healthy) + code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c0","input_sha256":"sha"}]}`) + if code != http.StatusCreated { + t.Fatalf("create: %d", code) + } + req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/ui/api/jobs/"+job["id"].(string)+"/cancel", nil) + req.SetBasicAuth("operator", uiToken) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("UI cancel = %d", resp.StatusCode) + } +} + func TestUIJobAndArtifactAreScopedToTheirJob(t *testing.T) { e := newEnv(t, healthy) code, job := e.do(t, "POST", "/jobs", `{"workload":"w","input_uri":"s3://in","chunks":[{"chunk_index":0,"input_uri":"s3://c","input_sha256":"sha"}]}`) @@ -392,6 +426,31 @@ func TestUploadDatasetChunksAndServesInput(t *testing.T) { } } +func TestUploadDatasetLimitsRows(t *testing.T) { + e := newEnv(t, healthy) + var buf bytes.Buffer + mw := multipart.NewWriter(&buf) + _ = mw.WriteField("workload", "w") + _ = 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")) + _ = mw.Close() + req, _ := http.NewRequestWithContext(context.Background(), "POST", e.ts.URL+"/jobs/upload", &buf) + req.Header.Set("Authorization", "Bearer "+token) + req.Header.Set("Content-Type", mw.FormDataContentType()) + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + var result map[string]any + _ = json.NewDecoder(resp.Body).Decode(&result) + if resp.StatusCode != http.StatusCreated || result["task_count"].(float64) != 2 { + t.Fatalf("limited upload = (%d, %v)", resp.StatusCode, result) + } +} + func TestErrorMappings(t *testing.T) { e := newEnv(t, healthy) zero := "00000000-0000-0000-0000-000000000000" diff --git a/coordinator/internal/transport/http/templates/dashboard.html b/coordinator/internal/transport/http/templates/dashboard.html index f4769ee..5e24f90 100644 --- a/coordinator/internal/transport/http/templates/dashboard.html +++ b/coordinator/internal/transport/http/templates/dashboard.html @@ -14,7 +14,7 @@

Local coordinator

SciMesh operator dashboard

See where a computation is and what should happen next.

Start a check
This screen currently diagnoses shard jobs.Workers upload partial CSVs to the coordinator. Until a reducer is implemented, those files are not one final scientific result.
1. Upload TSVThe coordinator splits the file into shard tasks.
2. Wait for a workerA worker claims a shard, calculates similarity, and returns a CSV.
3. Inspect artifactsDownload a partial result from the job page.

Recent jobs

-
{{range .Jobs}}{{else}}{{end}}
ComputationStateProgressCreated
{{workloadLabel .Workload}}
Open job details		{{statusLabel .Status}}
{{statusHint .Status}}		{{.Completed}} / {{.Total}} complete{{if gt .Failed 0}} · failed: {{.Failed}}{{end}}
{{time .CreatedAt}}
No jobs yet.
Click “Start a check”, upload a small TSV, and leave a worker running.
+
{{range .Jobs}}{{else}}{{end}}
ComputationStateProgressCreated
{{workloadLabel .Workload}}
Open job details		{{statusLabel .Status}}
{{statusHint .Status}}		{{.Completed}} / {{.Total}} complete{{if gt .Failed 0}} · failed: {{.Failed}}{{end}}{{if gt .Cancelled 0}} · stopped: {{.Cancelled}}{{end}}
{{time .CreatedAt}}
No jobs yet.
Click “Start a check”, upload a small TSV, and leave a worker running.

Workers

{{range .Workers}}
{{.Name}}
{{.ID}}
{{workerStatusLabel .Status}}
{{range .Capabilities}}{{.}} {{end}}
Last signal
{{time .LastHeartbeatAt}}
{{else}}
No worker is registered yet.
Run scimesh-worker with the coordinator URL and worker token.
{{end}}
diff --git a/coordinator/internal/transport/http/templates/job.html b/coordinator/internal/transport/http/templates/job.html index 3ce8e6b..f983e2d 100644 --- a/coordinator/internal/transport/http/templates/job.html +++ b/coordinator/internal/transport/http/templates/job.html @@ -6,13 +6,13 @@ SciMesh job
← Back to jobs

{{workloadLabel .Workload}}

Execution progress

-
{{statusLabel .Status}}

{{statusHint .Status}}

Summary refreshes automatically every two seconds.

{{.Completed}} of {{.Total}} tasks complete

{{.Total}}total shards
{{.Completed}}complete
{{.Pending}}waiting
{{add .Leased .Running}}with workers
{{.Failed}}failed
Technical details

Job ID: {{.ID}}
Workload: {{.Workload}}
Created: {{time .CreatedAt}}

+
{{statusLabel .Status}}

{{statusHint .Status}}

{{if cancellable .Status}}This cancels every shard that is not finished yet.{{else}}Summary refreshes automatically every two seconds.{{end}}

{{.Completed}} of {{.Total}} tasks complete

{{.Total}}total shards
{{.Completed}}complete
{{.Pending}}waiting
{{add .Leased .Running}}with workers
{{.Failed}}failed
{{.Cancelled}}stopped
Technical details

Job ID: {{.ID}}
Workload: {{.Workload}}
Created: {{time .CreatedAt}}

What can be downloaded now?
partial_result files come from individual shards. They are useful for checking the pipeline, but are not a merged final CSV because the reducer is not implemented yet.

Shard tasks

If a task fails, its code and message appear here. Refresh the page to update the detailed rows.

{{range .Tasks}}{{else}}{{end}}
ShardStateAttemptWorker / leaseError
#{{.ChunkIndex}}{{statusLabel .Status}}{{.Attempt}} / {{.MaxAttempts}}{{if .LeaseOwner}}{{.LeaseOwner}}{{if .LeaseExpiresAt}}
until {{time .LeaseExpiresAt}}{{end}}{{else}}—{{end}}		{{if .ErrorCode}}{{taskErrorLabel .ErrorCode}}
{{taskErrorHint .ErrorCode}}{{else}}—{{end}}
No tasks have appeared yet.
@@ -20,8 +20,9 @@
{{range .Artifacts}}{{else}}{{end}}
TypeFileSizeIntegrity check
{{if .Diagnostic}}Partial result
diagnostic{{else}}{{.Kind}}{{end}}		{{.Filename}}{{bytes .SizeBytes}}{{.SHA256}}{{if .Downloadable}}Download CSV{{else}}Unavailable{{end}}
No artifacts yet. The worker uploads a CSV after it completes a shard.
diff --git a/coordinator/internal/transport/http/templates/new-job.html b/coordinator/internal/transport/http/templates/new-job.html index 5a28bf2..365df86 100644 --- a/coordinator/internal/transport/http/templates/new-job.html +++ b/coordinator/internal/transport/http/templates/new-job.html @@ -18,12 +18,13 @@

CCO is ethanol. For gefitinib, use its SMILES here or the local CLI with --query-id.

This is the top-k within each shard, not a global top-k for the whole dataset yet.

Fewer rows mean more tasks and more visible progress; more rows mean fewer, longer tasks.

+

Useful for a quick check of a large TSV. The coordinator creates shards from only the first N data rows; it still stores the original upload.

diff --git a/coordinator/internal/transport/http/ui.go b/coordinator/internal/transport/http/ui.go index f7d233e..ddf1168 100644 --- a/coordinator/internal/transport/http/ui.go +++ b/coordinator/internal/transport/http/ui.go @@ -28,6 +28,7 @@ var uiTemplates = template.Must(template.New("ui").Funcs(template.FuncMap{ "workerStatusLabel": uiWorkerStatusLabel, "workloadLabel": uiWorkloadLabel, "progressPercent": uiProgressPercent, + "cancellable": uiCancellable, "bytes": uiBytes, "add": func(a, b int) int { return a + b }, }).ParseFS(uiFiles, "templates/*.html")) @@ -51,6 +52,8 @@ func uiStatusLabel(status string) string { return "Tasks complete" case "failed": return "Needs attention" + case "cancelled": + return "Stopped" default: return status } @@ -68,6 +71,8 @@ func uiStatusHint(status string) string { return "Every shard task is complete. Files below are still partial results." case "failed": return "One or more shard tasks failed. Open the task list below for details." + case "cancelled": + return "The operator stopped this job. No new shards can be claimed." default: return "Status reported by the coordinator." } @@ -79,6 +84,8 @@ func uiStatusClass(status string) string { return "success" case "failed": return "danger" + case "cancelled": + return "waiting" case "running", "leased": return "active" default: @@ -145,11 +152,15 @@ func uiWorkloadLabel(workload string) string { } } -func uiProgressPercent(completed, failed, total int) int { +func uiCancellable(status string) bool { + return status == "pending" || status == "running" +} + +func uiProgressPercent(completed, failed, cancelled, total int) int { if total <= 0 { return 0 } - percent := (completed + failed) * 100 / total + percent := (completed + failed + cancelled) * 100 / total if percent > 100 { return 100 } diff --git a/coordinator/internal/transport/http/ui_test.go b/coordinator/internal/transport/http/ui_test.go index 2d3b903..a6f82cb 100644 --- a/coordinator/internal/transport/http/ui_test.go +++ b/coordinator/internal/transport/http/ui_test.go @@ -26,10 +26,10 @@ func TestUIStatusPresentation(t *testing.T) { } func TestUIProgressPercent(t *testing.T) { - if got := uiProgressPercent(3, 1, 8); got != 50 { + if got := uiProgressPercent(3, 1, 0, 8); got != 50 { t.Errorf("progress = %d, want 50", got) } - if got := uiProgressPercent(1, 1, 0); got != 0 { + if got := uiProgressPercent(1, 1, 0, 0); got != 0 { t.Errorf("empty progress = %d, want 0", got) } } diff --git a/coordinator/internal/usecase/dto.go b/coordinator/internal/usecase/dto.go index 6491cbf..17e0498 100644 --- a/coordinator/internal/usecase/dto.go +++ b/coordinator/internal/usecase/dto.go @@ -53,9 +53,12 @@ type SubmitDatasetInput struct { Workload string Parameters map[string]any RowsPerShard int - Filename string - ContentType string - Body io.Reader + // MaxRows limits how many data rows are turned into shards. Zero means the + // whole uploaded dataset; the input artifact itself remains stored intact. + MaxRows int + Filename string + ContentType string + Body io.Reader } type SubmitDatasetResult struct { diff --git a/coordinator/internal/usecase/job.go b/coordinator/internal/usecase/job.go index 28e6c6e..748f021 100644 --- a/coordinator/internal/usecase/job.go +++ b/coordinator/internal/usecase/job.go @@ -62,6 +62,45 @@ type GetJobStatus struct { tasks TaskRepository } +// --- CancelJob ----------------------------------------------------------- + +type CancelJob struct { + jobs JobRepository + tasks TaskRepository + tx TxManager + clock Clock +} + +func NewCancelJob(jobs JobRepository, tasks TaskRepository, tx TxManager, clock Clock) *CancelJob { + return &CancelJob{jobs: jobs, tasks: tasks, tx: tx, clock: clock} +} + +// Execute stops a job atomically. Completed and finally failed tasks are kept +// as historical evidence; all other tasks are cancelled, including leased and +// running ones. A repeated cancel of an already cancelled job is idempotent. +func (uc *CancelJob) Execute(ctx context.Context, jobID uuid.UUID) (int64, error) { + now := uc.clock.Now() + var cancelled int64 + err := uc.tx.WithinTx(ctx, func(ctx context.Context) error { + job, err := uc.jobs.Get(ctx, jobID) + if err != nil { + return err + } + if job.Status == domain.JobCancelled { + return nil + } + if job.Status == domain.JobCompleted || job.Status == domain.JobFailed { + return domain.ErrJobNotCancellable + } + cancelled, err = uc.tasks.CancelByJob(ctx, jobID, now) + if err != nil { + return err + } + return uc.jobs.UpdateStatus(ctx, jobID, domain.JobCancelled, &now) + }) + return cancelled, err +} + func NewGetJobStatus(jobs JobRepository, tasks TaskRepository) *GetJobStatus { return &GetJobStatus{jobs: jobs, tasks: tasks} } @@ -144,9 +183,10 @@ func progressFrom(job domain.Job, counts map[domain.TaskStatus]int) domain.JobPr Job: job, Pending: counts[domain.TaskPending], // Leased and running are both "in flight" for progress purposes. - Leased: counts[domain.TaskLeased] + counts[domain.TaskRunning], - Done: counts[domain.TaskCompleted], - Failed: counts[domain.TaskFailed], + Leased: counts[domain.TaskLeased] + counts[domain.TaskRunning], + Done: counts[domain.TaskCompleted], + Failed: counts[domain.TaskFailed], + Cancelled: counts[domain.TaskCancelled], } for _, n := range counts { p.Total += n diff --git a/coordinator/internal/usecase/ports.go b/coordinator/internal/usecase/ports.go index a85481b..7b56211 100644 --- a/coordinator/internal/usecase/ports.go +++ b/coordinator/internal/usecase/ports.go @@ -56,6 +56,10 @@ type TaskRepository interface { // CountByStatus aggregates a job's tasks for progress reporting. CountByStatus(ctx context.Context, jobID uuid.UUID) (map[domain.TaskStatus]int, error) + // CancelByJob marks every non-terminal task as cancelled and invalidates its + // lease. It returns how many tasks changed. + CancelByJob(ctx context.Context, jobID uuid.UUID, now time.Time) (int64, error) + // ExpireLeases applies the lease-expiry rule to every elapsed task and // reports how many were affected. ExpireLeases(ctx context.Context, now time.Time) (int64, error) diff --git a/coordinator/internal/usecase/ui.go b/coordinator/internal/usecase/ui.go index f250480..c7283c3 100644 --- a/coordinator/internal/usecase/ui.go +++ b/coordinator/internal/usecase/ui.go @@ -30,6 +30,7 @@ type JobCard struct { Running int `json:"running"` Completed int `json:"completed"` Failed int `json:"failed"` + Cancelled int `json:"cancelled"` } type TaskCard struct { @@ -166,9 +167,11 @@ func jobCard(job domain.Job, tasks []domain.Task) JobCard { c.Completed++ case domain.TaskFailed: c.Failed++ + case domain.TaskCancelled: + c.Cancelled++ } } - p := domain.JobProgress{Job: job, Total: c.Total, Pending: c.Pending, Leased: c.Leased + c.Running, Done: c.Completed, Failed: c.Failed} + p := domain.JobProgress{Job: job, Total: c.Total, Pending: c.Pending, Leased: c.Leased + c.Running, Done: c.Completed, Failed: c.Failed, Cancelled: c.Cancelled} c.Status = string(p.DeriveStatus()) return c } diff --git a/coordinator/internal/usecase/upload.go b/coordinator/internal/usecase/upload.go index f04f582..20fb899 100644 --- a/coordinator/internal/usecase/upload.go +++ b/coordinator/internal/usecase/upload.go @@ -64,7 +64,7 @@ func (uc *SubmitDataset) Execute(ctx context.Context, in SubmitDatasetInput) (Su cleanup() return SubmitDatasetResult{}, err } - splitErr := chunk.SplitTSV(rc, in.RowsPerShard, func(index int, shard io.Reader) error { + splitErr := chunk.SplitTSVLimit(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 { diff --git a/coordinator/internal/usecase/usecase_test.go b/coordinator/internal/usecase/usecase_test.go index 4786b53..68b94ac 100644 --- a/coordinator/internal/usecase/usecase_test.go +++ b/coordinator/internal/usecase/usecase_test.go @@ -54,6 +54,7 @@ type harness struct { downloadArt *usecase.DownloadArtifact getInput *usecase.GetTaskInput expire *usecase.ExpireLeases + cancel *usecase.CancelJob } func newHarness() *harness { @@ -79,6 +80,7 @@ func newHarness() *harness { 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.cancel = usecase.NewCancelJob(h.jobs, h.tasks, tx, h.clk) return h } @@ -468,6 +470,44 @@ 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" + res, err := h.submit.Execute(ctx, usecase.SubmitDatasetInput{ + Workload: "w", RowsPerShard: 2, MaxRows: 3, Filename: "chembl.tsv", + ContentType: "text/tab-separated-values", Body: strings.NewReader(tsv), + }) + if err != nil { + t.Fatal(err) + } + if res.TaskCount != 2 { + t.Fatalf("task_count = %d, want 2", res.TaskCount) + } +} + +func TestCancelJobInvalidatesClaimedAndPendingTasks(t *testing.T) { + h := newHarness() + jobID := h.seedJob(t, "w", 3) + claimed, err := h.claim.Execute(ctx, usecase.ClaimTaskInput{WorkerID: "w1", Workloads: []string{"w"}}) + if err != nil || claimed == nil { + t.Fatalf("claim: %v", err) + } + cancelled, err := h.cancel.Execute(ctx, jobID) + if err != nil || cancelled != 3 { + t.Fatalf("cancel = (%d, %v), want (3, nil)", cancelled, err) + } + if _, err := h.renew.Execute(ctx, usecase.RenewLeaseInput{TaskID: claimed.TaskID, WorkerID: "w1", Attempt: claimed.Attempt}); !errors.Is(err, domain.ErrTaskNotLeased) { + t.Errorf("cancelled lease heartbeat = %v, want ErrTaskNotLeased", err) + } + progress, err := h.status.Execute(ctx, jobID) + if err != nil || progress.DeriveStatus() != domain.JobCancelled || progress.Cancelled != 3 { + t.Errorf("cancelled progress = %+v, err = %v", progress, err) + } + if cancelled, err := h.cancel.Execute(ctx, jobID); err != nil || cancelled != 0 { + t.Errorf("second cancel = (%d, %v), want (0, nil)", cancelled, err) + } +} + func TestGetTaskInputMissingForURITask(t *testing.T) { h := newHarness() h.seedJob(t, "w", 1) // URI-based task, no coordinator-stored input diff --git a/docs/api-contract.md b/docs/api-contract.md index ba16e84..50abf4d 100644 --- a/docs/api-contract.md +++ b/docs/api-contract.md @@ -52,9 +52,11 @@ Content-Type: multipart/form-data ``` Fields, in order (text fields first, file last — the file is streamed): -`workload`, `parameters` (JSON), `chunk_rows` (int, default 1000), and the file -part `file`. The coordinator stores the input, splits the TSV into shard -artifacts (header repeated per shard), and creates one task per shard. +`workload`, `parameters` (JSON), `chunk_rows` (int, default 1000), optional +`max_rows` (positive int), and the file part `file`. `max_rows` limits the +leading data rows that become shards; it does not change the stored source +artifact. The coordinator splits the selected TSV rows into shard artifacts +(header repeated per shard) and creates one task per shard. `201`: @@ -65,6 +67,24 @@ artifacts (header repeated per shard), and creates one task per shard. Each resulting task's claim response carries `input.uri = /tasks/{id}/input`, served by §5.4. +## Stop a job + +```http +POST /jobs/{job_id}/cancel +Authorization: Bearer +``` + +The coordinator transactionally marks every pending, leased, or running shard +as `cancelled`, invalidates its lease, and marks the job `cancelled`. Completed +and terminally failed shards remain as history. Repeating a cancellation of an +already cancelled job is safe. + +`200`: + +```json +{ "job_id": "uuid", "status": "cancelled", "cancelled_tasks": 12 } +``` + ## Register worker ```http diff --git a/docs/openapi.yaml b/docs/openapi.yaml index 4ae0557..7117a78 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -95,7 +95,7 @@ paths: summary: Upload a dataset; the coordinator chunks it into shard tasks description: > multipart/form-data. The text fields (`workload`, `parameters`, - `chunk_rows`) MUST precede the `file` part: the file is streamed, not + `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. requestBody: required: true @@ -130,6 +130,23 @@ paths: "401": { $ref: "#/components/responses/Unauthorized" } "404": { $ref: "#/components/responses/NotFound" } + /jobs/{job_id}/cancel: + post: + tags: [jobs] + summary: Cancel a job and invalidate all unfinished task leases + parameters: + - $ref: "#/components/parameters/JobID" + responses: + "200": + description: The job is cancelled. Completed and terminally failed tasks remain unchanged. + content: + application/json: + schema: { $ref: "#/components/schemas/CancelJobResponse" } + "400": { $ref: "#/components/responses/BadRequest" } + "401": { $ref: "#/components/responses/Unauthorized" } + "404": { $ref: "#/components/responses/NotFound" } + "409": { $ref: "#/components/responses/Conflict" } + /tasks/claim: post: tags: [tasks] @@ -429,6 +446,11 @@ components: type: integer description: Data rows per shard. Default 1000. example: 1000 + max_rows: + type: integer + minimum: 1 + description: Optional leading data-row limit for a small pipeline check. + example: 500 file: type: string format: binary @@ -441,6 +463,13 @@ components: task_count: { type: integer, example: 3 } input_artifact_id: { type: string, format: uuid } + CancelJobResponse: + type: object + properties: + job_id: { type: string, format: uuid } + status: { type: string, enum: [cancelled] } + cancelled_tasks: { type: integer } + JobProgress: type: object properties: @@ -451,6 +480,7 @@ components: leased: { type: integer } completed: { type: integer } failed: { type: integer } + cancelled: { type: integer } ClaimRequest: type: object @@ -549,4 +579,4 @@ components: TaskStatus: type: string - enum: [pending, leased, completed, failed, cancelled] + enum: [pending, leased, running, completed, failed, cancelled] diff --git a/docs/web-interface-plan.md b/docs/web-interface-plan.md index 429f0d2..48b29d7 100644 --- a/docs/web-interface-plan.md +++ b/docs/web-interface-plan.md @@ -173,6 +173,9 @@ Rules: - TSV file, required, streamed; show expected columns `chembl_id` and `canonical_smiles`. - `chunk_rows`: integer 1--100000, default 1000. +- `max_rows`: optional positive integer. The coordinator creates shards only + from the first N data rows, so a user can test a large upload without + creating thousands of tasks. It does not truncate the stored source blob. - optional human-readable run name is a later schema/API addition; v1 does not silently store it. - display file name and client-side size only as convenience; server limits and