From 6b67326c3bfcdd1cb960def4d3d6b12003d5e031 Mon Sep 17 00:00:00 2001 From: Emil Date: Mon, 3 Aug 2026 15:36:22 +0300 Subject: [PATCH] Run claim-execute-upload loops concurrently in the worker agent --- coordinator/internal/agent/config.go | 17 ++++ coordinator/internal/agent/configfile.go | 5 ++ coordinator/internal/agent/daemon.go | 54 +++++++++++-- coordinator/internal/agent/daemon_test.go | 79 +++++++++++++++++++ coordinator/internal/agent/setupui/server.go | 5 ++ .../internal/agent/setupui/template.html | 4 +- 6 files changed, 156 insertions(+), 8 deletions(-) diff --git a/coordinator/internal/agent/config.go b/coordinator/internal/agent/config.go index 8b6841f..4b03ab4 100644 --- a/coordinator/internal/agent/config.go +++ b/coordinator/internal/agent/config.go @@ -32,6 +32,10 @@ type Config struct { TaskRunner []string // command + args; defaults to python -m scimesh.worker.task MaxTasks int // 0 = unlimited ExitWhenIdle bool + // Concurrency is how many claim→execute→upload loops run in parallel + // under one worker id: N shards processed concurrently on one machine, + // using the coordinator's own task pipeline as the parallel unit. + Concurrency int } func envList(name string) ([]string, error) { @@ -145,9 +149,22 @@ func LoadConfig() (*Config, error) { TaskRunner: runner, MaxTasks: maxTasks, ExitWhenIdle: os.Getenv("EXIT_WHEN_IDLE") == "1", + Concurrency: envInt("WORKER_CONCURRENCY", 1), }, nil } +func envInt(name string, fallback int) int { + raw := os.Getenv(name) + if raw == "" { + return fallback + } + parsed, err := strconv.Atoi(raw) + if err != nil || parsed < 1 { + return fallback + } + return parsed +} + func durationEnv(name string, fallback time.Duration) (time.Duration, error) { raw := os.Getenv(name) if raw == "" { diff --git a/coordinator/internal/agent/configfile.go b/coordinator/internal/agent/configfile.go index 1637013..2a39b7b 100644 --- a/coordinator/internal/agent/configfile.go +++ b/coordinator/internal/agent/configfile.go @@ -21,6 +21,7 @@ type ConfigFile struct { WorkerName string `json:"worker_name,omitempty"` CPUCount int `json:"cpu_count"` MemoryMB int `json:"memory_mb"` + Concurrency int `json:"concurrency,omitempty"` TaskRunner []string `json:"task_runner,omitempty"` } @@ -103,6 +104,10 @@ func (f *ConfigFile) Config() (*Config, error) { if config.MemoryMB < 0 { config.MemoryMB = 0 } + config.Concurrency = f.Concurrency + if config.Concurrency < 1 { + config.Concurrency = 1 + } if len(f.TaskRunner) > 0 { config.TaskRunner = f.TaskRunner } diff --git a/coordinator/internal/agent/daemon.go b/coordinator/internal/agent/daemon.go index 1143fbf..bb6dd1a 100644 --- a/coordinator/internal/agent/daemon.go +++ b/coordinator/internal/agent/daemon.go @@ -37,15 +37,52 @@ func NewDaemon(config *Config, client *Client, runner *TaskRunner, log *slog.Log return &Daemon{config: config, client: client, runner: runner, log: log} } -// RunForever loops until interrupted, idle-exit, or max tasks. +// RunForever registers once, then runs the claim→execute→upload loop +// concurrently under one worker id. With Concurrency > 1, several shards are +// processed in parallel on this machine, using the coordinator's own task +// pipeline as the parallel unit. func (d *Daemon) RunForever() error { + if !d.registered { + if err := d.register(); err != nil { + return err + } + } + workers := d.config.Concurrency + if workers < 1 { + workers = 1 + } + if workers == 1 { + return d.loop() + } + d.log.Info("agent running concurrently", "loops", workers) + var wg sync.WaitGroup + errs := make(chan error, workers) + for i := 0; i < workers; i++ { + wg.Add(1) + go func(loop int) { + defer wg.Done() + if err := d.loop(); err != nil { + errs <- err + return + } + errs <- nil + }(i) + } + wg.Wait() + close(errs) + for err := range errs { + if err != nil { + return err + } + } + return nil +} + +// loop is one claim→execute→upload cycle until interrupted, idle-exit, or +// the shared max-tasks budget is consumed. +func (d *Daemon) loop() error { failures := 0 for { - if !d.registered { - if err := d.register(); err != nil { - return err - } - } d.cleanupExpiredDirectories() outcome, err := d.runOnce() if err != nil { @@ -63,8 +100,11 @@ func (d *Daemon) RunForever() error { } failures = 0 if outcome.Claimed && outcome.Completed { + d.mu.Lock() d.completed++ - if d.config.MaxTasks > 0 && d.completed >= d.config.MaxTasks { + done := d.config.MaxTasks > 0 && d.completed >= d.config.MaxTasks + d.mu.Unlock() + if done { d.log.Info("max tasks reached") return nil } diff --git a/coordinator/internal/agent/daemon_test.go b/coordinator/internal/agent/daemon_test.go index ef0a252..fbab249 100644 --- a/coordinator/internal/agent/daemon_test.go +++ b/coordinator/internal/agent/daemon_test.go @@ -11,6 +11,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "time" ) @@ -55,6 +56,7 @@ type fakeCoordinator struct { uploadSize int64 inputBytes []byte conflict bool // 409 on heartbeat/upload/result + mu sync.Mutex } func newFakeCoordinator(t *testing.T, task map[string]any) *fakeCoordinator { @@ -64,6 +66,8 @@ func newFakeCoordinator(t *testing.T, task map[string]any) *fakeCoordinator { fake.uploadSize = int64(len(fake.inputBytes)) var server *httptest.Server server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + fake.mu.Lock() + defer fake.mu.Unlock() switch { case r.Method == http.MethodPost && r.URL.Path == "/workers/register": writeJSON(w, http.StatusCreated, map[string]any{ @@ -294,3 +298,78 @@ func TestDaemonIdleClaimIsNotCompleted(t *testing.T) { t.Fatalf("outcome = %+v", outcome) } } + +func TestDaemonConcurrencyProcessesTasksInParallel(t *testing.T) { + t.Parallel() + marker := filepath.Join(t.TempDir(), "marker") + script := filepath.Join(t.TempDir(), "fake-runner.sh") + content := `#!/bin/sh +out="" +task_dir="" +while [ "$#" -gt 0 ]; do + case "$1" in + --output) out="$2"; shift 2;; + --task-dir) task_dir="$2"; shift 2;; + *) shift;; + esac +done +echo start >> ` + marker + ` +sleep 1 +echo end >> ` + marker + ` +printf 'id,score\n1,1\n' > "$task_dir/result.csv" +printf '{"artifact_path":"%s/result.csv","content_type":"text/csv","metrics":{"rows":1}}' "$task_dir" > "$out" +exit 0 +` + if err := os.WriteFile(script, []byte(content), 0o755); err != nil { + t.Fatal(err) + } + fake := newFakeCoordinator(t, validClaimedTaskPayload()) + defer fake.close() + config := &Config{ + CoordinatorURL: fake.server.URL, + WorkerName: "concurrent-worker", + WorkerID: "22222222-2222-4222-8222-222222222222", + WorkDir: t.TempDir(), + CPUCount: 1, + PollInterval: time.Millisecond, + RequestTimeout: 5 * time.Second, + Heartbeat: 15 * time.Second, + Capabilities: []string{"similarity-search"}, + TaskRunner: []string{script}, + MaxTasks: 3, + Concurrency: 3, + } + client := NewClient(fake.server.URL, &StaticToken{token: "test-token"}, 5*time.Second) + logger := slog.New(slog.NewTextHandler(io.Discard, nil)) + daemon := NewDaemon(config, client, NewTaskRunner(config.TaskRunner), logger) + if err := daemon.RunForever(); err != nil { + t.Fatalf("run: %v", err) + } + raw, err := os.ReadFile(marker) + if err != nil { + t.Fatalf("marker: %v", err) + } + starts := strings.Count(string(raw), "start\n") + ends := strings.Count(string(raw), "end\n") + if starts < 3 || ends < 3 { + t.Fatalf("marker: %d starts / %d ends, want at least 3/3", starts, ends) + } + // With three loops sleeping 1s each, the marker proves all three ran + // concurrently (three starts before the first end completes a 1s sleep). + lines := strings.Split(strings.TrimSpace(string(raw)), "\n") + concurrent := 0 + running := 0 + for _, line := range lines { + if line == "start" { + running++ + if running > concurrent { + concurrent = running + } + } else { + running-- + } + } + if concurrent < 3 { + t.Errorf("max concurrent executions = %d, want 3", concurrent) + } +} diff --git a/coordinator/internal/agent/setupui/server.go b/coordinator/internal/agent/setupui/server.go index 0077606..8aff709 100644 --- a/coordinator/internal/agent/setupui/server.go +++ b/coordinator/internal/agent/setupui/server.go @@ -377,6 +377,7 @@ type saveConfigRequest struct { WorkerName string `json:"worker_name"` CPUCount int `json:"cpu_count"` MemoryMB int `json:"memory_mb"` + Concurrency int `json:"concurrency"` TaskRunner []string `json:"task_runner"` } @@ -395,6 +396,7 @@ func (s *Server) handleSaveConfig(w http.ResponseWriter, r *http.Request) { WorkerName: strings.TrimSpace(req.WorkerName), CPUCount: req.CPUCount, MemoryMB: req.MemoryMB, + Concurrency: req.Concurrency, TaskRunner: req.TaskRunner, } if file.CoordinatorURL == "" { @@ -418,6 +420,9 @@ func (s *Server) handleSaveConfig(w http.ResponseWriter, r *http.Request) { if file.CPUCount < 1 { file.CPUCount = 1 } + if file.Concurrency < 1 { + file.Concurrency = 1 + } // The wizard UI bakes the venv python into the runner after an install; // an API-driven or scripted flow may not, so the server guarantees it: // workloads execute through scimesh's task runner, which lives in the venv. diff --git a/coordinator/internal/agent/setupui/template.html b/coordinator/internal/agent/setupui/template.html index dd91808..3cc4fae 100644 --- a/coordinator/internal/agent/setupui/template.html +++ b/coordinator/internal/agent/setupui/template.html @@ -141,6 +141,7 @@ code{font-family:var(--mono);font-size:.86em} +

Process this many shards in parallel on this machine. Each loop runs its own task runner subprocess.

@@ -228,7 +229,8 @@ function draftConfig(){ userservice_url:state.mode==='key'?$('in-users').value.trim():'', work_dir:$('in-dir').value.trim(), worker_name:$('in-name').value.trim(), - cpu_count:state.cpu==='custom'?parseInt($('in-cpu').value||'1',10):0 + cpu_count:state.cpu==='custom'?parseInt($('in-cpu').value||'1',10):0, + concurrency:parseInt($('in-conc').value||'1',10) }; if(state.venvPython)cfg.task_runner=[state.venvPython,'-m','scimesh.worker.task']; return cfg;